L3 · 专题文章

分布式训练

并行计算模型训练GPU集群

本页只讲 4 条最关键判断,不写百科

4核心点
先记住并行策略决定成败

选错并行类型会导致通信开销吞噬算力,显存不均衡引发OOM

LIVE
CORE 分布式训练
01 WHAT

是什么

将模型参数、数据或计算任务拆分到多个设备协同训练的技术体系

02 WHY

为什么

当单卡显存或算力无法承载模型规模时,工程师与算法研究员必须通过分布式训练突破硬件瓶颈

03 HOW

怎么做

通过框架配置并行策略(如数据并行/张量并行),启动多进程通信并同步梯度

04 WHEN

什么时候

模型参数量超过单卡容量或需缩短训练周期时使用;单卡可跑且时间可接受时不应用

FOCUS

先记住这些

只留最重要的判断,细节见下方实践
01 并行策略选择

不掌握就无法正确启动训练

数据并行适合小模型大数据,张量并行切分大矩阵,流水线并行按层拆分;选错会导致通信爆炸或显存浪费

02 通信瓶颈控制

决定实际加速比而非理论值

梯度同步频率、通信后端选择(NCCL/Gloo)与网络拓扑直接影响吞吐;高频AllReduce会拖慢迭代

03 显存均衡分配

最高频OOM根因

激活值峰值、优化器状态与参数分布不均会导致部分GPU提前耗尽显存;需启用梯度检查点或重计算

04 容错与恢复

长周期训练不可跳过

节点故障或网络抖动必发生;必须配置定期保存检查点与自动重启逻辑,否则前功尽弃

PROBLEM / POSITION / INTERFACE

先弄清它为什么存在,以及谁在使用

不从历史开始,从真实工作关系开始。
  1. 01 解决的问题

    模型参数量或激活值超出单GPU显存,或单卡训练耗时过长

    盲目多卡导致通信开销爆炸、OOM、梯度不同步或训练中断

    成功标准
    线性加速比接近理论值,显存均衡,训练稳定可恢复
  2. 02 AI 生态位

    连接模型架构与底层硬件的中间层,向上承接训练任务,向下调度GPU/NPU集群

    上游
    依赖模型结构定义、数据集划分与硬件拓扑
    下游
    为预训练、微调等任务提供可扩展的训练能力
  3. 03 人的生态位

    工程师负责策略选择、拓扑配置与异常排查

    适合使用
    模型显存需求>单卡容量,或需将训练时间压缩至可接受范围
    不必使用
    模型可单卡运行、数据量小、或团队缺乏集群运维能力时
  4. 04 独特价值

    唯一能突破单设备物理限制、实现千亿参数模型训练的工程路径

    通信与计算重叠难优化,拓扑感知配置复杂,故障恢复成本高

    复杂度判断
    复杂度不是功能数量,而是控制与验证成本。

INTERFACE FLOW

谁在操作,信息怎样流动

  1. WHO OPERATES

    HUMAN选择并行策略、配置集群拓扑、监控指标与干预故障

    SYSTEM执行进程间通信、梯度同步、检查点保存与设备调度

  2. INPUT

    训练数据批次、模型定义、并行策略配置

  3. CONTROL

    并行类型、通信后端、批大小、同步频率、容错策略

  4. OUTPUT

    更新后的模型权重、训练日志、检查点文件;不产生外部服务调用

CODE / PYTHONPyTorch 分布式最小启动
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup():
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

def train():
    setup()
    model = MyModel().cuda()
    model = DDP(model, device_ids=[int(os.environ["LOCAL_RANK"])])
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    for batch in dataloader:
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()
    dist.destroy_process_group()

if __name__ == "__main__":
    train()

最短闭环:初始化NCCL后端 → 包装DDP → 正常训练循环 → 销毁组。需通过torchrun启动。

JSON / RESPONSE典型训练日志与状态(示意)
{
  "rank": 0,
  "world_size": 4,
  "backend": "nccl",
  "status": "running",
  "gpu_memory_used_gb": 18.2,
  "gpu_memory_total_gb": 24.0,
  "step": 1250,
  "loss": 0.342,
  "throughput_samples_per_sec": 1280,
  "checkpoint_saved_at": "2024-06-15T10:30:00Z"
}

PRACTICE

启动前必查清单

确认所有节点GPU驱动与CUDA版本一致
验证NCCL通信可用(nccl-tests 带宽测试)
设置正确的 LOCAL_RANK 与 WORLD_SIZE 环境变量
启用梯度缩放(AMP)避免混合精度溢出
配置检查点保存路径与频率(至少每2小时一次)

FAQ

常见问题

选型、用法与失败,不复述定义。
01数据并行和张量并行该怎么选?

模型能放进单卡选数据并行;显存不够但层内矩阵大可加张量并行;两者可组合。

选错策略会导致通信爆炸或显存浪费,直接决定训练能否跑通。
02最小可用启动方式是什么?

用 torchrun 启动脚本,设置 NCCL 后端,用 DDP 包装模型,保持训练循环不变。

避免手动管理进程与通信,快速验证多卡可用性。
03训练中途节点挂了怎么办?

依赖定期保存的检查点重启;配置自动重试逻辑,从最近step恢复优化器状态。

长周期训练必遇故障,无容错机制将导致全部算力浪费。
04什么时候不需要分布式训练?

模型可单卡运行、数据量小、或训练时间可接受时,优先单卡简化调试。

分布式引入通信与运维成本,非必要不增加系统复杂度。

NEXT

下一步

预训练 量化