veScale-FSDP: FSDP Flexible y de Alto Rendimiento a Gran 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
Resumen
El Paralelismo de Datos Totalmente Fragmentado (FSDP), también conocido como ZeRO, es ampliamente utilizado para entrenar modelos a gran escala, destacándose por su flexibilidad y mínima intrusión en el código del modelo. Sin embargo, los sistemas FSDP actuales presentan dificultades con métodos de entrenamiento conscientes de la estructura (por ejemplo, entrenamiento cuantizado por bloques) y con optimizadores no elementales (como Shampoo y Muon) utilizados en modelos de vanguardia (por ejemplo, Gemini, Kimi K2). Los formatos de fragmentación fijos, ya sea por elemento o por fila, de FSDP entran en conflicto con los cálculos de estructura en bloques. Además, las implementaciones actuales adolecen de eficiencia en comunicaciones y memoria, lo que limita la escalabilidad a decenas de miles de GPUs. Presentamos veScale-FSDP, un sistema FSDP rediseñado que combina un formato de fragmentación flexible, RaggedShard, con un algoritmo de planificación consciente de la estructura para ofrecer tanto flexibilidad como rendimiento a escala. veScale-FSDP admite de forma nativa la colocación eficiente de datos requerida por FSDP, potenciando la cuantización por bloques y los optimizadores no elementales. Como resultado, veScale-FSDP logra un rendimiento entre un 5% y un 66% superior y un uso de memoria entre un 16% y un 30% menor que los sistemas FSDP existentes, mientras escala eficientemente a decenas de miles 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.