移动创新提出智算异构混合并行训练技术,攻克资源混池训练难题
为实现大模型训练任务在多种架构智算芯片上的计算负载均衡及高速通信,攻破异构智算集群软硬件合池壁垒,提升多样性算力利用率,中国移动创新提出智算异构混合并行训练技术,从非均匀计算任务拆解及训练数据高速协同等角度构建混训技术能力。图1为智算异构混合并行训练技术总体架构,技术关键层为混合训练框架层、混合训练通信层,其中框架层实现异构混合分布式并行训练计算策略拆解及最优性能预测策略执行,通信层实现异构智算芯片在训练任务执行过程中的数据传输协同
图1 异构混合分布式训练技术总体架构
异构混合分布式训练技术总体架构关键技术层介绍如下:
混合训练框架层
异构混合训练框架层用于实现模型训练任务并行化拆解,在数据、模型等维度进行任务切分,切分后的任务与适配算力映射,实现训练过程的分布式并行处理。
由于异构混合训练需将切分任务映射到不同类型、规模的异构算力集群,当前业界所使用的分布式训练框架仅适用于同构算力集群环境,无法复用在异构混合训练场景中,为解决行业技术空白,在框架层创新设计了可适配异构算力集群的通用异构混训框架,通过该框架实现数据、模型等维度的非均匀任务切分。同时为实现非均匀切分性能最优,在框架层设计了异构非均匀切分性能预测工具,在给定模型规模及集群规模条件下,通过策略生成算法以及性能模拟工具分析生成适用该模型、集群的最优非均匀切分策略。
(1)非均匀计算任务切分
混合分布式并行训练非均匀计算任务切分算法对齐同构集群训练方法,对原始并行训练策略算法进行非均匀切分优化。面对同一训练任务下的深度神经网络,由于混合分布式并行技术算力集群是异构的,需对隔离后不同的算力子集群放置不同的子神经网络(后统称“子网络”),基于不同子网络的计算量差异,对计算需求较大的子网络放置在算力较高的子集群中,对计算需求较小的子网络放置在算力较低的子集群中,最终实现混合分布式训练任务切分负载均衡。
图2.非均匀计算任务切分技术原理图
(2)性能预测技术
异构非均匀切分性能预测工具接收神经网络模型结构、单节点智算芯片硬件标称值,结合集群中实际的通信带宽和训练框架适配的切分算法,搜索多种并行策略组合,对每种可行的策略进行建模,处理为策略树,并将策略树解析为算子、张量执行拓扑图,通过分层拓扑感知执行器模拟整个模型的训练过程,并通过分析器实现不同非均匀切分策略下的性能分析,性能分析维度包括硬件显存利用率、计算吞吐率、通信带宽等,综合分析得到最优非均匀切分并行配置信息,并输出最优训练配置文件到异构混合训练框架中执行。
图3.性能预测技术原理图
混合训练通信层
异构混合训练通信机制协同异构混合训练框架层,在框架层非均匀切分基础上进行并行训练计算过程中间值传输同步,并基于异构混合训练任务调度映射策略实现调用智算硬件后端通信能力。
由于不同异构智算芯片在通信协议、互联方式、通信策略等层面存在较大区别,需要一套高速统一的数据协同策略,通过发现混合训练通信拓扑、实现通信初始化链接、配合非均匀切分算法制定通信策略、定义通信标准约束,解决各类智算芯片在通信机制层面的差异,实现异构混合训练通信能力。
图4.混合训练通信技术原理图
(1)通信拓扑发现
混合通信数据传输的首要步骤是识别智算资源池中的异构设备及设备间的连接方式、传输带宽,混合训练机制通过GPU Driver获取GPU拓扑信息,识别每个异构智算子集群节点的智算芯片及类型。
(2)数据传输建链
混合训练数据传输需要建立特定传输通道,包括物理互联方式定义及数据收发方式定义。在物理互联方式定义上,集群单节点内部采用PCIe或特定智算芯片卡间高速互联协议,节点间通过RDMA网络(IB/RoCE)进行消息传递,实现异构节点物理互联。
在数据收发方式上,数据传输通道由节点数据发送端、节点数据接收端、节点间连接通道组成,节点数据发送端以GPU最小计算处理单元为粒度进行传输数据分割,在发送端完成传输数据存储,并通过IB/RoCE网络进行消息传递,当节点数据接收端完成数据接收准备后,接收来自数据发送端的数据,建立数据传输通道。
以下图5为例,发送端为智算芯片A,接收端为智算芯片B,综合前面获取的通信拓扑信息、GPU最小计算处理单元SM(Streaming Multiprocesser,流式多处理器)计算能力等确定在SM0-SMn分别处理的数据微批量及buffer大小,通过对应的QP完成数据发送准备,实现建立数据传输通道Channel0-Channeln。
图5 混合训练数据传输通道说明
(3)传输策略选择
在建立数据传输通道和基本收发能力基础上,结合异构混合训练任务非均匀切分策略,进行数据传输策略设计。
在异构混合训练任务非均匀切分场景中,并行训练策略依据神经网络计算量及智算芯片硬件算力综合评估得出,异构硬件单节点计算能力和通信策略实现算法各异,不同并行训练策略对通信策略算法的需求也各不相同,为解决异构混合训练问题,创新设计了一套可适配不同并行训练策略的通信操作检索层,通过为特定并行训练策略下检测到的通信操作设计通信算法,包括点对点通信、集合通信策略等,使得该通信算法能够尽可能解决异构节点计算及通信能力不均衡的问题。
在架构上,系统单独设计了通信操作检索层(图6),通信操作检索层获取任务分发模型拆解的通信队列信息,例如在特定模型LLaMA预训练任务下,检索出流水线并行策略包含layer间点对点通信、初始化阶段Allgather同步各节点信息、GradNorm阶段Allreduce实现梯度正则化规约操作等,表示为并行策略::layer层::通信算子,如PP::GradNorm::Allreduce。
通信策略设计层将各涉及到的通信操作进行最优策略算法设计,完成特定模型训练任务下的异构混合训练通信实现。
图6 通信传输策略定义流程
(4)训练数据协同
定义好混合训练数据传输策略后,需经由通信接口完成训练任务实时数据协同。混训通信接口作为多厂商可识别的通信标准约束,其将数据传输策略、communicator创建及管理等进行接口定义,定义后的接口可在不同类型智算芯片上运行,通过智算芯片驱动自动安装,经由编译、运行实现在智算芯片上的通信操作。
多芯混训通过合池集群验证,可实现大模型低损耗训练
为验证技术功能及性能效果,中国移动研究团队在Nvidia GPU和天数智芯、壁仞、瀚博、海光等国产智算芯片所组成的混合算力资源池中,开展对LLaMA2等模型混合交叉训练技术试验,验证过程及结论如下:
验证结论一:多芯混训效率损耗较低,可满足大模型训练需求
为评估异构混合训练效率,采用并行系统集群加速比作为评价指标,其计算方式如下:
其中,C为混训实测集群吞吐TPS(tokens/cluster/s),为混训合池性能理论吞吐TPS,其中为特定智算芯片在单一集群上的吞吐值,混合集群包含α类智算芯片。
当实测合池集群吞吐C接近理论吞吐,说明采用异构混合训练后性能相比各集群单一性能总和来说,计算效率损耗较小,越接近理论吞吐表明异构混合训练并行策略优化效果越好。
根据实测验证,通过预设定LLaMA2等大模型策略生成模型特征参数,可实现针对不同异构算力的任务拆解及分发协同,且混合训练集群加速比可达95%以上,效率损耗较低,可基本接近单一类型智算集群训练效果。
验证结论二:模型收敛度及精度正常,合池不影响模型训练性能
为评估模型训练性能,采用Loss曲线、LLM语言模型困惑度PPL作为评价指标。其中,Loss曲线用于验证模型在混合训练条件下能否正常收敛,PPL用于验证模型是否有较好的语言数据预测能力,PPL值越低模型预测能力越好。
为体现模型收敛趋势,通过每隔n次迭代记录Loss值进行Loss曲线绘制;为体现模型预测能力,采用PPL曲线反映语言模型性能结果,试验过程将数据集按照特定比例分割为训练数据集和评估数据集,各n次迭代进行一次模型评估,绘制PPL曲线。
根据实测验证,混训条件下LLaMA2模型Loss曲线呈下降趋势,可正常收敛(如图7),且PPL曲线呈现下降趋势(如图8),合池训练对模型性能没有影响。
图7.LLaMA2模型混合训练Loss曲线
图8.LLaMA2模型混合训练PPL曲线
思考及展望
本文介绍了智算异构混合并行训练技术,该技术可通过非均匀计算任务切分算法、非均匀切分性能预测技术、混训通信技术等实现不同类型智算芯片在单一大模型任务下的混合分布式训练。
当前混合训练技术研究侧重基础通用能力构建,在并行策略优化、多类型模型支撑、大规模集群验证等方面仍需进一步开展相关工作。后续研究团队将持续深入探索智算异构混合并行训练机制,从新增功能、优化性能、提高效能三方面进行技术攻关和产品能力升级,同时构建混合训练系列标准,推动智算混训产业生态构建,并逐步拓展验证方案及模型场景,攻破大模型混合训练系列挑战,加速技术能力落地,助力万亿级参数大模型发展,为国家级人工智能训练场贡献移动方案。
审稿:张昊 | 网络与IT技术研究所
本文作者