AI/ML 详解：训练基础设施

一个 70B 模型单卡装不下，就像一张表太大单机存不下——你的本能是分片（sharding）。但训练有三个正交的切分维度：切数据（像多个只读副本）、切模型的层（像微服务流水线）、切单层里的大矩阵（像把一个大 JOIN 拆到多机）。理解分布式训练，本质是理解这三个轴怎么组合。

痛点是显存（VRAM），不是算力。一个 70B 模型用 fp16，光参数就 140GB；加上梯度和 Adam 优化器状态，训练时每个参数要占 ~16 字节（下张卡的 ZeRO 会拆解这个数字），总共 1TB 出头——一张 80GB 的 H100 根本放不下。三种并行各解决一块：

• 数据并行 Data Parallel（DP）：每张卡存一份完整模型，只把训练 batch 切开。各卡独立前向反向，再用 all-reduce 把梯度求平均同步。类比：多个只读副本各处理一批请求，定期对账。前提是模型本身塞得进单卡；
• 流水线并行 Pipeline Parallel（PP）：把模型按层切成几段，分到不同卡，像装配线 / CPU 流水线。代价是 bubble（气泡）——第 1 段算完在等第 2 段，靠把 batch 拆成 micro-batch 填满流水线来缓解；
• 张量并行 Tensor Parallel（TP）：把单层内的大矩阵乘法横竖切到多卡（第 4 张卡详解）。通信最密集，只在一台机器的高速 NVLink 内做。

# 最简单的数据并行：PyTorch DDP。用 torchrun 启动多进程
import torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")        # NCCL = NVIDIA 的 GPU 集合通信库
rank = dist.get_rank()
torch.cuda.set_device(rank)

model = MyTransformer().to(rank)
model = DDP(model, device_ids=[rank])  # 反向时自动 all-reduce 梯度

for batch in shard_loader(rank):     # 每个 rank 拿不同数据切片
    loss = model(batch).loss
    loss.backward()                  # ← 梯度在这一步跨卡求平均
    optimizer.step(); optimizer.zero_grad()
# 启动：torchrun --nproc_per_node=8 train.py

纯数据并行有个巨大浪费：每张卡都存一份完整的参数、梯度、优化器状态——这是 N 份冗余副本，等于分布式存储里"全量复制到每个节点"。ZeRO 的洞察是：这些状态大部分时间用不到，干嘛每张卡都存满？于是把它们切片（shard）到各卡，每卡只存 1/N，用到时临时 all-gather 拼回来。本质是用通信换显存——和用 erasure coding 替代全量复制是同一个 trade-off。

先把"每参数 16 字节"拆开看——这是理解 ZeRO 的关键。混合精度训练里，每个参数要存：

ZeRO 分三级，渐进式地切掉冗余，省得越多、通信越重：

• ZeRO-1：切优化器状态（那 12 字节）。最肥的部分先分片，几乎不增加通信，性价比最高；
• ZeRO-2：再切梯度。反向算完即把梯度 reduce-scatter 到负责的卡，本地只留自己那片；
• ZeRO-3：连参数本身也切。前向/反向走到某层时才 all-gather 出完整参数，用完立刻丢弃。每卡显存降到约 16/N 字节/参数——理论上加卡就能线性降显存。

这是微软 DeepSpeed 的 ZeRO（Rajbhandari et al. 2019, arXiv 1910.02054）。它最优雅的地方：在数据并行的编程模型下（每卡看到的还是"完整模型"的假象），偷偷把内存冗余消掉了。

# DeepSpeed 用一份 JSON config 开启 ZeRO，业务代码几乎不变
ds_config = {
    "train_micro_batch_size_per_gpu": 4,
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,                 # 切 优化器+梯度+参数
        "offload_optimizer": {"device": "cpu"},  # 还能下放到 CPU 内存
        "stage3_prefetch_bucket_size": 5e7        # 预取，藏通信延迟
    }
}
import deepspeed
engine, _, _, _ = deepspeed.initialize(
    model=model, config=ds_config, model_parameters=model.parameters())
for batch in loader:
    loss = engine(batch).loss
    engine.backward(loss)   # ZeRO 在这里做 reduce-scatter + all-gather
    engine.step()

FSDP 就是 PyTorch 官方把 ZeRO-3 写进了标准库。机制几乎等价，区别像第三方库 vs 内置实现：DeepSpeed/ZeRO 是外挂的（类似 requests），FSDP 是和 PyTorch 的 autograd、dispatcher、显存分配器深度同源的官方版本（类似内置的 urllib，但更顺手）。对大多数纯 PyTorch 项目，FSDP 是"不引入新生态就能拿到 ZeRO-3"的默认选择。

解决的痛点和 ZeRO-3 一样——消除数据并行的内存冗余——但要无缝融进原生 PyTorch。机制可以一句话概括：平时分片、用时聚合、用完释放。

• 包裹（wrap）成单元：把模型按层/模块切成若干 FSDP unit。平时每个 unit 的参数只在本卡留 1/N 分片；
• 前向到某 unit：先 all-gather 拼出这一层的完整参数 → 算这一层 → 立刻 free 掉非本地的那部分，只留回 1/N；
• 反向同理：再 all-gather 一次算梯度，然后 reduce-scatter 把梯度切片分发回各自负责的卡；
• 通信-计算重叠：算第 i 层时，后台预取第 i+1 层的参数，把 all-gather 的延迟藏在计算背后——这是它能跑快的关键。

权威参考是 PyTorch FSDP 论文（Zhao et al. 2023, arXiv 2304.11277），里面讲清了它怎么和 PyTorch 核心组件协同设计。auto_wrap_policy 决定切分粒度：unit 切太细 → all-gather 次数多、通信碎；切太粗 → 单次聚合的峰值显存高。粒度本身就是显存与通信的连续权衡。

# PyTorch 原生 FSDP，无需第三方库
import torch, functools
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

# 按 Transformer block 自动切分单元——粒度决定通信/显存权衡
policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock})

model = FSDP(
    model,
    auto_wrap_policy=policy,
    device_id=torch.cuda.current_device())

# 之后的训练循环和单卡几乎一模一样——分片对业务代码透明
for batch in loader:
    loss = model(batch).loss
    loss.backward()      # FSDP 在此处理 all-gather / reduce-scatter
    optimizer.step(); optimizer.zero_grad()

ZeRO 和 FSDP 都假设单层算得动——它们切的是"谁存参数"，但算的时候还是把整层拼回单卡。可如果单层的矩阵本身就大到一张卡放不下呢？那就得把单个矩阵乘法本身切开，让多张卡协同算一个 matmul。类比：一行记录大到单机存不下，你只能把这一行的字段也拆到多机（垂直拆单个算子）。这就是张量并行（Tensor Parallel），Megatron-LM 是它的经典实现。

Megatron-LM（Shoeybi et al. 2019, arXiv 1909.08053）的精妙之处：怎么切 Transformer 的两大算子，让切分边界刚好落在非线性函数之外，从而一次前向只需一次同步。看 MLP 层 Y = GeLU(X·A)·B：

• 第一个矩阵 A 按「列」切成 [A₁, A₂]：X·A₁ 和 X·A₂ 分到两卡独立算。因为 GeLU 是逐元素的，GeLU(X·A₁) 不需要看 A₂ 的结果——非线性可以在各卡本地完成，无需中间同步；
• 第二个矩阵 B 按「行」切成 [B₁; B₂]：两卡各算出部分和，最后用一次 all-reduce 求和得到完整 Y。

为什么是"列切 A、行切 B"而不是反过来？因为这个组合让唯一的非线性（GeLU）落在切分内部，整层前向只在末尾同步一次。换个切法就得在 GeLU 前后各同步一次——通信翻倍。自注意力则更自然：按 attention head 切，每个 head 本就独立计算，天生适合并行。

张量并行通信极重（每层都要 all-reduce），所以只在一台机器内的 NVLink 上做（带宽 ~900GB/s），绝不跨节点。最终的万卡训练是 3D 并行：

# 张量并行的核心思想，用纯 PyTorch 示意「列切 + 行切 + all-reduce」
import torch, torch.distributed as dist
import torch.nn.functional as F

def tp_mlp(X, A_shard, B_shard):
    # A 按列切：本卡只持有 A 的一部分列 → 算出部分隐藏维
    H = F.gelu(X @ A_shard)        # GeLU 逐元素，本地完成，无需同步
    # B 按行切：本卡持有对应的行 → 算出部分和
    Y_partial = H @ B_shard
    # 唯一一次同步：把各卡部分和加起来 = 完整输出
    dist.all_reduce(Y_partial, op=dist.ReduceOp.SUM)
    return Y_partial            # 整层前向只通信了一次
# 真实生产用 NVIDIA Megatron-LM / Megatron-Core，已封装好这套切分

• ZeRO 论文 (Rajbhandari et al. 2019) — 内存冗余分析与三级分片的奠基性工作
• Megatron-LM 论文 (Shoeybi et al. 2019) — 张量并行怎么切 Transformer，必读
• PyTorch FSDP 论文 (Zhao et al. 2023) — FSDP 与 PyTorch 核心的协同设计
• GPipe 论文 (Huang et al. 2018) — 流水线并行与 micro-batch、bubble 的原始论文
• DeepSpeed ZeRO 官方教程 — config 用法与各 stage 的工程取舍
• PyTorch FSDP 官方教程 — auto_wrap_policy、state_dict 等实操