实战国产千卡集群训练，要啥自行车！
实战国产千卡集群训练，要啥自行车！
背景
近年来，大规模语言模型（Large Language Model, 简称LLM）在对话、问答和文本总结等多种应用场景中展现出卓越性能，受到了学术界和工业界的广泛关注。然而，随着LLM逐渐朝着更多的参数量和更长的文本进行演化，对计算设备的存储和处理能力也提出了更高要求。当前，传统的流水线并行训练方法在处理百、千亿参数的模型时会产生存储和计算负载不均衡的现象，直接影响了资源利用率和整体训练效率；同时，由于现有国产算力卡存在高速内存容量和通信能力不足的情况，使得这一问题显得更加突出。
为解决上述问题，陈文光教授团队开发了AdaPipe技术，该技术根据具体的模型和硬件参数，通过细化重计算的粒度来优化重计算策略，并考虑到训练各阶段计算量的差异，进一步对重计算和流水线切分策略进行优化。该技术不仅将存储资源的利用率最大化，还确保不同计算节点上的计算负载得以平衡分配，训练效率显著提升了。
实验平台
我们在两个集群上评估 AdaPipe。
集群 A。8 台 英伟达（NVIDIA）GPU DGX-A100 服务器节点，每个节点配备 8 × A100 80GB 加速器 （GPU），双路 64 个CPU 内核和 2 TB 内存。八个 GPU 通过 NVLink 连接。节点互联通过 800 Gbps Infiniband HCA 。
集群 B。 32 台华为 Atlas 800 服务器节点，每个节点配备 8 个× Ascend 910 32GB 加速器 （NPU）、四路 192 个 CPU 内核和 2TB 内存。每台服务器中 8 个 NPU 安装在两个 NPU 板上，每个板上的 4 个 NPU 通过 30 GB/s 链路 Full Mesh 全互联。此外，每个 NPU 通过一个 PCIe 4.0 x16 链路 （32GB/s） 直接连接到 CPU，并自带一个用于节点间组网的 100 Gbps NIC。
分别在两个代表性模型即 GPT-3 （175B） 和 Llama 2 （70B）上评估 AdaPipe。在集群 A 上，GPT-3 和 Llama 2 的评估使用了 64 个和 32 个 GPU。在集群 B 上，我们分别使用 256 和 128 个 NPU 进行小规模实验，以及使用 2048 个和 1024 个 NPU 进行大规模GPT-3 和 Llama 2 实验。此外，在集群A上进行的所有实验都使用PyTorch框架，而在集群B上进行的实验则使用MindSpore框架。
论文显示，AdaPipe 在英伟达 GPU 和 Ascend NPU 上的端到端训练性能分别提高了 1.32 倍和 1.22 倍。
此外，AdaPipe在“鹏城·脑海”通用大模型（200B）4K窗口3456卡的实际训练中，取得了10%以上的效率提升。这些案例将为未来万卡国产集群上的优化训练提供技术储备和经验参考。