LLM高性能并行训练技术

LLM高性能并行训练技术

### 研究背景与意义

- **深度学习的重要性**：人工智能成为国际竞争焦点，深度学习是其核心技术，在众多领域取得突破，推动社会向智能化跃升。

- **面临的挑战**：数据、模型规模呈指数级增长，硬件算力发展滞后。单个 GPU 难以满足大规模模型训练需求，分布式训练面临通信和内存瓶颈。

相关技术与研究现状

- **分布式训练

- **数据并行**：每个计算节点持有模型完整副本，处理小批量训练样本子集，通过 All - Reduce 或参数服务器架构实现全局梯度通信。已有工作探索通信优化技术，如梯度稀疏化、通信调度等，但需平衡压缩率和精度损失，且不同互联架构影响性能。

- **模型并行**：将模型**参数张量划分到多个计算节**点，相应划分或复制特征映射张量和误差张量。可降低单个设备内存需求，但通信成本高，不同模型结构适用情况不同。在框架层面，有多种支持模型并行的框架，如 DistBelief、Mesh - Tensorflow、Megatron - LM 等。

- **流水线并行**：将模型纵向划分为多个子网络分配到计算节点，批量样本分微批以流水线方式输入。**Gpipe 是首个研究，但存在性能缺陷，如 GPU 利用率低、层间分区不均衡。异步流水线并行虽受关注，但算法设计复杂，需降低对模型收敛的影响**。

- **显存优化**

- **激活重计算**：在前向计算存储部分激活张量，反向传播重计算释放的张量。可分为线性化网络和一般图两类方法，不同方法有其特点和适用场景，部分方法已集成到主流框架。

- **张量卸载**：源**于虚拟内存技术，将暂时不需要的张量卸载到 CPU 内存**。针对卷积神经网络和模型参数有不同卸载策略，但存在确定卸载对象和时机的关键问题。

- **低精度训练**：用**低精度格式表示训练数据，降低内存占用提升性能**，但需引入误差补偿技术维持准确率。

混合并行训练策略

- **成本模型构建**：深入分析模型训练中层内（数据并行和模型并行）和层间（流水线并行）并行性，分别构**建前向计算与反向计算的成本模型**。层内并行成本模型考虑数据交换的**通信对象、时机和开销**；层间并行成本模型考虑微批量训练数据在各子网络的**执行时间和数据传输开销**。

- **混合并行分区算法**：基于动态规划算法搜索混合并行策略空间，将模型分区映射到计算设备。相比暴力搜索的指数级复杂度，该算法时间复杂度为\(O(K^{2}N^{2})\)，有效降低了策略搜索成本。

- **实验验证**：在多个深度学习模型和分布式环境下验证有效性。**与数据并行和 Gpipe 相比，分别实现高达 1.88 倍和 2 倍的加速比**，证明了该策略能根据不同环境和模型给出高效并行方案。

- **Transformer 层内模型并行优化策略**

- **细粒度流水线执行**：引入**细粒度层内流水线执行**，解耦层间计算和通信依赖。通过将每层操作沿样本维度划分，设计动态规划分区算法确定最优分区方案，使计算和通信开销最大程度重叠，搜索复杂度为\(O(L)\)。

- **混合通信压缩机制**：联合 Token 级的 **Top - k 稀疏化和分段量化方法**。Token 级 Top - k 稀疏化根据注意力分数评估 Token 重要性，减少通信开销，复杂度远小于元素级；分段量化根据张量数值分布设计，减少与原始张量的均方误差。

- **实验验证**：在不同规模 Transformer 模型上实验，与 Megatron - LM 相比，AccTFM 训练吞吐量**提高 2.08 倍**，且基本不影响模型收敛性。**消融实验表明各模块能提供显著性能收益。**

- **异构训练策略**

- **参数分桶与卸载策略**：引入参数分桶机制最大化 PCIe 带宽利用率。设计贪心算法确定卸载参数对象及卸载和预取时机，考虑 GPU 计算与内存拷贝异步执行，将参数卸载问题建模为优化问题求解。

- **参数预分配算法**：基于确定的卸载策略，分析训练过程中参数共存情况，设计参数预分配算法复用 GPU 内存空间，解决内存碎片化问题。

- **实验验证**：在不同规模 Bert 和 GPT 模型上验证。与 Zero - Offload 相比提升了单机单卡可训练模型规模，与 Gemini - CPU 相比在相同内存预算下实现 2.13 倍性能提升。