LLM Interview / Distributed Training: 8 Illustrated Articles

Distributed training section in LLM interviews. AiGC Interview Guide illustrated series, 9 PDFs.

1. Overview of Distributed Training (file_44)

Why do we need distributed training?

Single-GPU bottlenecks	Solutions
Insufficient GPU memory — A 65B model in fp16 is 130GB, which cannot fit in an 80GB A100	Model Parallelism / ZeRO
Computation takes too long — Pre-training a 13B model on a single GPU takes 100+ days	Data Parallelism
Too small a batch is unstable — Batch size = 8 on a single GPU has high noise	Distributed batch scaling
Multi-machine coordination — Inter-node communication	NCCL / Gloo / MPI

4 major parallelism strategies

Strategy	Idea	Communication
Data Parallel (DP)	Each GPU holds a full copy of the model, processing different data	AllReduce gradients
Model Parallel (MP)	Different parts of the model reside on different GPUs	Forward / backward activations
Pipeline Parallel (PP)	Different layers on different GPUs, executed in a pipeline	Inter-layer activations
Tensor Parallel (TP)	Single-layer tensors split across multiple GPUs	Intra-layer AllReduce

Practical combination (3D Parallelism)

Used by Megatron-Turing NLG 530B:

TP=8 (within a single 8-GPU machine)
PP=35 (across machines)
DP=replacement

2. Pipeline Parallelism (file_45)

Pipeline idea

模型: Layer1 → Layer2 → Layer3 → Layer4
分到 4 个 GPU:
GPU0: Layer1
GPU1: Layer2
GPU2: Layer3
GPU3: Layer4

mini-batch 切成 micro-batch:
  micro1, micro2, micro3, micro4

执行流水线:
  t=1: GPU0: micro1
  t=2: GPU0: micro2 | GPU1: micro1
  t=3: GPU0: micro3 | GPU1: micro2 | GPU2: micro1
  ...

Pain points of Pipeline parallelism

Bubble — GPUs sit idle waiting for data
Solutions:1F1B (one-forward-one-backward) scheduling / interleaved 1F1B

Mainstream PP implementations

GPipe (Google) — synchronous pipeline
PipeDream (Microsoft) — asynchronous pipeline
Megatron-LM PP — 1F1B + interleaved
DeepSpeed PP — integrates ZeRO

3. PyTorch Data Parallel Series

file_46 nn.DataParallel (DP, outdated but simple)

model = nn.DataParallel(model)

Features:

Single-process, multi-threaded
Limited by the GIL lock
High memory pressure on the primary GPU(aggregates all gradients)
Not recommended(outdated)

file_47 nn.parallel.DistributedDataParallel (DDP, mainstream)

import torch.distributed as dist
dist.init_process_group(backend='nccl')
model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

Features:

Multi-process, multi-GPU
Each GPU runs independently
Efficient gradient synchronization via AllReduce
Mainstream solution

file_48 torch.multiprocessing

import torch.multiprocessing as mp
mp.spawn(train_func, args=(world_size,), nprocs=world_size)

Purpose: launches multi-process training tasks, used in combination with DDP.

4. Mixed Precision Training AMP (file_49)

Why do we need mixed precision?

fp32 is slow and uses a large amount of GPU memory
fp16 is fast and memory-efficient
But fp16 has a small numeric range, prone to overflow / underflow

AMP idea

1. fp32 主权重保留
2. fp16 副本做前向 / 反向 / 优化器
3. Loss Scaling:fp16 loss × 大数 防下溢
4. 梯度反向时再除回

PyTorch AMP usage

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for data, target in loader:
    optimizer.zero_grad()
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

5. DeepSpeed (file_50)

What is DeepSpeed

Developed by Microsoft, a large-scale training framework that adds a layer on top of PyTorch.

Core features

Feature	Function
ZeRO	Partitioning of optimizer states / gradients / parameters
DeepSpeed Inference	Efficient inference
DeepSpeed-Chat	End-to-end RLHF
DeepSpeed-MII	Model inference
DeepSpeed Compression	Model compression
Pipeline Parallel	Pipeline parallelism
Mixture of Experts (MoE)	Sparse expert networks

Relationship with PyTorch

DeepSpeed is a "superset" of PyTorch
When using DeepSpeed, the model remains a PyTorch model, and DeepSpeed takes over the training loop

6. accelerate (file_51)

What is accelerate

Developed by Hugging Face,A lightweight library that simplifies distributed training.

Advantages of accelerate

The same code can run on single GPU / multi-GPU / TPU / DeepSpeed / FSDP run
Minimal code changes: just wrap the model / optimizer / dataloader

accelerate usage

from accelerate import Accelerator

accelerator = Accelerator()
model, optimizer, dataloader = accelerator.prepare(
    model, optimizer, dataloader
)
for batch in dataloader:
    output = model(batch)
    loss = output.loss
    accelerator.backward(loss)
    optimizer.step()

Launch:accelerate launch train.py(automatically detects environment configuration)

7. ZeRO Learning (file_52)

What is ZeRO

Zero Redundancy Optimizer — the core innovation of DeepSpeed.

Problem: In DDP, each GPU retains a full model copy + Adam states (8x the parameter size) + gradients, causing memory explosion.

ZeRO idea:Partition these across GPUs for storage, and communicate when needed.

3 stages of ZeRO

Stage	Partitioned content	Memory savings
ZeRO-1	Optimizer state partitioning	~4x
ZeRO-2	+ Gradient partitioning	~8x
ZeRO-3	+ Parameter partitioning	~Nx (N = number of GPUs)

ZeRO-Offload

Offload Adam states + gradients to CPU CPU
A single GPU can train 13B+ models
Tradeoff: CPU↔GPU transfer latency

ZeRO-Infinity

Offload to NVMe NVMe SSD
Train a 1T-parameter model on 1 V100(exaggerated but possible)

ZeRO communication overhead

Stage	Communication volume	Speed impact
ZeRO-1	Comparable to DDP	Almost no impact
ZeRO-2	Slightly more	10-20% slower
ZeRO-3	Significantly more (parameter AllGather)	30-50% slower(but enables training large models)

8. Comprehensive Comparison

Mainstream solution recommendations

Scenario	Recommendation
Learning entry	PyTorch DDP
Rapid experimentation	accelerate
Large models (>10B)	DeepSpeed ZeRO-2/3
Ultra-large models (>100B)	Megatron-LM (TP+PP+DP)
TPU	JAX + Flax

NVIDIA recommended stack

Megatron-LM (after Megatron-Core) + Apex / FlashAttention / Transformer Engine.

Industry production stack

OpenAI: in-house development + Triton
Anthropic: JAX
Meta: PyTorch + FSDP + Megatron
Microsoft: DeepSpeed
NVIDIA: NeMo + Megatron-LM
Major domestic (Chinese) companies: DeepSpeed + in-house development

High-frequency exam point priority

🔥🔥🔥 Three stages of ZeRO — Must-know
🔥🔥🔥 DP vs DDP — Must-know
🔥🔥 Pipeline Parallelism + 1F1B
🔥🔥 AMP Mixed Precision + Loss Scaling
🔥🔥 3D Parallelism (TP+PP+DP)
🔥 accelerate vs DeepSpeed
🔥 ZeRO-Offload / ZeRO-Infinity