Implementing Transformers in PyTorch

Chapter Overview

This chapter provides complete, production-ready PyTorch implementations of transformer models. We build from scratch: attention mechanisms, encoder/decoder blocks, position encodings, and full models (BERT, GPT, T5). Each implementation includes training loops, optimization, and best practices.

Learning Objectives

1. Implement multi-head attention from scratch
2. Build transformer encoder and decoder blocks
3. Create complete BERT and GPT models
4. Write efficient training loops with mixed precision
5. Apply gradient accumulation and checkpointing
6. Debug common implementation issues

Multi-Head Attention Implementation

Core Components

The implementation of multi-head attention in PyTorch requires careful attention to efficiency and memory usage. The standard approach involves projecting queries, keys, and values through linear layers, reshaping tensors to separate attention heads, computing scaled dot-product attention, and finally concatenating the results. However, several optimizations can significantly improve both speed and memory efficiency.

Key implementation considerations:

• Efficient batched matrix multiplications
• Proper dimension handling for multi-head split
• Memory-efficient attention computation
• Gradient flow through softmax

PyTorch multi-head attention structure:

1. Project Q, K, V: Linear layers
2. Reshape for multiple heads: view + transpose
3. Compute scaled dot-product attention
4. Concatenate heads and project output

Memory-Efficient Attention

The standard attention mechanism computes the full attention matrix of size $(B, h, n, n)$, which becomes prohibitively expensive for long sequences. For a sequence length of 512 with 12 heads and batch size 32, this requires approximately 400MB just for the attention scores. We can implement several optimizations to reduce this memory footprint.

The first optimization involves computing attention in chunks rather than materializing the entire attention matrix at once. This approach processes the attention computation in blocks, trading some computational efficiency for substantial memory savings. The chunked attention implementation divides the sequence into smaller segments and computes attention scores for each segment independently.

def memory_efficient_attention(Q, K, V, chunk_size=128):
    """
    Compute attention in chunks to reduce memory usage.
    Q, K, V: (batch, heads, seq_len, head_dim)
    """
    B, h, n, d = Q.shape
    output = torch.zeros_like(Q)
    
    for i in range(0, n, chunk_size):
        end_i = min(i + chunk_size, n)
        Q_chunk = Q[:, :, i:end_i, :]  # (B, h, chunk, d)
        
        # Compute attention scores for this chunk
        scores = torch.matmul(Q_chunk, K.transpose(-2, -1))
        scores = scores / math.sqrt(d)
        attn = F.softmax(scores, dim=-1)
        
        # Apply to values
        output[:, :, i:end_i, :] = torch.matmul(attn, V)
    
    return output

Another critical optimization is the use of PyTorch's scaled dot-product attention function introduced in PyTorch 2.0, which implements Flash Attention algorithms internally. This function provides significant speedups and memory reductions through kernel fusion and optimized memory access patterns.

import torch.nn.functional as F

def efficient_attention(Q, K, V, mask=None):
    """
    Use PyTorch's optimized scaled_dot_product_attention.
    Automatically uses Flash Attention when available.
    """
    # PyTorch 2.0+ provides optimized implementation
    output = F.scaled_dot_product_attention(
        Q, K, V,
        attn_mask=mask,
        dropout_p=0.1 if self.training else 0.0,
        is_causal=False
    )
    return output

This optimized implementation can reduce memory usage by up to 50\% and provide 2-3× speedups compared to the naive implementation, particularly for longer sequences.

Dimension Tracking Example

For BERT-base configuration ($d=768$, $h=12$):

Input: $(B, n, 768)$ where $B$ is batch size, $n$ is sequence length

After Q/K/V projection: $(B, n, 768)$

Reshape for heads: $(B, n, 12, 64) \to (B, 12, n, 64)$

Attention scores: $(B, 12, n, n)$

After applying to V: $(B, 12, n, 64)$

Concatenate heads: $(B, n, 12, 64) \to (B, n, 768)$

Output projection: $(B, n, 768)$

Position Encodings

Sinusoidal Encoding

Mathematical formula:

Implementation strategy:

1. Pre-compute position encoding matrix at initialization
2. Register as buffer (not parameter, doesn't need gradients)
3. Add to embeddings in forward pass

Learned Positional Embeddings

Alternative approach (BERT):

• Embedding layer: max\_len $\times$ d\_model
• Learn position representations during training
• More flexible but requires fixed max length

Masking Strategies

Causal Mask for GPT

Lower triangular mask:

causal_mask = torch.tril(torch.ones(seq_len, seq_len))
scores = scores.masked_fill(mask == 0, -1e9)

Padding Mask

# input_ids: (batch, seq_len)
# 0 indicates padding
pad_mask = (input_ids != 0).unsqueeze(1).unsqueeze(2)
# Shape: (batch, 1, 1, seq_len)

Training Optimizations

Fused Kernels for Layer Normalization

Standard PyTorch operations like layer normalization involve multiple kernel launches, each reading from and writing to global memory. Fusing these operations into a single kernel can provide substantial speedups by reducing memory bandwidth requirements. Modern deep learning frameworks provide fused implementations of common operations that should be used whenever possible.

The Apex library from NVIDIA provides highly optimized fused kernels for layer normalization and other operations. These implementations can be 2-3× faster than the standard PyTorch versions, particularly for smaller batch sizes where kernel launch overhead dominates.

# Standard PyTorch layer norm
layer_norm = nn.LayerNorm(d_model)

# Fused layer norm from Apex (faster)
try:
    from apex.normalization import FusedLayerNorm
    layer_norm = FusedLayerNorm(d_model)
except ImportError:
    # Fall back to standard implementation
    layer_norm = nn.LayerNorm(d_model)

Similarly, fused dropout and bias addition can be combined with other operations to reduce memory traffic. The key principle is to minimize the number of separate kernel launches and memory accesses by combining operations that naturally occur together in the computation graph.

Mixed Precision Training

Mixed precision training uses 16-bit floating point (FP16) for most operations while maintaining 32-bit precision for critical computations. This approach provides substantial benefits in terms of both memory usage and computational speed, particularly on modern GPUs with dedicated tensor cores optimized for FP16 operations.

Benefits:

• 2× memory reduction for activations and gradients
• 2-3× training speedup on modern GPUs with tensor cores
• Same final accuracy with proper loss scaling
• Enables training larger models or using larger batch sizes

PyTorch provides automatic mixed precision (AMP) through the torch.cuda.amp module, which automatically handles the conversion between FP16 and FP32 as needed. The implementation requires minimal code changes and provides automatic loss scaling to prevent gradient underflow.

from torch.cuda.amp import autocast, GradScaler

# Initialize gradient scaler for loss scaling
scaler = GradScaler()

# Training loop
for batch in dataloader:
    optimizer.zero_grad()
    
    # Forward pass in mixed precision
    with autocast():
        outputs = model(batch['input_ids'])
        loss = criterion(outputs, batch['labels'])
    
    # Backward pass with scaled loss
    scaler.scale(loss).backward()
    
    # Unscale gradients and clip
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    
    # Update weights
    scaler.step(optimizer)
    scaler.update()

The gradient scaler automatically adjusts the loss scaling factor to maintain numerical stability. It increases the scale when no overflow is detected and decreases it when overflow occurs, ensuring that gradients remain in a representable range for FP16 arithmetic.

For BERT-base training, mixed precision typically reduces memory usage from approximately 16GB to 8GB per GPU while maintaining the same final accuracy. The speedup varies depending on the GPU architecture, with Volta and newer architectures providing the largest benefits due to their tensor cores.

Gradient Accumulation

Gradient accumulation enables training with effective batch sizes larger than what fits in GPU memory by accumulating gradients over multiple forward-backward passes before updating weights. This technique is essential for training large models or when hardware constraints limit the physical batch size.

Purpose: Simulate large batch sizes on limited memory

Effective batch size:

The implementation requires careful handling of gradient normalization to ensure that the accumulated gradients have the correct scale. Each loss value should be divided by the number of accumulation steps so that the final gradient magnitude matches what would be obtained with a single large batch.

accumulation_steps = 8
optimizer.zero_grad()

for i, batch in enumerate(dataloader):
    # Forward pass
    with autocast():
        outputs = model(batch['input_ids'])
        loss = criterion(outputs, batch['labels'])
        # Scale loss by accumulation steps
        loss = loss / accumulation_steps
    
    # Backward pass
    scaler.scale(loss).backward()
    
    # Update weights every accumulation_steps
    if (i + 1) 
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

Example:

• Physical batch: 32 (fits in GPU)
• Accumulation steps: 8
• Effective batch: 256

This approach allows training with large effective batch sizes that would otherwise require multiple GPUs or be impossible due to memory constraints. The trade-off is increased training time, as the optimizer updates occur less frequently.

Gradient Checkpointing

Gradient checkpointing trades computation for memory by selectively storing only a subset of activations during the forward pass and recomputing the others during the backward pass. This technique can dramatically reduce memory usage, enabling training of much larger models or longer sequences at the cost of increased computation time.

Trade computation for memory:

• Don't store all activations during forward pass
• Recompute intermediate activations during backward pass
• Enables larger models or longer sequences
• Typically 20-30\% slower but saves 40-50\% memory

PyTorch provides gradient checkpointing through the torch.utils.checkpoint module. The key is to wrap transformer layers or blocks in checkpoint functions, which handle the recomputation automatically during the backward pass.

from torch.utils.checkpoint import checkpoint

class TransformerLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, use_checkpoint=False):
        super().__init__()
        self.use_checkpoint = use_checkpoint
        self.attention = MultiHeadAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
    
    def forward(self, x, mask=None):
        if self.use_checkpoint and self.training:
            # Use gradient checkpointing
            return checkpoint(self._forward, x, mask)
        else:
            return self._forward(x, mask)
    
    def _forward(self, x, mask):
        # Attention block
        attn_out = self.attention(x, x, x, mask)
        x = self.norm1(x + attn_out)
        
        # Feed-forward block
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)
        return x

For a 12-layer BERT model, gradient checkpointing can reduce peak memory usage from approximately 16GB to 9GB, allowing training with longer sequences or larger batch sizes. The computational overhead is typically 20-30\%, which is often an acceptable trade-off for the memory savings.

Model Initialization

Best Practices

Weight initialization:

• Linear layers: Xavier/Glorot normal
• Embeddings: Normal(0, 0.02)
• Layer norm: gamma=1, beta=0

Special considerations:

• Scale residual connections by $1/\sqrt{N_{\text{layers}}}$
• Weight tying: LM head shares embeddings
• Careful initialization prevents gradient issues

Memory Profiling and Optimization

Understanding Memory Usage

Memory consumption in transformer training comes from several sources: model parameters, optimizer states, activations, gradients, and temporary buffers. Understanding the breakdown of memory usage is essential for effective optimization. For a BERT-base model with 110M parameters, the memory requirements can be substantial even before considering batch data.

The model parameters themselves occupy relatively little memory compared to other components. With 110M parameters stored in FP32, the parameters require approximately 440MB. However, the Adam optimizer maintains two additional states per parameter (first and second moments), tripling the parameter memory to 1.3GB. Activations stored during the forward pass for gradient computation typically consume the largest portion of memory, scaling with both sequence length and batch size.

PyTorch provides comprehensive memory profiling tools through the torch.cuda module. The memory\_summary function provides detailed information about current memory allocation, including cached memory, allocated memory, and peak memory usage.

import torch

# Profile memory usage during training
def profile_memory(model, batch_size, seq_len, device='cuda'):
    torch.cuda.reset_peak_memory_stats(device)
    torch.cuda.empty_cache()
    
    # Create sample input
    input_ids = torch.randint(0, 30000, (batch_size, seq_len)).to(device)
    
    # Forward pass
    outputs = model(input_ids)
    loss = outputs.mean()
    
    print(f"After forward pass:")
    print(f"Allocated: {torch.cuda.memory_allocated(device) / 1e9:.2f} GB")
    print(f"Reserved: {torch.cuda.memory_reserved(device) / 1e9:.2f} GB")
    
    # Backward pass
    loss.backward()
    
    print(f"\nAfter backward pass:")
    print(f"Allocated: {torch.cuda.memory_allocated(device) / 1e9:.2f} GB")
    print(f"Peak: {torch.cuda.max_memory_allocated(device) / 1e9:.2f} GB")
    
    # Detailed summary
    print("\nDetailed memory summary:")
    print(torch.cuda.memory_summary(device))

Identifying Memory Bottlenecks

The first step in optimization is identifying where memory is being consumed. For transformer models, the attention mechanism typically dominates memory usage due to the quadratic scaling of attention scores with sequence length. A single attention layer with sequence length 512 and 12 heads requires approximately 12MB for attention scores alone, and this scales quadratically with sequence length.

Activation memory can be estimated using the formula:

where $B$ is batch size, $n$ is sequence length, $d$ is model dimension, $L$ is number of layers, and the factor of 2 accounts for both attention and feed-forward activations. For BERT-base with batch size 32 and sequence length 512, this amounts to approximately 9GB in FP32 or 4.5GB in FP16.

Optimization Strategies

Several strategies can dramatically reduce memory consumption. The most effective approach combines multiple techniques tailored to the specific bottlenecks identified through profiling.

Strategy 1: Reduce Sequence Length

The quadratic scaling of attention with sequence length makes this the most impactful optimization for long sequences. Reducing sequence length from 512 to 256 reduces attention memory by 4× and total activation memory by 2×. When possible, use techniques like sliding windows or hierarchical attention to process longer documents without materializing full attention matrices.

Strategy 2: Optimize Batch Size

Finding the optimal batch size requires balancing memory usage with computational efficiency. Larger batches improve GPU utilization but consume more memory. Use gradient accumulation to achieve large effective batch sizes while keeping physical batch sizes manageable.

def find_optimal_batch_size(model, seq_len, device='cuda'):
    """Binary search to find maximum batch size that fits in memory."""
    min_batch = 1
    max_batch = 256
    optimal_batch = 1
    
    while min_batch <= max_batch:
        batch_size = (min_batch + max_batch) // 2
        torch.cuda.empty_cache()
        
        try:
            # Test if this batch size fits
            input_ids = torch.randint(0, 30000, 
                                     (batch_size, seq_len)).to(device)
            outputs = model(input_ids)
            loss = outputs.mean()
            loss.backward()
            
            optimal_batch = batch_size
            min_batch = batch_size + 1
        except RuntimeError as e:
            if "out of memory" in str(e):
                max_batch = batch_size - 1
            else:
                raise e
    
    return optimal_batch

Strategy 3: Layer-wise Optimization

Different layers have different memory characteristics. Attention layers consume more memory than feed-forward layers due to the attention score matrix. Applying gradient checkpointing selectively to attention layers can provide most of the memory benefits with less computational overhead than checkpointing all layers.

Case Study: Optimizing BERT-base

Consider optimizing BERT-base training to reduce memory from 16GB to 8GB while maintaining training throughput. The baseline configuration uses batch size 32, sequence length 512, and FP32 precision.

Baseline measurements:

• Memory usage: 16.2 GB
• Training speed: 120 samples/second
• Parameters: 440 MB
• Optimizer states: 880 MB
• Activations: 9.1 GB
• Gradients: 4.5 GB
• Other: 1.3 GB

Optimization steps:

First, enable mixed precision training. This immediately reduces activation and gradient memory by 50\%, bringing total memory to approximately 10GB. The training speed increases to 280 samples/second due to tensor core utilization.

Second, apply gradient checkpointing to all transformer layers. This reduces activation memory by an additional 40\%, bringing total memory to 7.8GB. Training speed decreases to 220 samples/second due to recomputation overhead.

Third, optimize the batch size. With the memory savings, we can increase batch size to 48, improving GPU utilization. Final measurements show 7.9GB memory usage and 310 samples/second throughput.

Final measurements:

• Memory usage: 7.9 GB (51\% reduction)
• Training speed: 310 samples/second (2.6× faster)
• Same final accuracy after convergence

This optimization demonstrates that combining multiple techniques can achieve substantial improvements in both memory efficiency and training speed without sacrificing model quality.

Debugging Transformers

Common Issues

1. Dimension mismatches:

• Check shapes at each operation
• Use assertions for critical dimensions
• Print intermediate tensor shapes

2. NaN/Inf in training:

• Too high learning rate
• Gradient explosion (add clipping)
• Numerical instability in softmax (check mask values)

3. Slow convergence:

• Insufficient warmup
• Bad initialization
• Learning rate too low

4. Memory issues:

• Reduce batch size
• Use gradient checkpointing
• Enable mixed precision
• Reduce sequence length

Validation Checks

Sanity checks before full training:

1. Overfit single batch (should reach near-zero loss)
2. Check gradient norms are reasonable
3. Verify attention weights sum to 1
4. Test with different sequence lengths
5. Profile memory usage

Inference Optimization

TorchScript Compilation

Inference optimization is critical for deploying transformer models in production environments where latency and throughput requirements are stringent. TorchScript provides a way to serialize and optimize PyTorch models for inference, removing Python overhead and enabling additional optimizations.

The torch.jit.script function traces the model's execution and converts it to an intermediate representation that can be optimized and executed more efficiently. This process eliminates Python interpreter overhead and enables fusion of operations that would otherwise require multiple kernel launches.

import torch.jit as jit

class TransformerForInference(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transformer = TransformerModel(config)
    
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        # Type annotations required for TorchScript
        return self.transformer(input_ids)

# Create and script the model
model = TransformerForInference(config)
model.eval()

# Convert to TorchScript
scripted_model = jit.script(model)

# Save for deployment
scripted_model.save('model_scripted.pt')

# Load and use
loaded_model = jit.load('model_scripted.pt')
with torch.no_grad():
    output = loaded_model(input_ids)

TorchScript compilation typically provides 10-30\% speedup for transformer inference, with larger models seeing greater benefits. The compilation process also validates that the model can be executed without Python dependencies, which is essential for deployment in production environments.

KV Cache for Autoregressive Generation

Autoregressive generation in models like GPT requires computing attention over all previous tokens at each step. Without optimization, this results in redundant computation as the keys and values for previous tokens are recomputed at every step. Implementing a KV cache stores these values and reuses them, dramatically reducing computation. For systems-level KV cache management at scale (PagedAttention, memory fragmentation, batch scheduling), see Chapter~[ref].

class GPTWithKVCache(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerLayer(config) for _ in range(config.num_layers)
        ])
        self.embed = nn.Embedding(config.vocab_size, config.d_model)
    
    def forward(self, input_ids, past_key_values=None, use_cache=False):
        """
        Args:
            input_ids: (batch, seq_len) - new tokens to process
            past_key_values: List of (key, value) tuples from previous steps
            use_cache: Whether to return key-value cache
        """
        batch_size, seq_len = input_ids.shape
        
        # Embed input tokens
        hidden_states = self.embed(input_ids)
        
        # Initialize cache if not provided
        if past_key_values is None:
            past_key_values = [None] * len(self.layers)
        
        # Store new key-values if caching
        present_key_values = [] if use_cache else None
        
        # Process through layers
        for i, (layer, past_kv) in enumerate(
            zip(self.layers, past_key_values)):
            
            # Layer forward with cache
            hidden_states, new_kv = layer(
                hidden_states, 
                past_key_value=past_kv,
                use_cache=use_cache
            )
            
            if use_cache:
                present_key_values.append(new_kv)
        
        return hidden_states, present_key_values

class TransformerLayerWithCache(nn.Module):
    def forward(self, x, past_key_value=None, use_cache=False):
        # Compute Q, K, V
        Q = self.q_proj(x)
        K = self.k_proj(x)
        V = self.v_proj(x)
        
        # Use cached K, V if available
        if past_key_value is not None:
            past_K, past_V = past_key_value
            K = torch.cat([past_K, K], dim=1)
            V = torch.cat([past_V, V], dim=1)
        
        # Compute attention
        attn_output = self.attention(Q, K, V)
        
        # Return new cache if requested
        new_kv = (K, V) if use_cache else None
        return attn_output, new_kv

# Generation with KV cache
def generate_with_cache(model, input_ids, max_length=100):
    """Generate tokens using KV cache for efficiency."""
    past_key_values = None
    
    for _ in range(max_length):
        # Only process new token (or all tokens on first step)
        if past_key_values is None:
            current_input = input_ids
        else:
            current_input = input_ids[:, -1:]
        
        # Forward pass with cache
        logits, past_key_values = model(
            current_input,
            past_key_values=past_key_values,
            use_cache=True
        )
        
        # Sample next token
        next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
        input_ids = torch.cat([input_ids, next_token], dim=1)
        
        # Stop if EOS token
        if next_token.item() == eos_token_id:
            break
    
    return input_ids

KV caching reduces the computational complexity of generating $n$ tokens from $O(n^2)$ to $O(n)$, providing speedups of 5-10× for typical generation lengths. The memory overhead is proportional to the sequence length and number of layers, typically requiring 1-2GB for a GPT-2 sized model generating 1000 tokens.

ONNX Export

ONNX (Open Neural Network Exchange) provides a standardized format for representing neural networks, enabling deployment across different frameworks and hardware platforms. Exporting to ONNX allows using optimized inference engines like ONNX Runtime, which can provide substantial speedups.

import torch.onnx

def export_to_onnx(model, output_path, batch_size=1, seq_len=128):
    """Export PyTorch model to ONNX format."""
    model.eval()
    
    # Create dummy input
    dummy_input = torch.randint(
        0, model.config.vocab_size, 
        (batch_size, seq_len)
    )
    
    # Export to ONNX
    torch.onnx.export(
        model,
        dummy_input,
        output_path,
        input_names=['input_ids'],
        output_names=['logits'],
        dynamic_axes={
            'input_ids': {0: 'batch', 1: 'sequence'},
            'logits': {0: 'batch', 1: 'sequence'}
        },
        opset_version=14,
        do_constant_folding=True
    )

# Use ONNX Runtime for inference
import onnxruntime as ort

# Create inference session
session = ort.InferenceSession(
    'model.onnx',
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)

# Run inference
input_ids = torch.randint(0, 30000, (1, 128)).numpy()
outputs = session.run(
    ['logits'],
    {'input_ids': input_ids}
)

ONNX Runtime typically delivers 1.5--2$\times$ speedup over PyTorch for transformer inference through operator fusion, memory layout optimization, and hardware-specific kernel selection. Combined with INT8 quantization, 3--4$\times$ speedup is achievable (see Chapter~[ref] for quantization fundamentals).

TensorRT Optimization

NVIDIA TensorRT provides highly optimized inference for NVIDIA GPUs through aggressive kernel fusion, precision calibration, and dynamic tensor memory management. TensorRT can provide 2-5× speedup over standard PyTorch inference for transformer models.

# Convert ONNX to TensorRT
import tensorrt as trt

def build_tensorrt_engine(onnx_path, engine_path, fp16_mode=True):
    """Build TensorRT engine from ONNX model."""
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )
    parser = trt.OnnxParser(network, logger)
    
    # Parse ONNX model
    with open(onnx_path, 'rb') as f:
        if not parser.parse(f.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None
    
    # Configure builder
    config = builder.create_builder_config()
    config.max_workspace_size = 4 << 30  # 4GB
    
    if fp16_mode:
        config.set_flag(trt.BuilderFlag.FP16)
    
    # Build engine
    engine = builder.build_engine(network, config)
    
    # Save engine
    with open(engine_path, 'wb') as f:
        f.write(engine.serialize())
    
    return engine

# Use TensorRT for inference
def infer_with_tensorrt(engine_path, input_ids):
    """Run inference using TensorRT engine."""
    logger = trt.Logger(trt.Logger.WARNING)
    
    with open(engine_path, 'rb') as f:
        runtime = trt.Runtime(logger)
        engine = runtime.deserialize_cuda_engine(f.read())
    
    context = engine.create_execution_context()
    
    # Allocate buffers
    inputs, outputs, bindings = allocate_buffers(engine)
    
    # Copy input data
    inputs[0].host = input_ids.cpu().numpy()
    
    # Run inference
    outputs = do_inference(
        context, bindings, inputs, outputs, stream
    )
    
    return outputs[0]

TensorRT optimization is particularly effective for deployment scenarios where inference latency is critical. The optimization process includes layer fusion, precision calibration for INT8 quantization, and kernel auto-tuning for the specific GPU architecture.

Quantization

Quantization reduces model size and inference latency by using lower precision representations for weights and activations. PyTorch supports several quantization approaches, from simple dynamic quantization to full quantization-aware training. For the theoretical foundations of quantization (precision formats, scale factors, zero-points) and pruning/distillation techniques, see Chapter~[ref].

import torch.quantization as quant

# Dynamic quantization (easiest, good for LSTM/Transformer)
def dynamic_quantize(model):
    """Apply dynamic quantization to linear layers."""
    quantized_model = quant.quantize_dynamic(
        model,
        {nn.Linear},  # Quantize linear layers
        dtype=torch.qint8
    )
    return quantized_model

# Static quantization (requires calibration)
def static_quantize(model, calibration_dataloader):
    """Apply static quantization with calibration."""
    model.eval()
    
    # Specify quantization configuration
    model.qconfig = quant.get_default_qconfig('fbgemm')
    
    # Prepare model for quantization
    model_prepared = quant.prepare(model)
    
    # Calibrate with representative data
    with torch.no_grad():
        for batch in calibration_dataloader:
            model_prepared(batch['input_ids'])
    
    # Convert to quantized model
    model_quantized = quant.convert(model_prepared)
    return model_quantized

# Quantization-aware training
def quantization_aware_training(model, train_dataloader):
    """Train model with quantization simulation."""
    model.train()
    model.qconfig = quant.get_default_qat_qconfig('fbgemm')
    
    # Prepare for QAT
    model_prepared = quant.prepare_qat(model)
    
    # Train normally
    for epoch in range(num_epochs):
        for batch in train_dataloader:
            outputs = model_prepared(batch['input_ids'])
            loss = criterion(outputs, batch['labels'])
            loss.backward()
            optimizer.step()
    
    # Convert to quantized model
    model_quantized = quant.convert(model_prepared.eval())
    return model_quantized

Dynamic quantization typically provides 2-3× speedup and 4× model size reduction with minimal accuracy loss for transformer models. Static quantization and quantization-aware training can provide additional benefits but require more careful tuning and calibration data.

Inference Benchmarking

Comprehensive benchmarking is essential for understanding the trade-offs between different optimization techniques. The following framework measures latency, throughput, and memory usage across different configurations.

def benchmark_inference(model, batch_sizes, seq_lengths, num_runs=100):
    """Comprehensive inference benchmarking."""
    results = []
    model.eval()
    
    for batch_size in batch_sizes:
        for seq_len in seq_lengths:
            # Create input
            input_ids = torch.randint(
                0, 30000, (batch_size, seq_len)
            ).cuda()
            
            # Warmup
            with torch.no_grad():
                for _ in range(10):
                    _ = model(input_ids)
            
            # Benchmark
            torch.cuda.synchronize()
            start = time.time()
            
            with torch.no_grad():
                for _ in range(num_runs):
                    _ = model(input_ids)
            
            torch.cuda.synchronize()
            elapsed = time.time() - start
            
            # Calculate metrics
            latency_ms = (elapsed / num_runs) * 1000
            throughput = (batch_size * num_runs) / elapsed
            memory_mb = torch.cuda.max_memory_allocated() / 1e6
            
            results.append({
                'batch_size': batch_size,
                'seq_len': seq_len,
                'latency_ms': latency_ms,
                'throughput': throughput,
                'memory_mb': memory_mb
            })
            
            torch.cuda.reset_peak_memory_stats()
    
    return results

# Compare optimization techniques
def compare_optimizations(base_model, config):
    """Compare different optimization approaches."""
    models = {
        'baseline': base_model,
        'torchscript': jit.script(base_model),
        'quantized': dynamic_quantize(base_model),
        'fp16': base_model.half()
    }
    
    results = {}
    for name, model in models.items():
        print(f"Benchmarking {name}...")
        results[name] = benchmark_inference(
            model, 
            batch_sizes=[1, 8, 32],
            seq_lengths=[128, 512]
        )
    
    return results

Typical results for BERT-base inference optimization show that combining TorchScript, FP16, and dynamic quantization can achieve 5-8× speedup with less than 1\% accuracy degradation, making deployment feasible for latency-sensitive applications.

Complete Training Pipeline

Training Script Structure

config = {
    'd_model': 768,
    'num_heads': 12,
    'num_layers': 12,
    'd_ff': 3072,
    'vocab_size': 30000,
    'max_seq_len': 512,
    'dropout': 0.1,
    'batch_size': 32,
    'learning_rate': 1e-4,
    'warmup_steps': 10000,
    'max_steps': 1000000
}

model = BERTModel(**config)
model = model.to(device)

from torch.optim import AdamW
optimizer = AdamW(
    model.parameters(),
    lr=config['learning_rate'],
    betas=(0.9, 0.999),
    eps=1e-6,
    weight_decay=0.01
)

from transformers import get_linear_schedule_with_warmup
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=config['warmup_steps'],
    num_training_steps=config['max_steps']
)

5. Training loop with all optimizations:

• Mixed precision
• Gradient accumulation
• Gradient clipping
• Checkpointing
• Logging

Production Optimizations

The implementations shown above can be enhanced with several production optimizations:

• Combined QKV projection: Concatenating $\mW_Q$, $\mW_K$, $\mW_V$ into a single $d_{\text{model}} \times 3d_{\text{model}}$ matrix reduces kernel launch overhead and improves memory coalescing.
• Fused layer normalization: Libraries like NVIDIA Apex provide FusedLayerNorm that combines the normalization computation into a single GPU kernel, reducing memory traffic by $\sim$30\%.
• Pre-norm architecture: Placing layer normalization before (rather than after) attention and feed-forward sublayers improves training stability for deep models and is standard in GPT-2 and later architectures.
• FlashAttention: Using memory-efficient attention implementations (see Chapter~16) dramatically reduces memory consumption for long sequences.

Putting It Together

A complete training loop combines the components from this chapter: the model architecture with mixed precision via torch.cuda.amp, gradient accumulation for large effective batch sizes, gradient checkpointing for memory savings, and learning rate scheduling with warmup. The training pipeline from Section~[ref] demonstrates this integration. For production deployments, add gradient clipping, periodic checkpointing, and distributed training via torch.nn.parallel.DistributedDataParallel.

Comprehensive Benchmarks

The following benchmarks demonstrate the impact of various optimizations on memory usage and training speed for a BERT-base model.

Baseline Configuration:

• Model: BERT-base (110M parameters)
• Batch size: 32
• Sequence length: 512
• Precision: FP32
• Hardware: NVIDIA A100 40GB

Optimization Results:

Inference Optimization Results:

These benchmarks demonstrate that combining multiple optimization techniques can achieve substantial improvements in both training and inference performance. The key insight is that different optimizations address different bottlenecks, and the cumulative effect can be dramatic when applied systematically.

Distributed Training

Data Parallel

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

Effective batch size: $B \times N_{\text{GPUs}}$

Distributed Data Parallel (DDP)

More efficient than DataParallel:

• One process per GPU
• Gradient synchronization via all-reduce
• Better scaling to multiple nodes

Setup requires:

1. Initialize process group
2. Wrap model in DistributedDataParallel
3. Use DistributedSampler for data
4. Synchronize across processes

Performance Optimization

DataLoader Optimization

The PyTorch DataLoader is often a bottleneck in training pipelines, particularly when data preprocessing is complex or I/O is slow. Proper configuration of the DataLoader can significantly improve training throughput by ensuring that data loading does not become the limiting factor.

The num\_workers parameter controls how many subprocesses are used for data loading. Setting this too low results in the GPU waiting for data, while setting it too high can cause excessive CPU and memory usage. A good starting point is to use 4-8 workers per GPU, but the optimal value depends on the specific dataset and preprocessing pipeline.

from torch.utils.data import DataLoader

# Optimized DataLoader configuration
dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=8,           # Parallel data loading
    pin_memory=True,         # Faster GPU transfer
    persistent_workers=True, # Keep workers alive between epochs
    prefetch_factor=2        # Prefetch batches per worker
)

The pin\_memory option allocates data in pinned (page-locked) memory, which enables faster transfers to the GPU using asynchronous DMA transfers. This can provide 20-30\% speedup for data transfer, particularly beneficial when the model is small relative to the batch size.

Persistent workers keep the worker processes alive between epochs, avoiding the overhead of spawning new processes. This is particularly beneficial for datasets with expensive initialization or when using many workers.

Asynchronous Data Transfer

Overlapping data transfer with computation can hide data transfer latency. PyTorch supports non-blocking transfers that allow the CPU to continue executing while data is being copied to the GPU.

for batch in dataloader:
    # Non-blocking transfer to GPU
    input_ids = batch['input_ids'].to(device, non_blocking=True)
    labels = batch['labels'].to(device, non_blocking=True)
    
    # Computation can start while transfer completes
    with autocast():
        outputs = model(input_ids)
        loss = criterion(outputs, labels)

This technique is most effective when combined with pinned memory, as it enables true asynchronous transfers. The speedup depends on the ratio of transfer time to computation time, with larger models benefiting more as computation dominates.

Profiling with torch.profiler

Understanding where time is spent during training is essential for effective optimization. PyTorch's profiler provides detailed information about CPU and GPU operations, memory usage, and kernel execution times.

from torch.profiler import profile, ProfilerActivity, schedule

# Configure profiler
profiler_schedule = schedule(
    wait=1,      # Skip first batch
    warmup=1,    # Warmup for 1 batch
    active=3,    # Profile 3 batches
    repeat=2     # Repeat cycle twice
)

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=profiler_schedule,
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log'),
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    for step, batch in enumerate(dataloader):
        if step >= 10:  # Profile first 10 batches
            break
        
        # Training step
        outputs = model(batch['input_ids'].to(device))
        loss = outputs.mean()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
        prof.step()  # Signal end of iteration

# Print summary
print(prof.key_averages().table(
    sort_by="cuda_time_total", row_limit=10))

The profiler output identifies operations that consume the most time, enabling targeted optimization. Common bottlenecks include inefficient attention implementations, excessive memory allocations, and CPU-GPU synchronization points.

Batch Size Tuning

Batch size has a complex relationship with training speed and model quality. Larger batches improve GPU utilization and reduce the number of optimizer steps, but may require learning rate adjustments and can affect convergence.

The optimal batch size maximizes GPU utilization without causing memory overflow. For transformer models, GPU utilization typically plateaus at batch sizes where the GPU is fully occupied, with further increases providing diminishing returns.

def benchmark_batch_sizes(model, seq_len, device='cuda'):
    """Benchmark training speed for different batch sizes."""
    results = []
    
    for batch_size in [8, 16, 32, 64, 128]:
        try:
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()
            
            # Warmup
            for _ in range(5):
                input_ids = torch.randint(0, 30000, 
                                         (batch_size, seq_len)).to(device)
                outputs = model(input_ids)
                loss = outputs.mean()
                loss.backward()
            
            # Benchmark
            torch.cuda.synchronize()
            start = time.time()
            
            for _ in range(20):
                input_ids = torch.randint(0, 30000, 
                                         (batch_size, seq_len)).to(device)
                outputs = model(input_ids)
                loss = outputs.mean()
                loss.backward()
            
            torch.cuda.synchronize()
            elapsed = time.time() - start
            
            samples_per_sec = (20 * batch_size) / elapsed
            memory_gb = torch.cuda.max_memory_allocated() / 1e9
            
            results.append({
                'batch_size': batch_size,
                'samples_per_sec': samples_per_sec,
                'memory_gb': memory_gb
            })
            
        except RuntimeError as e:
            if "out of memory" in str(e):
                break
    
    return results

Compilation with torch.compile

PyTorch 2.0 introduces torch.compile, which uses TorchDynamo and TorchInductor to compile models into optimized kernels. This can provide substantial speedups with minimal code changes.

# Compile model for faster execution
model = torch.compile(model, mode='max-autotune')

# Training proceeds as normal
for batch in dataloader:
    outputs = model(batch['input_ids'])
    loss = outputs.mean()
    loss.backward()
    optimizer.step()

The compilation process analyzes the model's computation graph and generates optimized CUDA kernels. The first iteration is slow due to compilation overhead, but subsequent iterations benefit from the optimized code. Speedups of 20-50\% are common for transformer models, with larger models typically seeing greater benefits.

Distributed Training Implementation

Understanding Distributed Strategies

Distributed training enables training on multiple GPUs or machines, dramatically reducing training time for large models. PyTorch provides several distributed training strategies, each with different trade-offs and use cases.

Data parallelism replicates the model on each GPU and distributes different batches of data to each replica. Gradients are synchronized across replicas after the backward pass, ensuring all replicas maintain identical weights. This approach scales well when the model fits in a single GPU's memory and is the most commonly used distributed training strategy.

Model parallelism splits the model itself across multiple GPUs, with different layers or components on different devices. This is necessary when the model is too large to fit on a single GPU but is more complex to implement and can suffer from poor GPU utilization due to sequential dependencies.

DistributedDataParallel Setup

DistributedDataParallel (DDP) is PyTorch's recommended approach for multi-GPU training. It provides better performance than DataParallel through more efficient gradient synchronization and support for multi-node training.

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup_distributed(rank, world_size):
    """Initialize distributed training."""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    
    # Initialize process group
    dist.init_process_group(
        backend='nccl',  # Use NCCL for GPU training
        rank=rank,
        world_size=world_size
    )

def cleanup_distributed():
    """Clean up distributed training."""
    dist.destroy_process_group()

def train_distributed(rank, world_size, model, dataset):
    """Training function for each process."""
    setup_distributed(rank, world_size)
    
    # Move model to GPU
    device = torch.device(f'cuda:{rank}')
    model = model.to(device)
    
    # Wrap model in DDP
    model = DDP(model, device_ids=[rank])
    
    # Create distributed sampler
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )
    
    # Create dataloader with distributed sampler
    dataloader = DataLoader(
        dataset,
        batch_size=32,
        sampler=sampler,
        num_workers=4,
        pin_memory=True
    )
    
    optimizer = AdamW(model.parameters(), lr=1e-4)
    
    # Training loop
    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # Shuffle differently each epoch
        
        for batch in dataloader:
            input_ids = batch['input_ids'].to(device)
            labels = batch['labels'].to(device)
            
            outputs = model(input_ids)
            loss = criterion(outputs, labels)
            
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    
    cleanup_distributed()

# Launch training on multiple GPUs
if __name__ == '__main__':
    world_size = torch.cuda.device_count()
    torch.multiprocessing.spawn(
        train_distributed,
        args=(world_size, model, dataset),
        nprocs=world_size,
        join=True
    )

Gradient Synchronization

DDP automatically synchronizes gradients across all processes during the backward pass using efficient all-reduce operations. The synchronization happens in parallel with the backward pass through gradient bucketing, which groups gradients into buckets and overlaps communication with computation.

The effective learning rate in distributed training should typically be scaled with the number of GPUs to maintain the same optimization dynamics. If training with 8 GPUs, the learning rate should be multiplied by 8, or equivalently, the batch size per GPU should be kept constant and gradient accumulation used to achieve the same effective batch size.

Scaling Efficiency

Distributed training efficiency is measured by scaling efficiency, which compares actual speedup to ideal linear speedup. Perfect scaling would achieve 8× speedup with 8 GPUs, but communication overhead and synchronization typically reduce this.

def measure_scaling_efficiency(model, batch_size, seq_len):
    """Measure scaling efficiency across different GPU counts."""
    results = {}
    
    # Single GPU baseline
    single_gpu_time = benchmark_single_gpu(model, batch_size, seq_len)
    results[1] = {
        'time': single_gpu_time,
        'speedup': 1.0,
        'efficiency': 1.0
    }
    
    # Multi-GPU measurements
    for num_gpus in [2, 4, 8]:
        if num_gpus > torch.cuda.device_count():
            break
        
        multi_gpu_time = benchmark_multi_gpu(
            model, batch_size, seq_len, num_gpus)
        speedup = single_gpu_time / multi_gpu_time
        efficiency = speedup / num_gpus
        
        results[num_gpus] = {
            'time': multi_gpu_time,
            'speedup': speedup,
            'efficiency': efficiency
        }
    
    return results

For transformer models, scaling efficiency typically ranges from 85-95\% for 2-8 GPUs on a single node, with larger models achieving better efficiency due to higher computation-to-communication ratios. Multi-node training introduces additional communication overhead, with efficiency typically dropping to 70-85\% depending on network bandwidth and model size.

Exercises

Solutions

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

Prompt: "The future of artificial intelligence"
Output: "The future of artificial intelligence will be shaped by 
advances in deep learning and transformer architectures. These 
models have demonstrated remarkable capabilities in natural 
language understanding and generation..."