💻 computer science

MatrixFSDP: communication-free matrix optimizers under ZeRO-3 parameter sharding

MatrixFSDP permet un entraînement à grande échelle sans communication avec des optimiseurs de matrice comme Muon sous le partitionnement ZeRO-3 en réorganisant le placement des paramètres de sorte que chaque matrice de poids 2D réside entièrement sur un seul rang, éliminant ainsi le besoin de reconstruction de matrice coûteuse lors des étapes de l'optimiseur tout en maintenant l'efficacité de la mémoire et en atteignant des réductions de latence significatives.

Auteurs originaux : Ming Gao, Yanwu Xu, Hao Zhang

Publié 2026-07-08✓ Author reviewed ⓘ

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Ming Gao, Yanwu Xu, Hao Zhang

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'enseigner à une équipe massive de robots (un cluster d'ordinateurs) comment écrire un roman. Les robots travaillent ensemble pour apprendre à partir d'un immense livre de texte. Pour ce faire, ils utilisent une règle d'apprentissage spéciale appelée Muon.

Le Problème : La « Vue d'ensemble » vs Les « Pièces du Puzzle »

Normalement, quand les robots apprennent, ils découpent le livre massif en minuscules pièces de puzzle. Chaque robot ne détient que quelques pages (un « shard »). C'est excellent pour économiser la mémoire car aucun robot n'a besoin de porter tout le livre. Cette méthode est appelée ZeRO-3.

Cependant, la règle d'apprentissage Muon est un peu exigeante. Elle ne veut pas apprendre à partir de quelques pages à la fois. Pour faire son travail parfaitement, elle a besoin de voir la page entière en 2D (la matrice complète) d'un seul coup afin de comprendre les relations entre les mots.

Le Conflit :

ZeRO-3 dit : « Nous n'avons que des pièces de puzzle. »
Muon dit : « J'ai besoin de la page entière pour apprendre. »

Les Anciennes Solutions (Les Mauvaises Options) :

La Méthode de « Reconstruction » : Chaque fois que les robots ont besoin d'apprendre, ils s'arrêtent, rassemblent toutes les pièces du puzzle de chaque robot, les recollent pour reconstituer la page entière, Muon apprend, puis ils déchirent immédiatement la page à nouveau.
- L'Inconvénient : C'est comme un groupe de personnes qui s'arrête constamment pour assembler un immense puzzle géant, juste pour le démonter à nouveau. Cela gaspille énormément de temps et d'énergie (communication) à chaque étape.
La Méthode de la « Copie Complète » : Au lieu de partager des morceaux, chaque robot garde une copie complète de l'intégralité du livre. Muon peut apprendre instantanément car tout le monde possède l'image complète.
- L'Inconvénient : Cela nécessite tellement de mémoire que si le livre devient trop grand, les cerveaux des robots (GPU) explosent. Ils tombent en panne de place.

La Nouvelle Solution : MatrixFSDP

Les auteurs de ce papier, MatrixFSDP, ont trouvé une troisième voie ingénieuse. Ils n'ont pas changé la règle d'apprentissage (Muon) ni forcé tout le monde à porter le livre entier. Ils ont changé qui détient le livre.

L'Analogie : Le « Bibliothécaire Spécialisé »

Imaginez une bibliothèque où les livres sont habituellement découpés et distribués entre tous les bibliothécaires.

L'idée de MatrixFSDP : Pour chaque « page » (matrice) du livre, ils nomment un bibliothécaire spécifique pour être le « Propriétaire ».
- Ce Propriétaire détient la page entière et complète.
- Tous les autres bibliothécaires ne détiennent rien (un espace vide) pour cette page spécifique.
- Pour les parties du livre qui n'ont pas besoin de la règle spéciale Muon, ils s'en tiennent à l'ancienne méthode des « pièces de puzzle ».

Comment cela fonctionne en pratique :

Pendant l'Apprentissage (L'étape de l'Optimiseur) : Puisque le « Propriétaire » possède déjà la page entière, Muon peut apprendre immédiatement. Personne n'a besoin de rassembler les pièces ou de rien coller. C'est comme si le bibliothécaire lisait le livre directement sur son bureau. Aucune communication n'est nécessaire.
Pendant la Lecture/Écriture (Passages Forward/Backward) : Lorsque les robots ont besoin de lire ou d'écrire dans le livre, ils rassemblent temporairement les pièces, font leur travail, puis les remettent immédiatement dans leurs emplacements de « Propriétaire ».

Pourquoi c'est une avancée majeure

Le papier affirme que cette approche résout le plus gros goulot d'étranglement de l'entraînement des grands modèles d'IA :

Vitesse : Parce qu'ils ont arrêté le constant « collage et déchirement » des pages, l'étape d'apprentissage est devenue incroyablement rapide. Sur un seul nœud informatique, c'était 4,2 fois plus rapide. Sur un grand cluster de 8 nœuds, c'était 54,6 fois plus rapide, car l'ancienne méthode perdait du temps à envoyer des données à travers le réseau entre les ordinateurs, tandis que MatrixFSDP garde les données locales.
Mémoire : Contrairement à la méthode de la « Copie Complète », MatrixFSDP n'utilise toujours que la mémoire de la méthode des « pièces de puzzle ». Cela permet d'entraîner des modèles trop volumineux pour la méthode de la « Copie Complète ».
Précision : Ils ont prouvé que, puisque le « Propriétaire » reçoit exactement les mêmes données que s'il avait rassemblé tout le livre, les résultats d'apprentissage sont identiques à la méthode parfaite, bien que plus lente.

Résumé

MatrixFSDP revient à réorganiser une équipe de sorte que la personne qui a besoin du document complet pour faire son travail soit celle qui détient réellement le document complet. Tous les autres ne détiennent rien pour cette tâche spécifique. Cela élimine le besoin de transmettre constamment des documents d'un côté à l'autre, ce qui permet à l'équipe de travailler beaucoup plus vite sans avoir besoin de bureaux plus grands (plus de mémoire) pour tout contenir.

Résumé Technique : MatrixFSDP

Problématique

L'entraînement de modèles de langage à grande échelle repose de plus en plus sur des optimiseurs à structure matricielle comme Muon, Shampoo et SOAP. Contra irement aux optimiseurs par coordonnée (ex. AdamW), ces méthodes nécessitent la matrice de gradient 2D complète en entrée pour effectuer des opérations telles que l'orthogonalisation de Newton–Schulz ou le préconditionnement de matrice. Cela crée un décalage systémique fondamental avec FSDP (Fully Sharded Data Parallel) et ZeRO-3, qui fragmentent les paramètres, les gradients et les états de l'optimiseur entre les rangs pour maximiser l'efficacité de la mémoire.

Sous ZeRO-3 standard, chaque rang ne détient qu'une tranche du modèle. Pour appliquer un optimiseur matriciel, le système doit reconstruire la matrice de gradient 2D complète à chaque étape de l'optimiseur. Cela nécessite un all-gather de la taille des poids après la passe arrière (backward pass). Cette communication est un goulot d'étranglement sur le chemin critique qui ne peut être superposé au calcul des couches, entraînant une latence significative qui s'adapte mal au nombre de nœuds et à la taille du modèle.

Les solutions existantes sont incomplètes :

Reconstruction (Stock FSDP2-Muon) : Maintient les économies de mémoire de ZeRO-3 mais induit une lourde surcharge de communication à chaque étape de l'optimiseur.
Placement de type ZeRO-1 Owner : Assigne des matrices entières à des rangs spécifiques (propriétaires/owners) pour rendre l'étape d'optimisation locale, mais nécessite que les paramètres complets résident sur chaque rang, abandonnant ainsi l'efficacité mémoire de ZeRO-3.

L'objectif est d'obtenir des entrées d'optimiseur de matrice entière sans sacrifier la mémoire des paramètres de l'échelle ZeRO-3.

Méthodologie : MatrixFSDP

MatrixFSDP introduit une troisième voie : les fragments ZeRO-3 à forme de propriétaire (Owner-shaped ZeRO-3 shards). Au lieu de modifier l'algorithme de l'optimiseur ou d'abandonner le fractionnement, il redéfinit où les fragments résident au sein du groupe de parallélisme de données.

1. Placement de forme Propriétaire (Owner-Shaped Placement)

Pour chaque matrice de poids 2D du modèle, MatrixFSDP désigne un rang de parallélisme de données unique comme propriétaire (owner).

Rang Propriétaire : Détient la matrice 2D complète.
Rangs Non-Propriétaires : Détiennent un fragment vide pour ce paramètre spécifique.
Tenseurs Non-Matrices : Les vecteurs 1D et autres tenseurs non-2D sont regroupés dans une "queue" (tail) et assignés à un propriétaire distinct, où ils continuent d'utiliser l'optimiseur AdamW standard.

Cela garantit qu'après la passe arrière et la réduction des gradients, l'entrée complète de Muon (la matrice de gradient 2D complète) réside naturellement sur le rang propriétaire. Par conséquent, la mise à jour de Newton–Schulz s'exécute localement sur le propriétaire avec zéro communication de matrice lors de l'étape de l'optimiseur.

2. Défis et Solutions au Runtime

Cette disposition inégale (certains rangs détenant des matrices complètes, d'autres des fragments vides) viole les hypothèses de taille égale des collectifs standards de FSDP2. MatrixFSDP traite cela via quatre mécanismes clés au runtime :

Métadonnées MatrixShard & Planificateur Global : Un planificateur global assigne des propriétaires pour équilibrer les octets résidents et la charge de travail de l'optimiseur à travers le cluster. Ces métadonnées suivent quel rang possède quel rôle de tenseur, garantissant que le système reste une structure ZeRO-3 cohérente plutôt qu'une collection de cas particuliers.
Collectifs P2P Segment-Propriétaire : Les all-gathers standards sont remplacés par des collectifs segment-propriétaire déterministes. Pendant la matérialisation forward/backward, seuls les segments non vides sont transférés du propriétaire vers les rangs en ayant besoin. Pendant la réduction de gradient, les résultats sont écrits uniquement dans les segments non vides du propriétaire. Cela évite de remplir les rangs vides et élimine les all-gathers de matrice complète.
Verrouillage de Buffer Propriétaire (Owner-Buffer Pinning) : Pour gérer le cycle de vie FSDP2 (unshard $\to$ compute $\to$ reshard), MatrixFSDP introduit un état "épinglé" (pinned). Lorsqu'un buffer est réduit après la passe forward, le stockage n'est pas retourné au pool partagé mais est épinglé. Cela permet à la passe backward de "revivre" la matrice complète dans le même emplacement mémoire, empêchant les erreurs de vue autograd et garantissant que le propriétaire peut effectuer la mise à jour avant de libérer le buffer.
Checkpointing et Routage : Le checkpointing (DCP) et le routage des états de l'optimiseur utilisent les métadonnées MatrixShard. Les états de l'optimiseur sont routés vers les bons propriétaires (Muon pour les matrices, AdamW pour les queues), et les checkpoints sauvegardent la disposition spécifique du placement pour supporter le resharding à travers différentes tailles de monde.

Contributions Clés

ZeRO-3 de forme Propriétaire (Owner-Shaped ZeRO-3) : Une nouvelle stratégie de placement où les matrices 2D sont entières sur un seul propriétaire et vides ailleurs, permettant une optimisation de matrice locale dans un régime de mémoire fractionnée.
Étapes d'Optimiseur Sans Communication : En alignant la destination de la réduction de gradient avec l'exigence d'entrée de l'optimiseur, MatrixFSDP élimine le goulot d'étranglement de l'all-gather de la taille des poids précédemment inhérent aux optimiseurs matriciels dans ZeRO-3.
Infrastructure de Runtime : Une suite de composants runtime (planificateur, collectifs personnalisés, épinglage de buffer) qui rendent les dispositions de propriétaires inégales efficaces et correctes au sein du cadre FSDP2.
Généralité : L'approche est agnostique vis-à-vis de l'optimiseur, démontrant qu'elle fonctionne avec Muon, Shampoo et SOAP sans nécessiter de variantes distribuées des préconditionneurs.

Résultats Expérimentaux

Évalué sur 64 GPU NVIDIA A100 (8 nœuds) avec diverses tailles de modèles (jusqu'à 32B paramètres) :

Latence de l'Étape d'Optimiseur : MatrixFSDP réduit la latence de l'étape d'optimiseur de 4,2× sur un seul nœud et de 54,6× sur huit nœuds par rapport au FSDP2-Muon standard. L'accélération augmente avec l'échelle car la communication de base traverse la structure inter-nœuds, tandis que l'étape de MatrixFSDP reste locale.
Accélération End-to-End : Atteint jusqu'à 2,15× d'accélération d'entraînement end-to-end.
Efficacité Mémoire : MatrixFSDP maintient les niveaux de mémoire ZeRO-3. Pour un modèle de 32B, il utilise 10 Go/rang, alors qu'une approche de placement propriétaire ZeRO-1 dépasse la limite de 80 Go par GPU (projetée à >54 Go/rang pour 14B). Il est 3 à 7× plus efficace en mémoire que les approches "gather-once" de FSDP2-Muon.
Correction : La trajectoire d'entraînement correspond à une référence DDP Muon à matrice complète avec une haute précision (différence de perte imprimée $\leq 0$ ), confirmant que le propriétaire reçoit la matrice de gradient exacte requise pour la mise à jour.

Signification et Revendications

Le papier affirme que MatrixFSDP résout le "décalage systémique" entre les optimiseurs matriciels et le fractionnement des paramètres. Sa principale importance réside dans la capacité à permettre l'entraînement à grande échelle avec des optimiseurs matriciels (qui offrent une meilleure efficacité de jetons et de convergence) sans subir les coûts de communication prohibitifs ou les surcharges de mémoire des approches précédentes.

Les auteurs soulignent que MatrixFSDP est particulièrement efficace pour les exécutions ZeRO-3 multi-nœuds avec de grands étendues de partage (shard spans), où le coût de reconstruction des matrices à chaque étape dominerait autrement le temps d'entraînement. Le papier note modestement que les gains sont plus faibles sur les exécutions mono-nœuds ou sous un fort passage à l'échelle (strong scaling) à modèle fixe où le nombre de rôles de propriétaires devient épars, rendant potentiellement le fanout du propriétaire un nouveau goulot d'étranglement. Le travail ne prétend pas remplacer les optimiseurs par coordonnée pour les paramètres non-matrices, mais propose plutôt une voie unifiée et efficace pour gérer les composants de matrice 2D des modèles modernes.