arXiv:2501.04266v2 宣布类型: 替换-交叉 摘要：扩大大型语言模型（LLM）训练涉及在有限数量的 workers 中分布大量的训练参数。然而，像 ZeRO-3 这样的方法虽然大幅减轻了 GPU 内存压力，但经常需要大量的通信以确保全局同步和一致性。像 ZeRO++ 这样的既存努力通过使用次级分区来避免节点间通信，因为在节点内的 GPU-GPU 转移通常具有更多的带宽和更低的延迟。然而，随着如 Frontier 这样具备强大计算能力的基础设施的出现，配備了 AMD GPU 的 Frontier 具有显著的计算能力，因此有必要研究硬件拓扑并开发针对性策略以提高训练效率。在这项工作中，我们为 ZeRO++ 提出了几种通信和优化策略，以减少通信成本并提高内存利用率。在这篇论文中，我们针对目前排名第二的超级计算机集群 Frontier，提出了一种三层级的分层分区方法，旨在利用不同层级通信层（GCD-GCD、GPU-GPU 和节点间）的各种带宽来减少通信开销。对于我们测试的 20B GPT 模型，在使用多达 384 个 GCD 的情况下，我们观察到每 GPU 的 TFLOPS 增加了 1.71 倍，并且在使用多达 384 个 GCD 的情况下，扩展效率达到了 0.94。