Dissecting Llama 3.1: A Deep Dive

I. Introduction

The https://ai.meta.com/research/publications/the-llama-3-herd-of-models/ paper introduces Llama 3, a new family of language models that natively supports multilingualism, coding, reasoning, and tool usage.

This paper discusses the design choices that went into developing Llama 3, focusing on the levers of data, scale, and managing complexity. The authors emphasize the importance of high-quality data and large-scale training for achieving strong performance.

The paper also details the pre-training process, including data curation, model architecture, and scaling laws. The authors also describe the post-training process, which involves supervised fine-tuning (SFT), rejection sampling, and direct preference optimization (DPO).

Furthermore, the paper explores the integration of multimodal capabilities (images, video, and speech) into Llama 3, highlighting the development of separate encoders for each modality and the use of adapters to integrate these encoders into the language model.

Finally, the paper discusses the safety considerations for Llama 3, including the construction of safety benchmarks, the application of safety finetuning, and the development of a system-level safety classifier called Llama Guard.

Overall, this paper provides a detailed overview of the design, development, and evaluation of Llama 3, a powerful new family of language models.

II. Data & Preprocessing

Llama 3.1 pre-training leverages a diverse range of data sources, including:

• Web data: The model was trained on a vast corpus of text scraped from the internet, spanning a variety of domains and languages. This web data was carefully curated and filtered to ensure quality and safety.
• Code data: To enhance coding capabilities, the training dataset includes a significant amount of high-quality code from a variety of programming languages, such as Python, Java, Javascript, C/C++, Typescript, Rust, PHP, HTML/CSS, and SQL.
• Multilingual data: Llama 3.1 supports multiple languages, including English, German, French, Italian, Portuguese, Hindi, Spanish, and Thai. This was achieved by including multilingual text data in the training corpus.

To ensure data quality and safety, the authors applied various filtering and cleaning methods, including:

• PII and safety filtering: The training dataset was scrubbed for personally identifiable information (PII) and content that could be considered harmful, such as adult content.
• De-duplication: Duplicate and near-duplicate content was removed from the dataset to improve training efficiency and reduce the potential for bias. This was achieved through multiple levels of de-duplication: URL, document, and line-level.
• Heuristic filtering: Additional heuristics were applied to remove low-quality documents, such as those with excessive repetitions or those containing repetitive content like logging or error messages.
• Model-based quality filtering: Finally, the authors experimented with using various model-based quality classifiers to further refine the training data. These classifiers were trained to recognize high-quality text and were used to identify and remove low-quality content from the training corpus.

The authors also carefully considered the data mix and annealing strategy used for pre-training:

• Data Mix: To achieve the desired balance of capabilities, the authors carefully determined the proportion of different data sources in the training mix. The final data mix consisted of roughly 50% general knowledge, 25% mathematical and reasoning tokens, 17% code tokens, and 8% multilingual tokens.
• Annealing Strategy: Annealing was employed to further improve performance on key benchmarks. This involved gradually reducing the learning rate while up-sampling high-quality data from specific domains during the final stage of pre-training.

These data curation and processing strategies are crucial for ensuring the quality, safety, and effectiveness of the Llama 3.1 pre-training process.

III. Architecture & Training

Llama 3.1 is based on a standard dense Transformer model architecture, with minor adaptations for improved efficiency and scalability. Here's a breakdown of its key components and training process:

Architecture:

• Layers: Llama 3.1 models are composed of a varying number of transformer layers, depending on the model size. The 8B model has 32 layers, the 70B model has 80 layers, and the 405B model has 126 layers.
• Model Dimension: The model dimension, or the number of hidden units in each layer, also scales with model size. It is 4,096 for the 8B model, 8,192 for the 70B model, and 16,384 for the 405B model.
• Attention Heads: The number of attention heads in each layer also increases with model size. The 8B model uses 32 heads, the 70B model uses 64 heads, and the 405B model uses 128 heads.
• Key/Value Heads: The number of key/value heads, which control the attention mechanism, is kept constant across all models at 8.
• Peak Learning Rate: The peak learning rate for each model varies based on model size and is empirically determined. It is 3 × 10-4 for the 8B model, 1.5 × 10-4 for the 70B model, and 8 × 10-5 for the 405B model.
• Activation Function: The activation function used in Llama 3.1 is SwiGLU.
• Vocabulary Size: All Llama 3.1 models use a vocabulary of 128,000 tokens, combining 100,000 tokens from the tiktoken tokenizer with 28,000 additional tokens for better support of non-English languages.
• Positional Embeddings: The model uses the Rotary Positional Embedding (ROPE) method to encode positional information, with a base frequency of 500,000.

Training:

Scaling Laws: The authors conducted extensive scaling law experiments to determine the optimal model size and predict downstream performance based on training FLOPs. The results indicate that the 405B model is approximately compute-optimal for the given training budget.

The authors utilized scaling laws to determine the optimal model size for Llama 3.1 and to predict its performance on downstream tasks, given a specific training compute budget. Scaling laws, however, are often noisy and unreliable, especially when applied to small compute budgets.

To address these challenges, the authors implemented a two-stage methodology:

1. Correlating Compute and Loss: They established a correlation between the compute-optimal model's negative log-likelihood on downstream tasks and the training FLOPs.
2. Correlating Loss and Accuracy: They then correlated the negative log-likelihood on downstream tasks with task accuracy, leveraging data from both scaling law models and older models trained with higher compute FLOPs.

This two-stage methodology enabled the authors to predict downstream performance for compute-optimal models with reasonable accuracy, considering a wide range of compute budgets.

• Figure 2: This figure showcases the IsoFLOPs curves generated during scaling law experiments, demonstrating the relationship between compute budget and negative log-likelihood on a held-out validation set. The IsoFLOPs curves reveal a clear minimum point representing the compute-optimal model for each specific compute budget.
• Figure 3: This figure depicts the relationship between the training compute budget and the number of training tokens for the identified compute-optimal models. The authors used a power-law relationship to extrapolate this data and predict the optimal number of training tokens for the given compute budget (3.8 × 1025 FLOPs).

The findings of these experiments suggested that the performance of the flagship 405B parameter model was relatively robust to small changes in the trade-off between model size and training tokens.

4D Parallelism: To enable efficient training at scale, the authors implemented a 4D parallelism strategy combining tensor parallelism, pipeline parallelism, context parallelism, and data parallelism. This allowed them to efficiently distribute computation across 16,384 GPUs.

Let's break down each parallelism type in more detail:

• Tensor Parallelism (TP): This involves splitting individual weight tensors across multiple GPUs. This allows for parallel computation of the matrix multiplications in each layer, enabling the use of larger models with more parameters.
• Pipeline Parallelism (PP): This partitions the model vertically into stages, where each stage consists of multiple layers. Different GPUs process different stages of the model pipeline, enabling parallel processing of the entire model.
• Context Parallelism (CP): This technique divides the input context into segments, reducing memory bottlenecks for very long sequences. This is particularly useful for models trained on large documents or code repositories. The authors implemented a novel all-gather-based context parallelism, which allows for efficient computation of attention output for the local query tensor chunk.
• Data Parallelism (DP): This involves distributing the training data across multiple GPUs. The authors employed fully sharded data parallelism (FSDP), where model parameters, optimizer states, and gradients are sharded across GPUs, enabling efficient parallel processing of large datasets.

The 4D parallelism strategy introduced several challenges, including:

• Batch Size Constraint: Traditional implementations impose limitations on batch size per GPU, restricting the flexibility of model training. The authors addressed this by modifying their pipeline schedule to allow for a flexible number of micro-batches, enabling them to optimize batch size for specific training scenarios.
• Memory Imbalance: The different stages of the pipeline can consume varying amounts of memory, leading to inefficient resource allocation. The authors addressed this by employing an interleaved schedule and reducing the number of layers in the first and last stages, minimizing memory imbalances.
• Computation Imbalance: Certain stages of the pipeline, such as the last layer, can experience higher execution latency, leading to pipeline bubbles. The authors addressed this by incorporating asynchronous point-to-point communication and proactively deallocating tensors that are no longer needed for future computation.

The authors' careful design and optimization of the 4D parallelism strategy, coupled with their detailed understanding of the network topology, collective communication libraries, and model-specific requirements, enabled them to train the 405B parameter model efficiently and achieve remarkable results.

Training Recipe

The authors employed a multi-stage training recipe to achieve strong performance across various capabilities:

Training Llama 3.1 involved a multi-stage process to ensure strong performance across various capabilities. The authors employed a careful combination of optimization techniques, learning rate schedules, and data selection strategies to achieve optimal results:

Initial Pre-training:

The initial pre-training stage for Llama 3.1 involved training the model on a massive corpus of text tokens using a standard next-token prediction objective. The authors employed a combination of techniques to ensure efficient and stable training:

• Optimizer: AdamW, a popular optimizer for large language models, was used to update model parameters.
• Learning Rate: A cosine learning rate schedule was used, with a peak learning rate of 8 × 10-5 and a linear warmup phase of 8,000 steps. This learning rate gradually decayed to 8 × 10-7 over 1,200,000 steps.
• Batch Size: Initially, the batch size was set to 4M tokens with sequences of length 4,096. This was gradually increased to 8M tokens with sequences of length 8,192 after training on 252M tokens. Finally, the batch size was doubled again to 16M after pre-training on 2.87T tokens. This gradual increase in batch size ensured efficient and stable training.

The authors also carefully adjusted the data mix during this stage to improve performance on specific tasks. They increased the percentage of non-English data to enhance the model's multilingual capabilities. They up-sampled mathematical data to improve performance on mathematical reasoning tasks. They added more recent web data to advance the model's knowledge cut-off. They also down-sampled subsets of data that were later identified as being of lower quality.

Long Context Pre-training:

To enable the processing of long documents and complex reasoning tasks, the final stages of pre-training involved transitioning the model to longer sequences, up to 128K tokens. The authors gradually increased the context window length in increments, ensuring the model adapted successfully to each new length before proceeding. The success of the adaptation was evaluated by ensuring that:

• Performance on short-context tasks remained at an acceptable level.
• The model could successfully solve "needle in a haystack" tasks, demonstrating its ability to retrieve specific information from longer sequences.

This long-context pre-training stage involved training on approximately 800B training tokens.

These careful adjustments to the training strategy, data mix, and pre-training stages played a crucial role in the development of Llama 3.1, a high-performance and scalable language model.

This overview of Llama 3.1's architecture and training process highlights the key decisions and strategies employed by the authors to create a high-performance and scalable language model.

IV. Post-training

To align Llama 3.1 with human preferences and further enhance its capabilities, the authors employed a multi-round post-training approach, building on top of the pre-trained checkpoints.

This process involves three key steps:

1. Rejection Sampling:

• The authors leveraged rejection sampling to create more on-policy negative samples for training. This involves sampling multiple outputs from the model for a given prompt, using a reward model to select the best candidate, and using the remaining outputs as negative samples.
• To improve efficiency, the authors adopted Paged Attention, a technique that enhances memory efficiency through dynamic key-value cache allocation.
• Rejection sampling plays a crucial role in improving the model's ability to reason, understand complex instructions, and generate more helpful and engaging outputs.

2. Supervised Fine-tuning (SFT):

• The pre-trained model was further fine-tuned on a large dataset of human-annotated examples and synthetic data, aiming to improve its performance on specific tasks.
• The SFT data was curated from multiple sources:

• The authors carefully balanced the data mix to optimize performance across a wide range of capabilities and target specific areas where the model lagged behind.
• SFT with high-quality data is a critical step in aligning the model with human expectations and improving its overall performance.

3. Direct Preference Optimization (DPO):

• To further refine the model's alignment with human preferences, the authors employed DPO, a technique that directly optimizes the model's parameters based on human preference data.
• DPO training used recent batches of preference data collected during the previous rounds, ensuring the training data closely matched the model's current behavior.
• To stabilize DPO training and prevent undesired model behaviors, the authors introduced several algorithmic modifications:
• DPO significantly improves the model's ability to follow instructions, generate factually accurate outputs, and demonstrate overall helpfulness.

Improving Specific Capabilities

The authors invested significant effort in enhancing Llama 3.1's performance on specific capabilities:

Code:

• Expert Training: To improve code generation, documentation, debugging, and review capabilities, the authors trained a dedicated "code expert" model. This involved branching off the main pre-training run and continuing pre-training on a dataset primarily consisting of code data. This domain-specific pre-training has been shown to be effective for improving performance within a particular domain. The authors also performed long-context finetuning on a high-quality mix of repo-level code data to further enhance the model's capabilities.
• Synthetic Data Generation: The authors identified key challenges in code generation, such as difficulty following instructions, code syntax errors, incorrect code generation, and difficulty fixing bugs. To address these challenges, they generated a large amount of synthetic data for SFT, using three main approaches:
• Prompt Steering: To improve code formatting, the authors implemented prompt steering techniques, using system prompts to guide the model's output.
• Quality Filtering: The authors implemented quality filters to remove bad samples from their training data. This involved filtering out code samples that exhibited incorrect syntax, code style issues, or those that failed to pass unit tests.

Multilinguality:

• Multilingual Data Sourcing: To enhance the model's capabilities across multiple languages, the authors sourced high-quality multilingual data for German, French, Italian, Portuguese, Hindi, Spanish, and Thai.
• Multilingual Expert Training: They further trained a dedicated "multilingual expert" model by branching off the main pre-training run and continuing pre-training on a dataset primarily consisting of multilingual tokens.
• Language Steering: The authors addressed challenges related to language steering, ensuring consistent performance across various languages. This involved identifying and mitigating biases related to translationese, name bias, gender bias, and cultural bias. They also translated synthetic quantitative reasoning data to improve performance in non-English languages.

Math and Reasoning:

• Prompt Creation: The authors addressed challenges related to the lack of prompts and ground truth chains of thought for mathematical reasoning by actively sourcing prompts from humans and developing a taxonomy of mathematical skills.
• Step-by-step Solution Generation: They generated step-by-step solutions for training data, using the model to produce multiple solutions and filtering them based on correctness.
• Reward Model Training: They trained outcome and step-wise reward models to filter out data with incorrect intermediate reasoning steps, ensuring high-quality training data.
• Interleaving Code and Text Reasoning: They prompted the model to solve reasoning problems through a combination of textual reasoning and associated Python code, using code execution as a feedback signal to eliminate cases where the reasoning chain was not valid.
• Learning from Feedback: They simulated human feedback by prompting the model to generate correct solutions based on incorrect reasoning traces, helping the model learn from its mistakes.

Long Context:

• The authors extended the context window to 128K tokens, enabling the processing of long documents and complex reasoning tasks. This involved gradually increasing the context window length and ensuring the model adapted successfully to each new length before proceeding.
• They leveraged hierarchical summarization and question answering on long documents, prompting the model to summarize chunks of 8K tokens and then summarizing those summaries.
• They generated synthetic data for long-context code reasoning, prompting the model to identify missing code dependencies and generate the necessary code.

Tool Use:

• The authors trained the model to use tools such as search engines and code interpreters, creating datasets that encompass multi-step tool use scenarios.
• They designed prompts that encouraged tool use only when necessary, as well as prompted the model to call tools in sequence, reason about the outputs, and utilize zero-shot tool use capabilities.

Steerability:

• The authors introduced techniques to improve the model's ability to follow instructions, including careful system prompt design and selection of preference data.

These targeted efforts to improve Llama 3.1's capabilities across various domains demonstrate the authors' commitment to developing a versatile and powerful language model that can excel in a wide range of tasks.

V. Safety & Reliability

The authors of Llama 3.1 emphasize the importance of developing a safe and responsible AI system, focusing on mitigating potential risks while maximizing helpfulness. Their approach to safety encompasses various stages:

1. Safety Benchmarks:

• The authors created a comprehensive set of internal benchmarks to assess the model's safety across various capabilities. These benchmarks were heavily inspired by the ML Commons taxonomy of hazards and included a wide range of adversarial and borderline prompts.
• Adversarial prompts were designed to elicit harmful responses, while borderline prompts tested the model's ability to provide safe and helpful responses even when presented with challenging or potentially ambiguous requests.
• The benchmark encompasses various capabilities, such as English text generation, multilingual text generation, long-context document question answering, tool use (search), and vision & speech capabilities.

2. Safety Pre-training:

• The authors applied various filtering techniques during pre-training to minimize the potential for harmful content and reduce the risk of memorization.
• These techniques included:

3. Safety Finetuning:

• The authors introduced a dedicated safety finetuning stage, building on top of the general fine-tuning process, to further improve the model's ability to adhere to safety policies.
• Two primary metrics were used to evaluate the model's safety performance:
• To mitigate risks effectively, the authors focused on:
• The authors discovered that model size plays a significant role in safety performance, with larger models generally requiring a lower proportion of safety data relative to helpfulness data.

4. System Level Safety:

• To provide more flexibility and control for developers, the authors developed Llama Guard, a separate classifier trained to detect violations of safety policies on input prompts and output responses.
• Llama Guard can be used to filter out harmful content, either before or after model generation, and can be customized for specific harm categories.
• The authors also introduced two prompt-based system guards, Prompt Guard and Code Shield, designed to detect and mitigate prompt attacks and code generation vulnerabilities, respectively.

5. Safety Results:

• The authors conducted extensive evaluations of Llama 3.1's safety across various capabilities, comparing it to other models and systems.
• Overall, Llama 3.1 demonstrates strong safety performance, achieving significant reductions in violation rates while maintaining a low false refusal rate.
• The authors observed that the model's safety performance varies across languages, with English generally being easier to mitigate than non-English languages.
• They also found that long-context models are more susceptible to safety risks and require targeted mitigations, such as using long-context data in SFT and leveraging additional safety measures for tool use.
• The authors also conducted uplift testing for cybersecurity and chemical/biological weapons risks, demonstrating that Llama 3.1 does not significantly increase the risk of malicious actors leveraging the model for harmful purposes.

6. Red Teaming:

• Red teaming plays a crucial role in continuously discovering new risks and improving safety mitigation strategies.
• The authors have a dedicated red teaming team with expertise in various domains, including cybersecurity, adversarial machine learning, and multilingual content.
• Red teaming efforts focus on discovering and mitigating prompt-level attacks, identifying vulnerabilities in specific model capabilities, and exploring the potential for misuse of tools.

The authors' comprehensive approach to safety, encompassing various stages of development, thorough evaluation, and continuous improvement through red teaming, demonstrates their commitment to building a safe and responsible AI system.

VI. Inference & Efficiency

To enable efficient inference with the large Llama 3.1 405B parameter model, the authors employed two key techniques:

1. Pipeline Parallelism:

• Due to the model's size, the 405B parameters do not fit in the GPU memory of a single machine, even with high-performance GPUs like the Nvidia H100.
• To address this, the authors implemented pipeline parallelism, distributing the model across multiple GPUs on two machines.
• Within each machine, the high NVLink bandwidth enables the use of tensor parallelism, further accelerating inference.
• Across machines, the lower bandwidth and higher latency necessitate the use of pipeline parallelism.
• Micro-batching was employed to improve inference throughput while using pipeline parallelism.
• Micro-batching allows for concurrent execution of smaller batches within each stage of the pipeline, resulting in significant performance improvements.

2. FP8 Quantization:

• To further boost inference efficiency, the authors leveraged the FP8 quantization capabilities of the Nvidia H100 GPUs.
• This involved quantizing most parameters and activations in the feedforward network layers of the model, reducing the overall computational cost.
• To ensure accuracy and mitigate quantization errors, the authors implemented several strategies, including:
• Experimental evaluations demonstrate that FP8 quantization achieves significant throughput improvements (up to 50% in the pre-fill stage) while maintaining a negligible impact on model performance.

These optimizations significantly improve the efficiency of Llama 3.1 inference, making it possible to deploy and leverage this powerful language model for a wide range of applications.

VII. Vision & Speech Integration

Llama 3.1 goes beyond traditional text-based language modeling by integrating vision and speech capabilities through a compositional approach. This approach leverages separate encoders for each modality and uses adapters to integrate them into the language model.

Vision

Data: The image and video encoders were trained on a large dataset of image-text pairs and video-text pairs, respectively.

Architecture: The vision component consists of three main parts:

• Image Encoder: This is based on the Vision Transformer (ViT) architecture, pre-trained on a large dataset of image-text pairs. It is trained to align images and text, and uses the ViT-H/14 variant. It has 630M parameters and was trained on 2.5B image-text pairs for five epochs. The image encoder processes images with a resolution of 224 × 224, dividing them into 16 × 16 patches of equal size (i.e., a patch size of 14x14 pixels).
• Image Adapter: This module introduces cross-attention layers between the image encoder and the language model. This allows the model to process visual information.
• Video Adapter: This module is responsible for processing temporal information from videos, using a combination of temporal aggregators and video cross-attention layers. It merges 32 consecutive frames into one, and introduces additional video cross-attention layers.

Speech:

Data: The training data for the speech component can be categorized into two types:

• Pre-training Data: The speech encoder was pre-trained on a massive dataset of unlabeled speech, spanning a variety of languages. This unlabeled data, processed in a self-supervised manner, helps the model learn general acoustic and linguistic representations.
• Fine-tuning Data: Supervised fine-tuning data for speech understanding was sourced from speech recognition, speech translation, and spoken dialogue tasks. This labeled data enables the model to acquire specific speech understanding abilities, further enhancing its performance.

Architecture: The speech module consists of two components:

• Speech Encoder: This is a Conformer model, pre-trained on unlabeled speech data. It takes as input 80-dimensional mel-spectrogram features and processes them using 24 Conformer layers, each with a latent dimension of 1536. This encoder leverages a convolution module with a kernel size of 7 and a rotary attention module, ultimately yielding a token representation of speech signals. The Speech encoder boasts 1B parameters.
• Speech Adapter: This module maps the speech encoder's output to a dimension compatible with the language model embeddings, enabling direct interaction between speech and text. It consists of a convolution layer, a rotary transformer layer, and a linear layer.

Training: The training process for the speech module included two stages:

• Speech Encoder Pre-training: The authors utilized the self-supervised BEST-RQ algorithm to pre-train the speech encoder, leveraging unlabeled speech data.
• Speech Adapter Supervised Fine-tuning: The speech encoder and adapter were jointly trained with the language model on speech recognition, speech translation, and spoken dialogue data. This supervised fine-tuning further refines the model's capabilities and enables it to respond to speech input more effectively.

Speech Generation: Llama 3.1 also incorporates speech generation capabilities, leveraging the Llama 3 embeddings for text normalization and prosody modeling, which enhance the naturalness and expressiveness of generated speech.

Overall: The authors' compositional approach for integrating vision and speech into Llama 3.1 demonstrates the flexibility and scalability of language models, allowing for the development of powerful new capabilities without sacrificing existing text-based performance. This approach lays the foundation for future research in multi-modal language modeling and opens up exciting possibilities for developing more versatile and intelligent AI systems.

VIII. Conclusion

The development of Llama 3.1 suggests that high-quality foundation models are still in their infancy, with significant room for improvement. This paper highlights the crucial roles of high-quality data, scale, and simplicity in achieving strong model results. The authors' focus on these aspects, along with the consistent application of best practices, has resulted in a powerful model family that exhibits strong performance across a wide range of capabilities.

The key implementation details discussed in this outline demonstrate the authors' commitment to:

• Leveraging high-quality data: The use of carefully curated web data, code data, and multilingual data contributes significantly to model performance.
• Scaling training to massive scale: The authors utilized 4D parallelism and a two-stage training recipe to effectively leverage the available compute budget and achieve strong results for the 405B model.
• Maintaining architectural simplicity: While the model architecture relies on the standard Transformer architecture, the authors opted for a minimalist approach, making minimal changes to optimize for efficiency and scalability.
• Refining the model with human feedback: The use of multi-round post-training with rejection sampling, supervised finetuning, and direct preference optimization enables the model to align closely with human preferences and improve its overall performance and helpfulness.
• Integrating multimodal capabilities: The authors successfully incorporated vision and speech capabilities into the model through a compositional approach, demonstrating the flexibility and scalability of foundation models.
• Prioritizing safety and responsibility: The authors implemented various safety mitigations, including the construction of safety benchmarks, the application of safety finetuning, and the development of Llama Guard, to ensure that the model generates safe and responsible content.

The release of Llama 3.1 is a significant step forward in the development of foundation models, offering a powerful new tool for researchers and developers. The authors hope that this work will:

• Accelerate research in foundation models: The open release of the models and the detailed insights into their development process will encourage further exploration and innovation in this field.
• Promote responsible development of AGI: The authors believe that the open release of foundation models plays a key role in fostering responsible development and encouraging the industry to embrace safety and ethical considerations.

The future of foundation models is filled with exciting potential, and the work on Llama 3.1 represents a significant step towards building more powerful, versatile, and responsible AI systems. The ongoing research and development efforts in this area will continue to push the boundaries of what is possible, leading to even more impactful and beneficial applications of AI.