近日，在美国圣迭戈市举行的计算机体系结构重要国际会议ASPLOS上，由鹏城实验室智能计算研究部、清华大学计算机系教授陈文光团队自主研发的国产千卡集群训练优化技术AdaPipe正式发布。

近年来，大规模语言模型在对话、问答和文本总结等多种应用场景中展现出卓越性能，受到了学术界和工业界的广泛关注。然而，随着大规模语言模型逐渐朝着更多的参数量和更长的文本进行演化，对计算设备的存储和处理能力也提出了更高要求。

当前，传统的流水线并行训练方法在处理百、千亿参数的模型时会产生存储和计算负载不均衡的现象，直接影响了资源利用率和整体训练效率；同时，由于现有国产算力卡存在高速内存容量和通信能力不足的情况，使得这一问题显得更加突出。

为解决上述问题，陈文光团队开发了AdaPipe技术，该技术根据具体的模型和硬件参数，通过细化重计算的粒度来优化重计算策略，并考虑到训练各阶段计算量的差异，进一步对重计算和流水线切分策略进行优化。该技术不仅将存储资源的利用率最大化，还确保不同计算节点上的计算负载得以平衡分配，训练效率显著提升了。

该研究显示，AdaPipe支持GPU、NPU主流加速卡，该技术被应用在“鹏城云脑II”的国产千卡集群上训练多种模型（如Llama-2，GPT 3等）时，实现了20%以上的性能提升。此外，AdaPipe在“鹏城·脑海”通用大模型（200B）4K窗口3456卡的实际训练中，取得了10%以上的效率提升。这些案例将为未来万卡国产集群上的优化训练提供技术储备和经验参考。

上述技术的研发得到国家自然科学基金和鹏城实验室的支持和资助。

相关论文信息：https://doi.org/10.1145/3620666.3651359