分布式训练知识
基本概念
节点
节点是一个独立的物理服务器，通常包含：
- 多个CPU（如2颗Intel Xeon，共128个核心）
- 大容量内存（如1TB DDR4）
- 多个GPU（通常4或8张）
- 本地高速存储（如NVMe SSD）
- 网络接口卡（NIC）
一个节点的关键特征：内部的GPU之间通过超高速互联直接通信。
示例：
- NVIDIA DGX A100：一个节点 = 1台服务器 = 8张A100 GPU
- AWS p4d实例：一个节点 = 1台EC2实例 = 8张A100 GPU
集群的层次结构
在超算中心或云服务器中，计算资源是按层次组织的：
数据中心
↓
机架 (Rack)
↓
节点 (Node) ←─ 这是关键层级！
↓
GPU/加速卡
GPU之间的互联层次
-
层级1：节点内互联（极高速
- 技术：NVLink（NVIDIA）或 Infinity Fabric（AMD）
- 带宽：600 GB/s+（NVLink 4.0）
- 延迟：纳秒级
- 特点：GPU可以直接访问彼此的内存，像一个大GPU
-
层级2：节点间互联（高速网络）
- 技术：InfiniBand（IB）或 RoCE（RDMA over Ethernet）
- 带宽：200-400 Gb/s（约25-50 GB/s）
- 延迟：微秒级
- 特点：通过交换机连接不同节点
并行策略选择
正是因为节点内和节点间的通信性能相差一个数量级，所以并行策略需要精心设计：
原则：通信密集型操作放在节点内
-
张量并行必须放在节点内
- 原因：张量并行需要频繁交换大量中间激活值
- 在Transformer层中，每层的计算都需要多次All-Reduce通信
- 如果跨节点进行，带宽会成为致命瓶颈
-
流水线并行适合跨节点
- 原因：流水线并行通信的是层间的激活值
- 通信频率较低（每个微批次只通信一次）
- 对带宽要求相对较低，但对延迟敏感
-
数据并行相对灵活
- 可以在节点内，也可以跨节点
- 但通常为了管理方便，一个数据并行组内的GPU会在同一节点或相邻节点
训练范式
数据并行
- 思想：将训练数据划分到多个设备上，每个设备持有完整的模型副本，独立计算梯度，然后同步梯度并更新模型。
- 典型应用场景：适用于模型可以单卡存放，但数据量大的情况。
- 实现方式：
- All-Reduce同步：各卡计算梯度后，通过All-Reduce操作（如NCCL）求平均梯度，再各自更新参数。
- 参数服务器架构：Worker计算梯度，发送给Parameter Server聚合更新（在LLM时代较少用）。
假设我们有一个简单的神经网络：Y = W * X + b，其中：
W是权重矩阵，形状为[4, 2]X是输入数据，形状为[8, 4]（8个样本，每个样本4维特征）b是偏置- 我们使用2个GPU进行训练。
数据并行的过程：
-
模型复制：将完整的模型
{W, b}复制到 GPU-0 和 GPU-1 上。每个GPU上都有一个相同的模型副本。 -
数据分割：将输入数据
X和对应的标签Y沿着样本维度（第0维）平均分割。X被分割为X0（前4个样本，形状[4,4]）和X1（后4个样本，形状[4,4]）。- 将
(X0, Y0)发送到 GPU-0，将(X1, Y1)发送到 GPU-1。
-
独立前向传播：
- GPU-0 计算：
Output0 = W * X0 + b - GPU-1 计算：
Output1 = W * X1 + b
- GPU-0 计算：
-
独立计算损失和梯度：每个GPU根据自己的输出和标签，计算损失（Loss），并进行反向传播，得到针对自己那一半数据的梯度
Grad_W0和Grad_W1。 -
梯度同步（关键步骤）：这是数据并行的核心。两个GPU需要通信，将各自的梯度进行全局平均。
- 通过
All-Reduce集体通信操作，计算：Avg_Grad_W = (Grad_W0 + Grad_W1) / 2 - 最终，GPU-0和GPU-1都获得了完全相同的平均梯度
Avg_Grad_W。
- 通过
-
独立参数更新：每个GPU使用这个相同的平均梯度
Avg_Grad_W来更新自己本地的模型参数W和b。- GPU-0 更新：
W = W - learning_rate * Avg_Grad_W - GPU-1 执行完全相同的更新操作。
- GPU-0 更新：
结果：两个GPU上的模型参数在每次迭代后都保持完全一致，但它们处理的数据量是单个GPU的两倍。通信开销是梯度的大小。
模型并行
-
思想：将模型的不同部分放置在不同的设备上，每个设备只负责计算模型的一部分，前向和反向传播时需要跨设备通信。
-
细分类型：
张量并行
将单个层的矩阵运算拆分到多个设备上（例如，将一个大矩阵乘按行或列拆分）。
- 代表：Megatron-LM提出的Transformer层内拆分（如Attention头、FFN层分到不同设备）。
- 优势：通信主要在层内，通信量相对流水线并行较小，但需要设备间高带宽（如NVLink）。
现在假设我们的权重矩阵
W太大，单个GPU的内存放不下。W的形状为[4, 4]，我们决定使用2个GPU，以张量并行的方式拆分它。这里采用 Megatron-LM 中的列并行方式：
-
模型拆分（关键步骤）：将矩阵
W按列（第二维） 分成两半。W = [W_left, W_right]，其中W_left形状为[4, 2]，W_right形状为[4, 2]。- 将
W_left放置在 GPU-0 上，将W_right放置在 GPU-1 上。 - 每个GPU只持有模型的一部分。
-
数据广播：将完整的输入数据
X（形状[8,4]）广播到两个GPU上。两个GPU看到的是相同的数据。 -
并行计算（部分结果）：每个GPU用自己的那部分权重进行计算。
- GPU-0 计算：
Z0 = X * W_left（结果形状为[8, 2]） - GPU-1 计算：
Z1 = X * W_right（结果形状为[8, 2]） - 此时，
Z0和Z1分别是最终输出矩阵Z的左半部分和右半部分。
- GPU-0 计算：
-
通信与合并：为了进行后续的计算（例如激活函数），需要将两个部分结果合并。
- 通过
All-Gather通信操作，每个GPU都获得完整的Z。 Z = [Z0, Z1]，形状为[8, 4]。
- 通过
-
后续操作：对完整的
Z应用激活函数（如GELU），得到A = GELU(Z)。 -
反向传播的拆分：反向传播时，梯度也会按照正向传播的拆分方式进行流动和同步，最终每个GPU只更新自己持有的那部分参数（
W_left或W_right）。
关键点：在正向传播的每一步（如矩阵乘）中，两个GPU都在协作计算一个操作。通信发生在层内，传递的是中间激活值（如
Z0,Z1）。流水线并行
将模型按层切分到多个设备，形成流水线，不同设备处理不同的micro-batch。
- 代表：GPipe，PipeDream。
- 挑战：存在流水线气泡（bubble），设备利用率不足；需要精心安排调度策略减少空闲时间。
假设我们有一个简单的4层神经网络：
输入X → Layer1 → Layer2 → Layer3 → Layer4 → 输出Y假设我们有2个GPU，网络分为2部分：
- GPU1: Layer1, Layer2
- GPU2: Layer3, Layer4
正向传播流水线（假设batch_size=8，分成4个micro-batch，每个size=2）：
时间 GPU1操作 GPU2操作 ------------------------------------------------------ t1: F(L1+L2, mb0) → 输出A0 → 空闲 t2: 发送A0到GPU2 接收A0, F(L3+L4, A0)=Y0 t2: F(L1+L2, mb1) → 输出A1 → 空闲 t3: 发送A1到GPU2 接收A1, F(L3+L4, A1)=Y1 t3: F(L1+L2, mb2) → 输出A2 → 空闲 t4: 发送A2到GPU2 接收A2, F(L3+L4, A2)=Y2 t4: F(L1+L2, mb3) → 输出A3 → 空闲 t5: 发送A3到GPU2 接收A3, F(L3+L4, A3)=Y3反向传播流水线（与前向方向相反）：
时间 GPU1操作 GPU2操作 ------------------------------------------------------ t5: 计算Loss，B(L3+L4, Y3) t6: 接收梯度G3，B(L1+L2, G3) 发送梯度G3到GPU1 t6: 计算B(L3+L4, Y2) t7: 接收梯度G2，B(L1+L2, G2) 发送梯度G2到GPU1 ...以此类推
训练方法
单进程多GPU训练
- 特点
- 单进程控制多个GPU：一个Python进程管理所有GPU
- 主GPU收集梯度：在前向传播时，数据被拆分到不同GPU；反向传播时梯度汇总到主GPU进行参数更新
- 特征
- 程序内部使用
DataParallel包装模型 - 用普通的
python命令启动程序
- 程序内部使用
transformers库相关参数self.args.n_gpu等于GPU数
分布式训练
- 特点
- 多进程架构：每个GPU对应一个独立的进程
- 点对点通信：使用NCCL等通信库，进程间直接通信
- 内存效率高：每个GPU只存储自己的模型副本和优化器状态
- 真正的数据并行：每个进程独立处理数据，梯度通过all-reduce操作同步
- 特征
- 用DDP包装模型
- 使用
torchrun启动程序
transformers库相关参数self.args.n_gpu等于一self.accelerator.num_processes等于进程数（也就是GPU数）
- 只要使用了deepspeed，那么就会使用分布式训练
浙公网安备 33010602011771号