10 · 数据并行与 AllReduce

# PyTorch DDP：标准写法
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
model = DDP(model.to(local_rank), device_ids=[local_rank])

sampler = torch.utils.data.distributed.DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler, batch_size=micro_bs)

for epoch in range(E):
    sampler.set_epoch(epoch)      # ⭐ 必须！否则每轮的数据划分完全相同
    for batch in loader:
        loss = model(batch)
        loss.backward()           # ⭐ AllReduce 在这里自动发生
        opt.step(); opt.zero_grad(set_to_none=True)

   ❌ 朴素做法：等所有梯度算完，一次性 AllReduce
      → 反向传播时 GPU 在算，网络完全闲置
      → AllReduce 时网络在传，GPU 完全闲置

   ✅ DDP 的做法：
      反向传播是【从后往前】的
      → 最后一层的梯度【最先】算出来
      → 攒够一个"桶"（默认 25MB）就【立刻】开始 AllReduce
      → 通信和后面层的计算【重叠】⭐

   时间轴：
   计算  ████████████████████████
   通信      ▓▓▓▓  ▓▓▓▓  ▓▓▓▓  ▓▓▓▓     ← 藏在计算后面
             ↑ 桶1  ↑桶2  ↑桶3  ↑桶4

# ❌ 每个 micro batch 都 AllReduce → 通信量翻 N 倍
for i, batch in enumerate(loader):
    loss = model(batch) / accum
    loss.backward()

# ✅ 只在最后一步同步
for i, batch in enumerate(loader):
    is_last = (i + 1) % accum == 0
    ctx = nullcontext() if is_last else model.no_sync()    # ⭐
    with ctx:
        (model(batch) / accum).backward()
    if is_last:
        opt.step(); opt.zero_grad(set_to_none=True)

   N 张卡排成环，梯度切成 N 份

   阶段① ReduceScatter（转 N-1 步）：
     每一步，每张卡把手里的一份发给下家、从上家收一份并累加
     → 结束时，第 i 张卡持有【第 i 份的完整求和】

   阶段② AllGather（再转 N-1 步）：
     把各自那份完整结果传遍全环
     → 结束时，所有卡都有完整的求和结果

   ⭐ 每张卡收发的数据量 = 2 × (N-1)/N × 数据量 ≈ 2 × 数据量
      【和卡数几乎无关】

   对比朴素的"都发给 0 号卡"：0 号要收发 N-1 倍 💀

   问题：跨机带宽只有机内 NVLink 的 1/15（第 5 章）
        纯 Ring 会让慢链路成为瓶颈

   ✅ 分层：
   ① 机内 8 卡先做一次 ReduceScatter/AllReduce（NVLink，快）
   ② 各机的代表跨机做一次 AllReduce（慢链路，但数据量只有 1/8）⭐
   ③ 机内 Broadcast/AllGather 回去

   → 跨机通信量降到 1/8

   每步时间 = 计算时间 + max(0, 通信时间 − 可重叠部分)

   ⭐ 卡数增加时：
   · 计算时间 ↓（每卡分到的数据变少）
   · 通信量   ≈ 不变（Ring 的性质）
   → 通信占比【上升】→ 效率下降

   一个具体的例子（7B 模型，BF16 梯度 14GB）：

   通信量 = 2 × 14 GB = 28 GB/步

   机内 NVLink 900GB/s   → 0.03 秒
   跨机 IB 25GB/s        → 1.1 秒   ⭐ 差 35 倍

   如果单步计算是 0.5 秒：
   · 机内：通信可完全重叠 → 扩展效率 ~95%
   · 跨机：通信 1.1 秒 > 计算 0.5 秒 → 【通信成为瓶颈】💀