多芯异构混池训练技术及实践

        随着大模型参数量的不断增长，前期投入建设的大模型基础设施会面临资源不足、需要扩容的需求。智算芯片厂商不断推出新品，且从供应链安全角度考虑会引入多家GPU厂商，同一个集群中会出现同一厂商不同代际的GPU硬件或不同厂商的GPU硬件。这给智算集群的系统集成和优化管理带来了一系列挑战，需要探索异构GPU集群场景下的混池训练解决方案。

异构GPU混池训练的挑战

        异构GPU的算力、显存容量、显存带宽、数据传输速率各不相同，而目前主流AI框架多维并行模型拆分时，只支持按照AI模型基本要素均匀拆分，并行策略比较简单；不同厂家的芯片架构和软件生态差异大，异厂家软件栈独立无法互通。异构GPU混池训练存在如下挑战：

• 资源利用率低

        只能同构GPU下实现并行训练，模型均匀拆分时，需要按照低配GPU设计模型拆分参数。如A40只有48G显存，A100有80G显存，则A100会出现32G的显存浪费，因此模型均匀拆分会导致部分GPU资源浪费。

• 跨厂家通信不兼容

        异厂家GPU混池场景下，集合通信无法互通。

• 精度损失风险

        异构GPU加速器上算子的精度差异可能会使模型的精度难以达到期望同步的精度。

• 优化策略困难

        随着异构GPU加速器、模型层或运算符的数量增加，可选的分布式训练策略数量会呈指数级增长，这导致要找到最优的训练策略变得非常困难。

        异构GPU混池训练的主要挑战在于如何充分利用异构资源，保证跨厂家通信和模型精度，同时寻找最优的分布式训练策略。

异构GPU混池训练方案

        为应对上述挑战，需考虑在通用AI框架基础上增加插件包，支持对模型的DP（数据并行）/PP（流水线并行）非均匀拆分；资源管理平台感知底层GPU硬件能力，与AI框架协商，实现最优策略编排调度；基于统一的集合通信库实现异厂家GPU集合通信互通。

        中兴通讯异构GPU混池训练方案（见图1），主要包括PP非均质拆分、DP非均质拆分、统一集合通信互通以及异构混训超参自动寻优等技术。

PP非均质拆分

        PP非均质拆分通过训练框架和资源调度的双向协商，完成自动非均质拆分，根据不同GPU的显存以及算力灵活分配layer数，以达到最佳资源利用率；支持模型拆分到不同型号的GPU上，并且按照GPU规格拆分模型。具体策略如下：

• 在显存允许的情况下，不同PP实例之间的算力比值就约等于不同卡PP layer拆分的比值。
• GPU卡算力高，显存小，采用显存优化措施，如选择性/全量重计算、offload等方式降低显存使用。
• PP并行场景下，靠前的卡有显存峰值，所以通常把显存充裕的卡放到PP靠前stage进行编排。

DP非均质拆分

        DP非均质拆分按照每条DP的算力比例分配样本数量，以便实现梯度累计时间对齐。具体策略如下：

• 通过调整Micro-Batch Size（Micro-BS）的大小来平衡每路DP的前向与后向的计算时长，从而平衡AllReduce的时长。在训练时有梯度累积（GA），需要使每路DP最后一个Micro-BS计算完梯度的时长尽量接近，从而降低某一路DP的GPU空闲等待另一路DP GPU的时长，提升训练效率。
• 通过调整Mini-Batch Size（Mini-BS）的大小来平衡Allreduce的时长。对于每一路DP，分配不同Mini-BS的数据，结果表现为每路DP的GA次数不一样。例如：两种GPU算力比是70:58，无法通过调整Micro-BS实现，但是可以保持Micro-BS为1:1，通过调整GA为70:58，即Mini-BS为70:58，精细化实现算力比和数据比的严格对齐，进而达到平衡AllReduce时长的目的。

统一集合通信互通

        AI平台纳管异构GPU，统一集合通信库支持跨节点异构GPU集合通信互通能力。集合通信路径主要有如下几种方式：

• H2D（Host to Device）：主机1的GPU A显存数据经CPU送给NIC/RDMA NIC发送，主机2的NIC/RDMA NIC接收后经CPU送至GPU B，通过普通NIC发送和接收的为TCP方式，通过RDMA NIC发送和接收的为RDMA方式。
• GDR（GPU Direct RDMA）：主机1的GPU A显存数据采用GDR方式通过RDMA NIC至主机2的RDMA NIC，直接拷贝到GPU B显存。

异构混训超参自动寻优

        实现对单一任务多种芯片的自适应并行策略搜索功能，通过计算各种并行策略所需要的计算量、存储量、通信量，以及不同芯片的实测算子和通信性能，快速计算出最优的任务切分策略。

异构GPU混池训练实践

        中兴通讯联合实验室针对同厂家异代（英伟达）以及异厂家异构（壁仞+英伟达）两种场景，采用Llama2模型进行PP/DP非均质拆分验证，性能达到良好效果，尤其是H2D使用RDMA NIC进行异厂家集合通信时性能得到明显提升（训练效率达到理论上限的95%以上）。验证结果如表1所示。后续我们会基于GDR集合通信方式进一步提升混池训练的性能。

        此外，2024数字科技生态大会期间，中国电信联合中兴通讯及多个智算生态伙伴正式发布了面向大模型的“智算异构四芯混训解决方案”，该方案采用统一训练框架、统一集合通信库、统一RDMA网络。

        异构GPU混池训练虽面临诸多挑战，但中兴通讯异构GPU混池训练方案通过一系列创新技术有效应对，并展现出良好效果。未来，随着技术的不断发展与实践的深入推进，异构GPU混池训练方案有望解决异构芯片间生态竖井、算力孤岛难题。