Training Transformers

Chapter Overview

Training transformers requires specialized techniques beyond standard optimization. This chapter provides comprehensive coverage of transformer training procedures, from loss functions and backpropagation through the architecture to optimization algorithms, learning rate schedules, and hardware-efficient training strategies. We examine why transformers need warmup, how mixed precision training reduces memory consumption, when to use gradient accumulation and checkpointing, and how distributed training enables models that exceed single-GPU capacity. Throughout, we provide detailed hardware analysis, memory calculations, and practical guidance drawn from training state-of-the-art models like BERT, GPT-2, and GPT-3.

Learning Objectives

1. Understand training objectives and loss functions for different transformer architectures
2. Analyze gradient flow and backpropagation through transformer layers
3. Implement optimization algorithms (Adam, AdamW, LAMB) with appropriate hyperparameters
4. Apply learning rate schedules with warmup and decay
5. Use mixed precision training to reduce memory and accelerate training
6. Apply gradient accumulation and checkpointing for memory-constrained scenarios
7. Understand distributed training strategies for large-scale models
8. Select appropriate batch sizes and sequence lengths based on hardware constraints
9. Apply regularization techniques to prevent overfitting
10. Estimate training time and costs for transformer models

Training Objectives and Loss Functions

The training objective fundamentally shapes how a transformer learns and what capabilities it develops. Different transformer architectures employ distinct training objectives tailored to their intended use cases, from masked language modeling in BERT to causal language modeling in GPT to sequence-to-sequence learning in T5. Understanding these objectives in depth—including their mathematical formulations, computational requirements, and practical implications—is essential for training transformers effectively.

Masked Language Modeling

Masked language modeling, introduced by BERT, trains the model to predict randomly masked tokens based on bidirectional context. This objective enables the model to learn rich representations that capture relationships in both directions, making it particularly effective for tasks requiring understanding of complete sentences or documents.

The masking strategy is more sophisticated than simply replacing tokens with a special [MASK] symbol. BERT's approach selects 15\% of tokens for prediction, but handles them in three different ways: 80\% are replaced with [MASK], 10\% are replaced with random tokens from the vocabulary, and 10\% are left unchanged. This strategy prevents the model from simply memorizing that [MASK] tokens need prediction and forces it to maintain representations for all tokens, since any token might need to be predicted. The random token replacement encourages the model to use context to correct errors, while leaving some tokens unchanged helps the model learn that not all tokens are corrupted.

The loss function for masked language modeling is cross-entropy computed only over the masked positions. For a sequence $\mathbf{x} = (x_1, \ldots, x_n)$ with masked positions $M \subseteq \{1, \ldots, n\}$, the loss is:

The computational and memory implications of this loss are significant. For vocabulary size $V = 30{,}000$, sequence length $n = 512$, and batch size $B = 32$, the output logits tensor has shape $\R^{32 \times 512 \times 30000}$, requiring $32 \times 512 \times 30{,}000 \times 4 = 1{,}966{,}080{,}000$ bytes, or approximately 1.97 GB of memory just for the logits in FP32. This massive memory footprint explains why the output projection and softmax computation often become bottlenecks during training. The memory requirement can be reduced by computing the loss in chunks (processing subsets of positions at a time) or by using mixed precision training where logits are computed in FP16, though care must be taken to maintain numerical stability in the softmax operation.

In practice, BERT-base masks approximately 77 tokens per sequence (15\% of 512), so the loss is computed over $32 \times 77 = 2{,}464$ predictions per batch. The cross-entropy computation requires exponentiating 30,000 logits for each prediction to compute the softmax denominator, then taking the logarithm of the target class probability. Modern implementations optimize this by fusing the softmax and cross-entropy operations and by using numerically stable implementations that subtract the maximum logit before exponentiation to prevent overflow.

Causal Language Modeling

Causal language modeling, used in GPT and other decoder-only models, trains the model to predict the next token given all previous tokens. Unlike masked language modeling, which uses bidirectional context, causal language modeling uses only left-to-right context, enforced through causal attention masks that prevent positions from attending to future positions.

The training objective is to maximize the likelihood of each token given its preceding context. For a sequence $\mathbf{x} = (x_1, \ldots, x_n)$, the loss is:

This formulation means that every position in the sequence contributes to the loss, unlike masked language modeling where only 15\% of positions contribute. For a batch of 32 sequences of length 512, we compute loss over $32 \times 512 = 16{,}384$ predictions, compared to only 2,464 for BERT's masked language modeling. This makes causal language modeling more sample-efficient in terms of predictions per sequence, though the unidirectional context may be less informative than bidirectional context for some tasks.

A crucial distinction exists between training and inference for causal language models. During training, we use teacher forcing: the model receives the ground-truth previous tokens as input, even if it would have predicted different tokens. This enables parallel computation of the loss across all positions in a sequence, since we can compute $P(x_i | x_1, \ldots, x_{i-1})$ for all $i$ simultaneously using causal masking. During inference, however, generation is autoregressive: the model generates one token at a time, using its own predictions as input for subsequent positions. This sequential generation process is much slower than parallel training, which motivates optimizations like KV caching (discussed in Chapter 12).

The memory requirements for causal language modeling are similar to masked language modeling: the output logits tensor for batch size 32, sequence length 512, and vocabulary size 50,257 (GPT-2's vocabulary) requires $32 \times 512 \times 50{,}257 \times 4 = 3{,}296{,}019{,}456$ bytes, or approximately 3.3 GB in FP32. However, since we compute loss over all positions rather than just 15\%, the gradient computation is more expensive. The backward pass through the output projection receives gradients from all 16,384 predictions rather than just 2,464, increasing the gradient computation cost proportionally.

Sequence-to-Sequence Training

Sequence-to-sequence models like T5 and BART use encoder-decoder architectures where the encoder processes the input sequence bidirectionally and the decoder generates the output sequence autoregressively. The training objective combines aspects of both masked and causal language modeling: the encoder can use bidirectional attention over the input, while the decoder uses causal attention over the output sequence and cross-attention to the encoder's representations.

The loss function for sequence-to-sequence training is computed over the target sequence. For input sequence $\mathbf{x} = (x_1, \ldots, x_n)$ and target sequence $\mathbf{y} = (y_1, \ldots, y_m)$:

Like causal language modeling, sequence-to-sequence training uses teacher forcing during training: the decoder receives the ground-truth previous target tokens as input, enabling parallel computation of the loss. This differs from inference, where the decoder must generate tokens sequentially using its own predictions.

The memory requirements for sequence-to-sequence models are higher than encoder-only or decoder-only models because both encoder and decoder activations must be stored. For T5-base with input length 512, target length 512, and batch size 32, we must store encoder activations ($32 \times 512 \times 768$ per layer), decoder activations ($32 \times 512 \times 768$ per layer), and cross-attention activations ($32 \times 12 \times 512 \times 512$ for attention matrices between decoder and encoder). The total activation memory is roughly 1.5-2× that of an encoder-only model of the same size.

Different sequence-to-sequence models use different input corruption strategies. T5 uses span corruption, where contiguous spans of tokens are replaced with sentinel tokens and the model must predict the original spans. BART uses a variety of corruption strategies including token masking, token deletion, sentence permutation, and document rotation. These diverse corruption strategies help the model learn robust representations that generalize across different types of noise and transformations.

Backpropagation Through Transformers

Understanding how gradients flow through the transformer architecture is essential for diagnosing training issues, designing better architectures, and implementing custom training procedures. The transformer's combination of attention mechanisms, residual connections, layer normalization, and feed-forward networks creates a complex gradient flow pattern that differs fundamentally from simpler architectures like MLPs or CNNs.

Gradient Flow Analysis

Backpropagation through a transformer begins at the output and flows backward through each component. For a language modeling task, the loss $L$ is computed from the output logits, and we must compute gradients with respect to all parameters in the model. The gradient flow follows the reverse path of the forward computation, with each operation contributing its Jacobian to the chain rule.

The output projection layer maps the final transformer layer's output to vocabulary logits. For output $\mathbf{h}_n \in \R^{d_{\text{model}}}$ at position $n$ and output weight matrix $\mW^{\text{out}} \in \R^{d_{\text{model}} \times V}$, the logits are $\mathbf{z}_n = \mW^{\text{out}\transpose} \mathbf{h}_n$. The gradient of the loss with respect to the output weights is:

The gradient with respect to the output representations is:

This gradient then flows backward through each transformer layer. Within a layer, the gradient must flow through the feed-forward network, the second residual connection and layer normalization, the attention mechanism, and the first residual connection and layer normalization.

Gradients Through Residual Connections

Residual connections are crucial for training deep transformers because they provide "gradient highways" that allow gradients to flow directly through many layers without vanishing. Consider a residual block with function $F$:

The gradient with respect to the input is:

The first term $\frac{\partial L}{\partial \mathbf{y}}$ is the direct gradient path that bypasses the function $F$ entirely. This ensures that even if $\frac{\partial F(\mathbf{x})}{\partial \mathbf{x}}$ becomes very small (vanishing gradients) or very large (exploding gradients), the gradient $\frac{\partial L}{\partial \mathbf{x}}$ still receives the direct contribution $\frac{\partial L}{\partial \mathbf{y}}$. This is why transformers can be trained with many layers (BERT-large has 24 layers, GPT-3 has 96 layers) without suffering from vanishing gradients that plagued early deep networks.

For a transformer with $L$ layers, the gradient from the output to the input has $2^L$ paths through the network: at each layer, the gradient can either flow through the residual connection (direct path) or through the attention/FFN (indirect path). This exponential number of paths creates a rich gradient flow that helps training, though in practice most gradient flows through the shorter paths that use more residual connections.

Gradients Through Layer Normalization

Layer normalization normalizes activations across the feature dimension, computing mean and variance for each position independently. For input $\mathbf{x} \in \R^{d}$, layer normalization computes:

The gradient computation for layer normalization is complex because the normalization couples all dimensions: changing one input element affects the mean and variance, which affects all output elements. The gradient with respect to the input is:

This gradient has three terms: the direct gradient scaled by the normalization factor, a mean-centering term, and a variance-correction term. The complexity of this gradient is why layer normalization is sometimes replaced with simpler alternatives like RMSNorm in some recent models, though layer normalization generally provides better training stability.

The learned parameters $\gamma$ and $\beta$ have simple gradients:

Layer normalization helps gradient flow by preventing activations from becoming too large or too small, which would cause gradients to vanish or explode. By maintaining normalized activations throughout the network, layer normalization ensures that gradients remain in a reasonable range, facilitating stable training.

Gradients Through Attention

The attention mechanism involves several matrix multiplications and a softmax operation, each contributing to the gradient computation. For self-attention with queries $\mQ$, keys $\mK$, and values $\mV$:

Working backward, the gradient with respect to the values is:

The gradient with respect to the attention matrix is:

This has shape $(n \times d_v) \times (d_v \times n) = (n \times n)$, matching the attention matrix shape.

The gradient must then flow through the softmax operation. For softmax output $\mathbf{a} = \text{softmax}(\mathbf{s})$, the Jacobian is:

This computation must be performed for each row of the attention matrix independently, since softmax is applied row-wise.

Finally, gradients flow to the query and key projections. The gradient with respect to queries is:

These gradients then flow through the projection matrices $\mW^Q$, $\mW^K$, and $\mW^V$. For the query projection $\mQ = \mX \mW^Q$:

This gradient has shape $(d_{\text{model}} \times n) \times (n \times d_k) = (d_{\text{model}} \times d_k)$, matching $\mW^Q$. For BERT-base with $d_{\text{model}} = 768$ and $d_k = 64$, this requires $768 \times 64 \times 4 = 196{,}608$ bytes (197 KB) per head, or $12 \times 197 = 2.4$ MB for all 12 heads.

Gradients Through Feed-Forward Networks

The feed-forward network consists of two linear transformations with a non-linear activation (typically GELU) in between:

The gradient with respect to the second layer weights is:

The gradient flows through the GELU activation. GELU is defined as:

Finally, the gradient with respect to the first layer weights is:

This has shape $(d_{\text{model}} \times d_{ff}) = (768 \times 3072)$, also requiring 9.4 MB in FP32.

Computational Cost of Backpropagation

The backward pass through a transformer requires approximately twice the FLOPs of the forward pass. This factor of two arises because each matrix multiplication $\mathbf{Y} = \mathbf{X} \mW$ in the forward pass requires two matrix multiplications in the backward pass: $\frac{\partial L}{\partial \mW} = \mathbf{X}\transpose \frac{\partial L}{\partial \mathbf{Y}}$ and $\frac{\partial L}{\partial \mathbf{X}} = \frac{\partial L}{\partial \mathbf{Y}} \mW\transpose$. Each of these backward matrix multiplications has similar computational cost to the forward multiplication.

For BERT-base with 96.6 GFLOPs per forward pass, the backward pass requires approximately $2 \times 96.6 = 193.2$ GFLOPs. A complete training step (forward pass + backward pass) thus requires approximately $96.6 + 193.2 = 289.8$ GFLOPs, or roughly three times the forward pass cost. This 3× factor is a useful rule of thumb for estimating training costs from inference costs.

The memory requirements for backpropagation are substantial because all intermediate activations from the forward pass must be stored to compute gradients. For BERT-base with batch size 32 and sequence length 512, the activations require approximately 12 GB as analyzed in Chapter 12. This activation memory often dominates the total memory consumption during training, which motivates techniques like gradient checkpointing that trade computation for memory by recomputing activations during the backward pass.

Optimization Algorithms

The choice of optimization algorithm significantly impacts transformer training dynamics, convergence speed, and final model quality. While stochastic gradient descent (SGD) with momentum works well for many deep learning tasks, transformers benefit particularly from adaptive learning rate methods that adjust the learning rate for each parameter based on gradient statistics. The Adam family of optimizers has become the de facto standard for transformer training, with variants like AdamW and LAMB addressing specific challenges in large-scale training.

Adam Optimizer

Adam (Adaptive Moment Estimation) maintains exponential moving averages of both the gradient (first moment) and the squared gradient (second moment) for each parameter. These statistics enable adaptive per-parameter learning rates that automatically adjust based on the gradient history, helping with the varying scales of gradients across different layers and components of the transformer.

The Adam algorithm maintains two state vectors for each parameter $\mathbf{w}$: the first moment $\mathbf{m}$ (exponential moving average of gradients) and the second moment $\mathbf{v}$ (exponential moving average of squared gradients). At each training step $t$ with gradient $\mathbf{g}_t$:

where $\beta_1$ and $\beta_2$ are decay rates (typically $\beta_1 = 0.9$ and $\beta_2 = 0.999$). The squared gradient $\mathbf{g}_t^2$ is computed element-wise.

Because $\mathbf{m}$ and $\mathbf{v}$ are initialized to zero, they are biased toward zero, especially in early training steps. Adam corrects this bias by computing bias-corrected estimates:

The parameter update is then:

where $\eta$ is the learning rate and $\epsilon$ is a small constant (typically $10^{-8}$) for numerical stability.

The adaptive learning rate $\frac{\eta}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon}$ is larger for parameters with small historical gradients and smaller for parameters with large historical gradients. This adaptation is particularly beneficial for transformers because different components have vastly different gradient scales. Embedding layers, which are updated sparsely (only for tokens present in the batch), benefit from larger effective learning rates, while frequently updated parameters in the attention and FFN layers benefit from smaller effective learning rates that prevent overshooting.

The memory requirements for Adam are substantial: for each parameter, we must store the parameter itself, the gradient, the first moment, and the second moment. For a model with $P$ parameters in FP32, Adam requires:

• Parameters: $P \times 4$ bytes
• Gradients: $P \times 4$ bytes
• First moments: $P \times 4$ bytes
• Second moments: $P \times 4$ bytes
• Total: $16P$ bytes

For BERT-base with 110 million parameters, Adam requires $110{,}000{,}000 \times 16 = 1{,}760{,}000{,}000$ bytes, or 1.76 GB, just for the optimizer state. This is four times the memory required for the parameters alone, and this overhead grows linearly with model size. For GPT-3 with 175 billion parameters, Adam would require $175{,}000{,}000{,}000 \times 16 = 2{,}800$ GB just for parameters and optimizer states, necessitating distributed training strategies that shard the optimizer state across multiple GPUs.

AdamW: Decoupled Weight Decay

AdamW modifies Adam by decoupling weight decay from the gradient-based update. In standard Adam with L2 regularization, the weight decay is incorporated into the gradient: $\mathbf{g}_t = \nabla L(\mathbf{w}_t) + \lambda \mathbf{w}_t$, where $\lambda$ is the regularization coefficient. This means the weight decay is affected by the adaptive learning rate, which can lead to unexpected behavior.

AdamW instead applies weight decay directly to the parameters after the adaptive update:

This decoupling means that weight decay acts as a true regularizer, shrinking parameters toward zero at a rate proportional to the learning rate, independent of the gradient statistics. In practice, this leads to better generalization, particularly for transformers where different parameters have very different gradient scales.

The typical weight decay coefficient for transformer training is $\lambda = 0.01$. However, weight decay is usually not applied to all parameters. Biases and layer normalization parameters (the scale $\gamma$ and shift $\beta$ parameters) are typically excluded from weight decay, as regularizing these parameters can hurt performance. The exclusion is implemented by maintaining separate parameter groups in the optimizer, with different weight decay settings for each group.

AdamW has become the standard optimizer for training transformers, used in BERT, GPT-2, GPT-3, T5, and most other modern models. The improved generalization from decoupled weight decay often allows training with higher learning rates, which can accelerate convergence. The memory requirements are identical to Adam: $16P$ bytes for a model with $P$ parameters in FP32.

LAMB: Large Batch Training

LAMB (Layer-wise Adaptive Moments optimizer for Batch training) extends Adam to enable training with very large batch sizes, up to 64,000 or more. Large batch training is desirable because it improves hardware utilization and reduces training time by processing more examples in parallel, but naive scaling of the batch size often hurts convergence and final model quality.

The key insight of LAMB is to compute layer-wise learning rates that adapt based on the ratio of parameter norm to gradient norm within each layer. For layer $l$ with parameters $\mathbf{w}^{(l)}$ and Adam update $\mathbf{u}^{(l)} = \frac{\hat{\mathbf{m}}^{(l)}}{\sqrt{\hat{\mathbf{v}}^{(l)}} + \epsilon} + \lambda \mathbf{w}^{(l)}$, LAMB computes:

The parameter update is then:

This layer-wise adaptation ensures that the update magnitude is proportional to the parameter magnitude within each layer, preventing some layers from being updated too aggressively while others are updated too conservatively. This is particularly important for large batch training because large batches produce more accurate gradient estimates, which can lead to overly aggressive updates without proper scaling.

LAMB enabled training BERT-large to the same accuracy as the original paper in just 76 minutes using a batch size of 65,536 on 1,024 TPU v3 chips, compared to several days with standard batch sizes. The ability to use such large batches dramatically reduces training time for large-scale models, though it requires access to substantial computational resources to realize the benefits.

The memory requirements for LAMB are similar to Adam and AdamW: $16P$ bytes for a model with $P$ parameters in FP32. The additional computation for layer-wise norm calculations is negligible compared to the forward and backward passes.

Optimizer Memory Comparison

Different optimizers have different memory footprints, which can be a critical consideration for large models:

• SGD (no momentum): $8P$ bytes (parameters + gradients in FP32)
• SGD with momentum: $12P$ bytes (parameters + gradients + momentum in FP32)
• Adam/AdamW/LAMB: $16P$ bytes (parameters + gradients + first moment + second moment in FP32)

For BERT-base with 110 million parameters:

• SGD: $110\text{M} \times 8 = 880$ MB
• SGD with momentum: $110\text{M} \times 12 = 1{,}320$ MB
• Adam/AdamW/LAMB: $110\text{M} \times 16 = 1{,}760$ MB

The additional memory overhead of Adam-family optimizers (880 MB compared to SGD) is usually worthwhile because the adaptive learning rates lead to faster convergence and better final performance. However, for very large models where memory is at a premium, techniques like ZeRO (Zero Redundancy Optimizer) can shard the optimizer state across multiple GPUs to reduce per-GPU memory requirements.

Learning Rate Schedules

Learning rate schedules are critical for transformer training, perhaps more so than for other architectures. Transformers are sensitive to the learning rate, and using a constant learning rate throughout training typically leads to poor results. The standard approach combines a warmup phase, where the learning rate increases from zero to a maximum value, with a decay phase, where the learning rate gradually decreases. This schedule helps stabilize early training and enables continued improvement in later training.

The Necessity of Warmup

Learning rate warmup is essential for stable transformer training. Without warmup, using the full learning rate from the beginning often causes training to diverge or get stuck in poor local minima. The instability arises from the interaction between large initial gradients and Adam's adaptive learning rates.

In the first few training steps, Adam's second moment estimates $\mathbf{v}$ are very small because they are initialized to zero and have not yet accumulated gradient statistics. This means the effective learning rate $\frac{\eta}{\sqrt{\mathbf{v}} + \epsilon}$ is very large, potentially much larger than the nominal learning rate $\eta$. When combined with large gradients that are common early in training (when the model's predictions are random and the loss is high), these large effective learning rates can cause parameter updates that are far too aggressive, leading to numerical instability or divergence.

Warmup solves this problem by starting with a very small learning rate and gradually increasing it over the first $W$ steps (typically 10\% of total training steps). During warmup, the learning rate at step $t$ is:

This linear increase gives Adam's moment estimates time to accumulate meaningful statistics while preventing overly aggressive updates. By the time the learning rate reaches its maximum value $\eta_{\max}$, the optimizer has stabilized and can handle the full learning rate safely.

The warmup period also serves another purpose: it allows the model to learn basic patterns before attempting more complex optimization. In the first few steps, the model learns simple statistics like token frequencies and basic co-occurrence patterns. These foundational patterns provide a stable base for learning more complex relationships later in training.

Warmup Plus Linear Decay

The warmup plus linear decay schedule, used in BERT and many other models, combines linear warmup with linear decay to zero. For total training steps $T$ and warmup steps $W$:

The decay phase gradually reduces the learning rate to zero over the remaining training steps. This decay is beneficial because it allows the model to make large updates early in training when far from a good solution, then make progressively smaller updates as it approaches a good solution. The smaller learning rate in late training helps the model settle into a sharper minimum, which often generalizes better.

For BERT-base, the typical configuration is $\eta_{\max} = 1 \times 10^{-4}$, $W = 10{,}000$ steps, and $T = 1{,}000{,}000$ steps. This means the learning rate increases linearly from 0 to $10^{-4}$ over the first 10,000 steps (1\% of training), then decreases linearly from $10^{-4}$ to 0 over the remaining 990,000 steps. The warmup period is relatively short, but it is crucial for stable training.

Different models use different maximum learning rates based on their size and architecture. GPT-2 uses $\eta_{\max} = 2.5 \times 10^{-4}$, slightly higher than BERT. GPT-3 uses $\eta_{\max} = 6 \times 10^{-5}$, lower than smaller models, reflecting the general trend that larger models require smaller learning rates for stable training. The warmup period for GPT-3 is 375 million tokens, which corresponds to a different number of steps depending on the batch size and sequence length.

Inverse Square Root Decay

The original "Attention is All You Need" paper used a different schedule that combines warmup with inverse square root decay:

This schedule has two phases. During warmup ($t \leq W$), the learning rate increases linearly:

After warmup ($t > W$), the learning rate decays as the inverse square root of the step number:

The inverse square root decay is slower than linear decay, maintaining a higher learning rate for longer. This can be beneficial for very long training runs where continued exploration is desirable. The original Transformer used $W = 4{,}000$ warmup steps and $d_{\text{model}} = 512$, giving a peak learning rate of $512^{-0.5} \cdot 4000^{-0.5} \approx 0.00070$.

The inverse square root schedule is less commonly used than linear decay in modern transformers, but it remains popular for some applications, particularly in machine translation where the original Transformer architecture is still widely used.

Cosine Annealing

Cosine annealing provides a smooth decay curve that starts slowly, accelerates in the middle, and slows again near the end. After warmup, the learning rate follows a cosine curve:

where $\eta_{\min}$ is the minimum learning rate (often 0 or $0.1 \eta_{\max}$). At the start of decay ($t = W$), the cosine term is $\cos(0) = 1$, giving $\eta_W = \eta_{\max}$. At the end of training ($t = T$), the cosine term is $\cos(\pi) = -1$, giving $\eta_T = \eta_{\min}$.

The smooth decay of cosine annealing can provide better final performance than linear decay, particularly for tasks where the model benefits from extended fine-tuning at low learning rates. The slower initial decay allows the model to continue exploring, while the accelerated decay in the middle helps the model converge, and the slow final decay allows careful refinement.

Cosine annealing is popular in computer vision (where it was originally developed) and has been adopted for some transformer training, particularly in vision transformers and multimodal models. However, linear decay remains more common for language models.

Mixed Precision Training

Mixed precision training is one of the most impactful optimizations for transformer training, reducing memory consumption and accelerating computation by leveraging lower-precision arithmetic. The technique uses 16-bit floating point (FP16 or BF16) for most operations while maintaining 32-bit floating point (FP32) master weights for numerical stability. This combination achieves substantial speedups on modern hardware while preserving training dynamics and final model quality.

FP16 Training Algorithm

Mixed precision training with FP16 maintains two copies of the model parameters: an FP16 copy used for forward and backward passes, and an FP32 master copy used for parameter updates. The algorithm proceeds as follows:

1. Forward pass: Convert FP32 master weights to FP16, perform all forward computations in FP16, producing FP16 activations
2. Loss computation: Compute loss in FP16, then scale the loss by a large factor $S$ (typically 1024 or dynamically adjusted)
3. Backward pass: Compute gradients in FP16 using the scaled loss, producing FP16 gradients that are also scaled by $S$
4. Gradient unscaling: Divide FP16 gradients by $S$ to recover the true gradient scale
5. Gradient conversion: Convert unscaled FP16 gradients to FP32
6. Parameter update: Update FP32 master weights using FP32 gradients and the optimizer
7. Repeat: Copy updated FP32 weights to FP16 for the next iteration

The loss scaling step is crucial for preventing gradient underflow. FP16 has a much smaller representable range than FP32: the smallest positive normal number in FP16 is approximately $6 \times 10^{-5}$, compared to $1.2 \times 10^{-38}$ in FP32. Gradients in deep networks are often very small, particularly in later layers or after many training steps. Without scaling, these small gradients would underflow to zero in FP16, preventing the corresponding parameters from being updated.

By scaling the loss by a factor $S$ before backpropagation, all gradients are also scaled by $S$ (due to the chain rule). This shifts the gradient values into the representable range of FP16. After the backward pass, we divide by $S$ to recover the true gradient values. The scaling and unscaling operations are mathematically equivalent to computing gradients in FP32, but they allow the actual gradient computation to occur in FP16, leveraging faster FP16 hardware.

The scaling factor $S$ can be fixed (typically 1024 or 2048) or dynamic. Dynamic loss scaling starts with a large scaling factor and reduces it if gradient overflow is detected (indicated by NaN or Inf values in the gradients). If training proceeds without overflow for a certain number of steps, the scaling factor is increased. This adaptive approach maximizes the use of FP16's range while preventing overflow.

Memory Savings

Mixed precision training reduces memory consumption primarily through smaller activations. The memory breakdown for mixed precision training is:

• FP16 parameters (forward/backward): $2P$ bytes
• FP32 master parameters: $4P$ bytes
• FP32 gradients: $4P$ bytes
• FP32 optimizer states (Adam): $8P$ bytes (first and second moments)
• FP16 activations: $A/2$ bytes (where $A$ is FP32 activation memory)

The total is $18P + A/2$ bytes, compared to $16P + A$ bytes for FP32 training. Surprisingly, mixed precision uses slightly more memory for parameters and optimizer states ($18P$ vs $16P$) because we maintain both FP16 and FP32 copies of the parameters. However, the activation memory is halved ($A/2$ vs $A$), and since activations typically dominate memory consumption, mixed precision usually provides substantial overall savings.

For BERT-base with 110 million parameters, batch size 32, and sequence length 512:

FP32 training:

Mixed precision training:

Mixed precision saves $13.8 - 8.0 = 5.8$ GB, a 42\% reduction. This memory saving enables larger batch sizes or longer sequences on the same hardware, directly improving training efficiency.

Hardware Acceleration

Modern GPUs provide dedicated hardware for accelerated FP16 computation. NVIDIA's Tensor Cores, available on Volta (V100), Turing (RTX 20xx), Ampere (A100, RTX 30xx), and newer architectures, can perform FP16 matrix multiplications at twice the throughput of FP32 operations.

For the NVIDIA A100 GPU:

• FP32 performance: 156 TFLOPS (teraflops)
• FP16 performance (Tensor Cores): 312 TFLOPS
• Theoretical speedup: 2×

In practice, the speedup is typically 1.5-1.8× rather than the full 2× because:

• Not all operations benefit from FP16 (e.g., layer normalization, softmax, and other element-wise operations may still run in FP32 for numerical stability)
• Memory bandwidth limitations can bottleneck performance, particularly for small batch sizes
• Overhead from data type conversions and loss scaling
• Non-matrix operations (activations, normalizations) don't use Tensor Cores

For BERT-base training on an A100 GPU, mixed precision typically provides a 1.6× speedup, reducing training time from approximately 4 days to 2.5 days on the same hardware. This speedup, combined with the memory savings that enable larger batch sizes, makes mixed precision training essential for efficient transformer training.

BF16: An Alternative to FP16

BF16 (bfloat16) is an alternative 16-bit format that maintains the same exponent range as FP32 (8 bits) while reducing the mantissa precision (7 bits, compared to 10 bits in FP16). This design choice provides better numerical stability than FP16 at the cost of slightly lower precision.

The key advantage of BF16 is that it can represent the same range of values as FP32, from approximately $10^{-38}$ to $10^{38}$. This eliminates the need for loss scaling because gradients are unlikely to underflow in BF16's range. The training algorithm simplifies to:

1. Forward pass in BF16
2. Loss computation in BF16 (no scaling needed)
3. Backward pass in BF16
4. Convert BF16 gradients to FP32
5. Update FP32 master weights

BF16 is supported on Google's TPUs (v2, v3, v4), NVIDIA A100 GPUs, and newer hardware. For transformers, BF16 often provides similar or slightly better results than FP16 with less tuning required, since the loss scaling factor doesn't need to be adjusted. However, FP16 remains more widely supported across different hardware platforms.

The memory savings and computational speedups for BF16 are similar to FP16: activations are halved, and Tensor Cores provide approximately 2× theoretical speedup (1.5-1.8× in practice). The choice between FP16 and BF16 often depends on hardware availability and whether loss scaling tuning is problematic for a particular training setup.

Gradient Accumulation

Gradient accumulation is a technique for achieving large effective batch sizes when GPU memory limits the actual batch size that can be processed in a single forward-backward pass. The technique accumulates gradients over multiple mini-batches before updating parameters, mathematically equivalent to training with a larger batch but with lower memory requirements.

Algorithm and Implementation

The gradient accumulation algorithm processes $K$ mini-batches of size $B_{\text{mini}}$, accumulating their gradients, then performs a single parameter update. The effective batch size is $B_{\text{eff}} = K \times B_{\text{mini}}$.

The loss scaling by $1/K$ ensures that the accumulated gradient has the correct magnitude. Without this scaling, the accumulated gradient would be $K$ times larger than the gradient from a single batch of size $B_{\text{eff}}$, leading to overly aggressive parameter updates.

In PyTorch, gradient accumulation is implemented by simply not calling optimizer.zero\_grad() after each mini-batch. Gradients accumulate automatically because PyTorch adds new gradients to existing gradients by default:

optimizer.zero_grad()
for k in range(accumulation_steps):
    batch = next(dataloader)
    loss = model(batch) / accumulation_steps
    loss.backward()  # Accumulates gradients
    
optimizer.step()  # Update parameters

Trade-offs and Considerations

Gradient accumulation is mathematically equivalent to training with a larger batch size, but it has different computational characteristics. The key trade-offs are:

Memory: Gradient accumulation requires only the memory for a single mini-batch of size $B_{\text{mini}}$, not the full effective batch size $B_{\text{eff}}$. This is the primary benefit—it enables training with large effective batch sizes on memory-constrained hardware.

Computation time: Gradient accumulation is slower than true large-batch training because the mini-batches are processed sequentially rather than in parallel. For $K$ accumulation steps, we perform $K$ forward passes and $K$ backward passes before a single parameter update. If we could fit the full batch in memory, we would perform 1 forward pass and 1 backward pass, processing $K$ times more data in parallel.

The time overhead is typically 10-20\% compared to true large-batch training, arising from:

• Reduced parallelism: processing mini-batches sequentially rather than in parallel
• Increased overhead: $K$ forward-backward passes have more overhead than 1 pass
• Memory bandwidth: loading model parameters $K$ times rather than once

Batch normalization incompatibility: Gradient accumulation is incompatible with batch normalization because batch normalization computes statistics over the mini-batch, not the effective batch. Each mini-batch has different statistics, leading to incorrect normalization. Fortunately, transformers use layer normalization rather than batch normalization, so this is not a concern for transformer training.

Practical Example

Consider training BERT-base where we want an effective batch size of 512, but GPU memory only allows batch size 32. We use gradient accumulation with $K = 512 / 32 = 16$ steps.

Memory requirements:

• Without accumulation (batch 512): $\approx 220$ GB (exceeds any single GPU)
• With accumulation (batch 32): $\approx 13.8$ GB (fits on V100 16GB)

Training time comparison:

• True batch 512 (if it fit): 1 forward + 1 backward = 2 passes
• Gradient accumulation: 16 forward + 16 backward = 32 passes

The gradient accumulation approach requires 16× more passes, but each pass is faster because it processes less data. The total time is approximately 15\% longer than true batch 512 would be, but it's feasible on available hardware.

When to use gradient accumulation:

• When the desired batch size exceeds GPU memory capacity
• When trying to match published training recipes that use large batches
• When larger batches improve convergence (common for transformers)
• When training time is less critical than achieving good final performance

When not to use gradient accumulation:

• When the mini-batch size is already optimal for convergence
• When training time is critical and larger batches don't improve convergence
• When the overhead (15-20\%) is unacceptable

For BERT-base, gradient accumulation is commonly used to achieve effective batch sizes of 256-512, which provide better convergence than smaller batches. The time overhead is acceptable given the improved final performance.

Gradient Checkpointing

Gradient checkpointing, also called activation checkpointing, is a memory-computation trade-off technique that dramatically reduces activation memory at the cost of increased training time. Instead of storing all intermediate activations during the forward pass for use in backpropagation, gradient checkpointing stores only a subset of activations (typically at layer boundaries) and recomputes the remaining activations during the backward pass as needed.

The Memory-Computation Trade-off

Standard backpropagation requires storing all intermediate activations from the forward pass because computing gradients requires both the gradients flowing backward and the activations from the forward pass. For a transformer with $L$ layers, batch size $B$, and sequence length $n$, the activation memory scales as $O(LBnd_{\text{model}})$ for linear terms and $O(LBhn^2)$ for attention matrices. As analyzed in Chapter 12, this activation memory often dominates total memory consumption, particularly for large batch sizes or long sequences.

Gradient checkpointing reduces activation memory by storing only activations at layer boundaries (the input to each transformer layer) and discarding all intermediate activations within layers. During the backward pass, when gradients need to flow through a layer, the forward computation for that layer is re-executed to reconstruct the intermediate activations needed for gradient computation. This recomputation happens on-the-fly during backpropagation, so the intermediate activations are used immediately and then discarded.

The memory savings are substantial. Without checkpointing, we store activations for every operation: QKV projections, attention scores, attention outputs, FFN intermediate activations, layer norm outputs, and residual connections. With checkpointing, we store only the layer inputs. For a typical transformer layer, this reduces activation memory by approximately 80\%, storing only 1-2 tensors per layer instead of 8-10 tensors.

The computational cost is the price for these memory savings. Each layer's forward computation must be executed twice: once during the forward pass (with activations discarded) and once during the backward pass (to reconstruct activations for gradient computation). This doubles the forward computation cost, but the backward pass cost remains the same. Since the backward pass already costs approximately 2× the forward pass, the total cost increases from 3× to 4× the forward pass, a 33\% increase in training time. In practice, the overhead is typically 20-30\% due to optimizations and the fact that some operations (like attention softmax) are relatively cheap to recompute.

Implementation Strategies

The most common checkpointing strategy is to checkpoint at transformer layer boundaries. For a model with $L$ layers, we store $L+1$ activation tensors (the input to each layer plus the final output), rather than storing all intermediate activations within layers.

In PyTorch, gradient checkpointing is implemented using torch.utils.checkpoint.checkpoint, which wraps a function and handles the recomputation automatically:

from torch.utils.checkpoint import checkpoint

class TransformerLayer(nn.Module):
    def forward(self, x):
        # Use checkpointing for this layer
        return checkpoint(self._forward, x)
    
    def _forward(self, x):
        # Actual layer computation
        # Attention
        attn_out = self.attention(x)
        x = x + self.dropout(attn_out)
        x = self.layer_norm1(x)
        
        # Feed-forward
        ffn_out = self.ffn(x)
        x = x + self.dropout(ffn_out)
        x = self.layer_norm2(x)
        
        return x

During the forward pass, PyTorch executes \_forward but doesn't store intermediate activations. During the backward pass, when gradients reach this layer, PyTorch re-executes \_forward with the saved input x, reconstructing the intermediate activations needed for gradient computation.

An alternative strategy is selective checkpointing, where only some layers are checkpointed. This provides a middle ground between memory and computation. For example, checkpointing every other layer reduces activation memory by approximately 50\% while increasing training time by only 10-15\%. This can be optimal when memory is tight but not critically constrained.

Practical Impact

The impact of gradient checkpointing is best illustrated with concrete examples. For GPT-2 (small) with 12 layers, $d_{\text{model}} = 768$, sequence length 1024, and batch size 32:

Without checkpointing:

This exceeds the memory of most GPUs when combined with parameters and optimizer states.

With checkpointing:

This dramatic reduction enables training with much larger batch sizes or longer sequences on the same hardware. For GPT-2 on an NVIDIA V100 (16 GB), checkpointing enables increasing the batch size from approximately 4 to 20, a 5× improvement.

Training time impact:

• Without checkpointing: 100\% (baseline)
• With checkpointing: 125\% (25\% slower)

The 25\% time increase is usually acceptable given the 5× increase in batch size, which often improves convergence and reduces the total number of steps needed for training.

When to Use Gradient Checkpointing

Gradient checkpointing is most beneficial in specific scenarios:

Use checkpointing when:

• Training with long sequences (e.g., $n > 1024$) where activation memory dominates
• GPU memory is the limiting factor preventing larger batch sizes
• The model is very deep (many layers) and activation memory scales linearly with depth
• Training time is less critical than maximizing batch size or sequence length
• Combined with mixed precision, checkpointing enables training that would otherwise be impossible

Avoid checkpointing when:

• Memory is not constrained and the 20-30\% time overhead is unacceptable
• Training with short sequences and small batch sizes where activation memory is already manageable
• Optimizing for minimum training time rather than maximum throughput
• The model is shallow enough that activation memory is not the bottleneck

For most transformer training, particularly for models with more than 12 layers or sequences longer than 512 tokens, gradient checkpointing is beneficial. The memory savings enable configurations that would otherwise be impossible, and the time overhead is modest compared to the benefits.

Distributed Training Strategies

As transformer models grow beyond the capacity of single GPUs, distributed training becomes essential. Different distributed training strategies partition the model, data, or optimizer state across multiple GPUs, each with distinct trade-offs in terms of memory reduction, communication overhead, and implementation complexity. Understanding these strategies is crucial for training large-scale models efficiently.

Data Parallelism

Data parallelism is the simplest and most widely used distributed training strategy. The model is replicated on each GPU, and each GPU processes a different subset of the training batch. After computing gradients locally, the GPUs synchronize their gradients using an AllReduce operation, then each GPU updates its local copy of the model with the averaged gradients.

The algorithm proceeds as follows:

1. Each GPU has a complete copy of the model
2. The global batch is split across GPUs: GPU $i$ processes mini-batch $\mathcal{B}_i$
3. Each GPU performs forward and backward passes independently, computing local gradients $\mathbf{g}_i$
4. AllReduce operation computes the average gradient: $\bar{\mathbf{g}} = \frac{1}{N}\sum_{i=1}^{N} \mathbf{g}_i$ where $N$ is the number of GPUs
5. Each GPU updates its model using $\bar{\mathbf{g}}$
6. All GPUs now have identical models (up to floating-point precision)

Data parallelism scales efficiently to 8-16 GPUs on a single node (connected via NVLink or PCIe) because the communication overhead is relatively small compared to computation. For BERT-base with 110M parameters, the AllReduce operation must communicate $110\text{M} \times 4 = 440$ MB of gradients. On NVLink (300 GB/s bandwidth), this takes approximately $440\text{ MB} / 300\text{ GB/s} \approx 1.5$ ms, which is small compared to the forward-backward computation time of 10-20 ms per batch.

However, data parallelism does not reduce memory requirements per GPU—each GPU still stores the complete model, optimizer states, and activations for its mini-batch. This limits the size of models that can be trained with data parallelism alone. For GPT-3 with 175B parameters requiring 700 GB in FP32, data parallelism is insufficient because no single GPU has enough memory for the complete model.

Model Parallelism

Model parallelism splits the model across multiple GPUs, with different layers residing on different devices. For a model with $L$ layers split across $N$ GPUs, each GPU stores approximately $L/N$ layers. This reduces per-GPU memory proportionally to the number of GPUs.

The forward pass proceeds sequentially: GPU 1 processes the input through its layers, sends activations to GPU 2, which processes through its layers, and so on. The backward pass proceeds in reverse: GPU $N$ computes gradients for its layers, sends gradients to GPU $N-1$, which computes gradients for its layers, and so on.

The primary challenge with model parallelism is the pipeline bubble problem. While GPU 1 is processing the next batch, GPUs 2 through $N$ are idle, waiting for activations from GPU 1. Similarly, during the backward pass, GPU $N$ finishes first and sits idle while earlier GPUs complete their backward passes. This sequential execution leads to poor GPU utilization, with each GPU active only $1/N$ of the time in the worst case.

Model parallelism is necessary when a single layer or the complete model exceeds single-GPU memory, but it should be combined with other strategies to improve utilization. For GPT-3, model parallelism alone would require hundreds of GPUs and would have terrible utilization due to pipeline bubbles.

Pipeline Parallelism

Pipeline parallelism improves upon model parallelism by splitting each batch into micro-batches and pipelining their execution across GPUs. Instead of processing one batch completely before starting the next, pipeline parallelism processes multiple micro-batches concurrently, with different micro-batches at different stages of the pipeline.

For example, with 4 GPUs and 4 micro-batches:

• Time 1: GPU 1 processes micro-batch 1 (forward)
• Time 2: GPU 1 processes micro-batch 2 (forward), GPU 2 processes micro-batch 1 (forward)
• Time 3: GPU 1 processes micro-batch 3 (forward), GPU 2 processes micro-batch 2 (forward), GPU 3 processes micro-batch 1 (forward)
• Time 4: All GPUs are active, processing different micro-batches

This pipelining significantly reduces idle time. The pipeline bubble (time when some GPUs are idle) is proportional to the number of GPUs divided by the number of micro-batches. With $N$ GPUs and $M$ micro-batches, the bubble fraction is approximately $N/M$. Using $M = 4N$ micro-batches reduces the bubble to 25\%, achieving 75\% utilization.

Pipeline parallelism implementations like GPipe and PipeDream differ in how they handle gradient computation and weight updates. GPipe uses synchronous updates, accumulating gradients from all micro-batches before updating weights. PipeDream uses asynchronous updates, updating weights after each micro-batch, which can improve throughput but requires careful handling of weight versions.

Tensor Parallelism

Tensor parallelism, pioneered by Megatron-LM, splits individual layers across multiple GPUs rather than splitting the model layer-wise. For attention and feed-forward layers, the computation can be partitioned across GPUs with minimal communication.

For the attention mechanism, the heads can be split across GPUs. With $h$ heads and $N$ GPUs, each GPU computes $h/N$ heads independently. The only communication required is an AllReduce after computing the attention output, to sum the contributions from all heads.

For the feed-forward network, the first linear layer $\mW_1 \in \R^{d_{\text{model}} \times d_{ff}}$ can be column-partitioned across GPUs. Each GPU computes a subset of the $d_{ff}$ intermediate activations. The GELU activation is applied independently on each GPU. The second linear layer $\mW_2 \in \R^{d_{ff} \times d_{\text{model}}}$ is row-partitioned, and an AllReduce sums the outputs from all GPUs.

Tensor parallelism achieves $N\times$ memory reduction with only two AllReduce operations per layer (one for attention, one for FFN). The communication volume is $O(Bnd_{\text{model}})$ per layer, which is much smaller than the $O(P)$ communication required for data parallelism (where $P$ is the number of parameters).

Tensor parallelism is particularly effective for very large layers. For GPT-3 with $d_{\text{model}} = 12{,}288$ and $d_{ff} = 49{,}152$, a single FFN layer has $2 \times 12{,}288 \times 49{,}152 \approx 1.2$B parameters, requiring 4.8 GB in FP32. Splitting across 8 GPUs reduces this to 600 MB per GPU, making the layer tractable.

ZeRO: Zero Redundancy Optimizer

ZeRO (Zero Redundancy Optimizer) is a family of optimizations that reduce memory by sharding optimizer states, gradients, and parameters across GPUs while maintaining the computational efficiency of data parallelism. ZeRO has three stages, each providing progressively more memory reduction:

ZeRO Stage 1: Optimizer State Partitioning

Each GPU stores only $1/N$ of the optimizer states (first and second moments for Adam). During the optimizer step, each GPU updates only its partition of the parameters. This reduces optimizer memory by $N\times$ with minimal communication overhead.

For BERT-base with 110M parameters and 8 GPUs:

• Without ZeRO: Each GPU stores 880 MB of optimizer states
• With ZeRO-1: Each GPU stores $880 / 8 = 110$ MB of optimizer states
• Memory saved: 770 MB per GPU

ZeRO Stage 2: Gradient Partitioning

In addition to optimizer states, gradients are also partitioned. Each GPU computes gradients for all parameters during backpropagation but only retains the gradients for its partition, discarding the rest. This reduces gradient memory by $N\times$.

For BERT-base with 8 GPUs:

• Without ZeRO: Each GPU stores 440 MB of gradients
• With ZeRO-2: Each GPU stores $440 / 8 = 55$ MB of gradients
• Total memory saved: $770 + 385 = 1{,}155$ MB per GPU

ZeRO Stage 3: Parameter Partitioning

The most aggressive stage partitions the parameters themselves. Each GPU stores only $1/N$ of the parameters. During the forward pass, each GPU gathers the parameters it needs from other GPUs, computes its portion of the forward pass, then discards the gathered parameters. The backward pass proceeds similarly.

For BERT-base with 8 GPUs:

• Without ZeRO: Each GPU stores 440 MB of parameters
• With ZeRO-3: Each GPU stores $440 / 8 = 55$ MB of parameters
• Total memory saved: $770 + 385 + 385 = 1{,}540$ MB per GPU

ZeRO-3 enables training models that wouldn't fit on any single GPU by distributing all memory across the cluster. For GPT-3 with 175B parameters requiring 700 GB in FP32, ZeRO-3 across 64 A100 GPUs (80 GB each) reduces per-GPU memory to $700 / 64 \approx 11$ GB, making training feasible.

The communication overhead of ZeRO increases with each stage. ZeRO-1 has minimal overhead (only during optimizer step). ZeRO-2 adds gradient communication (similar to data parallelism). ZeRO-3 adds parameter communication during forward and backward passes, which can be significant but is often acceptable given the memory savings.

Comparison of Strategies

In practice, large-scale training often combines multiple strategies. GPT-3 training used a combination of data parallelism, model parallelism, and pipeline parallelism across thousands of GPUs. Modern frameworks like DeepSpeed and Megatron-LM provide implementations of these strategies that can be combined flexibly based on model size and available hardware.

Batch Size and Sequence Length Selection

Selecting appropriate batch sizes and sequence lengths is crucial for efficient transformer training. These choices directly impact memory consumption, training throughput, convergence behavior, and final model quality. The optimal configuration depends on the interplay between hardware constraints, model architecture, and training objectives.

Batch Size Considerations

Batch size affects both computational efficiency and optimization dynamics. Larger batches improve GPU utilization by amortizing the cost of loading model parameters and by providing more parallelism for matrix operations. Modern GPUs achieve peak performance with large matrix multiplications, and larger batches create larger matrices that better utilize the hardware.

For BERT-base on an NVIDIA A100, throughput (tokens processed per second) increases significantly with batch size:

• Batch size 8: $\approx 15{,}000$ tokens/sec (30\% GPU utilization)
• Batch size 32: $\approx 50{,}000$ tokens/sec (80\% GPU utilization)
• Batch size 64: $\approx 70{,}000$ tokens/sec (90\% GPU utilization)
• Batch size 128: $\approx 75{,}000$ tokens/sec (95\% GPU utilization)

Beyond batch size 64, the throughput gains diminish because the GPU is already well-utilized. The optimal batch size for throughput is typically where GPU utilization reaches 85-95\%, which depends on the model size and sequence length.

However, larger batches are not always better for optimization. Very large batches can hurt generalization, a phenomenon known as the "generalization gap." The intuition is that large batches provide very accurate gradient estimates, which can lead the optimizer to sharp minima that don't generalize well. Smaller batches provide noisier gradients that help the optimizer find flatter minima with better generalization.

The relationship between batch size and generalization is complex and depends on the learning rate schedule and total training budget. Research has shown that the generalization gap can be mitigated by:

• Scaling the learning rate proportionally with batch size (linear scaling rule)
• Extending the warmup period for larger batches
• Training for more steps to compensate for fewer parameter updates

For transformer training, batch sizes of 256-2048 are typical. BERT-base uses an effective batch size of 256 (32 per GPU × 8 GPUs). GPT-2 uses batch sizes of 512-1024. GPT-3 uses batch sizes up to 3.2 million tokens (approximately 1600 sequences of length 2048), enabled by LAMB optimizer and massive parallelism.

Memory Scaling with Batch Size

Memory consumption scales linearly with batch size for most components. For BERT-base with sequence length 512:

The linear scaling means that doubling the batch size doubles the memory requirement. This quickly exceeds single-GPU capacity, necessitating either gradient accumulation (to simulate large batches with small physical batches) or distributed training (to split the batch across multiple GPUs).

The memory breakdown for batch size 32 is approximately:

• Parameters + optimizer: 1.76 GB (independent of batch size)
• Activations: 12 GB (scales linearly with batch size)

Since activations dominate, techniques that reduce activation memory (mixed precision, gradient checkpointing) have a large impact on the maximum feasible batch size.

Sequence Length Considerations

Sequence length has a more complex impact on memory and computation than batch size. The attention mechanism's quadratic scaling means that memory and computation grow as $O(n^2)$ for sequence length $n$, while other components grow linearly as $O(n)$.

For BERT-base with batch size 32, memory consumption varies dramatically with sequence length:

Doubling the sequence length from 512 to 1024 roughly triples the memory (not quadruples, because some components scale linearly). The attention matrices grow quadratically: for 12 heads, the attention memory is $32 \times 12 \times n^2 \times 4$ bytes. At $n=512$, this is 403 MB; at $n=1024$, this is 1.6 GB; at $n=2048$, this is 6.4 GB.

The quadratic scaling limits practical sequence lengths. BERT uses $n=512$, GPT-2 uses $n=1024$, GPT-3 uses $n=2048$. Longer sequences require either:

• Efficient attention mechanisms (sparse attention, linear attention) that reduce the $O(n^2)$ complexity
• Gradient checkpointing to reduce activation memory
• Smaller batch sizes to fit within memory constraints
• More powerful GPUs with larger memory

The choice of sequence length depends on the task. For tasks requiring long-range dependencies (document classification, long-form generation), longer sequences are beneficial despite the computational cost. For tasks with local dependencies (named entity recognition, part-of-speech tagging), shorter sequences may suffice.

Dynamic Batching

Dynamic batching groups sequences of similar length together to minimize padding waste. In a typical batch, sequences have varying lengths, and all sequences are padded to the length of the longest sequence in the batch. This padding wastes computation and memory on padding tokens that don't contribute to learning.

For example, if a batch contains sequences of lengths [128, 256, 512, 512], all sequences are padded to 512, wasting:

Out of $4 \times 512 = 2048$ total tokens, 640 (31\%) are padding.

Dynamic batching sorts sequences by length and groups similar lengths together. This reduces padding significantly. If we instead batch [128, 128, 128, 128] and [512, 512, 512, 512] separately, there's no padding waste within each batch.

The throughput improvement from dynamic batching can be substantial:

• Without dynamic batching: 50,000 tokens/sec (including padding)
• With dynamic batching: 70,000 tokens/sec (40\% improvement)

The improvement depends on the length distribution in the dataset. For datasets with highly variable lengths, dynamic batching can provide 2-3× throughput improvements. For datasets with uniform lengths, the benefit is minimal.

Dynamic batching is implemented by sorting the dataset by sequence length before creating batches, or by using a bucketing strategy that assigns sequences to length buckets and samples batches from within buckets. Most modern training frameworks (Hugging Face Transformers, fairseq) support dynamic batching.

Practical Guidelines

Based on the analysis above, practical guidelines for batch size and sequence length selection are:

For batch size:

• Start with the largest batch size that fits in GPU memory
• If memory-constrained, use gradient accumulation to achieve larger effective batch sizes
• For BERT-base on V100 (16 GB): batch size 16-32 with sequence length 512
• For BERT-base on A100 (40 GB): batch size 32-64 with sequence length 512
• Scale learning rate proportionally when increasing batch size
• Extend warmup period for very large batches (>1024)

For sequence length:

• Use the longest sequence length that fits in memory and is relevant for the task
• For memory-constrained scenarios, reduce batch size rather than sequence length if long context is important
• Use gradient checkpointing to enable longer sequences
• Consider efficient attention mechanisms for sequences longer than 2048
• Use dynamic batching to reduce padding waste

Memory-constrained optimization:

1. Enable mixed precision training (FP16/BF16): 40-50\% memory reduction
2. Enable gradient checkpointing: 80\% activation memory reduction
3. Use gradient accumulation: simulate large batches with small physical batches
4. Reduce sequence length if task permits
5. Use dynamic batching to reduce padding waste

These techniques can be combined. For example, BERT-base with mixed precision + gradient checkpointing can train with batch size 128 and sequence length 512 on a V100 (16 GB), compared to batch size 16 without these optimizations.

Regularization Techniques

Regularization prevents overfitting by constraining the model's capacity or adding noise during training. Transformers, with their large parameter counts, are particularly susceptible to overfitting on small datasets. Effective regularization enables transformers to generalize well from training data to unseen examples.

Dropout

Dropout randomly sets activations to zero during training with probability $p$, forcing the model to learn robust features that don't rely on any single activation. During inference, dropout is disabled, and activations are scaled by $(1-p)$ to maintain the expected magnitude.

In transformers, dropout is applied at multiple locations:

Attention dropout: Applied to the attention weights after softmax, before multiplying by values:

This prevents the model from relying too heavily on specific attention patterns, encouraging it to learn diverse attention strategies.

Residual dropout: Applied to the output of each sub-layer before adding to the residual connection:

This regularizes the transformations learned by attention and feed-forward layers.

Embedding dropout: Applied to the sum of token embeddings and positional encodings:

This prevents overfitting to specific token representations.

Typical dropout rates for transformers are relatively low compared to other architectures. BERT uses $p = 0.1$ (10\% dropout) for all dropout locations. GPT-2 also uses $p = 0.1$. Larger models sometimes use even lower dropout rates ($p = 0.05$ or less) because their increased capacity provides implicit regularization.

The dropout rate should be tuned based on the dataset size and model capacity. For small datasets (thousands of examples), higher dropout rates ($p = 0.2$ or $p = 0.3$) may be beneficial. For large datasets (millions of examples), lower dropout rates ($p = 0.1$ or less) are typically sufficient.

Weight Decay

Weight decay adds an L2 penalty to the loss function, encouraging parameters to remain small. In the context of AdamW (the standard optimizer for transformers), weight decay is applied directly to parameters rather than through the gradient:

The weight decay coefficient $\lambda$ controls the strength of regularization. Typical values for transformer training are $\lambda = 0.01$ or $\lambda = 0.001$. BERT uses $\lambda = 0.01$, which provides moderate regularization without overly constraining the model.

Weight decay is not applied uniformly to all parameters. Biases and layer normalization parameters (scale $\gamma$ and shift $\beta$) are typically excluded from weight decay. The reasoning is that these parameters control the scale and offset of activations rather than the complexity of learned features, and regularizing them can hurt performance. In practice, this exclusion is implemented by creating separate parameter groups in the optimizer with different weight decay settings.

The interaction between weight decay and learning rate is important. Because weight decay is applied with coefficient $\eta \lambda$, the effective regularization strength increases with the learning rate. During warmup, when the learning rate is small, weight decay has minimal effect. As the learning rate increases, weight decay becomes stronger. During decay, as the learning rate decreases, weight decay weakens. This dynamic regularization schedule often works well in practice.

Label Smoothing

Label smoothing replaces hard one-hot targets with soft targets that assign small probabilities to incorrect classes. For a classification problem with vocabulary size $V$ and true class $y$, the smoothed target distribution is:

where $\epsilon$ is the smoothing parameter, typically $\epsilon = 0.1$.

Label smoothing prevents the model from becoming overconfident in its predictions. Without smoothing, the model is trained to assign probability 1 to the correct class and probability 0 to all other classes. This can lead to overconfident predictions that don't reflect the model's true uncertainty. With smoothing, the model is trained to assign high probability to the correct class but also small probabilities to other classes, leading to better-calibrated predictions.

For language modeling with vocabulary size 30,000 and $\epsilon = 0.1$:

The smoothed target assigns 90\% probability to the correct class and distributes the remaining 10\% uniformly across all classes.

Label smoothing is particularly beneficial for tasks with ambiguous labels or where multiple outputs could be considered correct. In machine translation, for example, multiple translations may be valid, and label smoothing encourages the model to consider alternatives rather than committing entirely to the reference translation.

The cross-entropy loss with label smoothing is:

The second term is the negative entropy of the predicted distribution, which encourages the model to maintain some uncertainty rather than collapsing to a single prediction.

Gradient Clipping

Gradient clipping prevents exploding gradients by limiting the norm of the gradient vector. If the gradient norm exceeds a threshold $\theta$, the gradient is scaled down:

The typical threshold for transformer training is $\theta = 1.0$. This value is chosen empirically and works well across different model sizes and tasks.

Gradient clipping is essential for training stability, particularly in the early stages of training when gradients can be very large. Without clipping, occasional large gradients can cause the parameters to jump to regions of the loss landscape with poor gradients, derailing training. With clipping, these large gradients are tamed, allowing training to proceed smoothly.

The clipping threshold should be tuned based on the typical gradient norms observed during training. If gradients are frequently clipped, the threshold may be too low, preventing the model from making necessary large updates. If gradients are rarely clipped, the threshold may be too high, providing insufficient protection against exploding gradients. Monitoring the fraction of steps where clipping occurs (typically 1-5\%) helps tune the threshold.

Gradient clipping interacts with the learning rate: with a lower learning rate, gradients have less impact, so clipping is less necessary. With a higher learning rate, clipping becomes more important. The combination of learning rate warmup and gradient clipping provides robust training stability.

Training Time and Cost Estimates

Training costs for transformers scale dramatically with model size, spanning several orders of magnitude from models accessible to academic labs to those requiring industrial-scale infrastructure. Table~[ref] summarizes the approximate costs for representative models, assuming cloud GPU pricing of \$3--4 per V100 GPU-hour.

The cost increases by 3--4 orders of magnitude from BERT-base to GPT-3, while performance improvements, though substantial, exhibit diminishing returns. For many applications, smaller models fine-tuned on task-specific data provide excellent performance at a fraction of the cost.

Exercises

Solutions

Full solutions for all exercises are available at \url{https://deeplearning.hofkensvermeulen.be}.

import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler
import time

class SmallTransformer(nn.Module):
    def __init__(self, vocab_size=10000, d_model=256, n_heads=8, 
                 n_layers=4, d_ff=1024):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model, n_heads, d_ff, batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, n_layers)
        self.output = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        x = self.embedding(x)
        x = self.transformer(x)
        return self.output(x)

# Training function with FP32
def train_fp32(model, data_loader, optimizer, epochs=5):
    model.train()
    start_time = time.time()
    
    for epoch in range(epochs):
        for batch_idx, (data, target) in enumerate(data_loader):
            optimizer.zero_grad()
            output = model(data)
            loss = nn.functional.cross_entropy(
                output.view(-1, output.size(-1)), target.view(-1)
            )
            loss.backward()
            optimizer.step()
    
    return time.time() - start_time

# Training function with mixed precision
def train_mixed_precision(model, data_loader, optimizer, epochs=5, 
                         loss_scale=2**16):
    model.train()
    scaler = GradScaler(init_scale=loss_scale)
    start_time = time.time()
    
    for epoch in range(epochs):
        for batch_idx, (data, target) in enumerate(data_loader):
            optimizer.zero_grad()
            
            # Forward pass in FP16
            with autocast():
                output = model(data)
                loss = nn.functional.cross_entropy(
                    output.view(-1, output.size(-1)), target.view(-1)
                )
            
            # Backward pass with scaled loss
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
    
    return time.time() - start_time

# Memory profiling
def profile_memory(model, data_loader, use_mixed_precision=False):
    torch.cuda.reset_peak_memory_stats()
    
    if use_mixed_precision:
        scaler = GradScaler()
        with autocast():
            for data, target in data_loader:
                output = model(data)
                loss = nn.functional.cross_entropy(
                    output.view(-1, output.size(-1)), target.view(-1)
                )
        scaler.scale(loss).backward()
    else:
        for data, target in data_loader:
            output = model(data)
            loss = nn.functional.cross_entropy(
                output.view(-1, output.size(-1)), target.view(-1)
            )
        loss.backward()
    
    return torch.cuda.max_memory_allocated() / 1024**3  # GB

import torch
from torch.utils.checkpoint import checkpoint

# Without checkpointing
torch.cuda.reset_peak_memory_stats()
output = model(input_ids)
loss = output.loss
loss.backward()
memory_without = torch.cuda.max_memory_allocated() / 1024**3
print(f"Memory without checkpointing: {memory_without:.2f} GB")

# With checkpointing
torch.cuda.reset_peak_memory_stats()
output = checkpoint(model, input_ids)
loss = output.loss
loss.backward()
memory_with = torch.cuda.max_memory_allocated() / 1024**3
print(f"Memory with checkpointing: {memory_with:.2f} GB")

import torch
import torch.nn as nn
import time

def train_with_accumulation(model, data_loader, optimizer, 
                           physical_batch_size=32, 
                           effective_batch_size=512):
    accumulation_steps = effective_batch_size // physical_batch_size
    model.train()
    optimizer.zero_grad()
    
    losses = []
    start_time = time.time()
    
    for batch_idx, (data, target) in enumerate(data_loader):
        # Forward pass
        output = model(data)
        loss = nn.functional.cross_entropy(
            output.view(-1, output.size(-1)), target.view(-1)
        )
        
        # Scale loss by accumulation steps
        loss = loss / accumulation_steps
        loss.backward()
        
        # Update weights every accumulation_steps
        if (batch_idx + 1) 
            optimizer.step()
            optimizer.zero_grad()
            losses.append(loss.item() * accumulation_steps)
    
    training_time = time.time() - start_time
    return losses, training_time

def train_true_batch(model, data_loader, optimizer, batch_size=512):
    model.train()
    losses = []
    start_time = time.time()
    
    for data, target in data_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = nn.functional.cross_entropy(
            output.view(-1, output.size(-1)), target.view(-1)
        )
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    
    training_time = time.time() - start_time
    return losses, training_time

import matplotlib.pyplot as plt
import numpy as np

# Compare loss curves
losses_true = train_true_batch(model, loader_512, optimizer)
losses_accum = train_with_accumulation(model, loader_32, optimizer)

plt.figure(figsize=(10, 6))
plt.plot(losses_true, label='True batch 512', alpha=0.7)
plt.plot(losses_accum, label='Gradient accumulation', alpha=0.7)
plt.xlabel('Update step')
plt.ylabel('Loss')
plt.legend()
plt.title('Training Dynamics: True Batch vs Gradient Accumulation')
plt.grid(True)

# Compute correlation
correlation = np.corrcoef(losses_true, losses_accum)[0, 1]
print(f"Loss correlation: {correlation:.4f}")  # Expected: > 0.99

import torch
import torch.nn as nn
import math

# (a) Warmup + Linear Decay
def linear_schedule(step, warmup_steps=4000, total_steps=100000):
    if step < warmup_steps:
        return step / warmup_steps
    else:
        return max(0.0, (total_steps - step) / (total_steps - warmup_steps))

# (b) Warmup + Inverse Square Root Decay
def inverse_sqrt_schedule(step, warmup_steps=4000, d_model=256):
    return min(step ** (-0.5), step * warmup_steps ** (-1.5))

# (c) Warmup + Cosine Annealing
def cosine_schedule(step, warmup_steps=4000, total_steps=100000):
    if step < warmup_steps:
        return step / warmup_steps
    else:
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        return 0.5 * (1 + math.cos(math.pi * progress))

# Training function
def train_with_schedule(model, data_loader, base_lr=1e-3, 
                       schedule_fn=linear_schedule, epochs=50):
    optimizer = torch.optim.Adam(model.parameters(), lr=base_lr)
    
    losses = []
    lrs = []
    step = 0
    
    for epoch in range(epochs):
        for data, target in data_loader:
            # Update learning rate
            lr_scale = schedule_fn(step)
            for param_group in optimizer.param_groups:
                param_group['lr'] = base_lr * lr_scale
            
            # Training step
            optimizer.zero_grad()
            output = model(data)
            loss = nn.functional.cross_entropy(
                output.view(-1, output.size(-1)), target.view(-1)
            )
            loss.backward()
            optimizer.step()
            
            losses.append(loss.item())
            lrs.append(optimizer.param_groups[0]['lr'])
            step += 1
    
    return losses, lrs

import matplotlib.pyplot as plt

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

# Plot learning rate schedules
steps = range(10000)
ax1.plot([linear_schedule(s) for s in steps], label='Linear')
ax1.plot([inverse_sqrt_schedule(s) for s in steps], label='Inverse sqrt')
ax1.plot([cosine_schedule(s) for s in steps], label='Cosine')
ax1.set_xlabel('Training step')
ax1.set_ylabel('LR multiplier')
ax1.set_title('Learning Rate Schedules')
ax1.legend()
ax1.grid(True)

# Plot loss curves
ax2.plot(losses_linear, label='Linear', alpha=0.7)
ax2.plot(losses_inverse, label='Inverse sqrt', alpha=0.7)
ax2.plot(losses_cosine, label='Cosine', alpha=0.7)
ax2.set_xlabel('Training step')
ax2.set_ylabel('Loss')
ax2.set_title('Training Loss Curves')
ax2.legend()
ax2.grid(True)
plt.tight_layout()

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import time

def setup_distributed(rank, world_size):
    """Initialize distributed training"""
    dist.init_process_group(
        backend='nccl',
        init_method='env://',
        world_size=world_size,
        rank=rank
    )
    torch.cuda.set_device(rank)

def train_distributed(rank, world_size, model, data_loader, epochs=10):
    setup_distributed(rank, world_size)
    
    # Wrap model with DDP
    model = model.to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=1e-3)
    
    # Track timing
    compute_time = 0
    comm_time = 0
    
    for epoch in range(epochs):
        for data, target in data_loader:
            data, target = data.to(rank), target.to(rank)
            
            # Computation phase
            start_compute = time.time()
            optimizer.zero_grad()
            output = ddp_model(data)
            loss = nn.functional.cross_entropy(
                output.view(-1, output.size(-1)), target.view(-1)
            )
            loss.backward()
            compute_time += time.time() - start_compute
            
            # Communication phase (AllReduce)
            start_comm = time.time()
            optimizer.step()  # Includes gradient synchronization
            comm_time += time.time() - start_comm
    
    return compute_time, comm_time

import torch
import torch.nn as nn

def train_with_regularization(model, train_loader, val_loader, 
                             dropout=0.0, weight_decay=0.0, 
                             label_smoothing=0.0, epochs=100):
    # Apply dropout to model
    for module in model.modules():
        if isinstance(module, nn.Dropout):
            module.p = dropout
    
    optimizer = torch.optim.AdamW(
        model.parameters(), 
        lr=1e-3, 
        weight_decay=weight_decay
    )
    
    criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing)
    
    train_losses, val_losses = [], []
    
    for epoch in range(epochs):
        # Training
        model.train()
        train_loss = 0
        for data, target in train_loader:
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output.view(-1, output.size(-1)), 
                           target.view(-1))
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
        
        # Validation
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for data, target in val_loader:
                output = model(data)
                loss = criterion(output.view(-1, output.size(-1)), 
                               target.view(-1))
                val_loss += loss.item()
        
        train_losses.append(train_loss / len(train_loader))
        val_losses.append(val_loss / len(val_loader))
    
    return train_losses, val_losses

import torch
from torch.nn.utils.rnn import pad_sequence
import time

def static_batching(dataset, batch_size=32, max_length=512):
    """Traditional batching with fixed max length"""
    batches = []
    total_tokens = 0
    padding_tokens = 0
    
    for i in range(0, len(dataset), batch_size):
        batch = dataset[i:i+batch_size]
        
        # Pad all sequences to max_length
        padded = []
        for seq in batch:
            if len(seq) < max_length:
                padded.append(torch.cat([
                    seq, 
                    torch.zeros(max_length - len(seq), dtype=torch.long)
                ]))
            else:
                padded.append(seq[:max_length])
        
        batch_tensor = torch.stack(padded)
        batches.append(batch_tensor)
        
        # Count tokens
        total_tokens += batch_size * max_length
        for seq in batch:
            padding_tokens += max(0, max_length - len(seq))
    
    padding_fraction = padding_tokens / total_tokens
    return batches, padding_fraction

def dynamic_batching(dataset, batch_size=32, max_tokens=16384):
    """Dynamic batching: group similar lengths, minimize padding"""
    # Sort by length
    sorted_data = sorted(enumerate(dataset), key=lambda x: len(x[1]))
    
    batches = []
    total_tokens = 0
    padding_tokens = 0
    
    i = 0
    while i < len(sorted_data):
        batch = []
        batch_length = 0
        
        # Fill batch up to max_tokens
        while i < len(sorted_data) and len(batch) < batch_size:
            idx, seq = sorted_data[i]
            seq_len = len(seq)
            
            # Check if adding this sequence exceeds max_tokens
            if len(batch) > 0:
                new_batch_length = max(batch_length, seq_len)
                if new_batch_length * (len(batch) + 1) > max_tokens:
                    break
            
            batch.append(seq)
            batch_length = max(batch_length, seq_len)
            i += 1
        
        # Pad batch to max length in batch
        padded = pad_sequence(batch, batch_first=True, padding_value=0)
        batches.append(padded)
        
        # Count tokens
        actual_tokens = sum(len(seq) for seq in batch)
        total_tokens += padded.numel()
        padding_tokens += padded.numel() - actual_tokens
    
    padding_fraction = padding_tokens / total_tokens
    return batches, padding_fraction

def measure_throughput(model, batches, device='cuda'):
    model.eval()
    total_tokens = 0
    
    torch.cuda.synchronize()
    start_time = time.time()
    
    with torch.no_grad():
        for batch in batches:
            batch = batch.to(device)
            output = model(batch)
            total_tokens += (batch != 0).sum().item()
    
    torch.cuda.synchronize()
    elapsed = time.time() - start_time
    
    throughput = total_tokens / elapsed
    return throughput

# Compare methods
static_batches, static_padding = static_batching(dataset)
dynamic_batches, dynamic_padding = dynamic_batching(dataset)

static_throughput = measure_throughput(model, static_batches)
dynamic_throughput = measure_throughput(model, dynamic_batches)

print(f"Static batching:")
print(f"  Padding fraction: {static_padding:.2
print(f"  Throughput: {static_throughput:.0f} tokens/sec")

print(f"\nDynamic batching:")
print(f"  Padding fraction: {dynamic_padding:.2
print(f"  Throughput: {dynamic_throughput:.0f} tokens/sec")

speedup = dynamic_throughput / static_throughput
print(f"\nSpeedup: {speedup:.2f}x")