分布式训练基础#

1. 数据并行（DDP）#

DDP（DistributedDataParallel）是最基本的分布式训练策略：

• 每张 GPU 持有一份完整的模型副本
• 训练数据被均匀分割到各 GPU（通过 DistributedSampler）
• 各 GPU 独立做前向、反向传播
• 反向传播完成后，通过 AllReduce 聚合所有 GPU 的梯度（取平均）
• 各 GPU 用相同的平均梯度更新参数（保证参数同步）

$$g_{\text{avg}} = \frac{1}{N}\sum_{i=1}^{N} g_i \quad \xrightarrow{\text{AllReduce}} \quad \text{每张 GPU 都获得 } g_{\text{avg}}$$

N 张 GPU，每张贡献梯度 g_i，AllReduce 后每张都拿到平均梯度

关键点：DDP 的通信开销主要在 AllReduce 阶段。PyTorch DDP 通过梯度分桶（gradient bucketing）和通信-计算重叠优化：反向传播时一旦某个桶的梯度全部算完就立即开始通信，而不是等所有梯度算完再统一通信。

# PyTorch DDP 最小示例
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
model = MyModel().to(local_rank)
model = DDP(model, device_ids=[local_rank])

# 训练循环和单卡完全相同
for batch in dataloader:
    loss = model(batch).loss
    loss.backward()  # DDP 自动在 backward 中做 AllReduce
    optimizer.step()
    optimizer.zero_grad()

2. ZeRO / FSDP（全分片数据并行）#

DDP 的问题：每张 GPU 存一份完整模型 + 完整优化器状态，对于大模型，单卡放不下。ZeRO（Zero Redundancy Optimizer）的核心思想是分片存储，消除冗余。

2.1 显存组成分析（以 FP16 混合精度 + Adam 为例）#

假设模型有 $\Psi$ 个参数（单位：元素数）：

显存瓶颈

7B 参数模型仅存储参数和优化器状态就需要 112 GB！这还不包括激活值（activation）。而 A100 80GB 单卡显然放不下。DDP 要求每卡一份完整副本，更加不可行。

2.2 ZeRO 三个阶段#

数值示例：7B 模型，8 张 A100：

• DDP：每卡 112 GB → 单卡放不下
• ZeRO-1：每卡 $28 + 84/8 = 38.5$ GB
• ZeRO-2：每卡 $14 + 98/8 = 26.25$ GB
• ZeRO-3：每卡 $112/8 = 14$ GB

以上是理想常驻模型状态，不含激活、参数 AllGather 峰值、通信缓冲与碎片。ZeRO-3 在计算时仍需临时聚合对应参数；通信比值依具体调度与基线而变，不能由常驻存储公式直接推出实际峰值或耗时。ZeRO，§3–4。

2.3 PyTorch FSDP#

FSDP（Fully Sharded Data Parallel）是 PyTorch 原生的 ZeRO-3 实现：

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision

bf16_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

model = FSDP(
    model,
    mixed_precision=bf16_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3
    auto_wrap_policy=transformer_auto_wrap_policy,
)

3. 模型并行#

3.1 Tensor Parallel（张量并行）#

将单个权重矩阵按列或按行切分到多张 GPU 上。例如一个线性层 $Y = XW$，若 $W \in \mathbb{R}^{d \times 4d}$：

• 将 $W$ 按列切为 $W = [W_1, W_2]$（每块 $d \times 2d$）
• GPU 0 计算 $Y_1 = XW_1$，GPU 1 计算 $Y_2 = XW_2$
• 最终 $Y = [Y_1, Y_2]$（拼接）或 AllReduce（取决于下一层的切分方式）

3.2 Pipeline Parallel（流水线并行）#

将模型按层分割到不同 GPU 上：GPU 0 放前 $L/k$ 层，GPU 1 放下 $L/k$ 层……输入数据像"流水线"一样逐级通过各 GPU。

问题：朴素流水线有严重的"气泡"（bubble），即 GPU 空闲等待时间。解决方案：

• GPipe：将 mini-batch 分为多个 micro-batch，减少气泡
• 1F1B（One Forward One Backward）：交替执行前向和反向，进一步减少气泡

并行策略总结

数据并行：模型复制，数据分割 → 扩展 batch size
张量并行：单层内部分割 → 扩展模型宽度
流水线并行：层间分割 → 扩展模型深度
实际训练超大模型（如 LLaMA 70B）通常三种并行组合使用（3D 并行）。

4. DeepSpeed vs FSDP 选择#

选型需固定框架版本、模型结构、offload 与并行需求，分别验证正确性、峰值显存和吞吐；表中的配置入口不构成跨项目的性能优劣结论。

5. 训练显存与激活估算#

总量可按 $M_{\text{train}}=M_{\text{state}}+M_{\text{act}}+M_{\text{temporary}}$ 记账。第 2.1 节的 $16N$ 使用低精度参数与梯度、FP32 主副本及两份 FP32 Adam 状态。若梯度改为 FP32，或常规 autocast 保留 FP32 参数，必须重新加总；$16N$ 不是任意 AMP 的定义。

5.1 激活增长的示意模型#

设层数 $L$、microbatch 大小 $B$、序列长度 $S$、隐藏宽度 $D$、注意力头数 $H$、每元素字节数 $b$。用系数 $a,c$ 表示实际保存张量的常数，粗略模型为：

$$M_{\text{act}}\approx Lb\left(aBSD+cBHS^2\right).$$

它用于区分线性项与显式 attention 矩阵的平方项，不是所有 Transformer 的精确公式。系数随 FFN 宽度、dropout、算子融合和保存策略变化；Transformer 激活分析 §4采用特定模型假设，不能脱离配置搬用。

保留旧笔记的手算设置 $L=32,D=4096,H=32,S=2048,B=1,b=2,a=12,c=2$：线性项为 $6.442$ GB，平方项为 $17.180$ GB，共 $23.622$ GB。这里 GB 按 $10^9$ 字节计；只是所选系数下的示意估算，不是 LLaMA 实测。

FlashAttention 避免保存完整 $S\times S$ attention 矩阵，但仍有其他激活和临时张量，不能直接把线性项当作精确峰值。激活重算的分段模型见重算篇。

5.2 模型状态的量级#

在这一状态约定下，未分片的 7B 全参训练仅状态已超过单卡 80GB；更换优化器、精度或分片后必须重新核算。LoRA 冻结主体可以省掉对应梯度和优化器状态，但仍保留基础权重、adapter 状态与传播所需的激活。

5.3 各技巧减少哪一项#

旧笔记中“7B LoRA 约 22GB、QLoRA 约 8GB”的简单相加没有明确层数、插入模块、长度与量化缓冲，不能作为设备容量结论。规划时先做上述分项估算，再以目标配置的实际峰值验证。