Architecting Transformers from Scratch in Python
Optimized Training Loops
Production-Grade Training Loops
With the model architecture defined, the next critical step is implementing a training loop that is not just functional but also efficient and stable for large-scale models. A naive training loop will fail due to numerical instability, slow convergence, or excessive memory consumption. We will construct a production-grade loop by incorporating several advanced techniques specific to training modern Transformers.
Precision, Performance, and Stability
Modern GPUs contain specialized hardware, like Tensor Cores, that provide massive throughput gains for half-precision (float16) matrix multiplications. Automatic Mixed Precision (AMP) leverages this by performing operations in float16 where possible, while maintaining float32 for stability in other areas, such as weight updates.
PyTorch's torch.cuda.amp module simplifies this process. The autocast context manager automatically selects the appropriate precision for each operation within its scope. However, using float16 for gradients can lead to underflow, where small gradient values become zero. To prevent this, a GradScaler is used to scale the loss, which in turn scales the gradients into a range that float16 can represent. The optimizer step then unscales these gradients before updating the weights.
import torch
# Assume model, dataloader, optimizer, and loss_fn are defined
scaler = torch.cuda.amp.GradScaler()
for batch in dataloader:
optimizer.zero_grad()
# Use autocast for the forward pass
with torch.cuda.amp.autocast():
logits = model(batch['input_ids'])
loss = loss_fn(logits, batch['labels'])
# Scale the loss and perform a backward pass
scaler.scale(loss).backward()
# Unscale gradients and call optimizer.step()
scaler.step(optimizer)
# Update the scale for next iteration
scaler.update()
While AMP boosts performance, deep networks are prone to exploding gradients, where gradients grow exponentially and destabilize training. Gradient clipping addresses this by capping the norm of the gradients before the optimizer step. A common practice is to clip the L2 norm of all gradients concatenated together.
# ... inside the training loop, after scaler.scale(loss).backward()
# First, unscale the gradients that the scaler has scaled
scaler.unscale_(optimizer)
# Clip the gradients to a max norm (e.g., 1.0)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# Optimizer step and scaler update follow
scaler.step(optimizer)
scaler.update()
Advanced Optimization Strategies
The standard Adam optimizer couples weight decay with the gradient calculation. This can be suboptimal because the adaptive learning rates from Adam also scale the weight decay, making it less effective for weights with large gradients. AdamW decouples weight decay, treating it as true L2 regularization by subtracting a value proportional to the weight itself directly from the weight, after the gradient update. This leads to better generalization.
A further refinement is to apply weight decay selectively. Parameters like biases and normalization layer weights (LayerNorm, RMSNorm) generally do not benefit from weight decay and can even be harmed by it. The optimal approach is to separate model parameters into two groups: one for weights that receive decay, and one for biases and normalization parameters that do not.
def get_optimizer_param_groups(model):
# Create two parameter groups
decay_params = []
no_decay_params = []
for name, param in model.named_parameters():
if not param.requires_grad:
continue
# Check for weight tensors of linear layers
if 'weight' in name and isinstance(param, torch.nn.Linear):
decay_params.append(param)
# Check for biases and normalization layer parameters
elif 'bias' in name or 'norm' in name:
no_decay_params.append(param)
else:
# Default to decay for other parameters like embeddings
decay_params.append(param)
return [
{'params': decay_params, 'weight_decay': 0.01},
{'params': no_decay_params, 'weight_decay': 0.0}
]
param_groups = get_optimizer_param_groups(model)
optimizer = torch.optim.AdamW(param_groups, lr=3e-4, betas=(0.9, 0.95))
Learning rate scheduling is also crucial. A common and effective strategy for Transformers is a linear warmup followed by a cosine decay. The warmup phase gradually increases the learning rate from a small value to its target over the first few thousand steps. This prevents early instability when weights are randomly initialized. After the warmup, the learning rate anneals following a cosine curve, smoothly decreasing to a minimum value by the end of training.
Simulating Large Batches
Training LLMs benefits from very large batch sizes, often in the millions of tokens. However, this is rarely feasible on a single GPU due to memory constraints. Gradient accumulation is a technique to simulate a larger batch size. It works by performing the forward and backward passes for several smaller "micro-batches" and accumulating their gradients. The optimizer step is only performed after a specified number of accumulation steps, at which point the accumulated gradients are averaged and used to update the weights.
This effectively decouples the batch size used for the weight update from the batch size that fits in memory. An effective_batch_size is achieved by multiplying the micro_batch_size by the number of gradient_accumulation_steps.
micro_batch_size = 16
gradient_accumulation_steps = 8
# effective_batch_size = 16 * 8 = 128
for i, batch in enumerate(dataloader):
# Forward and backward pass
with torch.cuda.amp.autocast():
logits = model(batch['input_ids'])
loss = loss_fn(logits, batch['labels'])
# Scale loss to average over accumulation steps
loss = loss / gradient_accumulation_steps
scaler.scale(loss).backward()
# Perform optimizer step only after N accumulation steps
if (i + 1) % gradient_accumulation_steps == 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
# Update learning rate scheduler here
By combining these techniques—mixed precision, gradient clipping, decoupled weight decay, learning rate scheduling, and gradient accumulation—we can construct a robust and efficient training loop capable of handling large Transformer models.
What is the primary role of the GradScaler when using Automatic Mixed Precision (AMP) in PyTorch?
How does the AdamW optimizer improve upon the standard Adam optimizer?