下文来自于零一万物。
前言
随着 ChatGPT 的发布和不断升级，生成式 AI 已成为时代的热点。Gemini 和 Sora 的发布更进一步的引领从文本生成到多模态的理解和内容生成。伴随算力和数据规模指数级扩大，神经网络出现了“涌现”智能，此刻通往 AGI 的道路正在被 Scaling Law 点亮。在未来一到两年内，我们将看到模型参数量逐渐从千亿迈向万亿，模型能力更加泛化。大模型对底层算力的诉求进一步升级，万卡集群建设和运营将成为大模型训练的基础门槛。
零一万物的愿景是让通用的人工智能普惠各地，人人受益。针对大规模 AI 集群的建设和运维面临的多方面挑战，我们以多位行业高级技术专家为班底，打造了一支顶尖的 AI Infra 团队。这个团队具备万卡集群的设计、建设和运营经验，且成功训练过多个百亿和千亿参数模型。团队成立后，仅耗时 2 个月就完成了数千张最新 GPU 集群和配套高性能分布式存储服务的设计、选型、施工、交付和验收工作，成功应用于千亿参数的 Yi 系列模型生产训练中。为了实现我们的愿景和使命，我们将分享 AI Infra 团队在面向万卡集群基建、计算、网络、存储、调度以及运维管理方面的设计思路和经验。
基建
电力供应
传统的数据中心面向 X86 服务器设计，单机柜电力负载通常低于 20KW。然而，目前主流的 GPU 8 卡服务器的功率超过 10KW，这意味着单机柜（Rack）仅能容纳 1-2 台 GPU 服务器。按照英伟达官方推荐的单机柜 4 服务器布局，单机柜需要提供超过 40KW 的电力供应。英伟达在 GTC2024 推出的 GB200-NVL72 单柜功率需求超过 100 KW，这使得电力供应将成为一个更大的挑战。考虑到 GPU 型号、功率、电力供应和机房占地面积等因素，要合理设计 Rack 内设备布局和机房内所有 Rack 的空间布局，以满足电力约束条件。我们将该问题转化为多目标约束问题，从而得到满足电力约束条件的 Rack 设计和机房合理布局。
液冷技术
在未来几年，传统的风冷系统将逐渐被液冷系统所取代。使用液冷系统主要解决以下几个挑战：
冷却效率：随着单机柜计算密度和功率的提升，风冷系统无法满足热量移除的要求。液冷系统能够提供更高效的冷却，确保设备在高负载情况下的稳定运行。
功率需求：相比风冷系统，液冷系统需要的功率远低。这意味着在数据中心的整体运行中，液冷系统可以节约大量的能源成本和大幅度降低 PUE。
环境影响：液冷系统能够大幅度降低对水资源的浪费，并且减少数据中心的噪音污染。这有助于提高数据中心的环境友好性，并符合可持续发展的趋势。
空间布局：引入液冷系统可以节省大量的数据中心空间，因为它不需要像风冷系统那样大量的空气流动空间。这样可以优化数据中心的布局，提高空间利用效率。
长期稳定性：高温和灰尘是芯片等电子设备的克星，液冷的引入降低了空气流动和灰尘附着，并且能更好的控制温度，有利于芯片和服务器的长期稳定性。
然而，液冷系统也带来了布局设计、设备安装和日常运维等方面的挑战，需要更高水平的技术和管理来应对。
计算
对于千亿到万亿参数的模型，训练通常需要数月完成。我们针对模型训练的加速主要采用以下四个方面的思路
采用最新的硬件支持特性，譬如使用 FP8 进行端到端训练，相对当前流行的 BF16 大幅提高训练吞吐量
通过算子融合、自动并行，提升单卡的有效计算利用率 MFU
算法和工程联合优化，在设计模型结构和并行策略考虑计算拓扑，提高训练效率
训练性能线性 Scale 增长，即随着训练规模的扩大，保证训练性能线性扩展
端到端的 FP8 训练实现
训练容错方案：由于没有 BF16 的 baseline 来检查千亿模型 FP8 训练的 loss下降是否正常，于是，每间隔一定的步数，同时使用 FP8 和 BF16 进行训练，并根据 BF16 和 FP8 训练的 loss diff 和评测指标的差异，决定是否用 BF16 训练修正 FP8 训练。在这个过程中，我们基于 NVIDIA 软硬结合的技术栈，在功能开发、调试和性能层面，完成了在大模型的 FP8 训练和验证，训练吞吐相对 BF16 得到 1.3 倍的性能提升。
图 2：FP8 训练范式
算子优化融合和热点定位
训练框架中丰富的小算子为算法开发带来了便利性，但是却引入了性能的瓶颈。在模型训练中，如果算子过于琐碎，计算过程中会引入过多的 cuda kernel 启动的开销，而且反复的显存读写效率也非常低，除了对性能的影响，也会引入过多的激活显存空间的占用。
算子融合优化是 GPU 计算中非常高效的优化手段，将小算子融合能解决 kernel 启动的开销、提升算子计算效率和优化显存占用。除了使用目前比较通用的开源 flash-atten 实现外，在训练过程中，还针对训练中其余算子使用编译优化方式和手写算子方式实现算子融合，同时也开发了针对 optimizer 的融合实现。
为了分析训练过程中的瓶颈，我们使用自定义系列度的 nvtx 打点方式，将代码和 kernel timeline 对应起来，找到有优化空间的代码模块进行优化。同时也结合 pytorch 的 execution trace 和 kineto trace 进行联合仿真分析，找到训练过程中的通信和计算瓶颈，进行针对优化。
4D/5D 分布式并行
在大模型训练中，模型、数据、序列、上下文等维度并行是非常重要的分布式训练并行方式，帮助我们更加灵活的分配我们的 GPU 达到一个最佳的训练吞吐。多维度的并行都会引入通信开销，针对不同维度的并行进行优化，将通信和计算并行起来，提升训练的效率。
图 3：from huggingface.co
MOE 训练优化
MoE 模型在 MLP 模块中引入了多个 expert 单元，并在训练中激活其中的部分 expert。不同的 tokens 会根据 router 的计算分发到了不同的 expert 单元中，这些 expert 又通过 expert parallel 分布到了不同的 rank 上，引入了额外的 all2all 通信开销。除了将多个 expert 的计算通过 pipeline 的方式进行多 cuda stream 并行，同时将通信和 local/shared expert 计算并行起来进行隐藏，提升训练的效率。MoE 训练中的 permution/unpermution 也进行了融合算子优化，优化显存占用，提高了计算效率。
网络
网络协议选型
RoCE 和 IB 网络在高性能计算中各有优劣，两者都能满足 AGI 场景的训练和推理需求。对比两种网络类型，IB 网络具备如下三个方面的优点：
高吞吐，低延迟的性能优势：IB 网络在高性能计算中表现出色，具有高速数据传输和低延迟的优势，能够满足对性能要求极高的场景。
具备 SDN 特点的日常运维：IB 网络在日常运维中具备软件定义网络（SDN）的特点，子网管理器和 UFM 能够全局视角纵览网络的抖动、报错和拥塞情况，有利于网络管理和故障排查。
简单的路由算法：IB 网络的路由算法相对简单，日常使用不需要过多考虑流控和拥塞控制，运维管理相对便利。
然而，IB 网络也存在弊端：
生态封闭：IB 网络的生态相对封闭，对第三方生态系统的集成和扩展能力有限。
价格昂贵：IB 网络的硬件设备和部署成本较高，对于一些预算有限的场景可能不太合适。
多租和虚拟化：IB 网络对多租和虚拟化的支持不够友好，限制了其在公有云上的推广和使用。
我们认为 RoCE 的不足都可以通过软件手段来优化解决，而 IB 网络的生态封闭和价格昂贵则相对难以克服。因此，未来 RoCE 方案可能会占据主导地位。博通的 TH4/TH5 方案将会占有大部分市场份额。英伟达自身也意识到了这个问题，并提供了 RoCE 解决方案 Spectrum-X。
网络拓扑设计
针对大模型训练特征，我们采用的网络拓扑设计原则为局部强连通，全局非对称。大模型训练对通信带宽的要求 TP >> PP > DP。在合理划分训练集群后，可以将全联接网络降级成非 1:1 对称收敛比网络，从而在保障训练性能的同时降低网络设备成本。然而，这其中存在着两个主要的技术挑战：
降级和网络性能的权衡：在进行网络降级时，需要权衡网络性能和成本之间的关系，确保降级后的网络能够满足训练性能的需求。同时，需要充分考虑降级对训练场景的实际影响，以确保训练结果的准确性和可靠性。
容错保障：在进行网络降级时，需要确保系统具备良好的容错机制，能够有效应对可能出现的网络故障和错误。这包括对通信链路的监控和检测，以及及时的故障恢复和数据重传机制，以保障训练过程的连续性和稳定性。
解决这些技术挑战需要综合考虑网络设计、算法优化和系统架构等方面的因素，以确保降级后的网络能够在保证性能的同时降低实现成本，并且具备良好的容错保障机制。
图 4：三层无阻塞 IB 网络
存储
存储在 GPU 集群中占据着非常重要的位置，但其中的技术挑战通常容易被忽略。一方面，随着多模态数据的加入，万卡集群的存储容量需要进行合理的规划。另一方面，在数据加载和checkpoint 写入/加载的场景下，存储的分布式读写性能至关重要。如何在容量、性能和成本三者之间取得平衡，以下是我们的实践和思考。
基于 NVME 和 RDMA 的高性能分布式存储
GPU 服务器通常配备高规格的 NVMe 存储盘，但由于功率、负载和温度的关系，其稳定性远远低于专业的 X86 存储服务器。因此，GPU 服务器配备的 NVMe 存储通常只适用于作为本地缓存。通常情况下，我们需要配备高性能的分布式存储解决方案。然而，高性能分布式存储的硬件价格和软件成本通常较高，因此适用于对实时性要求较高的场景，如 CKPT 加载/写入和数据集加载等。针对容量配比，我们的建议是一卡 1TB。换言之，每千卡计算资源需要配置 1PB 容量的高性能分布式存储。
多级缓存技术和读写分离
如前所述，高性能分布式存储由于造价和维护成本高昂，仅适合在大模型训练核心业务场景中使用。除此之外，我们还需要建设一定规模的冷存储用于日常的数据存储和归档。在这种业务场景中，通常对读写带宽要求不高，但对稳定性要求很高。因此，可以选择使用普通的 SATA 磁盘，并基于对象存储的形式来提供服务。
图 5：对象、分布式和本地三层存储架构
开源方案和商业方案对比
为了选取最适合大模型场景的存储方案，我们对主流的开源方案和商业方案做了详尽的对比测试。最终得到的结论是开源方案在万卡场景下的读写吞吐延迟可以达到商业方案的水平。在单点或者少量节点业务场景里，存储系统的瓶颈在 FUSE 实现里的内核态和用户态数据拷贝中，比商业方案性能弱 50% 左右。基于此结论我们设计并实现了开源方案和商业方案混合存储实现。
调度系统
拓扑亲和调度
在实践过程中，我们发现对于小规模的集群比如 256 卡，集合通信通常只在一层网络发生，因此通常具备不错的开箱性能。然而，对于千卡乃至万卡的训练任务，集合通信在二层和三层网络完全随机会导致训练性能下降。因此，调度系统需要感知网络拓扑和训练框架的并行策略，以确保流水线并行控制在一层网络，而数据并行配合梯度累积则分布在二层网络。这样一来，我们可以保持大规模网络和小规模网络的性能一致性。
图 6：零一万物自研万卡集群训推任务调度系统
Goodput 优化
Goodput 是由谷歌提出的衡量集群有效使用的参数，即集群有效训练时间与实际用时的比值。Goodput 是衡量 Infra 技术的最重要指标。基于实践我们得出影响 Goodput 的最主要因素包括：1）任务由于 GPU Xid 错误或 IB 网络抖动错误异常中断，2）任务启动和网络建链耗时，3）Checkpoint 加载和存储耗时。针对这些因素，我们分别进行了如下优化，实现集群 Goodput > 99%。
智能看护：启用智能看护功能后，对于常见的物理硬件异常中断，系统可以自动隔离故障并重启任务，无需人工干预。
Fast Checkpoint：将 checkpoint 的写入方式由默认的同步磁盘写入优化成异步内存写入，减少对训练流程的中断。
任务不中断重启：在实际训练中，任务异常退出通常是由单点故障引起的。因此，我们致力于实现单点故障的替换，而不是重启整个训练任务。
优化 GPU 利用率和分配率
在大模型的场景中，通常将八卡视为一个整体调度单元，因此资源碎片对调度的挑战并不十分显著。更多的挑战在于实现任务优先级的管理，即低优先级任务的驱逐和高优先级任务的抢占。
故障监控与定位
在构建复杂系统时，可观测性是不可或缺的一环。在大模型时代的基础设施中，可扩展性的主要挑战在于如何在数以万计的硬件设备中实现秒级定位问题设备，并触发调度系统的故障隔离。根据国内公有云大厂实测数据，一个千卡集群如果缺少有效的可观测性系统辅助，则集群的实际有效使用比率只能在 70% 左右，30% 的时间将会被浪费在故障问题定位、修复和任务重启等环节。
图 7：零一万物万卡集群可观测性系统架构
海量设备故障快速定位
在万卡集群中，需要搭配小十万个光模块和光纤。在这数以万计的设备中，任何一个发生故障都可能导致训练中断。随着集群规模的扩大，这个问题会变得更加严峻。为解决此挑战，我们建立了综合快速定位机制。例如，利用系统日志定位 Xid 或者 SXid 故障， 基于子网管理器的日志快速检测网卡的 Flapping/LinkDown 等问题。若缺乏合适的定位方案，这类问题的定位通常需要很长时间。我们基于 Prometheus + Grafana + FluentBit 实现了上述异常的实时检测和可视化展示，从而达到对硬件故障导致的训练异常的秒级定位。
故障可预测
可预测性是大模型训练不可或缺的能力。基于我们的实践经验，硬件设备损坏往往并非突然发生，而是伴随着许多先兆事件。通过对硬件设备历史信息进行建模，我们能够提前预测可能损坏的硬件设备并发送告警通知。系统管理员可以根据预测结果，结合业务资源使用安排、备件采购周期和数据中心运维窗口等合理安排硬件替换时间。避免故障硬件造成业务异常中断，进一步提升集群的有效使用率。这部分工作我们正在撰写论文，敬请关注。
