mirrors/wanshuiyin--auto-claude-code-research-in-sleep
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
wanshuiyin--auto-claude-cod…/docs/tutorials/distributed_training_tutorial_en.html
T
wehub-resource-sync a302add0a2 chore: import upstream snapshot with attribution
2026-07-13 13:37:02 +08:00

1430 lines
108 KiB
HTML
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Distributed Training Tutorial En</title>
<meta name="generator" content="ARIS render-html (academic, v1)">
<meta name="aris:source-path" content="docs/tutorials/distributed_training_tutorial_en.md">
<meta name="aris:source-sha256" content="9ee6b9a3603bcd089bca124fa44f3bdb41d7b5604595f3a506cb9ff601bcd999">
<meta name="aris:generated-at" content="2026-05-19 19:00 UTC">
<!-- MathJax 3 -->
<script>
window.MathJax = {
tex: { inlineMath: [['$', '$'], ['\\(', '\\)']], displayMath: [['$$', '$$'], ['\\[', '\\]']], processEscapes: true },
options: { skipHtmlTags: ['script', 'noscript', 'style', 'textarea', 'pre', 'code'] }
};
</script>
<script src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js" async></script>
<!-- highlight.js -->
<link rel="stylesheet" href="https://cdn.jsdelivr.net/gh/highlightjs/cdn-release@11.9.0/build/styles/atom-one-light.min.css">
<script src="https://cdn.jsdelivr.net/gh/highlightjs/cdn-release@11.9.0/build/highlight.min.js"></script>
<script>document.addEventListener('DOMContentLoaded', () => hljs.highlightAll());</script>
<style>
:root {
--bg: #fdfcf7;
--bg-soft: #f4f1ea;
--bg-code: #f8f5ec;
--ink: #1a1a1a;
--ink-soft: #4a4a4a;
--ink-muted: #6b6b6b;
--primary: #1a4a8c;
--primary-soft: #2d6cb8;
--accent: #b8390e;
--warn: #b45309;
--warn-bg: #fef3c7;
--info-bg: #dbeafe;
--good-bg: #d1fae5;
--good: #065f46;
--bad-bg: #fee2e2;
--bad: #991b1b;
--border: #d6d0c0;
--border-soft: #e8e3d5;
}
* { box-sizing: border-box; }
html { scroll-behavior: smooth; }
body {
font-family: "Source Serif Pro", "Source Serif 4", "Crimson Pro", "Georgia", "Songti SC", "STSong", serif;
line-height: 1.65;
color: var(--ink);
background: var(--bg);
margin: 0;
padding: 0;
font-size: 16px;
}
.layout {
max-width: 1280px;
margin: 0 auto;
display: grid;
grid-template-columns: 260px 1fr;
gap: 48px;
padding: 40px 32px;
}
nav.toc {
position: sticky;
top: 24px;
align-self: start;
font-size: 13px;
max-height: calc(100vh - 48px);
overflow-y: auto;
border-right: 1px solid var(--border-soft);
padding-right: 16px;
}
nav.toc h3 {
margin: 0 0 12px;
font-size: 12px;
text-transform: uppercase;
letter-spacing: 0.08em;
color: var(--ink-muted);
font-weight: 600;
}
nav.toc ol { list-style: none; padding: 0; margin: 0; counter-reset: toc; }
nav.toc ol li { margin: 5px 0; counter-increment: toc; }
nav.toc ol li::before { content: counter(toc) ". "; color: var(--ink-muted); margin-right: 4px; }
nav.toc a {
color: var(--ink-soft);
text-decoration: none;
border-bottom: 1px dotted transparent;
}
nav.toc a:hover { color: var(--primary); border-bottom-color: var(--primary); }
nav.toc ul { list-style: none; padding-left: 14px; margin: 3px 0; font-size: 12px; }
nav.toc ul li::before { content: "→ "; color: var(--border); }
main { min-width: 0; }
header.hero {
border-bottom: 3px double var(--primary);
padding-bottom: 24px;
margin-bottom: 32px;
}
header.hero .eyebrow {
color: var(--accent);
font-size: 13px;
text-transform: uppercase;
letter-spacing: 0.12em;
font-weight: 600;
margin-bottom: 8px;
}
header.hero h1 {
font-size: 32px;
line-height: 1.2;
margin: 0 0 12px;
color: var(--ink);
font-weight: 700;
letter-spacing: -0.01em;
}
header.hero .subtitle {
font-size: 16px;
color: var(--ink-soft);
margin: 0 0 8px;
font-style: italic;
}
header.hero .byline {
font-size: 14px;
color: var(--ink-soft);
margin: 0 0 20px;
}
header.hero .byline strong {
color: var(--ink);
font-weight: 600;
}
header.hero .meta {
display: flex;
gap: 20px;
flex-wrap: wrap;
font-size: 12px;
color: var(--ink-muted);
border-top: 1px solid var(--border-soft);
padding-top: 14px;
}
header.hero .meta span strong { color: var(--ink-soft); }
header.hero .meta code {
font-family: "JetBrains Mono", "SF Mono", "Menlo", "Consolas", monospace;
font-size: 11px;
background: var(--bg-soft);
padding: 1px 5px;
border-radius: 3px;
border: 1px solid var(--border-soft);
}
h2 {
font-size: 24px;
margin: 44px 0 14px;
padding-bottom: 8px;
border-bottom: 1px solid var(--border);
color: var(--ink);
font-weight: 700;
}
h2 .num { color: var(--primary); font-weight: 600; margin-right: 8px; }
h3 { font-size: 19px; margin: 28px 0 10px; color: var(--primary); font-weight: 600; }
h4 { font-size: 16px; margin: 20px 0 8px; color: var(--ink); font-weight: 600; }
p { margin: 10px 0; }
ul, ol { padding-left: 22px; margin: 10px 0; }
ul li, ol li { margin: 4px 0; }
ul li::marker { color: var(--primary); }
strong { color: var(--accent); font-weight: 600; }
em { color: var(--ink-soft); }
a { color: var(--primary); }
a:hover { color: var(--accent); }
code:not(.hljs) {
font-family: "JetBrains Mono", "SF Mono", "Menlo", "Consolas", monospace;
font-size: 0.86em;
background: var(--bg-code);
padding: 1px 5px;
border-radius: 3px;
border: 1px solid var(--border-soft);
color: var(--accent);
}
pre {
background: #fafaf6;
border: 1px solid var(--border);
border-left: 4px solid var(--primary);
padding: 0;
overflow-x: auto;
border-radius: 4px;
margin: 14px 0;
}
pre code, pre code.hljs {
background: transparent !important;
display: block;
padding: 14px 18px !important;
font-size: 13px;
line-height: 1.55;
font-family: "JetBrains Mono", "SF Mono", "Menlo", monospace;
color: var(--ink);
}
pre.diagram {
background: #f9f6ed;
border-left: 4px solid var(--accent);
font-size: 12.5px;
line-height: 1.4;
}
.callout {
margin: 16px 0;
padding: 12px 16px;
border-radius: 4px;
border-left: 4px solid;
font-size: 15px;
}
.callout-title {
font-weight: 600;
margin-bottom: 6px;
font-size: 12px;
text-transform: uppercase;
letter-spacing: 0.06em;
}
.callout-info { background: var(--info-bg); border-left-color: var(--primary); }
.callout-info .callout-title { color: var(--primary); }
.callout-warn { background: var(--warn-bg); border-left-color: var(--warn); }
.callout-warn .callout-title { color: var(--warn); }
.callout-good { background: var(--good-bg); border-left-color: var(--good); }
.callout-good .callout-title { color: var(--good); }
.callout-bad { background: var(--bad-bg); border-left-color: var(--bad); }
.callout-bad .callout-title { color: var(--bad); }
table {
width: 100%;
border-collapse: collapse;
margin: 16px 0;
font-size: 14px;
border: 1px solid var(--border);
border-radius: 4px;
overflow: hidden;
}
thead { background: var(--primary); color: white; }
th, td {
text-align: left;
padding: 9px 12px;
border-bottom: 1px solid var(--border-soft);
vertical-align: top;
}
th { font-weight: 600; font-size: 13px; letter-spacing: 0.02em; }
tr:last-child td { border-bottom: none; }
tbody tr:nth-child(even) { background: var(--bg-soft); }
details.qa, details {
background: white;
border: 1px solid var(--border-soft);
border-radius: 6px;
margin: 10px 0;
padding: 0;
}
details summary {
cursor: pointer;
padding: 10px 14px;
font-weight: 600;
font-size: 14px;
color: var(--primary);
list-style: none;
user-select: none;
}
details summary::-webkit-details-marker { display: none; }
details summary::before {
content: "▸ ";
margin-right: 4px;
display: inline-block;
transition: transform 0.15s;
}
details[open] summary::before { transform: rotate(90deg); }
details[open] summary { border-bottom: 1px solid var(--border-soft); }
details > :not(summary) { padding: 10px 14px; }
details p:first-of-type { margin-top: 8px; }
mjx-container[display="true"] { margin: 12px 0 !important; }
footer.aris-footer {
margin-top: 60px;
padding-top: 20px;
border-top: 1px solid var(--border);
font-size: 12px;
color: var(--ink-muted);
}
footer.aris-footer a { color: var(--ink-muted); border-bottom: 1px dotted var(--border); }
@media (max-width: 900px) {
.layout { grid-template-columns: 1fr; gap: 20px; padding: 20px 16px; }
nav.toc {
position: static;
max-height: none;
border-right: none;
border-bottom: 1px solid var(--border-soft);
padding-right: 0;
padding-bottom: 14px;
}
header.hero h1 { font-size: 24px; }
h2 { font-size: 20px; }
}
@media print {
nav.toc { display: none; }
.layout { grid-template-columns: 1fr; padding: 0; }
body { background: white; }
header.hero { border-bottom-color: var(--ink); }
}
</style>
</head>
<body>
<div class="layout">
<nav class="toc">
<h3>Contents</h3>
<ol>
<li><a href="#0-tldr-cheat-sheet">§0 TL;DR Cheat Sheet</a>
</li>
<li><a href="#1-intuition--why-one-gpu-cannot-hold-a-large-model">§1 Intuition — why one GPU cannot hold a large model</a>
</li>
<li><a href="#2-nccl-communication-primitives-and-topology">§2 NCCL communication primitives and topology</a>
<ul>
<li><a href="#21-the-five-collective-primitives">2.1 The five collective primitives</a></li>
<li><a href="#22-nvlink--ib--topology">2.2 NVLink / IB / topology</a></li>
<li><a href="#23-nccl-calls-in-pytorch">2.3 NCCL calls in PyTorch</a></li>
</ul>
</li>
<li><a href="#3-ddp--distributeddataparallel">§3 DDP — DistributedDataParallel</a>
<ul>
<li><a href="#31-algorithm-skeleton">3.1 Algorithm skeleton</a></li>
<li><a href="#32-bucket-fusion--overlap-ddp-engineering-essence">3.2 Bucket fusion + overlap (DDP engineering essence)</a></li>
<li><a href="#33-pytorch-code-with-overlap">3.3 PyTorch code (with overlap)</a></li>
<li><a href="#34-complexity">3.4 Complexity</a></li>
</ul>
</li>
<li><a href="#4-zero-123--zero-redundancy-optimizer">§4 ZeRO 1/2/3 — Zero Redundancy Optimizer</a>
<ul>
<li><a href="#41-three-stage-memory-math">4.1 Three-stage memory math</a></li>
<li><a href="#42-zero-3-workflow-most-commonly-used-forward--backward--optimize">4.2 ZeRO-3 workflow (most commonly used, forward / backward / optimize)</a></li>
<li><a href="#43-zero-123-vs-ddp-communication-comparison-important-interview-question">4.3 ZeRO-1/2/3 vs DDP communication comparison (important interview question)</a></li>
<li><a href="#44-zero-offload--zero-infinity">4.4 ZeRO-Offload / ZeRO-Infinity</a></li>
</ul>
</li>
<li><a href="#5-fsdp--fsdp2--pytorch-native-zero-3">§5 FSDP / FSDP2 — PyTorch native ZeRO-3</a>
<ul>
<li><a href="#51-fsdp1-vs-fsdp2-core-differences">5.1 FSDP1 vs FSDP2 core differences</a></li>
<li><a href="#52-fsdp-wrap-policy-the-most-important-design-decision">5.2 FSDP wrap policy (the most important design decision)</a></li>
<li><a href="#53-fsdp2--mixed-precision--activation-checkpoint">5.3 FSDP2 + mixed precision + activation checkpoint</a></li>
<li><a href="#54-fsdp-vs-zero-3-communication-difference-l3-question">5.4 FSDP vs ZeRO-3 communication difference (L3 question)</a></li>
<li><a href="#55-hsdp--hybrid-sharded-dp">5.5 HSDP — Hybrid Sharded DP</a></li>
</ul>
</li>
<li><a href="#6-zero--communication-optimizations-2023">§6 ZeRO++ — communication optimizations (2023)</a>
<ul>
<li><a href="#61-qwz-quantized-weight-all-gather-quantize-forward-comms">6.1 qwZ: Quantized Weight all-gather (quantize forward comms)</a></li>
<li><a href="#62-hpz-hierarchical-partition">6.2 hpZ: Hierarchical Partition</a></li>
<li><a href="#63-qgz-quantized-gradient-reduce">6.3 qgZ: Quantized Gradient reduce</a></li>
<li><a href="#64-combined-effect">6.4 Combined effect</a></li>
</ul>
</li>
<li><a href="#7-tensor-parallel--megatron-lm">§7 Tensor Parallel — Megatron-LM</a>
<ul>
<li><a href="#71-core-idea-column-parallel--row-parallel-pairing">7.1 Core idea: column-parallel → row-parallel pairing</a></li>
<li><a href="#72-communication-analysis">7.2 Communication analysis</a></li>
<li><a href="#73-tp-sharding-of-attention">7.3 TP sharding of attention</a></li>
<li><a href="#74-tp-code-skeleton">7.4 TP code skeleton</a></li>
<li><a href="#75-tp-must-fit-inside-a-node-nvlink">7.5 TP must fit inside a node (NVLink)</a></li>
</ul>
</li>
<li><a href="#8-sequence-parallel--saving-activation-memory">§8 Sequence Parallel — saving activation memory</a>
<ul>
<li><a href="#81-motivation-activation-memory-is-mostly-element-wise-ops-like-layernorm--dropout">8.1 Motivation: activation memory is mostly element-wise ops like LayerNorm / Dropout</a></li>
<li><a href="#82-sp-solution-shard-along-the-sequence-dim">8.2 SP solution: shard along the sequence dim</a></li>
<li><a href="#83-how-much-activation-memory-does-sp-save-l3-question">8.3 How much activation memory does SP save (L3 question)</a></li>
<li><a href="#84-selective-activation-recompute">8.4 Selective activation recompute</a></li>
</ul>
</li>
<li><a href="#9-pipeline-parallel">§9 Pipeline Parallel</a>
<ul>
<li><a href="#91-naive-pp-and-the-bubble">9.1 Naive PP and the bubble</a></li>
<li><a href="#92-gpipe-huang-et-al-neurips-2019">9.2 GPipe (Huang et al., NeurIPS 2019)</a></li>
<li><a href="#93-1f1b--pipedream-narayanan-sosp-2019-megatron-lm-2-sc-2021">9.3 1F1B / PipeDream (Narayanan SOSP 2019, Megatron-LM-2 SC 2021)</a></li>
<li><a href="#94-interleaved-1f1b-megatron-lm-2-sc-2021">9.4 Interleaved 1F1B (Megatron-LM-2, SC 2021)</a></li>
<li><a href="#95-1f1b-schedule-pseudo-code">9.5 1F1B schedule pseudo-code</a></li>
<li><a href="#96-pp-communication-characteristics">9.6 PP communication characteristics</a></li>
<li><a href="#97-pps-flaws-load-imbalance">9.7 PP&#x27;s flaws: load imbalance</a></li>
</ul>
</li>
<li><a href="#10-context-parallel--sharding-long-sequences">§10 Context Parallel — sharding long sequences</a>
<ul>
<li><a href="#101-ring-attention-liu-et-al-arxiv-231001889-2024">10.1 Ring Attention (Liu et al., arXiv 2310.01889, 2024)</a></li>
<li><a href="#102-llama-3-cp-implementation">10.2 Llama 3 CP implementation</a></li>
<li><a href="#103-cp-and-tp-orthogonality">10.3 CP and TP orthogonality</a></li>
</ul>
</li>
<li><a href="#11-expert-parallel--moe-routing">§11 Expert Parallel — MoE routing</a>
<ul>
<li><a href="#111-basic-moe-structure">11.1 Basic MoE structure</a></li>
<li><a href="#112-expert-parallel-experts-across-gpus">11.2 Expert Parallel: experts across GPUs</a></li>
<li><a href="#113-ep-all-to-all-code-skeleton-pseudo-code">11.3 EP all-to-all code skeleton (pseudo-code)</a></li>
</ul>
</li>
<li><a href="#12-activation-memory-optimization">§12 Activation memory optimization</a>
<ul>
<li><a href="#121-gradient-checkpointing-chen-et-al-arxiv160406174-2016">12.1 Gradient Checkpointing (Chen et al., arXiv:1604.06174, 2016)</a></li>
<li><a href="#122-selective-recompute-korthikanti-2023">12.2 Selective recompute (Korthikanti 2023)</a></li>
<li><a href="#123-offload-zero-infinity">12.3 Offload (ZeRO-Infinity)</a></li>
<li><a href="#124-activation-memory-formula-must-memorize">12.4 Activation memory formula (must memorize)</a></li>
</ul>
</li>
<li><a href="#13-synthesis-3d--4d--5d-parallelism">§13 Synthesis: 3D / 4D / 5D Parallelism</a>
<ul>
<li><a href="#131-llama-3-405b-training-topology-public">13.1 Llama 3 405B training topology (public)</a></li>
<li><a href="#132-deepseek-v3-training-topology-public">13.2 DeepSeek-V3 training topology (public)</a></li>
</ul>
</li>
<li><a href="#14-dualpipe--the-2024-pipeline-frontier">§14 DualPipe — the 2024 pipeline frontier</a>
<ul>
<li><a href="#141-core-idea">14.1 Core idea</a></li>
<li><a href="#142-dualpipe-vs-vanilla-1f1b-properties">14.2 DualPipe vs vanilla 1F1B properties</a></li>
<li><a href="#143-dualpipe-costs">14.3 DualPipe costs</a></li>
</ul>
</li>
<li><a href="#15-torchtitan--pytorch-native-4d-platform">§15 TorchTitan — PyTorch-native 4D platform</a>
<ul>
<li><a href="#151-design-goals">15.1 Design goals</a></li>
<li><a href="#152-code-style-vs-deepspeed-monkey-patch">15.2 Code style (vs DeepSpeed monkey patch)</a></li>
<li><a href="#153-float8-training-hopper--blackwell">15.3 Float8 training (Hopper / Blackwell)</a></li>
<li><a href="#154-async-checkpointing">15.4 Async checkpointing</a></li>
</ul>
</li>
<li><a href="#16-communication-primitives-summary">§16 Communication primitives summary</a>
</li>
<li><a href="#17-25-frequently-asked-interview-questions">§17 25 frequently-asked interview questions</a>
<ul>
<li><a href="#l1-must-know-any-ml-engineer-role-will-ask">L1 must-know (any ML engineer role will ask)</a></li>
<li><a href="#l2-advanced-research-oriented-roles">L2 advanced (research-oriented roles)</a></li>
<li><a href="#l3-advanced-variants-top-labs--in-house-infra-roles">L3 advanced variants (top labs / in-house infra roles)</a></li>
</ul>
</li>
<li><a href="#a-appendix-complete-4d-wrap-code-skeleton">§A Appendix: complete 4D wrap code skeleton</a>
</li>
<li><a href="#b-references">§B References</a>
</li>
</ol>
</nav>
<main>
<header class="hero">
<h1>Distributed Training Tutorial En</h1>
<div class="meta">
<span><strong>Source:</strong> <code>docs/tutorials/distributed_training_tutorial_en.md</code></span>
<span><strong>SHA256:</strong> <code>9ee6b9a3603b</code></span>
<span><strong>Rendered:</strong> 2026-05-19 19:00 UTC</span>
</div>
</header>
<h2 id="0-tldr-cheat-sheet">§0 TL;DR Cheat Sheet</h2>
<div class="callout callout-info"><div class="callout-title">Distributed training in 7 dimensions, one page</div><p>DP (data split) × TP (tensor split) × PP (layer split) × SP (sequence split) × CP (context split) × EP (expert split) × activation recompute. See §1–§11 for derivations.</p></div>
<ol><li><strong>DDP</strong> (PyTorch): every GPU holds a full model copy; gradients are synced via NCCL <strong>all-reduce</strong> during backward. Bucket fusion + computation/communication overlap is the key engineering optimization.</li><li><strong>ZeRO 1/2/3</strong> (Rajbhandari et al., SC 2020): shard <strong>optimizer state / gradient / parameter</strong> respectively across $N$ GPUs, reducing per-GPU memory from $\Phi (2+2+12) = 16\Phi$ bytes to $16\Phi/N$ (fp16 training, Adam, parameter count $\Phi$).</li><li><strong>FSDP / FSDP2</strong> (Zhao et al., VLDB 2023; PyTorch 2.4+, 2024): PyTorch-native ZeRO-3. FSDP2 replaces FSDP1's flat-parameter mode with <strong>per-parameter DTensor sharding</strong>, composing more naturally with TP / PP.</li><li><strong>TP</strong> (Megatron-LM, Shoeybi 2019; Narayanan SC 2021): <strong>column-parallel → row-parallel</strong> pairing, with each layer requiring only 2× all-reduce per forward; attention is split by head, FFN uses col-parallel first layer + row-parallel second layer.</li><li><strong>PP</strong>: GPipe (Huang NeurIPS 2019) → 1F1B / PipeDream (Narayanan SOSP 2019) → Megatron-LM <strong>interleaved 1F1B</strong> (Narayanan SC 2021). Bubble ratio $\approx (P-1)/M$ ($P$ stages, $M$ micro-batches), with interleaved compressing it by a further factor of $V$.</li><li><strong>SP / CP / EP</strong>: SP (Korthikanti et al., MLSys 2023) shards LayerNorm/Dropout activations outside TP along the sequence; CP (Megatron-LM 2024 / Ring Attention Liu 2024) shards long context; EP partitions MoE experts across GPUs with <strong>all-to-all</strong> routing in forward.</li><li><strong>2024 frontier</strong>: <strong>DualPipe</strong> (DeepSeek-V3, Dec 2024) bidirectional pipeline fully overlaps forward/backward computation with communication; <strong>Llama 3 405B</strong> (Meta 2024) uses 16K H100s, 3.8e25 FLOPs, 54 days with 466 interruptions (419 unexpected + 47 planned maintenance); <strong>TorchTitan</strong> (Liang et al., ICLR 2025) integrates FSDP2 + TP + PP + SP + Float8.</li></ol>
<h2 id="1-intuition--why-one-gpu-cannot-hold-a-large-model">§1 Intuition — why one GPU cannot hold a large model</h2>
<p><strong>A softball question that few people fully answer</strong>: when training a model, what exactly is stored in a single GPU's memory?</p>
<p>Consider fp16 mixed-precision training + Adam optimizer with parameter count $\Phi$:</p>
<table><thead><tr><th>Item</th><th>Per-GPU usage</th><th>Notes</th></tr></thead><tbody><tr><td>Parameter (fp16)</td><td>$2\Phi$</td><td>used in forward / backward</td></tr><tr><td>Gradient (fp16)</td><td>$2\Phi$</td><td>accumulated in backward</td></tr><tr><td>Optimizer state</td><td>$12\Phi$</td><td>fp32 master copy ($4\Phi$) + Adam $m$, $v$ ($4\Phi + 4\Phi$)</td></tr><tr><td><strong>Subtotal (model states)</strong></td><td><strong>$16\Phi$</strong></td><td>independent of batch size</td></tr><tr><td>Activation</td><td>$O(B \cdot L \cdot D \cdot \text{depth})$</td><td>linear in batch size, seq len, and depth</td></tr><tr><td>Workspace / temp buffer</td><td>varies</td><td>NCCL / cuDNN workspace</td></tr></tbody></table>
<p>A 7B model has $\approx 112$ GB of model states alone — too large even for an 80GB H100. So distributed training is <strong>first and foremost about sharding model states + activations</strong>.</p>
<div class="callout callout-info"><div class="callout-title">Three orthogonal sharding dimensions</div><p>training tasks can be sharded along these three axes, theoretically independently and in practice combined:</p></div>
<ul><li><strong>Data dimension</strong> (DP / ZeRO / FSDP): shard the batch; replicas exchange gradients</li><li><strong>Model dimension (depth)</strong> (PP): shard layers; stages pass activations</li><li><strong>Model dimension (width)</strong> (TP / SP / CP / EP): shard tensors within a single layer; communicate intermediate results each step</li></ul>
<p>3D / 4D / 5D parallelism is the Cartesian product of these axes. Llama 3 / DeepSeek-V3 both use 4D (DP × TP × PP × a subset of CP/SP/EP).</p>
<h2 id="2-nccl-communication-primitives-and-topology">§2 NCCL communication primitives and topology</h2>
<p>99% of distributed-training communication goes through NCCL; you must be familiar with five primitives' semantics and traffic.</p>
<h3 id="21-the-five-collective-primitives">2.1 The five collective primitives</h3>
<p>Suppose there are $N$ GPUs, each holding a buffer of size $S$.</p>
<table><thead><tr><th>Primitive</th><th>Input → output</th><th>Equivalent semantics</th><th>Ring-algorithm traffic / GPU</th></tr></thead><tbody><tr><td><strong>all-reduce</strong></td><td>$N$ × $S$ → $N$ identical $S$</td><td>sum + broadcast</td><td>$2(N-1)/N \cdot S \approx 2S$</td></tr><tr><td><strong>reduce-scatter</strong></td><td>$N$ × $S$ → $N$ different $S/N$</td><td>sum then slice</td><td>$(N-1)/N \cdot S \approx S$</td></tr><tr><td><strong>all-gather</strong></td><td>$N$ × $S/N$ → $N$ × $S$</td><td>concatenate shards across ranks</td><td>$(N-1)/N \cdot S \approx S$</td></tr><tr><td><strong>broadcast</strong></td><td>rank0's $S$ → $N$ × $S$</td><td>single source</td><td>$S$ (tree)</td></tr><tr><td><strong>all-to-all</strong></td><td>$N \times N$ block transpose</td><td>shuffle</td><td>$(N-1)/N \cdot S \approx S$</td></tr></tbody></table>
<p><strong>Key identity</strong>:</p>
<p>$$\boxed{\;\text{all-reduce} = \text{reduce-scatter} + \text{all-gather}\;}$$</p>
<p>Each ring step transmits $S/N$, total $2(N-1)$ steps → per-GPU total $2(N-1)S/N \approx 2S$ bytes (almost independent of $N$ — this is the elegance of the ring algorithm).</p>
<div class="callout callout-warn"><div class="callout-title">Interview bonus: NCCL is not a single algorithm</div><p>NCCL uses <strong>tree all-reduce</strong> for small messages (latency-bound, $O(\log N)$ hops); <strong>ring all-reduce</strong> for large messages (bandwidth-bound, throughput-optimal). On NVLink topologies it also offers <strong>NVLS (NVLink SHARP)</strong> — hardware-side reduction (H100/H200 NVSwitch supports this).</p></div>
<h3 id="22-nvlink--ib--topology">2.2 NVLink / IB / topology</h3>
<table><thead><tr><th>Link</th><th>Unidirectional bandwidth (H100-generation)</th><th>Use</th></tr></thead><tbody><tr><td>NVLink 4.0</td><td>900 GB/s (per GPU, 18 lanes aggregated)</td><td>intra-node GPU↔GPU</td></tr><tr><td>PCIe 5.0 x16</td><td>64 GB/s</td><td>GPU↔CPU, slow path</td></tr><tr><td>InfiniBand NDR 400G</td><td>50 GB/s (per port)</td><td>inter-node</td></tr></tbody></table>
<p><strong>Rule of thumb</strong>: intra-node communication is <strong>10-20×</strong> faster than inter-node. So TP must fit inside a node; DP / PP can cross nodes. Llama 3 training topology: TP=8 (intra-node NVLink) × PP=16 (inter-node IB) × DP=128.</p>
<h3 id="23-nccl-calls-in-pytorch">2.3 NCCL calls in PyTorch</h3>
<pre><code class="language-python">import torch
import torch.distributed as dist
dist.init_process_group(backend=&quot;nccl&quot;) # backend fixed to nccl
rank = dist.get_rank()
world_size = dist.get_world_size()
# all-reduce (default SUM; can also be AVG / MIN / MAX / PRODUCT)
buf = torch.ones(1024, device=f&quot;cuda:{rank % 8}&quot;) * rank
dist.all_reduce(buf, op=dist.ReduceOp.SUM)
# now buf == sum(0..world_size-1) * ones(1024)
device = torch.device(f&quot;cuda:{rank % 8}&quot;)
# reduce-scatter
input_list = [torch.full((1024,), float(rank + i), device=device) for i in range(world_size)]
output = torch.empty(1024, device=device)
dist.reduce_scatter(output, input_list, op=dist.ReduceOp.SUM)
# all-gather
local = torch.full((1024,), float(rank), device=device)
gathered = [torch.empty(1024, device=device) for _ in range(world_size)]
dist.all_gather(gathered, local)
# all-to-all (essential for MoE)
in_split = list(torch.randn(world_size, 1024, device=device).unbind(0))
out_split = [torch.empty(1024, device=device) for _ in range(world_size)]
dist.all_to_all(out_split, in_split)</code></pre>
<h2 id="3-ddp--distributeddataparallel">§3 DDP — DistributedDataParallel</h2>
<h3 id="31-algorithm-skeleton">3.1 Algorithm skeleton</h3>
<p>DDP is the simplest data-parallel implementation:</p>
<ol><li><strong>Replicate</strong>: every rank holds a complete model copy (parameters / gradients / optimizer states are all full)</li><li><strong>Shard batch</strong>: the global batch $B$ is split into $N$ micro-batches; each rank computes its forward + backward on its share</li><li><strong>Sync gradients</strong>: after backward, perform <strong>all-reduce</strong> on all gradients (SUM then divide by $N$ for averaging, equivalent to AVG)</li><li><strong>Local optimizer step</strong>: each rank runs the same optimizer on the same gradients; parameters stay consistent</li></ol>
<p>Mathematically (loss $\mathcal{L}$ as a mean over the mini-batch):</p>
<p>$$g_\text{global} = \frac{1}{N}\sum_{i=1}^N \nabla_\theta \mathcal{L}(\theta; \mathcal{B}_i) = \text{all-reduce-mean}(g_1, \dots, g_N)$$</p>
<p>Each rank gets the same $g_\text{global}$, so the $\theta$'s on $N$ ranks always stay synchronized (same init + same gradient + same optimizer).</p>
<h3 id="32-bucket-fusion--overlap-ddp-engineering-essence">3.2 Bucket fusion + overlap (DDP engineering essence)</h3>
<p>Naive implementation: after backward completes → concatenate all gradient tensors → one all-reduce. This <strong>idles the GPU waiting for communication</strong>, an enormous waste.</p>
<p>What PyTorch DDP actually does:</p>
<ul><li><strong>Bucket</strong>: pack multiple gradient tensors into a <strong>bucket</strong> (default 25 MB) in reverse-computation order</li><li><strong>Hook</strong>: trigger a hook when each parameter's gradient is computed</li><li><strong>Overlap</strong>: when a bucket fills (all its grads are computed), <strong>immediately launch an all-reduce on a background stream</strong>, while the main stream continues backward on earlier layers</li><li><strong>Result</strong>: communication is fully or partially hidden by the backward of earlier layers</li></ul>
<pre class="diagram"><code>Backward time axis →
Layer N: [grad N]──┐
Layer N-1: [grad N-1]┼─bucket_N─[all-reduce N]
Layer N-2: [grad N-2]┘ ↓ (background stream)
Layer N-3: [grad N-3]──┐
Layer N-4: [grad N-4]──┼─bucket_N-1─[all-reduce N-1]
... ↓
Layer 1: [grad 1]──────────────[all-reduce 1]
Main stream (compute): ████████████████████████████████
Background (NCCL): ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓
↑ overlap with compute</code></pre>
<div class="callout callout-warn"><div class="callout-title">Bucket too small / too large are both bad</div><p>too small: communication launch overhead dominates, bandwidth utilization is low; too large: filling takes too long, the first all-reduce starts late. PyTorch's default 25 MB is the common sweet spot for large models. Tune via <code>DDP(bucket_cap_mb=...)</code>.</p></div>
<h3 id="33-pytorch-code-with-overlap">3.3 PyTorch code (with overlap)</h3>
<pre><code class="language-python">import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
def setup_ddp(rank, world_size, local_rank):
dist.init_process_group(&quot;nccl&quot;, rank=rank, world_size=world_size)
torch.cuda.set_device(local_rank) # multi-node: local_rank != global rank
def main(rank, world_size, local_rank):
setup_ddp(rank, world_size, local_rank)
device = torch.device(f&quot;cuda:{local_rank}&quot;)
model = MyModel().to(device)
model = DDP(
model,
device_ids=[local_rank],
bucket_cap_mb=25, # bucket size
gradient_as_bucket_view=True, # memory opt: grad is a bucket view
static_graph=False, # if graph is static, enables more fusion
)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
loader = make_distributed_loader(rank, world_size)
for batch in loader:
x, y = batch[0].to(device), batch[1].to(device)
loss = nn.functional.cross_entropy(model(x), y)
optimizer.zero_grad(set_to_none=True)
loss.backward() # ← bucket hooks trigger all-reduce here
optimizer.step()</code></pre>
<h3 id="34-complexity">3.4 Complexity</h3>
<div class="callout callout-info"><div class="callout-title">Notation convention (used throughout)</div><p>$\Phi$ denotes <strong>parameter count</strong> (number). In fp16 / bf16 training, one parameter is 2 bytes, so the fp16 weight buffer is $2\Phi$ bytes. In the table below, "$2\Phi$ bytes" means byte count, "$\Phi$ params" means count.</p></div>
<table><thead><tr><th>Item</th><th>Quantity</th><th>Note</th></tr></thead><tbody><tr><td>Per-GPU memory</td><td>$16\Phi$ bytes + activation</td><td>model states not sharded (fp16 params/grads + fp32 master/Adam)</td></tr><tr><td>Per-step communication</td><td>$\approx 2 \cdot 2\Phi = 4\Phi$ bytes (ring all-reduce on fp16 gradient)</td><td>equivalent to reduce-scatter + all-gather</td></tr><tr><td>Scalability</td><td>compute scales linearly; communication bandwidth $O(1)$ (ring independent of $N$)</td><td>but latency-bound for small models, still affected by $N$</td></tr></tbody></table>
<p><strong>DDP's fatal flaw</strong>: model states are not sharded. With Adam training a 70B model, the per-GPU model states alone are 1.12 TB — DDP cannot fit. So 7B+ models must use ZeRO / FSDP.</p>
<h2 id="4-zero-123--zero-redundancy-optimizer">§4 ZeRO 1/2/3 — Zero Redundancy Optimizer</h2>
<p>Rajbhandari et al. <strong>"ZeRO: Memory Optimizations Toward Training Trillion Parameter Models"</strong> (SC 2020) is the second epochal work in distributed training (the first being Megatron-TP). Core idea: DDP wastes memory by replicating model states $N$ times — shard them into $N$ pieces so each GPU only holds $1/N$, with only slightly more communication.</p>
<h3 id="41-three-stage-memory-math">4.1 Three-stage memory math</h3>
<p>Let $\Phi$ be the parameter count and $N$ the DP world size. With fp16 mixed precision + Adam, per-GPU model states:</p>
<p>$$\text{DDP}: \quad 2\Phi + 2\Phi + 12\Phi = 16\Phi$$</p>
<p>The three ZeRO stages differ in <strong>"what is sharded"</strong>:</p>
<table><thead><tr><th>Stage</th><th>Sharded content</th><th>Per-GPU model states</th><th>Communication (per step)</th></tr></thead><tbody><tr><td><strong>ZeRO-1</strong></td><td>optimizer state</td><td>$2\Phi + 2\Phi + 12\Phi/N$</td><td>$2\Phi$ (all-reduce = reduce-scatter + all-gather; ZeRO-1 only reduce-scatters grad)</td></tr><tr><td><strong>ZeRO-2</strong></td><td>optimizer state + gradient</td><td>$2\Phi + 2\Phi/N + 12\Phi/N$</td><td>$2\Phi$ (same as ZeRO-1)</td></tr><tr><td><strong>ZeRO-3</strong></td><td>opt state + grad + <strong>parameter</strong></td><td>$2\Phi/N + 2\Phi/N + 12\Phi/N = 16\Phi/N$</td><td>$3\Phi$ (forward + backward each do one all-gather, backward does one reduce-scatter)</td></tr></tbody></table>
<div class="callout callout-good"><div class="callout-title">In the limit</div><p>when $N$ is large enough (e.g. 1024 H100s), ZeRO-3 per-GPU model states $16\Phi / 1024 = 0.0156\Phi$. For a 65B model that's $\approx 1$ GB / GPU; with activation checkpointing a single H100 can fully fit it.</p></div>
<h3 id="42-zero-3-workflow-most-commonly-used-forward--backward--optimize">4.2 ZeRO-3 workflow (most commonly used, forward / backward / optimize)</h3>
<p>ZeRO-3 shards parameters themselves, so parameters must be <strong>all-gathered</strong> before use.</p>
<p><strong>Forward</strong> (per layer / module):</p>
<pre class="diagram"><code>1. all-gather: collect full W^(ℓ) from N GPUs ──[comm: φ_ℓ bytes]
2. compute: y = f(x; W^(ℓ)) ──[compute]
3. release: drop the shards that don&#x27;t belong locally ──[free memory]</code></pre>
<p><strong>Backward</strong>:</p>
<pre class="diagram"><code>1. all-gather: re-fetch W^(ℓ) (forward released it) ──[comm: φ_ℓ]
2. compute: grad_W^(ℓ), grad_x
3. reduce-scatter: reduce grad_W^(ℓ) and slice back to per-shard ──[comm: φ_ℓ]
4. release: drop non-local shards</code></pre>
<p>Pseudo-code:</p>
<pre><code class="language-python"># ZeRO-3 forward (single-layer abstraction, simplified)
def zero3_forward(layer_idx, x, sharded_W):
# 1. all-gather full weight
full_W = all_gather(sharded_W) # [Φ_ℓ / N, ...] × N -&gt; [Φ_ℓ, ...]
# 2. compute
y = layer_forward(x, full_W)
# 3. release full_W (keep only sharded_W)
del full_W
return y
def zero3_backward(layer_idx, dy, sharded_W, cached_input):
# 1. all-gather (already released after forward)
full_W = all_gather(sharded_W)
# 2. compute local gradients
dW_local, dx = layer_backward(dy, full_W, cached_input)
del full_W
# 3. reduce-scatter gradient to the corresponding shard
dW_sharded = reduce_scatter(dW_local) # [Φ_ℓ, ...] / N -&gt; [Φ_ℓ/N, ...]
return dW_sharded, dx</code></pre>
<h3 id="43-zero-123-vs-ddp-communication-comparison-important-interview-question">4.3 ZeRO-1/2/3 vs DDP communication comparison (important interview question)</h3>
<p>Total parameter count $\Phi$ (count). Below, the unit is <strong>"fp16 weight buffer equivalents"</strong> (i.e. $\Phi$ in the traffic column represents $2\Phi$ bytes of fp16-buffer traffic). Per-step <strong>per-GPU</strong> communication (ring assumption) for one forward+backward+update:</p>
<table><thead><tr><th>Mode</th><th>Forward</th><th>Backward</th><th>Optim</th><th>Total (fp16 buffer equiv.)</th></tr></thead><tbody><tr><td>DDP</td><td>0</td><td>$2\Phi$ (all-reduce grad)</td><td>0</td><td>$2\Phi$ (i.e. $4\Phi$ bytes)</td></tr><tr><td>ZeRO-1</td><td>0</td><td>$\Phi$ (reduce-scatter grad)</td><td>$\Phi$ (all-gather updated params)</td><td>$2\Phi$</td></tr><tr><td>ZeRO-2</td><td>0</td><td>$\Phi$ (reduce-scatter grad)</td><td>$\Phi$ (all-gather)</td><td>$2\Phi$</td></tr><tr><td>ZeRO-3</td><td>$\Phi$ (all-gather params, on-the-fly)</td><td>$\Phi$ (all-gather) + $\Phi$ (reduce-scatter grad)</td><td>0 (already in backward)</td><td>$3\Phi$</td></tr></tbody></table>
<div class="callout callout-info"><div class="callout-title">Key conclusion</div><p>ZeRO-1/2 communication is the same as DDP ($2\Phi$ buffer) but memory drops sharply; ZeRO-3 has $1.5\times$ communication for $N\times$ memory reduction. In practice ZeRO-3's communication can also be partially hidden by <strong>prefetch + overlap</strong>, making it the mainstream choice for 70B+ models. In bytes: DDP $\approx 4\Phi$ bytes, ZeRO-3 $\approx 6\Phi$ bytes.</p></div>
<h3 id="44-zero-offload--zero-infinity">4.4 ZeRO-Offload / ZeRO-Infinity</h3>
<p><strong>ZeRO-Offload</strong> (Ren et al., USENIX ATC 2021): offload optimizer state + part of gradient to <strong>CPU</strong>; CPU runs the Adam update. Cost: CPU↔GPU PCIe traffic + slow CPU compute. Best for <strong>small clusters + large models</strong> (e.g. 13B on a single 8-GPU node).</p>
<p><strong>ZeRO-Infinity</strong> (Rajbhandari et al., SC 2021): further offload parameters / optimizer to <strong>NVMe</strong>. Theoretically lets a single machine train 1T parameters (in practice throughput is extremely low; mostly for inference / fine-tuning).</p>
<div class="callout callout-warn"><div class="callout-title">Offload is a trade-off</div><p>CPU offload typically increases per-step time 1.5-3×; NVMe offload is 5-10× slower. Use only when "doesn't fit + can't afford more GPUs". Production training prefers scaling out GPU count.</p></div>
<h2 id="5-fsdp--fsdp2--pytorch-native-zero-3">§5 FSDP / FSDP2 — PyTorch native ZeRO-3</h2>
<p>Zhao et al. <strong>"PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel"</strong> (VLDB 2023) brings the ZeRO-3 idea into the PyTorch mainline. FSDP2 (PyTorch 2.4+, 2024) is a major rewrite: <strong>per-parameter DTensor sharding</strong> replaces FSDP1's flat-parameter mode.</p>
<h3 id="51-fsdp1-vs-fsdp2-core-differences">5.1 FSDP1 vs FSDP2 core differences</h3>
<table><thead><tr><th>Dimension</th><th>FSDP1 (2022-2023)</th><th>FSDP2 (2024+)</th></tr></thead><tbody><tr><td>Data structure</td><td><strong>FlatParameter</strong>: concatenate all params in the wrap unit into one 1D buffer, then chunk</td><td><strong>DTensor per-parameter</strong>: each param sharded independently along dim-0</td></tr><tr><td>State dict</td><td>needs all-gather to produce unflattened state dict</td><td>sharded state dict, communication-free</td></tr><tr><td>Frozen parameters</td><td>within a unit, all must be frozen or all trainable</td><td>each param independent, mixed frozen/trainable is natural</td></tr><tr><td>TP composition</td><td>difficult (flat-buffer conflicts with TP's different shard axis)</td><td><strong>natively compatible</strong>: DTensor describes multi-axis placement (<code>Shard(0)</code>, <code>Replicate</code>, <code>Shard(1)</code> combos)</td></tr><tr><td>API</td><td><code>FullyShardedDataParallel</code></td><td><code>fully_shard()</code> functional wrap</td></tr></tbody></table>
<h3 id="52-fsdp-wrap-policy-the-most-important-design-decision">5.2 FSDP wrap policy (the most important design decision)</h3>
<p>FSDP does not treat the whole model as one unit. It shards by <strong>custom unit boundaries</strong>; params in a unit all-gather / reduce-scatter together.</p>
<pre><code class="language-python">import torch
import torch.nn as nn
from torch.distributed.fsdp import fully_shard # FSDP2 API
from torch.distributed.fsdp import MixedPrecisionPolicy
class TransformerBlock(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.attn = MultiHeadAttention(d_model, n_heads)
self.mlp = FeedForward(d_model)
self.ln1 = nn.LayerNorm(d_model)
self.ln2 = nn.LayerNorm(d_model)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
# Treat each TransformerBlock as one FSDP unit
model = MyLLM(n_layers=32, d_model=4096, n_heads=32).cuda()
mp_policy = MixedPrecisionPolicy(
param_dtype=torch.bfloat16, # forward uses bf16
reduce_dtype=torch.float32, # gradient reduce uses fp32 to avoid error accumulation
)
for block in model.blocks:
fully_shard(block, mp_policy=mp_policy)
fully_shard(model, mp_policy=mp_policy) # root also needs wrapping</code></pre>
<div class="callout callout-warn"><div class="callout-title">Engineering trade-off of wrap granularity</div><p>smaller unit (e.g. wrap each linear): each all-gather only fetches one layer, peak memory low, but many communication calls and prefetch is hard; larger unit (e.g. a whole block or several blocks): fewer comms, easy to overlap, but higher peak memory. <strong>TransformerBlock granularity is the standard for LLaMA / GPT-class models</strong>.</p></div>
<h3 id="53-fsdp2--mixed-precision--activation-checkpoint">5.3 FSDP2 + mixed precision + activation checkpoint</h3>
<pre><code class="language-python">from torch.distributed.fsdp import CPUOffloadPolicy, OffloadPolicy
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
checkpoint_wrapper, CheckpointImpl
)
# 1. Activation checkpoint (gradient checkpointing) — recompute trades memory
def apply_ac(model):
for i, block in enumerate(model.blocks):
model.blocks[i] = checkpoint_wrapper(
block,
checkpoint_impl=CheckpointImpl.NO_REENTRANT,
)
apply_ac(model)
# 2. FSDP2 wrap
mp_policy = MixedPrecisionPolicy(
param_dtype=torch.bfloat16,
reduce_dtype=torch.float32,
)
offload_policy = OffloadPolicy() # or CPUOffloadPolicy() to offload to CPU
for block in model.blocks:
fully_shard(block, mp_policy=mp_policy, offload_policy=offload_policy)
fully_shard(model, mp_policy=mp_policy)
# 3. Training (note: optimizer must be torch._foreach_ friendly)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=True)
for batch in loader:
loss = compute_loss(model, batch)
loss.backward()
optimizer.step()
optimizer.zero_grad(set_to_none=True)</code></pre>
<h3 id="54-fsdp-vs-zero-3-communication-difference-l3-question">5.4 FSDP vs ZeRO-3 communication difference (L3 question)</h3>
<p><strong>Conclusion</strong>: <strong>completely the same</strong> (at the wrap-unit = layer granularity). FSDP is the PyTorch-native implementation of ZeRO-3; the algorithms are equivalent. The differences are engineering:</p>
<table><thead><tr><th>Dimension</th><th>DeepSpeed ZeRO-3</th><th>FSDP2</th></tr></thead><tbody><tr><td>Integration</td><td>external library + monkey patch</td><td>PyTorch mainline</td></tr><tr><td>state_dict</td><td>needs a separate API to stitch</td><td>native <code>set_model_state_dict</code> / DCP</td></tr><tr><td>TP composition</td><td>DeepSpeed-Megatron bridge</td><td>DTensor native</td></tr><tr><td>Freezing / LoRA</td><td>complex (flat-param conflicts)</td><td>natural</td></tr><tr><td>Compilation</td><td>restricted</td><td><code>torch.compile(fullgraph=False)</code> friendly</td></tr></tbody></table>
<p>Practical choice: <strong>new projects use FSDP2; DeepSpeed is still used heavily in 70B+ MoE / multi-framework scenarios</strong>.</p>
<h3 id="55-hsdp--hybrid-sharded-dp">5.5 HSDP — Hybrid Sharded DP</h3>
<p><strong>Problem</strong>: 1024 GPUs all on ZeRO-3 means a single all-gather across 1024 GPUs → per-GPU traffic $(N-1)/N \cdot \phi \to \phi$ stays the same, but <strong>latency rises sharply</strong> (IB cross-node is expensive).</p>
<p><strong>HSDP</strong>: ZeRO-3 within a group (e.g. 8 GPUs in a node), DDP across groups. Equivalent to dividing 1024 GPUs into 128 ZeRO-3 groups (8 each).</p>
<table><thead><tr><th>Mode</th><th>Shard scope</th><th>Per-GPU model states</th><th>Inter-node communication</th></tr></thead><tbody><tr><td>Pure ZeRO-3</td><td>full 1024</td><td>$16\Phi / 1024$</td><td>many</td></tr><tr><td>HSDP (8)</td><td>8 GPUs per group</td><td>$16\Phi / 8 = 2\Phi$</td><td>few (inter-group is grad all-reduce)</td></tr></tbody></table>
<p>Trade-off: <strong>memory traded for communication efficiency</strong>. Llama 3 uses HSDP at some stages of training.</p>
<h2 id="6-zero--communication-optimizations-2023">§6 ZeRO++ — communication optimizations (2023)</h2>
<p>Wang et al. <strong>"ZeRO++: Extremely Efficient Collective Communication for Giant Model Training"</strong> (NeurIPS MLSys workshop 2023, arXiv 2306.10209). Three additions on top of ZeRO-3:</p>
<h3 id="61-qwz-quantized-weight-all-gather-quantize-forward-comms">6.1 qwZ: Quantized Weight all-gather (quantize forward comms)</h3>
<p>ZeRO-3 forward all-gathers fp16 weight per layer. <code>qwZ</code> converts fp16 → int8 before all-gather, <strong>halving traffic</strong>.</p>
<pre class="diagram"><code>forward (vanilla): weight_shard (fp16) ──all-gather──&gt; full_weight (fp16) ──compute
forward (qwZ): weight_shard (fp16) ──quant──&gt; shard (int8)
──all-gather──&gt; full (int8)
──dequant──&gt; full (fp16) ──compute</code></pre>
<p>Cost: do block-wise quantization / dequantization before and after each all-gather (typical block size 2048-4096 elements with per-block scale). Block-quant is much more precise than per-tensor quant, with &lt; 1% effect on training loss.</p>
<h3 id="62-hpz-hierarchical-partition">6.2 hpZ: Hierarchical Partition</h3>
<p>Observation: <strong>all-gather during backward is more expensive than during forward</strong> (deeper layers backprop first). <code>hpZ</code> <strong>replicates weight within a node</strong> (all NVLink), and shards across nodes:</p>
<ul><li>Intra-node: 8 GPUs each hold the full weight (hpZ = intra-node DDP)</li><li>Inter-node: weight is sharded (still ZeRO-3 mode)</li></ul>
<p>Backward's weight all-gather only goes intra-node (NVLink, cheap), not across IB. Cost: per-node memory grows 8× — but model states are already sharded across 1024 GPUs, so 8× is still much less than DDP.</p>
<h3 id="63-qgz-quantized-gradient-reduce">6.3 qgZ: Quantized Gradient reduce</h3>
<p>Backward's tail reduce-scatter on gradients also uses int8 (fp16 → int8 + quantized reduce). The vanilla reduce-scatter SUM cannot be naïvely quantized (quant + sum accumulates error); ZeRO++ instead uses <strong>all-to-all + dequant + local sum</strong>.</p>
<h3 id="64-combined-effect">6.4 Combined effect</h3>
<p>ZeRO++ paper reports <strong>2.16× throughput improvement at 384 GPU scale</strong>, with 4× lower traffic (fp16→int8 saves 2× × 2 collective primitives). Cost: complex implementation, precision requires careful ablation.</p>
<div class="callout callout-warn"><div class="callout-title">Quantized communication ≠ quantized training</div><p>Here we only quantize <strong>transient buffers in the collective communication path</strong>; weight storage and compute remain fp16 / bf16. So loss impact is usually negligible, and this is different from fp8 training (e.g. fp8 GEMM on Hopper).</p></div>
<h2 id="7-tensor-parallel--megatron-lm">§7 Tensor Parallel — Megatron-LM</h2>
<p>Shoeybi et al. <strong>"Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism"</strong> (arXiv 2019) + Narayanan et al. <strong>"Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM"</strong> (SC 2021).</p>
<h3 id="71-core-idea-column-parallel--row-parallel-pairing">7.1 Core idea: column-parallel → row-parallel pairing</h3>
<p>Consider an MLP layer: $Y = \text{GELU}(X W_1) W_2$, with $X \in \mathbb{R}^{B \times D}$, $W_1 \in \mathbb{R}^{D \times 4D}$, $W_2 \in \mathbb{R}^{4D \times D}$.</p>
<p><strong>Column Parallel</strong> (shard $W_1$'s output dim $4D$):</p>
<p>$$W_1 = [W_1^{(1)} \mid W_1^{(2)} \mid \cdots \mid W_1^{(T)}], \quad W_1^{(i)} \in \mathbb{R}^{D \times 4D/T}$$</p>
<p>Each rank $i$ holds $W_1^{(i)}$ and independently computes $Y_1^{(i)} = X W_1^{(i)} \in \mathbb{R}^{B \times 4D/T}$. <strong>Input $X$ is fully replicated; output is sharded along columns</strong>.</p>
<p><strong>Row Parallel</strong> (shard $W_2$'s input dim $4D$):</p>
<p>$$W_2 = \begin{bmatrix} W_2^{(1)} \\ W_2^{(2)} \\ \vdots \\ W_2^{(T)} \end{bmatrix}, \quad W_2^{(i)} \in \mathbb{R}^{4D/T \times D}$$</p>
<p>Input is sharded along rows: $Z^{(i)} = \text{GELU}(Y_1^{(i)}) W_2^{(i)} \in \mathbb{R}^{B \times D}$. Each rank computes its partial sum, then <strong>all-reduce</strong> to sum them:</p>
<p>$$Y = \sum_{i=1}^T Z^{(i)} = \text{all-reduce}(Z^{(1)}, \dots, Z^{(T)})$$</p>
<h3 id="72-communication-analysis">7.2 Communication analysis</h3>
<p>Column parallel: input $X$ is fully replicated, output sharded along columns → <strong>forward has no communication</strong> (output is distributed across ranks).</p>
<p>Row parallel: input is row-sharded, output requires all-reduce → <strong>forward: 1× all-reduce $(BD)$</strong>.</p>
<p><strong>Col + row pairing</strong>: the column-parallel tail has no communication (its result feeds directly into row-parallel), and the row-parallel tail does 1× all-reduce. The full MLP block forward involves <strong>only 1× all-reduce</strong>.</p>
<p>Backward mirrors: each MLP block backward also performs 1× all-reduce (gradient flow through col-row mirrors the communication).</p>
<p><strong>Total</strong>: a single transformer block (MLP + Attention, where attention is also col-row) forward + backward = <strong>4× all-reduce, each of size $BLD$</strong> ($L$ = seq len).</p>
<h3 id="73-tp-sharding-of-attention">7.3 TP sharding of attention</h3>
<p>Multi-head attention has $W_Q, W_K, W_V \in \mathbb{R}^{D \times D}$. Shard <strong>directly by head dim</strong> (each head is independent, naturally column-parallel):</p>
<pre><code>H heads, T-way TP:
- Each rank holds W_Q^(i), W_K^(i), W_V^(i) for H/T heads
- Each rank computes its head_h(Q W_Q^h, ...) independently
- Output W_O uses row-parallel → tail all-reduce</code></pre>
<div class="callout callout-info"><div class="callout-title">Why shard attention by head rather than by dim</div><p>the head dim is naturally independent (heads don't interact), so head-sharding requires no cross-rank communication of intermediates; sharding hidden dim requires extra communication around softmax. Code-wise it's also simpler — heads map directly to ranks.</p></div>
<h3 id="74-tp-code-skeleton">7.4 TP code skeleton</h3>
<p>Below is a pedagogical col-parallel / row-parallel implementation, with explicit communication pairs in forward and backward (key: col-parallel forward has no comm, backward all-reduces the input grad; row-parallel forward all-reduces the output, backward has no comm):</p>
<pre><code class="language-python">import math
import torch
import torch.nn as nn
import torch.distributed as dist
class _CopyToTPRegion(torch.autograd.Function):
&quot;&quot;&quot; Identity in forward, all-reduce in backward (col-parallel entry) &quot;&quot;&quot;
@staticmethod
def forward(ctx, x, tp_group):
ctx.tp_group = tp_group
return x
@staticmethod
def backward(ctx, grad_out):
dist.all_reduce(grad_out, op=dist.ReduceOp.SUM, group=ctx.tp_group)
return grad_out, None
class _ReduceFromTPRegion(torch.autograd.Function):
&quot;&quot;&quot; All-reduce in forward, identity in backward (row-parallel exit) &quot;&quot;&quot;
@staticmethod
def forward(ctx, x, tp_group):
dist.all_reduce(x, op=dist.ReduceOp.SUM, group=tp_group)
return x
@staticmethod
def backward(ctx, grad_out):
return grad_out, None
class ColumnParallelLinear(nn.Module):
&quot;&quot;&quot; Y = X W, W ∈ R^{in×out}, shard out dim into T pieces;
forward: input replicated -&gt; output [B, out/T] sharded along last dim
backward: input grad must be all-reduced (each TP rank computed partial dX) &quot;&quot;&quot;
def __init__(self, in_features, out_features, tp_group, tp_size):
super().__init__()
assert out_features % tp_size == 0
self.tp_group = tp_group
self.out_per_rank = out_features // tp_size
self.weight = nn.Parameter(torch.empty(self.out_per_rank, in_features))
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
def forward(self, x):
# x: [B, in] (replicated). Entry _CopyToTPRegion lets backward auto all-reduce dX
x = _CopyToTPRegion.apply(x, self.tp_group)
return torch.nn.functional.linear(x, self.weight) # [B, out/T]
class RowParallelLinear(nn.Module):
&quot;&quot;&quot; Y = X W, W ∈ R^{in×out}, shard in dim into T pieces;
forward: input [B, in/T] (sharded), output partial sum -&gt; all-reduce
backward: dW/dX are computed locally, no comm needed &quot;&quot;&quot;
def __init__(self, in_features, out_features, tp_group, tp_size):
super().__init__()
assert in_features % tp_size == 0
self.tp_group = tp_group
self.in_per_rank = in_features // tp_size
self.weight = nn.Parameter(torch.empty(out_features, self.in_per_rank))
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
def forward(self, x):
# x: [B, in/T] (sharded along last dim)
local_out = torch.nn.functional.linear(x, self.weight) # [B, out] partial sum
return _ReduceFromTPRegion.apply(local_out, self.tp_group)
class TPTransformerMLP(nn.Module):
&quot;&quot;&quot; Megatron-LM standard col+row pairing &quot;&quot;&quot;
def __init__(self, d_model, tp_group, tp_size):
super().__init__()
d_ff = 4 * d_model
self.fc1 = ColumnParallelLinear(d_model, d_ff, tp_group, tp_size) # out [B, L, 4D/T]
self.fc2 = RowParallelLinear(d_ff, d_model, tp_group, tp_size) # in [B, L, 4D/T] -&gt; [B, L, D]
def forward(self, x):
# x: [B, L, D] (replicated)
h = torch.nn.functional.gelu(self.fc1(x)) # [B, L, 4D/T]
return self.fc2(h) # [B, L, D] after all-reduce</code></pre>
<div class="callout callout-info"><div class="callout-title">TP communication pairs (must memorize)</div><p>col-parallel: <code>forward = identity, backward = all-reduce(dX)</code>; row-parallel: <code>forward = all-reduce(Y), backward = identity</code>. A transformer block (col + row × 2 sets: MLP + attention) has 4 all-reduces total — 2 in forward, 2 in backward. This is the source of the 4× count in §7.2.</p></div>
<h3 id="75-tp-must-fit-inside-a-node-nvlink">7.5 TP must fit inside a node (NVLink)</h3>
<p>Each transformer block forward + backward = 4× all-reduce on activation tensors of size $\approx BLD$. For a 7B model with $D = 4096$ and $B \times L = 2048$ → each is $\approx 32$ MB; per-block × 4 × 32 layers = <strong>128 calls per step, $\approx 4$ GB per step</strong>.</p>
<p>At NVLink 900 GB/s that's $\approx 4.4$ ms; on IB at 50 GB/s it's $\approx 80$ ms. So <strong>TP must fit inside a node</strong> (NVLink domain). This is why LLaMA / Llama 3 / GPT-3 all use TP=8 (one 8-GPU node) rather than TP=16.</p>
<h2 id="8-sequence-parallel--saving-activation-memory">§8 Sequence Parallel — saving activation memory</h2>
<p>Korthikanti et al. <strong>"Reducing Activation Recomputation in Large Transformer Models"</strong> (MLSys 2023).</p>
<h3 id="81-motivation-activation-memory-is-mostly-element-wise-ops-like-layernorm--dropout">8.1 Motivation: activation memory is mostly element-wise ops like LayerNorm / Dropout</h3>
<p>TP shards the intermediate activations of attention / MLP (along head / hidden dim), but the <strong>portion outside TP</strong> (LayerNorm input/output, Dropout, residual) is still <strong>fully replicated</strong> — each TP rank stores a full $[B, L, D]$ activation.</p>
<p>For a 7B model with $B \times L \times D = 2048 \times 4096$, fp16 = 16 MB per activation. A layer has about 4-6 such full-replica activations. 32 layers × 5 → 2.5 GB / GPU wasted.</p>
<h3 id="82-sp-solution-shard-along-the-sequence-dim">8.2 SP solution: shard along the sequence dim</h3>
<p>Shard the non-TP activations along $L$ (sequence) across TP ranks too.</p>
<pre><code>TP-only: [B, L, D] [B, L, D]
full replica full replica
↓ ↑
LayerNorm → split → Attention → all-gather → LayerNorm
(TP)
TP + SP: [B, L/T, D] [B, L/T, D]
sharded along L sharded along L
↓ ↑
LayerNorm → all-gather → Attention → reduce-scatter → LayerNorm
(TP)</code></pre>
<p><strong>Key operation changes</strong>:</p>
<ul><li>TP-only: forward inside attention/MLP requires broadcast / no-op at entry, all-reduce at exit</li><li>TP + SP: at entry <strong>all-gather</strong> (assemble the SP-sharded $L$ back to full length), at exit <strong>reduce-scatter</strong> (reduce the full length back to $L$ shard)</li></ul>
<p>Communication volume: each transformer block under TP+SP needs 1× all-gather + 1× reduce-scatter (replacing pure-TP's 1× all-reduce); these have equal volume, so total is unchanged from pure TP. <strong>$\boxed{\text{Same total communication as pure TP}}$</strong> — but LayerNorm / Dropout activations drop from $BLD$ to $BLD/T$.</p>
<div class="callout callout-info"><div class="callout-title">Why SP has the same traffic as pure TP</div><p>because $\text{all-reduce} = \text{reduce-scatter} + \text{all-gather}$, pure TP's tail 1× all-reduce $= $ SP-mode 1× reduce-scatter + the following 1× all-gather (at next block's entry). <strong>Communication is redistributed; total is unchanged, but activation memory is saved</strong>.</p></div>
<h3 id="83-how-much-activation-memory-does-sp-save-l3-question">8.3 How much activation memory does SP save (L3 question)</h3>
<p>Let the <strong>non-sharded</strong> single transformer block total activation memory be $A = A_\text{in} + A_\text{out}$, where:</p>
<ul><li>$A_\text{in}$: the TP-<strong>internal</strong> part (attention's Q/K/V/scores, MLP's intermediate [B, L, 4D]) — TP shards it along the hidden dim</li><li>$A_\text{out}$: the TP-<strong>external</strong> part (LayerNorm/Dropout/residual [B, L, D]) — fully replicated under pure TP</li></ul>
<p><strong>Pure-TP per-GPU activation</strong>: $A^\text{TP}_\text{per-card} = A_\text{in}/T + A_\text{out}$.</p>
<p><strong>TP + SP per-GPU activation</strong>: SP also shards the TP-outer part along seq dim (each GPU holds $[B, L/T, D]$) → $A^\text{TP+SP}_\text{per-card} = A_\text{in}/T + A_\text{out}/T = A/T$.</p>
<p><strong>SP's savings over pure TP</strong>:</p>
<p>$$\boxed{\;A^\text{TP}_\text{per-card} - A^\text{TP+SP}_\text{per-card} = A_\text{out} \cdot \left(1 - \frac{1}{T}\right)\;}$$</p>
<p>For LLaMA-class models, $A_\text{out}$ is about 30-50% of the unsharded total activation. With TP=8, this saves $7/8 \cdot A_\text{out}$ ≈ <strong>a 25-40% reduction in total activation</strong> (consistent with Korthikanti 2023).</p>
<h3 id="84-selective-activation-recompute">8.4 Selective activation recompute</h3>
<p>Korthikanti et al. in the same paper proposes <strong>selective recompute</strong>: only recompute certain ops (e.g. attention's $QK^\top$ and softmax — the quadratic-memory ones), leaving the rest stored. Compared to full activation checkpointing it reduces compute overhead by 30-40% while saving the same memory.</p>
<h2 id="9-pipeline-parallel">§9 Pipeline Parallel</h2>
<h3 id="91-naive-pp-and-the-bubble">9.1 Naive PP and the bubble</h3>
<p>Shard $L$ layers along depth into $P$ stages ($L/P$ layers per stage). Naive PP runs each mini-batch sequentially through all stages:</p>
<pre><code>Stage 0 (GPU0): [F1] [B1]
Stage 1 (GPU1): [F1] [B1]
Stage 2 (GPU2): [F1] [B1]
Stage 3 (GPU3): [F1] [B1]
↑ many GPUs idle ↑ ↑ many idle ↑</code></pre>
<p>GPU utilization = $1/P$, completely unusable.</p>
<h3 id="92-gpipe-huang-et-al-neurips-2019">9.2 GPipe (Huang et al., NeurIPS 2019)</h3>
<p>Split the mini-batch into $M$ <strong>micro-batches</strong>; pipeline them:</p>
<pre><code>Stage 0: [F1][F2][F3][F4] [B4][B3][B2][B1]
Stage 1: [F1][F2][F3][F4] [B4][B3][B2][B1]
Stage 2: [F1][F2][F3][F4] [B4][B3][B2][B1]
Stage 3: [F1][F2][F3][F4] [B4][B3][B2][B1]
|←─ &quot;warm-up&quot; ─→| |← all micro-batch backward ─→| &quot;cool-down&quot;</code></pre>
<p><strong>Bubble</strong> (GPU-idle fraction):</p>
<p>$$\boxed{\;\text{bubble ratio} = \frac{P - 1}{M + P - 1}\;}$$</p>
<p>Derivation: each micro-batch traverses $P$ stages in $P$ steps; the warm-up phase (stage $i$ waits for $i$ previous micro-batches) has $P-1$ idle steps; the cool-down phase mirrors it, $P-1$ idle steps; total steps = $M + P - 1$ (forward) + $M + P - 1$ (backward) = $2(M+P-1)$, of which idle = $2(P-1)$. Idle fraction = $(P-1)/(M+P-1)$.</p>
<div class="callout callout-warn"><div class="callout-title">GPipe&#x27;s flaw</div><p>backward can only start after all $M$ micro-batches have completed forward; all micro-batches' activations must be held, so <strong>activation memory grows linearly with $M$</strong>. 1F1B solves this.</p></div>
<h3 id="93-1f1b--pipedream-narayanan-sosp-2019-megatron-lm-2-sc-2021">9.3 1F1B / PipeDream (Narayanan SOSP 2019, Megatron-LM-2 SC 2021)</h3>
<p><strong>1F1B = 1 Forward 1 Backward</strong>: each stage, after forwarding a micro-batch, <strong>immediately</strong> backwards the previously completed one:</p>
<pre><code>Stage 0: [F1][F2][F3][F4][B1][F5][B2][F6][B3][F7][B4]...
↑ once micro-batch 1&#x27;s backward is ready (its forward reached stage P)
do B1 immediately, freeing micro-batch 1&#x27;s activation memory</code></pre>
<p><strong>Key properties</strong>:</p>
<ul><li>Bubble ratio is <strong>still $(P-1)/(M+P-1)$</strong> (warm-up + cool-down unchanged)</li><li>But the number of simultaneously alive activations per stage is <strong>$P$</strong> (not GPipe's $M$) — saving significant activation memory</li></ul>
<pre><code>Stage i in steady state has the following activations in memory:
- micro-batches that finished forward but are awaiting backward
- because there are P-i stages after i, each running forward / backward one step
- so stage i holds P-i forwarded-but-not-backwarded activations
- first stage (i=0) holds the most: P; last stage holds the least: 1</code></pre>
<h3 id="94-interleaved-1f1b-megatron-lm-2-sc-2021">9.4 Interleaved 1F1B (Megatron-LM-2, SC 2021)</h3>
<p>Further <strong>compress the bubble</strong>. Each GPU does not hold $L/P$ consecutive layers but rather $V$ segments of <strong>non-consecutive layers</strong> (<strong>virtual stages</strong>):</p>
<pre><code>Original 1F1B (P=4, L=8 layers):
GPU0: layers 0,1 GPU1: layers 2,3 GPU2: layers 4,5 GPU3: layers 6,7
Interleaved (P=4, V=2, L=8):
GPU0: layers 0,4 GPU1: layers 1,5 GPU2: layers 2,6 GPU3: layers 3,7
(each GPU holds V=2 segments, total L/(PV) = 1 layer/segment)</code></pre>
<p>Each micro-batch flows through stage 0 (layer 0) → stage 1 (layer 1) → ... → stage 3 (layer 3) → back to stage 0 (layer 4) → ... → stage 3 (layer 7). <strong>A single micro-batch passes through the stage column $V$ times</strong>.</p>
<p><strong>Bubble ratio</strong> (Narayanan et al., SC 2021, Eq. 4):</p>
<p>$$\boxed{\;\text{interleaved bubble} = \frac{P-1}{V \cdot M + P - 1} \approx \frac{P-1}{V \cdot M}\;}$$</p>
<p>Replace $M$ in the denominator with $V \cdot M$ (the pipeline now has $V$× more virtual-stage passes), and warm-up / cool-down's $P-1$ remains. Cost: <strong>send/recv across GPUs per micro-batch becomes $V$× more</strong> (each micro-batch traverses the stage column $V$ times), so $V$ should not be too large (typically $V = 2$ to $4$).</p>
<div class="callout callout-info"><div class="callout-title">Intuitive interleaved bubble derivation</div><p>Vanilla 1F1B: bubble = $(P-1)/(M+P-1)$, denominator is micro-batch count $M$ plus warm/cool $P-1$; interleaved: each micro-batch traverses the stage column $V$ times, so the <strong>effective micro-batch count is $V \cdot M$</strong>, giving bubble = $(P-1)/(V M + P - 1) \approx (P-1)/(VM)$. The core intuition is "$V$× more micro-batches".</p></div>
<h3 id="95-1f1b-schedule-pseudo-code">9.5 1F1B schedule pseudo-code</h3>
<pre><code class="language-python">def one_f_one_b_schedule(P, M, stage_rank, num_warmup_microbatches):
&quot;&quot;&quot;
stage_rank: pipeline stage index held by current GPU (0..P-1)
num_warmup_microbatches = P - 1 - stage_rank
&quot;&quot;&quot;
# ===== Warm-up: stage_rank runs (P-1-stage_rank) forwards
for i in range(num_warmup_microbatches):
recv_activation_from_prev_stage()
out = forward(model, activation)
send_activation_to_next_stage(out)
# ===== Steady state: F1B1 alternating
num_microbatches_remaining = M - num_warmup_microbatches
for i in range(num_microbatches_remaining):
# forward
recv_activation_from_prev_stage()
out = forward(model, activation)
send_activation_to_next_stage(out)
# backward (the one previously forwarded, now arriving)
recv_grad_from_next_stage()
grad_in = backward(model, grad_out)
send_grad_to_prev_stage(grad_in)
# ===== Cool-down: remaining backwards
for i in range(num_warmup_microbatches):
recv_grad_from_next_stage()
grad_in = backward(model, grad_out)
send_grad_to_prev_stage(grad_in)</code></pre>
<h3 id="96-pp-communication-characteristics">9.6 PP communication characteristics</h3>
<p>PP only transmits <strong>activations / gradients</strong> at stage boundaries, with each transmission about $B/M \cdot L \cdot D$ bytes per micro-batch. <strong>Traffic is very small</strong> (much less than TP / DP), so PP can cross nodes (IB is enough).</p>
<h3 id="97-pps-flaws-load-imbalance">9.7 PP's flaws: load imbalance</h3>
<ul><li>Layers / compute per stage need manual balancing (embedding + LM head are heavy and often need special handling)</li><li>One stage's OOM / slowdown blocks the entire pipeline (weakest-link effect)</li><li>Too few micro-batches $M$ → large bubble; too many $M$ → activation memory grows</li></ul>
<h2 id="10-context-parallel--sharding-long-sequences">§10 Context Parallel — sharding long sequences</h2>
<p>As context windows grow from 4K → 128K → 1M, a single GPU cannot hold the full attention computation. <strong>Context Parallel (CP)</strong> shards along the sequence dim.</p>
<h3 id="101-ring-attention-liu-et-al-arxiv-231001889-2024">10.1 Ring Attention (Liu et al., arXiv 2310.01889, 2024)</h3>
<p>Shard $Q, K, V$ along seq dim across $C$ ranks: each rank holds $L/C$ tokens of Q/K/V.</p>
<pre><code>Rank 0: Q[0:L/C], K[0:L/C], V[0:L/C]
Rank 1: Q[L/C:2L/C], K[L/C:2L/C], V[L/C:2L/C]
...</code></pre>
<p>But attention requires every query to see all keys (causal: all past keys). The <strong>ring attention solution</strong>:</p>
<ol><li>Each rank computes a partial attention with its local Q and local K, V</li><li>Forward its local K, V along the ring to the next rank</li><li>The next rank computes partial attention with its local Q × the previous rank's K, V, accumulating into output (online-softmax style)</li><li>After $C$ rotations, every Q has seen all K, V — attention is complete</li></ol>
<p><strong>Communication</strong>: each rank holds $L/C$ tokens of K, V, size $2 \cdot L/C \cdot D$ (× 2 bytes in fp16). The ring rotates $C-1$ times so every K, V shard visits every rank. <strong>Per-rank total transmission</strong> $\approx (C-1) \cdot 2 L D / C \approx 2 L D$ (almost independent of $C$ — a typical ring property).</p>
<p><strong>Key</strong>: use <strong>online softmax</strong> (same as FlashAttention) so partial attention can be accumulated without materializing the full $L \times L$ score matrix.</p>
<div class="callout callout-info"><div class="callout-title">Relationship between Ring attention and FlashAttention</div><p>FlashAttention does block tiling within a single GPU (blocks in SRAM). Ring attention pushes this tiling to <strong>multi-GPU / multi-node level</strong>: each rank holds a K, V block and computes partial attention in ring order, accumulating. The two share the same math — online softmax + block-wise accumulation.</p></div>
<h3 id="102-llama-3-cp-implementation">10.2 Llama 3 CP implementation</h3>
<p>The Llama 3 paper (arXiv 2407.21783) reports using <strong>CP=16</strong> during the <strong>128K long-context</strong> phase (short-context 8K phase doesn't need CP, redirected to DP). Combined with FlashAttention v3, 128K context per-step time goes from untrainable → seconds.</p>
<h3 id="103-cp-and-tp-orthogonality">10.3 CP and TP orthogonality</h3>
<ul><li>TP shards hidden / head dim → intra-node</li><li>CP shards sequence dim → can cross nodes (traffic $\propto L$, not $L^2$)</li><li>Composed: each GPU holds $L/C$ tokens of a $D/T$-dim sub-tensor</li></ul>
<h2 id="11-expert-parallel--moe-routing">§11 Expert Parallel — MoE routing</h2>
<h3 id="111-basic-moe-structure">11.1 Basic MoE structure</h3>
<p>Mixture-of-Experts replaces a single FFN with $E$ expert FFNs + a gate / router:</p>
<p>$$y = \sum_{e=1}^E G_e(x) \cdot \text{Expert}_e(x), \quad G \in \mathbb{R}^E, \quad \sum_e G_e = 1$$</p>
<p>In practice, use <strong>top-K routing</strong>: only pick the top $K$ experts (typically $K = 1, 2$), with other experts not computed. <strong>Compute is independent of expert count $E$; only depends on $K$</strong> — this is what lets MoE scale parameters.</p>
<h3 id="112-expert-parallel-experts-across-gpus">11.2 Expert Parallel: experts across GPUs</h3>
<table><thead><tr><th>Mode</th><th>Experts per GPU</th><th>Communication</th></tr></thead><tbody><tr><td>Not parallel (replicate)</td><td>$E$</td><td>0</td></tr><tr><td>TP-style expert split</td><td>shard each expert, all GPUs compute each expert</td><td>many all-reduces</td></tr><tr><td><strong>EP</strong>: each GPU holds $E/N$ experts</td><td>$E/N$</td><td><strong>all-to-all</strong> dispatch + combine</td></tr></tbody></table>
<p>EP forward:</p>
<pre><code>1. Each GPU computes gates for its token batch → routing decision per token (assigned to e_1, e_2, ..., e_K)
2. all-to-all dispatch: send tokens to the GPU of their assigned expert
- Input: each GPU holds B/N tokens, each carrying K (expert_id, token_data)
- Output: each GPU receives the tokens destined for its local experts
3. Each GPU runs FFN on its local experts
4. all-to-all combine: send expert outputs back to the token&#x27;s original GPU
5. Each GPU merges K expert outputs by gate weight</code></pre>
<div class="callout callout-warn"><div class="callout-title">EP&#x27;s two flaws</div><p>(1) <strong>load imbalance</strong>: gates favor a few experts, overloading some GPUs; fix: load balancing loss (Switch Transformer / GShard). (2) <strong>all-to-all traffic is heavy</strong>: proportional to token count × hidden, with inter-node IB as the bottleneck. DeepSeek-V3 uses <strong>node-limited routing</strong> to cap each token's dispatch to at most $M$ nodes, reducing IB traffic.</p></div>
<h3 id="113-ep-all-to-all-code-skeleton-pseudo-code">11.3 EP all-to-all code skeleton (pseudo-code)</h3>
<p>Below is the <strong>pseudo-code skeleton</strong> of EP forward — the focus is the two all-to-all calls for dispatch / combine, with routing / bucketing engineering details omitted (production code: Megatron-Core MoE or DeepSpeed-MoE):</p>
<pre><code class="language-python">def expert_parallel_forward(
x, # [B, L, D]
gate, # nn.Module: tokens [BL, D] -&gt; (top_k_ids, top_k_w) each [BL, K]
experts_local, # E_local local experts (nn.ModuleList)
ep_group, ep_size, ep_rank,
E_total, K,
):
&quot;&quot;&quot; Pedagogical: each GPU holds E_local = E_total / ep_size experts &quot;&quot;&quot;
B, L, D = x.shape
tokens = x.reshape(B * L, D)
# 1. routing: each token picks K experts
top_k_ids, top_k_w = gate(tokens) # both [BL, K]
# 2. expand: top-K duplicates each token K times, one expert_id each
expanded = tokens.unsqueeze(1).expand(-1, K, -1).reshape(B * L * K, D)
expand_ids = top_k_ids.reshape(B * L * K) # [BL*K]
expand_w = top_k_w.reshape(B * L * K)
# 3. compute which EP rank each duplicated token should go to
target_rank = expand_ids // (E_total // ep_size) # [BL*K]
# 4. sort by target_rank + count send_count per rank
perm = torch.argsort(target_rank)
sorted_tokens = expanded[perm]
send_counts = torch.bincount(target_rank, minlength=ep_size).tolist()
# 5a. exchange send_counts -&gt; recv_counts (each rank tells others how many to expect)
send_t = torch.tensor(send_counts, dtype=torch.int64, device=x.device)
recv_t = torch.empty_like(send_t)
dist.all_to_all_single(recv_t, send_t, group=ep_group) # one small a2a to sync counts
recv_counts = recv_t.tolist()
# 5b. sync token expert_ids (for assigning to local experts)
sorted_ids = expand_ids[perm]
received_ids = torch.empty(sum(recv_counts), dtype=sorted_ids.dtype, device=x.device)
dist.all_to_all_single(received_ids, sorted_ids,
output_split_sizes=recv_counts,
input_split_sizes=send_counts,
group=ep_group)
# 5c. sync token data
received_tokens = torch.empty(sum(recv_counts), D, device=x.device, dtype=x.dtype)
dist.all_to_all_single(received_tokens, sorted_tokens,
output_split_sizes=recv_counts,
input_split_sizes=send_counts,
group=ep_group)
# 6. local expert compute
received_out = torch.zeros_like(received_tokens)
for local_eid in range(len(experts_local)):
global_eid = ep_rank * (E_total // ep_size) + local_eid
mask = (received_ids == global_eid)
if mask.any():
received_out[mask] = experts_local[local_eid](received_tokens[mask])
# 7. all-to-all combine: reverse, swap split sizes
combined = torch.empty_like(sorted_tokens)
dist.all_to_all_single(combined, received_out,
output_split_sizes=send_counts, # reversed
input_split_sizes=recv_counts,
group=ep_group)
# 8. inverse permute + gate weight merge
inv_perm = torch.empty_like(perm)
inv_perm[perm] = torch.arange(perm.numel(), device=perm.device)
out_expanded = combined[inv_perm] # [BL*K, D]
out_expanded = out_expanded * expand_w.unsqueeze(-1)
out = out_expanded.view(B * L, K, D).sum(dim=1) # [BL, D]
return out.view(B, L, D)</code></pre>
<div class="callout callout-warn"><div class="callout-title">Production implementations are far more complex</div><p>real code handles (a) capacity factor (preventing expert overload); (b) <strong>dropping tokens</strong> when an expert is full; (c) NVL/IB hierarchical all-to-all (DeepSeek's node-limited routing); (d) backward's mirror all-to-all + gate gradient. This skeleton only illustrates the dispatch / combine bidirectional flow.</p></div>
<h2 id="12-activation-memory-optimization">§12 Activation memory optimization</h2>
<h3 id="121-gradient-checkpointing-chen-et-al-arxiv160406174-2016">12.1 Gradient Checkpointing (Chen et al., arXiv:1604.06174, 2016)</h3>
<p>Don't store intermediate activations; recompute on backward. Space $O(\sqrt{L})$ (with optimal segmentation), time +33% (one extra forward).</p>
<pre><code class="language-python">from torch.utils.checkpoint import checkpoint
def block_forward(x):
return transformer_block(x)
# Backward recomputes block_forward; no intermediate activation stored
y = checkpoint(block_forward, x, use_reentrant=False)</code></pre>
<h3 id="122-selective-recompute-korthikanti-2023">12.2 Selective recompute (Korthikanti 2023)</h3>
<p>Only recompute <strong>quadratic-memory ops</strong> (attention's $QK^\top$ and softmax); store everything else. Reduces 30-40% compute overhead compared to full recompute while saving the same memory. Megatron-LM enables this by default.</p>
<h3 id="123-offload-zero-infinity">12.3 Offload (ZeRO-Infinity)</h3>
<p>Offload activations / optimizer state to CPU RAM or NVMe. CPU offload is suitable for 13B-30B single-machine training; NVMe offload has very low throughput and is mainly used for inference or trillion-scale model exploration.</p>
<h3 id="124-activation-memory-formula-must-memorize">12.4 Activation memory formula (must memorize)</h3>
<p>A single transformer block's activations are roughly (fp16, full save):</p>
<p>$$A_\text{block} \approx \underbrace{34 \cdot B \cdot L \cdot D}_{\text{LayerNorm/QKV/output/MLP residual}} + \underbrace{5 \cdot B \cdot L^2 \cdot H}_{\text{attention intermediates}}$$</p>
<p>See Korthikanti et al. 2023 Table 2 for details. When $L \gg D$, the second term dominates (FlashAttention kills this part); when $L \approx D$, the first term dominates (SP reduces it by a factor of $T$).</p>
<h2 id="13-synthesis-3d--4d--5d-parallelism">§13 Synthesis: 3D / 4D / 5D Parallelism</h2>
<p>Combine the dimensions. With world size $W$:</p>
<p>$$W = D_{DP} \times T_{TP} \times P_{PP} \times C_{CP} \times E_{EP}$$</p>
<p>Roles of each axis:</p>
<table><thead><tr><th>Axis</th><th>Shards what</th><th>Where</th><th>Communication primitive</th></tr></thead><tbody><tr><td>DP / FSDP</td><td>batch + model states</td><td>OK across nodes via IB</td><td>all-reduce / reduce-scatter + all-gather</td></tr><tr><td>TP</td><td>hidden / head within a layer</td><td><strong>intra-node NVLink</strong></td><td>all-reduce (4× per block)</td></tr><tr><td>SP</td><td>LayerNorm / Dropout activations outside TP</td><td>intra-node (shares group with TP)</td><td>reduce-scatter + all-gather</td></tr><tr><td>PP</td><td>layer depth</td><td>OK across nodes via IB</td><td>point-to-point send/recv</td></tr><tr><td>CP</td><td>sequence dim</td><td>either intra-node or across nodes (traffic $\propto L$)</td><td>ring K/V</td></tr><tr><td>EP</td><td>MoE experts</td><td>OK across nodes via IB (heavy all-to-all)</td><td>all-to-all</td></tr></tbody></table>
<h3 id="131-llama-3-405b-training-topology-public">13.1 Llama 3 405B training topology (public)</h3>
<p>Meta 2024 (<a href="https://arxiv.org/abs/2407.21783">2407.21783</a>):</p>
<ul><li>16K H100 GPUs (16384 total)</li><li><p>Parallelism (total GPU count constant; CP up means DP down):</p>
<ul><li><strong>Short-context phase (8K)</strong>: TP=8 × CP=1 × PP=16 × DP=128 = 16384</li><li><strong>Long-context phase (128K)</strong>: TP=8 × CP=16 × PP=16 × DP=8 = 16384</li></ul></li><li>Training precision: <strong>BF16</strong> (paper Table reports BF16 MFU); FP8 is used for inference quantization, <strong>not for 405B training</strong></li><li><strong>54 days, 466 interruptions</strong> (419 unexpected + 47 planned/maintenance), 78% of unexpected ones from hardware causes</li><li>Effective training time &gt; 90%</li></ul>
<h3 id="132-deepseek-v3-training-topology-public">13.2 DeepSeek-V3 training topology (public)</h3>
<p>DeepSeek 2024 (<a href="https://arxiv.org/abs/2412.19437">2412.19437</a>):</p>
<ul><li>2048 H800 GPUs</li><li>Parallelism: TP=1 (<strong>no TP</strong>! offset by ZeRO + EP + PP) × PP=16 × EP=64 (across 8 nodes) × ZeRO-1 DP</li><li><strong>DualPipe</strong> bidirectional pipeline (see next section) + all-to-all overlap</li><li>fp8 GEMM + bf16 accumulation</li></ul>
<div class="callout callout-info"><div class="callout-title">Why DeepSeek-V3 doesn&#x27;t use TP</div><p>V3 uses MLA (multi-head latent attention) + many MoE experts; TP-on-head returns are small (the latent attention head dim is already small); and EP's all-to-all overlap with DualPipe hides communication completely. The combination <strong>PP × EP × ZeRO</strong> is enough to fit 671B parameters.</p></div>
<h2 id="14-dualpipe--the-2024-pipeline-frontier">§14 DualPipe — the 2024 pipeline frontier</h2>
<p>DeepSeek 2024's <strong>DualPipe</strong> algorithm, published in the V3 paper and an independent repo (arXiv 2412.19437 + github.com/deepseek-ai/DualPipe).</p>
<h3 id="141-core-idea">14.1 Core idea</h3>
<p>1F1B's bubble comes from the <strong>warm-up + cool-down</strong> phases. DualPipe runs <strong>two directions of the pipeline simultaneously</strong> — one group of micro-batches goes stage 0 → P, the other group goes P → 0; when they meet in the middle, they exactly fill the warm-up / cool-down gaps.</p>
<pre><code>Traditional 1F1B (P=4):
Stage 0: [F1][F2][F3][F4][B1][F5][B2]...
Stage 1: [F1][F2][F3][F4][B1][F5][B2]...
(bubbles at both ends)
DualPipe (P=4):
Forward-direction micro-batch: [F1][F2][F3][F4]...
Reverse-direction micro-batch: ...[F4&#x27;][F3&#x27;][F2&#x27;][F1&#x27;]
Stage 0: simultaneously processes forward&#x27;s F + reverse&#x27;s F&#x27; + corresponding B ← compute / comms fully overlap</code></pre>
<p>More precisely: DualPipe designs a <strong>bidirectional schedule</strong> in which every GPU at every time has two micro-batches in overlapped forward / backward execution; expert-parallel all-to-all communication is also hidden in the gaps of the two micro-batch streams.</p>
<h3 id="142-dualpipe-vs-vanilla-1f1b-properties">14.2 DualPipe vs vanilla 1F1B properties</h3>
<table><thead><tr><th>Dimension</th><th>1F1B</th><th>Interleaved 1F1B</th><th>DualPipe</th></tr></thead><tbody><tr><td>Bubble</td><td>$(P-1)/(M+P-1)$</td><td>$(P-1)/(VM+P-1)$</td><td><strong>ideally 0</strong> (warm-up / cool-down complementary)</td></tr><tr><td>Activation memory</td><td>$P$ × per stage</td><td>$V \cdot P$ ×</td><td>$\approx 2 \times P$</td></tr><tr><td>Communication overlap</td><td>partial</td><td>same as 1F1B</td><td><strong>near-100% all-to-all overlap</strong></td></tr><tr><td>Implementation complexity</td><td>medium</td><td>high</td><td>very high (needs forward &amp; reverse scheduling)</td></tr></tbody></table>
<h3 id="143-dualpipe-costs">14.3 DualPipe costs</h3>
<ul><li><strong>Massive code complexity</strong>: a stage simultaneously runs forward (direction 1) + forward (direction 2) + backward; CUDA stream management is very hard</li><li><strong>2× activation memory</strong>: both directions must store activations, $\approx 2\times$ over 1F1B</li><li><strong>Stages must be balanced</strong>: any stuck stage blocks both directions</li></ul>
<p>DeepSeek uses DualPipe because EP all-to-all traffic is so heavy it must be hidden. <strong>For ordinary training, vanilla 1F1B is sufficient</strong>.</p>
<h2 id="15-torchtitan--pytorch-native-4d-platform">§15 TorchTitan — PyTorch-native 4D platform</h2>
<p>Liang et al. (ICLR 2025, arXiv 2410.06511) <strong>"TorchTitan: One-stop PyTorch Native Solution for Production-Ready LLM Pre-training"</strong>.</p>
<h3 id="151-design-goals">15.1 Design goals</h3>
<ul><li><strong>No more monkey-patching</strong>: integrate FSDP2 / TP / PP / SP / CP / Float8 / <code>torch.compile</code> into PyTorch mainline</li><li><strong>DTensor as the unified language</strong>: all sharding described by DTensor placement (<code>Shard(d)</code>, <code>Replicate</code>, <code>Partial</code>)</li><li><strong>Composable</strong>: FSDP2 composes naturally with TP (FSDP1 was nearly impossible to compose with TP)</li></ul>
<h3 id="152-code-style-vs-deepspeed-monkey-patch">15.2 Code style (vs DeepSpeed monkey patch)</h3>
<pre><code class="language-python"># TorchTitan style: declarative DTensor placement
import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
parallelize_module, RowwiseParallel, ColwiseParallel,
SequenceParallel, PrepareModuleInput,
)
from torch.distributed.fsdp import fully_shard
# 1. Build the 4D device mesh
mesh = init_device_mesh(
&quot;cuda&quot;,
mesh_shape=(2, 8, 4, 8), # PP=2, FSDP=8, CP=4, TP=8
mesh_dim_names=(&quot;pp&quot;, &quot;fsdp&quot;, &quot;cp&quot;, &quot;tp&quot;),
)
# 2. Apply TP + SP (declare sharding strategy per sub-module)
parallelize_module(
model,
mesh[&quot;tp&quot;],
{
&quot;attn.wq&quot;: ColwiseParallel(),
&quot;attn.wk&quot;: ColwiseParallel(),
&quot;attn.wv&quot;: ColwiseParallel(),
&quot;attn.wo&quot;: RowwiseParallel(),
&quot;mlp.fc1&quot;: ColwiseParallel(),
&quot;mlp.fc2&quot;: RowwiseParallel(),
&quot;norm1&quot;: SequenceParallel(),
&quot;norm2&quot;: SequenceParallel(),
},
)
# 3. FSDP2 wrap (auto-detects mesh[&quot;fsdp&quot;])
for block in model.blocks:
fully_shard(block, mesh=mesh[&quot;fsdp&quot;])
fully_shard(model, mesh=mesh[&quot;fsdp&quot;])
# 4. PP (pipeline schedule, 1F1B interleaved, pseudo-illustration)
# Real ScheduleInterleaved1F1B needs List[PipelineStage] (each rank holds V virtual stages)
from torch.distributed.pipelining import PipelineStage, ScheduleInterleaved1F1B
stages = [PipelineStage(submod_v0, ...), PipelineStage(submod_v1, ...)] # V=2 virtual stages
schedule = ScheduleInterleaved1F1B(stages, n_microbatches=32, loss_fn=loss_fn)</code></pre>
<h3 id="153-float8-training-hopper--blackwell">15.3 Float8 training (Hopper / Blackwell)</h3>
<p>TorchTitan integrates <strong>Float8 training</strong> (H100/B100 fp8 GEMM) — weights / activations in fp8, accumulation in fp32. Typical torchao usage (refer to current torchao docs for exact API):</p>
<pre><code class="language-python">import torch.nn as nn
from torchao.float8 import convert_to_float8_training, Float8LinearConfig
convert_to_float8_training(
model,
config=Float8LinearConfig(), # default dynamic scaling
module_filter_fn=lambda m, n: isinstance(m, nn.Linear) and &quot;lm_head&quot; not in n,
)</code></pre>
<p>Effect: H100 throughput +20-40%; loss is nearly identical (block-wise scaling is the key).</p>
<h3 id="154-async-checkpointing">15.4 Async checkpointing</h3>
<pre><code class="language-python">import torch.distributed.checkpoint as DCP
# Async save: returns a Future, does not block training
future = DCP.async_save(model.state_dict(), checkpoint_id=&quot;step_10000&quot;)
# ... continue training
future.result() # wait if needed</code></pre>
<p>DCP (Distributed Checkpoint) combined with FSDP2's sharded state dict makes <strong>single checkpoint writes take only a few seconds</strong> (each rank writes its portion to distributed storage), without blocking training.</p>
<h2 id="16-communication-primitives-summary">§16 Communication primitives summary</h2>
<p>To recall quickly in interviews, a closing table.</p>
<table><thead><tr><th>Operation</th><th>When the all-rank total is $S$ (each rank has $S$)</th><th>Per-rank traffic (ring)</th><th>Used in</th></tr></thead><tbody><tr><td><code>all-reduce(buf, SUM)</code></td><td>$N \cdot S$ → all ranks $S$</td><td>$2(N-1)/N \cdot S$</td><td>DDP gradient sync, TP block tail</td></tr><tr><td><code>reduce-scatter(buf)</code></td><td>$N \cdot S$ → each rank $S/N$</td><td>$(N-1)/N \cdot S$</td><td>ZeRO grad reduce, SP exit</td></tr><tr><td><code>all-gather(shard)</code></td><td>$N \cdot S/N$ → all ranks $S$</td><td>$(N-1)/N \cdot S$</td><td>ZeRO-3 forward param, SP entry</td></tr><tr><td><code>broadcast(buf, src)</code></td><td>rank src's $S$ → all ranks $S$</td><td>$S$ (tree) / $(N-1)/N \cdot S$ (ring)</td><td>model init broadcast</td></tr><tr><td><code>all-to-all(buf)</code></td><td>$N \cdot N$ blocks → transpose</td><td>$(N-1)/N \cdot S$</td><td>MoE EP routing, sequence sharding transform</td></tr><tr><td><code>point-to-point send/recv</code></td><td>1 → 1</td><td>$S$</td><td>PP stage boundary</td></tr></tbody></table>
<h2 id="17-25-frequently-asked-interview-questions">§17 25 frequently-asked interview questions</h2>
<p>Ordered by L1 / L2 / L3, with collapsible answer key points.</p>
<h3 id="l1-must-know-any-ml-engineer-role-will-ask">L1 must-know (any ML engineer role will ask)</h3>
<details>
<summary>Q1. Difference between DDP and DP (DataParallel)?</summary>
<ul><li>DP (<code>nn.DataParallel</code>): single-process multi-GPU; main rank scatters input + gathers output; <strong>GIL + main-rank bottleneck</strong>, deprecated</li><li>DDP (<code>DistributedDataParallel</code>): multi-process multi-GPU, one GPU per process, NCCL all-reduce for grad sync</li><li>DDP is 1.5-3× faster than DP and scales far better</li></ul>
<p>Footgun: saying DP is "fast" — wrong, DP is a legacy API; production uses DDP.</p>
</details>
<details>
<summary>Q2. NCCL all-reduce traffic?</summary>
<ul><li>Ring algorithm: per-GPU total traffic $2(N-1)/N \cdot S \approx 2S$ bytes</li><li>Nearly independent of $N$ (the ring's essence)</li><li>Equivalent to reduce-scatter ($S$) + all-gather ($S$)</li></ul>
<p>Footgun: saying $N \cdot S$ or $S/N$; forgetting the ring step count and per-step traffic.</p>
</details>
<details>
<summary>Q3. Per-GPU model states under Adam mixed-precision training?</summary>
<ul><li>Parameters fp16: $2\Phi$</li><li>Gradients fp16: $2\Phi$</li><li>Optimizer (fp32 master + Adam m + v): $4\Phi + 4\Phi + 4\Phi = 12\Phi$</li><li><strong>Total $16\Phi$ bytes</strong></li></ul>
<p>Footgun: forgetting the fp32 master copy; or treating Adam $m, v$ as fp16 (they're fp32).</p>
</details>
<details>
<summary>Q4. What do ZeRO 1/2/3 shard?</summary>
<ul><li>ZeRO-1: shard optimizer state</li><li>ZeRO-2: + gradient</li><li>ZeRO-3: + parameter (most aggressive)</li></ul>
<p>Per-GPU model states drop from $16\Phi$ to $16\Phi/N$ (ZeRO-3).</p>
<p>Footgun: getting the order wrong; not knowing ZeRO-1/2 traffic is the same as DDP ($2\Phi$).</p>
</details>
<details>
<summary>Q5. Difference between FSDP and ZeRO-3?</summary>
<ul><li>Algorithmically: <strong>completely equivalent</strong> (FSDP is PyTorch's ZeRO-3 implementation)</li><li>Engineering: FSDP is integrated in the PyTorch mainline; ZeRO-3 is in the DeepSpeed library</li><li>FSDP2 uses per-parameter DTensor sharding, composing naturally with TP, with simpler state_dict</li></ul>
<p>Footgun: saying "FSDP has less traffic than ZeRO" or vice versa — wrong, the traffic is the same.</p>
</details>
<details>
<summary>Q6. Which dim does Tensor Parallel shard attention along?</summary>
<ul><li>Shard the <strong>head dim</strong> (each GPU holds $H/T$ heads)</li><li>$W_Q, W_K, W_V$ are column-parallel (shard output dim)</li><li>$W_O$ is row-parallel (shard input dim)</li><li>Each attention block forward 1× all-reduce, backward 1× all-reduce</li></ul>
<p>Footgun: saying it shards hidden_dim; or forgetting that col + row pairing makes the middle communication-free.</p>
</details>
<details>
<summary>Q7. Where must TP be placed?</summary>
<ul><li><strong>Intra-node NVLink domain</strong> (900 GB/s bandwidth)</li><li>Going across nodes via IB (50 GB/s) crushes TP performance</li><li>This is why TP=8 is the golden size (one node, 8 GPUs)</li></ul>
<p>Footgun: saying TP can cross nodes arbitrarily; or confusing intra-node / inter-node bandwidths.</p>
</details>
<details>
<summary>Q8. How is Pipeline Parallel's bubble computed?</summary>
<ul><li>Naive PP: $M = 1$, bubble = $(P-1)/P$ (huge)</li><li>GPipe / 1F1B with $M$ micro-batches: $(P-1)/(M+P-1)$</li><li>Typically need $M \geq 4P$ for bubble &lt; 20%</li></ul>
<p>Footgun: forgetting $M$ in the denominator; saying only $1/P$ without mentioning $M$.</p>
</details>
<details>
<summary>Q9. What advantage does 1F1B have over GPipe?</summary>
<ul><li><strong>Same bubble ratio</strong> $(P-1)/(M+P-1)$</li><li>But 1F1B's simultaneously-alive activation count per stage = $P$ (not GPipe's $M$)</li><li>Saves significant <strong>activation memory</strong></li></ul>
<p>Footgun: saying 1F1B reduces the bubble — wrong, 1F1B doesn't reduce bubble; it saves activations.</p>
</details>
<details>
<summary>Q10. Cost of activation checkpointing (gradient checkpointing)?</summary>
<ul><li>Memory: $O(L) \to O(\sqrt{L})$</li><li>Time: +33% (one extra forward)</li><li>Production almost always enables this (70B+ models OOM without it)</li></ul>
<p>Footgun: saying "memory is halved" — imprecise, theory is $\sqrt{L}$; or saying "time is halved".</p>
</details>
<h3 id="l2-advanced-research-oriented-roles">L2 advanced (research-oriented roles)</h3>
<details>
<summary>Q11. Why does ZeRO-3 forward need all-gather?</summary>
<ul><li>Parameters are sharded across $N$ GPUs → each GPU holds $1/N$</li><li>A layer forward needs the full $W^{(\ell)}$ → temporarily <strong>all-gather</strong> to all ranks</li><li>Forward completes → <strong>release</strong> immediately, keep only the shard</li><li>Backward needs to all-gather again (released after forward)</li></ul>
<p>Footgun: thinking parameters are resident on all ranks; or forgetting backward also all-gathers.</p>
</details>
<details>
<summary>Q12. Derive: DDP vs ZeRO-3 traffic difference.</summary>
<ul><li>DDP: backward 1× all-reduce = $2\Phi$, total $2\Phi$</li><li>ZeRO-3: forward 1× all-gather ($\Phi$) + backward 1× all-gather ($\Phi$) + 1× reduce-scatter ($\Phi$) = $3\Phi$</li><li><strong>ZeRO-3 has 50% more traffic than DDP, in exchange for $N\times$ memory reduction</strong></li></ul>
<p>Footgun: thinking ZeRO reduces traffic — wrong, ZeRO reduces memory, may increase traffic (depending on the stage).</p>
</details>
<details>
<summary>Q13. NCCL ring all-reduce single-step traffic?</summary>
<ul><li>Total steps: $2(N-1)$ (reduce-scatter $N-1$ + all-gather $N-1$)</li><li>Per step: each rank sends $S/N$ bytes</li><li>Per-GPU total: $2(N-1) \cdot S/N \approx 2S$</li><li>Bandwidth utilization $\to 1$ as $N \to \infty$</li></ul>
<p>Footgun: saying $S$ instead of $2S$; or saying $N \cdot S$ (forgets ring's essence).</p>
</details>
<details>
<summary>Q14. What does SP (Sequence Parallel) save?</summary>
<ul><li>Does not save traffic (same total as pure TP)</li><li>Saves <strong>TP-external activation memory</strong> (LayerNorm / Dropout)</li><li>Shards the fully-replicated $[B, L, D]$ to $[B, L/T, D]$</li><li>Total activation memory drops 25-40%</li></ul>
<p>Footgun: saying "SP reduces traffic" — wrong, SP only redistributes the all-reduce into reduce-scatter + all-gather (same total), but activation is sharded.</p>
</details>
<details>
<summary>Q15. How does interleaved 1F1B reduce the bubble?</summary>
<ul><li>Each stage holds $V$ virtual stages (V discontinuous layer segments)</li><li>Bubble: $(P-1)/(VM+P-1) \approx (P-1)/(VM)$</li><li>For the same $M$, bubble drops $V$×</li><li>Cost: communication count × $V$</li></ul>
<p>Footgun: saying "interleaved reduces compute" — wrong, it only redistributes the time axis; or forgetting the communication × V cost.</p>
</details>
<details>
<summary>Q16. EP all-to-all traffic in MoE?</summary>
<ul><li>Each token picks $K$ experts</li><li>Dispatch: each GPU sends its $B/N$ tokens to the appropriate expert GPU</li><li>Combine: reverse</li><li>Per-rank total traffic $\approx 2 \cdot K \cdot B/N \cdot D$ (bidirectional all-to-all)</li></ul>
<p>Footgun: forgetting top-K (not $E$); forgetting that dispatch + combine = 2 calls.</p>
</details>
<details>
<summary>Q17. HSDP vs FSDP difference?</summary>
<ul><li>FSDP: all GPUs in the same sharding group</li><li>HSDP: FSDP / ZeRO-3 within groups, DDP across groups</li><li>HSDP has less intra-group communication (intra-node NVLink); inter-group uses grad all-reduce (no weight all-gather)</li><li>Trade-off: more model states intra-group (not sharded to full world), exchanged for less inter-node communication</li></ul>
<p>Footgun: thinking HSDP reduces total traffic — strictly, it reduces inter-node traffic, with intra-group / inter-group being a trade-off.</p>
</details>
<details>
<summary>Q18. What parallelism did Llama 3 405B use?</summary>
<ul><li>16K H100s (16384 total, constant; CP up means DP down)</li><li>Short context (8K): TP=8 × CP=1 × PP=16 × DP=128</li><li>Long context (128K): TP=8 × CP=16 × PP=16 × DP=8</li><li>54 days, 466 interruptions (419 unexpected + 47 planned), 90%+ effective training time</li><li>Training in BF16 (not FP8 — FP8 is for inference quantization)</li><li>Meta 2024, arXiv 2407.21783</li></ul>
<p>Footgun: saying it used EP — wrong, Llama 3 is dense, no EP; or forgetting the CP phase exists.</p>
</details>
<details>
<summary>Q19. What are ZeRO++'s three tricks?</summary>
<ul><li><strong>qwZ</strong>: forward all-gather uses int8 quantization (block-wise quant)</li><li><strong>hpZ</strong>: weight replicated within a node (NVLink domain), sharded across nodes — backward all-gather runs over NVLink</li><li><strong>qgZ</strong>: backward gradient reduce-scatter also uses int8</li><li>Total traffic drops 4×, throughput on 384 GPUs +116% (Wang 2023)</li></ul>
<p>Footgun: thinking ZeRO++ changes training precision — it only quantizes buffers in the communication path; weights and compute remain fp16/bf16.</p>
</details>
<details>
<summary>Q20. Can a 7B model be trained in fp16 on 8 A100-40Gs?</summary>
<ul><li>Model states (Adam): $16 \times 7$B $= 112$ GB</li><li>DDP: per-GPU 112 GB, <strong>does not fit</strong> (A100 40G)</li><li>ZeRO-3 / FSDP: per-GPU $112/8 = 14$ GB ✓</li><li><ul><li>activation (with checkpoint): several GB ✓</li></ul></li><li><ul><li>workspace: several GB</li></ul></li><li><strong>Conclusion: FSDP/ZeRO-3 + activation checkpointing can run it</strong></li></ul>
<p>Footgun: saying "DDP also works"; or only counting parameter $\Phi$ without model states; or forgetting fp32 master copy.</p>
</details>
<h3 id="l3-advanced-variants-top-labs--in-house-infra-roles">L3 advanced variants (top labs / in-house infra roles)</h3>
<details>
<summary>Q21. Derive 1F1B bubble ratio + how interleaved further reduces it.</summary>
<ul><li>1F1B: warm-up $P-1$ steps (forward fill) + cool-down $P-1$ steps (backward drain) + steady $M-P+1$ steps</li><li>Total steps $= 2M$ (forward + backward) takes $2(M + P - 1)$ time slots (including $P-1$ idle slots each at warm/cool)</li><li>Bubble = $2(P-1) / [2(M+P-1)] = (P-1)/(M+P-1)$</li><li><strong>Interleaved 1F1B</strong> (Narayanan SC 2021, Eq. 4): each GPU holds $V$ non-contiguous layer segments (virtual stages); each micro-batch traverses the stage column $V$ times → pipeline has effectively $V \cdot M$ micro-batches in queue</li><li>Bubble = $(P-1) / (V \cdot M + P - 1) \approx (P-1)/(V \cdot M)$</li><li><strong>For the same $M$, bubble drops $V$×; cost: cross-GPU send/recv ×$V$</strong></li></ul>
<p>Footgun: forgetting interleaved's communication × V cost; saying "interleaved reduces forward compute" — it only redistributes time.</p>
</details>
<details>
<summary>Q22. FSDP vs ZeRO-3 traffic + applicable scenarios.</summary>
<ul><li><strong>Traffic is completely the same</strong>: forward all-gather $\Phi$ + backward all-gather $\Phi$ + reduce-scatter $\Phi$ = $3\Phi$</li><li><p>Engineering differences:</p>
<ul><li>FSDP2: uses DTensor to describe per-param sharding, composes naturally with TP / PP, <code>torch.compile</code> friendly</li><li>DeepSpeed ZeRO-3: uses flat-parameter, needs monkey-patching, but has a complete ZeRO++ / Offload / Infinity ecosystem</li></ul></li><li><p>Choice:</p>
<ul><li>New project (dense + 4D parallelism) → <strong>FSDP2 / TorchTitan</strong></li><li>MoE + cross-framework + offload → <strong>DeepSpeed</strong></li></ul></li></ul>
<p>Footgun: saying "FSDP has less traffic than ZeRO" — algorithms are the same; or saying "FSDP doesn't support offload" — FSDP2 OffloadPolicy supports it.</p>
</details>
<details>
<summary>Q23. How much activation memory does TP + SP save vs pure TP?</summary>
<ul><li>Let transformer-block activation be $A_\text{block}$</li><li>TP-internal (attention intermediates / MLP intermediates): about $A_\text{block} \times 0.5-0.7$, TP shards it</li><li>TP-external (LayerNorm / Dropout / residual): about $A_\text{block} \times 0.3-0.5$, <strong>TP does not shard (full replica)</strong></li><li>Pure-TP per-GPU activation = $A_\text{TP-in}/T + A_\text{TP-out}$</li><li>TP+SP per-GPU activation = $A_\text{TP-in}/T + A_\text{TP-out}/T = A_\text{block}/T$</li><li><strong>Saves $A_\text{TP-out} \cdot (1 - 1/T)$</strong>, about <strong>25-40% of total activation</strong> (depends on model)</li></ul>
<p>Footgun: saying "SP reduces activation by $T$×" — imprecise, only for the TP-outer part; or forgetting that SP traffic is the same as pure TP.</p>
</details>
<details>
<summary>Q24. Core improvement of DualPipe over 1F1B?</summary>
<ul><li><strong>Bidirectional pipeline</strong>: one group of micro-batches goes stage 0 → P, another P → 0</li><li>Meeting in the middle, they exactly fill 1F1B's warm-up / cool-down bubble</li><li><strong>Theoretical bubble = 0</strong> (ideal case, two sides complementary)</li><li>Key gain: <strong>all-to-all comms fully overlap</strong> (important for EP)</li><li>Cost: activation memory × 2 (both directions store activations); extremely complex implementation</li><li>DeepSeek-V3 December 2024 (arXiv 2412.19437) + github.com/deepseek-ai/DualPipe</li></ul>
<p>Footgun: saying "DualPipe is a new way to shard PP" — it's a new schedule; or forgetting the 2× activation cost.</p>
</details>
<details>
<summary>Q25. How is 4D / 5D parallelism composed in frontier training (e.g. DeepSeek-V3 / Llama 3)?</summary>
<ul><li><strong>5 orthogonal dimensions</strong>: DP / FSDP × TP × PP × CP × EP</li><li>World size $W = D \times T \times P \times C \times E$</li><li><p>Rules of thumb:</p>
<ul><li>TP=8 (must fit in one NVLink domain)</li><li>PP=8-32 (OK across nodes via IB; bubble controlled via $M \geq 4P$)</li><li>CP depends on context length (Llama 3 uses CP=16 at 128K; 1M may need CP=64+)</li><li>EP depends on MoE expert count (DeepSeek-V3 uses EP=64)</li><li>FSDP / DP uses the remaining world size</li></ul></li><li><p><strong>Llama 3 405B</strong> (16K GPUs constant, CP up means DP down):</p>
<ul><li>8K context: TP=8 × CP=1 × PP=16 × DP=128</li><li>128K context: TP=8 × CP=16 × PP=16 × DP=8</li></ul></li><li><strong>DeepSeek-V3 671B</strong>: TP=1 + PP=16 × EP=64 × ZeRO-1 DP=2, 2048 H800s total</li><li><p><strong>Key engineering points</strong>:</p>
<ul><li>Place communication primitives on the right topology (TP on NVLink, PP / DP on IB)</li><li>DualPipe / Interleaved 1F1B to compress PP bubble</li><li>FlashAttention v3 + SP / CP to compress activations</li><li>fp8 GEMM (H100 / B100) + bf16 accumulation to lift throughput</li></ul></li></ul>
<p>Footgun: reversing the dimension order; forgetting TP must be intra-node; thinking DeepSeek uses TP (V3 actually has TP=1).</p>
</details>
<h2 id="a-appendix-complete-4d-wrap-code-skeleton">§A Appendix: complete 4D wrap code skeleton</h2>
<p>Below is a minimal end-to-end 4D-parallelism wrap example (FSDP2 + TP + SP + PP), in TorchTitan style.</p>
<pre><code class="language-python">import torch
import torch.nn as nn
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
parallelize_module, RowwiseParallel, ColwiseParallel,
SequenceParallel, PrepareModuleInput, PrepareModuleOutput,
)
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.pipelining import pipeline, SplitPoint, ScheduleInterleaved1F1B
from torch.utils.checkpoint import checkpoint
def build_4d_model(model, world_size):
&quot;&quot;&quot; 4D parallelism wrap (PP × FSDP × CP × TP) &quot;&quot;&quot;
# Step 1. Build 4D device mesh
# e.g. 64 GPUs = PP=2 × FSDP=4 × CP=1 × TP=8
mesh = init_device_mesh(
&quot;cuda&quot;,
mesh_shape=(2, 4, 1, 8),
mesh_dim_names=(&quot;pp&quot;, &quot;fsdp&quot;, &quot;cp&quot;, &quot;tp&quot;),
)
# Step 2. TP + SP wrap (mesh[&quot;tp&quot;])
tp_plan = {
# Attention
&quot;attn.wq&quot;: ColwiseParallel(),
&quot;attn.wk&quot;: ColwiseParallel(),
&quot;attn.wv&quot;: ColwiseParallel(),
&quot;attn.wo&quot;: RowwiseParallel(),
# MLP
&quot;mlp.fc1&quot;: ColwiseParallel(),
&quot;mlp.fc2&quot;: RowwiseParallel(),
# SP: norm / residual sharded along seq dim
&quot;ln1&quot;: SequenceParallel(),
&quot;ln2&quot;: SequenceParallel(),
}
for block in model.blocks:
parallelize_module(block, mesh[&quot;tp&quot;], tp_plan)
# Step 3. Activation checkpoint (one per block)
for i, block in enumerate(model.blocks):
model.blocks[i] = _ckpt_wrap(block)
# Step 4. PP split (mesh[&quot;pp&quot;])
# Split model into 2 segments (PP=2)
split_spec = {&quot;blocks.16&quot;: SplitPoint.BEGINNING}
pipe = pipeline(
model,
mb_args=(torch.randn(1, 4096, device=&quot;cuda&quot;),),
split_spec=split_spec,
)
pp_stage = pipe.build_stage(
stage_index=mesh[&quot;pp&quot;].get_local_rank(),
device=torch.device(f&quot;cuda:{torch.cuda.current_device()}&quot;),
group=mesh[&quot;pp&quot;].get_group(),
)
# Step 5. FSDP2 wrap (mesh[&quot;fsdp&quot;])
mp_policy = MixedPrecisionPolicy(
param_dtype=torch.bfloat16,
reduce_dtype=torch.float32,
)
for module in pp_stage.submod.modules():
if isinstance(module, TransformerBlock):
fully_shard(module, mesh=mesh[&quot;fsdp&quot;], mp_policy=mp_policy)
fully_shard(pp_stage.submod, mesh=mesh[&quot;fsdp&quot;], mp_policy=mp_policy)
return pp_stage, mesh
def _ckpt_wrap(block):
&quot;&quot;&quot; Wrap a TransformerBlock with activation checkpointing &quot;&quot;&quot;
class Wrapped(nn.Module):
def __init__(self, inner):
super().__init__()
self.inner = inner
def forward(self, *args, **kw):
return checkpoint(self.inner, *args, use_reentrant=False, **kw)
return Wrapped(block)
# ============================================================
# Training loop with PP schedule
# ============================================================
def train_step_4d(pp_stage, schedule, optimizer, batch):
&quot;&quot;&quot; 1 step of 4D parallel training &quot;&quot;&quot;
optimizer.zero_grad(set_to_none=True)
# PP schedule runs forward + backward across all stages
losses = []
schedule.step(batch, losses=losses) # triggers FSDP all-gather / TP all-reduce / etc.
optimizer.step()
return torch.stack(losses).mean()</code></pre>
<p><strong>Note</strong>: the code above is a pedagogical skeleton; for production, use the TorchTitan repo directly (pytorch/torchtitan) — it includes a complete <code>Trainer</code>, checkpointing, profiling, loss / lr schedules. This section only sketches concepts.</p>
<h2 id="b-references">§B References</h2>
<ul><li><strong>DDP</strong>: PyTorch documentation, <code>nn.parallel.DistributedDataParallel</code></li><li><strong>ZeRO</strong>: Rajbhandari et al., "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models", SC 2020 (<a href="https://arxiv.org/abs/1910.02054">arXiv:1910.02054</a>)</li><li><strong>ZeRO-Offload</strong>: Ren et al., USENIX ATC 2021</li><li><strong>ZeRO-Infinity</strong>: Rajbhandari et al., SC 2021</li><li><strong>ZeRO++</strong>: Wang et al., "ZeRO++: Extremely Efficient Collective Communication for Giant Model Training", arXiv:2306.10209, 2023</li><li><strong>FSDP</strong>: Zhao et al., "PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel", VLDB 2023</li><li><strong>FSDP2 / TorchTitan</strong>: Liang et al., "TorchTitan: One-stop PyTorch Native Solution", ICLR 2025 (<a href="https://arxiv.org/abs/2410.06511">arXiv:2410.06511</a>)</li><li><strong>Megatron-LM TP</strong>: Shoeybi et al., "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism", arXiv:1909.08053, 2019</li><li><strong>Megatron-LM-2 (Interleaved 1F1B)</strong>: Narayanan et al., "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM", SC 2021</li><li><strong>GPipe</strong>: Huang et al., "GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism", NeurIPS 2019</li><li><strong>PipeDream / 1F1B</strong>: Narayanan et al., "PipeDream: Generalized Pipeline Parallelism for DNN Training", SOSP 2019</li><li><strong>Sequence Parallel + Selective Recompute</strong>: Korthikanti et al., "Reducing Activation Recomputation in Large Transformer Models", MLSys 2023</li><li><strong>Ring Attention</strong>: Liu et al., "Ring Attention with Blockwise Transformers for Near-Infinite Context", ICLR 2024 (<a href="https://arxiv.org/abs/2310.01889">arXiv:2310.01889</a>)</li><li><strong>Gradient Checkpointing</strong>: Chen et al., "Training Deep Nets with Sublinear Memory Cost", arXiv:1604.06174, 2016</li><li><strong>DualPipe</strong>: DeepSeek-V3 Technical Report, <a href="https://arxiv.org/abs/2412.19437">arXiv:2412.19437</a>, December 2024</li><li><strong>Llama 3</strong>: Meta AI, "The Llama 3 Herd of Models", <a href="https://arxiv.org/abs/2407.21783">arXiv:2407.21783</a>, 2024</li><li><strong>GShard / Switch Transformer (MoE EP)</strong>: Lepikhin et al. 2020 / Fedus et al. JMLR 2022</li></ul>
<footer class="aris-footer">
Generated by <a href="https://github.com/wanshuiyin/Auto-claude-code-research-in-sleep/blob/main/skills/render-html/SKILL.md">ARIS <code>/render-html</code></a> ·
source path <code>docs/tutorials/distributed_training_tutorial_en.md</code> ·
SHA256 <code>9ee6b9a3603b</code> ·
generated at 2026-05-19 19:00 UTC.
This is a generated view — edit the source Markdown, then re-render.
</footer>
</main>
</div>
</body>
</html>
Reference in New Issue View Git Blame Copy Permalink