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

> Optimize LLM pretraining speed with PyTorch. Discover techniques like torch.compile, bfloat16, DDP, and throughput aggregation for faster model training.

- Repository: [Sebastian Raschka/LLMs-from-scratch](https://github.com/rasbt/LLMs-from-scratch)
- Tags: performance
- Published: 2026-05-12

---

**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`](https://github.com/rasbt/LLMs-from-scratch/blob/main/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:

```python
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:

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

```

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

```python
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`](https://github.com/rasbt/LLMs-from-scratch/blob/main/01_opt_single_gpu.py) uses both `pin_memory=True` and `num_workers=4` to overlap tokenization with GPU training:

```python
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:

```python
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`](https://github.com/rasbt/LLMs-from-scratch/blob/main/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:

```python
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:

```python
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:

```python
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:

```python
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:

```python
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):

```python
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

```python
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`:

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

```

### DDP Training Setup

```python
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`](https://github.com/rasbt/LLMs-from-scratch/blob/main/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`](https://github.com/rasbt/LLMs-from-scratch/blob/main/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`](https://github.com/rasbt/LLMs-from-scratch/blob/main/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`](https://github.com/rasbt/LLMs-from-scratch/blob/main/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.