Google5万卡集群 & Meta2.4万卡集群管理

 各家万卡集群

最近Meta分享其两个2.4万卡集群管理经验。https://engineering.fb.com/2024/03/12/data-center-engineering/building-metas-genai-infrastructure/

加上之前Google介绍了他们的5万卡训练集群管理经验，

https://cloud.google.com/blog/products/compute/the-worlds-largest-distributed-llm-training-job-on-tpu-v5e

以及OpenAI介绍他们6万卡（7500节点）的集群优化经验：

https://openai.com/index/scaling-kubernetes-to-7500-nodes/

对比上次字节的1.2万卡训练集群管理经验：https://arxiv.org/html/2402.15627v1

见：从《字节万卡训练论文》看ModelArts的异曲同工

一起简单汇总：

一起看下各家共同提及，并值得关注的一些技术点：

1. 集群规模

基本上当前的AI集群都以Kubernetes为主，除了Meta，他依赖HPC厂商: penguin solutions。参见此文

大家都提到了随着芯片数据的增加，单个集群管理这么多节点的难度开始触及上限。如何做到单个K8s集群管理这么多芯片成为挑战，因为单个K8s集群的节点数量，一般也就2000节点（1万6卡），再往上就需要优化了，这是各家都必须面对的首要问题。

OpenAI发文说他们搞到单K8s集群7500节点，里面列举了很多需要克服的瓶颈点，比较偏技术细节一些。

由于单个分布式训练任务，目前都是在单个集群内运行管理的。所以要想训练任务规模的扩大，必须首先扩大集群规模。所以昇腾12.8万卡的单集群，对我们来说，几乎不可能。要想单训练任务128K卡，必须得下点功夫才行。

2. 分布式数据存储

大量的训练数据的访问，需要极大的带宽，特别是多个epoch还会访问数据多轮。正如《跑AI大模型的K8s与普通K8s有什么不同》文中所说的，分布式数据存储，是整个AI-Infrastructure中非常重要的一个环节。因此，构建一个性能足够强悍的分布式数据存储系统，必须十分重视。

Google在文章中还提到了一些数据加载的“技巧”，比如大家同时访问数据时，带宽冲击大，改进为先让部分主机加载的方式来提高效率。

3. Checkpoint的保存和恢复

故障在AI训练中不可避免已经是大家的共识，因此Checkpoint的保存和恢复，就是一个必须优化的选项。因为CKPT的恢复，会严重影响训练恢复速度。

各家都提到了，训练恢复时，不用所有卡同时去读CKPT，而是让部分实例去读，然后同步给其他GPU，这样可以降低存储带宽的冲击。其他例如每个人只读CKPT的一小段，然后一起拼成完整的。保存CKPT时，可以先写入内存，然后异步保存到对象存储中来提高效率。

总之，这是一个绝对值得投资的地方。

4. 训练进程组初始化时间

有3家都提到了：随着训练规模的增大，训练进程组初始化的耗时不断扩大，因为大家要等“兄弟们”都初始化ready，才能开始互相通信。这个地方是一个需要优化的点，并且字节和Meta都在这块下了功夫并获得不错的改进。

5. 参数面网络的优化编排

训练的Rank成员在训练过程中，互相之间需要进行如AllReduce类的集合通信来交换参数权限信息，而这种通信关系（即谁会和谁通信）是可以提前计算出来的，因此对参数面网络的编排优化也是必须考虑的一环。如RANK-ID排序，源端口动态hash，NCCL优化等措施来减少流量冲突。

因此拓扑感知的任务的编排优化，是大规模AI训练需要考虑的一个方向。

一个小结

Meta在文章中强调了构建大规模AI集群的一个原则：性能和易用性缺一不可。

这一点在字节的文章中也有类似的描述，因为AI集群中故障是不可避免的，所以如何快速定界定位故障，并提供快速调试&恢复能力，真的是帮程序员自己拯救自己的好特性。就跟我们写代码一样，“可维护性”比“极致性能”优先级更高。