LLM 面试 / 分布式训练 8 篇图解专题

LLM 面试中分布式训练板块。AiGC 面试宝典图解系列,9 份 PDF。

一、分布式训练总览(file_44)

为什么需要分布式训练?

单卡瓶颈	解法
显存不够 — 65B 模型 fp16 = 130GB,A100 80GB 装不下	模型并行 / ZeRO
计算时间太长 — 13B 模型预训练单卡 100+ 天	数据并行
batch 太小不稳 — 单卡 batch=8 噪声大	分布式扩 batch
多机协调 — 节点间通信	NCCL / Gloo / MPI

4 大并行策略

策略	思路	通信
Data Parallel (DP)	各 GPU 复制完整模型,分不同数据	AllReduce 梯度
Model Parallel (MP)	模型不同部分在不同 GPU	前向 / 反向激活
Pipeline Parallel (PP)	不同层在不同 GPU,流水线执行	层间 activation
Tensor Parallel (TP)	单层 tensor 切分多 GPU	层内 AllReduce

实际组合(3D Parallelism)

Megatron-Turing NLG 530B 用的:

TP=8(单机 8 卡内)
PP=35(跨机)
DP=替换

二、流水线并行(file_45)

流水线思想

模型: 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
  ...

Pipeline 流水线痛点

气泡(Bubble) — GPU 等数据空闲
解决:1F1B (one-forward-one-backward) 调度 / interleaved 1F1B

主流 PP 实现

GPipe (Google) — 同步流水
PipeDream (Microsoft) — 异步流水
Megatron-LM PP — 1F1B + interleaved
DeepSpeed PP — 集成 ZeRO

三、PyTorch 数据并行系列

file_46 nn.DataParallel(DP,过时但简单)

model = nn.DataParallel(model)

特点:

单进程多线程
GIL 锁限制
主 GPU 显存压力大(聚合所有梯度)
不推荐(已过时)

file_47 nn.parallel.DistributedDataParallel(DDP,主流)

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

特点:

多进程多 GPU
每 GPU 独立运行
AllReduce 高效同步梯度
主流方案

file_48 torch.multiprocessing

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

用途:启动多进程训练任务,与 DDP 结合。

四、混合精度训练 AMP(file_49)

为什么需要混合精度?

fp32 慢且占显存大
fp16 快且省显存
但 fp16 数值范围小,容易溢出 / 下溢

AMP 思路

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

PyTorch AMP 用法

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()

五、DeepSpeed(file_50)

DeepSpeed 是什么

Microsoft 出品,大规模训练框架,在 PyTorch 之上加层。

核心特性

特性	作用
ZeRO	优化器状态 / 梯度 / 参数分散
DeepSpeed Inference	高效推理
DeepSpeed-Chat	一站式 RLHF
DeepSpeed-MII	模型推理
DeepSpeed Compression	模型压缩
Pipeline Parallel	流水线并行
Mixture of Experts (MoE)	稀疏专家网络

与 PyTorch 关系

DeepSpeed 是 PyTorch 的"超集"
用 DeepSpeed 时,模型仍是 PyTorch 模型,DeepSpeed 接管训练 loop

六、accelerate(file_51)

accelerate 是什么

Hugging Face 出品,简化分布式训练的轻量库。

accelerate 优势

同一份代码可在单 GPU / 多 GPU / TPU / DeepSpeed / FSDP 跑
极简改动:把 model / optimizer / dataloader 包一层即可

accelerate 用法

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()

启动:accelerate launch train.py(自动检测环境配置)

七、ZeRO 学习(file_52)

ZeRO 是什么

Zero Redundancy Optimizer — DeepSpeed 核心创新。

问题:DDP 每 GPU 保留完整模型副本 + Adam 状态(8x 参数大小)+ 梯度,显存爆炸。

ZeRO 思路:把这些跨 GPU 分散保存,需要时通信。

ZeRO 3 个阶段

阶段	分散内容	显存节省
ZeRO-1	优化器状态分散	~4x
ZeRO-2	+ 梯度分散	~8x
ZeRO-3	+ 参数分散	~Nx (N=GPU 数)

ZeRO-Offload

把 Adam 状态 + 梯度 offload 到 CPU
单卡也能训 13B+ 模型
代价:CPU↔GPU 传输延迟

ZeRO-Infinity

offload 到 NVMe SSD
1 张 V100 训 1T 参数模型(夸张但可能)

ZeRO 通信开销

Stage	通信量	速度影响
ZeRO-1	与 DDP 相当	几乎无影响
ZeRO-2	略多	10-20% 慢
ZeRO-3	显著多(参数 AllGather)	30-50% 慢(但能训大模型)

八、综合对比

主流方案推荐

场景	推荐
学习入门	PyTorch DDP
快速实验	accelerate
大模型(>10B)	DeepSpeed ZeRO-2/3
超大模型(>100B)	Megatron-LM (TP+PP+DP)
TPU	JAX + Flax

NVIDIA 推荐栈

Megatron-LM (Megatron-Core 后) + Apex / FlashAttention / Transformer Engine。

业界生产栈

OpenAI: 自研 + Triton
Anthropic: JAX
Meta: PyTorch + FSDP + Megatron
Microsoft: DeepSpeed
NVIDIA: NeMo + Megatron-LM
国内大厂: DeepSpeed + 自研

高频考点优先级

🔥🔥🔥 ZeRO 三阶段 — 必考
🔥🔥🔥 DP vs DDP — 必考
🔥🔥 流水线并行 + 1F1B
🔥🔥 AMP 混合精度 + Loss Scaling
🔥🔥 3D Parallelism (TP+PP+DP)
🔥 accelerate vs DeepSpeed
🔥 ZeRO-Offload / ZeRO-Infinity