3.7 Megatron-LM, DeepSpeed & Parallelism in Practice¶

Knowing the theory of tensor, pipeline, and data parallelism is one thing. Knowing how to wire them together on a 512-GPU cluster, pick the right degrees, and then confirm that your hardware is actually doing useful work is another. This chapter bridges that gap. We take the parallelism primitives introduced in Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP and Distributed Training II: Tensor, Pipeline, Sequence & Expert Parallelism, and show how Megatron-LM and DeepSpeed compose them into production training runs.

By the end of this chapter you will understand the Megatron-Core abstraction layer, the full ZeRO hierarchy and its offload variants, the 4-D (DP × TP × PP × EP) parallelism space, how to reason about Model FLOP Utilization (MFU) and Hardware FLOP Utilization (HFU), and exactly which configuration levers to pull for a 70B-parameter run.

Megatron-LM: A Framework Built Around 3-D Parallelism¶

Megatron-LM, developed at NVIDIA, was the first framework to train models beyond 100B parameters in a systematic way. The 2021 paper by Narayanan et al. (“Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM”) introduced the idea of combining tensor parallelism (TP), pipeline parallelism (PP), and data parallelism (DP) in a principled way. The current codebase ships as two packages:

• Megatron-LM — the outer loop: launch scripts, training harness, logging, checkpointing.
• Megatron-Core (megatron.core) — a library of parallelism-aware transformer building blocks that other frameworks (NVIDIA NeMo, Databricks MosaicML, Aleph Alpha) can import.

The 3-D Parallelism Layout¶

Given a cluster with \(N\) GPUs, Megatron partitions them into a 3-D grid:

Every GPU belongs to exactly one data-parallel replica, one tensor-parallel group, and one pipeline stage. The layout is typically described as:

The rule of thumb that emerges from practice: place TP within a node (so the all-reduce inside a TP group can use NVLink rather than crossing InfiniBand) and use PP to span nodes (pipeline sends are point-to-point and latency-tolerant). DP can be at any granularity but is often the outermost dimension.

Megatron-Core Parallel State¶

Megatron-Core tracks the parallelism topology in a global parallel_state module. Understanding this is essential for debugging.

# megatron/core/parallel_state.py  (simplified)
import torch
import torch.distributed as dist

_TP_GROUP = None  # tensor-parallel process group
_PP_GROUP = None  # pipeline-parallel process group
_DP_GROUP = None  # data-parallel process group

def initialize_model_parallel(
    tensor_model_parallel_size: int = 1,
    pipeline_model_parallel_size: int = 1,
):
    """
    Build the 3-D process group topology.
    Assumes dist.init_process_group() has already been called.
    world_size = TP * PP * DP is enforced implicitly.
    """
    global _TP_GROUP, _PP_GROUP, _DP_GROUP

    world_size = dist.get_world_size()
    rank = dist.get_rank()

    dp_size = world_size // (tensor_model_parallel_size * pipeline_model_parallel_size)

    # --- Tensor-parallel groups ---
    # Each group of TP consecutive ranks forms one TP group.
    for i in range(pipeline_model_parallel_size * dp_size):
        ranks = list(range(i * tensor_model_parallel_size,
                           (i + 1) * tensor_model_parallel_size))
        group = dist.new_group(ranks)
        if rank in ranks:
            _TP_GROUP = group

    # --- Pipeline-parallel groups ---
    # Stride by TP; each group of PP elements (spaced TP apart) is a PP group.
    for i in range(tensor_model_parallel_size * dp_size):
        ranks = list(range(i, world_size, tensor_model_parallel_size))[:pipeline_model_parallel_size]
        group = dist.new_group(ranks)
        if rank in ranks:
            _PP_GROUP = group

    # --- Data-parallel groups ---
    for i in range(tensor_model_parallel_size * pipeline_model_parallel_size):
        ranks = list(range(i, world_size, tensor_model_parallel_size * pipeline_model_parallel_size))
        group = dist.new_group(ranks)
        if rank in ranks:
            _DP_GROUP = group

def get_tensor_model_parallel_group(): return _TP_GROUP
def get_pipeline_model_parallel_group(): return _PP_GROUP
def get_data_parallel_group(): return _DP_GROUP

Every Megatron-Core layer that performs a collective operation (e.g., the column/row linear layers in the transformer) calls get_tensor_model_parallel_group() to target the right process group. This design keeps the distributed logic out of user code.

Sequence Parallelism in Megatron¶

The original Megatron TP implementation broadcasts the input activations to all TP ranks before computing. This means the activations (LayerNorm inputs, dropout outputs) are replicated across TP ranks and waste memory proportional to TP.

Megatron-LM’s sequence parallelism (introduced in “Reducing Activation Recomputation in Large Transformer Models,” Korthikanti et al. 2022) avoids this. Outside the TP-sharded GEMM blocks, the sequence dimension is sharded across TP ranks. The transition in/out of sequence-parallel regions uses all-gather before the column parallel GEMM and reduce-scatter after the row parallel GEMM — replacing an all-reduce with two smaller collectives.

Net effect: activation memory at steady state is H * B * S / TP rather than H * B * S. For TP=8 this is an 8× activation reduction — significant at long sequences.

DeepSpeed: ZeRO and Beyond¶

DeepSpeed (Microsoft) complements Megatron by providing the optimizer-side solution. The core abstraction is ZeRO (Zero Redundancy Optimizer), which eliminates the redundant copies of optimizer states, gradients, and parameters that vanilla DDP maintains.

ZeRO Stages Recap¶

The three ZeRO stages correspond to increasingly aggressive sharding across DP ranks. We cover the theory in Distributed Training I; here we focus on the implementation details that matter for a real run.

That last row is the whole trade: ZeRO-3 buys a DP-fold reduction in model-state memory at 1.5× the communication of DDP, and unlike ZeRO-½ the extra all-gather sits on the critical path of the forward pass, so it must be prefetched (stage3_prefetch_bucket_size) to be hidden.

For a model with \(P\) parameters stored in fp16 (2 bytes) and Adam optimizer states in fp32:

The first \(2P\) is the fp16 parameters, the second \(2P\) the fp16 gradients, and the \(12P\) comes from the fp32 master weights (4), Adam first moment (4), and second moment (4) — the same 16 bytes per parameter tallied in Step 1 below. ZeRO-3 shards all 16 bytes; each DP rank holds \(16P / \text{DP}\) bytes of “owned” state plus the activations for its pipeline stage.

ZeRO-Offload and ZeRO-Infinity¶

For clusters with more CPU memory or NVMe than GPU memory, DeepSpeed provides offload variants:

• ZeRO-Offload: moves optimizer states (and optionally gradients) to CPU RAM. The optimizer step runs on CPU, which is fine because optimizer steps are memory-bandwidth-bound (read param + states, write updated param) and do not require GPU arithmetic throughput.
• ZeRO-Infinity: extends offload to NVMe storage using heterogeneous memory management. A bandwidth-aware scheduler overlaps NVMe reads with GPU compute.

# DeepSpeed config JSON for ZeRO-3 with CPU offload
import json

zero3_config = {
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",       # optimizer states live on CPU RAM
            "pin_memory": True     # page-locked for fast DMA transfer
        },
        "offload_param": {
            "device": "cpu",       # fp16 params also offloaded
            "pin_memory": True
        },
        "overlap_comm": True,      # overlap reduce-scatter with backward pass
        "contiguous_gradients": True,
        "sub_group_size": 1e9,     # process params in 1B-element chunks
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
    },
    "fp16": {
        "enabled": True,
        "loss_scale": 0,           # dynamic loss scaling
        "loss_scale_window": 1000
    },
    "gradient_clipping": 1.0,
    "train_micro_batch_size_per_gpu": 2,
    "gradient_accumulation_steps": 8,
}

with open("ds_config_zero3.json", "w") as f:
    json.dump(zero3_config, f, indent=2)

Initializing a DeepSpeed Engine¶

import deepspeed
import torch
import torch.nn as nn

class TinyTransformerBlock(nn.Module):
    """A minimal transformer block for demonstration."""
    def __init__(self, d_model: int, n_heads: int, ffn_mult: int = 4):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ffn_mult * d_model),
            nn.GELU(),
            nn.Linear(ffn_mult * d_model, d_model),
        )

    def forward(self, x):
        # Pre-norm residual style (GPT-style)
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.ffn(self.norm2(x))
        return x

model = nn.Sequential(*[TinyTransformerBlock(1024, 16) for _ in range(24)])

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)

# deepspeed.initialize wraps the model, optimizer, and dataloader
# into a DeepSpeedEngine that handles ZeRO sharding transparently.
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    config="ds_config_zero3.json",
)

# Training step — identical to vanilla PyTorch apart from backward()/step()
# being called on the engine rather than on the loss and the optimizer.
import torch.nn.functional as F

for batch in dataloader:                       # yields (inputs, labels) on the right device
    inputs, labels = batch
    hidden = model_engine(inputs)              # (B, S, d_model)
    logits = lm_head(hidden)                   # your untied/tied output projection
    loss = F.cross_entropy(
        logits.view(-1, logits.size(-1)).float(),
        labels.view(-1),
    )
    model_engine.backward(loss)    # gradient reduce-scatter, accumulation-aware
    model_engine.step()            # triggers all-gather, optimizer step, re-shard

Note that model_engine.backward() and .step() are gradient-accumulation aware: DeepSpeed reads gradient_accumulation_steps from the config and only fires the reduce-scatter and the optimizer update on boundary micro-steps, so your loop never calls zero_grad() or counts micro-batches itself.

The 4-D Parallelism Space: Adding Expert Parallelism¶

Modern MoE (Mixture-of-Experts) models add a fourth dimension. See Mixture-of-Experts (MoE) Architectures for the architecture. The full parallelism space becomes:

Expert parallelism (EP) shards the experts across EP ranks. Within a single MoE layer, tokens are dispatched to experts on different GPUs via all-to-all collectives. EP communicates token activations rather than parameters, so the all-to-all volume is proportional to sequence length and hidden size, not parameter count.

Megatron-Core’s MoELayer handles the EP dimension natively (--expert-model-parallel-size, with --moe-token-dispatcher-type alltoall). The key constraint: each EP group must see enough tokens to keep all experts loaded. An expert that processes very few tokens is wasted capacity — the load imbalance problem that auxiliary loss terms (introduced by Switch Transformer, Fedus et al.) address.

The Fifth Axis: Context Parallelism¶

For long-context runs there is one more independent axis. Context parallelism (CP) shards the sequence across CP ranks and computes attention with a ring exchange of K/V blocks (Ring Attention), so activation memory and attention FLOPs per GPU both fall by CP× while the model weights stay replicated across the CP group. Megatron-Core exposes it as --context-parallel-size, and the full grid becomes \(N = \text{DP} \times \text{TP} \times \text{PP} \times \text{CP} \times \text{EP}\).

Do not confuse CP with the sequence parallelism described above: Megatron’s SP is a memory optimization strictly inside a TP group (it shards the norm/dropout regions along sequence and is free, so it is always on whenever TP > 1), whereas CP is a genuine extra dimension of the GPU grid that you spend GPUs on. The rule of thumb is to leave CP=1 until sequence length pushes activation memory past what recomputation can absorb — typically 32K tokens and beyond — then raise CP before raising TP, because the ring exchange is point-to-point and overlappable while the TP all-reduce is not. The mechanism is developed in Long-Context Pretraining & Context Extension and Distributed Training II.

Choosing Parallelism Degrees: A Systematic Approach¶

The parallelism config is the single highest-leverage decision when launching a large training run. Getting it wrong can cost 30-50% of hardware efficiency. Here is a principled workflow.

Step 1 — Fit the Model on One Node¶

Start with the model’s parameter count \(P\). For mixed-precision training with bf16 parameters and fp32 optimizer states:

For a 70B model: \(70 \times 10^9 \times 16 = 1{,}120 \text{ GB}\) of pure model state, before activations or intermediate buffers.

With 8× H100 80 GB per node (640 GB HBM), you need at least \(\lceil 1120 / 80 \rceil = 14\) GPUs just for model state. In practice, activation memory can double this, so TP=8 × PP=4 = 32 GPUs minimum before any DP replication.

Step 2 — Pick TP¶

TP is constrained by intra-node bandwidth (NVLink). The all-reduce inside a TP column-parallel GEMM must finish before the next GEMM begins; it sits on the critical path.

• TP=1: no communication, maximum arithmetic intensity.
• TP=2: doubles memory for attention and FFN weight distribution; all-reduce is 2× 25 GB/s NVLink streams.
• TP=4 or TP=8: recommended for nodes with 4 or 8 GPUs respectively and NVLink.
• TP > 8: crosses PCIe/InfiniBand; avoid unless forced.

Rule: TP = number of GPUs per node (or a divisor of it) and never span node boundaries.

Step 3 — Pick PP¶

Pipeline parallelism introduces a bubble overhead. The bubble fraction for the 1F1B (one-forward-one-backward) schedule is approximately:

where \(m\) is the number of micro-batches in flight. To keep bubble under 5%:

For PP=4 this means \(m \geq 57\) micro-batches, which is achievable with large global batch sizes.

Interleaved pipeline schedules (Megatron’s virtual pipeline parallelism) split each stage into \(v\) chunks, reducing the bubble to:

at the cost of additional pipeline communication per micro-batch.

Step 4 — Set DP from what remains¶

DP is “free” communication if you use ZeRO-1 or ZeRO-2 (the reduce-scatter / all-gather can be overlapped with the backward pass). ZeRO-3 adds synchronous all-gather per forward pass but eliminates parameter redundancy.

Step 5 — Tune global batch size and gradient accumulation¶

Global batch size (GBS) drives convergence. There is a critical batch size \(B_{\text{crit}}\) — well approximated by the gradient noise scale of McCandlish et al. (An Empirical Model of Large-Batch Training, 2018) — below which increasing the batch size buys a near-linear reduction in the number of optimizer steps, and above which the returns diminish sharply. \(B_{\text{crit}}\) is not a function of vocabulary size and has no simple closed form; it grows over the course of training as the loss falls. Rather than a formula, practitioners target empirical global batch sizes on the order of 1M–4M tokens per step for runs above 10B parameters, then tune the learning rate against batch size as described in Learning Rate Schedules, Warmup, Batch Size & Hyperparameters.

Given a fixed GBS and a per-GPU micro-batch size (MBS), the number of gradient accumulation steps (GAS) is:

Running the Funnel at 100M: the Answer Is Usually DP-Only¶

Apply the same five steps to the capstone’s ~100M-parameter model and every model-parallel degree collapses to 1. Step 1: \(16P = 16 \times 10^8 \approx 1.6\) GB of model state — it fits on a 16 GB T4 with room to spare, so nothing must be sharded. Step 2: TP exists only to make a model fit or to cut per-GPU activation memory; with 1.6 GB of state there is nothing to split, and TP=2 would add an all-reduce per layer to a model whose GEMMs (\(d_{\text{model}}=512\)) are too small to amortize it. Step 3: PP has the same answer, plus a bubble you cannot pay for. Steps 4–5: DP = number of GPUs, and you spend the whole global-batch budget on gradient accumulation.

The practical consequence is worth stating plainly, because it saves readers weeks: you do not need Megatron-LM or DeepSpeed to train a 100M model. Plain DistributedDataParallel — or FSDP2 if you want the optimizer-state saving for free — over torchrun --nproc_per_node=8 is the correct tool, and The Pretraining Run: A Complete Single-GPU Training Loop shows that even one GPU suffices for the whole capstone. What does transfer directly from this chapter at 100M scale is the measurement discipline: the \(6P\) FLOPs/token rule, the MFU calculation, and the Nsight-before-you-tune habit. Small models typically land at lower MFU than 70B ones (they are launch-latency- and memory-bandwidth-bound rather than GEMM-bound), so calibrate expectations against similar-sized runs, not against the 40–55% frontier-scale band. The techniques here become mandatory somewhere around 7B–13B, the transition described in Retrospective: Cost Accounting, Reproducibility, and the Path to 1B.

MFU and HFU: Measuring Real Hardware Utilization¶

You have launched the run. Now you want to know if you are getting good value from your hardware. Two metrics matter.

Model FLOP Utilization (MFU)¶

MFU measures what fraction of the GPU’s peak throughput is being used for forward + backward arithmetic of the model itself.

The number of FLOPs per token for a standard decoder-only Transformer is approximately:

where the \(6P\) term comes from \(\sim 2P\) for the forward pass (each parameter participates in roughly 2 multiply-adds) times 3 for the full backward pass, and the second term is the attention quadratic cost (often secondary for moderate sequence lengths). Following the PaLM/Megatron convention, the attention term counts the full \(S \times S\) score matrix even though causal masking means a FlashAttention kernel actually computes only half of it — MFU is deliberately a model-FLOP metric, so everyone must count the same nominal FLOPs for the numbers to be comparable across papers.

For practical purposes, practitioners often use the simplified rule:

MFU is then:

Target MFU for dense models on A100/H100 clusters is 35–55%. Values below 25% suggest a misconfiguration (bubble too high, TP spans nodes, small micro-batch sizes hitting latency bottlenecks, checkpointing overhead).

Hardware FLOP Utilization (HFU)¶

HFU counts all FLOPs actually issued to the GPU, including those in recomputed activations (gradient checkpointing). If you recompute one third of layers:

With full activation recomputation, you pay for an extra forward pass just before each backward, giving forward \(2P\) + recomputed forward \(2P\) + backward \(4P\) = \(8P\) per token, so:

HFU ≥ MFU. A good cluster should see HFU of 55–70% on H100s with modern frameworks.

def compute_mfu(
    model_params: int,        # number of parameters
    tokens_per_second: float, # observed training throughput
    peak_flops_per_sec: float, # e.g., 5.06e17 for dense bf16 on 512 H100 SXM5 GPUs
    n_layers: int = None,
    d_model: int = None,
    seq_len: int = None,
) -> float:
    """
    Compute Model FLOP Utilization.

    Uses the simplified 6P rule for forward+backward FLOPs per token.
    If n_layers, d_model, seq_len are provided, also adds attention cost.
    """
    flops_per_token = 6 * model_params

    # Attention cost (forward + backward): 12 * n_layers * d_model * seq_len per token.
    # Forward is 4 * d * S per layer (QK^T and AV, each 2 * d * S); x3 for the
    # backward pass, matching the fwd+bwd convention of the 6P base term.
    if all(v is not None for v in [n_layers, d_model, seq_len]):
        attn_flops = 12 * n_layers * d_model * seq_len
        flops_per_token += attn_flops

    achieved_flops = flops_per_token * tokens_per_second
    mfu = achieved_flops / peak_flops_per_sec
    return mfu


# Example: 70B model, 512 H100 SXM5 GPUs
# H100 SXM5 dense bf16 peak: ~989 TFLOP/s = 9.89e14 FLOP/s per GPU.
# (The 1979 TFLOP/s figure NVIDIA quotes is the 2:4-sparse rate; dense
#  training runs against half of it.)
H100_BF16_TFLOPS = 9.89e14
n_gpus = 512
peak_cluster_flops = H100_BF16_TFLOPS * n_gpus  # ~5.06e17 FLOP/s

# Observed: 1200 tokens/second per GPU = 614.4K tokens/second total
tokens_per_sec = 1200 * n_gpus

mfu = compute_mfu(
    model_params=70e9,
    tokens_per_second=tokens_per_sec,
    peak_flops_per_sec=peak_cluster_flops,
    n_layers=80,
    d_model=8192,
    seq_len=4096,
)
print(f"MFU: {mfu:.2%}")  # prints ~54.9% (~55%) for a well-configured run, incl. attention term

A Complete Worked Configuration: 70B Pretraining Run¶

This section presents a concrete, production-style launch for a 70B dense model using Megatron-LM + DeepSpeed ZeRO-1.

Cluster topology¶

Cluster: 64 nodes × 8 H100-SXM5-80GB = 512 GPUs
Network: 400 Gb/s InfiniBand NDR (inter-node), 900 GB/s NVLink (intra-node)
Parallelism: TP=8, PP=8, DP=8  →  512 = 8 × 8 × 8

Model config¶

# model_config.py — Llama-style 70B architecture
MODEL_CONFIG = {
    "num_layers": 80,
    "hidden_size": 8192,
    "ffn_hidden_size": 28672,   # ~3.5x hidden_size for SwiGLU
    "num_attention_heads": 64,
    "num_key_value_heads": 8,   # GQA with 8 KV heads
    "max_position_embeddings": 8192,
    "vocab_size": 128256,       # Llama-3 vocabulary
    "activation_function": "swiglu",
    "normalization": "rmsnorm",
    "tie_embeddings": False,
}

Megatron launch script¶

#!/bin/bash
# launch_70b.sh — SLURM-based Megatron-LM 70B launch

#SBATCH --nodes=64
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=12
#SBATCH --mem=960G           # enough for ZeRO-Offload CPU tensors if needed

# ---- Parallelism degrees ----
TP=8
PP=8
DP=8  # implicit: 512 / (8*8)

# ---- Batch configuration ----
# Global batch size: ~4M tokens per step
# Seq len 4096, MBS=1 per GPU, GAS=128 → GBS = 1 * 512 * 128 * 4096 = 268M tokens... too large
# More typically: MBS=2, GAS=32 → GBS = 2 * 512 * 32 * 4096 = ~134M tokens/step — still large
# In practice: GBS set to 1M tokens = 244 sequences of 4096 tokens
# With DP=8, MBS=1, GAS=ceil(244/8/GAS_factor): tune per run
SEQ_LEN=4096
GLOBAL_BATCH_SIZE=2048      # sequences per step = 2048 × 4096 = 8.4M tokens
MICRO_BATCH_SIZE=2
GAS=$((GLOBAL_BATCH_SIZE / (DP * MICRO_BATCH_SIZE)))  # = 128

# ---- Training config ----
TRAIN_ITERS=500000
LR=3e-4
MIN_LR=3e-5
LR_WARMUP_ITERS=2000
LR_DECAY_STYLE=cosine
CLIP_GRAD=1.0
WEIGHT_DECAY=0.1

# ---- Paths ----
DATA_PATH=/mnt/storage/tokenized/llama3_merged
CHECKPOINT_PATH=/mnt/checkpoints/70b-run

torchrun \
  --nnodes=$SLURM_NNODES \
  --nproc_per_node=8 \
  --rdzv_backend=c10d \
  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
  pretrain_gpt.py \
  --tensor-model-parallel-size $TP \
  --pipeline-model-parallel-size $PP \
  --num-layers 80 \
  --hidden-size 8192 \
  --ffn-hidden-size 28672 \
  --num-attention-heads 64 \
  --group-query-attention \
  --num-query-groups 8 \
  --seq-length $SEQ_LEN \
  --max-position-embeddings $SEQ_LEN \
  --micro-batch-size $MICRO_BATCH_SIZE \
  --global-batch-size $GLOBAL_BATCH_SIZE \
  --train-iters $TRAIN_ITERS \
  --lr $LR \
  --min-lr $MIN_LR \
  --lr-warmup-iters $LR_WARMUP_ITERS \
  --lr-decay-style $LR_DECAY_STYLE \
  --weight-decay $WEIGHT_DECAY \
  --clip-grad $CLIP_GRAD \
  --bf16 \
  --use-flash-attn \
  --recompute-activations \
  --recompute-granularity selective \
  --use-distributed-optimizer \
  --overlap-grad-reduce \
  --overlap-param-gather \
  --use-rope-scaling \
  --normalization RMSNorm \
  --swiglu \
  --tokenizer-type TikTokenizer \
  --data-path $DATA_PATH \
  --save $CHECKPOINT_PATH \
  --load $CHECKPOINT_PATH \
  --save-interval 1000 \
  --eval-interval 500 \
  --log-interval 10 \
  --tensorboard-dir $CHECKPOINT_PATH/tb \
  --wandb-project llm-stack-70b

The four performance flags are the ones worth memorizing. --recompute-activations selects the selective recomputation policy of Korthikanti et al. (attention softmax/dropout only); --use-distributed-optimizer turns on Megatron’s own optimizer-state sharding across the DP group — that is ZeRO-1, implemented natively inside Megatron-Core; --overlap-grad-reduce hides the DP reduce-scatter behind the backward pass; and --overlap-param-gather hides the distributed optimizer’s parameter all-gather behind the forward pass. Add --tp-comm-overlap (which requires Transformer Engine’s userbuffers) to additionally overlap the TP all-gather/reduce-scatter with the GEMMs they bracket.

DeepSpeed config for this run¶

{
  "zero_optimization": {
    "stage": 1,
    "overlap_comm": true,
    "allgather_partitions": true,
    "reduce_scatter": true,
    "allgather_bucket_size": 500000000,
    "reduce_bucket_size": 500000000
  },
  "bf16": {
    "enabled": true
  },
  "gradient_clipping": 1.0,
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 128,
  "steps_per_print": 10,
  "wall_clock_breakdown": false
}

Monitoring the run¶

# parse_megatron_logs.py — extract MFU from Megatron-LM stdout
import re
import sys

LOG_LINE_RE = re.compile(
    r"iteration\s+(\d+)/\s*\d+.*?elapsed time per iteration \(ms\): ([\d.]+).*?"
    r"tokens-per-second-per-gpu: ([\d.]+)",
    re.DOTALL,
)

H100_BF16_PEAK_TFLOPS = 989.0  # dense bf16 Tensor Core, per GPU (1979 is the 2:4-sparse rate)
MODEL_PARAMS = 70e9

def toks_per_sec_to_mfu(tps_per_gpu: float) -> float:
    flops_per_tok = 6 * MODEL_PARAMS
    achieved_tflops = flops_per_tok * tps_per_gpu / 1e12
    return achieved_tflops / H100_BF16_PEAK_TFLOPS

for line in sys.stdin:
    m = LOG_LINE_RE.search(line)
    if m:
        iteration = int(m.group(1))
        ms_per_iter = float(m.group(2))
        tps = float(m.group(3))
        mfu = toks_per_sec_to_mfu(tps)
        print(f"iter {iteration:6d} | {ms_per_iter:6.0f} ms/it | {tps:5.0f} tok/s/gpu | MFU {mfu:.1%}")

Practical Pitfalls and Tuning Knobs¶

The TP Communication Bottleneck¶

Tensor parallelism sits on the critical path of the forward pass. If TP all-reduces are slow (e.g., because TP spans InfiniBand instead of NVLink), you can lose 20-40% of throughput. Always profile with nsys profile and check that ncclAllReduce calls within a TP group run at NVLink speed (≈ 600 GB/s aggregate bidirectional).

Pipeline Bubble vs. Memory Tradeoff¶

Increasing PP reduces per-GPU memory but increases the bubble fraction. For PP=8 and \(m=32\) micro-batches, the bubble is \((8-1)/(8-1+32) \approx 18\%\). Doubling micro-batches (increasing global batch or reducing MBS) drops this to 9%. Interleaved schedules (virtual PP) halve it again but increase inter-stage communication by a factor of \(v\).

Activation Recomputation Granularity¶

Megatron-LM offers three granularity levels:

For large models, selective is the sweet spot — it eliminates the expensive-to-store softmax activations (which grow as \(O(S^2)\) in sequence length) while retaining the cheaper MLP activations.

Choosing Between Megatron and FSDP¶

PyTorch’s Fully Sharded Data Parallel (Distributed Training I) covers a similar use case to ZeRO-3. Rule of thumb:

• Megatron-Core (+ optionally DeepSpeed) for runs on dedicated clusters with real TP/PP requirements (>30B parameters), and wherever Transformer Engine’s FP8 and TP-comm-overlap kernels matter.
• FSDP2 for runs up to ~30B parameters where single-framework PyTorch is preferred and TP is unnecessary.
• FSDP2 + TP/PP/CP via DTensor is the PyTorch-native path for very large models without a Megatron dependency.

You rarely wire that last option by hand. The reference implementation is pytorch/torchtitan, PyTorch’s own pretraining platform, which composes FSDP2, DTensor tensor parallelism, pipeline parallelism (torch.distributed.pipelining), context parallelism, torch.compile, Float8 training via torchao, and Distributed Checkpoint behind a single TOML config — the PyTorch-native answer to Megatron-LM. Its ND-parallel mesh is built with init_device_mesh, the same primitive you would use yourself:

from torch.distributed.device_mesh import init_device_mesh

# 512 GPUs as DP=8 x PP=8 x TP=8 — the direct analogue of Megatron's
# parallel_state.initialize_model_parallel(), but as a first-class PyTorch object.
mesh = init_device_mesh("cuda", (8, 8, 8), mesh_dim_names=("dp", "pp", "tp"))
tp_group = mesh["tp"].get_group()      # what get_tensor_model_parallel_group() returns
dp_mesh = mesh["dp"]                   # pass straight to fully_shard(module, mesh=dp_mesh)

The other stacks you will meet in the wild: NVIDIA NeMo (a full framework built on Megatron-Core, adding data/recipe management and multimodal), HuggingFace nanotron (a compact, readable 3D-parallel pretrainer), allenai/OLMo-core (the fully open OLMo training stack), and HuggingFace accelerate, which is not a parallelism engine itself but a launcher that plugs a Trainer-style loop into DeepSpeed or FSDP configs.

For mixed-precision choices and the role of bf16 vs fp8 in these runs, see Mixed Precision, bf16 & FP8 Training.

nsys profile \
    --trace cuda,nvtx \
    --output profile_iter \
    --capture-range cudaProfilerApi \
    python pretrain_gpt.py --profile-step 5 ...

Hyperparameter Sensitivity and Scaling the Config¶

Not all hyperparameters are scale-invariant. When you double the cluster and increase GBS, you typically need to:

1. Scale the learning rate with the square root of GBS (linear scaling works empirically up to a point, but for very large batches \(\sqrt{\text{GBS}}\) scaling is safer). See Learning Rate Schedules, Warmup, Batch Size & Hyperparameters.
2. Extend warmup proportionally to GBS — a common heuristic is to warm up over 1–2B tokens regardless of batch size.
3. Reduce gradient clipping threshold as model depth grows to avoid spurious gradient explosions (see Training Stability, Loss Spikes & Debugging Large Runs).

Virtual Pipeline Parallelism¶

Megatron’s interleaved pipeline (enabled with --num-layers-per-virtual-pipeline-stage) assigns each GPU multiple non-contiguous chunks of layers:

This halves the bubble at the cost of sending twice as many pipeline messages per micro-batch. In practice, for TP=8 where the pipeline messages are relatively small (one micro-batch of activations), interleaving is almost always worth it above PP=4.

Combining Megatron-Core with External Libraries¶

Megatron-Core is designed to be embedded. The typical NeMo or Databricks Mosaic setup looks like:

# Pattern: Megatron-Core layer inside a custom training loop
from megatron.core import parallel_state
from megatron.core.transformer.transformer_config import TransformerConfig
from megatron.core.models.gpt.gpt_model import GPTModel
# Layer "specs" say which implementation backs each sublayer. The Transformer
# Engine spec uses TE's fused kernels (and is required for FP8); there is also a
# pure-PyTorch `get_gpt_layer_local_spec()` for debugging without TE installed.
from megatron.core.models.gpt.gpt_layer_specs import (
    get_gpt_layer_with_transformer_engine_spec,
)

# 1. Initialize distributed environment
import torch.distributed as dist
dist.init_process_group(backend="nccl")

# 2. Configure the model parallel topology
parallel_state.initialize_model_parallel(
    tensor_model_parallel_size=8,
    pipeline_model_parallel_size=4,
)

# 3. Build the model using TransformerConfig
config = TransformerConfig(
    num_layers=80,
    hidden_size=8192,
    num_attention_heads=64,
    num_query_groups=8,         # GQA
    ffn_hidden_size=28672,
    use_cpu_initialization=False,
    bf16=True,
    tensor_model_parallel_size=8,
    pipeline_model_parallel_size=4,
)

model = GPTModel(
    config=config,
    transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(),
    vocab_size=128256,
    max_sequence_length=8192,
    position_embedding_type="rope",
)

# 4. Optimizer-state sharding. Either Megatron's own distributed optimizer
#    (config.use_distributed_optimizer = True, ZeRO-1) or DeepSpeed — not both.
import deepspeed
model_engine, _, _, _ = deepspeed.initialize(model=model, config="ds_config_zero1.json")

This pattern — Megatron-Core for the TP/PP topology, DeepSpeed for the optimizer-side ZeRO sharding — is sometimes called 3D + ZeRO and remains the dominant approach for frontier model training runs in 2026, now increasingly running on NVIDIA Blackwell (B200 / GB200 NVL72) systems alongside the large H100/H200 fleets of the previous generation.

For inference serving after training, the parallelism story shifts toward pure TP (no PP, since autoregressive decode cannot pipeline) and often requires weight resharding from the training checkpoint format. See Multi-GPU & Multi-Node Inference for the inference-side parallelism story.

Further Reading¶

• Shoeybi, Patwary, et al. “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism” (2019). The original Megatron paper introducing TP for transformers.
• Narayanan, Shoeybi, et al. “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM” (NeurIPS 2021). Introduces the 3-D parallelism framework and the 1F1B pipeline schedule.
• Korthikanti, Casper, et al. “Reducing Activation Recomputation in Large Transformer Models” (MLSys 2023). Introduces sequence parallelism and selective activation recomputation.
• Rajbhandari, Rasley, Ruwase, He. “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models” (SC 2020). The ZeRO paper.
• Rajbhandari, et al. “ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning” (SC 2021). ZeRO-Offload and NVMe offloading.
• Ren, et al. “ZeRO-Offload: Democratizing Billion-Scale Model Training” (USENIX ATC 2021).
• NVIDIA Megatron-Core GitHub: NVIDIA/Megatron-LM — the canonical reference implementation.
• DeepSpeed GitHub: deepspeedai/DeepSpeed (formerly microsoft/DeepSpeed) — ZeRO implementation and tutorials.
• PyTorch torchtitan GitHub: pytorch/torchtitan — the PyTorch-native ND-parallel pretraining reference stack.
• Chowdhery, et al. “PaLM: Scaling Language Modeling with Pathways” (2022). Describes the 4D parallelism and MFU analysis methodology used at Google.

Exercises¶

1. The chapter’s rule of thumb is to place TP within a node and use PP to span nodes. Explain the hardware reason for each half of this rule. What specifically goes wrong if you set TP=16 on a cluster of 8-GPU nodes connected by InfiniBand?

2. A model has \(P = 13 \times 10^9\) parameters and is trained with Adam in mixed precision (bf16 params + bf16 grads + fp32 master + fp32 Adam moments), i.e. the 16-bytes-per-parameter accounting from Step 1. Using the ZeRO-3 formula \(\text{Memory per GPU} = 16P/\text{DP}\) (model state only, ignoring activations), what is the minimum DP degree needed to fit the model state within 40 GB per GPU? What is the per-GPU model state at that DP degree?

3. You launch a run with pipeline degree \(PP = 8\) and \(m = 24\) micro-batches per optimizer step, using the plain 1F1B schedule. (a) Compute the bubble fraction. (b) Your target is a bubble under 5%. Using the chapter’s rule \(m \ge 19(PP-1)\), how many micro-batches would the plain schedule need? © Instead you enable an interleaved (virtual pipeline) schedule with \(v = 4\) chunks per stage, keeping \(m = 24\). Estimate the new bubble fraction.

4. A 30B dense model is trained on 256 H100 GPUs. The observed throughput is 900 tokens/second/GPU. Using the simplified \(6P\) rule and the chapter’s dense bf16 H100 peak of \(9.89 \times 10^{14}\) FLOP/s per GPU, compute the MFU. Is this in the healthy range the chapter cites? Then, if the run uses full activation recomputation, estimate the corresponding HFU.

5. Extend the chapter’s compute_mfu function into a compute_hfu helper that accepts a recompute_mode argument ("none", "selective", or "full") and returns the HFU. Use the chapter’s overhead figures: "none" adds 0%, "selective" adds ~10% (a representative value from the chapter’s 5-15% range), and "full" multiplies model FLOPs by \(4/3\) (the \(8P\) vs \(6P\) ratio). Keep the style consistent with the chapter’s code.

def compute_hfu(
    model_params: int,
    tokens_per_second: float,
    peak_flops_per_sec: float,
    recompute_mode: str = "none",   # "none" | "selective" | "full"
    n_layers: int = None,
    d_model: int = None,
    seq_len: int = None,
) -> float:
    """
    Compute Hardware FLOP Utilization.

    HFU = MFU * (total FLOPs issued / model FLOPs), where the multiplier
    depends on how much activation recomputation is enabled:
      - "none":      +0%   -> factor 1.00 (HFU == MFU)
      - "selective": ~+10% -> factor 1.10 (recompute attention softmax/dropout)
      - "full":      +33%  -> factor 4/3  (an extra full forward pass, 8P vs 6P)
    """
    mfu = compute_mfu(
        model_params=model_params,
        tokens_per_second=tokens_per_second,
        peak_flops_per_sec=peak_flops_per_sec,
        n_layers=n_layers,
        d_model=d_model,
        seq_len=seq_len,
    )

    recompute_factor = {
        "none": 1.0,
        "selective": 1.10,
        "full": 4.0 / 3.0,
    }
    if recompute_mode not in recompute_factor:
        raise ValueError(f"unknown recompute_mode: {recompute_mode!r}")

    return mfu * recompute_factor[recompute_mode]


# Example: 30B model, 256 H100s, 900 tok/s/gpu, full recompute
H100_BF16_TFLOPS = 9.89e14
hfu = compute_hfu(
    model_params=30e9,
    tokens_per_second=900,          # per-GPU rate vs per-GPU peak
    peak_flops_per_sec=H100_BF16_TFLOPS,
    recompute_mode="full",
)
print(f"HFU: {hfu:.1%}")  # ~21.8% (Exercise 4's 21.9% rounds MFU to 0.164 first)

6. In initialize_model_parallel (the parallel_state code), consider a world of 8 GPUs configured with tensor_model_parallel_size=2 and pipeline_model_parallel_size=2. (a) Work out dp_size. (b) List the exact rank membership of every TP group, PP group, and DP group produced by the three loops. © A colleague proposes calling get_data_parallel_group() to target the collective for a tensor-parallel all-reduce inside a column-parallel GEMM. Explain why that is wrong and which accessor is correct.

i=0 -> ranks [0, 1]
i=1 -> ranks [2, 3]
i=2 -> ranks [4, 5]
i=3 -> ranks [6, 7]

i=0 -> range(0,8,2)[:2] = [0, 2]
i=1 -> range(1,8,2)[:2] = [1, 3]
i=2 -> range(2,8,2)[:2] = [2, 4]
i=3 -> range(3,8,2)[:2] = [3, 5]

i=0 -> [0, 4]
i=1 -> [1, 5]
i=2 -> [2, 6]
i=3 -> [3, 7]