Megatron-LM: what the tensor parallelism paper actually says

Reading Shoeybi et al. (NVIDIA, 2019) and Narayanan et al. (2021) after a training job that OOMed inside a single transformer layer.

The ZeRO job ran fine at 7B parameters. We moved to 13B and the optimizer state partitioning held up. Then we tried 30B, and the OOM came from somewhere ZeRO doesn't touch: the feedforward layer weights themselves. A single FFN layer in a 30B transformer has weight matrices large enough that they don't fit on one GPU even before you consider activations or optimizer state. ZeRO partitions the redundant copies of state across GPUs — but if the weight matrix is unique and larger than your device memory, partitioning copies doesn't help.

That's the problem Megatron-LM is solving. "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism", Shoeybi, Patwary, Puri, LeGresley, Casper, and Catanzaro, NVIDIA 2019 — trained 8.3B parameter models on 512 GPUs achieving 76% scaling efficiency and 15.1 petaFLOPs system-wide. The follow-up paper, "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM," Narayanan et al. 2021, adds pipeline parallelism and demonstrates training of a 1 trillion parameter model on 3072 GPUs at 502 petaFLOP/s. The specific mechanism — how weight matrices are split, how many communication operations this requires, and what determines the practical limit — is worth understanding before you configure a distributed training job.

The problem ZeRO doesn't solve

ZeRO's insight is that data parallel training replicates the optimizer state identically across all GPUs, and you can instead partition it: each GPU holds 1/N of the Adam moments, 1/N of the FP32 master weights. This works well because the optimizer update is independent per parameter — the moment estimate for parameter 42 depends only on gradients at parameter 42, never on what's happening at parameter 43.

This breaks down when the parameter count per layer exceeds device memory. Consider a 30B parameter transformer. The FFN sublayer typically expands the hidden dimension by 4x: if hidden dimension is 7168 (common at that scale), the FFN has weight matrices of shape [7168, 28672] and [28672, 7168]. At FP16 that's 7168 × 28672 × 2 bytes ≈ 411 MB per matrix, or about 822 MB just for the FFN weights per layer. Add the forward activations, attention weights, QKV projections — a single layer easily exceeds 2-3 GB. On an 80 GB A100, you can stack maybe 20-25 such layers before you're out of room, but the model has 60.

ZeRO Stage 3 partitions even the parameters themselves (each GPU holds 1/N of all parameters), which reduces the per-device parameter footprint. But it introduces an all-gather before every forward pass to reconstitute the parameters — sequential layers require sequential all-gathers, which adds latency on the critical path. It's workable but creates a different bottleneck.

Megatron takes a different approach: don't reconstitute the full weight matrix at all. Split the computation itself so that each GPU only ever needs to store and compute with its shard of the matrix, and design the communication pattern so that the results are correct without a full gather.

How transformer layers actually decompose

To understand why tensor parallelism works cleanly for transformers, you need to look at how the computation is structured inside each layer.

A transformer block has two sublayers: multi-head self-attention and a position-wise feedforward network. Both are dominated by matrix multiplications. Consider the FFN first because it's simpler.

The FFN computes:

output = GeLU(xW1) W2

Where x is the input of shape [batch, seq, hidden], W1 is [hidden, 4×hidden], and W2 is [4×hidden, hidden]. This is two matrix multiplications with a nonlinearity in between.

Now observe: W1 can be partitioned column-wise. Split it into two halves [W1_left | W1_right] where each half has shape [hidden, 2×hidden]. GPU 0 computes GeLU(xW1_left) and GPU 1 computes GeLU(xW1_right). Each GPU has a partial result of shape [batch, seq, 2×hidden] — the left or right half of the intermediate activations.

For the second matrix multiplication, W2 is partitioned row-wise. Split it into [W2_top; W2_bottom] where each half has shape [2×hidden, hidden]. GPU 0 applies W2_top to its partial activations, GPU 1 applies W2_bottom to its. Each GPU now has a [batch, seq, hidden] result that represents a partial sum of the final output.

To get the correct final output, you sum the partial results. One all-reduce across GPUs.

The key insight: you've replaced one large matrix multiplication with two smaller ones — each using 1/N the memory — and the only communication cost is a single all-reduce at the end. The nonlinearity (GeLU) is applied independently on each GPU because it's element-wise and each GPU has the correct portion of the intermediate state.

The attention heads split

The self-attention sublayer parallelizes even more naturally. Multi-head attention with H heads processes H independent attention computations in parallel — that's the entire design of multi-head attention. Each head has its own Q, K, V projections and output projection.

With tensor parallelism degree N, you assign H/N heads to each GPU. The Q, K, V weight matrices (each of shape [hidden, hidden]) are split column-wise — each GPU holds columns corresponding to its heads. The output projection matrix (also [hidden, hidden]) is split row-wise. Each GPU computes the full attention for its assigned heads and produces a partial output. One all-reduce at the end to sum partial outputs.

The communication structure is identical to the FFN: one all-reduce per sublayer, two all-reduces per transformer block total. This is the result that makes the approach practical: regardless of tensor parallel degree N, every transformer block requires exactly two all-reduces. You can split across 2 GPUs or 8 GPUs and the communication count stays constant. Only the message size changes.

The paper reports that with this scheme, training an 8.3B parameter model across 512 GPUs achieves 76% scaling efficiency relative to a single-GPU baseline, with each GPU sustaining 39 TeraFLOPs — about 30% of the A100's theoretical peak. The 30% number looks low but is typical for attention-heavy workloads; the bound is bandwidth, not compute.

Why tensor parallelism must stay within a node

The all-reduce in tensor parallelism is on the critical path of the forward pass. Unlike the gradient all-reduce in data parallelism (which happens after the backward pass and can overlap with optimization), the tensor parallel all-reduce happens between layers. You can't proceed to the next layer until the all-reduce completes.

This makes bandwidth the dominant constraint. An NVLink connection between GPUs within an A100 node runs at ~600 GB/s bidirectional. An InfiniBand connection between nodes runs at ~25-100 GB/s depending on generation. Putting the tensor parallel all-reduces over InfiniBand would increase per-all-reduce latency by 6-24x, directly slowing the forward pass.

The practical rule: tensor parallelism degree should match the number of GPUs within a single node. For 8-GPU nodes, tensor parallel degree 8. For 4-GPU nodes, degree 4. Don't split tensor parallel groups across nodes if you can avoid it.

This constraint also gives you the decomposition for larger clusters:

Within a node: tensor parallelism (fast NVLink, on critical path)
Across nodes: pipeline parallelism (model stages, communication only at stage boundaries)
Across node groups: data parallelism (gradient sync, off critical path)

This combination — tensor parallel × pipeline parallel × data parallel — is what the 2021 paper calls 3D parallelism.

Pipeline parallelism and the bubble problem

Pipeline parallelism divides the model's layers into sequential stages, each assigned to a different device (or group of devices). Stage 0 handles layers 1-8, stage 1 handles layers 9-16, and so on. The forward pass flows from stage 0 to the last stage; the backward pass flows back.

The fundamental problem with naïve pipeline parallelism: only one stage is active at a time during each forward or backward pass. If you have 4 pipeline stages, 3 of them are idle while one is computing. The fraction of time spent idle — the "pipeline bubble" — is (p-1)/p where p is the number of stages. With 8 stages, you're wasting 87.5% of compute.

The standard mitigation is microbatching: split each batch into smaller microbatches and pipeline them. Stage 0 processes microbatch 1 and immediately starts microbatch 2 while stage 1 processes microbatch 1. With m microbatches and p stages, the bubble fraction reduces to (p-1)/(m + p - 1). With 8 stages and 32 microbatches, you're wasting about 18% — a real cost, but manageable.

The 2021 paper introduces an interleaved schedule that goes further. Instead of each GPU handling a contiguous block of layers (layers 1-8 on GPU 0, layers 9-16 on GPU 1), each GPU handles multiple non-contiguous chunks (layers 1-4 and 17-20 on GPU 0, layers 5-8 and 21-24 on GPU 1). This reduces the bubble fraction by a factor equal to the number of chunks per GPU, at the cost of doubling or tripling the number of pipeline communication operations (because each stage boundary now requires a send/receive). The paper reports this interleaved schedule improves throughput by more than 10% in their configurations.

The 3D parallelism decomposition in practice

For a large training job, the configuration looks something like this. Suppose you have 256 GPUs in 32 nodes of 8 GPUs each:

tensor parallel degree (t) = 8   # within each node, NVLink
pipeline stages (p) = 8          # across 8 nodes  
data parallel replicas (d) = 4   # 4 groups of 64 GPUs each

total GPUs = t × p × d = 8 × 8 × 4 = 256 ✓

Each of the 4 data-parallel replicas processes a different slice of the global batch. Gradients from all 4 replicas are averaged after the backward pass — this all-reduce is off the critical path and happens over InfiniBand between replica groups.

Within each replica, the 64 GPUs are organized as 8 pipeline stages × 8 tensor-parallel GPUs per stage. The pipeline partitions the model layers across 8 node groups; within each node, tensor parallelism splits each layer's weight matrices across 8 GPUs.

The 2021 paper demonstrates this running a 1 trillion parameter model on 3072 GPUs (384 nodes), achieving 502 petaFLOP/s at 52% of theoretical peak throughput per GPU. The 52% peak utilization is notably higher than the 30% single-GPU number from the 2019 paper — the system efficiency improves because the larger batch sizes enabled by more GPUs increase arithmetic intensity.

What determines your tensor parallel degree

The limiting factor on tensor parallel degree is not memory or bandwidth — it's the hidden dimension of the model.

In attention, you're splitting H heads across N GPUs. H must be divisible by N. Most models are designed with H a multiple of 32 or 64; with N=8 you need H ≥ 8, which all practical models satisfy. With N=16 you need H ≥ 16, fine. With N=32 you're starting to reduce heads per GPU to the point where the attention computation per head becomes small relative to the communication overhead.

In the FFN, you're splitting the intermediate dimension (4 × hidden) across N GPUs. The intermediate dimension per GPU must be large enough that the matrix multiplication stays compute-bound rather than overhead-bound. Roughly, you want at least 512-1024 intermediate neurons per GPU. Below that, the all-reduce dominates and you're slower than not parallelizing.

The 2019 paper doesn't recommend going above tensor parallel degree 8 for the architectures they tested. The 2021 paper uses degree 8 as the standard configuration for 8-GPU nodes. Increasing beyond the node boundary requires cross-node NVLink or accepting InfiniBand latency on the critical path — typically not worth it.

What actually limits throughput in production

The numbers in the papers are measured under ideal conditions: synthetic workloads, no checkpoint saving, no evaluation, no dynamic sequence lengths. In real training runs, several factors reduce effective throughput:

Activation memory vs. recomputation trade-off. The paper uses activation checkpointing (recomputing activations during the backward pass rather than storing them). This halves the activation memory footprint but roughly doubles the computation cost of the backward pass. The papers account for this in their throughput numbers, but your job may have different trade-offs depending on sequence length and batch size.

Microbatch size and pipeline efficiency. To minimize the pipeline bubble, you want many microbatches. But each microbatch must be large enough to keep the compute units busy. Small microbatches reduce arithmetic intensity and drop GPU utilization. The optimal microbatch count depends on your model dimensions, sequence length, and hardware — the papers tune this empirically.

Communication-computation overlap. Data parallel gradient all-reduces can overlap with computation during the backward pass if the implementation supports it. Megatron-LM implements this. But tensor parallel and pipeline parallel communications are harder to overlap because they're on the critical path. The reported 76% scaling efficiency partly reflects the cost of this unavoidable communication.

Load imbalance across pipeline stages. If the embedding layers (stage 0) are significantly cheaper to compute than the transformer layers, stage 0 finishes early and sits idle while downstream stages catch up. The 2021 paper addresses this by placing embedding tables on the first and last stages so that their compute is distributed — a detail that matters for very large vocabulary sizes.

When NOT to use Megatron-style tensor parallelism

Model fits on one device. If your model parameters, activations, and optimizer state fit within the device memory budget using ZeRO (stages 1 or 2), tensor parallelism adds communication overhead for no memory benefit. Tensor parallelism is specifically for when layers themselves don't fit — not just the optimizer state.

You don't have high-bandwidth intra-node interconnect. Tensor parallelism puts all-reduces on the forward-pass critical path. Without NVLink (or equivalent), those all-reduces go over PCIe or Ethernet and your effective throughput can be worse than running a smaller model without parallelism. The 2019 paper specifically notes this as a requirement.

You're doing inference, not training. Inference has different memory and latency constraints. Tensor parallelism can reduce latency for a single request by parallelizing the compute, but it increases the minimum batch size needed to saturate all GPUs and adds communication overhead that matters more when generating token by token. For serving, the tradeoffs are different enough that you should look at inference-specific solutions (vLLM, TensorRT-LLM) rather than directly applying Megatron's training configuration.

Small tensor parallel degree with large layers. If your hidden dimension per GPU drops below ~512-1024, the matrix multiplications become memory-bandwidth bound rather than compute-bound, and the all-reduce overhead dominates. You'll often get better throughput by not parallelizing those layers and instead increasing pipeline stages or data parallel degree.

Very large pipeline depth. Pipeline parallelism with many stages requires large batch sizes to keep the pipeline full (to minimize the bubble fraction). If your task requires small batch sizes — low latency requirements, fine-tuning on small datasets — pipeline depth is hard to justify. The bubble fraction (p-1)/(m+p-1) only improves with large m (microbatch count), which requires large total batches.

The layernorm placement the paper flags as critical

The 2019 paper notes, somewhat cryptically, that "careful attention to layer normalization placement proved critical to achieving increased performance as the model size grows." This refers to the pre-norm vs. post-norm distinction.

Original transformer architecture places LayerNorm after the residual addition (post-norm). Megatron switches to pre-norm — LayerNorm before the sublayer, applied to the residual stream before it goes into attention or FFN. At large scale, pre-norm produces substantially more stable training: gradients don't vanish as deeply, and you can train without learning rate warmup periods that were required to stabilize post-norm at scale.

GPT-2, GPT-3, and most subsequent large models adopted pre-norm. It's now the default. If you're reading codebases that use Megatron or are derived from it (LLaMA, Falcon, Mistral), the LayerNorm is before the attention and FFN, not after.

The paper's actual claims vs. what people assume

A few things the Megatron papers do and don't claim:

They claim: 76% scaling efficiency from 1 to 512 GPUs using tensor parallelism for 8.3B parameters. This is measured as throughput-per-GPU relative to a single-GPU baseline, not as absolute efficiency relative to perfect linear scaling.

They don't claim: that tensor parallelism is always the right first choice. The 2021 paper is explicit that data parallelism alone should be used first (it has the lowest communication overhead), pipeline parallelism second (communication only at stage boundaries), and tensor parallelism last (communication on critical path). The recommendation is to scale out in that order.

They claim: two all-reduces per transformer block, constant regardless of tensor parallel degree. This is the core efficiency result. The count is fixed; only the message size changes.

They don't claim: that 3D parallelism is always better than 2D. The optimal configuration depends on model architecture, hardware topology, and batch size constraints. The papers provide guidelines and empirical results, not universal prescriptions.

They claim: the interleaved pipeline schedule reduces the bubble by a factor equal to the number of interleaved chunks, while approximately doubling communication operations at stage boundaries. Whether this is net positive depends on the relative cost of computation vs. stage-boundary communication for your hardware.

If you're sizing a training job beyond what ZeRO handles: keep tensor parallelism within the node, use pipeline parallelism to span nodes, scale with data parallelism after those are configured, and tune microbatch count to keep the pipeline bubble below 10-15% of compute time. The 2021 paper's appendix has a cost model that lets you estimate bubble fraction given your p and m values — worth working through before you provision the cluster.