09 分布式训练框架

资源层和调度层把「卡」管好了。这一章进入负载层：训练任务本身怎么写、怎么把一个模型切到多卡上跑，以及在 K8s 上用什么编排、多机之间靠什么通信。这是 AI Infra 工程师和算法团队协作的接口层，也是通向 12 平台实战的桥。

💡 给你的话：这章是 AI 味最浓的一章。你不需要会手写训练代码，但要懂「模型怎么切、多机怎么通信、为什么慢」——这样算法方说「训练扩展性差」时你能定位是调度问题还是通信问题。

---

9.1 训练任务本质（回顾 01）

一次训练 = 前向 → 算 Loss → 反向（算梯度） → 更新参数，反复迭代。全是矩阵运算，吃显存（参数+梯度+优化器+激活）。

⭐ 显存红线（已核实）：7B 全参微调 ≈ 112 GB（16 字节/参数），单张 A100 80G 放不下，必须 ZeRO-3/FSDP 或多卡；想单卡就上 LoRA。详见 01/03。

---

9.2 PyTorch 训练基本流程

import torch
from torch import nn

model = MyModel().cuda()          # 模型放 GPU
optimizer = torch.optim.AdamW(model.parameters())
loss_fn = nn.CrossEntropyLoss()

for x, y in dataloader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad()
    out = model(x)                # 前向
    loss = loss_fn(out, y)
    loss.backward()               # 反向（算梯度）
    optimizer.step()              # 更新参数

名词解释：

• PyTorch：当下最主流的深度学习框架（Meta 开源）。.cuda() 把数据/模型搬到 GPU。
• AdamW：一种优化器（见 01），带权重衰减的 Adam。
• CrossEntropyLoss：交叉熵损失函数，分类任务常用。
• DataLoader：PyTorch 里加载数据的迭代器。

你作为 AI Infra 工程师不一定要能手写训练循环，但要懂：模型/数据怎么上 GPU、显存里放了什么、为什么多卡要通信。

---

9.3 为什么需要分布式训练

• 放不下：大模型参数+梯度+优化器超单卡显存（如 70B 微调需多卡）。
• 跑太慢：数据量大，单卡训练数月。
• 解法：多卡/多机并行。

---

9.4 并行策略（⭐核心知识）

数据并行（Data Parallelism）

每张卡都有一份完整模型，把数据切分给各卡分别算，反向后同步梯度。

方式	说明	特点
DP（旧）	单机多卡，Parameter Server 模式	已淘汰
DDP（Distributed Data Parallel）	每卡一份完整模型副本，梯度 AllReduce 同步	主流、易用
FSDP（Fully Sharded Data Parallel）	参数分片到各卡（ZeRO 思想）	省显存，能跑更大模型

名词解释：

• Parameter Server（参数服务器）：一种老的分布式训练架构，有专门节点存参数。已被 AllReduce 架构取代。
• AllReduce：一种集合通信原语，所有卡把各自梯度求和并广播给所有人（见 9.6）。

DDP：每张卡有完整模型，batch 切分，反向后梯度 AllReduce 求平均 → 等效大 batch。 FSDP：模型参数也分片，用时 AllGather，显存最优。

流水线并行（Pipeline Parallelism, PP）

• 模型按层切到不同卡（前几层卡 A，后几层卡 B）。
• micro-batch 流水执行，提升设备利用率。
  • 名词解释：micro-batch：把一个 batch 再切成小份，让流水线各阶段轮流处理，减少空等。
• 难点：气泡（bubble）——流水线各阶段空等的时间。

张量并行（Tensor Parallelism, TP）

• 单层的矩阵乘法切分（列并行/行并行）。
• 需要卡间高带宽（NVLink）—— Megatron 的核心。
• 适合超大模型单层都放不下。

💡 类比：数据并行是「8 个人各抄一份作业的不同题」；张量并行是「一道大题拆成 8 块各算一块再拼」；流水线并行是「流水线 8 道工序，每人负责一段」。

ZeRO（DeepSpeed）

• 把优化器状态 / 梯度 / 参数分片到各卡，三阶段（ZeRO-1/2/3）。
• 显存随卡数近线性下降，是大模型训练标配。
• DeepSpeed 库实现，微软开源。
• ZeRO-3 = 全分片，等同 FSDP 思想。

---

9.5 Megatron-LM（大模型训练事实标准）

• NVIDIA 开源，结合 TP + PP + DP 三维并行。
• 把超大规模模型（如千亿）切到千卡集群。
• 你的训练任务若上规模，底层多半是 Megatron / DeepSpeed + 调度器。

⭐ 重点：算法方写模型，AI Infra 方负责「把 Megatron/DeepSpeed 任务以 Gang Job 形式跑在 K8s 多卡上、保证通信拓扑最优」——这是你工作的真实切面。

---

9.6 集群互联与集合通信（⭐AI Infra 核心难点）

多卡好办，多机才是难点。卡间/机间怎么通信，直接决定扩展效率。

互联层级

层级	技术	带宽量级	说明
卡间（同节点）	NVLink / NVSwitch	数百 GB/s~TB/s	最快；NVSwitch 做节点内全互联交叉
机间（同机房）	InfiniBand (IB) / RoCE	数十~数百 GB/s	IB 低延迟高带宽，AI 集群主流
机间（以太网）	普通以太网 / 以太网+拥塞控制	较低	成本低，需优化才能跑训练

名词解释：

• InfiniBand (IB)：Mellanox（被 NVIDIA 收购）的高速互联协议，低延迟高带宽，AI 集群机间主流。
• RoCE（RDMA over Converged Ethernet）：在以太网上跑 RDMA 的协议，比普通以太网快、比 IB 便宜。
• RDMA（Remote Direct Memory Access）：远程直接内存访问，绕过 CPU 直接读写远端内存，延迟极低。

经验：节点内走 NVLink，跨节点走 IB/RoCE；用普通以太网跑大规模训练会严重掉速。

NCCL 与集合通信原语

• NCCL（NVIDIA Collective Communications Library） 是 GPU 间通信的事实标准库，PyTorch distributed 默认后端。
• 关键集合通信原语：
  • AllReduce：所有卡梯度求和并广播（DDP 核心）。
  • AllGather：收集所有卡的分片（FSDP/ZeRO-3 取参数）。
  • ReduceScatter：先求和再分散（ZeRO 分片梯度）。
  • Broadcast / Scatter / Gather：点对多/多对一。
• 拓扑感知：NCCL 会探测 NVLink/IB 拓扑，优先走最快路径；AI Infra 要配置好 NCCL_* 环境变量（如 NCCL_IB_DISABLE、NCCL_TOPO_FILE）做调优。
• 通信-计算 overlap：用 torch.distributed 的异步通信 + 梯度分桶（bucketing），让通信和计算重叠，隐藏延迟——这是训练吞吐优化的关键。

⭐ 重点：你会回答「为什么多机训练慢」时，先查通信是不是瓶颈（NCCL 日志、IB 利用率、是否走了以太网而非 IB）。这是 GPU 利用率低的第二大常见原因（第一是显存/数据搬运）。

⚠️ 内幕：大规模训练「扩展效率」往往卡在通信而非算力。千卡训练时通信开销可能占 30-50% 时间。这就是为什么 NVIDIA 拼命推 NVSwitch/NVLink 和 IB——互联是命脉。Ultra Ethernet Consortium（超以太网联盟） 正在搞开源高速以太网替代 IB，因为 IB 被 NVIDIA 垄断且贵。

---

9.7 在 K8s 上编排训练（呼应前面所有章节）

# 用 Volcano Gang Job 提交 8 机 DDP 训练（注意 API 已是 v1beta1）
kubectl apply -f - <<EOF
apiVersion: batch.volcano.sh/v1beta1
kind: Job
metadata: {name: megatron-train}
spec:
  minAvailable: 8          # Gang：8 机同时
  tasks:
  - name: worker
    replicas: 8
    template:
      spec:
        containers:
        - name: train
          image: megatron:latest
          command: ["bash","launch.sh"]
          resources: {limits: {nvidia.com/gpu: 8}}  # 每机 8 卡
EOF

训练进程通过环境变量感知集群：WORLD_SIZE=64、RANK、MASTER_ADDR、MASTER_PORT（通常由 launchers 如 torchrun / deepspeed 注入）。

名词解释：

• WORLD_SIZE：参与训练的总进程数。
• RANK：当前进程的全局编号。
• MASTER_ADDR/PORT：主进程地址，用于初始化通信组。
• torchrun：PyTorch 内置的分布式启动器，自动注入这些环境变量。

---

9.8 ⭐ K8s 训练编排：Kubeflow / Ray / MPI Operator（通往「平台」的桥）

手写 Volcano YAML 能跑，但工业上几乎都用训练编排框架——它们才是「迷你 TI-ONE」（见 12）真正要组装的组件：

Kubeflow Training Operator（事实标准）

• 用 CRD 描述训练作业，自动拉起 Pod、注入分布式环境变量、做容错：
  • PyTorchJob：跑 torchrun 风格分布式训练。
  • TFJob：TensorFlow 分布式。
  • MPIJob：基于 MPI（horovod 风格）的机间集合通信。
• 官网：https://www.kubeflow.org/docs/components/training/
• 意义：算法方只写 PyTorchJob YAML，不用管 Gang/环境变量注入——这正是平台该提供的抽象。

名词解释：MPI（Message Passing Interface） 是 HPC 老牌的进程间通信标准。horovod 是基于 MPI 的分布式训练库（Uber 开源）。

Ray（Ray Train / Ray Serve）

• Ray 是分布式计算框架，2025–2026 在 AI 训练和推理服务上极度重要。
• Ray Train：统一多种并行策略（DDP/FSDP/DeepSpeed）的训练 API，跨节点调度。
• Ray Serve：模型推理服务化（多模型、弹性、流式批处理）。
• 特点：一套框架贯通「训练 + 服务 + 调参 + 强化学习」，适合做统一 AI 平台底座。
• 官网：https://www.ray.io/

MPI Operator（HPC 风格）

• 针对传统 HPC/MPI 工作负载，在 K8s 上跑 mpirun 风格的分布式任务（如 horovod）。
• 适合已经用 MPI 体系、需要做机间集合通信的场景。

⭐ 重点：你搭平台（12）时，不要让人手写 Volcano YAML，而是提供 Kubeflow Training Operator / Ray 这样的上层抽象。AI Infra 工程师的价值 = 把「难用的底层」封装成「算法方能自助的接口」。

---

9.9 全链路回顾

选型(03) → 接入(04) → 切分/共享(05) → 调度(06/07/08) → 训练编排(Kubeflow/Ray，09) → 通信底座(NCCL/IB，09.6) → 训练运行，一条龙打通。下一站 10 推理服务化，把训好的模型变成「能对外服务」的产品。

---

9.10 实战：在 MacBook 上跑分布式训练（CPU 模拟多卡）

没有 GPU？完全没问题——这一节用 PyTorch 的 DDP（DistributedDataParallel） 在 CPU 上模拟多卡训练。机制和真实多 GPU 完全一样，只是用 CPU 进程代替 GPU。

9.10.1 安装 PyTorch（MacBook / Linux）

# 创建虚拟环境（避免污染系统 Python）
python3 -m venv ~/ai-lab-venv
source ~/ai-lab-venv/bin/activate

# 安装 PyTorch（CPU 版本，不需要 GPU）
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

# 验证安装
python3 -c "import torch; print('PyTorch', torch.__version__); print('CUDA available:', torch.cuda.is_available())"
# 应输出：CUDA available: False（正常，我们用 CPU 模拟）

9.10.2 写一个 DDP 训练脚本

# 保存为 ~/ai-lab/ddp_demo.py
import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    """初始化进程组"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    # ★ 这就是 9.5 讲的「会合点(rdzv)」——所有进程连到这个地址协调
    dist.init_process_group(
        backend='gloo',       # CPU 用 gloo（GPU 用 nccl）
        rank=rank,            # 当前进程编号
        world_size=world_size # 总进程数（模拟 GPU 数量）
    )

def cleanup():
    dist.destroy_process_group()

class SimpleModel(nn.Module):
    """一个极简模型：2 层全连接"""
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

def train(rank, world_size):
    setup(rank, world_size)

    # 创建模型，放到当前「设备」（CPU 模拟）
    model = SimpleModel().to(rank)
    # ★ DDP 包装：这就是「数据并行」——每个进程各有一份模型副本
    ddp_model = DDP(model, device_ids=None)  # CPU 用 None

    # 模拟数据
    data = torch.randn(32, 784)  # batch=32
    target = torch.randint(0, 10, (32,))

    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.001)

    # 训练 5 步
    for step in range(5):
        optimizer.zero_grad()
        output = ddp_model(data)
        loss = loss_fn(output, target)
        loss.backward()
        # ★ 梯度同步（AllReduce）发生在 backward() 内部——这就是 9.6 讲的 NCCL/gloo 通信
        optimizer.step()

        if rank == 0:
            print(f"  Step {step}, Loss: {loss.item():.4f}")

    cleanup()

if __name__ == "__main__":
    world_size = 4  # 模拟 4 张 GPU
    # 用 torchrun 启动 4 个进程（模拟 4 卡）
    os.system(f"torchrun --nproc_per_node={world_size} ddp_demo.py")

每段的目的：

• setup()：初始化进程组，所有进程连到 MASTER_ADDR:PORT 协调——这就是「会合点」
• DDP(model)：把模型用 DDP 包装，梯度自动在进程间同步（AllReduce）
• torchrun --nproc_per_node=4：启动 4 个进程，模拟 4 张 GPU

9.10.3 运行并观察

cd ~/ai-lab
python3 ddp_demo.py

# 输出类似：
#   Step 0, Loss: 2.3145
#   Step 1, Loss: 2.2876
#   Step 2, Loss: 2.2534
#   Step 3, Loss: 2.2102
#   Step 4, Loss: 2.1678

# ★ 观察 4 个进程的 CPU 使用率
# 在另一个终端运行：
ps aux | grep python | grep ddp_demo
# 应看到 4 个 Python 进程

目的：你现在跑了一个真正的分布式训练——4 个进程各有一份模型副本，各自处理不同数据，梯度通过 AllReduce 同步。真实多 GPU 的唯一区别是 backend='nccl'（而非 gloo）和 device_ids=[gpu_id]。

9.10.4 体验 ZeRO/FSDP 分片

# 保存为 ~/ai-lab/fsdp_demo.py
# 体验 FSDP（完全分片数据并行）——把模型参数也切分到各进程
import torch
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
import torch.distributed as dist
import os

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12356'
    dist.init_process_group(backend='gloo', rank=rank, world_size=world_size)

class BigModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(*[nn.Linear(1000, 1000) for _ in range(10)])

    def forward(self, x):
        return self.layers(x)

def train(rank, world_size):
    setup(rank, world_size)

    model = BigModel()
    # 参数量
    n_params = sum(p.numel() for p in model.parameters())
    if rank == 0:
        print(f"模型参数量: {n_params:,} ({n_params * 4 / 1024 / 1024:.1f} MB FP32)")

    # ★ FSDP 包装：参数/梯度/优化器状态全部分片到各进程
    # 这就是 9.3 讲的 ZeRO-3 / FSDP
    fsdp_model = FSDP(model)

    data = torch.randn(8, 1000)
    output = fsdp_model(data)
    loss = output.sum()
    loss.backward()

    if rank == 0:
        print(f"  FSDP 前向+反向完成，Loss: {loss.item():.4f}")

    dist.destroy_process_group()

if __name__ == "__main__":
    world_size = 4
    os.system(f"torchrun --nproc_per_node={world_size} fsdp_demo.py")

python3 fsdp_demo.py
# 应看到模型参数量和训练完成信息

目的：FSDP 把模型参数也分片到各进程——每个进程只存 1/4 的参数。前向时聚合（AllGather），反向后分片。这就是 ZeRO-3 的原理。

9.10.5 在 K8s 上用 Kubeflow 提交训练任务

# 复用之前的 kind 集群
# 安装 Kubeflow 训练 Operator（只装 PyTorchJob CRD + controller）
kubectl apply -k "github.com/kubeflow/training-operator/manifests/overlays/standalone?ref=v1.7.0"

# 等待 controller 就绪
kubectl get pods -n kubeflow -w

# 提交一个 PyTorchJob（4 worker，CPU 训练）
cat <<EOF | kubectl apply -f -
apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: ddp-train-demo
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
          - name: pytorch
            image: pytorch/pytorch:2.2.0-cpu
            command:
            - python
            - -c
            - |
              import torch
              import torch.distributed as dist
              dist.init_process_group(backend='gloo')
              print(f"Master ready, world_size={dist.get_world_size()}")
              import time; time.sleep(60)
    Worker:
      replicas: 3
      restartPolicy: OnFailure
      template:
        spec:
          containers:
          - name: pytorch
            image: pytorch/pytorch:2.2.0-cpu
            command:
            - python
            - -c
            - |
              import torch
              import torch.distributed as dist
              dist.init_process_group(backend='gloo')
              print(f"Worker ready, rank={dist.get_rank()}")
              import time; time.sleep(60)
EOF

# 观察 4 个 Pod（1 Master + 3 Worker）同时启动
kubectl get pods -l pytorch-job-name=ddp-train-demo -w

目的：这就是在 K8s 上编排分布式训练——PyTorchJob CRD + Kubeflow Training Operator 自动创建 Master+Worker Pod 并配置环境变量。真实多 GPU 场景只是加 nvidia.com/gpu: 1 资源请求。

9.10.6 清理

kubectl delete pytorchjob ddp-train-demo --ignore-not-found
deactivate

---

9.11 自测题

1. 训练时显存里放哪四类东西？7B 全参微调真实显存量级？
2. DDP 和 FSDP 的区别？用「抄作业」类比解释三种并行。
3. 张量并行为什么需要高带宽互联（NVLink）？
4. ZeRO 分片了什么？ZeRO-3 和 FSDP 是什么关系？
5. Megatron-LM 结合了哪三种并行？
6. 在 K8s 上跑 8 机训练为什么要用 Volcano Gang Job？
7. ⭐ 节点内/跨节点的互联分别用什么技术？为什么以太网跑训练慢？什么是 RDMA/RoCE？
8. ⭐ AllReduce / AllGather / ReduceScatter 分别做什么？NCCL 是什么？
9. ⭐ Kubeflow Training Operator、Ray、MPI Operator 各自解决什么？为什么平台要用它们而不是手写 YAML？
10. ⭐ 大规模训练「扩展效率」常卡在通信还是算力？为什么？

答上即可进 10——训练完的模型怎么变成「能对外服务」的推理。