š Ultra-Guide to Scaling LLM training
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
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.
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:
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.
| 1B (1) | 7B | 70B | 340B (2) | 400B (3) | |
|---|---|---|---|---|---|
| N Layers | 24 | 32 | 80 | 96 | 126 |
| N Heads | 32 | 32 | 64 | 96 | 128 |
| Dimension | 2048 | 4096 | 8192 | 18432 | 16384 |
(1) FineWeb ablation models
(2) Nemotron-340B architecture (without GQA)
(3) Llama-400B, ffn dim = 1.2 hidden dim (without GQA)
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
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:
As weāll see in the future, these steps may be repeated or intertwined but for now weāll start simple:
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:
(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.
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:
import torch; torch.ones((1, 1)).to("cuda") and then checking the GPU memory with nvidia-smiScaling 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:
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):
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:
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 parameters | Memory requirements |
|---|---|
| 1B | 16 GB |
| 7B | 112 GB |
| 70B | 1120 GB |
| 405B | 6480 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:
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:
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.
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:
We can select these key activations according to several strategies and modern frameworks usually choose among the following three strategies:
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.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:
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:
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!
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!
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:
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.
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:
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!
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:
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:
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 comprise 2 successive MLPs with a non-linearity in-between. Here is the first part of it:
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:
In row schema:
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:
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:
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
def example_column_row_linear():
# torchrun --nproc_per_node=2 tp_all_reduce.py
group = dist.distributed_c10d._get_default_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)
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()
# 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)
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)
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)
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"]))
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).
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.
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
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.
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.
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]
Even though TP-SP mode helps reduce the memory used by activation values, it has two main drawbacks:
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.
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:
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:
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.
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:
The rotation matrix has the following properties:
RoPE in 2D space
Assuming q and k are 2D column vectors, we can show that:
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)
transformers library:The angle of rotation, Īøi is defined as follows, where d is the dimension of the attention head:
How does this look? When moving the same distance, vectors in some dimensions rotate faster than vectors in other dimensions.
@Phuc Nguyen?
@Haojun Zhao
@Phuc Nguyen
@Haojun Zhao
@Haojun Zhao maybe?
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 š¤.
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},
eprint={2406.17557},
archivePrefix={arXiv},
primaryClass={cs.CL}
url={https://arxiv.org/abs/2406.17557},
}