Production LLM Training Systems: Megatron, Elastic Training, and Fault Tolerance

"Training a 405-billion-parameter model is not a research problem. It is a systems engineering problem with research inside it."

Scale, Industrially Scaled AI Agent

1. Megatron-LM and Multi-Axis Parallelism

Section 6.6 introduced tensor parallelism (TP), pipeline parallelism (PP), and data parallelism (DP) as independent strategies. In practice, training a frontier LLM requires composing all of these simultaneously. Megatron-LM, developed by NVIDIA (Shoeybi et al., 2020; Narayanan et al., 2021), is the most widely adopted framework for multi-axis parallelism in LLM pretraining. Its successor, Megatron-Core, provides a modular library that can be integrated into custom training loops.

1.1 The Four Parallelism Axes

A typical Megatron-LM configuration specifies four parallelism degrees: TP, PP, DP, and (optionally) CP. The total GPU count equals $\text{TP} \times \text{PP} \times \text{DP} \times \text{CP}$. Choosing the right combination depends on model size, sequence length, cluster topology, and interconnect bandwidth.

Tensor Parallelism (TP) splits individual matrix multiplications across GPUs. For a transformer's attention and MLP layers, TP partitions the weight matrices along the hidden dimension. This requires an all-reduce after each layer, so it works best within a node where NVLink provides 900 GB/s bandwidth (on H100 nodes). Typical TP degree: 2, 4, or 8 within a single node.

Pipeline Parallelism (PP) assigns contiguous groups of transformer layers to different devices. Megatron uses an interleaved 1F1B (one-forward-one-backward) schedule to minimize pipeline bubbles. The bubble fraction is approximately $\frac{p - 1}{m}$ where $p$ is the number of pipeline stages and $m$ is the number of microbatches. With enough microbatches, bubble overhead drops below 5%.

Data Parallelism (DP) replicates the model (or model shard) and distributes data batches across replicas. When combined with ZeRO Stage 1 (optimizer state sharding), DP scales throughput linearly with minimal memory overhead. Megatron-Core integrates distributed data parallelism with optional gradient compression.

Context Parallelism (CP) splits the input sequence across GPUs along the sequence dimension, distributing the attention computation. This is essential for models trained on sequences longer than 8K tokens, where the quadratic attention memory becomes a bottleneck even with FlashAttention. CP uses ring attention or similar algorithms to exchange KV blocks between GPUs.

1.2 Practical Configurations for Common Model Sizes

The following table shows recommended parallelism configurations for training common model sizes on clusters of NVIDIA H100 GPUs. These represent starting points; actual configurations require profiling on your specific hardware and interconnect.

1.3 Megatron-Core Configuration

Code Fragment 6.8.1 shows how to configure multi-axis parallelism in Megatron-Core. The TransformerConfig object specifies the model architecture, and parallel_state initializes the process groups for each parallelism axis.

# Code Fragment 6.8.1: Megatron-Core multi-axis parallelism setup
import megatron.core as mcore
from megatron.core.transformer.transformer_config import TransformerConfig
from megatron.core import parallel_state

# Initialize distributed process groups
parallel_state.initialize_model_parallel(
 tensor_model_parallel_size=8, # TP: split layers within node
 pipeline_model_parallel_size=4, # PP: distribute layers across nodes
 context_parallel_size=1, # CP: sequence-dimension splitting
 # DP is inferred: world_size / (TP * PP * CP)
)

# Configure transformer architecture
config = TransformerConfig(
 num_layers=80, # 70B-class model
 hidden_size=8192,
 num_attention_heads=64,
 num_query_groups=8, # GQA: 8 KV heads
 ffn_hidden_size=28672,
 max_position_embeddings=8192,
 bf16=True, # BF16 mixed precision
 # Pipeline parallelism schedule
 pipeline_dtype="bfloat16",
 num_microbatches=16, # Reduce bubble fraction
 overlap_p2p_comm=True, # Overlap pipeline communication
 # Activation checkpointing to reduce memory
 recompute_granularity="selective",
 recompute_method="uniform",
)

# Build the model with Megatron-Core's GPT architecture
model = mcore.models.gpt.GPTModel(
 config=config,
 transformer_layer_spec=mcore.models.gpt.gpt_layer_specs.get_gpt_layer_spec(),
 vocab_size=128256,
 max_sequence_length=8192,
)

1.4 MoE Load Balancing

Mixture-of-Experts (MoE) models (such as Mixtral 8x7B, DeepSeek-V3, and Arctic) add another layer of parallelism complexity. In an MoE transformer, each token is routed to a subset of experts (typically 2 out of 8 or 16). Without load balancing, some experts receive disproportionately more tokens, creating stragglers that slow the entire training step.

Megatron-Core implements Expert Parallelism (EP), where experts are distributed across GPUs and tokens are communicated via all-to-all collectives. Load balancing is enforced through an auxiliary loss term added to the training objective:

$$\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot p_i$$

where $N$ is the number of experts, $f_i$ is the fraction of tokens routed to expert $i$, $p_i$ is the average routing probability for expert $i$, and $\alpha$ is a balancing coefficient (typically 0.01). This loss encourages uniform token distribution across experts without interfering with the primary language modeling objective.

2. Compiler and Kernel Optimization for LLM Training

Raw parallelism strategy determines how work is distributed across GPUs. Compiler and kernel optimization determine how efficiently each GPU executes its share of the work. The gap between theoretical peak FLOPS and actual training throughput (MFU) is where compiler and kernel optimizations make their impact.

2.1 torch.compile for LLM Training Loops

PyTorch 2.x introduced torch.compile, which uses TorchDynamo (a Python bytecode analyzer) and TorchInductor (a code generation backend) to automatically fuse operations and generate optimized GPU kernels. For LLM training, torch.compile can provide 10-30% throughput improvements by fusing elementwise operations, reducing kernel launch overhead, and optimizing memory access patterns.

# Code Fragment 6.8.2: Applying torch.compile to an LLM training loop
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
 "meta-llama/Llama-3.1-8B",
 torch_dtype=torch.bfloat16,
 attn_implementation="flash_attention_2",
)

# Compile the model for training
# mode="reduce-overhead" minimizes kernel launch overhead
# fullgraph=True ensures the entire forward+backward is captured
compiled_model = torch.compile(
 model,
 mode="reduce-overhead",
 fullgraph=True,
 backend="inductor", # TorchInductor: generates Triton kernels
)

# Training loop proceeds normally; compiled model is a drop-in replacement
optimizer = torch.optim.AdamW(compiled_model.parameters(), lr=3e-4)
for batch in dataloader:
 outputs = compiled_model(**batch)
 loss = outputs.loss
 loss.backward()
 optimizer.step()
 optimizer.zero_grad()

2.2 FlashAttention Kernel Selection

FlashAttention (Dao et al., 2022; Dao, 2023) rewrites the attention computation to minimize HBM (High Bandwidth Memory) accesses by tiling the attention matrix and fusing softmax into the GEMM loop. FlashAttention-2 improved throughput by 2x over the original through better work partitioning across GPU thread blocks. FlashAttention-3 (2024) further optimizes for H100 Hopper architecture features, including FP8 tensor cores and asynchronous operations.

The choice of attention kernel depends on the hardware generation and sequence length:

2.3 FP8 and FP4 Training with Transformer Engine

NVIDIA's Transformer Engine library enables mixed-precision training with FP8 (8-bit floating point) on Hopper and later architectures. FP8 training nearly doubles the throughput of BF16 training by exploiting the H100's FP8 tensor cores, which deliver 3958 TFLOPS compared to 1979 TFLOPS for BF16 (Micikevicius et al., 2022).

FP8 training uses two formats: E4M3 (4 exponent bits, 3 mantissa bits) for forward activations and weights, and E5M2 (5 exponent bits, 2 mantissa bits) for gradients, which need the larger dynamic range. A per-tensor scaling factor is maintained and updated each iteration to keep values within the FP8 representable range.

# Code Fragment 6.8.3: Enabling FP8 training with Transformer Engine
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import DelayedScaling, Format

# FP8 recipe: delayed scaling updates the scale factor
# based on the maximum absolute value from the previous iteration
fp8_recipe = DelayedScaling(
 margin=0,
 fp8_format=Format.HYBRID, # E4M3 for fwd, E5M2 for bwd
 amax_history_len=16, # Track max values over 16 iterations
 amax_compute_algo="max", # Use max of history for scaling
)

# Replace nn.Linear with Transformer Engine's FP8-aware version
# in a transformer block
class TransformerBlock(te.pytorch.module.TransformerLayer):
 def __init__(self, hidden_size, num_heads):
 super().__init__(
 hidden_size=hidden_size,
 ffn_hidden_size=4 * hidden_size,
 num_attention_heads=num_heads,
 # FP8 is enabled via the context manager, not here
 )

# During training, wrap the forward pass in the fp8 context
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 output = model(input_ids, labels=labels)
 loss = output.loss
loss.backward()

3. Elastic and Fault-Tolerant Training

A multi-week LLM pretraining run on thousands of GPUs will experience hardware failures. At Meta's scale (16,384 GPUs for Llama 3.1), the team reported encountering over 400 unexpected interruptions during a 54-day training run. At this scale, fault tolerance is not optional; it is a fundamental requirement.

3.1 Torch Distributed Elastic (TorchElastic)

TorchElastic (accessed via torchrun) provides elastic training capabilities for PyTorch distributed jobs. Unlike standard torch.distributed.launch, TorchElastic can handle node failures and additions without requiring a complete restart of the training job. Key features include:

• Membership changes: Workers can join or leave the training group. When a node fails, TorchElastic detects the failure, removes the node from the group, and restarts the remaining workers with a new process group.
• Rendezvous: A coordination protocol (backed by etcd or C10d) that workers use to discover each other and agree on group membership. After a membership change, all surviving workers re-rendezvous before resuming.
• Restart semantics: On failure, TorchElastic kills all workers in the group, loads the latest checkpoint, and restarts with the new (possibly smaller) group. This "fail-stop-restart" approach is simpler and more robust than in-place recovery for LLM training.

# Code Fragment 6.8.4: Launching an elastic LLM training job with torchrun
# --nproc_per_node: GPUs per node
# --nnodes: min:max node range for elasticity
# --rdzv_backend: rendezvous backend (c10d or etcd)
# --rdzv_endpoint: rendezvous coordinator address
# --max_restarts: number of times to restart on failure

torchrun \
 --nproc_per_node=8 \
 --nnodes=4:8 \
 --rdzv_backend=c10d \
 --rdzv_endpoint=master-node:29400 \
 --max_restarts=3 \
 --monitor_interval=5 \
 train_llm.py \
 --model_size=7b \
 --checkpoint_dir=/shared/checkpoints \
 --resume_from_latest

The --nnodes=4:8 flag specifies that the job can run with anywhere from 4 to 8 nodes. If a node fails, training continues (after a brief restart) with the remaining nodes, as long as at least 4 are available. The data parallelism degree adjusts automatically.

3.2 Failure Detection and Recovery Strategies

Production LLM training systems employ multiple layers of failure detection and recovery:

4. Distributed Checkpointing

Checkpointing is the primary mechanism for both fault recovery and experiment management. For large LLM training runs, naive checkpointing (gathering all state to a single writer) is prohibitively slow. A 405B model's full state (parameters + optimizer + gradients) can exceed 8 TB. Writing this to storage as a single file would take tens of minutes, during which training is blocked.

4.1 PyTorch Distributed Checkpoint (DCP)

PyTorch Distributed Checkpoint (DCP), introduced in PyTorch 2.1, provides a native solution for saving and loading sharded model state. Each rank writes its own shard independently, and DCP handles the coordination and metadata management. Key properties include:

• Parallel I/O: All ranks write simultaneously. With 512 GPUs writing to a parallel file system, a 70B model checkpoint completes in seconds rather than minutes.
• Reshardable format: Checkpoints can be loaded with a different parallelism configuration than they were saved with. This means you can save a checkpoint from a TP=8, PP=4 configuration and load it into TP=4, PP=2 for fine-tuning on a smaller cluster.
• Async staging: DCP supports asynchronous checkpointing where state is first copied to CPU memory (non-blocking) and then written to storage in the background while training continues on the next iteration.

# Code Fragment 6.8.5: Saving and loading with PyTorch DCP
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
 get_model_state_dict,
 get_optimizer_state_dict,
 set_model_state_dict,
 set_optimizer_state_dict,
 StateDictOptions,
)

# --- Saving a sharded checkpoint ---
def save_checkpoint(model, optimizer, step, checkpoint_dir):
 """Save a reshardable distributed checkpoint."""
 state_dict = {
 "model": get_model_state_dict(model),
 "optimizer": get_optimizer_state_dict(model, optimizer),
 "step": step,
 }
 dcp.save(
 state_dict=state_dict,
 storage_writer=dcp.FileSystemWriter(
 f"{checkpoint_dir}/step_{step}",
 single_file_per_rank=True, # One file per rank for parallel I/O
 ),
 )

# --- Loading with potentially different parallelism ---
def load_checkpoint(model, optimizer, checkpoint_dir, step):
 """Load checkpoint; DCP handles resharding automatically."""
 state_dict = {
 "model": get_model_state_dict(model),
 "optimizer": get_optimizer_state_dict(model, optimizer),
 }
 dcp.load(
 state_dict=state_dict,
 storage_reader=dcp.FileSystemReader(
 f"{checkpoint_dir}/step_{step}"
 ),
 )
 # Apply loaded state back to model and optimizer
 set_model_state_dict(
 model, state_dict["model"],
 options=StateDictOptions(strict=True),
 )
 set_optimizer_state_dict(
 model, optimizer, state_dict["optimizer"],
 )
 return state_dict.get("step", 0)

4.2 Async Checkpointing and Staging

For multi-week training runs, even a few minutes of blocking per checkpoint adds up to hours of wasted GPU time. Async checkpointing eliminates this bottleneck by overlapping the storage write with the next training iteration:

1. Snapshot to CPU: Copy the model and optimizer state from GPU to pinned CPU memory. This takes seconds for a 70B model.
2. Continue training: The GPU immediately begins the next training iteration.
3. Background write: A separate thread writes the CPU snapshot to the parallel file system or object storage (S3, GCS).
4. Verification: Once the write completes, the checkpoint metadata is atomically updated to mark the new checkpoint as valid.

4.3 Checkpoint Frequency Trade-offs

Choosing checkpoint frequency involves a trade-off between recovery time and storage cost. If you checkpoint every $C$ iterations and a failure occurs, you lose on average $C/2$ iterations of work. For a training run with $N$ total iterations and a mean time between failures (MTBF) of $F$ iterations, the expected wasted compute is:

$$\text{Wasted fraction} \approx \frac{C}{2F} + \frac{t_{\text{ckpt}}}{t_{\text{iter}} \cdot C}$$

where $t_{\text{ckpt}}$ is the checkpoint duration and $t_{\text{iter}}$ is the iteration time. The first term represents lost work from rollback; the second represents time spent checkpointing. Minimizing their sum gives the optimal checkpoint interval. With async checkpointing, $t_{\text{ckpt}}$ drops to near zero (only the GPU-to-CPU copy time), making more frequent checkpointing practical.

5. Streaming Dataset Pipelines

LLM pretraining datasets are too large to fit in local storage on training nodes. The Llama 3 training corpus was approximately 15 trillion tokens, and the full tokenized dataset can exceed 30 TB. Streaming dataset libraries solve this by loading data on-the-fly from cloud object storage, eliminating the need to pre-download the entire dataset.

5.1 Mosaic Streaming (StreamingDataset)

MosaicML's streaming library (now part of the Databricks ecosystem) provides a streaming dataloader purpose-built for distributed LLM training. Key properties include:

• Deterministic sharding: Each rank sees a deterministic, non-overlapping subset of the data, regardless of the number of workers or the order of data arrival. This ensures reproducibility across runs with different parallelism configurations.
• Resumable iteration: The library tracks which samples each rank has consumed. On restart after a failure, iteration resumes exactly where it left off, without replaying or skipping any samples.
• Local caching: Frequently accessed shards are cached on local NVMe storage, reducing repeated downloads from object storage.
• Data mixing: Multiple data sources (web text, code, books, conversations) can be mixed with configurable proportions that can change during training (curriculum learning).

# Code Fragment 6.8.6: Streaming dataloader for distributed LLM training
from streaming import StreamingDataset, StreamingDataLoader

# Define the dataset with multiple data sources and mixing proportions
dataset = StreamingDataset(
 streams=[
 {
 "remote": "s3://training-data/web-text/",
 "local": "/local-nvme/cache/web-text/",
 "proportion": 0.50, # 50% web text
 },
 {
 "remote": "s3://training-data/code/",
 "local": "/local-nvme/cache/code/",
 "proportion": 0.25, # 25% code
 },
 {
 "remote": "s3://training-data/books/",
 "local": "/local-nvme/cache/books/",
 "proportion": 0.15, # 15% books
 },
 {
 "remote": "s3://training-data/conversations/",
 "local": "/local-nvme/cache/conversations/",
 "proportion": 0.10, # 10% conversational data
 },
 ],
 shuffle=True,
 shuffle_seed=42, # Deterministic shuffling
 num_canonical_nodes=64, # Controls sharding granularity
 batch_size=4, # Per-device micro-batch size
)

# StreamingDataLoader handles distributed rank assignment automatically
dataloader = StreamingDataLoader(
 dataset=dataset,
 batch_size=4,
 num_workers=8,
 pin_memory=True,
 prefetch_factor=4, # Prefetch 4 batches ahead
)

# Resume from a specific sample index after failure
dataloader.dataset.state_dict() # Save state
# ... after restart ...
dataloader.dataset.load_state_dict(saved_state) # Resume exactly

5.2 WebDataset

WebDataset takes a different approach: it stores data as tar archives containing individual samples, and streams them sequentially. This format is particularly efficient for multimodal training (image-text pairs) where each sample consists of multiple files. For text-only LLM training, WebDataset works well but offers less sophisticated mixing and resumability compared to Mosaic Streaming.

5.3 Data Integrity and Validation

At the scale of trillions of tokens, data corruption can silently degrade model quality. Production training pipelines include integrity checks at multiple levels:

• Shard-level checksums: Each data shard has an associated MD5 or SHA-256 hash verified on download.
• Sample-level validation: Each sample is checked for basic structural integrity (correct token IDs, valid sequence lengths) before entering the training batch.
• Epoch-level auditing: After each pass through the data, statistics on token counts, sequence lengths, and data source proportions are logged and compared against expected values.
• Deduplication verification: Near-duplicate detection (MinHash or suffix arrays) is run as a preprocessing step, and the streaming pipeline includes checks that no sample appears twice within a configurable window.

6. Putting It All Together: Production Training Architecture

A production LLM training system integrates all of the components discussed in this section. The following diagram shows how they fit together for a typical 70B+ model pretraining run:

Figure 6.8.1: Production LLM training architecture. Data flows from object storage through streaming loaders. Megatron-Core handles multi-axis parallelism with compiler optimizations. TorchElastic provides fault tolerance. DCP manages distributed checkpoints. The monitoring layer tracks training health across all components.