[Paper Reading] MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs

TL;DR
大规模高效训练
容错设计(Fault Tolerance)
Experiment
相关链接

MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs

link
时间：2024.02
单位：ByteDance, Peking University
相关领域：分布式机器学习/大语言模型训练
作者相关工作：字节跳动大模型训练基础设施团队
被引次数： 168
主页：https://github.com/volcengine/veScale

TL;DR

本文提出了首个支持10,000+GPU规模的大语言模型训练系统MegaScale，通过算法-系统协同设计和深度可观测性两大创新，在12,288块GPU上训练175B模型达到55.2%的MFU（模型FLOPs利用率），相比Megatron-LM提升1.34倍效率。系统已部署在字节跳动生产环境，支持数百B参数模型在数万亿token数据上的稳定训练，自动处理超过100次故障恢复。该工作为超大规模LLM训练提供了系统级解决方案，相关技术部分已开源。

大规模高效训练

算法优化

三项关键算法改进：并行Transformer块、滑动窗口注意力(SWA)、LAMB优化器
并行Transformer块
传统串行公式：
\(y=x+MLP(LN(x+Attention(LN(x))))\)
重构后并行版本：
\(y=x+MLP(LN(x))+Attention(LN(x))\)
注意：两者只是近似等价，重构的好处是MLP与Attention可以并行计算

LAMB优化器
支持4倍批量大小而不损失精度
减少87.5%的流水线气泡

通信重叠优化

针对不同并行策略的通信优化：

数据并行
预取首个all-gather操作
按计算依赖顺序确定通信优先级

流水线并行
解耦send/receive操作
异步启动通信与计算重叠

张量/序列并行
融合all-gather/reduce-scatter属于通信密集操作，FFN中Linear属于计算密集操作，并两类操作时间上并行

其他优化

高效算子

采用FlashAttention-2
融合LayerNorm和GeLU内核

数据流水线

异步数据预处理
树状数据加载消除冗余读取

通信组初始化

用Redis替代TCPStore
优化屏蔽操作使10,000+GPU初始化时间<30秒

网络调优

定制CLOS网络拓扑
减少ECMP哈希冲突
混合Swift/DCQCN拥塞控制算法
自适应重传超时设置

容错设计(Fault Tolerance)

关键组件

Driver进程：协调资源分配和故障恢
Executor：每个节点运行训练进程和守护进程
Heartbeat机制：实时异常检测

什么是Kubernetes？

Kubernetes（源自希腊语"舵手"）是Google开源的容器编排系统，用于自动化部署、扩展和管理容器化应用。在MegaScale等AI训练系统中，它负责： - 资源调度：将训练任务分配到GPU节点 - 故障恢复：自动重启崩溃的容器 - 扩缩容：根据负载动态调整资源

数据收集与分析

训练进程状态
stdout/stderr日志
RDMA流量指标（关键隐性故障指标）

诊断测试

主机内网络测试

环回测试测量RNIC到各端点的带宽
RNIC间连通性测试

什么是RNIC？

RNIC（RDMA Network Interface Card，远程直接内存访问网卡）是一种高性能网络硬件，专为低延迟、高吞吐量的数据传输设计。它通过RDMA（Remote Direct Memory Access）技术，允许计算机在不经过CPU的情况下直接访问其他计算机的内存，显著提升了分布式系统的效率。

NCCL测试

节点内all-to-all测试
同ToR交换机节点间的all-reduce测试

Fast Checkpointing and Recovery

保存checkpoint

主机内存：秒级完成，不阻塞训练
异步持久化：后台写入HDFS

恢复checkpoint

数据并行组内单节点读取+广播
避免HDFS带宽竞争

Experiment

12,288 GPU上55.2% MFU
相比Megatron-LM提升1.34x

Ablation Study

fariver