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

Le papier présente veScale-FSDP, un système FSDP repensé qui intègre un format de partitionnement flexible (RaggedShard) et un algorithme de planification adapté à la structure pour prendre en charge les méthodes d'entraînement avancées comme la quantification par blocs et les optimiseurs non élémentaires, tout en offrant une meilleure efficacité mémoire et un débit supérieur à l'échelle de dizaines de milliers de GPU.

L'essentiel

🚀 Le Problème : Construire un gratte-ciel avec des briques cassées

Imaginez que vous essayez de construire le plus grand gratte-ciel du monde (un modèle d'intelligence artificielle géant) avec des milliers d'ouvriers (des puces graphiques ou GPU).

Pour que la construction soit rapide, vous devez diviser les tâches entre tous les ouvriers. C'est ce qu'on appelle le FSDP (Fully Sharded Data Parallel). C'est comme si chaque ouvrier ne gardait que sa propre partie du plan de l'étage, au lieu de porter tout le plan entier.

Mais il y a deux gros problèmes avec les systèmes actuels :

1. Le problème de la "coupe" (Flexibilité) :
   Imaginez que vous devez couper un gâteau en parts égales pour les ouvriers. Les systèmes actuels coupent le gâteau au couteau, n'importe où, même au milieu d'une fraise ou d'un morceau de chocolat (les "blocs" de données).

   • Conséquence : Si un ouvrier a besoin d'un bloc entier de chocolat pour faire une recette spéciale (comme les nouveaux algorithmes d'optimisation Muon ou la quantification), il doit attendre que les autres lui envoient les morceaux manquants. C'est lent et compliqué. Les ouvriers passent plus de temps à se passer des morceaux de gâteau qu'à construire.

2. Le problème du "gaspillage" (Performance) :
   Pour que les ouvriers travaillent ensemble, ils doivent s'aligner parfaitement. Les systèmes actuels forcent parfois les ouvriers à remplir des espaces vides avec du carton (du "padding") pour que tout soit carré.

   • Conséquence : Cela prend de la place dans les camions de livraison (la mémoire du GPU) et ralentit le transport. De plus, les ouvriers doivent souvent copier les plans d'un camion à l'autre, ce qui crée des embouteillages.

---

💡 La Solution : veScale-FSDP (Le Chef d'Orchestre Intelligent)

L'équipe de ByteDance a créé veScale-FSDP, un nouveau système qui résout ces problèmes avec deux innovations majeures.

1. La "Coupe sur Mesure" (RaggedShard)

Au lieu de couper le gâteau au couteau n'importe où, imaginez que vous avez un couteau magique (RaggedShard) qui peut couper exactement là où le gâteau est fait.

• L'analogie : Si un bloc de données est un carré de 10x10, le système ne le coupe pas en deux. Il donne le carré entier à un ouvrier, ou deux carrés entiers à un autre.
• Le résultat : Les ouvriers peuvent maintenant utiliser des recettes spéciales (comme les optimiseurs Muon ou la quantification 8-bit) sans avoir à attendre des pièces détachées. Tout est prêt à l'emploi. C'est comme si chaque ouvrier avait exactement les ingrédients dont il a besoin, sans avoir à courir les chercher.

2. Le "Plan de Livraison Optimisé" (Algorithme de Planification)

Une fois que les morceaux sont coupés proprement, il faut les ranger dans les camions pour les livrer.

• L'analogie : Imaginez que vous devez charger des boîtes de tailles différentes dans un camion. Si vous les empilez au hasard, il y aura plein d'espaces vides et ça va prendre du temps.
• Le résultat : veScale utilise un algorithme intelligent (un peu comme un super-organisateur) qui réarrange les boîtes pour qu'elles s'emboîtent parfaitement, sans espaces vides inutiles. Il calcule le meilleur agencement en une fraction de seconde.
• Le "DBuffer" (Le Camion Magique) : C'est une nouvelle façon de stocker les données. Au lieu de décharger et recharger les camions à chaque étape (ce qui prend du temps), le système garde les données dans un espace de stockage unique et partagé. Les ouvriers y accèdent directement, sans avoir à déplacer les cartons. C'est du "zéro copie".

---

🏆 Les Résultats : Pourquoi c'est génial ?

Grâce à cette nouvelle organisation :

• Plus rapide : Les travaux avancent 5 % à 66 % plus vite. C'est comme si vos ouvriers construisaient un étage entier en moins de temps.
• Moins de gaspillage : Ils utilisent 16 % à 30 % de mémoire en moins. C'est comme si vous pouviez construire le même gratte-ciel avec la moitié des camions de livraison.
• Évolutivité : Le système fonctionne parfaitement même avec 10 000 ouvriers (GPU) qui travaillent ensemble. Les systèmes précédents commençaient à s'effondrer ou à ramer avec autant de monde.

🎯 En résumé

veScale-FSDP, c'est passer d'une équipe de construction où tout le monde se bouscule, perd du temps à couper des briques et remplit des camions de carton, à une équipe de chefs d'orchestre qui :

1. Découpent les tâches exactement comme il faut (RaggedShard).
2. Organisent le transport pour qu'il n'y ait aucun espace vide (Planification).
3. Utilisent un système de stockage ultra-rapide (DBuffer).

Le résultat ? On peut entraîner des intelligences artificielles plus grandes, plus vite, et avec moins de matériel. C'est une révolution pour l'avenir de l'IA.

Résumé technique

Titre : veScale-FSDP : FSDP Flexible et Haute Performance à Grande Échelle

1. Problématique

Le Fully Sharded Data Parallel (FSDP), également connu sous le nom de ZeRO, est la méthode de référence pour l'entraînement de modèles de grande taille (LLM) sur des milliers de GPU. Cependant, les systèmes FSDP actuels (DeepSpeed ZeRO, PyTorch FSDP1/FSDP2, Megatron-FSDP) rencontrent deux obstacles majeurs face aux techniques d'entraînement modernes :

• Inflexibilité face aux méthodes "Structure-Aware" : Les modèles de pointe (ex: Gemini, DeepSeek-V3) utilisent des optimiseurs non élémentaires (comme Shampoo et Muon) et des techniques de quantification par blocs (Block-wise Quantization). Ces méthodes nécessitent que les tenseurs soient traités par blocs atomiques (matrices 2D complètes ou blocs de poids).
  • Les formats de fessage (sharding) actuels sont soit élément par élément, soit par ligne (row-wise). Cela crée des frontières de fessage qui ne s'alignent pas avec les blocs requis, obligeant les développeurs à modifier intrusivement le code des modèles/optimiseurs ou à gérer des vérifications de limites complexes, du padding et une communication supplémentaire.
• Efficacité limitée à grande échelle : Les implémentations actuelles souffrent d'inefficacités en matière de communication et de gestion de la mémoire.
  • Mémoire : La fragmentation mémoire et les allocations non déterministes limitent la capacité à scaler au-delà de quelques milliers de GPU, entraînant des échecs (OOM) ou un sur-provisionnement coûteux.
  • Communication : Les opérations collectives (AllGather, ReduceScatter) sont souvent fragmentées ou nécessitent des copies de données intercalées (interleaved copy), réduisant le débit (throughput).

2. Méthodologie et Architecture

L'équipe de ByteDance propose veScale-FSDP, un système redessiné qui combine flexibilité et performance via trois innovations principales :

A. RaggedShard : Un format de fessage flexible

Pour résoudre le problème d'alignement, les auteurs introduisent RaggedShard, un nouveau format de placement pour les tenseurs distribués (DTensor) de PyTorch.

• Granularité arbitraire : Contrairement aux formats rigides, RaggedShard permet de définir des blocs atomiques de taille personnalisée (ex: blocs 2D pour la quantification, lignes entières pour les optimiseurs matriciels).
• Distribution irrégulière : Il permet d'assigner un nombre différent de blocs à chaque appareil, respectant ainsi les frontières structurelles des tenseurs.
• Composabilité : Il s'intègre nativement avec les stratégies de parallélisme existantes (TP, EP) et les API PyTorch (fully_shard), évitant toute modification intrusive du code du modèle.

B. Algorithme de Planification (Planning Algorithm)

Pour optimiser la communication avec des tenseurs RaggedShard, un algorithme de planification est nécessaire pour regrouper les tenseurs dans des buffers de communication.

• Problème d'optimisation : Le regroupement optimal est formulé comme un problème NP-difficile (réductible au problème de partition). L'objectif est de minimiser la taille du buffer tout en respectant trois contraintes : pas de blocs fragmentés, mémoire contiguë pour les tenseurs, et équilibrage de charge strict.
• Solution Heuristique : Les auteurs proposent un algorithme dynamique (Dynamic Programming) guidé par des heuristiques de permutation. Il exploite la régularité des modèles Transformer (ex: poids linéaires dominants) pour trouver des solutions quasi-optimales en temps polynomial, évitant les temps de calcul excessifs des solveurs ILP.

C. Distributed Buffer (DBuffer)

Pour soutenir RaggedShard, un nouveau primitif matériel/logiciel, le DBuffer, est introduit.

• Accès Zero-Copy : Il mappe directement les données des tenseurs dans un buffer global, éliminant les copies intercalées (copy-in/copy-out) coûteuses.
• Opérateurs de Groupe : Au lieu d'exécuter des noyaux CUDA par tenseur, le DBuffer fusionne les opérations identiques (add, scale, copy) à l'échelle du groupe avant la communication, réduisant la surcharge de calcul et les blocages.
• Gestion Mémoire : Il utilise des allocations par lots (batched) et une gestion explicite des dépendances de flux pour réduire la fragmentation et permettre une libération prévisible de la mémoire.

3. Contributions Clés

1. Support Natif des Optimiseurs et Quantifications Avancés : veScale-FSDP permet l'utilisation directe d'optimiseurs comme Muon et de techniques comme la quantification 8-bit par blocs sans modifier le code du modèle, grâce à RaggedShard.
2. Performance Supérieure : Le système atteint un débit 5 % à 66 % supérieur aux systèmes FSDP existants (DeepSpeed, FSDP1/2, Megatron-FSDP).
3. Réduction de la Mémoire : Il réduit l'utilisation de la mémoire GPU de 16 % à 30 %, permettant d'entraîner des modèles plus grands ou d'utiliser moins de GPU pour une même charge.
4. Scalabilité Massive : Le système a été validé pour scaler efficacement jusqu'à 10 000+ GPU et des modèles de 2,4 billions de paramètres.
5. Intégration PyTorch Native : Il fonctionne comme un module plug-and-play transparent, remplaçant le backend de FSDP2 tout en conservant la même API.

4. Résultats Expérimentaux

Les évaluations ont été menées sur des clusters GPU avec des modèles variés (LLaMA-3-70B, GPT-OSS-120B, et un modèle MoE interne de 160B/800B/2.4T).

• Débit (Throughput) :
  • Sur les modèles MoE, veScale-FSDP est 11 % à 66 % plus rapide que les concurrents.
  • Sur LLaMA-3-70B, il est 5 % plus rapide que DeepSpeed et FSDP1/2.
  • L'amélioration provient de l'évitement du padding excessif, de la fusion des noyaux CUDA via DBuffer et de l'absence de copies de données inutiles.
• Mémoire :
  • Réduction de 16 % à 30 % de la mémoire pic réservée.
  • Évite les erreurs OOM (Out-Of-Memory) qui affectent FSDP2 sur des configurations spécifiques (ex: GPT-OSS sur 256 GPU).
• Scalabilité :
  • Weak Scaling : Scalabilité quasi-linéaire jusqu'à 8 000 GPU.
  • Strong Scaling : Scalabilité linéaire jusqu'à 10 000 GPU avec des tailles de batch globales fixes.
  • Model Scaling : Capacité à entraîner des modèles de 2,4T de paramètres sur 1 000 GPU sans dégradation de performance (MFU légèrement améliorée).
• Cas d'usage spécifiques :
  • 8-bit Adam : Implémenté avec quelques lignes de code, montrant une convergence similaire à DDP mais avec une économie massive de mémoire.
  • Muon Optimizer : Permet d'exécuter des mises à jour Newton-Schulz sur des matrices complètes sur un seul nœud racine (via redistribution asynchrone), atteignant 47,3 % d'utilisation des FLOPS (MFU) sur 256 GPU.

5. Signification et Impact

veScale-FSDP représente une avancée majeure pour l'entraînement de modèles de fondation à l'échelle industrielle.

• Découplage Architecture/Système : Il permet aux chercheurs de concevoir des architectures complexes (MoE, optimiseurs exotiques) sans se soucier des contraintes de bas niveau du système de parallélisme.
• Efficacité Économique : En réduisant la consommation mémoire et en augmentant le débit, il diminue le coût de l'entraînement et permet d'utiliser des clusters existants plus efficacement.
• Adoption Future : Le concept de RaggedShard est déjà intégré dans la feuille de route officielle de PyTorch, indiquant une adoption potentielle large dans l'écosystème open-source.
• Production Ready : Le système a été testé en production chez ByteDance sur des milliers de GPU et est portable sans dépendre d'infrastructures internes spécifiques.

En résumé, veScale-FSDP résout le conflit historique entre la flexibilité requise par les nouvelles architectures de modèles et l'efficacité requise par l'entraînement à grande échelle, offrant une solution unifiée, performante et évolutive.