直击算力焦虑，这家清华系 AI 创企搞了个千卡异构混训平台，算力利用率最高达 97.6%
7 月 4 日，在 2024 年世界人工智能大会 AI 基础设施论坛上，无问芯穹联合创始人兼 CEO 夏立雪发布了全球首个千卡规模异构芯片混训平台，和同构训练相比，千卡异构混合训练利用率最高达到了 97.6%。这也是全球首个可进行单任务千卡规模异构芯片混合训练的平台，具备万卡扩展性，支持包括 AMD、华为昇腾、天数智芯、沐曦、摩尔线程、NVIDIA 六种异构芯片在内的大模型混合训练。
当前，全球范围内在大模型训练领域，主流做法是结合 PyTorch、Megatron 等先进的计算框架与英伟达（NVIDIA）的 GPU 芯片进行高效分布式训练。尽管英伟达 GPU 在大模型训练所需的算力供给中占据主导地位，近年来中国本土芯片制造商也逐渐崭露头角，研发出一系列高能效的 AI 处理器，包括华为昇腾、寒武纪 MLU、燧原科技、海光 DCU、摩尔线程等品牌。这些国产芯片不仅丰富了市场选择，也为大模型的训练提供了多样化的算力支持。
然而，在实际操作中，大模型训练往往需要数百乃至数千张 GPU 卡的并行运算，以 Llama3-70B 这一特定模型为例，其庞大的 700 亿参数意味着训练过程至少需要 900 张 H100 GPU 卡持续工作 10 个月。但由于多种原因，可能无法集齐如此大量的同型号 GPU 卡，这种情况的出现可能性更大：仅有 500 张 NV H100（每张提供 969 TFLOPS 算力）、400 张华为 Ascend 910B（每张 256 TFLOPS）以及 700 张 AMD MI250X（每张 383 TFLOPS）。虽然任何单一类型的 GPU 卡数量均不足以独立完成 Llama3-70B 模型的训练，但若将这三种不同品牌、不同规格的 GPU 卡进行混合使用，其总体算力则足以覆盖该大模型的训练需求。
然而由于不同芯片厂商的芯片体系结构及软件生态的差异性，导致异构卡混合训练大模型存在非常大的挑战。7 月 4 日，在 2024 年世界人工智能大会上，无问芯穹发布全球首个千卡规模异构芯片混训平台，和同构训练相比，千卡异构混合训练利用率最高达到了 97.6%，是全球首个可进行单任务千卡规模异构芯片混合训练的平台，具备万卡扩展性，支持包括 AMD、华为昇腾、天数智芯、沐曦、摩尔线程、NVIDIA 六种异构芯片在内的大模型混合训练。
本文将深入解读这一重大进展背后的技术思路，由无问芯穹与清华、上交的联合研究团队呈现。
异构芯片间的混训主要面临两大挑战，一是异构卡通信库差异，导致异构卡之间通信难，二是异构卡之间性能差异，导致模型分布式训练低效。
多芯片互通互联的复杂性：如何实现不同芯片间高效、稳定的通信；
算力不均衡导致性能损失：不同芯片的计算能力存在天然的差异，这种不均衡性会直接影响到模型训练的整体效率，如何根据合理的并行策略进行任务的切分。
异构卡通信库的差异导致异构卡之间通信难，通常会成为大模型训练的瓶颈。不同厂商的 GPU 或 AI 芯片为了优化通信，通常会构建自己的通信库，在内部架构和通信接口上存在差异，导致不同通信库之间无法通信。这要求在构建大规模混合训练系统时，必须设计出高度兼容且高效的互联互通方案，确保多芯片间的高效、容错及稳定性通信，这需要复杂的软件栈支持和细致的网络配置，以最小化通信延迟和带宽瓶颈。
为了解决不同类型信显卡之间的通信问题，我们实现了无穹集合通信库 (IHCCM)，它支持基于 CPU 或者基于 GPU 的两种通信方式。
图注：IHCCM 通信库架构
如下图所展示，IHCCM 通信架构注重了异构集群的可扩展性，能够支持新型芯片快速、低成本的部署到集群，为异构计算环境下的高性能训练系统打下基础。
在此架构中，同构节点内部，GPU 与 CPU 之间通过高速 PCIe（Peripheral Component Interconnect Express）交换机进行通信。跨节点通信则采取了更为灵活的策略，不同节点的 CPU 之间，无论是通过 IPoIB（IP over InfiniBand）适配器或者以太网，均可实现高效的数据传输。这一设计充分考虑了异构环境的复杂性与多样性，确保了在架构迥异的芯片间维持稳定而高效的通信链路。
具体通信流程为：当模型参数需要从源 GPU 迁移到目标节点的异构 GPU 时，首先会经由 PCIe 通道，将数据从 GPU 复制到源节点的 CPU，随后通过 IPoIB 或以太网跨越节点边界，抵达目标节点的 CPU。在这里，数据再次借助 PCIe 的高速能力，被复制至目标 GPU 上，整个过程中，CPU 起到了至关重要的中转与协调作用。通过这一系列步骤，有效地规避了不同 GPU 间通信库差异带来的潜在障碍，尽管引入了额外的 CPU-GPU 数据拷贝开销，但这在很大程度上换取了异构混训系统的整体稳定性和兼容性，为大规模、跨架构的模型训练提供了坚实的基础。
图注：基于 CPU 的通信流程
针对异构 GPU 环境下大规模模型训练所面临的通信瓶颈，我们设计了一个基于 RDMA（Remote Direct Memory Access）技术的创新通信架构，将不同 GPU 芯片间通过 InfiniBand（IB）网络进行高效数据交换。鉴于大模型训练中计算与通信的强耦合特性，通信效率直接关系到整体训练性能，因而各大芯片制造商纷纷研发专属通信库，旨在提升自家 GPU 的通信效能。然而，这些通信库在 API、数据类型等层面的差异化，构成了异构 GPU 间直接通信的主要障碍。
为破局这一困境，我们制定了一套标准化的分布式通信接口规范，该协议涵盖了基础的点对点通信功能，如异步发送（iSend）、异步接收（iReceive），以及集合通信机制，包括异步全规约（iAllReduce）、异步全播送（iAll-to-All）等高级功能，确保了数据在异构集群内能够顺畅流通。这套统一的通信协议，不仅显著提升了通信效率，还为异构 GPU 间的协作开辟了全新路径。
值得注意的是，这一方法离不开对各厂商集合通信库的深度适配与优化。通过将无穹集合通信库（IHCCM）应用到训练过程，我们不仅实现了异构 GPU 间的通信，更通过时间重叠策略（即将计算过程与通信过程互相重叠）最大程度地减少了通信延迟对整体训练性能的影响，从而显著提升了大规模模型训练的效率。
图注：基于 GPU 的通信流程
在异构分布式混合训练系统中，硬件之间具有不同的特性，系统如何在算力分布不均匀的情况下，根据模型特征和系统的实时训练状态对任务进行均匀或非均匀的切分，保证算力的最大化利用，也是十分具有挑战的。
在大语言模型的训练中，采用基于流水线并行的异构训练方案，通过非均匀拆分 transformer 层，可有效应对算力不均衡的难题。可以根据各芯片的实际算力，让算力较高的芯片处理更多的层，算力较低的芯片处理更少的层，从模型结构角度，每层的计算量是一致的，所以按算力比作为层拆分比例可以达到最好的效果。例如，假设芯片 A 的算力是芯片 B 的 4 倍，则可以让芯片 A 计算 4 层 transformer 模块，让芯片 B 计算 1 层 transformer 模块，以此达到资源的最优配置，理论上可实现最高的吞吐性能。
基于数据并行的异构训练，同样可通过调整不同芯片上的数据批处理规模，来应对算力差异。例如，在芯片 A 算力为芯片 B 的 4 倍情形下，可以让芯片 A 每轮迭代计算 4 个 batch，而芯片 B 仅计算 1 个 batch，以此平衡各芯片的工作负载，达到理论上的最优吞吐率。
然而，理想与现实之间总存在差距。在实际操作中，还需综合考量诸多复杂因素：
芯片在流水线并行中的位置顺序：不同位置的芯片可能需要前后节点的数据作为输入，进而影响其实际的计算效率。
模型前后 embedding 层的计算量：这些层的计算强度可能与 transformer 层存在差异，需额外关注。
显存大小与重算策略：根据芯片显存容量，可能需要启用重算机制，以节省显存占用，但这可能会影响整体性能。
单款芯片算子优化程度：不同芯片对特定算子的执行效率不尽相同，需针对性地进行优化，以发挥最大效能。
以上种种，均可能间接改变两款芯片在实际训练中的性能比例，对模型非均匀拆分后异构集群的整体训练吞吐量产生影响。
图注：两种训练并行策略
此外在大模型高效训练的过程中，合理选择并组合数据并行、流水线并行及张量并行等并行策略，也是突破性能瓶颈的关键。
数据并行可以将庞大的数据集分割到多个GPU，允许并行处理，从而加速训练速度；
流水线并行则针对结构较深的模型，通过将模型结构划分为多个段落，在不同GPU上依次处理，实现计算资源的高效利用；
张量并行则聚焦于模型参数层面的切分，进一步细化并行粒度，提升计算密集型任务的处理效率。
本项工作目前的系统框架可以在任意两种不同硬件组成的异构混训集群上实现 llama2-7B/70B 的训练，通过自动并行策略的选择和性能模拟器筛选，在千卡异构集群上找到了最优的分布式策略，实测得到的算力利用率为 97.6%。
无问芯穹与清华、上交的联合研究团队表示，希望能够通过整合更多异构算力，继续推高大模型技术能力的上限，打破单一芯片品牌训练资源瓶颈。在当前取得的成果基础之上，持续探索如何进一步提升异构集群的算力整合效率，构建适应多模型与多芯片格局的 AI Native 基础设施。
AIGC技术正以惊人的速度重塑着创新的边界，InfoQ 首期《大模型领航者AIGC实践案例集锦》电子书，深度对话30位国内顶尖大模型专家，洞悉大模型技术前沿与未来趋势，精选10余个行业一线实践案例，全面展示大模型在多个垂直行业的应用成果，同时，揭秘全球热门大模型效果，为创业者、开发者提供决策支持和选型参考。关注「AI前线」，回复「领航者」免费获取电子书。
在主题演讲环节，我们已经邀请到了「蔚来创始人 李斌」，分享基于蔚来汽车 10 年来创新创业过程中的思考和实践，聚焦 SmartEV 和 AI 结合的关键问题和解决之道。大会火热报名中，7 月 31 日前可以享受 9 折优惠，单张门票节省 480 元（原价 4800 元），详情可联系票务经理 13269078023 咨询。
今日荐文
微信扫码关注该文公众号作者