05 — Distributed Papers (Training + Serving at Scale)

Phase alignment: Months 13–18 (Phase 6). Read after your mini-engine works and your Triton FA2 lands.

Half of these are training papers, half are serving papers. Understand both because modern inference-at-scale (RL rollouts, PD-disaggregation, MoE serving) reuses training-scale communication patterns.

Training-scale parallelism (read in order)

1. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

arxiv:<phone_number_or_numberic_id_or_random_id_47> · Shoeybi et al., 2019 · [HARD] · Prereqs: transformer, matrix decompositions

Key trick: Shard the two linears of a transformer block — the attention Q/K/V/O projections and the two FFN matmuls — across N GPUs via matched column-parallel then row-parallel decompositions, adding just two all-reduces per block.

What to extract:

  • Derive the column/row split on paper. Q/K/V column-split, O row-split. FFN up column-split, FFN down row-split. Two all-reduces per block.

  • Attention heads split cleanly: each GPU owns a subset of heads.

  • Why TP prefers NVLink and stays within a node: two all-reduces per layer per token is bandwidth-brutal; must be on the fast fabric.

  • Sequence Parallelism (Megatron-2, arxiv:<phone_number_or_numberic_id_or_random_id_48>): shard LayerNorm and Dropout across sequence dim; overlap comm with compute.

  • Megatron-3 (arxiv:<phone_number_or_numberic_id_or_random_id_49>): expert parallelism + interleaved schedules.

2. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

arxiv:<phone_number_or_numberic_id_or_random_id_50> · Rajbhandari et al., 2019 · [HARD] · Prereqs: Adam, mixed precision

Key trick: Data parallelism replicates the model; ZeRO shards the optimizer states, gradients, and (stage 3) parameters across DP ranks, gathering them just in time for the compute step — near-linear memory scaling.

What to extract:

  • The three stages:

    • ZeRO-1: shard optimizer states (4× memory savings for Adam in fp16).

    • ZeRO-2: also shard gradients.

    • ZeRO-3: also shard parameters (gather-forward, gather-backward). This is what PyTorch FSDP2 implements.

  • Communication cost analysis Table 2. ZeRO-3 doubles the DP bandwidth cost.

  • The offload variants (ZeRO-Offload, ZeRO-Infinity): swap to CPU RAM or NVMe. Enables 20B on a single GPU. Rarely used at inference scale.

  • In modern practice: PyTorch FSDP2 has largely replaced DeepSpeed for open-model training. Read the FSDP2 docs alongside.

3. GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism

arxiv:<phone_number_or_numberic_id_or_random_id_51> · Huang et al., 2018 · [MEDIUM] · Prereqs: transformer training loop

Key trick: Split the model layers into stages across N GPUs; split each mini-batch into M micro-batches; pipeline them so multiple micro-batches are in flight simultaneously, reducing the pipeline bubble to O(N/M).

What to extract:

  • Draw the pipeline schedule (Figure 2). Understand the bubble.

  • Bubble fraction = (N−1)/(N+M−1). To keep it small, M ≫ N.

  • Activation recomputation to make M large without blowing memory.

  • PP vs TP: PP is cross-node friendly (few, coarse comms); TP is intra-node (many, fine comms). Real systems use both.

4. PipeDream / 1F1B scheduling

arxiv:<phone_number_or_numberic_id_or_random_id_52> · Narayanan et al., 2018 (+SOSP ‘19 follow-up) · [MEDIUM] · Prereqs: #3

Key trick: Instead of GPipe’s fill-then-drain schedule, use 1-forward-1-backward: each stage alternates fwd on micro-batch i, bwd on micro-batch i−N+1, keeping activations at steady state.

What to extract:

  • 1F1B halves peak activation memory vs GPipe.

  • Interleaved 1F1B (Megatron-3): each rank owns multiple interleaved stages of the model. Further reduces bubble, at cost of more comms.

  • Zero-bubble scheduling (arxiv:<phone_number_or_numberic_id_or_random_id_53>): splits backward into activation-gradient and weight-gradient phases, exploits the weight-gradient’s ability to be delayed. Currently the state of the art.

5. GShard: Scaling Giant Models with Conditional Computation

arxiv:<phone_number_or_numberic_id_or_random_id_54> · Lepikhin et al., 2020 · [MEDIUM] · Prereqs: transformer

Key trick: Sparse MoE with a top-k router, expert parallelism (each expert lives on one device), and dispatch/combine via all-to-all communication.

What to extract:

  • Top-k routing (usually k=1 or k=2). Load balancing via aux loss.

  • Dispatch (all-to-all send tokens to their assigned experts) → local FFN → combine (all-to-all back). Two all-to-alls per MoE layer.

  • Capacity factor: cap on how many tokens any single expert accepts per batch. Overflow tokens are dropped (or routed to a second choice).

  • Modern serving of MoE (DeepSeek-V3, Mixtral, etc.) is dominated by these mechanics.

6. Switch Transformer: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

arxiv:<phone_number_or_numberic_id_or_random_id_55> · Fedus, Zoph, Shazeer, JMLR ‘22 · [MEDIUM] · Prereqs: #5

Key trick: k=1 routing (“switch”) — send each token to exactly one expert. Simpler, less comms, still competitive.

What to extract:

  • Why k=1: halves communication vs k=2, simpler load balancing, easier auxiliary loss design.

  • Capacity factor tuning: trade drop rate against comms cost.

  • Distillation-to-dense: interesting but tangential to inference.

  • Companion: DeepSeek-MoE (arxiv:<phone_number_or_numberic_id_or_random_id_56>) — fine-grained expert segmentation + shared experts — the design pattern that reappears in V3.

7. Ring Attention with Blockwise Transformers for Near-Infinite Context

arxiv:<phone_number_or_numberic_id_or_random_id_57> · Liu, Zaharia, Abbeel, 2023 · [HARD] · Prereqs: FA (kernels #3–5)

Key trick: Shard the sequence across N devices; each holds its slice of Q/K/V; rotate K/V blocks in a ring so every device sees every K/V slice exactly once and computes its partial attention output. Overlap ring comm with compute.

What to extract:

  • Context-parallelism (CP) is sequence-dimension parallelism.

  • Communication cost: (N−1) block-sized sends per device per attention layer. Overlapped with compute if block is big enough.

  • Enables 1M+ context on hardware that couldn’t hold the KV alone — Gemini/Claude-level context relies on this class of technique.

  • Related: Striped Attention, Blockwise Attention, Sequence Packing for load balance.

Serving-scale distribution

8. DistServe: Disaggregating Prefill and Decoding for Goodput-optimized LLM Serving

arxiv:<phone_number_or_numberic_id_or_random_id_58> · Zhong et al., OSDI ‘24 · [HARD] · Prereqs: engines #2, #3

Key trick: Prefill (compute-bound, latency-sensitive) and decode (memory-bound, throughput-sensitive) have opposite resource profiles and interfere on the same GPU. Split them onto different GPU pools; transfer KV cache from prefill node to decode node once per request.

What to extract:

  • The interference argument: co-located prefill+decode causes ITL spikes (Sarathi/chunked-prefill mitigates but doesn’t eliminate). Full disaggregation removes it.

  • KV transfer cost: must fit in the model’s own decode latency budget. Uses NVLink (intra-node), RDMA/RoCE/InfiniBand (inter-node), sometimes CXL.

  • Optimal split of prefill:decode replicas depends on workload (input:output length ratio). DistServe formulates as an optimization.

  • Goodput = throughput that meets SLO. This paper popularized the metric — memorize it.

9. Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving

arxiv:<phone_number_or_numberic_id_or_random_id_59> · Qin et al., Kimi/Moonshot, FAST ‘25 · Best Paper · [HARD] · Prereqs: #8

Key trick: Elevate the KV cache to a first-class distributed store (Mooncake Store) accessed by prefill nodes, decode nodes, and even other clusters via RDMA. Enables cross-request KV reuse at data-center scale, tiered storage (HBM → DRAM → SSD), and prefill offloading.

What to extract:

  • Prefix cache at fleet scale. Not per-node prefix cache (SGLang RadixAttention) but a shared distributed cache. Hit rates >70% on production chat.

  • Tiered storage: hot KV in HBM, warm in DRAM, cold in SSD. LRU + prefetch on scheduler routing.

  • Prefill offloading: when prefill queue is saturated, ship the request (or its intermediate state) to a spare prefill node with matching prefix.

  • Overload-oriented scheduling: admission control on prefill queue length, KV cache utilization.

  • Kimi’s production numbers: this is not theoretical.

10. LMCache / NIXL / vLLM KV-connector architecture

No single paper. LMCache: https://github.com/LMCache/LMCache · NIXL (NVIDIA Inference Xfer Library): https://github.com/ai-dynamo/nixl · [MEDIUM] · Prereqs: #8, #9

Not a paper, but the reference implementations of PD disaggregation that ship in vLLM V1 and NVIDIA Dynamo. Read the LMCache paper (arxiv:<phone_number_or_numberic_id_or_random_id_60>) and the NIXL docs.

What to extract:

  • The KVConnector abstraction in vLLM: SharedStorageConnector, LMCacheConnector, NixlConnector. Uniform API to “where does KV go/come from.”

  • NIXL: GPUDirect RDMA between prefill and decode, hides transfer under decode.

  • Practical config: --kv-transfer-config in vLLM 0.7+.

11. DeepSeek-V3 Technical Report

arxiv:2412.19437 · DeepSeek-AI, Dec 2024 · [HARD] · Prereqs: papers #1–8, MLA (foundations #11), MoE (#5–6)

Key trick: 671B total / 37B active. MLA + fine-grained MoE with 256 experts + 1 shared + aux-free load balancing (per-expert bias) + MTP + fp8 mixed-precision training — architecture co-designed with inference economics.

What to extract:

  • The fp8 training recipe: per-block scaling, high-precision accumulation, careful loss scaling. Section 3.3. Study this even if you never train in fp8; the accuracy engineering is universal.

  • Aux-free load balancing: instead of an aux loss (which distorts gradients), maintain a per-expert bias updated each step to nudge underutilized experts. Section 2.1.2. Elegant.

  • MTP heads for training-integrated speculative decoding (see engines #8).

  • Cost transparency: the report gives dollar/hour costs for training. 2.788M H800-hours, ~$5.5M. Sets the modern price tag for a frontier model.

  • Inference infra section: wide expert parallelism (32 experts per node), all-to-all optimization (DeepEP), prefill/decode disaggregation.

Companion: the DeepSeek-V2 report (arxiv:2405.04434) for MLA + MoE derivation. Read V2 first if any V3 concept feels ungrounded.

12. DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning

arxiv:<phone_number_or_numberic_id_or_random_id_61> · DeepSeek-AI, Jan 2025 · [MEDIUM] · Prereqs: #11

Key trick: Pure-RL (GRPO) fine-tuning on math/code with rule-based rewards, no SFT bootstrap for R1-Zero, produces long chain-of-thought reasoning from a base model — and generates massive inference workloads (reasoning traces are 10–20× normal output length).

What to extract for inference:

  • Reasoning models make decode 10–20× longer. Everything you learned about long-context decode (FlashDecoding, MLA, KV quant) becomes 10× more important. This is the reason for the 2025 inference-hardware buildout.

  • The RL loop uses inference at massive scale for rollouts — GRPO fires N=16–64 completions per prompt. Rollout throughput is an inference-optimization problem inside training. This is where verl/OpenRLHF pipe vLLM into training.

  • Read Section 2 for GRPO. Skim the rest.

13. Llama-3 report — infrastructure sections (reread)

arxiv:2407.21783 · [HARD] · Prereqs: papers #1–12

Refer back to foundations 01_foundations_papers.md paper #8. Reread Section 3 (Pre-Training) in this phase. You now have the parallelism + NCCL background to understand the 4D parallelism grid, the failure rates, the checkpoint strategy, and the load-balancing decisions.

Awareness-level (read the abstract + skim figures, come back if you deploy this)

  • DeepEP (arxiv:<phone_number_or_numberic_id_or_random_id_62>) — DeepSeek’s high-throughput expert-parallel comm library.

  • PD-Disaggregation in Dynamo / llm-d — NVIDIA’s productionized version.

  • Star Attention / Mixture of Depths / Jamba — alternative attention/architecture families you should recognize by name.

  • RingFlashAttention (implementation of Ring Attention on top of FA2) — for the day you serve 1M-context.

  • Verl / OpenRLHF / TRL — the RL trainers-with-embedded-inference-engine. Study one seriously in Phase 6 (verl is the current standard).

The whiteboard test

At the end of Phase 6 you should be able to draw:

  1. A 4D parallelism grid (TP=8, PP=4, DP=2, EP=8) on 512 H100s. Label which parallelism uses NVLink vs IB. Compute the bytes/step for each all-reduce/all-to-all.

  2. A disaggregated PD deployment (say 4 prefill nodes + 8 decode nodes, each with its own TP degree). Label the KV transfer path and its bandwidth requirement.

  3. The DeepSeek-V3 MoE dispatch pattern for one MoE layer across 32 nodes.

  4. The GPipe bubble and how 1F1B, interleaved, and zero-bubble each shrink it.

  5. Ring attention — 4 devices, block rotation schedule.

If you can draw these five diagrams, you can talk to any staff-level infra person on the planet.