Sequence Parallelism

Tensor parallelism splits weight matrices but leaves LayerNorm, dropout, and other non-GEMM operations replicated on every GPU — wasting activation memory. Sequence parallelism fixes this by splitting these operations along the sequence dimension.

---

Prerequisite: Read Tensor Parallelism first — SP extends TP.

1. The Problem: Replicated Activations

With tensor parallelism, GEMM outputs are split across GPUs. But before and after each GEMM, there are element-wise operations (LayerNorm, dropout, residual add) that operate on the full activation tensor:

TP GEMM: each GPU computes 1/P of the output → split
LayerNorm: needs the full tensor → AllGather → replicated on all GPUs
Dropout: operates on full tensor → replicated
Next TP GEMM: needs split input → ReduceScatter → split again

Between each pair of TP GEMMs, the full activation is replicated on all GPUs. For a 70B model with $B=1, S=4096, d=8192$ in BF16: each replicated activation = $4096 \times 8192 \times 2 = 64$ MB. With 80 layers × 2 replication points = 10.2 GB wasted on redundant copies.

---

2. The Fix: Split Non-GEMM Ops Along Sequence

Instead of replicating activations for LayerNorm/dropout, keep them split along the sequence dimension:

Without SP: [LayerNorm on full S] → [TP GEMM, split] → [AllGather to full S] → [LayerNorm]
With SP:    [LayerNorm on S/P] → [AllGather] → [TP GEMM, split] → [ReduceScatter to S/P] → [LayerNorm on S/P]

Each GPU holds only $S/P$ tokens’ activations for non-GEMM ops. Activation memory per GPU drops by $P\times$ for these operations.

Communication change: The AllGather and ReduceScatter that were already needed for TP are simply repositioned — total communication volume stays the same.

---

Worked Example

70B model, TP=8, B=1, S=4096, d=8192, BF16:

• Per activation tensor: $4096 × 8192 × 2 = 64$ MB
• Without SP: LayerNorm activations replicated on all 8 GPUs → 64 MB × 2 (per block) × 80 layers = 10.2 GB wasted per GPU
• With SP: each GPU holds S/8 = 512 tokens → 8 MB per activation × 2 × 80 = 1.3 GB per GPU
• Savings: 8.9 GB per GPU — enough for a larger batch size or avoiding activation checkpointing

3. Memory Savings

Component	Without SP	With SP
GEMM activations	Split (1/P per GPU)	Split (1/P) — same
LayerNorm activations	Full (replicated)	1/P per GPU
Dropout activations	Full (replicated)	1/P per GPU
Residual activations	Full (replicated)	1/P per GPU

For the 70B example with TP=8: non-GEMM activation memory drops from 10.2 GB to 1.3 GB — saving ~9 GB per GPU.

This saved memory can be used for larger batch sizes (higher throughput) or to avoid activation checkpointing (fewer recomputation FLOPs).

---

4. Sequence Parallelism in Practice

SP is always used together with TP — it’s not an independent parallelism dimension. In Megatron-LM, it’s enabled by default when TP is used:

# Megatron-LM: SP is automatic with TP
args.sequence_parallel = True  # default when tp_size > 1

SP is essentially “free” — same communication volume, less memory. Always enable SP when using TP.

---

Self-Check

• Explain why LayerNorm/dropout activations are replicated without SP
• Calculate the activation memory savings from SP for a given model
• Explain why SP doesn’t increase communication volume
• Explain why SP is always used together with TP (not independently)

---

References

• Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022)
• Megatron-LM — Sequence Parallelism