Building Large Language Models: From Pre-training to Post-training

Introduction

Large language models (LLMs) have become a cornerstone of modern artificial intelligence, powering chatbots and AI assistants like ChatGPT, Claude, and Gemini. This article provides a comprehensive overview of the process of building LLMs, from pre-training on massive datasets to post-training alignment and optimization.

Pre-training

The Task of Language Modeling

At its core, pre-training an LLM involves teaching the model to predict the next word in a sequence given the previous words. This task is known as language modeling. Mathematically, we can express this as modeling the probability distribution P(X1, X2, ..., XL) where Xi represents the ith word in a sequence.

For example, given the partial sentence "The mouse ate the", the model should assign a high probability to "cheese" as the next word, a lower probability to grammatically incorrect continuations, and an even lower probability to semantically nonsensical completions like "The cheese ate the mouse".

Autoregressive Language Models

Most modern LLMs use an autoregressive approach, decomposing the joint probability of a sequence into a product of conditional probabilities:

P(X1, X2, ..., XL) = P(X1) * P(X2|X1) * P(X3|X1,X2) * ... * P(XL|X1,...,XL-1)

This allows the model to generate text one token at a time, conditioning each new token on all previous tokens.

Tokenization

Before training, text must be converted into a format the model can process. This is done through tokenization, which breaks text into smaller units called tokens. Tokens are typically subword units, striking a balance between character-level and word-level representations.

Popular tokenization algorithms like Byte Pair Encoding (BPE) work by:

1. Starting with individual characters as tokens
2. Iteratively merging the most frequent adjacent token pairs
3. Stopping when a desired vocabulary size is reached

This results in common words or subwords having their own tokens, while rarer words are split into multiple tokens.

Model Architecture

While this article won't delve deeply into model architectures, it's worth noting that most modern LLMs are based on the Transformer architecture. Key components include:

• Token embeddings
• Self-attention layers
• Feed-forward neural networks
• Layer normalization

Training Objective

The primary training objective for autoregressive language models is to minimize the cross-entropy loss between the model's predicted probability distribution for the next token and the actual next token in the training data.

Evaluation: Perplexity

During pre-training, models are typically evaluated using perplexity, which is related to the average cross-entropy loss:

Perplexity = 2^(average cross-entropy loss)

Lower perplexity indicates better performance, with a perfect model achieving a perplexity of 1. In recent years, state-of-the-art models have reduced perplexity on standard datasets from around 70 to less than 10.

Data for Pre-training

Pre-training data typically comes from web crawls of the internet, supplemented with high-quality sources like books and academic papers. Preparing this data involves several steps:

1. Web crawling (e.g., using Common Crawl)
2. Content extraction from HTML
3. Filtering undesirable content (e.g., explicit material, personal information)
4. Deduplication
5. Quality filtering
6. Domain classification and weighting

The scale of pre-training data has grown dramatically, with current top models training on 15+ trillion tokens.

Scaling Laws

A key insight in LLM development is the existence of predictable scaling laws. As models grow larger and are trained on more data, their performance improves in a log-linear fashion. This allows researchers to extrapolate performance and make informed decisions about resource allocation.

Scaling laws help answer questions like:

• How large should the model be?
• How much data should be used?
• What's the optimal ratio of model size to dataset size?

For example, the "Chinchilla" scaling laws suggest using about 20 tokens of training data per model parameter for optimal performance.

Compute Requirements

Training large language models requires enormous computational resources. For example, training a state-of-the-art open-source model like Llama 3 (400B parameters) involves:

• ~3.8e25 floating point operations (FLOPs)
• 16,000 H100 GPUs running for ~70 days
• An estimated cost of $75 million (including hardware and personnel)

Post-training Alignment

While pre-training gives models broad knowledge and capabilities, additional steps are needed to turn them into helpful and safe AI assistants.

Supervised Fine-tuning (SFT)

The first step in alignment is supervised fine-tuning (SFT). This involves further training the model on a smaller dataset of high-quality human-written responses to prompts. SFT helps the model learn the desired format and style for responses.

Key points about SFT:

• Uses the same loss function as pre-training (next token prediction)
• Typically only needs 2,000-50,000 examples
• Can be augmented with synthetic data generated by other language models

Reinforcement Learning from Human Feedback (RLHF)

The next step is reinforcement learning from human feedback (RLHF). This process aims to align the model's outputs with human preferences. The basic RLHF pipeline involves:

1. Collecting human preferences between model outputs
2. Training a reward model to predict human preferences
3. Using reinforcement learning to optimize the language model against the reward model

Two main approaches to RLHF are:

1. Proximal Policy Optimization (PPO): A reinforcement learning algorithm that directly optimizes the model as a policy.
2. Direct Preference Optimization (DPO): A simpler approach that frames the problem as maximum likelihood training, avoiding some of the complexities of RL.

RLHF helps models learn to:

• Follow instructions more precisely
• Avoid generating harmful or biased content
• Improve the overall quality and helpfulness of responses

Challenges in Alignment

Some key challenges in the alignment process include:

• Expense and time required for human feedback
• Potential biases in human annotators
• Difficulty in specifying complex human values and preferences
• Unintended consequences (e.g., models becoming overly verbose)

Evaluation of Aligned Models

Evaluating aligned language models presents unique challenges:

• Traditional metrics like perplexity are no longer applicable
• Responses are open-ended and subjective
• Models need to be evaluated on a wide range of tasks and capabilities

Some approaches to evaluation include:

• Human evaluation (e.g., ChatBot Arena)
• Automated evaluation using other language models as judges
• Specialized benchmarks for specific capabilities

System Optimizations

Given the enormous computational requirements of LLMs, system-level optimizations are crucial. Some key areas of focus include:

GPU Utilization

Modern LLM training relies heavily on GPUs. Maximizing GPU utilization involves:

• Efficient data loading and preprocessing
• Optimizing memory usage
• Minimizing communication overhead between GPUs

Low Precision Training

Using lower precision number formats (e.g., 16-bit floats instead of 32-bit) can significantly speed up training and reduce memory usage with minimal impact on model quality.

Operator Fusion

Combining multiple operations into a single GPU kernel can reduce memory bandwidth requirements and improve performance. Tools like PyTorch's torch.compile() can automatically apply these optimizations.

Distributed Training

Efficiently scaling training across multiple GPUs and multiple machines requires careful consideration of:

• Data parallelism vs. model parallelism
• Communication protocols
• Load balancing

Conclusion

Building large language models is a complex process involving massive datasets, enormous computational resources, and intricate training procedures. From the initial pre-training on internet-scale data to the final alignment steps ensuring safety and usefulness, each stage presents unique challenges and opportunities for innovation.

As the field continues to advance, we can expect to see:

• Even larger models and datasets
• More efficient training techniques
• Improved alignment methods
• Novel applications across various domains

Understanding the full pipeline of LLM development is crucial for researchers, engineers, and policymakers as these models become increasingly central to our technological landscape.

Further Reading

For those interested in diving deeper into LLM development, consider exploring these Stanford courses:

• CS224N: Natural Language Processing with Deep Learning
• CS324: Large Language Models
• CS336: Building and Understanding Large Language Models

These courses offer in-depth coverage of the topics discussed in this article, as well as hands-on experience with building and optimizing language models.

Article created from: https://www.youtube.com/watch?v=9vM4p9NN0Ts