arXiv:2501.04266cs.DCcs.AI2025-01被引 2

针对前沿超算的低带宽分区策略,提升大模型训练效率。

Scaling Large Language Model Training on Frontier with Low-Bandwidth Partitioning

  • 设计三级分层分区,利用不同通信层级的带宽差异降低开销。
  • 200亿参数模型在384个计算单元上实现1.71倍每卡浮点性能提升。
  • 适合超大规模模型在先进异构集群上的高效训练部署。

扩大大语言模型(LLM)训练需将海量参数分布在有限工作节点中。现有方法如ZeRO-3虽显著降低显存压力,但常伴随高额通信成本以保证全局同步与一致性。传统方案如ZeRO++通过二级分区避免节点间通信,因节点内GPU-GPU传输通常具有更高带宽和更低延迟。随着具备强大算力的前沿超算(Frontier)等新型基础设施出现,亟需针对其硬件拓扑特性优化训练策略。本文提出一套专为当前排名第二的超级计算机Frontier设计的通信与优化策略,引入针对该架构的三级分层分区机制,旨在利用各层级通信(GCD-GCD、GPU-GPU、跨节点)间的带宽差异,有效减少通信开销。实验表明,在200亿参数的GPT模型上,相较于原始ZeRO++,使用384个计算单元(GCDs)时,每卡浮点性能(TFLOPS)提升1.71倍,且在384个计算单元下仍保持0.94的高扩展效率。

原文摘要 · Abstract (English)

Scaling up Large Language Model(LLM) training involves fitting a tremendous amount of training parameters across a limited number of workers. However, methods like ZeRO-3 that drastically reduce GPU memory pressure often incur heavy communication to ensure global synchronization and consistency. Established efforts such as ZeRO++ use secondary partitions to avoid inter-node communications, given that intra-node GPU-GPU transfer generally has more bandwidth and lower latency than inter-node connections. However, as more capable infrastructure like Frontier, equipped with AMD GPUs, emerged with impressive computing capability, there is a need for investigations on the hardware topology and to develop targeted strategies to improve training efficiency. In this work, we propose a collection of communication and optimization strategies for ZeRO++ to reduce communication costs and improve memory utilization. In this paper, we propose a 3-level hierarchical partitioning specifically for the current 2nd ranked supercomputing cluster, Frontier, which aims at leveraging various bandwidths across layers of communications (GCD-GCD, GPU-GPU, and inter-node) to reduce communication overhead. For a 20B GPT model, we observe a 1.71x increase in TFLOPS per GPU when compared with ZeRO++ up to 384 GCDs and a scaling efficiency of 0.94 for up to 384 GCDs.

大模型训练超算优化通信效率分布式训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。