veScale-FSDP: FSDP Flessibile e ad Alte Prestazioni su Larga Scala

veScale-FSDP: Flexible and High-Performance FSDP at Scale

February 25, 2026

Autori: 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

Abstract

Fully Sharded Data Parallel (FSDP), noto anche come ZeRO, è ampiamente utilizzato per l'addestramento di modelli su larga scala, caratterizzato dalla sua flessibilità e dalla minima intrusività nel codice del modello. Tuttavia, gli attuali sistemi FSDP faticano a gestire metodi di addestramento *structure-aware* (ad esempio, l'addestramento quantizzato a blocchi) e ottimizzatori non *element-wise* (come Shampoo e Muon) utilizzati in modelli all'avanguardia (ad esempio, Gemini, Kimi K2). I formati di partizionamento (*sharding*) fissi, basati su elementi o righe, dell'FSDP entrano in conflitto con i calcoli a struttura a blocchi. Inoltre, le implementazioni odierne sono carenti in termini di efficienza delle comunicazioni e della memoria, limitando il *scaling* a decine di migliaia di GPU. Introduciamo veScale-FSDP, un sistema FSDP riprogettato che combina un formato di partizionamento flessibile, RaggedShard, con un algoritmo di pianificazione *structure-aware* per offrire sia flessibilità che prestazioni su larga scala. veScale-FSDP supporta nativamente l'efficiente posizionamento dei dati richiesto dall'FSDP, abilitando la quantizzazione a blocchi e gli ottimizzatori non *element-wise*. Di conseguenza, veScale-FSDP raggiunge un *throughput* superiore dal 5% al 66% e un utilizzo della memoria inferiore dal 16% al 30% rispetto ai sistemi FSDP esistenti, scalando efficientemente fino a decine di migliaia di 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.