diff --git "a/src/index.html" "b/src/index.html" --- "a/src/index.html" +++ "b/src/index.html" @@ -1,5 +1,6 @@ + @@ -9,45 +10,77 @@ FineWeb: decanting the web for the finest text data at scale - - - - -

🔭 Ultra-Guide to Scaling LLM training

-
-
- FineWeb -
- -
- - - - - - -

The performance of a large language model (LLM) depends heavily on the quality and size of the LLMs. - However, the pretraining datasets for state-of-the-art open LLMs like Llama 3 and Mixtral are not publicly available and very little is known about how they were created.

+ + + + + + + +

The performance of a large language model (LLM) depends heavily on the quality and size of the LLMs. + However, the pretraining datasets for state-of-the-art open LLMs like Llama 3 and Mixtral are + not publicly available and very little is known about how they were created.

-

Recently, we released 🍷 FineWeb, a new, large-scale - (15-trillion tokens, 44TB disk space) dataset for LLM pretraining. FineWeb is derived from 96 CommonCrawl snapshots and produces better-performing LLMs than other open pretraining datasets. +

Recently, we released 🍷 + FineWeb, a new, large-scale + (15-trillion tokens, 44TB disk space) dataset for LLM pretraining. FineWeb is derived from + 96 CommonCrawl snapshots and produces better-performing LLMs + than other open pretraining datasets. -

+
-
-
+
+
-
+
+ + + +
+
+
+
-
- - - -
-
- - - -
-
- - - -
-
- - - -
-
- - - -
-
- - - -
-
- - - -
-
- - - +
+ + +
-
-
- - - +
+
+ + +
-
+
+ +
+
+
+
+ + +
-
+
-
+
+
+
+ + + +
+
-
+
+
+
+ + + +
+
+ +
+
+ + + +
+
+
+
+ + + +
+
+ + + +
+
- -

TLDR: This blog covers a discussion on processing and evaluating data quality at scale, the 🍷 FineWeb - recipe (listing and explaining all of our design choices), and the process followed to create its 📚 FineWeb-Edu subset.

- -

Scaling Models and Hardware

- -

Now that we know the basics of distributed communication and computations it's time to apply this to training LLMs at scale. Here's the plan of action: we'll go through increasingly complex distribution strategies, namely data, then tensor and finally pipeline parallelism, and show three things:

- -
    -
  1. conceptual explanations with diagrams
    2. -
  2. a minimal coding example illustrating how to implement said strategy
    3. -
  3. scaling experiments show casing strengths and limits of the method with real data
    4. -
- -

For the experiments we scale across two dimensions: we make the models larger and larger and add more and more compute nodes and measure how throughput changes.

- -

So this is a good point to get ☕ #2 and we'll have a look at the setup for the practical experiments.

- -

Experiment setup

- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
1B (1)7B70B340B (2)400B (3)
N Layers24328096126
N Heads32326496128
Dimension2048409681921843216384
- -

(1) FineWeb ablation models

-

(2) Nemotron-340B architecture (without GQA)

-

(3) Llama-400B, ffn dim = 1.2 hidden dim (without GQA)

- - -

Distribution Methods

- -

Efficiently training LLMs now requires amounts of compute which exceed in most case single GPUs or machine. Large distributed clusters are thus used to train these models and can range from hundreds to thousands of nodes each usually equipped with up to 8 GPUs. To make the best use of such an expensive hardware, a range of distributed training methods have been developed with the goal of ensuring that GPUs are highly utilized at all times and not waiting for data/synchronization/etc.

- -

Several methods can be used to distribute training and we'll start with 4D parallelism followed-up by DeepSpeed stages. While we explain these strategies we'll also run experiments to determine the trade-offs and understand the optimal settings.

-

The name "4D parallelism" originates from the fact that it involves combining up to 4 distribution methods: data, tensor, pipeline, and sequence parallelism (each of these techniques can be used independently of the other). You may thus ask "So which one should I use?".

- -

Unfortunately, there is no universal answer as the response will actually depend on the cluster setup as well as the model architecture. But do not despair for in this section we'll develop strategies to figure out the best setting experimentally!

- -

In addition to 4D parallelism we'll also take a look at "DeepSpeed", a method developed by Microsoft which is generally complimentary to 4D parallelism and can be leveraged on top of it.

- -

Idea: show two things in every section

-
    -
  1. a small toy model (e.g. 4 layer FFN) we can interactively show with every approach
    2. -
  2. a benchmark showing the improvement/limits of the approach (e.g. when you cross 1 node with TP)
    3. -
- -

No Parallelism

- -

Let's quickly go over the basics before going into distributed training. When a model is trained on a single GPU, the training consists of 3 steps in the simplest case:

-
    -
  1. one forward pass,
    2. -
  2. one backward pass to compute the gradients, and
    3. -
  3. an optimization step using the gradients to update the parameters
    4. -
- -

As we'll see in the future, these steps may be repeated or intertwined but for now we'll start simple:

-

As we'll see in the future, these steps may be repeated or intertwined but for now we'll start simple:

- -Training Steps - -

In this figure the successive blue boxes on the top line can be seen as successive layers inside a model (same for the last line). The red boxes are the associated gradients for each of these layers.

- -

The batch size (bs) is one of the most important hyper-parameters in machine learning, affecting both model convergence and throughput.

- -

If the batch size is too small, gradients will tend to be noisy and the model may not be able to converge to optimal performances while a batch size too large can make the convergence of the model slower and waste compute. You can find a nice discussion of this topic in OpenAI's paper on large batch training (https://arxiv.org/pdf/1812.06162).

- -

The batch size also affects the throughput: a small batch size will require more optimizer steps to train on a given amount of samples. Optimizer steps are costly (in compute time) and the throughput will thus be lower than when using a larger batch size. On the other hand, larger batches, while leading to higher throughput may suffer from slow convergence in the limits as we've just seen. There is generally an optimal batch size from a convergence/performance point of view (note that the batch size can usually still be changed around the optimal batch size without major impact to the performance of the model).

- -

Note that in the LLM community, batch sizes are commonly reported in terms of tokens instead of number of samples (BST - Batch Size Tokens) as each token has a label and thus a loss term and can thus be considered individual (although highly correlated) samples.

- -

A sweet spot for LLM training is usually on the order of 4-20 million tokens per batch (links GPT-3, DeepSeek, Llama). In the simplest case, training on a single machine, the BS and BST can be computed from the model input sequence length as follows:

- - -bst=bs *seq - - -

(note that from here on forward we'll show the formulas for the batch size in number of samples but you can always get its token-unit counterpart by multiplying it with the sequence length)

- -

And we're now hitting our first scaling problem:

- -
-

what if we can't fit the model into GPU memory even with BS=1?

-
- -

Good question, reader!

- -

Let's start by understanding what led to our out-of-memory issue in the first place.

- -

A brief overview of memory usage in Transformers

-

To train a neural network model, one needs to store many elements in memory besides the weights themselves. Generally, the memory usage is made up from the following elements:

-
    -
  • model weights
    • -
  • model gradients
    • -
  • optimizer states
    • -
  • activations computed during the forward pass and which are needed to compute the backward pass
    • -
  • also CUDA Kernels require 1-2GB of GPU memory which you can quickly check yourself by running import torch; torch.ones((1, 1)).to("cuda") and then checking the GPU memory with nvidia-smi
    • -
  • lower rest memory usage from buffers, intermediate results and some memory that can't be used due to fragmentation
    • -
+

TLDR: This blog covers a discussion on processing and evaluating data quality at scale, the + 🍷 FineWeb + recipe (listing and explaining all of our design choices), and the process followed to create its 📚 + FineWeb-Edu subset.

+ +

Scaling Models and Hardware

+ +

Now that we know the basics of distributed communication and computations it's time to apply this to training + LLMs at scale. Here's the plan of action: we'll go through increasingly complex distribution strategies, + namely data, then tensor and finally pipeline parallelism, and show three things:

+ +
    +
  1. conceptual explanations with diagrams
    2. +
  2. a minimal coding example illustrating how to implement said strategy
    3. +
  3. scaling experiments show casing strengths and limits of the method with real data
    4. +
+ +

For the experiments we scale across two dimensions: we make the models larger and larger and add more and + more compute nodes and measure how throughput changes.

+ +

So this is a good point to get ☕ #2 and we'll have a look at the setup for the practical experiments.

+ +

Experiment setup

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
1B (1)7B70B340B (2)400B (3)
N Layers24328096126
N Heads32326496128
Dimension2048409681921843216384
+ +

(1) FineWeb ablation models

+

(2) Nemotron-340B architecture (without GQA)

+

(3) Llama-400B, ffn dim = 1.2 hidden dim (without GQA)

+ + +

Distribution Methods

+ +

Efficiently training LLMs now requires amounts of compute which exceed in most case single GPUs or machine. + Large distributed clusters are thus used to train these models and can range from hundreds to thousands of + nodes each usually equipped with up to 8 GPUs. To make the best use of such an expensive hardware, a range + of distributed training methods have been developed with the goal of ensuring that GPUs are highly utilized + at all times and not waiting for data/synchronization/etc.

+ +

Several methods can be used to distribute training and we'll start with 4D parallelism followed-up by + DeepSpeed stages. While we explain these strategies we'll also run experiments to determine the trade-offs + and understand the optimal settings.

+

The name "4D parallelism" originates from the fact that it involves combining up to 4 distribution methods: + data, tensor, pipeline, and sequence parallelism (each of these techniques can be used independently of the + other). You may thus ask "So which one should I use?".

+ +

Unfortunately, there is no universal answer as the response will actually depend on the cluster setup as well + as the model architecture. But do not despair for in this section we'll develop strategies to figure out the + best setting experimentally!

+ +

In addition to 4D parallelism we'll also take a look at "DeepSpeed", a method developed by Microsoft which is + generally complimentary to 4D parallelism and can be leveraged on top of it.

+ +

Idea: show two things in every section

+
    +
  1. a small toy model (e.g. 4 layer FFN) we can interactively show with every approach
    2. +
  2. a benchmark showing the improvement/limits of the approach (e.g. when you cross 1 node with TP)
    3. +
+ +

No Parallelism

+ +

Let's quickly go over the basics before going into distributed training. When a model is trained on a single + GPU, the training consists of 3 steps in the simplest case:

+
    +
  1. one forward pass,
    2. +
  2. one backward pass to compute the gradients, and
    3. +
  3. an optimization step using the gradients to update the parameters
    4. +
+ +

As we'll see in the future, these steps may be repeated or intertwined but for now we'll start simple:

+

As we'll see in the future, these steps may be repeated or intertwined but for now we'll start simple:

+ + Training Steps + +

In this figure the successive blue boxes on the top line can be seen as successive layers inside a model + (same for the last line). The red boxes are the associated gradients for each of these layers.

+ +

The batch size (bs) is one of the most important hyper-parameters in machine learning, affecting + both model convergence and throughput.

+ +

If the batch size is too small, gradients will tend to be noisy and the model may not be able to converge to + optimal performances while a batch size too large can make the convergence of the model slower and waste + compute. You can find a nice discussion of this topic in OpenAI's paper on large batch training (https://arxiv.org/pdf/1812.06162).

+ +

The batch size also affects the throughput: a small batch size will require more optimizer steps to train on + a given amount of samples. Optimizer steps are costly (in compute time) and the throughput will thus be + lower than when using a larger batch size. On the other hand, larger batches, while leading to higher + throughput may suffer from slow convergence in the limits as we've just seen. There is generally an optimal + batch size from a convergence/performance point of view (note that the batch size can usually still be + changed around the optimal batch size without major impact to the performance of the model).

+ +

Note that in the LLM community, batch sizes are commonly reported in terms of tokens instead of number of + samples (BST - Batch Size Tokens) as each token has a label and thus a loss term and can thus be considered + individual (although highly correlated) samples.

+ +

A sweet spot for LLM training is usually on the order of 4-20 million tokens per batch (links GPT-3, + DeepSeek, Llama). In the simplest case, training on a single machine, the BS and BST can + be computed from the model input sequence length as follows:

+ + + bst=bs *seq + + +

(note that from here on forward we'll show the formulas for the batch size in number of samples but you can + always get its token-unit counterpart by multiplying it with the sequence length)

+ +

And we're now hitting our first scaling problem:

+ +
+

what if we can't fit the model into GPU memory even with BS=1?

+
+ +

Good question, reader!

-

Scaling up training is usually a question of playing with those constituents to keep memory low while not impacting performance too much. We'll neglect the last two contributors as there's usually not that much you can do about them unless you dive deep in the code.

+

Let's start by understanding what led to our out-of-memory issue in the first place.

-

For the rest, they are usually different types of tensors that can have various sizes (usually multiples of one or several of batch size, sequence length, model hidden dimension and some potential sharding) and various precisions (with optimizer states and weights copy being often kept in full FP32 precision while activations can be of lower precision like BF16 or FP8). Let's try to get some intuition for the memory requirement of these various elements.

+

A brief overview of memory usage in Transformers

-

Let's first look at the weights, gradients and optimizer states. They are all dependent on the number of parameters in a model. For a simple LLM the number of parameters is given by the following formula:

- - -N = h*v + L * (12 * h^2 + 13*h) + 2*h - - -

In that equation, h corresponds to the hidden dimension, v to the vocabulary size, and L the number of layers in the model. Note that looking at the equation we can see that the term that will dominate at large model scales is the one with h^2 since it's the only term growing quadratically as we scale the models.

- -

Let's see how the number of parameters translates to memory usage. The memory requirements for the parameters and gradients are the number of parameters multiplied by the number of bytes per parameter. Mixed precision training with BF16 is the default nowadays which requires 2 bytes per parameter. In addition, there are a number of values necessary for the optimizer states: for ADAM it requires the momentum and the variance in FP32, each using 4 bytes, and an additional copy of the model weights in FP32, thus 12 bytes per parameter (ref: ZeRO):

- - -m_{params} = 2 * N -m_{grad} = 2 * N -m_{opt} = (4+4+4) * N - - -

In old-fashioned full precision training both parameters and gradients would require 4 bytes each but the optimizer on the other hand wouldn't need to store an extra full precision copy of the weights:

- - -m_{params} = 4 * N -m_{grad} = 4 * N -m_{opt} = (4+4) * N - - -

So we can easily see that mixed precision itself doesn't save memory as it just distributes the memory differently across the three components. So by multiplying the number of parameters by 16 (=2+2+12) you can quickly get a sense of how much GPU memory we need for a model:

-

So we can easily see that mixed precision itself doesn't save memory as it just distributes the memory differently across the three components. So by multiplying the number of parameters by 16 (=2+2+12) you can quickly get a sense of how much GPU memory we need for a model:

- - - - - - - - - - - - - - - - - - - - - - - - - - -
Model parametersMemory requirements
1B16 GB
7B112 GB
70B1120 GB
405B6480 GB
- -

We can further decrease the memory usage if we choose FP8 training instead of BF16 but it is much less stable and a very active research topic (see here) thus we won't go in details here.

- -

But we are not done yet, we'll also need to store the forward pass activations which are used during the backward pass to compute the gradients. The total memory required for the activations in mixed precision (which contributes the leading factor of 2 below) is given by the following equation:

- - -m_{act} = 2 * L* seq * bs * h * (34 + \frac{5*n_{heads}*seq}{h}) - - -

You can follow this NVIDIA paper for a complete derivation, it essentially requires you to do some accounting of all the sizes of intermediate activations between each operation. What's interesting here is that the memory is not static for a given model but depends critically on the sequence length. We can use the memory formulas and have a look how the memory usage changes for a model for various sequence lengths:

- -Memory Usage Graph 1 -Memory Usage Graph 2 - -

This graph tells a striking story: for short sequences, activations are almost negligible, but starting at around 2-4k tokens they start to take up a significant amount of memory while parameter, gradient and optimizer state are roughly independent of the sequence length and batch size. For large batch/sequence, activations however become by far the largest memory burden.

- -

Is there a way to tame this "activation explosion"?

- -

Good question, reader! I see you're following well and you're lucky as the answer is "Yes"! Let's talk about a technique called gradient checkpointing or more frequently activation recomputation which can help us cap activation memory footprint and is an essential tool in today's large model training toolbox.

- -

Activation recomputation

- -

The general idea behind gradient checkpointing is to discard some activations to save memory if we are willing to spend some extra compute to recompute them when needed. Typically we will save activations at some key points in memory and discard the rest and recompute them during the backward pass from the nearest activations:

- -Activation Recompute - -

We can select these key activations according to several strategies and modern frameworks usually choose among the following three strategies:

-
    -
  • None: We don't recompute activations during the backward pass and keep all activations in memory. While this is the fastest and thus computationally cheapest option, it also requires the most memory.
    • -
  • Full: The simplest strategy from a conceptual point of view is to checkpoint activations between each Transformer layer. This is usually called the full strategy since it requires a forward pass through each layer essentially adding a full forward pass during the backward pass. This strategy saves the most memory but is the most expensive one in terms of compute. This increases the compute cost by up to 30-40% which is very noticeable.
    • -
  • Selective: In general we can do better than full. The authors of this paper did a detailed analysis studying which activations grow the largest and have the cheapest recomputation cost in terms of FLOPs. Turns out that the attention computations fall in that category, and thus we can usually discard them and focus on checkpointing expensive feedforward computations. Note: for a GPT-3 (175B) model this means 70% activation memory reduction at a 2.7% compute cost.
    • -
- -

Let's see how recomputation strategies can drastically reduce the memory footprint while selective recomputation strikes a nice balance between memory saving and recomputation cost:

-

Let's see how recomputation strategies can drastically reduce the memory footprint while selective recomputation strikes a nice balance between memory saving and recomputation cost:

+

To train a neural network model, one needs to store many elements in memory besides the weights themselves. + Generally, the memory usage is made up from the following elements:

+
    +
  • model weights
    • +
  • model gradients
    • +
  • optimizer states
    • +
  • activations computed during the forward pass and which are needed to compute the backward pass
    • +
  • also CUDA Kernels require 1-2GB of GPU memory which you can quickly check yourself by running + import torch; torch.ones((1, 1)).to("cuda") and then checking the GPU memory with + nvidia-smi +
    • +
  • lower rest memory usage from buffers, intermediate results and some memory that can't be used due to + fragmentation
    • +
-Recomputation Strategies +

Scaling up training is usually a question of playing with those constituents to keep memory low while not + impacting performance too much. We'll neglect the last two contributors as there's usually not that much you + can do about them unless you dive deep in the code.

+ +

For the rest, they are usually different types of tensors that can have various sizes (usually multiples of + one or several of batch size, sequence length, model hidden dimension and some potential sharding) and + various precisions (with optimizer states and weights copy being often kept in full FP32 precision while + activations can be of lower precision like BF16 or FP8). Let's try to get some intuition for the memory + requirement of these various elements.

+ +

Let's first look at the weights, gradients and optimizer states. They are all dependent on the number of + parameters in a model. For a simple LLM the number of parameters is given by the following formula:

+ + + N = h*v + L * (12 * h^2 + 13*h) + 2*h + + +

In that equation, h corresponds to the hidden dimension, v to the vocabulary size, and + L the number of layers in the model. Note that looking at the equation we can see that the term + that will dominate at large model scales is the one with h^2 since it's the only term growing + quadratically as we scale the models. +

+ +

Let's see how the number of parameters translates to memory usage. The memory requirements for the parameters + and gradients are the number of parameters multiplied by the number of bytes per parameter. Mixed precision + training with BF16 is the default nowadays which requires 2 bytes per parameter. In addition, there are a + number of values necessary for the optimizer states: for ADAM it requires the momentum and the variance in + FP32, each using 4 bytes, and an additional copy of the model weights in FP32, thus 12 bytes per parameter + (ref: ZeRO):

+ + + m_{params} = 2 * N + m_{grad} = 2 * N + m_{opt} = (4+4+4) * N + + +

In old-fashioned full precision training both parameters and gradients would require 4 bytes each but the + optimizer on the other hand wouldn't need to store an extra full precision copy of the weights:

+ + + m_{params} = 4 * N + m_{grad} = 4 * N + m_{opt} = (4+4) * N + + +

So we can easily see that mixed precision itself doesn't save memory as it just distributes the memory + differently across the three components. So by multiplying the number of parameters by 16 (=2+2+12) you can + quickly get a sense of how much GPU memory we need for a model:

+

So we can easily see that mixed precision itself doesn't save memory as it just distributes the memory + differently across the three components. So by multiplying the number of parameters by 16 (=2+2+12) you can + quickly get a sense of how much GPU memory we need for a model:

+ + + + + + + + + + + + + + + + + + + + + + + + + + +
Model parametersMemory requirements
1B16 GB
7B112 GB
70B1120 GB
405B6480 GB
+ +

We can further decrease the memory usage if we choose FP8 training instead of BF16 but it is much less stable + and a very active research topic (see here) + thus we won't go in details here.

+ +

But we are not done yet, we'll also need to store the forward pass activations which are used during the + backward pass to compute the gradients. The total memory required for the activations in mixed precision + (which contributes the leading factor of 2 below) is given by the following equation:

+ + + m_{act} = 2 * L* seq * bs * h * (34 + \frac{5*n_{heads}*seq}{h}) + + +

You can follow this NVIDIA paper for a complete derivation, it + essentially requires you to do some accounting of all the sizes of intermediate activations between each + operation. What's interesting here is that the memory is not static for a given model but depends critically + on the sequence length. We can use the memory formulas and have a look how the memory usage changes for a + model for various sequence lengths:

+ + Memory Usage Graph 1 + Memory Usage Graph 2 + +

This graph tells a striking story: for short sequences, activations are almost negligible, but starting at + around 2-4k tokens they start to take up a significant amount of memory while parameter, gradient and + optimizer state are roughly independent of the sequence length and batch size. For large batch/sequence, + activations however become by far the largest memory burden.

+ +

Is there a way to tame this "activation explosion"?

+ +

Good question, reader! I see you're following well and you're lucky as the answer is "Yes"! Let's talk about + a technique called gradient checkpointing or more frequently activation + recomputation which can help us cap activation memory footprint and is an essential tool in + today's large model training toolbox.

+ +

Activation recomputation

+ +

The general idea behind gradient checkpointing is to discard some activations to save memory if we are + willing to spend some extra compute to recompute them when needed. Typically we will save activations at + some key points in memory and discard the rest and recompute them during the backward pass from the nearest + activations:

+ + Activation Recompute + +

We can select these key activations according to several strategies and modern frameworks usually choose + among the following three strategies:

+
    +
  • None: We don't recompute activations during the backward pass and keep all activations + in memory. While this is the fastest and thus computationally cheapest option, it also requires the most + memory.
    • +
  • Full: The simplest strategy from a conceptual point of view is to checkpoint + activations between each Transformer layer. This is usually called the full strategy since + it requires a forward pass through each layer essentially adding a full forward pass during the backward + pass. This strategy saves the most memory but is the most expensive one in terms of compute. This + increases the compute cost by up to 30-40% which is very noticeable.
    • +
  • Selective: In general we can do better than full. The authors of this paper did a detailed analysis studying which + activations grow the largest and have the cheapest recomputation cost in terms of FLOPs. Turns out that + the attention computations fall in that category, and thus we can usually discard them and focus on + checkpointing expensive feedforward computations. Note: for a GPT-3 (175B) model this means 70% + activation memory reduction at a 2.7% compute cost.
    • +
-

Note: Hardware vs Model flops.

+

Let's see how recomputation strategies can drastically reduce the memory footprint while selective + recomputation strikes a nice balance between memory saving and recomputation cost:

+

Let's see how recomputation strategies can drastically reduce the memory footprint while selective + recomputation strikes a nice balance between memory saving and recomputation cost:

+ + Recomputation Strategies + +

Note: Hardware vs Model flops.

+ +

Most frameworks these days use FlashAttention (TODO: see later) which makes the attention computation less + memory intensive through kernel fusion, thus most trainings use the full settings.

+ +

We can save some GPU memory with activation recomputation but this only delays by a bit the next bottleneck: + as hinted earlier for LLM training there is usually a sweet spot for the GBST and we need to work out the + training configuration backward from there. However, you can't choose MBS to be an arbitrary large number on + your GPU; at some point you will run out of GPU memory again since you need to store at least some of the + activations in memory.

+ +

There is a useful trick to compensate for that: gradient accumulation (GradAcc). + With gradient accumulation we will split our batch in micro-batch, do forward and backward passes repeatedly + on each micro-batch, compute the gradients, and, as the name suggests, sum the gradients step by step before + doing a final optimizer step.

+ +

We call the micro batch size (MBS) the batch size for each forward pass on a single node (the + number of samples flowing through the model in one forward pass). We'll refer to the overall batch size + between each optimizer step as the global batch size (GBS). If we do one optimizer step each 8 + forward/backward pass, the global batch size will be 8 times the micro batch size. +

+ +

What we now call global batch size thus corresponds to what we've called up to now just + batch size for simplicity (we now make the terms more precise to avoid ambiguity). +

+ +

With gradient accumulation the global batch size can be computed as follows:

+ + + BS = GBS=MBS * GradAcc + + +

Gradient accumulation allows us to effectively increase our batch size up to infinity (!) while the memory + footprint stays constant. Gradient accumulation is also compatible with activation recomputation for further + memory reduction. One drawback however, is that gradient accumulation requires multiple consecutive + forward/backward passes per optimization step thereby increasing the compute overhead and slowing down + training. No free lunch!

+ + Gradient Accumulation + +

This is actually a bummer since the forward/backward passes for each micro-batch could actually totally be + run in parallel. They are independent from each other and the only changing parameter are the input samples. +

+ +

Here comes data parallelism to solve exactly this problem! Let's take a look, you say? Okay sure!

+ +

Data Parallelism

+ +

The idea behind data parallelism (DP) is to parallelize forward and backward passes across GPUs, passing + different batches of data per GPU (or groups of GPUs) to the same model instance. Just like for gradient + accumulation, we need to average gradients across instances before we do the optimization step. The GBS + equation can then be extended to:

+ + + GBS=MBS * GradAcc * DP + + +

This means that we can reduce the number of gradient accumulation steps in favor of data parallel processes + which speeds up training. In practice, people will tend to max out the number of data parallel nodes (the DP + above) as much as possible as it's inherently parallel versus the sequential Gradient Accumulation. Gradient + accumulation is then added only to achieve a target batch size if DP alone is not sufficient. One exception + to that is pipeline parallelism which we'll discuss later.

+ + Data Parallelism + +

As you can see on the figure above, some gradients can already be gathered and summed (red boxes) even before + gradients down the line (red boxes on the left of the current gradient) are still being computed. This + significantly speeds up data parallelism. For instance, as soon as the backward pass of the last layer is + done (last boxes on the right) those gradients can already be gathered/summed while the backward pass + computations move to earlier layers, aka to the left. This lowers the communication/bandwidth pressure to + sync gradients of the full model as it can be performed in part in parallel to the computation of said + gradients. See this article for more + information.

+ +

A general recipe to determine an optimal data-parallel setup can be as follows:

+
    +
  1. Determine the best (global) batch size in tokens to use either by consulting literature or running + experiments? This determines the GBST.
    2. +
  2. Select a sequence length for training, again by either consulting literature or running experiments. + Generally 2-8k tokens works reliably well.
    3. +
  3. You now know the batch size (GBS=GBST/SeqLen). Find the maximum MBS on a single GPU by increasing the + local batch size until you run out of memory. This determines the MBS.
    4. +
  4. Finally, the number of available GPUs corresponds to the potential DP. The ratio of GPT to DP determines + the remaining number of gradient accumulation steps needed for the desired GBS.
    5. +
+ +

If the gradient accumulation ratio is lower than one, i.e. you have too many GPUs (!), you can either choose + to not use all your GPUs or test if a lower MBS will speed up training. In these cases, you may want to + prioritize throughput over the individual GPU utilization, you can then choose DP first and use a smaller + MBS than possible in order to speed up training.

+ +

Time to take a concrete example: We want to train a model with a GBS of 4M tokens and a sequence length of + 4k. This means our batch size will be 1024 samples (we pick powers of two). We observe that a single of our + GPU can fit MBS=2 in memory and we have 128 GPUs available for training. This means with 4 gradient + accumulation steps we'll achieve our goal of 1024 samples or 4M tokens per training step. Now what if we + suddenly have 1024 GPUs available? We can achieve the same GBS and thus identical training by setting both + MBS and gradient accumulation to 1 speeding up training significantly.

+ +

[EXPERIMENTS WHERE WE INCREASE DP AND SHOW THROUGHPUT FOR SEVERAL MODELS]

+ +

We've explored data parallelism, a simple strategy to scale training across more GPUs and gives consistent + speed improvements. The keen reader might have noticed however that it rests on the assumption that we can + fit at least one input sample forward pass (MBS=1) into our GPU memory. This is not always the + case! In particular for larger models which often don't fit into a single GPU anymore even with activation + recomputations activated.

+ +

In such case, we need to shard the model across devices! We'll now study two complementary sharding methods, + tensor and pipeline parallelism which are doing that. Let's start by the simplest, tensor parallelism!

+ +

Tensor Parallelism

+ +

So you've exhausted all the previous textbook tricks to try to fit your model on a single GPU but it still + doesn't fit? Let's try to distribute this model across several GPUs. Unlike DP we will not simply duplicate + the model but various parts of the model instance will be living on various GPUs.

+ +

If we take a look at a typical matrix multiplication (the core of a neural network), we can get an idea about + how we could split the model:

+ + Matrix Multiplication Example + +

Tensor parallelism is a technique in which a tensor is split into N shards along a particular dimension + across N GPUs. Matrices can be split either on the column part or row part leading to row and column + parallelism. Depending on which splitting strategy we choose will require different communications + primitives.

+ +

Column linear:

+
    +
  • Splitting by column or row involves different synchronization primitives: +
      +
    • column: +
        +
      • A Broadcast operation is used to send the same input to different GPUs, +
        • +
      • Multiplications are done independently on the GPUs, and finally
        • +
      • An All-gather operation is used to gather the output results.
        • +
      +
      • +
    • Row: +
        +
      • A Scatter operation is used to split the input and send it to different + GPUs (we split the weight row-wise),
        • +
      • Multiplications are done independently on the GPUs, and finally
        • +
      • An All-reduce operation is used to add the results together and the + full output results.
        • +
      +
      • +
    +
    • +
-

Most frameworks these days use FlashAttention (TODO: see later) which makes the attention computation less memory intensive through kernel fusion, thus most trainings use the full settings.

+

This was for an example matrix multiplication. How do we apply this in practice to a real model? In the + Transformer, there are 2 basic building blocks where tensor parallel can be applied:

+
    +
  • Feedforward layers (MLP)
    • +
  • Multi-Head Attention (MHA)
    • +
-

We can save some GPU memory with activation recomputation but this only delays by a bit the next bottleneck: as hinted earlier for LLM training there is usually a sweet spot for the GBST and we need to work out the training configuration backward from there. However, you can't choose MBS to be an arbitrary large number on your GPU; at some point you will run out of GPU memory again since you need to store at least some of the activations in memory.

+

Feedforward layers comprise 2 successive MLPs with a non-linearity in-between. Here is the first part of it: +

-

There is a useful trick to compensate for that: gradient accumulation (GradAcc). With gradient accumulation we will split our batch in micro-batch, do forward and backward passes repeatedly on each micro-batch, compute the gradients, and, as the name suggests, sum the gradients step by step before doing a final optimizer step.

+ Feedforward Layers -

We call the micro batch size (MBS) the batch size for each forward pass on a single node (the number of samples flowing through the model in one forward pass). We'll refer to the overall batch size between each optimizer step as the global batch size (GBS). If we do one optimizer step each 8 forward/backward pass, the global batch size will be 8 times the micro batch size.

+

Should we use row or column parallelization for the first MLP?

-

What we now call global batch size thus corresponds to what we've called up to now just batch size for simplicity (we now make the terms more precise to avoid ambiguity).

+

Well it turns out parallelized GeLU only works in Column schema:

-

With gradient accumulation the global batch size can be computed as follows:

+

In column schema:

+ + GeLU(cat([XW1, XW2])) = cat([GeLU(XW1), GeLU(XW2)]) + - -BS = GBS=MBS * GradAcc - +

In row schema:

+ + GeLU(XW1 + XW2) \neq GeLU(XW1) + GeLU(XW2) + -

Gradient accumulation allows us to effectively increase our batch size up to infinity (!) while the memory footprint stays constant. Gradient accumulation is also compatible with activation recomputation for further memory reduction. One drawback however, is that gradient accumulation requires multiple consecutive forward/backward passes per optimization step thereby increasing the compute overhead and slowing down training. No free lunch!

+

If you rather like code, note that we can prove this with the following snippet as well:

-Gradient Accumulation + + ``` + + def example_gelu(): + from torch.nn.functional import gelu -

This is actually a bummer since the forward/backward passes for each micro-batch could actually totally be run in parallel. They are independent from each other and the only changing parameter are the input samples.

+ X = torch.randn(4, 2, device="cuda", dtype=torch.float32) + W = torch.randn(2, 2, device="cuda", dtype=torch.float32) -

Here comes data parallelism to solve exactly this problem! Let's take a look, you say? Okay sure!

+ W_0, W_1 = W.chunk(2, dim=1) -

Data Parallelism

+ # Column linear + y_col_1 = torch.cat([gelu(X @ W_0), gelu(X @ W_1)], dim=1) + y_col_2 = gelu(torch.cat([X @ W_0, X @ W_1], dim=1)) -

The idea behind data parallelism (DP) is to parallelize forward and backward passes across GPUs, passing different batches of data per GPU (or groups of GPUs) to the same model instance. Just like for gradient accumulation, we need to average gradients across instances before we do the optimization step. The GBS equation can then be extended to:

+ # All match + torch.testing.assert_close(y_col_1, y_col_2, rtol=1e-5, atol=1e-5) - -GBS=MBS * GradAcc * DP - + # Row linear + X_0, X_1 = X.chunk(2, dim=1) + W_0, W_1 = W.chunk(2, dim=0) + y_row_1 = gelu(X_0 @ W_0) + gelu(X_1 @ W_1) + y_row_2 = gelu(X_0 @ W_0 + X_1 @ W_1) -

This means that we can reduce the number of gradient accumulation steps in favor of data parallel processes which speeds up training. In practice, people will tend to max out the number of data parallel nodes (the DP above) as much as possible as it's inherently parallel versus the sequential Gradient Accumulation. Gradient accumulation is then added only to achieve a target batch size if DP alone is not sufficient. One exception to that is pipeline parallelism which we'll discuss later.

+ # Mismatch + torch.testing.assert_close(y_row_1, y_row_2, rtol=1e-5, atol=1e-5) + -Data Parallelism +

To avoid a synchronization step directly after the first MLP, we'll thus start with Column Parallel and be + able to directly perform parallel GELU.

-

As you can see on the figure above, some gradients can already be gathered and summed (red boxes) even before gradients down the line (red boxes on the left of the current gradient) are still being computed. This significantly speeds up data parallelism. For instance, as soon as the backward pass of the last layer is done (last boxes on the right) those gradients can already be gathered/summed while the backward pass computations move to earlier layers, aka to the left. This lowers the communication/bandwidth pressure to sync gradients of the full model as it can be performed in part in parallel to the computation of said gradients. See this article for more information.

+

Now, what about the second MLP? Should it be column or row parallel? Let's draft both options:

+
    +
  • Column Parallel followed by Column Parallel
    • + Column Parallel Schema 1 +
  • Column Parallel followed by Row Parallel
    • + Column Parallel Schema 2 +
-

A general recipe to determine an optimal data-parallel setup can be as follows:

-
    -
  1. Determine the best (global) batch size in tokens to use either by consulting literature or running experiments? This determines the GBST.
    2. -
  2. Select a sequence length for training, again by either consulting literature or running experiments. Generally 2-8k tokens works reliably well.
    3. -
  3. You now know the batch size (GBS=GBST/SeqLen). Find the maximum MBS on a single GPU by increasing the local batch size until you run out of memory. This determines the MBS.
    4. -
  4. Finally, the number of available GPUs corresponds to the potential DP. The ratio of GPT to DP determines the remaining number of gradient accumulation steps needed for the desired GBS.
    5. -
+

We see that the "Column Parallel followed by Row Parallel" schema only involves two communications instead of + four. It's thus the most efficient schema in terms of communications.

-

If the gradient accumulation ratio is lower than one, i.e. you have too many GPUs (!), you can either choose to not use all your GPUs or test if a lower MBS will speed up training. In these cases, you may want to prioritize throughput over the individual GPU utilization, you can then choose DP first and use a smaller MBS than possible in order to speed up training.

+

Let's take a quick look at the backward pass:

+ Backward Pass 1 + Backward Pass 2 -

Time to take a concrete example: We want to train a model with a GBS of 4M tokens and a sequence length of 4k. This means our batch size will be 1024 samples (we pick powers of two). We observe that a single of our GPU can fit MBS=2 in memory and we have 128 GPUs available for training. This means with 4 gradient accumulation steps we'll achieve our goal of 1024 samples or 4M tokens per training step. Now what if we suddenly have 1024 GPUs available? We can achieve the same GBS and thus identical training by setting both MBS and gradient accumulation to 1 speeding up training significantly.

+ + def column_linear_forward(X, local_W, group): + Y_local = X @ local_W.t() + return Y_local -

[EXPERIMENTS WHERE WE INCREASE DP AND SHOW THROUGHPUT FOR SEVERAL MODELS]

+ def column_linear_backward(local_grad_Y, X, local_W, group): + local_grad_X = local_grad_Y @ local_W + grad_W = local_grad_Y.t() @ X + return local_grad_X, grad_W -

We've explored data parallelism, a simple strategy to scale training across more GPUs and gives consistent speed improvements. The keen reader might have noticed however that it rests on the assumption that we can fit at least one input sample forward pass (MBS=1) into our GPU memory. This is not always the case! In particular for larger models which often don't fit into a single GPU anymore even with activation recomputations activated.

+ def row_linear_forward(local_X, local_W, group): + Y_local = local_X @ local_W.t() + dist.all_reduce(Y_local, group=group) + Y = Y_local + return Y -

In such case, we need to shard the model across devices! We'll now study two complementary sharding methods, tensor and pipeline parallelism which are doing that. Let's start by the simplest, tensor parallelism!

+ def row_linear_backward(grad_Y, X, local_W, group): + local_grad_X = grad_Y @ local_W + grad_W = grad_Y.t() @ X + return local_grad_X, grad_W -

Tensor Parallelism

+ def example_column_row_linear(): + # torchrun --nproc_per_node=2 tp_all_reduce.py + group = dist.distributed_c10d._get_default_group() -

So you've exhausted all the previous textbook tricks to try to fit your model on a single GPU but it still doesn't fit? Let's try to distribute this model across several GPUs. Unlike DP we will not simply duplicate the model but various parts of the model instance will be living on various GPUs.

+ X_ref = torch.arange(4 * 2, device="cuda", dtype=torch.float32, requires_grad=True).reshape(4, 2) + W_ref_layer1 = torch.arange(1, 5, device="cuda", dtype=torch.float32, requires_grad=True).reshape(2, 2) * 10 + W_ref_layer2 = torch.arange(1, 5, device="cuda", dtype=torch.float32, requires_grad=True).reshape(2, 2) -

If we take a look at a typical matrix multiplication (the core of a neural network), we can get an idea about how we could split the model:

+ X_ref.retain_grad() + W_ref_layer1.retain_grad() + W_ref_layer2.retain_grad() -Matrix Multiplication Example + dist.broadcast(X_ref, src=0, group=group) + dist.broadcast(W_ref_layer1, src=0, group=group) + dist.broadcast(W_ref_layer2, src=0, group=group) -

Tensor parallelism is a technique in which a tensor is split into N shards along a particular dimension across N GPUs. Matrices can be split either on the column part or row part leading to row and column parallelism. Depending on which splitting strategy we choose will require different communications primitives.

+ X = X_ref.clone() + W_layer1 = W_ref_layer1.clone() + W_layer2 = W_ref_layer2.clone() -

Column linear:

-
    -
  • Splitting by column or row involves different synchronization primitives: -
      -
    • column: -
        -
      • A Broadcast operation is used to send the same input to different GPUs,
        • -
      • Multiplications are done independently on the GPUs, and finally
        • -
      • An All-gather operation is used to gather the output results.
        • -
      -
      • -
    • Row: -
        -
      • A Scatter operation is used to split the input and send it to different GPUs (we split the weight row-wise),
        • -
      • Multiplications are done independently on the GPUs, and finally
        • -
      • An All-reduce operation is used to add the results together and the full output results.
        • -
      -
      • -
    -
    • -
- -

This was for an example matrix multiplication. How do we apply this in practice to a real model? In the Transformer, there are 2 basic building blocks where tensor parallel can be applied:

-
    -
  • Feedforward layers (MLP)
    • -
  • Multi-Head Attention (MHA)
    • -
- -

Feedforward layers comprise 2 successive MLPs with a non-linearity in-between. Here is the first part of it:

- -Feedforward Layers - -

Should we use row or column parallelization for the first MLP?

- -

Well it turns out parallelized GeLU only works in Column schema:

- -

In column schema:

- -GeLU(cat([XW1, XW2])) = cat([GeLU(XW1), GeLU(XW2)]) - - -

In row schema:

- -GeLU(XW1 + XW2) \neq GeLU(XW1) + GeLU(XW2) - - -

If you rather like code, note that we can prove this with the following snippet as well:

- - -``` - -def example_gelu(): - from torch.nn.functional import gelu - - X = torch.randn(4, 2, device="cuda", dtype=torch.float32) - W = torch.randn(2, 2, device="cuda", dtype=torch.float32) - - W_0, W_1 = W.chunk(2, dim=1) - - # Column linear - y_col_1 = torch.cat([gelu(X @ W_0), gelu(X @ W_1)], dim=1) - y_col_2 = gelu(torch.cat([X @ W_0, X @ W_1], dim=1)) - - # All match - torch.testing.assert_close(y_col_1, y_col_2, rtol=1e-5, atol=1e-5) - - # Row linear - X_0, X_1 = X.chunk(2, dim=1) - W_0, W_1 = W.chunk(2, dim=0) - y_row_1 = gelu(X_0 @ W_0) + gelu(X_1 @ W_1) - y_row_2 = gelu(X_0 @ W_0 + X_1 @ W_1) - - # Mismatch - torch.testing.assert_close(y_row_1, y_row_2, rtol=1e-5, atol=1e-5) - - -

To avoid a synchronization step directly after the first MLP, we'll thus start with Column Parallel and be able to directly perform parallel GELU.

- -

Now, what about the second MLP? Should it be column or row parallel? Let's draft both options:

-
    -
  • Column Parallel followed by Column Parallel
    • - Column Parallel Schema 1 -
  • Column Parallel followed by Row Parallel
    • - Column Parallel Schema 2 -
- -

We see that the "Column Parallel followed by Row Parallel" schema only involves two communications instead of four. It's thus the most efficient schema in terms of communications.

- -

Let's take a quick look at the backward pass:

-Backward Pass 1 -Backward Pass 2 - - -def column_linear_forward(X, local_W, group): - Y_local = X @ local_W.t() - return Y_local - -def column_linear_backward(local_grad_Y, X, local_W, group): - local_grad_X = local_grad_Y @ local_W - grad_W = local_grad_Y.t() @ X - return local_grad_X, grad_W - -def row_linear_forward(local_X, local_W, group): - Y_local = local_X @ local_W.t() - dist.all_reduce(Y_local, group=group) - Y = Y_local - return Y - -def row_linear_backward(grad_Y, X, local_W, group): - local_grad_X = grad_Y @ local_W - grad_W = grad_Y.t() @ X - return local_grad_X, grad_W + # Forward + Y_ref_linear1 = X_ref @ W_ref_layer1.t() + Y_ref_linear1.retain_grad() -def example_column_row_linear(): - # torchrun --nproc_per_node=2 tp_all_reduce.py - group = dist.distributed_c10d._get_default_group() + # We will transpose for matrix multiplication. As a result, we need to split row-wise + Y_local_linear1 = column_linear_forward(X, split_tensor(W_layer1, dim=0), group) - X_ref = torch.arange(4 * 2, device="cuda", dtype=torch.float32, requires_grad=True).reshape(4, 2) - W_ref_layer1 = torch.arange(1, 5, device="cuda", dtype=torch.float32, requires_grad=True).reshape(2, 2) * 10 - W_ref_layer2 = torch.arange(1, 5, device="cuda", dtype=torch.float32, requires_grad=True).reshape(2, 2) + torch.testing.assert_close(Y_local_linear1, split_tensor(Y_ref_linear1, dim=1), rtol=1e-5, atol=1e-5) - X_ref.retain_grad() - W_ref_layer1.retain_grad() - W_ref_layer2.retain_grad() - - dist.broadcast(X_ref, src=0, group=group) - dist.broadcast(W_ref_layer1, src=0, group=group) - dist.broadcast(W_ref_layer2, src=0, group=group) - - X = X_ref.clone() - W_layer1 = W_ref_layer1.clone() - W_layer2 = W_ref_layer2.clone() - - # Forward - Y_ref_linear1 = X_ref @ W_ref_layer1.t() - Y_ref_linear1.retain_grad() + Y_local_linear2 = row_linear_forward(Y_local_linear1, split_tensor(W_ref_layer2, dim=1), group) + Y_ref_linear2 = Y_ref_linear1 @ W_ref_layer2.t() + torch.testing.assert_close(Y_local_linear2, Y_ref_linear2, rtol=1e-5, atol=1e-5) - # We will transpose for matrix multiplication. As a result, we need to split row-wise - Y_local_linear1 = column_linear_forward(X, split_tensor(W_layer1, dim=0), group) + # Backward + Y_ref_linear2.sum().backward() - torch.testing.assert_close(Y_local_linear1, split_tensor(Y_ref_linear1, dim=1), rtol=1e-5, atol=1e-5) - - Y_local_linear2 = row_linear_forward(Y_local_linear1, split_tensor(W_ref_layer2, dim=1), group) - Y_ref_linear2 = Y_ref_linear1 @ W_ref_layer2.t() - torch.testing.assert_close(Y_local_linear2, Y_ref_linear2, rtol=1e-5, atol=1e-5) - - # Backward - Y_ref_linear2.sum().backward() - - grad_Y = torch.ones_like(Y_ref_linear2) - grad_X_linear2, grad_W_linear2 = row_linear_backward(grad_Y, Y_local_linear1, split_tensor(W_layer2, dim=1), group) + grad_Y = torch.ones_like(Y_ref_linear2) + grad_X_linear2, grad_W_linear2 = row_linear_backward(grad_Y, Y_local_linear1, split_tensor(W_layer2, dim=1), + group) - torch.testing.assert_close(grad_X_linear2, split_tensor(Y_ref_linear1.grad, dim=1), rtol=1e-5, atol=1e-5) - torch.testing.assert_close(grad_W_linear2, split_tensor(W_ref_layer2.grad, dim=1), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grad_X_linear2, split_tensor(Y_ref_linear1.grad, dim=1), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grad_W_linear2, split_tensor(W_ref_layer2.grad, dim=1), rtol=1e-5, atol=1e-5) - grad_X, grad_W = column_linear_backward(grad_X_linear2, X, split_tensor(W_layer1, dim=0), group) + grad_X, grad_W = column_linear_backward(grad_X_linear2, X, split_tensor(W_layer1, dim=0), group) - torch.testing.assert_close(grad_X, X_ref.grad, rtol=1e-5, atol=1e-5) - torch.testing.assert_close(grad_W, split_tensor(W_ref_layer1.grad, dim=0), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grad_X, X_ref.grad, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grad_W, split_tensor(W_ref_layer1.grad, dim=0), rtol=1e-5, atol=1e-5) -if __name__ == "__main__": - dist.init_process_group("nccl", rank=int(os.environ["RANK"]), world_size=int(os.environ["WORLD_SIZE"])) - torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + if __name__ == "__main__": + dist.init_process_group("nccl", rank=int(os.environ["RANK"]), world_size=int(os.environ["WORLD_SIZE"])) + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) - example_column_row_linear() - + example_column_row_linear() + -

Now that we've found the most efficient schema for the Feedforward part of the transformer, let's take a look at the multi-head attention block (MHA).

+

Now that we've found the most efficient schema for the Feedforward part of the transformer, let's take a look + at the multi-head attention block (MHA).

-

We can generally follow a similar approach where the Q, K, V will be split in a Column Parallel fashion and the output projection will be split along the Row dimension.

+

We can generally follow a similar approach where the Q, K, V will be split in a Column Parallel fashion and + the output projection will be split along the Row dimension.

-Multi-Head Attention Block + Multi-Head Attention Block -

To dive in further particularities, a nice reference paper detailing TP is for instance Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.

+

To dive in further particularities, a nice reference paper detailing TP is for instance Megatron-LM: Training Multi-Billion Parameter Language Models + Using Model Parallelism.

-

Note: Sequence Parallel

+

Note: Sequence Parallel

-

Sequence Parallelism

+

Sequence Parallelism

-

Tensor parallelism has been a great help to parallelize some of our computation on several GPU nodes with the limited cost of a few communication operations.

+

Tensor parallelism has been a great help to parallelize some of our computation on several GPU nodes with the + limited cost of a few communication operations.

-

It also had the additional benefit of reducing memory usage by splitting intermediate activations inside the feedforward elements across GPUs and thereby reducing the activations to store on each node.

+

It also had the additional benefit of reducing memory usage by splitting intermediate activations inside the + feedforward elements across GPUs and thereby reducing the activations to store on each node.

-

Could we push this approach further?

+

Could we push this approach further?

-

Sequence parallelism applies this same idea to other parts of our model. We've applied tensor parallelism to two main parts in our models where combination of MLP allowed to naturally split the weights along major axis.

+

Sequence parallelism applies this same idea to other parts of our model. We've applied tensor parallelism to + two main parts in our models where combination of MLP allowed to naturally split the weights along major + axis.

-

The rest of the model mostly comprises layer norms, dropout and various summation of residuals, these contribute little to the computation but come with rather large forward activations to store.

+

The rest of the model mostly comprises layer norms, dropout and various summation of residuals, these + contribute little to the computation but come with rather large forward activations to store.

-

[Add some illustration of the forward activations to store for each part]

+

[Add some illustration of the forward activations to store for each part]

-

Context Parallelism

+

Context Parallelism

-

Even though TP-SP mode helps reduce the memory used by activation values, it has two main drawbacks:

-
    -
  1. Internode connections are usually slow, so the TP degree shouldn't typically exceed 8
    2. -
  2. The TP degree is limited by the number of Key/Value heads, which is 8 for LLaMA 3 8B.
    3. -
+

Even though TP-SP mode helps reduce the memory used by activation values, it has two main drawbacks:

+
    +
  1. Internode connections are usually slow, so the TP degree shouldn't typically exceed 8
    2. +
  2. The TP degree is limited by the number of Key/Value heads, which is 8 for LLaMA 3 8B.
    3. +
-

An empirical estimation is that with TP=8, you can only train an 8B model with a 20K context length. However, LLaMA 3.1 has managed to scale the context length to 128K by using context parallelism.

+

An empirical estimation is that with TP=8, you can only train an 8B model with a 20K context length. However, + LLaMA 3.1 has managed to scale the context length to 128K by using context parallelism.

-

There are several ways to implement sequence parallelism. We used ring attention, which overlaps communication and computation. LLaMA3.1 uses all-gather along the sequence dimension because it is easier and more flexible to support different types of attention masks in all-gather based CP attention, such as the document mask.

+

There are several ways to implement sequence parallelism. We used ring attention, which overlaps + communication and computation. LLaMA3.1 uses all-gather along the sequence dimension because it is easier + and more flexible to support different types of attention masks in all-gather based CP attention, such as + the document mask.

-

Pipeline Parallelism

+

Pipeline Parallelism

-

Overlapping computation and communication

+

Overlapping computation and communication

-

ZeRO

+

ZeRO

-

II – Architecture

+

II – Architecture

-

Transformers

+

Transformers

-

Choosing the right dimensions

+

Choosing the right dimensions

-

Positional Embeddings (Learned, RoPE, ALiBi)

+

Positional Embeddings (Learned, RoPE, ALiBi)

-

RoPE

+

RoPE

-

In the transformer model, tokens have no inherent information about their positional information. For these reasons, we need to use a positional encoding function.

+

In the transformer model, tokens have no inherent information about their positional information. For these + reasons, we need to use a positional encoding function.

-

Assuming that in the multi-head attention layer, q_m is the "position-aware" query vector corresponding to a token at position m, k_n the "position-aware" key vector corresponding to the token at position n and f is our position embedding function, we would like our position vector to be a function of the input vectors and absolute positions like this:

+

Assuming that in the multi-head attention layer, q_m is the "position-aware" query vector + corresponding to a token at position m, k_n the "position-aware" key vector corresponding + to the token at position n and f is our position embedding function, we would like our + position vector to be a function of the input vectors and absolute positions like this:

- -q_m = f(q,m) -k_n = f(k,n) - + + q_m = f(q,m) + k_n = f(k,n) + -

We may also want the positional encoding to model relative positional information between two input tokens. Relative positions help the model to operate across longer context spans and even context lengths not seen during training. The attention operation is generally a dot product operation between "position-aware" vectors q and k, so for a positional encoding that contains relative positional information, we'll want to have:

+

We may also want the positional encoding to model relative positional information between two input tokens. + Relative positions help the model to operate across longer context spans and even context lengths not seen + during training. The attention operation is generally a dot product operation between "position-aware" + vectors q and k, so for a positional encoding that contains relative positional + information, we'll want to have:

- - = g(q, k, m-n) - + + = g(q, k, m-n) + -

In other words, we want the result of ⟨ 𝑞_𝑚 , 𝑘_𝑛 ⟩ to depend on the values of q and k themselves, as well as their relative position m − n, but not m and n. This way, the model can focus on the relative difference between two tokens rather than their absolute positions.

+

In other words, we want the result of ⟨ 𝑞_𝑚 , 𝑘_𝑛 ⟩ to depend on the values of q and + k themselves, as well as their relative position m − n, but not m and n. + This way, the model can focus on the relative difference between two tokens rather than their absolute + positions. +

-

Let's show that the RoPE positional embedding formulation satisfies the above formula.

+

Let's show that the RoPE positional embedding formulation satisfies the above formula.

-

Rotation matrix

+

Rotation matrix

-

RoPE are based on rotation matrices which have simple and interesting properties for us. In a 2D space, a rotation matrix has the following form:

+

RoPE are based on rotation matrices which have simple and interesting properties for us. In a 2D space, a + rotation matrix has the following form:

- -R(θ) = -\begin{pmatrix} -\cosθ & -\sinθ \\ -\sinθ & \cosθ -\end{pmatrix} - + + R(θ) = + \begin{pmatrix} + \cosθ & -\sinθ \\ + \sinθ & \cosθ + \end{pmatrix} + -

The rotation matrix has the following properties:

-
    -
  • R(θ)T = R(-θ)
    • -
  • R(θ1)R(θ2) = R(θ1+θ2)
    • -
- -Rotation Matrix - -

RoPE in 2D space

- -

Assuming q and k are 2D column vectors, we can show that:

- - - = (R(θ_1)q)T (R(θ_2)k) = qTR(-θ_1)R(θ_2)k = qTR(θ_2-θ_1)k = (R(θ_1-θ_2)q)Tk = - - -

Therefore, if we define our position embedding like this: f(x, m) = R(mθ)x where R is a 2D rotation matrix, we have q_m = R(mθ)q and k_n = R(nθ)k and then:

- - - = = - - -

We can see that a multiplication with a rotation matrix is exactly the positional encoding we were looking for. The result of ⟨ 𝑞_𝑚 , 𝑘_𝑛 ⟩ only depends on q, k and m-n.

- -

Implementation

- -

In our case, our internal vectors (the activations in our model) have much more than two elements. Let's pair elements to get 2D vectors and apply the 2D rotation operation on these pairs.

- -

There are combinatorially many ways we can pair elements but generally two options are the most popular for implementing RoPE: we call them the interleaved and non-interleaved versions. (It's still rather unfortunate to have two popular options)

- -
    -
  1. In the interleaved version, we pair consecutive elements (x0, x1),(x2,x3),… before applying the rotation matrix:
    2. - -Rd_{θ,m}x=\begin{pmatrix} -x_0 \\ -x_1 \\ -x_2 \\ -x_3 \\ -\vdots \\ -x_{d-2} \\ -x_{d-1} -\end{pmatrix} -\odot -\begin{pmatrix} -\cos mθ_0 \\ -\cos mθ_0 \\ -\cos mθ_1 \\ -\cos mθ_1 \\ -\vdots \\ -\cos mθ_{d/2-1} \\ -\cos mθ_{d/2-1} -\end{pmatrix} -+ -\begin{pmatrix} --x_1 \\ -x_0 \\ --x_3 \\ -x_2 \\ -\vdots \\ --x_{d-1} \\ -x_{d-2} -\end{pmatrix} -\odot -\begin{pmatrix} -\sin mθ_0 \\ -\sin mθ_0 \\ -\sin mθ_1 \\ -\sin mθ_1 \\ -\vdots \\ -\sin mθ_{d/2-1} \\ -\sin mθ_{d/2-1} -\end{pmatrix} - - -Rd_{θ,m}x=\begin{pmatrix} -x_0\cos mθ_0 - x_1\sin mθ_0 \\ -x_1\cos mθ_0 + x_0\sin mθ_0 \\ -x_2\cos mθ_1 - x_3\sin mθ_1 \\ -x_3\cos mθ_1 + x_2\sin mθ_1 \\ -\vdots \\ -x_{d-2}\cos mθ_{d/2-1} - x_{d-1}\sin mθ_{d/2-1} \\ -x_{d-1}\cos mθ_{d/2-1} + x_{d-2}\sin mθ_{d/2-1} -\end{pmatrix} - -
  2. In the non-interleaved version, we split the vector in two to pair elements as follows: (x0, xd/2),(x1,xd/2+1),… This is the implementation used in the transformers library:
    3. - -Rd_{θ,m}x=\begin{pmatrix} -x_0 \\ -x_1 \\ -\vdots \\ -x_{d/2-1} \\ -x_{d/2} \\ -x_{d/2+1} \\ -\vdots \\ -x_{d-1} -\end{pmatrix} -\odot -\begin{pmatrix} -\cos mθ_0 \\ -\cos mθ_1 \\ -\vdots \\ -\cos mθ_{d/2-1} \\ -\cos mθ_{0} \\ -\cos mθ_{1} \\ -\vdots \\ -\cos mθ_{d/2-1} -\end{pmatrix} -+ -\begin{pmatrix} --x_{d/2} \\ --x_{d/2+1} \\ -\vdots \\ --x_{d-1} \\ -x_{0} \\ -x_{1} \\ -\vdots \\ -x_{d/2-1} -\end{pmatrix} -\odot -\begin{pmatrix} -\sin mθ_0 \\ -\sin mθ_1 \\ -\vdots \\ -\sin mθ_{d/2-1} \\ -\sin mθ_{0} \\ -\sin mθ_{1} \\ -\vdots \\ -\sin mθ_{d/2-1} -\end{pmatrix} - - -Rd_{θ,m}x=\begin{pmatrix} -x_0\cos mθ_0 - x_{d/2}\sin mθ_0 \\ -x_1\cos mθ_1 - x_{d/2+1}\sin mθ_1 \\ -\vdots \\ -x_{d/2-1}\cos mθ_{d/2-1} - x_{d-1}\sin mθ_{d/2-1} \\ -x_{d/2}\cos mθ_0 + x_0\sin mθ_0 \\ -x_{d/2+1}\cos mθ_1 + x_0\sin mθ_1 \\ -\vdots \\ -x_{d-1}\cos mθ_{d/2-1} + x_{d-1}\sin mθ_{d/2-1} \\ -\end{pmatrix} - -

    The angle of rotation, θi is defined as follows, where d is the dimension of the attention head:

    - -θi = base-2(i-1)/d, i \in [1,2,...,d/2] - -

    How does this look? When moving the same distance, vectors in some dimensions rotate faster than vectors in other dimensions.

    - Rotation Speed -
- -

Attention (MHA, MQA, GQA)

- -

Optimized Operations

- -

Flash Attention 1&2&3

- -

Fused Kernels

- -

III – Training Recipe

- -

Batch Size

- -

Initialization + rescaling activations inside the model

- -

Numerical Precision

- -

FP16/BF16/FP8

- -

@Phuc Nguyen?

- -

Long Context Training

- -

Evaluation

- -

@Haojun Zhao

- -

Infini-Attention

- -

@Phuc Nguyen

- -

Ring Attention

- -

@Haojun Zhao

- -

RoPE scaling / Yarn

- -

@Haojun Zhao maybe?

- -

References

- -
    -
  • Harm's posts: +

    The rotation matrix has the following properties:

    -
    • -
  • Stas' guides: + + Rotation Matrix + +

    RoPE in 2D space

    + +

    Assuming q and k are 2D column vectors, we can show that:

    + + + = (R(θ_1)q)T (R(θ_2)k) = qTR(-θ_1)R(θ_2)k = + qTR(θ_2-θ_1)k = (R(θ_1-θ_2)q)Tk = + + +

    Therefore, if we define our position embedding like this: f(x, m) = R(mθ)x where R is a 2D + rotation matrix, we have q_m = R(mθ)q and k_n = R(nθ)k and then:

    + + + = = + + +

    We can see that a multiplication with a rotation matrix is exactly the positional encoding we were looking + for. The result of ⟨ 𝑞_𝑚 , 𝑘_𝑛 ⟩ only depends on q, k and m-n.

    + +

    Implementation

    + +

    In our case, our internal vectors (the activations in our model) have much more than two elements. Let's pair + elements to get 2D vectors and apply the 2D rotation operation on these pairs.

    + +

    There are combinatorially many ways we can pair elements but generally two options are the most popular for + implementing RoPE: we call them the interleaved and non-interleaved versions. (It's still + rather unfortunate to have two popular options)

    + +
      +
    1. In the interleaved version, we pair consecutive elements (x0, + x1),(x2,x3),… before applying the rotation matrix:
      2. + + Rd_{θ,m}x=\begin{pmatrix} + x_0 \\ + x_1 \\ + x_2 \\ + x_3 \\ + \vdots \\ + x_{d-2} \\ + x_{d-1} + \end{pmatrix} + \odot + \begin{pmatrix} + \cos mθ_0 \\ + \cos mθ_0 \\ + \cos mθ_1 \\ + \cos mθ_1 \\ + \vdots \\ + \cos mθ_{d/2-1} \\ + \cos mθ_{d/2-1} + \end{pmatrix} + + + \begin{pmatrix} + -x_1 \\ + x_0 \\ + -x_3 \\ + x_2 \\ + \vdots \\ + -x_{d-1} \\ + x_{d-2} + \end{pmatrix} + \odot + \begin{pmatrix} + \sin mθ_0 \\ + \sin mθ_0 \\ + \sin mθ_1 \\ + \sin mθ_1 \\ + \vdots \\ + \sin mθ_{d/2-1} \\ + \sin mθ_{d/2-1} + \end{pmatrix} + + + Rd_{θ,m}x=\begin{pmatrix} + x_0\cos mθ_0 - x_1\sin mθ_0 \\ + x_1\cos mθ_0 + x_0\sin mθ_0 \\ + x_2\cos mθ_1 - x_3\sin mθ_1 \\ + x_3\cos mθ_1 + x_2\sin mθ_1 \\ + \vdots \\ + x_{d-2}\cos mθ_{d/2-1} - x_{d-1}\sin mθ_{d/2-1} \\ + x_{d-1}\cos mθ_{d/2-1} + x_{d-2}\sin mθ_{d/2-1} + \end{pmatrix} + +
    2. In the non-interleaved version, we split the vector in two to pair elements as follows: + (x0, xd/2),(x1,xd/2+1),… This is the implementation + used in the transformers library: +
      3. + + Rd_{θ,m}x=\begin{pmatrix} + x_0 \\ + x_1 \\ + \vdots \\ + x_{d/2-1} \\ + x_{d/2} \\ + x_{d/2+1} \\ + \vdots \\ + x_{d-1} + \end{pmatrix} + \odot + \begin{pmatrix} + \cos mθ_0 \\ + \cos mθ_1 \\ + \vdots \\ + \cos mθ_{d/2-1} \\ + \cos mθ_{0} \\ + \cos mθ_{1} \\ + \vdots \\ + \cos mθ_{d/2-1} + \end{pmatrix} + + + \begin{pmatrix} + -x_{d/2} \\ + -x_{d/2+1} \\ + \vdots \\ + -x_{d-1} \\ + x_{0} \\ + x_{1} \\ + \vdots \\ + x_{d/2-1} + \end{pmatrix} + \odot + \begin{pmatrix} + \sin mθ_0 \\ + \sin mθ_1 \\ + \vdots \\ + \sin mθ_{d/2-1} \\ + \sin mθ_{0} \\ + \sin mθ_{1} \\ + \vdots \\ + \sin mθ_{d/2-1} + \end{pmatrix} + + + Rd_{θ,m}x=\begin{pmatrix} + x_0\cos mθ_0 - x_{d/2}\sin mθ_0 \\ + x_1\cos mθ_1 - x_{d/2+1}\sin mθ_1 \\ + \vdots \\ + x_{d/2-1}\cos mθ_{d/2-1} - x_{d-1}\sin mθ_{d/2-1} \\ + x_{d/2}\cos mθ_0 + x_0\sin mθ_0 \\ + x_{d/2+1}\cos mθ_1 + x_0\sin mθ_1 \\ + \vdots \\ + x_{d-1}\cos mθ_{d/2-1} + x_{d-1}\sin mθ_{d/2-1} \\ + \end{pmatrix} + +

      The angle of rotation, θi is defined as follows, where d is the dimension of + the attention head:

      + + θi = base-2(i-1)/d, i \in [1,2,...,d/2] + +

      How does this look? When moving the same distance, vectors in some dimensions rotate faster than vectors + in other dimensions.

      + Rotation Speed +
    + +

    Attention (MHA, MQA, GQA)

    + +

    Optimized Operations

    + +

    Flash Attention 1&2&3

    + +

    Fused Kernels

    + +

    III – Training Recipe

    + +

    Batch Size

    + +

    Initialization + rescaling activations inside the model

    + +

    Numerical Precision

    + +

    FP16/BF16/FP8

    + +

    @Phuc Nguyen?

    + +

    Long Context Training

    + +

    Evaluation

    + +

    @Haojun Zhao

    + +

    Infini-Attention

    + +

    @Phuc Nguyen

    + +

    Ring Attention

    + +

    @Haojun Zhao

    + +

    RoPE scaling / Yarn

    + +

    @Haojun Zhao maybe?

    + +

    References

    + -
    • -
  • data parallel: https://siboehm.com/articles/22/data-parallel-training
    • -
  • ZeRO: https://arxiv.org/abs/1910.02054
    • -
  • TP/SP + Selective Recomputation: https://arxiv.org/abs/2205.05198
    • -
-

Conclusion and looking forward

-

Through our open science efforts we hope to keep shining a light on the black box that is the training of high performance large language models as well as to give every model trainer the ability to create state-of-the-art LLMs. We are excited to continue iterating on FineWeb and to release increasingly better filtered subsets of web data, in a fully open and reproducible manner.

-

In the short term, we are looking forward to applying the learnings from (English) FineWeb to other languages. While English currently dominates the LLM landscape, we believe that making high quality web data in other languages as accessible as possible would be incredibly impactful.

-

In a nutshell: the future is bright and exciting for studying the science of creating datasets at scale and in the open 🤗.

- - - - - - -

Citation

-

For attribution in academic contexts, please cite this work as

- 
Penedo, et al., "The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale", 2024.
-

BibTeX citation

- 
@misc{penedo2024finewebdatasetsdecantingweb,
+        
Conclusion and looking forward

+        
Through our open science efforts we hope to keep shining a light on the black box that is the training of
+            high performance large language models as well as to give every model trainer the ability to create
+            state-of-the-art LLMs. We are excited to continue iterating on FineWeb and to release increasingly better
+            filtered subsets of web data, in a fully open and reproducible manner.

+        
In the short term, we are looking forward to applying the learnings from (English) FineWeb to other
+            languages. While English currently dominates the LLM landscape, we believe that making high quality web data
+            in other languages as accessible as possible would be incredibly impactful.

+        
In a nutshell: the future is bright and exciting for studying the science of creating datasets at scale and
+            in the open 🤗.

+    
+
+    
+        
+        
+
+        
Citation

+        
For attribution in academic contexts, please cite this work as

+        
Penedo, et al., "The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale", 2024.

+        
BibTeX citation

+        
@misc{penedo2024finewebdatasetsdecantingweb,
       title={The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale},
       author={Guilherme Penedo and Hynek Kydlíček and Loubna Ben allal and Anton Lozhkov and Margaret Mitchell and Colin Raffel and Leandro Von Werra and Thomas Wolf},
       year={2024},
@@ -989,71 +1255,72 @@ x_{d-1}\cos mθ_{d/2-1} + x_{d-1}\sin mθ_{d/2-1}  \\
       primaryClass={cs.CL}
       url={https://arxiv.org/abs/2406.17557},
 }

-
-
-
+            });
+        }
+    
 
 
+
 
\ No newline at end of file