veScale-FSDP : FSDP Flexible et Hautes Performances à Grande Échelle
veScale-FSDP: Flexible and High-Performance FSDP at Scale
February 25, 2026
Auteurs: 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
Résumé
Le parallélisme de données entièrement partitionnées (FSDP), également connu sous le nom de ZeRO, est largement utilisé pour l'entraînement de modèles à grande échelle, se distinguant par sa flexibilité et son intrusion minimale dans le code du modèle. Cependant, les systèmes FSDP actuels peinent à gérer les méthodes d'entraînement sensibles à la structure (par exemple, l'entraînement quantifié par blocs) et les optimiseurs non élément par élément (comme Shampoo et Muon) utilisés dans les modèles de pointe (tels que Gemini et Kimi K2). Les formats de partitionnement fixes, élément par élément ou ligne par ligne, de FSDP entrent en conflit avec les calculs à structure bloc. De plus, les implémentations actuelles présentent des lacunes en matière d'efficacité des communications et de la mémoire, limitant la montée en charge à des dizaines de milliers de GPU. Nous présentons veScale-FSDP, un système FSDP repensé qui associe un format de partitionnement flexible, RaggedShard, à un algorithme de planification sensible à la structure pour offrir à la fois flexibilité et performance à grande échelle. veScale-FSDP prend nativement en charge le placement efficace des données requis par FSDP, permettant la quantification par blocs et les optimiseurs non élément par élément. En conséquence, veScale-FSDP atteint un débit de 5 à 66 % supérieur et une utilisation mémoire de 16 à 30 % inférieure aux systèmes FSDP existants, tout en montant efficacement en charge sur des dizaines de milliers de GPU.
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.