七、分布式框架
·
分布式训练框架
当模型规模扩大到单卡无法容纳，或训练速度无法满足需求时，分布式训练成为必然选择。分布式训练框架负责协调多卡、多节点之间的计算与通信，隐藏底层复杂性，让用户能够高效地训练大模型。本章将深入介绍分布式训练的核心组件：通信后端、数据并行实现、大规模训练工具以及集群调度系统。
1 通信后端
分布式训练的基础是多卡之间的高效通信。通信后端提供了底层的通信原语，如点对点发送/接收、集体通信（broadcast, all-reduce, all-gather 等）。
1.1 NCCL（NVIDIA Collective Communications Library）
NCCL 是 NVIDIA 开发的针对 GPU 的高性能集体通信库，专为多卡和多节点通信优化。
┌─────────────────────────────────────────────────────────────────┐
│ NCCL 在深度学习栈中的位置 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 应用层: PyTorch / TensorFlow / JAX / DeepSpeed │
│ │ │
│ ▼ │
│ 框架层: torch.distributed / tf.distribute │
│ │ │
│ ▼ |
│ 通信层: ┌─────────────────────────────┐ │
│ │ ★ NCCL (GPU 专用) │ │
│ │ • Ring/Tree 算法 │ │
│ │ • NVLink/RDMA 优化 │ │
│ │ • 自动拓扑感知 │ │
│ └─────────────────────────────┘ │
│ │ │
│ ▼ │
│ 驱动层: CUDA Driver + GPU Hardware │
│ │ │
│ ▼ │
│ 硬件层: [GPU0]══NVLink══[GPU1]══NVLink══[GPU2]══NVLink══[GPU3] │
│ │ │ │
│ InfiniBand/RoCE (多节点互联) │ │
│ └──────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
NCCL 支持的通信原语（Collective Operations）
┌─────────────────┬──────────────────────────────────────────────────┐
│ 原语名称 │ 功能说明 │
├─────────────────┼──────────────────────────────────────────────────┤
│ All-Reduce │ 所有卡输入数据 → 归约(求和/平均) → 所有卡得到相同结果 │
│ All-Gather │ 每卡有部分数据 → 收集所有数据 → 所有卡得到完整数据 │
│ Reduce-Scatter │ 所有卡输入数据 → 归约后按分片分发 → 每卡得部分结果 │
│ Broadcast │ 根卡有数据 → 广播到所有卡 → 所有卡得到相同数据 │
│ Reduce │ 所有卡输入数据 → 归约 → 仅根卡得到结果 │
│ All-to-All │ 每卡发送不同数据给其他卡 → 每卡接收来自各卡的数据 │
└─────────────────┴──────────────────────────────────────────────────┘
特点
1、利用 GPU 直连技术（NVLink、InfiniBand、RoCE）实现高带宽、低延迟通信；
2、实现了 Ring All-Reduce、Tree All-Reduce 等算法，自动选择最优路径；
3、支持多节点通信，通过 RDMA（Remote Direct Memory Access）绕过 CPU，降低延迟；
1.1.1 Ring All-Reduce 原理
1、将所有 GPU 视为一个环，数据分块后在环上依次传递，每个 GPU 累加接收到的数据；
2、分为两个阶段：Reduce-Scatter（各 GPU 累加部分数据）和 All-Gather（各 GPU 广播累加后的完整数据）；
3、通信量随 GPU 数量线性增长，但带宽利用率高，适合同构集群；
Ring All-Reduce 原理图解（核心算法）
================================================================
场景：4 张 GPU，每张卡有 4 个数据块 [A, B, C, D]
目标：所有卡最终都得到 [A₀+A₁+A₂+A₃, B₀+B₁+B₂+B₃, C₀+... , D₀+...]
初始状态：
GPU 0: [A₀, B₀, C₀, D₀]
GPU 1: [A₁, B₁, C₁, D₁]
GPU 2: [A₂, B₂, C₂, D₂]
GPU 3: [A₃, B₃, C₃, D₃]
期望结果（每张卡）：
[ΣA, ΣB, ΣC, ΣD] 其中 ΣA = A₀+A₁+A₂+A₃
Ring All-Reduce 两阶段流程
═══════════════════════════════════════════════════════
**阶段一：Reduce-Scatter（归约 - 散射）**
═══════════════════════════════════════════════════════
目标：每张卡最终只保留"自己负责"的那部分数据的归约结果
拓扑：环状连接 (0→1→2→3→0)
轮次 1 (共需 P-1=3 轮，P=卡数):
─────────────────────────────────────
每张卡发送"下一张卡负责"的数据块，接收"上一张卡"的数据
GPU 0: 发送 [B₀] → GPU1, 接收 [B₃] ← GPU3
本地累加: B₀ = B₀ + B₃
状态: [A₀, B₀+B₃, C₀, D₀]
GPU 1: 发送 [C₁] → GPU2, 接收 [C₀] ← GPU0
本地累加: C₁ = C₁ + C₀
状态: [A₁, B₁, C₁+C₀, D₁]
GPU 2: 发送 [D₂] → GPU3, 接收 [D₁] ← GPU1
本地累加: D₂ = D₂ + D₁
状态: [A₂, B₂, C₂, D₂+D₁]
GPU 3: 发送 [A₃] → GPU0, 接收 [A₂] ← GPU2
本地累加: A₃ = A₃ + A₂
状态: [A₃+A₂, B₃, C₃, D₃]
轮次 2:
─────────────────────────────────────
继续传递累加
GPU 0: 发送 [B₀+B₃] → GPU1, 接收 [B₁] ← GPU3(实际是GPU1的上一轮)
累加: B₀ = B₀+B₃ + B₁
...
轮次 3 (最后一轮):
─────────────────────────────────────
每张卡完成自己负责数据块的完整归约
最终状态 (Reduce-Scatter 结束):
GPU 0: [ΣA, -, -, -] ← 只保留 A 的归约结果
GPU 1: [ -, ΣB, -, -] ← 只保留 B 的归约结果
GPU 2: [ -, -, ΣC, -] ← 只保留 C 的归约结果
GPU 3: [ -, -, -, ΣD] ← 只保留 D 的归约结果
═══════════════════════════════════════════════════════
**阶段二：All-Gather（全收集）**
═══════════════════════════════════════════════════════
目标：将各卡归约好的分片广播给所有卡
轮次 1:
─────────────────────────────────────
每张卡发送自己持有的归约结果
GPU 0: 发送 [ΣA] → GPU1, 接收 [ΣD] ← GPU3
状态: [ΣA, ΣD, -, -]
GPU 1: 发送 [ΣB] → GPU2, 接收 [ΣA] ← GPU0
状态: [ΣA, ΣB, -, -]
...
轮次 3 (最后一轮):
─────────────────────────────────────
所有卡收集到全部归约结果
最终状态 (All-Gather 结束):
GPU 0: [ΣA, ΣB, ΣC, ΣD]
GPU 1: [ΣA, ΣB, ΣC, ΣD]
GPU 2: [ΣA, ΣB, ΣC, ΣD]
GPU 3: [ΣA, ΣB, ΣC, ΣD]
Ring All-Reduce 通信量分析
符号定义：
- P: GPU 数量
- N: 每张卡的数据量 (Bytes)
- BW: 链路带宽 (Bytes/sec)
Ring All-Reduce 通信特性：
─────────────────────────────────────────
每轮通信量: N / P (每张卡发送/接收一个分片)
总轮次数: 2 × (P - 1) (Reduce-Scatter + All-Gather 各 P-1 轮)
总通信量: 2 × (P - 1) × (N / P) ≈ 2N (当 P 较大时)
理论时间: T = 2N × (P-1)/P / BW ≈ 2N / BW
关键结论：
√ 通信时间与卡数 P 无关！(带宽饱和时)
√ 带宽利用率接近 100% (环上所有链路同时工作)
× 延迟随卡数线性增长 (轮次 = 2(P-1))
适用场景：
- 大数据量传输 (带宽瓶颈 > 延迟瓶颈)
- 同构集群 (所有链路带宽一致)
- GPU 数量适中 (P < 64 效果最佳)
1.1.2 Tree All-Reduce 原理图解
二叉树归约流程
场景：8 张 GPU，数据归约求和
拓扑：二叉树结构
Root(GPU 0)
/ \
GPU 1 GPU 2
/ \ / \
GPU 3 GPU 4 GPU 5 GPU 6
|
GPU 7
═══════════════════════════════════════════════════════
阶段一：Reduce (自底向上归约)
═══════════════════════════════════════════════════════
步骤 1 (叶子→父节点):
GPU 3 → GPU 1: 发送 data₃, GPU 1 计算 data₁+data₃
GPU 4 → GPU 1: 发送 data₄, GPU 1 计算 (data₁+data₃)+data₄
GPU 5 → GPU 2: 发送 data₅, GPU 2 计算 data₂+data₅
GPU 6 → GPU 2: 发送 data₆, GPU 2 计算 (data₂+data₅)+data₆
GPU 7 → GPU 3 → GPU 1: 链式传递
步骤 2 (中间→根):
GPU 1 → GPU 0: 发送 subtree_sum_1
GPU 2 → GPU 0: 发送 subtree_sum_2
GPU 0: 计算最终结果 = data₀ + sum₁ + sum₂
═══════════════════════════════════════════════════════
阶段二：Broadcast (自顶向下广播)
═══════════════════════════════════════════════════════
步骤 1 (根→子节点):
GPU 0 → GPU 1, GPU 2: 广播最终结果
步骤 2 (子节点→叶子):
GPU 1 → GPU 3, GPU 4: 广播
GPU 2 → GPU 5, GPU 6: 广播
...
最终：所有 8 张卡都得到相同的归约结果
1.1.3 Tree vs Ring 对比
┌──────────────────┬─────────────────┬─────────────────┐
│ 特性 │ Ring All-Reduce │ Tree All-Reduce │
├──────────────────┼─────────────────┼─────────────────┤
│ 通信轮次 │ 2×(P-1) │ 2×log₂(P) │
│ 每轮通信量 │ N/P │ N │
│ 总通信量 │ ≈2N │ ≈2N×log₂(P) │
│ 延迟敏感度 │ 高 (轮次多) │ 低 (轮次少) │
│ 带宽利用率 │ 高 (~100%) │ 中 (~50-70%) │
│ 拓扑要求 │ 环状/任意 │ 树状/层次化 │
│ 适合场景 │ 大数据量 │ 小数据量/高延迟│
└──────────────────┴─────────────────┴─────────────────┘
NCCL 的自适应策略：
- 数据量 > 阈值: 自动选择 Ring (带宽优先)
- 数据量 < 阈值: 自动选择 Tree (延迟优先)
- 异构网络: 混合使用或降级
使用
PyTorch 等框架默认使用 NCCL 作为 GPU 通信后端。
1.2 Gloo
Gloo 是 Facebook 开发的集体通信库，支持 CPU 和 GPU，但 GPU 通信性能不如 NCCL。
Gloo 核心架构与定位
┌─────────────────────────────────────────────────────────────────┐
│ Gloo 在深度学习栈中的位置 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 应用层：PyTorch / Caffe2 / 其他框架 │
│ │ │
│ ▼ │
│ 通信层： ┌─────────────────────────────┐ │
│ │ ★ Gloo │ │
│ │ • CPU + GPU 支持 │ │
│ │ • TCP/IB/UDP 传输 │ │
│ │ • 跨平台 (Linux/Windows) │ │
│ │ • 丰富的集体算法 │ │
│ └─────────────────────────────┘ │
│ │ │
│ ▼ │
│ 传输层：TCP Socket / InfiniBand / UDP / CUDA │
│ │ │
│ ▼ │
│ 硬件层：[CPU0]──以太网──[CPU1] [GPU0]──PCIe──[GPU1] │
└─────────────────────────────────────────────────────────────────┘
Gloo vs NCCL 定位对比
┌─────────────────────────────────────────────────────────────────┐
│ Gloo 与 NCCL 对比 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ NCCL (NVIDIA) Gloo (Facebook) │
│ ┌─────────────────┐ ┌─────────────────┐ │
│ │ GPU 专用优化 │ │ CPU + GPU 通用 │ │
│ │ NVLink/RDMA │ │ TCP/IB/UDP │ │
│ │ 高性能 │ │ 跨平台 │ │
│ │ Linux 为主 │ │ Linux/Windows │ │
│ │ 生产环境首选 │ │ 调试/备选方案 │ │
│ └─────────────────┘ └─────────────────┘ │
│ │
│ 选择建议： │
│ - GPU 训练 (生产): 优先 NCCL │
│ - CPU 训练：使用 Gloo │
│ - 调试/开发：Gloo 更友好 │
│ - Windows 环境：只能用 Gloo │
│ - NCCL 不可用：Gloo 作为备选 │
└─────────────────────────────────────────────────────────────────┘
1.2.1 Gloo 核心特性
1、跨平台，支持 TCP、InfiniBand 等；
Gloo 传输层架构：
═══════════════════════════════════════════════════════
┌─────────────────────────────────────────────────────┐
│ Gloo Context (通信上下文) │
├─────────────────────────────────────────────────────┤
│ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ TCP │ │ IB │ │ CUDA │ │
│ │ Socket │ │ Verbs │ │ Direct │ │
│ └────┬─────┘ └────┬─────┘ └────┬─────┘ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ 以太网网络 InfiniBand GPU 显存直接访问 │
│ (通用) (高性能) (GPU-GPU 通信) │
│ │
└─────────────────────────────────────────────────────┘
传输层选择优先级：
1. CUDA Direct (GPU 间直接通信)
2. InfiniBand Verbs (高性能网络)
3. TCP Socket (通用兼容)
2、提供丰富的集体算法实现，包括 All-Reduce、Broadcast 等。
Gloo 内置算法策略：
═══════════════════════════════════════════════════════
All-Reduce 算法选择：
─────────────────────────────────────────
数据量小 (< 1MB): Ring 算法 (低延迟)
数据量中 (1-10MB): Tree 算法 (平衡)
数据量大 (> 10MB): Hierarchical Tree (分层树)
进程数影响：
─────────────────────────────────────────
进程数少 (< 8): 直接算法
进程数中 (8-64): Ring/Tree 混合
进程数多 (> 64): 分层树 + 本地归约
自适应策略：
Gloo 根据数据大小、进程数、网络拓扑自动选择最优算法
3、通常用于 CPU 训练或作为 NCCL 不可用时的备选。
使用场景：调试、小规模训练或混合 CPU/GPU 环境。
┌─────────────────────────────────────────────────────────────────┐
│ Gloo 使用场景决策树 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 开始训练 │
│ │ │
│ ▼ │
│ ┌─────────────────┐ │
│ │ 有 GPU 吗？ │ │
│ └────────┬────────┘ │
│ Yes │ No │
│ ┌─┴─────────┐ │
│ ▼ ▼ │
│ ┌──────────┐ ┌──────────┐ │
│ │ GPU 训练 │ │ CPU 训练 │ │
│ └────┬─────┘ └────┬─────┘ │
│ │ │ │
│ ▼ ▼ │
│ ┌───────────┐ ┌───────────┐ │
│ │ 生产环境？ │ │ 使用 Gloo │ │
│ └─────┬─────┘ └───────────┘ │
│ Yes │ No │
│ ┌─┴────────┐ │
│ ▼ ▼ │
│ ┌──────────┐ ┌──────────┐ │
│ │ 用 NCCL │ │ 用 Gloo │ │
│ │ (高性能) │ │ (调试用) │ │
│ └──────────┘ └──────────┘ │
│ │
└─────────────────────────────────────────────────────────────────┘
1.3 其他后端
1.3.1 MPI（Message Passing Interface）
传统 HPC 通信标准，可通过 OpenMPI 等实现，但深度学习框架较少直接使用；
┌─────────────────────────────────────────────────────────────────┐
│ MPI 核心特性 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 定义：消息传递接口标准 (1994 年发布，当前主流 MPI-3.1/4.0) │
│ │
│ 核心思想： │
│ ┌─────────────────────────────────────────┐ │
│ │ • 进程间通过"发送/接收"消息通信 │ │
│ │ • 支持点对点 (P2P) 和集体通信 (Collective)│ │
│ │ • 语言无关 (C/C++/Fortran/Python 绑定) │ │
│ │ • 平台无关 (超算/集群/云) │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 典型实现： │
│ ┌─────────────┬─────────────────────────────────┐ │
│ │ OpenMPI │ 开源，最流行，支持多种网络 │ │
│ │ MPICH │ 参考实现，稳定可靠 │ │
│ │ Intel MPI │ 商业优化，Intel 硬件加速 │ │
│ │ MVAPICH2 │ 专为 InfiniBand 优化 │ │
│ │ Spectrum MPI│ IBM 商业实现 │ │
│ └─────────────┴─────────────────────────────────┘ │
│ │
│ 在深度学习中的定位： │
│ ✅ 超算中心训练大模型 │
│ ✅ 已有 MPI 基础设施的集群 │
│ ✅ 需要精细控制通信的场景 │
│ ❌ 普通 GPU 集群 (优先 NCCL) │
│ ❌ 快速原型开发 (优先 Gloo/NCCL) │
└─────────────────────────────────────────────────────────────────┘
MPI 通信原语图解
点对点通信 (P2P)：
═══════════════════════════════════════════════════════
进程 0 进程 1
┌─────────┐ ┌─────────┐
│ 数据 │ │ │
│ [1,2,3] │ │ │
└────┬────┘ └────┬────┘
│ │
│ MPI_Send │
│─────────────────────>│
│ │
│ │ MPI_Recv
│ │<─────────────────────
│ │
│ │ 收到 [1,2,3]
阻塞 vs 非阻塞：
- MPI_Send/Recv: 阻塞，发送/接收完成才返回
- MPI_Isend/Irecv: 非阻塞，立即返回，需 MPI_Wait 等待完成
集体通信 (Collective)：
═══════════════════════════════════════════════════════
All-Reduce (归约广播):
[P0:1] [P0:4]
[P1:2] ──Sum──> 5 ──Broadcast──> [P1:4]
[P2:3] [P2:4]
[P3:4] [P3:4]
All-Gather (全收集):
P0:[A] P0:[A,B,C,D]
P1:[B] ──Gather──> P1:[A,B,C,D]
P2:[C] P2:[A,B,C,D]
P3:[D] P3:[A,B,C,D]
Reduce-Scatter (归约分片):
P0:[A₀,B₀,C₀,D₀] P0:[ΣA]
P1:[A₁,B₁,C₁,D₁] ──> P1:[ΣB]
P2:[A₂,B₂,C₂,D₂] P2:[ΣC]
P3:[A₃,B₃,C₃,D₃] P3:[ΣD]
1.3.2 高速网络
InfiniBand 和 RoCE（RDMA over Converged Ethernet）是现代分布式训练集群标配，NCCL 可直接利用。
为什么要高速网络？
分布式训练的通信瓶颈分析：
═══════════════════════════════════════════════════════
大模型训练通信量估算 (175B 参数，8 节点):
─────────────────────────────────────────
每步梯度同步数据量:
参数量: 175B × 2B (FP16) = 350 GB
每步同步: 350 GB (梯度) + 700 GB (优化器状态) = 1.05 TB
网络带宽需求:
目标训练速度: 1 step / 1 second
所需带宽: 1.05 TB/s = 8.4 Tbps
单链路带宽对比:
┌────────────────┬─────────┬─────────┐
│ 网络类型 │ 带宽 │ 实际吞吐 │
├────────────────┼─────────┼─────────┤
│ 10GbE 以太网 │ 1.25 GB/s│ ~1 GB/s │
│ 25GbE 以太网 │3.125 GB/s│~2.5 GB/s│
│ 100GbE 以太网 │ 12.5 GB/s│ ~10 GB/s│
│ InfiniBand HDR │ 25 GB/s │ ~23 GB/s│
│ InfiniBand NDR │ 50 GB/s │ ~46 GB/s│
│ NVLink 3.0 │ 50 GB/s │ ~48 GB/s│
└────────────────┴─────────┴─────────┘
结论：
- 普通以太网无法满足大模型训练需求
- 必须使用高速网络 (IB/RoCE) + 通信优化 (梯度压缩/累积)
InfiniBand 详解
┌─────────────────────────────────────────────────────────────────┐
│ InfiniBand 架构 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 核心特性： │
│ ┌─────────────────────────────────────────┐ │
│ │ • 专用高性能互连技术 (非以太网) │ │
│ │ • 用户态直接访问 (绕过内核) │ │
│ │ • 零拷贝 + RDMA (远程直接内存访问) │ │
│ │ • 硬件级拥塞控制和流量管理 │ │
│ │ • 微秒级延迟，100+GB/s 带宽 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 硬件组件： │
│ ┌─────────────────────────────────────────┐ │
│ │ HCA (Host Channel Adapter) │ │
│ │ ┌─────────────────────┐ │ │
│ │ │ GPU/CPU 内存 │ │ │
│ │ │ │ │ │ │
│ │ │ [RDMA Engine] │ ← 硬件卸载 │ │
│ │ │ │ │ │ │
│ │ │ [InfiniBand Port] │ │ │
│ │ └─────────────────────┘ │ │
│ │ │ │
│ │ Switch (交换机) │ │
│ │ ┌─────────────────────┐ │ │
│ │ │ 多端口线速转发 │ │ │
│ │ │ 自适应路由 │ │ │
│ │ │ 硬件多播支持 │ │ │
│ │ └─────────────────────┘ │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 代际演进： │
│ ┌───────┬─────────┬─────────┬─────────┐ │
│ │ 代际 │ 每通道 │ 4×带宽 │ 发布年份 │ │
│ ├───────┼─────────┼─────────┼─────────┤ │
│ │ SDR │ 2.5 Gb/s│ 10 Gb/s │ 2003 │ │
│ │ DDR │ 5 Gb/s │ 20 Gb/s │ 2006 │ │
│ │ QDR │ 10 Gb/s │ 40 Gb/s │ 2008 │ │
│ │ FDR │ 14 Gb/s │ 56 Gb/s │ 2011 │ │
│ │ EDR │ 25 Gb/s │ 100 Gb/s│ 2014 │ │
│ │ HDR │ 50 Gb/s │ 200 Gb/s│ 2019 │ │
│ │ NDR │ 100 Gb/s│ 400 Gb/s│ 2022 │ │
│ │ XDR │ 200 Gb/s│ 800 Gb/s│ 2024+ │ │
│ └───────┴─────────┴─────────┴─────────┘ │
└─────────────────────────────────────────────────────────────────┘
RoCE (RDMA over Converged Ethernet) 详解
┌─────────────────────────────────────────────────────────────────┐
│ RoCE 技术架构 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 核心思想：在以太网上实现 RDMA 功能 │
│ │
│ 两种版本： │
│ ┌─────────────────────────────────────────┐ │
│ │ RoCE v1 (2010): │ │
│ │ • 链路层协议 (L2) │ │
│ │ • 只能在同一广播域内使用 │ │
│ │ • 已被 RoCE v2 取代 │ │
│ │ │ │
│ │ RoCE v2 (2014): │ │
│ │ • 网络层协议 (L3, UDP/IP) │ │
│ │ • 支持路由，可跨子网 │ │
│ │ • 当前主流版本 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ RoCE v2 协议栈： │
│ ┌─────────────────────────────────────────┐ │
│ │ 应用层: MPI / NCCL / UCX │ │
│ │ │ │ │
│ │ 传输层: [RoCE v2 Header] │ │
│ │ │ • BTH: Base Transport Header│ │
│ │ │ • RETH: RDMA Extended Header│ │
│ │ │ • AETH: Ack Extended Header │ │
│ │ │ │ │
│ │ 网络层: [UDP Header] + [IP Header] │ │
│ │ │ • UDP Port: 4791 │ │
│ │ │ │ │
│ │ 链路层: [Ethernet Header + FCS] │ │
│ │ │ • 优先级标记 (PFC/ECN) │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 关键依赖： │
│ ┌─────────────────────────────────────────┐ │
│ │ • PFC (Priority-based Flow Control) │ │
│ │ 无损以太网，避免丢包重传 │ │
│ │ │ │
│ │ • ECN (Explicit Congestion Notification)│ │
│ │ 拥塞通知，提前降速 │ │
│ │ │ │
│ │ • DCQCN (Data Center Quantized Congestion│ │
│ │ Notification) │ │
│ │ RoCE 专用拥塞控制算法 │ │
│ └─────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
InfiniBand vs RoCE 对比
┌──────────────────┬─────────────────┬───────────────┐
│ 特性 │ InfiniBand │ RoCE v2 │
├──────────────────┼─────────────────┼───────────────┤
│ 协议类型 │ 专用协议 │ 以太网 + UDP │
│ 延迟 (典型) │ 0.5-1.5 μs │ 1.5-3 μs │
│ 带宽 (主流) │ 200-400 Gb/s │ 100-200 Gb/s │
│ 无损传输 │ 硬件保证 │ 需 PFC/ECN 配置│
│ 路由能力 │ 原生支持 │ 依赖以太网路由 │
│ 部署成本 │ 高 (专用硬件) │ 中 (兼容以太网)│
│ 运维复杂度 │ 中 │ 高 (需调优网络)│
│ 与现有设施兼容 │ 低 │ 高 │
│ NCCL 支持 │ 原生优化 │ 原生支持 │
│ 适用场景 │ 超算/顶级集群 │ 企业/云集群 │
└──────────────────┴─────────────────┴─────────────────┘
选择建议：
┌─────────────────────────────────────────────────────────────────┐
│ • 预算充足 + 极致性能 → InfiniBand (HDR/NDR) │
│ • 已有以太网设施 + 性价比 → RoCE v2 (100/200GbE) │
│ • 小规模/开发环境 → 普通以太网 + 梯度累积 │
│ • 云环境 → 优先选择提供 IB 的实例 (AWS p4d, Azure NDv4) │
└─────────────────────────────────────────────────────────────────┘
MPI 与 NCCL/Gloo 对比
┌──────────────────┬───────────┬─────────────┬────────────┐
│ 特性 │ MPI │ NCCL │ Gloo │
├──────────────────┼───────────┼─────────────┼────────────┤
│ 定位 │ 通用标准 │ GPU 专用 │ CPU+GPU │
│ GPU 通信性能 │ ★★★☆☆ │ ★★★★★ │ ★★☆☆☆ │
│ CPU 通信性能 │ ★★★★★ │ 不支持 │ ★★★★☆ │
│ 网络支持 │ 多种 │ NVLink/IB │ TCP/IB │
│ 易用性 │ ★★☆☆☆ │ ★★★★★ │ ★★★★☆ │
│ 调试友好度 │ ★★☆☆☆ │ ★★★☆☆ │ ★★★★★ │
│ 超算集成 │ ★★★★★ │ ★★☆☆☆ │ ★★☆☆☆ │
│ PyTorch 支持 │ 需配置 │ 默认 │ 默认 │
│ 跨平台 │ ★★★★★ │ ★★☆☆☆ │ ★★★★★ │
└──────────────────┴────────────┴────────────┴────────────┘
关键结论：
- 深度学习生产环境：优先 NCCL (GPU) / Gloo (CPU)
- MPI 适合：超算中心、已有 MPI 设施、需要精细控制的场景
- 混合方案：MPI 做进程管理 + NCCL 做 GPU 通信 (最佳实践)
2 数据并行实现
数据并行是最常用的分布式策略：每张 GPU 持有完整模型副本，处理不同数据分片，通过梯度同步保持模型一致。
2.1 PyTorch DistributedDataParallel（DDP）
DDP 是 PyTorch 官方提供的数据并行实现，基于 torch.distributed 包。
2.1.1 工作原理
1、每个进程（通常对应一张 GPU）独立加载模型和数据。
2、前向传播和反向传播在本地计算，得到梯度。
3、反向传播完成后，DDP 自动注册一个 All-Reduce hook，对所有进程的梯度进行求和平均。
4、平均后的梯度用于更新本地模型参数（所有副本保持相同）。
┌─────────────────────────────────────────────────────────────────┐
│ 数据并行 (Data Parallelism) │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 核心思想： │
│ ┌─────────────────────────────────────────┐ │
│ │ • 每张 GPU 保存完整模型副本 │ │
│ │ • 每张 GPU 处理不同数据分片 │ │
│ │ • 反向传播后同步梯度 (All-Reduce) │ │
│ │ • 所有副本保持参数一致 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 工作流程： │
│ │
│ GPU 0 GPU 1 GPU 2 │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │ Model W │ │ Model W │ │ Model W │ │
│ │ (完整) │ │ (完整) │ │ (完整) │ │
│ └────┬────┘ └────┬────┘ └────┬────┘ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │ Data 0 │ │ Data 1 │ │ Data 2 │ │
│ │ Batch 0 │ │ Batch 1 │ │ Batch 2 │ │
│ └────┬────┘ └────┬────┘ └────┬────┘ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │Forward │ │Forward │ │Forward │ │
│ │Backward │ │Backward │ │Backward │ │
│ │ Grad g₀ │ │ Grad g₁ │ │ Grad g₂ │ │
│ └────┬────┘ └────┬────┘ └────┬────┘ │
│ │ │ │ │
│ └────────────┬────────────┴─────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────┐ │
│ │ All-Reduce │ ← 梯度同步 (求和平均) │
│ │ (NCCL/NCCL) │ g_avg = (g₀+g₁+g₂)/3 │
│ └────────┬────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────┐ │
│ │ 更新参数 W │ ← 所有 GPU 使用相同梯度更新 │
│ │ W = W - lr*g_avg│ │
│ └─────────────────┘ │
│ │
│ 结果：所有 GPU 的模型参数保持完全一致 │
└─────────────────────────────────────────────────────────────────┘
DDP 工作原理图解
┌─────────────────────────────────────────────────────────────────┐
│ DDP 工作流程 (单步训练) │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 初始化阶段： │
│ ┌─────────────────────────────────────────┐ │
│ │ 1. 启动 N 个进程 (每进程 1 卡) │ │
│ │ 2. 每个进程加载完整模型 │ │
│ │ 3. 初始化分布式环境 (dist.init_process_group)│ │
│ │ 4. 用 DDP 包装模型 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 训练阶段 (每步)： │
│ ┌─────────────────────────────────────────┐ │
│ │ GPU 0 GPU 1 GPU 2 │ │
│ │ │ │ │ │ │
│ │ ▼ ▼ ▼ │ │
│ │ [Forward] [Forward] [Forward] │ ← 独立前向 │
│ │ │ │ │ │ │
│ │ ▼ ▼ ▼ │ │
│ │ [Loss] [Loss] [Loss] │ │
│ │ │ │ │ │ │
│ │ ▼ ▼ ▼ │ │
│ │ [Backward] [Backward] [Backward] │ ← 独立反向 │
│ │ │ │ │ │ │
│ │ ▼ ▼ ▼ │ │
│ │ Grad₀ Grad₁ Grad₂ │ │
│ │ │ │ │ │ │
│ │ └──────┬──────┴──────┬───────┘ │ │
│ │ ▼ ▼ │ │
│ │ ┌──────────────────────────┐ │ │
│ │ │ All-Reduce (NCCL) │ │ ← 梯度同步 │
│ │ │ g_avg = Σgᵢ / N │ │ │
│ │ └────────────┬─────────────┘ │ │
│ │ ▼ │ │
│ │ ┌──────────────────────────┐ │ │
│ │ │ Optimizer Step │ │ ← 独立更新 │
│ │ │ W = W - lr * g_avg │ │ │
│ │ └──────────────────────────┘ │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 关键：所有 GPU 在每步后参数完全一致 │
└─────────────────────────────────────────────────────────────────┘
2.1.2 关键优化：
1、梯度分桶（Bucketization）：将梯度按大小分桶，桶内梯度一次性 All-Reduce，减少通信次数。
2、异步通信：计算与通信重叠，提高效率。
3、支持混合精度：与 AMP（自动混合精度）无缝集成。
┌─────────────────────────────────────────────────────────────────┐
│ DDP 核心优化 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 1. 梯度分桶 (Gradient Bucketization) │
│ ═══════════════════════════════════════════════════════ │
│ 问题：逐参数通信开销大 │
│ 解决：将梯度分组打包，批量通信 │
│ │
│ 示例： │
│ ┌─────────────────────────────────────────┐ │
│ │ 参数：[p₁, p₂, p₃, p₄, p₅, p₆, ...] │ │
│ │ ↓ 分组 │ │
│ │ 桶 0: [p₁, p₂, p₃] → 一次 All-Reduce │ │
│ │ 桶 1: [p₄, p₅, p₆] → 一次 All-Reduce │ │
│ │ ... │ │
│ │ │ │
│ │ 优势：减少通信次数，提高带宽利用率 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 2. 通信与计算重叠 (Overlap Comm & Compute) │
│ ═══════════════════════════════════════════════════════ │
│ 传统方式： │
│ [Backward 全部完成] → [All-Reduce] → [Optimizer Step] │
│ │
│ DDP 优化方式： │
│ [Backward 层 N] → [All-Reduce 层 N] ─┐ │
│ [Backward 层 N-1] → [All-Reduce 层 N-1] ─┤ 重叠执行 │
│ [Backward 层 N-2] → [All-Reduce 层 N-2] ─┤ │
│ ... │ │
│ ↓ │
│ [Optimizer Step] │
│ │
│ 优势：隐藏通信延迟，提升吞吐量 │
│ │
│ 3. 混合精度支持 (AMP Integration) │
│ ═══════════════════════════════════════════════════════ │
│ DDP 与自动混合精度无缝集成： │
│ - 前向：FP16 计算 │
│ - 反向：FP16 梯度 → All-Reduce → FP32 主权重更新 │
│ - 优势：减少 50% 显存，加速计算 │
└─────────────────────────────────────────────────────────────────┘
使用方式：通过 torchrun 或 torch.distributed.launch 启动多进程，调用 model = DDP(model) 包装模型。
2.2 Horovod
Horovod 是 Uber 开源的分布式训练框架，支持 TensorFlow、PyTorch、MXNet 等，以易用性和高性能著称。
Horovod 架构与工作流程
┌─────────────────────────────────────────────────────────────────┐
│ Horovod 架构 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 核心组件： │
│ ┌─────────────────────────────────────────┐ │
│ │ • 初始化模块 - 建立通信上下文 │ │
│ │ • 尺寸查询 - 获取 worker 数 │ │
│ │ • 排名查询 - 获取当前 rank │ │
│ │ • 分布式优化器 - 包装原生优化器 │ │
│ │ • 参数广播 - 确保初始一致 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 工作流程： │
│ ┌─────────────────────────────────────────┐ │
│ │ 1. 初始化通信后端 │ │
│ │ 2. 创建模型和优化器 │ │
│ │ 3. 广播初始参数 (从 Rank 0) │ │
│ │ 4. 包装优化器 (注入 All-Reduce 逻辑) │ │
│ │ 5. 训练循环 (自动梯度同步) │ │
│ │ 6. 广播优化器状态 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 与 DDP 对比： │
│ ┌─────────────────────────────────────────┐ │
│ │ DDP: 包装模型 (模型层拦截) │ │
│ │ Horovod: 包装优化器 (优化器层拦截) │ │
│ └─────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
2.2.1 特点
1、基于 MPI 和 NCCL 实现 All-Reduce，抽象简洁；
2、提供 hvd.DistributedOptimizer 包装优化器，自动处理梯度平均；
3、支持动态缩放学习率（根据 worker 数量）；
4、兼容弹性训练（动态增删 worker）；
┌─────────────────────────────────────────────────────────────────┐
│ Horovod 核心特性 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 1. 跨框架支持： │
│ ═══════════════════════════════════════════════════════ │
│ • 统一 API 支持 TensorFlow, PyTorch, MXNet, Keras │
│ • 代码迁移成本低，逻辑一致 │
│ │
│ 2. 梯度压缩 (Gradient Compression)： │
│ ═══════════════════════════════════════════════════════ │
│ • 支持 FP16 压缩：梯度传输前转为 FP16 │
│ • 支持量化压缩：进一步减少通信量 │
│ • 适用场景：网络带宽受限的多节点训练 │
│ │
│ 3. 自适应学习率缩放： │
│ ═══════════════════════════════════════════════════════ │
│ • 线性缩放规则：LearningRate = BaseLR * WorkerCount │
│ • 自动根据并行度调整，保持收敛稳定性 │
│ │
│ 4. 弹性训练 (Elastic Training)： │
│ ═══════════════════════════════════════════════════════ │
│ • 支持动态增删 Worker │
│ • 训练过程中无需中断，自动重新平衡数据 │
│ • 适合云环境_spot 实例或资源波动场景 │
│ │
│ 5. 后端灵活性： │
│ ═══════════════════════════════════════════════════════ │
│ • 默认使用 MPI 进行进程管理 │
│ • 使用 NCCL 进行 GPU 通信 │
│ • 支持 Gloo 作为 CPU 后备 │
└─────────────────────────────────────────────────────────────────┘
2.2.3 对比
1、Horovod 更早支持跨框架，且在某些场景通信效率更高；
2、DDP 与 PyTorch 生态结合更紧密，社区支持更好；
┌─────────────────────────────────────────────────────────────────┐
│ DDP vs Horovod 对比 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 架构对比： │
│ ┌─────────────────────────────────────────┐ │
│ │ DDP │ │
│ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
│ │ │ Model │ │ Model │ │ Model │ │ │
│ │ │ (DDP) │ │ (DDP) │ │ (DDP) │ │ │
│ │ └────┬────┘ └────┬────┘ └────┬────┘ │ │
│ │ │ │ │ │ │
│ │ └────────────┴────────────┘ │ │
│ │ │ │ │
│ │ torch.distributed │ │
│ │ (NCCL Backend) │ │
│ └─────────────────────────────────────────────────┘ │
│ │
│ ┌─────────────────────────────────────────────────┐ │
│ │ Horovod │ │
│ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
│ │ │ Model │ │ Model │ │ Model │ │ │
│ │ └────┬────┘ └────┬────┘ └────┬────┘ │ │
│ │ │ │ │ │ │
│ │ ┌────▼────┐ ┌────▼────┐ ┌────▼────┐ │ │
│ │ │ Opt │ │ Opt │ │ Opt │ │ │
│ │ │(hvd.Opt)│ │(hvd.Opt)│ │(hvd.Opt)│ │ │
│ │ └────┬────┘ └────┬────┘ └────┬────┘ │ │
│ │ └────────────┴────────────┘ │ │
│ │ │ │ │
│ │ MPI + NCCL │ │
│ └─────────────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────┘
详细对比表：
┌──────────────────┬─────────────────┬─────────────────┐
│ 特性 │ DDP │ Horovod │
├──────────────────┼─────────────────┼─────────────────┤
│ 开发者 │ PyTorch 团队 │ Uber/LinkedIn │
│ 框架支持 │ PyTorch 专用 │ 多框架 (TF/PT/MX)│
│ 集成度 │ PyTorch 原生 │ 外部库 │
│ 学习曲线 │ 中等 │ 较低 │
│ 代码改动 │ 包装模型 │ 包装优化器 │
│ 启动方式 │ torchrun │ horovodrun │
│ 弹性训练 │ 有限支持 │ 原生支持 │
│ 梯度压缩 │ 需自定义 │ 内置支持 │
│ 社区支持 │ ★★★★★ │ ★★★☆☆ │
│ 文档质量 │ ★★★★★ │ ★★★★☆ │
│ 生产采用率 │ 非常高 │ 中等 │
│ 维护状态 │ 活跃 │ 维护模式 │
└──────────────────┴─────────────────┴─────────────────┘
性能对比 (8×A100, ResNet-50):
┌──────────────────┬─────────┬─────────┬─────────┐
│ 指标 │ DDP │ Horovod │ 差异 │
├──────────────────┼─────────┼─────────┼─────────┤
│ 训练吞吐量 │ 100% │ 98% │ -2% │
│ 启动时间 │ 快 │ 中 │ DDP 快 │
│ 通信效率 │ 高 │ 高 │ 相当 │
│ 显存占用 │ 相同 │ 相同 │ 无差异 │
│ 调试友好度 │ 高 │ 中 │ DDP 好 │
└─────────────────┴─────────┴─────────┴─────────┘
使用：初始化 hvd.init()，广播初始参数，包装优化器，然后 hvd.BroadcastGlobalVariablesHook 确保所有 worker 初始一致。
选型决策树
┌─────────────────────────────────────────────────────────────────┐
│ 数据并行框架选型决策树 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 开始选型 │
│ │ │
│ ▼ │
│ ┌─────────────────────────┐ │
│ │ 使用 PyTorch 吗？ │ │
│ └────────┬────────────────┘ │
│ No │ Yes │
│ ┌─┴─┐ │
│ ▼ ▼ │
│ ┌──────────┐ ┌──────────┐ │
│ │ 多框架 │ │ PyTorch │ │
│ │ 需求 │ │ 专用 │ │
│ └────┬─────┘ └────┬─────┘ │
│ │ │ │
│ ▼ ▼ │
│ ┌─────────┐ ┌─────────────┐ │
│ │Horovod │ │ 生产环境？ │ │
│ │(跨框架) │ └──────┬──────┘ │
│ └─────────┘ 是 │ 否 │
│ ┌───┴───┐ │
│ ▼ ▼ │
│ ┌──────────┐ ┌──────────┐ │
│ │ DDP │ │ 都可以 │ │
│ │ (推荐) │ │ DDP 优先 │ │
│ └──────────┘ └──────────┘ │
│ │
│ 特殊情况： │
│ • 需要弹性训练 → Horovod Elastic │
│ • 已有 MPI 设施 → Horovod + MPI │
│ • TensorFlow 项目 → Horovod │
│ • PyTorch 新项目 → DDP │
└─────────────────────────────────────────────────────────────────┘
3 大规模训练工具
单靠数据并行无法解决超大模型（百亿/千亿参数）的显存问题，因此出现了集成多种并行策略的高级训练工具。
3.1 DeepSpeed
DeepSpeed 是微软开源的深度学习优化库，基于 PyTorch 构建，提供 ZeRO 优化器、流水线并行、稀疏注意力等高级特性。
DeepSpeed 整体架构
┌─────────────────────────────────────────────────────────────────┐
│ DeepSpeed 架构全景 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ ┌─────────────────────────────────────────┐ │
│ │ 用户代码层 │ │
│ │ • PyTorch 原生训练脚本 │ │
│ │ • 最小改动即可接入 (通常 < 10 行) │ │
│ └────────────────┬────────────────────────┘ │
│ │ │
│ ┌────────────────▼────────────────────────┐ │
│ │ DeepSpeed Engine │ │
│ │ ┌─────────────────────────────────┐ │ │
│ │ │ • 初始化: initialize() │ │ │
│ │ │ • 前向: engine(inputs) │ │ │
│ │ │ • 反向: engine.backward(loss) │ │ │
│ │ │ • 更新: engine.step() │ │ │
│ │ └─────────────────────────────────┘ │ │
│ └────────────────┬────────────────────────┘ │
│ │ │
│ ┌────────────────▼────────────────────────┐ │
│ │ 优化策略层 │ │
│ │ ┌─────────┬─────────┬─────────┐ │ │
│ │ │ ZeRO │ PP │ TP │ │ │
│ │ │ 优化器 │ 流水线 │ 张量并行│ │ │
│ │ └─────────┴─────────┴─────────┘ │ │
│ │ ┌─────────┬─────────┬─────────┐ │ │
│ │ │ AMP │ Offload│ Sparse │ │ │
│ │ │ 混合精度│ CPU卸载│ 稀疏注意力│ │ │
│ │ └─────────┴─────────┴─────────┘ │ │
│ └────────────────┬────────────────────────┘ │
│ │ │
│ ┌────────────────▼────────────────────────┐ │
│ │ 通信与系统层 │ │
│ │ • NCCL: GPU 集体通信 │ │
│ │ • MPI: 多节点进程管理 │ │
│ │ • 异步通信: 重叠计算与通信 │ │
│ │ • 内存管理: 智能分配与复用 │ │
│ └────────────────┬────────────────────────┘ │
│ │ │
│ ┌────────────────▼────────────────────────┐ │
│ │ 硬件层 │ │
│ │ [GPU]──NVLink──[GPU]──IB/RoCE──[Node] │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 配置驱动：所有策略通过 JSON 配置启用，无需修改模型代码 │
└─────────────────────────────────────────────────────────────────┘
3.1.1 核心特性：
1、ZeRO 优化器：如前章所述，支持 ZeRO-1/2/3，极大降低显存占用，可训练万亿参数模型。
2、流水线并行：实现类似 GPipe 的流水线调度，支持自动切分模型。
3、混合精度训练：集成 AMP 和 BF16。
4、梯度累积与裁剪：内置支持。
5、通信优化：使用 NCCL 和异步通信重叠。
易用性
通过配置文件（JSON）指定优化策略，代码改动极小。例如，启用 ZeRO-3 只需在配置中设置 “zero_optimization”: {“stage”: 3}。
生态
与 HuggingFace Transformers 深度集成，成为训练大模型的事实标准之一。
3.2 Megatron-LM
Megatron-LM 由 NVIDIA 开发，专注于张量并行（层内并行）和流水线并行的实现，最早用于训练 GPT-3 等超大规模模型。
┌─────────────────────────────────────────────────────────────────┐
│ Megatron-LM 3D 并行架构 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 三维并行策略： │
│ ┌─────────────────────────────────────────┐ │
│ │ │ │
│ │ 维度 1: 数据并行 (Data Parallelism) │ │
│ │ • 切分：Batch 维度 │ │
│ │ • 作用：增加吞吐，复制模型 │ │
│ │ • 通信：梯度 All-Reduce │ │
│ │ │ │
│ │ 维度 2: 张量并行 (Tensor Parallelism) │ │
│ │ • 切分：矩阵内维度 (隐藏层/头数) │ │
│ │ • 作用：解决单层过大问题 │ │
│ │ • 通信：层内 All-Reduce (高频) │ │
│ │ │ │
│ │ 维度 3: 流水线并行 (Pipeline Parallelism)│ │
│ │ • 切分：层间维度 (Transformer Blocks) │ │
│ │ • 作用：解决层数过多问题 │ │
│ │ • 通信：层间 P2P (激活值/梯度) │ │
│ │ │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 拓扑结构示例 (64 卡，DP=2, TP=4, PP=8): │
│ ┌─────────────────────────────────────────┐ │
│ │ DP Group 0 (32 卡) │ │
│ │ ┌─────────────────────────────────┐ │ │
│ │ │ PP Stage 0 (TP=4, GPU 0-3) │ │ │
│ │ │ PP Stage 1 (TP=4, GPU 4-7) │ │ │
│ │ │ ... │ │ │
│ │ │ PP Stage 7 (TP=4, GPU 28-31) │ │ │
│ │ └─────────────────────────────────┘ │ │
│ │ │ │
│ │ DP Group 1 (32 卡) │ │
│ │ ┌─────────────────────────────────┐ │ │
│ │ │ 结构与 Group 0 相同，处理不同数据 │ │
│ │ └─────────────────────────────────┘ │ │
│ │ │ │
│ │ 总卡数 = DP × TP × PP = 2 × 4 × 8 = 64 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 通信域隔离： │
│ • 每个并行维度有独立的通信组 (Process Group) │
│ • TP 组内通信频繁，需高速互联 (NVLink) │
│ • PP 组间通信较少，可跨节点 │
│ • DP 组间通信仅在步末，可跨节点 │
└─────────────────────────────────────────────────────────────────┘
核心特性
1、张量并行：将 Transformer 层的权重矩阵按列/行切分，通过通信合并结果，支持高效计算。
2、流水线并行：结合张量并行，实现 3D 并行（数据+张量+流水线）。
3、序列并行：沿序列维度切分 LayerNorm 和 Dropout，减少激活显存。
4、优化器：支持 Adam、SGD 等，与分布式结合。
使用方式
Megatron-LM 提供了一套脚本和代码库，用户需按其规范定义模型，但核心并行逻辑已封装。
与 DeepSpeed 结合
常将 Megatron-LM 的张量并行与 DeepSpeed 的 ZeRO 结合（如 Megatron-DeepSpeed），实现极致扩展。
┌─────────────────────────────────────────────────────────────────┐
│ Megatron-DeepSpeed 融合架构 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 融合动机： │
│ ┌─────────────────────────────────────────┐ │
│ │ • Megatron 优势：TP/PP 性能极致，通信优化好 │ │
│ │ • Megatron 劣势：显存优化不如 ZeRO，代码侵入性强│ │
│ │ │ │
│ │ • DeepSpeed 优势：ZeRO 显存优化极强，易用性高 │ │
│ │ • DeepSpeed 劣势：TP 支持较弱 (主要靠 ZeRO) │ │
│ │ │ │
│ │ • 目标：结合 TP 的计算优势 + ZeRO 的显存优势 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 集成方案 (Megatron-DeepSpeed)： │
│ ┌─────────────────────────────────────────┐ │
│ │ 1. 模型结构： │ │
│ │ • 使用 Megatron 定义模型 (TP/PP/SP) │ │
│ │ • 保留 Megatron 的通信逻辑 │ │
│ │ │ │
│ │ 2. 优化器状态： │ │
│ │ • 使用 DeepSpeed ZeRO 优化器 │ │
│ │ • 切分优化器状态、梯度、甚至参数 │ │
│ │ │ │
│ │ 3. 通信协调： │ │
│ │ • TP 通信：Megatron 负责 (层内) │ │
│ │ • ZeRO 通信：DeepSpeed 负责 (层间/全局) │ │
│ │ • 需避免通信冲突和冗余 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 显存优化效果： │
│ ┌─────────────────────────────────────────┐ │
│ │ 纯 Megatron: │ │
│ │ • 每卡显存 = 模型/TP/PP + 优化器 + 梯度 │ │
│ │ │ │
│ │ Megatron + ZeRO-3: │ │
│ │ • 每卡显存 = 模型/TP/PP/ZeRO + 激活 │ │
│ │ • 优化器/梯度/参数进一步切分 │ │
│ │ │ │
│ │ 结果： │ │
│ │ • 可训练模型规模提升 2-4 倍 │ │
│ │ • 相同硬件下支持更大 Batch Size │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 代表性项目： │
│ • NVIDIA Megatron-DeepSpeed (官方融合版) │
│ • DeepSpeed-Megatron (社区维护版) │
│ • 用于训练：MT-NLG (530B), 以及各种开源千亿模型 │
└─────────────────────────────────────────────────────────────────┘
3.3 Megatron-LM vs DeepSpeed 对比
┌─────────────────────────────────────────────────────────────────┐
│ Megatron-LM vs DeepSpeed 对比 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 详细对比表： │
│ ┌──────────────────┬─────────────────┬─────────────────┐ │
│ │ 特性 │ Megatron-LM │ DeepSpeed │ │
│ ├──────────────────┼─────────────────┼─────────────────┤ │
│ │ 开发者 │ NVIDIA │ Microsoft │ │
│ │ 核心优势 │ 张量并行 (TP) │ ZeRO 优化器 │ │
│ │ 显存优化 │ ★★★★☆ │ ★★★★★ │ │
│ │ 计算性能 │ ★★★★★ │ ★★★★☆ │ │
│ │ 易用性 │ ★★☆☆☆ │ ★★★★★ │ │
│ │ 代码侵入性 │ 高 (需改模型) │ 低 (配置驱动) │ │
│ │ 模型定义 │ 专有格式 │ 兼容原生 │ │
│ │ 社区生态 │ 中等 │ 庞大 │ │
│ │ 文档质量 │ 专业但难懂 │ 友好 │ │
│ │ 推理支持 │ Megatron-Core │ DeepSpeed Inf │ │
│ │ 适用场景 │ 超大规模预训练 │ 微调/预训练 │ │
│ └──────────────────┴─────────────────┴─────────────────┘ │
│ │
│ 选择建议： │
│ ┌─────────────────────────────────────────┐ │
│ │ 选择 Megatron-LM: │ │
│ │ • 训练千亿参数以上模型 │ │
│ │ • 需要极致计算性能 │ │
│ │ • 团队有较强工程能力 │ │
│ │ • 硬件为 NVIDIA 全套 (GPU+NVLink+IB) │ │
│ │ │ │
│ │ 选择 DeepSpeed: │ │
│ │ • 快速验证想法，迭代速度快 │ │
│ │ • 基于 HuggingFace 模型微调 │ │
│ │ • 显存受限，需 ZeRO-Offload │ │
│ │ • 团队工程资源有限 │ │
│ │ │ │
│ │ 选择 Megatron-DeepSpeed: │ │
│ │ • 超大模型预训练 │ │
│ │ • 需要 TP 性能 + ZeRO 显存 │ │
│ │ • 愿意承担集成复杂度 │ │
│ └─────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
3.4 并行策略配置指南
┌─────────────────────────────────────────────────────────────────┐
│ 并行策略配置指南 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 1. 确定总卡数： │
│ ┌─────────────────────────────────────────┐ │
│ │ • 根据可用硬件资源确定总 GPU 数量 │ │
│ │ • 例如：8 节点 × 8 卡 = 64 GPU │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 2. 配置张量并行 (TP)： │
│ ┌─────────────────────────────────────────┐ │
│ │ • 原则：单机内优先 (利用 NVLink) │ │
│ │ • 建议：TP = 1, 2, 4, 8 (2 的幂次) │ │
│ │ • 场景： │ │
│ │ - 隐藏层大 (>10K): TP=4 或 8 │ │
│ │ - 隐藏层中 (<5K): TP=1 或 2 │ │
│ │ • 限制：TP 越大，通信开销越大 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 3. 配置流水线并行 (PP)： │
│ ┌─────────────────────────────────────────┐ │
│ │ • 原则：跨节点使用，解决层数过多问题 │ │
│ │ • 建议：PP = 总卡数 / (DP × TP) │ │
│ │ • 场景： │ │
│ │ - 层数多 (>100): PP=4, 8, 16 │ │
│ │ - 层数少 (<50): PP=1 或 2 │ │
│ │ • 限制：PP 越大，气泡率越高 │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 4. 配置数据并行 (DP)： │
│ ┌─────────────────────────────────────────┐ │
│ │ • 原则：剩余卡数全部用于 DP │ │
│ │ • 公式：DP = 总卡数 / (TP × PP) │ │
│ │ • 优势：DP 通信开销最低，扩展性最好 │ │
│ │ • 建议：优先增大 DP，再调整 TP/PP │ │
│ └─────────────────────────────────────────┘ │
│ │
│ 5. 典型配置示例 (64 卡): │
│ ┌─────────────────────────────────────────┐ │
│ │ 方案 A (侧重显存): TP=1, PP=8, DP=8 │ │
│ │ 方案 B (侧重性能): TP=4, PP=2, DP=8 │ │
│ │ 方案 C (均衡): TP=2, PP=4, DP=8 │ │
│ │ │ │
│ │ 选择依据： │ │
│ │ • 显存瓶颈 → 增加 PP 或 TP │ │
│ │ • 通信瓶颈 → 减少 TP，增加 DP │ │
│ │ • 计算瓶颈 → 增加 DP │ │
│ └─────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。
更多推荐
8
0- 0
已为社区贡献10条内容
所有评论(0)