12. Distributed Training: Parallel Engineering for LLMs

Chapter focus: This chapter teaches distributed training through minimal reproducible experiments: first understand NCCL communication and autograd with hand-written code, then progressively enter tensor parallelism, Flash Attention, mixed precision and checkpoints, ending with loss debugging and network tuning in real training.

🧭 Reading Map

• Prerequisites: 11. Pre-trained Models + 03. Network — parallelism is bounded by both model structure and network ceiling
• Role: core chapter of Part 4, L4 — 3D parallelism, hands-on LLM Lab, mixed precision/NCCL tuning and debugging
• Order: parallelism principles → LLM Lab 01→07 hands-on (DDP→TP→Megatron→FlashAttention) → tuning & debugging
• Next: 13. Inference Engines — how to serve the trained model efficiently

Articles

Distributed Parallelism Overview: From Single GPU to 10K

Partition logic and selection decisions for data, tensor, pipeline, sequence parallelism, ZeRO and MoE expert parallelism, with typical 7B / 70B / 1000-GPU LLM training configurations and communication cost estimation.

阅读文章 →

LLM Lab Notes 01: Distributed Training 101 — From init_process_group to all_reduce

Walk through an 18-line NCCL multi-GPU communication experiment: process group initialization, the rank/local_rank distinction, all_reduce collective semantics, and the first set of distributed training debugging lessons.

阅读文章 →

LLM Lab Notes 02: DDP Minimal Training — Multi-GPU in One Line

Wrap a Linear model with DistributedDataParallel and run a 5-step multi-GPU training loop. Break down DDP's parameter broadcast, gradient synchronization, and optimizer update mechanics — what one line of wrapper code actually does.

阅读文章 →

LLM Lab Notes 03: Writing Collective Communication Autograd — How AllReduce Joins Backpropagation

Wrap all_reduce in a custom torch.autograd.Function, verify that collective communication runs in both forward and backward, and derive the theoretical basis for Row Parallel gradient synchronization in tensor parallelism.

阅读文章 →

LLM Lab Notes 04: Tensor Parallelism 101 — Column/Row Parallel and the Megatron FFN

Three minimal experiments on tensor parallelism: Column Parallel (split columns + AllGather), Row Parallel (split rows + AllReduce), and a reproduction of Megatron-LM's two-layer FFN split, with communication cost analysis.

阅读文章 →

LLM Lab Notes 05: Attention Tensor Parallelism — Splitting QKV by Head, Megatron Style

Multi-head attention is naturally parallelizable by head: split QKV weights by head count, compute attention locally, AllGather to reassemble the full output. A minimal implementation and communication-point analysis of Megatron-LM's Self-Attention TP.

阅读文章 →

LLM Lab Notes 06: Building a Megatron-Style Transformer Block from Scratch

Assemble LayerNorm, multi-head attention, FFN, residual connections and an AllReduce autograd function into a trainable, segment-profiled Megatron-style Transformer Block, with per-segment timing for Attention/FFN/Backward.

阅读文章 →

LLM Lab Notes 07: Why FlashAttention Is Fast — An On-Box Comparison

Compare standard Attention (SDPA) against flash_attn on the same GPU in terms of latency and memory behavior. Deconstruct FlashAttention's three designs — tiling, IO-awareness, online softmax — to build first-hand knowledge of long-context bottlenecks.

阅读文章 →

Loss Spike Debugging: Locating Training Divergence

A systematic framework for loss spikes, divergence, and creeping loss during training: four layers (data, numeric, communication, hardware), with tools and a production case.

阅读文章 →

Mixed Precision Training & Checkpointing: Stable and Recoverable

BF16/FP16/FP8 mixed precision principles and configuration, checkpoint content design and async saving, resume workflow and consistency verification, with a production checklist.

阅读文章 →

NCCL Tuning in Practice: Locating Communication Bottlenecks in a 10K-GPU Cluster

A complete tuning path from nccl-tests baselines and NCCL_* environment variables to InfiniBand-side PFC/ECN, ECMP hashing and topology congestion, with a real 32-node AllReduce bottleneck case study.

阅读文章 →

ZeRO & FP8 Training Ledger: Sharding Memory

Four-state ledger, ZeRO-1/2/3 savings, combining with TP/PP/DP, FP8 gains and precision management.

Read →

Post-Training Ledger: SFT, LoRA & RL

Full vs LoRA, PPO/DPO/GRPO comparison, three evaluation gates, with engineering tables and reproducible Python.

Read →

## Suggested Reading Path

Run LLM Lab 01-03 first to build communication intuition, then tensor parallelism (04-06) and Flash Attention (07) as needed; production training starts with mixed precision/checkpoints, loss debugging, and nccl tuning.

Key Takeaways

The "minimal unit" of distributed training is NCCL collectives plus autograd; parallel-strategy selection is driven by three questions: does it fit in memory, what is the communication share, and is the load balanced?