Model Parallelism Explained: Tensor, Pipeline, and Sequence Parallelism

Note: This is a review article (no first-hand experimental data), based on Megatron-LM and public large-model training practice.

Experiment Objective

When one card cannot hold the model, beyond data parallelism (ch09 DDP/FSDP) there is model parallelism: tensor parallelism slices one layer across cards, pipeline parallelism assigns layers to stages, sequence parallelism slices the sequence dimension. Goal: clarify the principles (how to slice) → memory and communication trade-offs → Megatron practice → the engineering matrix for combining with data parallelism. After reading you can answer: which parallel for which scenario, how costly the communication, and how to combine.

Code Walkthrough

Part 1: Tensor parallelism — slicing one layer

# Tensor Parallelism (TP): matrix ops inside a layer sliced across cards
# Linear layer example (Megatron column/row split):
#   Column: W split by columns W=[W1,W2]; each card computes Y_i = X·W_i; AllReduce merge
#   Row: W split by rows; partial sums per card; AllReduce accumulate
# Communication: one AllReduce per layer (proportional to layer count)
# Memory: per card 1/TP weights + 1/TP activations
def tp_linear(X, W, tp=8):
    W_split = split_col(W, tp)            # each card holds 1/8
    Y_local = X @ W_split[rank]           # local part
    Y = all_reduce(Y_local)               # one in-layer communication
    return Y

Key points:

1. TP is "intra-layer slicing": matrix ops split across cards, partial results AllReduced — communication scales with layer count (sync every layer).
2. Memory gain is direct: weights and activations at 1/TP — TP primarily solves "one layer doesn't fit" (huge hidden sizes).
3. TP has the heaviest communication: every layer AllReduces, so TP needs high-speed interconnect (ch03) and is generally confined to one node (8 cards NVLink).

Part 2: Pipeline parallelism — assigning layers to stages

# Pipeline Parallelism (PP): layers sliced into stages; cards assigned by stage
#   Model layers → PP stages; stage i's cards compute only its own layers
# Data flows through the pipeline as micro-batches (1F1B scheduling reduces bubbles)
# Communication: only stage-boundary activations/gradients (low-frequency, << TP)
# Memory: per card only this stage's weights+activations (some pipeline buffering)
# Bubble cost: larger PP → bigger bubbles (if poorly scheduled)
def pp_forward(model, micro_batches, pp=4):
    for mb in micro_batches:
        out = stage[rank](mb)      # each card runs its own stage
        send(out, rank+1)          # low-frequency inter-stage communication

Key points:

• PP is "inter-layer slicing": layers into stages; each card holds its stage — solves "the whole model doesn't fit," not "one layer doesn't fit."
• Low-frequency communication: only stage-boundary data transfer, far less than TP — PP is interconnect-friendly (cross-node OK).
• The cost is bubbles: pipeline ramp-up/drain leaves idle cards; 1F1B mitigates but over-large PP increases bubbles — PP gains diminish with scale.

Principles

Three parallelisms compared (estimated)

Dimension	Tensor TP	Pipeline PP	Sequence SP
slicing	intra-layer matrix	layers into stages	sequence dimension
solves	one layer too big	whole model too big	long-sequence activations
communication	AllReduce per layer (heavy)	stage boundary (light)	attention region (medium)
interconnect need	high (intra-node)	low (cross-node)	medium
memory gain	weights+activations 1/TP	weights 1/PP	activations 1/SP

Core insight: the three solve different problems — TP for "layer too big," PP for "model too long," SP for "sequence too long" (SP links to ch10's long-context pretraining). Communication and memory gain are inverse: TP gains directly but communicates heavily; PP communicates lightly but bubbles; SP only compresses activations.

Combination matrix: data parallel × model parallel (3D parallelism)

Large-model training = Data Parallel (DP/FSDP) × TP × PP
  DP handles "data replicas" (throughput) → memory-friendly (FSDP shards)
  TP handles "intra-layer" (one layer) → intra-node 8 cards
  PP handles "inter-layer" (whole model) → cross-node
Combination rules (empirical):
  TP stays in-node (NVLink bandwidth); PP crosses nodes (light communication);
  DP maximizes (highest parallelism); deepen TP/PP when memory is short.

Combination judgment: start DP/FSDP to saturate data parallelism; add TP when one layer doesn't fit (intra-node); add PP when the whole model doesn't fit (cross-node) — deepen "data first, then intra-layer, then inter-layer" (linking ch11's Megatron practice).

Results

(Estimated, not measured; illustrative trends)

Option	Trainable scale	Communication pressure	Bubbles
DP/FSDP	medium (sharded)	low-medium	none
+ TP(8)	large (one layer)	high (intra-node)	none
+ PP	larger (whole model)	low (inter-stage)	some (controllable)

Engineering Takeaways

1. Decide first "what doesn't fit": one layer → TP; whole model → PP; long sequence → SP — prescribe for the symptom; don't start with 3D parallelism.
2. TP stays in-node: TP AllReduces every layer; must pair with intra-node high-speed interconnect (ch03/ch01 NVLink); cross-node TP communication explodes.
3. PP communicates lightly but bubbles: low-frequency inter-stage suits cross-node; 1F1B controls bubbles; don't over-size PP.
4. Combine "data first, then intra-layer, then inter-layer": DP/FSDP base → TP for layers → PP for the model (linking ch11's Megatron practice).

Next article preview: model parallelism done; one more pretraining piece — "Post-Training Alignment: The Compute and Engineering of SFT, RLHF/DPO."