【LLM学习笔记】第二篇：分布式训练介绍

1. 分布式训练概述

1.1 分布式训练的概念

• 定义: 分布式训练是指将机器学习或深度学习模型训练任务分解成多个子任务，并在多个计算设备上并行地进行训练。
• 计算设备: 包括中央处理器（CPU）、图形处理器（GPU）、张量处理器（TPU）和神经网络处理器（NPU）。
• 系统架构: 即使在同一服务器内部，多个计算设备之间内存也可能不共享，因此无论这些计算设备是否处于一个服务器还是多个服务器中，其系统架构都属于分布式系统范畴。
  

2. 分布式训练的必要性

• 算力需求: 单个计算设备的算力已经不足以支撑大规模模型的训练。从2013年的AlexNet到2022年的PalM模型，机器学习模型的参数量和数据量迅速增长，加剧了对算力的需求。
• 算力增长: CPU的算力增长远低于摩尔定律，尽管GPU、TPU等加速设备提供了大量算力，但其增长速度仍未能突破摩尔定律。因此，分布式训练系统是满足模型不断增长的算力需求的关键。

3. 分布式训练的目标

• 提升训练速度: 总训练速度可以通过以下公式简略估计：
  $$\text{总训练速度}\propto \text{单设备计算速度}\times \text{计算设备总量}\times \text{多设备加速比}$$
• 优化手段:
  • 单设备计算速度: 主要由单块计算加速芯片的运算速度和数据I/O能力决定。优化手段包括混合精度训练、算子融合、梯度累加等。
  • 计算设备数量: 计算设备数量越多，理论峰值计算速度越高，但通信效率会影响加速比。
  • 多设备加速比: 由计算和通信效率决定，需要结合算法和网络拓扑结构进行优化。

2. 大语言模型的分布式训练案例

2.1 OPT模型

• 硬件: 使用了992块NVIDIA A100 80G GPU。
• 并行策略: 采用全分片数据并行（Fully Sharded Data Parallel）和Megatron-LM张量并行（Tensor Parallelism）。
• 训练时间: 近2个月。
  

2. BLOOM模型

• 硬件: 使用48个计算节点，每个节点包含8块NVIDIA A100 80G GPU（总计384个GPU）。
• 通信: 节点内部使用4*NVLink通信，节点之间采用四个Omni-Path 100 Gbps网卡构建的增强8维超立方体全局拓扑网络。
• 训练时间: 3.5个月。
  

3. LLaMA模型

• 硬件: 使用NVIDIA A100-80GB GPU。
• 训练时间:
  • LLaMA-7B模型: 82432 GPU小时
  • LLaMA-13B模型: 135168 GPU小时
  • LLaMA-33B模型: 530432 GPU小时
  • LLaMA-65B模型: 1022362 GPU小时

LLaMA的全程是： Large Language Model Meta AI（羊驼 ）

3. 分布式训练的挑战

3.1 计算墙

• 问题: 单个计算设备的计算能力与大语言模型所需的总计算量之间存在巨大差异。
• 例子: NVIDIA H100 SXM的单卡FP16算力为2000 TFLOPs，而GPT-3需要314 ZFLOPs的总计算量，相差8个数量级。

3.2 显存墙

• 问题: 单个计算设备无法完整存储一个大语言模型的参数。
• 例子: GPT-3包含1750亿参数，采用FP32格式存储需要700GB内存，而NVIDIA H100 GPU只有80GB显存。

3.3 通信墙

• 问题: 各计算设备之间需要频繁地进行参数传输和同步，通信延迟和带宽限制可能成为瓶颈。
• 例子: GPT-3训练过程中，如果分布式系统中存在128个模型副本，每次迭代过程中至少需要传输89.6TB的梯度数据。截至2023年8月，单个InfiniBand链路仅能提供不超过800Gb/s带宽。

关于通信墙，此前我在 存算一体化系统（Processing-in-Memory, PIM）深入解析 文章中有介绍过。

4. 分布式训练的并行策略

1. 数据并行（Data Parallelism）

• 定义: 将数据切分，每个计算设备拥有完整的模型副本，并行处理不同的数据分片。
• 过程:
  • 每个计算设备分配到一批次数据的一部分，进行前向计算。
  • 计算设备根据本地数据计算梯度，并将梯度进行广播。
  • 所有计算设备聚合梯度，使用平均梯度更新模型参数。

2. 模型并行（Model Parallelism）

• 定义: 将模型中的算子分发到多个设备分别完成。
• 过程:
  • 模型的不同部分分布在不同的计算设备上。
  • 每个设备负责计算模型的一部分，并与其他设备通信以同步中间结果。

3. 混合并行（Hybrid Parallelism）

• 定义: 同时对数据和模型进行切分，实现更高程度的并行。
• 过程:
  • 结合数据并行和模型并行的优势，适用于大规模模型的训练。

5. 分布式训练的集群架构

5.1参数服务器架构

• 角色:
  • 参数服务器: 提供充足的内存和通信资源，负责参数同步。
  • 训练服务器: 提供大量的计算资源，进行模型训练。
• 训练模式:
  • 同步训练: 训练服务器完成一个小批次的训练后，将梯度推送给参数服务器，参数服务器聚合梯度并更新参数。
  • 异步训练: 训练服务器完成一个小批次的训练后，直接将梯度推送给参数服务器，参数服务器立即更新参数。
    

5.2 去中心化架构

• 定义: 没有中央服务器或控制节点，节点之间直接通信和协调。
• 优势: 减少通信瓶颈，提高系统的可扩展性。
• 通信原语:
  • Broadcast: 主节点将数据发送到其他节点。
  • Scatter: 主节点将数据切分并发送到其他节点。
  • Reduce: 将不同节点上的计算结果进行聚合。
  • All Reduce: 在所有节点上应用相同的Reduce操作。
  • Gather: 将多个节点上的数据收集到单个节点。
  • All Gather: 每个节点都收集所有其他节点上的数据。
  • Reduce Scatter: 将每个节点中的张量切分为多个块，每个块分配给不同的节点，并进行特定操作。

6. 高性能计算集群的硬件组成

• 单个服务器: 通常由2至16个计算加速设备组成，这些设备之间的通信带宽是影响分布式训练的重要因素。
• 通信带宽:
  • PCIe 5.0: 提供128GB/s的带宽。
  • NVIDIA H100: 采用高带宽内存（HBM），提供3350GB/s的带宽。
• 服务器内部架构: 服务器内部通常采用异构网络架构，如NVIDIA HGX H100 8-GPU服务器，采用NVLink和NVSwitch技术。

7. 总结

分布式训练通过将模型训练任务分解并在多个计算设备上并行执行，显著提升了训练速度和效率。然而，分布式训练面临计算墙、显存墙和通信墙等挑战，需要通过优化并行策略和集群架构来克服这些挑战。具体的应用案例展示了不同模型在分布式训练中的实际配置和效果，为未来的分布式训练提供了宝贵的参考。