veScale-FSDP: FSDP Flexível e de Alto Desempenho em Escala
veScale-FSDP: Flexible and High-Performance FSDP at Scale
February 25, 2026
Autores: Zezhou Wang, Youjie Li, Zhiqi Lin, Jiacheng Yang, Cong Xie, Guanyu Feng, Zheng Zhong, Ziyue Huang, Hongyu Zhu, Zhi Zhang, Yanghua Peng, Xin Liu
cs.AI
Resumo
O Fully Sharded Data Parallel (FSDP), também conhecido como ZeRO, é amplamente utilizado para treinar modelos de grande escala, destacando-se pela sua flexibilidade e intrusão mínima no código do modelo. No entanto, os sistemas FSDP atuais enfrentam dificuldades com métodos de treinamento com consciência estrutural (por exemplo, treinamento quantizado em blocos) e com otimizadores não elementares (por exemplo, Shampoo e Muon) usados em modelos de ponta (por exemplo, Gemini, Kimi K2). Os formatos de fragmentação fixos, por elemento ou por linha, do FSDP entram em conflito com os cálculos de estrutura em blocos. Além disso, as implementações atuais são insuficientes em termos de eficiência de comunicação e memória, limitando a escalabilidade para dezenas de milhares de GPUs. Apresentamos o veScale-FSDP, um sistema FSDP redesenhado que combina um formato de fragmentação flexível, o RaggedShard, com um algoritmo de planeamento com consciência estrutural para oferecer flexibilidade e desempenho em escala. O veScale-FSDP suporta nativamente o posicionamento eficiente de dados exigido pelo FSDP, permitindo a quantização em blocos e otimizadores não elementares. Como resultado, o veScale-FSDP alcança um throughput 5~66% superior e um uso de memória 16~30% inferior aos sistemas FSDP existentes, escalando eficientemente para dezenas de milhares de GPUs.
English
Fully Sharded Data Parallel (FSDP), also known as ZeRO, is widely used for training large-scale models, featuring its flexibility and minimal intrusion on model code. However, current FSDP systems struggle with structure-aware training methods (e.g., block-wise quantized training) and with non-element-wise optimizers (e.g., Shampoo and Muon) used in cutting-edge models (e.g., Gemini, Kimi K2). FSDP's fixed element- or row-wise sharding formats conflict with the block-structured computations. In addition, today's implementations fall short in communication and memory efficiency, limiting scaling to tens of thousands of GPUs. We introduce veScale-FSDP, a redesigned FSDP system that couples a flexible sharding format, RaggedShard, with a structure-aware planning algorithm to deliver both flexibility and performance at scale. veScale-FSDP natively supports efficient data placement required by FSDP, empowering block-wise quantization and non-element-wise optimizers. As a result, veScale-FSDP achieves 5~66% higher throughput and 16~30% lower memory usage than existing FSDP systems, while scaling efficiently to tens of thousands of GPUs.