Training Speed Optimization Techniques for LLM Pretraining: A PyTorch Performance Guide

Use torch.compile, bfloat16 mixed precision, and pin-memory DataLoaders for single-GPU speedups, then scale to multiple GPUs with DistributedDataParallel (DDP), DistributedSampler, and global token-throughput aggregation to maximize training throughput.

Optimizing training throughput is critical when pretraining large language models from scratch, as even small efficiency gains compound into significant time and cost savings. The rasbt/LLMs-from-scratch repository demonstrates concrete training speed optimization techniques for LLM pretraining, implementing both single-GPU acceleration methods and multi-GPU distributed training strategies in its Chapter 5 training scripts. By leveraging PyTorch 2.0 compilation features and Distributed Data Parallel (DDP) patterns, you can achieve substantial speedups with minimal changes to your transformer training loop.

Single-GPU Training Optimizations

The repository's ch05/10_llm-training-speed/01_opt_single_gpu.py demonstrates how to maximize throughput on a single accelerator before scaling out. These optimizations require minimal code changes but deliver significant performance improvements.

PyTorch 2.0 Compilation and Mixed Precision

torch.compile generates a fused, optimized execution graph that removes Python overhead and merges kernels, typically yielding 10-30% speedups on modern GPUs. In the repository, this is applied to the GPTModel immediately after instantiation:

model = GPTModel(gpt_config)      # ↝ ch05/10_llm-training-speed/01_opt_single_gpu.py L13

model = torch.compile(model)      # ↝ L13-L14

BFloat16 mixed precision reduces memory traffic and accelerates matrix multiplications on Tensor Core-equipped GPUs (A100, H100). The implementation converts the model after compilation:

model = model.to(torch.bfloat16)  # ↝ L15-L16

For optimal Tensor Core utilization, the script also sets high-precision FP32 matmul mode when available:

torch.set_float32_matmul_precision("high")  # ↝ L83-L89

Data Loading and Memory Optimization

CPU-GPU transfer bottlenecks are eliminated through page-locked memory and parallel data preprocessing. The DataLoader configuration in 01_opt_single_gpu.py uses both pin_memory=True and num_workers=4 to overlap tokenization with GPU training:

dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    drop_last=True,
    pin_memory=True,    # Enables page-locked host memory ↝ L55

    num_workers=4,      # Parallelizes tokenization ↝ L55

)

The dataset implementation in GPTDatasetV1.__init__ (lines 58-71) uses single-pass encoding with sliding-window chunking to avoid repeated tokenizer calls and maintain contiguous tensor storage.

Timing and Throughput Monitoring

Precise performance measurement relies on CUDA events rather than time.time() to exclude GPU kernel launch overhead. The train_model_simple_with_timing() function (lines 78-107) implements this pattern:

if use_cuda:
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    torch.cuda.synchronize()
    start.record()
    
    # ... training steps ...

    
    end.record()
    torch.cuda.synchronize()
    elapsed = start.elapsed_time(end) / 1000  # Convert to seconds

    tps = tokens_processed / elapsed

The script also monitors GPU memory usage at the end of each epoch using torch.cuda.memory_allocated and torch.cuda.memory_reserved (lines 127-132), enabling informed batch-size tuning.

Multi-GPU Scaling with Distributed Data Parallel

When scaling beyond a single GPU, ch05/10_llm-training-speed/02_opt_multi_gpu_ddp.py implements PyTorch's DistributedDataParallel (DDP) while retaining all single-GPU optimizations.

Process Initialization and Device Management

The ddp_setup() function (lines 27-48) initializes the NCCL backend on Linux (or GLOO on Windows) and binds each process to a dedicated GPU to prevent context-switching overhead:

def ddp_setup(rank, world_size):
    backend = "nccl" if torch.cuda.is_available() else "gloo"
    dist.init_process_group(backend, rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)  # ↝ L50

A barrier synchronization ensures all ranks wait for rank 0 to download the dataset before proceeding:

if rank == 0:
    # Download data

    pass
torch.distributed.barrier(device_ids=[device.index])  # ↝ L77

Distributed Sampling and Data Sharding

The DistributedSampler guarantees each GPU processes a unique, non-overlapping subset of the dataset, eliminating redundant work. This replaces standard shuffling:

sampler = DistributedSampler(dataset)
train_loader = DataLoader(
    dataset,
    batch_size=batch_size,
    sampler=sampler,      # ↝ L99

    shuffle=False,        # Shuffling handled by sampler ↝ L95

    pin_memory=True,
)

Critical for correctness, you must call set_epoch() each training epoch to ensure different shuffling across epochs:

for epoch in range(num_epochs):
    train_loader.sampler.set_epoch(epoch)  # ↝ L39-L40

The model is wrapped with DDP after compilation and precision conversion:

model = DDP(model, device_ids=[rank])  # ↝ L91-L92

Global Throughput Aggregation

To measure global tokens-per-second across all GPUs, the script uses all_reduce to sum token counts from each rank. This pattern appears in the evaluation loop (lines 71-84):

local_interval = total_tokens - last_tokens
local_tensor = torch.tensor([local_interval], device=device, dtype=torch.float)
global_tensor = local_tensor.clone()
dist.all_reduce(global_tensor, op=dist.ReduceOp.SUM)  # Sum across ranks

global_tps = global_tensor.item() / elapsed

Rank-0 logging prevents console spam by guarding print statements: if rank == 0: (lines 94-100, 120-124). The learning rate scales with world size (comment at line 71), typically multiplying the base rate by the number of GPUs to maintain effective batch-size scaling.

Implementation Examples

Single-GPU Optimized Training Loop

import torch
from torch.utils.data import DataLoader

# Model setup with compilation and BF16

model = GPTModel(gpt_config)
model = torch.compile(model)                    # Graph optimization

model = model.to("cuda", torch.bfloat16)      # Mixed precision

# Optimizer with fused kernels

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=5e-4,
    fused=True,                                 # Fused CUDA kernels

)

# Efficient data loading

train_loader = DataLoader(
    dataset,
    batch_size=32,
    pin_memory=True,                            # Fast host-to-GPU transfer

    num_workers=4,                              # Parallel preprocessing

)

Multi-GPU DDP Launch Configuration

Execute distributed training using torchrun:

torchrun --nproc_per_node=4 \
    ch05/10_llm-training-speed/02_opt_multi_gpu_ddp.py

DDP Training Setup

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

def main(rank, world_size):
    # Initialize process group

    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    
    # Build and wrap model

    model = GPTModel(config)
    model = torch.compile(model)
    model = model.to(rank, torch.bfloat16)
    model = DDP(model, device_ids=[rank])
    
    # Distributed sampler ensures unique data per rank

    sampler = DistributedSampler(train_dataset)
    loader = DataLoader(
        train_dataset,
        batch_size=32,
        sampler=sampler,
        pin_memory=True,
    )
    
    for epoch in range(epochs):
        sampler.set_epoch(epoch)  # Required for shuffling

        # Training loop...

Summary

• torch.compile and bfloat16 provide immediate single-GPU speedups by fusing kernels and reducing memory bandwidth.
• pin_memory=True and num_workers > 0 eliminate data loading bottlenecks by overlapping CPU preprocessing with GPU computation.
• CUDA events deliver accurate throughput measurements by properly synchronizing GPU kernels for timing.
• DistributedSampler and DDP enable linear scaling across multiple GPUs while preventing duplicate data processing.
• Global token aggregation via all_reduce provides true cluster-wide throughput metrics, while rank-0 logging keeps output readable.
• Barrier synchronization ensures data availability across all ranks before training begins.

Frequently Asked Questions

What is torch.compile and why does it speed up LLM training?

torch.compile is a PyTorch 2.0 feature that traces the model into optimized C++ graphs, eliminating Python interpreter overhead and fusing compatible operations into single CUDA kernels. In 01_opt_single_gpu.py (lines 13-14), wrapping the GPTModel with torch.compile() allows the compiler to optimize the transformer blocks statically, typically improving throughput by 10-30% on modern hardware without changing model behavior.

How does bfloat16 improve training speed compared to float32?

BFloat16 (Brain Floating Point) uses 16 bits instead of 32, halving memory bandwidth requirements for matrix multiplications while maintaining the same exponent range as FP32. As implemented in lines 15-16 of 01_opt_single_gpu.py, converting the model to torch.bfloat16 enables Tensor Core acceleration on A100/H100 GPUs, allowing more matrix operations per clock cycle and supporting larger batch sizes due to reduced memory footprint.

When should I use DistributedSampler versus regular shuffling?

Use DistributedSampler exclusively when training with DistributedDataParallel across multiple GPUs. As shown in 02_opt_multi_gpu_ddp.py (lines 95, 99), the sampler ensures each rank processes a disjoint subset of the dataset, eliminating redundant computation. Regular shuffling (shuffle=True) is appropriate only for single-GPU training; in DDP, you must set shuffle=False on the DataLoader and call sampler.set_epoch(epoch) each epoch to maintain proper randomization across the distributed dataset.

How do I calculate global tokens-per-second across multiple GPUs?

Calculate local tokens processed, then use torch.distributed.all_reduce to sum across all ranks. According to lines 71-84 in 02_opt_multi_gpu_ddp.py, you create a tensor of the local token count, call dist.all_reduce(global_tensor, op=dist.ReduceOp.SUM), then divide the summed value by the elapsed time. This yields the true global throughput rather than per-GPU statistics, accounting for the aggregate processing capacity of your entire cluster.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}