突破传统FSDP限制,实现超大规模训练的灵活高效通信。
veScale-FSDP: Flexible and High-Performance FSDP at Scale
- 采用可变分块的RaggedShard格式与结构感知规划算法
- 支持块量化和非逐元素优化器,吞吐提升5%~66%
- 适用于超万卡规模训练,内存降低16%~30%
全分片数据并行(FSDP),又称零冗余优化器(ZeRO),因其内存效率高且对模型代码侵入小,被广泛用于大规模模型训练。然而现有FSDP系统依赖固定的逐元素或逐行分片格式,与块结构计算冲突,难以支持现代结构感知训练方法,如块级量化和非逐元素优化器(如Shampoo、Muon)。此外,当前实现存在通信与内存开销,在数万张GPU规模下效率下降。本文提出veScale-FSDP,结合灵活的RaggedShard分片格式与结构感知规划算法,实现零拷贝通信,并原生支持块级量化与非逐元素优化器。在超大规模场景下,相较现有FSDP系统,吞吐提升5%至66%,内存占用降低16%至30%,且可高效扩展至数万张GPU。
原文摘要 · Abstract (English)
Fully Sharded Data Parallel (FSDP), also known as Zero Redundancy Optimizer (ZeRO), is widely used for large-scale model training, because of its memory efficiency and minimal intrusion on model code. However, existing FSDP systems rely on fixed element-wise or row-wise sharding formats that conflict with block-structured computations. As a result, they struggle to support modern structure-aware training methods, including block-wise quantization and non-element-wise optimizers such as Shampoo and Muon. In addition, today's implementations incur communication and memory overheads that degrade efficiency at the scale of tens of thousands of GPUs. We introduce veScale-FSDP, a novel FSDP system that combines RaggedShard, a flexible sharding format, with a structure-aware planning algorithm to deliver both flexibility and performance. veScale-FSDP enables zero-copy FSDP communications and natively supports block-wise quantization and non-element-wise optimizers, achieving 5% to 66% higher throughput and 16% to 30% lower memory usage than existing FSDP systems, while scaling efficiently to tens of thousands of GPUs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。