💻 computer science

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

MatrixFSDP maakt communicatievrije grootschalige training mogelijk met matrix-optimizers zoals Muon onder ZeRO-3 sharding door de plaatsing van parameters te reorganiseren zodat elke 2D-gewichtsmatrix volledig op een enkele rank verblijft, waardoor de noodzaak voor kostbare matrixreconstructie tijdens optimizer-stappen wordt geëlimineerd terwijl de geheugenefficiëntie behouden blijft en significante latentiereducties worden bereikt.

Oorspronkelijke auteurs: Ming Gao, Yanwu Xu, Hao Zhang

Gepubliceerd 2026-07-08✓ Author reviewed ⓘ

📖 4 min leestijd☕ Koffiepauze-leesvoer

CC BY 4.0

Oorspronkelijke auteurs: Ming Gao, Yanwu Xu, Hao Zhang

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een enorm team robots (een computercluster) probeert te leren hoe ze een roman moeten schrijven. De robots werken samen om te leren van een gigantisch boek vol tekst. Om dit efficiënt te doen, gebruiken ze een speciale "leerregel" genaamd Muon.

Het Probleem: Het "Volledige Beeld" vs. De "Puzzelstukjes"

Normaal gesproken, wanneer robots leren, breken ze het enorme boek op in kleine puzzelstukjes. Elke robot houdt slechts een paar pagina's vast (een "shard"). Dit is geweldig voor het besparen van geheugen, omdat geen enkele robot het hele boek hoeft te dragen. Deze methode wordt ZeRO-3 genoemd.

Echter, de Muon leerregel is een beetje veeleisend. Hij wil niet leren van slechts een paar pagina's tegelijk. Om zijn werk perfect te doen, moet hij de volledige 2D-pagina (de hele matrix) in één keer zien om de relaties tussen woorden te begrijpen.

Het Conflict:

ZeRO-3 zegt: "Wij hebben alleen maar puzzelstukjes."
Muon zegt: "Ik heb de hele pagina nodig om te leren."

De Oude Oplossingen (De Slechte Opties):

De "Reconstructie"-methode: Elke keer wanneer de robots moeten leren, stoppen ze, verzamelen alle puzzelstukjes van elke robot, plakken ze weer aan elkaar om de volledige pagina te maken, Muon leert, en vervolgens scheuren ze de pagina onmiddellijk weer uit elkaar.
- Het Nadeel: Dit is alsof een groep mensen constant stopt met werken om een enorme legpuzzel in elkaar te zetten, alleen maar om hem daarna direct weer uit elkaar te halen. Dit verspilt een enorme hoeveelheid tijd en energie (communicatie) bij elke stap.
De "Volledige Kopie"-methode: In plaats van puzzelstukjes te delen, houdt elke robot een volledige kopie van het hele boek bij. Muon kan direct leren omdat iedereen het volledige beeld heeft.
- Het Nadeel: Dit vereist zoveel geheugen dat als het boek te groot wordt, de hersenen (GPU's) van de robots ontploffen. Ze raken door hun ruimte heen.

De Nieuwe Oplossing: MatrixFSDP

De auteurs van dit paper, MatrixFSDP, hebben een slimme derde weg gevonden. Ze hebben de leerregel (Muon) niet veranderd, en ze hebben ook niet gedwongen dat iedereen het hele boek draagt. In plaats daarvan hebben ze veranderd wie het boek vasthoudt.

De Analogie: De "Gespecialiseerde Bibliothecaris"

Stel je een bibliotheek voor waar de boeken normaal gesproken in stukken worden gehakt en verdeeld over alle bibliothecarissen.

MatrixFSDP's Idee: Voor elke afzonderlijke "pagina" (matrix) van het boek, benoemen ze één specifieke bibliothecaris als de "Eigenaar".
- Deze Eigenaar houdt de volledige, complete pagina vast.
- Alle andere bibliothecarissen houden niets (een lege ruimte) vast voor die specifieke pagina.
- Voor de delen van het boek die de speciale Muon-regel niet nodig hebben, houden ze vast aan de oude "puzzelstukjes"-methode.

Hoe het in de praktijk werkt:

Tijdens het Leren (De Optimizer Stap): Omdat de "Eigenaar" al de volledige pagina heeft, kan Muon onmiddellijk leren. Niemand hoeft stukjes te verzamelen of iets aan elkaar te plakken. Het is alsof de bibliothecaris het boek gewoon rechtstreeks van zijn bureau leest. Er is nul communicatie nodig.
Tijdens het Lezen/Schrijven (Forward/Backward Passes): Wanneer de robots de boeken moeten lezen of schrijven, brengen ze de stukjes tijdelijk weer samen, doen hun werk, en plaatsen ze ze dan onmiddellijk terug in hun "Eigenaar"-slots.

Waarom dit een Groot Ding is

Het paper beweert dat deze aanpak de grootste flessenhals bij het trainen van grote AI-modellen oplost:

Snelheid: Omdat ze stopten met het constante "plakken en verscheuren" van pagina's, werd de leerstap ongelooflijk snel. Op een enkele computer node was het 4,2 keer sneller. Op een grote cluster van 8 nodes was het 54,6 keer sneller, omdat de oude methode veel tijd verspilde aan het versturen van gegevens over het netwerk tussen computers, terwijl MatrixFSDP de gegevens lokaal houdt.
Geheugen: In tegen tegenstelling tot de "Volledige Kopie"-methode, gebruikt MatrixFSDP nog steeds slechts het geheugen van de "puzzelstukjes"-methode. Het stelt hen in staat om modellen te trainen die te groot zijn voor de "Volledige Kopie"-methode om te verwerken.
Nauwkeurigheid: Ze hebben bewezen dat omdat de "Eigenaar" exact dezelfde gegevens krijgt als wanneer hij de hele boeken zou verzamelen, de leerresultaten identiek zijn aan de perfecte, trage methode.

Samenvatting

MatrixFSDP is als het reorganiseren van een team zodat de persoon die het volledige document nodig heeft om zijn werk te doen, ook daadwerkelijk het volledige document vasthoudt. Iedereen anders houdt niets vast voor die specifieke taak. Dit elimineert de noodzaak om constant documenten heen en weer te sturen, waardoor het team veel sneller werkt zonder dat ze grotere bureaus (meer geheugen) nodig hebben om alles vast te houden.

Technische Samenvatting: MatrixFSDP

Probleemstelling

De training van grootschalige taalmodellen leunt steeds vaker op matrix-gestructureerde optimizers zoals Muon, Shampoo en SOAP. In tegenstelling tot coördinaat-gewijze optimizers (bijv. AdamW), vereisen deze methoden de volledige 2D-gradiëntmatrix als input om operaties uit te voeren zoals Newton–Schulz orthogonalisatie of matrix-preconditionering. Dit creëert een fundamentele systeemmismatch met FSDP (Fully Sharded Data Parallel) en ZeRO-3, die parameters, gradiënten en optimizer-states over verschillende ranks sharden om de geheugenefficiëntie te maximaliseren.

Onder standaard ZeRO-3 bezit elke rank slechts een deel van het model. Om een matrix-optimizer toe te passen, moet het systeem bij elke optimizer-stap de volledige 2D-gradiëntmatrix reconstrueren. Dit vereist een all-gather communicatie ter grootte van het gewicht na de backward pass. Deze communicatie is een kritieke pad-bottleneck die niet kan worden overlapt met de berekening van lagen, wat leidt tot significante latentie die slecht schaalt met het aantal nodes en de modelgrootte.

Bestaande oplossingen zijn incompleet:

Reconstructie (Stock FSDP2-Muon): Behoudt de ZeRO-3 geheugenbesparingen, maar veroorzaakt zware communicatie-overhead bij elke optimizer-stap.
ZeRO-1 Owner Placement: Wint volledige matrices toe aan specifieke ranks (owners) om de optimizer-stap lokaal te maken, maar vereist dat volledige parameters op elke rank aanwezig zijn, waardoor de ZeRO-3 geheugenefficiëntie wordt opgeofferd.

Het doel is om volledige matrix-inputs te verkrijgen zonder de ZeRO-3-schaal van parametergeheugen op te offeren.

Methodologie: MatrixFSDP

MatrixFSDP introduceert een derde pad: Owner-shaped ZeRO-3 shards. In plaats van de optimizer-algoritmen te veranderen of de sharding op te geven, herdefinieert het waar de shards zich bevinden binnen de data-parallel groep.

1. Owner-Shaped Placement

Voor elke 2D-gewichtmatrix in het model wijst MatrixFSDP een enkele data-parallel rank aan als de owner.

Owner Rank: Bezit de volledige 2D-matrix.
Niet-Owner Ranks: Houden een lege shard vast voor die specifieke parameter.
Niet-Matrix Tensors: 1D-vectoren en andere niet-2D tensors worden samengepakt in een "tail" en toegewezen aan een aparte owner, waar ze doorgaan met het gebruik van de standaard AdamW optimizer.

Dit zorgt ervoor dat na de backward pass en de gradiëntreductie, de volledige Muon-input (de volledige 2D-gradiënt) van nature op de owner-rank terechtkomt. Bijgevolg vindt de Newton–Schulz update lokaal plaats op de owner met nul matrix-communicatie tijdens de optimizer-stap.

2. Runtime Uitdagingen en Oplossingen

Deze ongelijkmatige lay-out (sommige ranks houden volledige matrices vast, andere lege shards) schendt de aannames van gelijke grootte bij collectieve operaties in standaard FSDP2. MatrixFSDP pakt dit aan via vier cruciale runtime-mechanismen:

MatrixShard Metadata & Global Planner: Een globale planner wijst owners toe om de residentie van bytes en de optimizer-workload over de cluster te balanceren. Deze metadata houdt bij welke rank welke tensorrol bezit, waardoor het systeem een enkel coherent ZeRO-3 layout blijft in plaats van een verzameling speciale gevallen.
Owner-Segment P2P Collectives: Standaard all-gathers worden vervangen door deterministische, owner-segment collectieven. Tijdens de forward/backward materialisatie worden alleen niet-lege segmenten van de owner naar de ranks getransfereerd die ze nodig hebben. Tijdens de gradiëntreductie worden resultaten alleen geschreven naar de niet-lege owner-segmenten. Dit voorkomt het opvullen van lege ranks en elimineert full-matrix all-gathers.
Owner-Buffer Pinning: Om de FSDP2-levenscyclus (unshard $\to$ compute $\to$ reshard) af te handelen, introduceert MatrixFSDP een "gepinde" staat. Wanneer een buffer wordt verkleind na de forward pass, wordt de opslag niet teruggegeven aan de gedeelde pool, maar gepind. Dit stelt de backward pass in staat de volledige matrix op dezelfde locatie te "reviven", wat autograd view-fouten voorkomt en ervoor zorgt dat de owner de update kan uitvoeren voordat de buffer wordt vrijgegeven.
Checkpointing en Routing: Checkpointing (DCP) en optimizer state routing gebruiken de MatrixShard metadata. Optimizer states worden gerouteerd naar de juiste owners (Muon voor matrices, AdamW voor tails), en checkpoints slaan de specifieke plaatsingslay-out op om resharding over verschillende wereldgroottes te ondersteunen.

Belangrijkste Bijdragen

Owner-Shaped ZeRO-3: Een nieuwe plaatsingsstrategie waarbij 2D-matrices volledig zijn op één owner en elders leeg, wat lokale matrix-optimalisatie mogelijk maakt binnen een gescharde geheugenregime.
Communicatievrije Optimizer-stappen: Door de bestemming van de gradiëntreductie af te stemmen op de inputvereisten van de optimizer, elimineert MatrixFSDP de weight-sized all-gather bottleneck die voorheen inherent was aan matrix-optimizers in ZeRO-3.
Runtime Infrastructuur: Een suite van runtime-componenten (planner, custom collectieven, buffer pinning) die ongelijke owner-lay-outs efficiënt en correct maken binnen het FSDP2-framework.
Generaliteit: De aanpak is optimizer-agnostisch en werkt bewezen met Muon, Shampoo en SOAP zonder dat er gedistribueerde varianten van de preconditioners nodig zijn.

Experimentele Resultaten

Geëvalueerd op 64 NVIDIA A100 GPU's (8 nodes) met diverse modelgroottes (tot 32B parameters):

Optimizer-Step Latentie: MatrixFSDP vermindert de optimizer-stap latentie met 4,2× op een enkele node en 54,6× op acht nodes vergeleken met stock FSDP2-Muon. De versnelling groeit met de schaal omdat de baseline-communicatie de inter-node fabric kruist, terwijl de MatrixFSDP-stap lokaal blijft.
End-to-End Versnelling: Bereikt tot 2,15× end-to-end trainingsversnelling.
Geheugenefficiëntie: MatrixFSDP behoudt ZeRO-3 geheugenniveaus. Voor een 32B model gebruikt het 10 GB/rank, terwijl een ZeRO-1 owner placement benadering de 80 GB GPU-limiet overschrijdt (geprojecteerd >54 GB/rank bij 14B). Het is 3–7× meer geheugenefficiënt dan "gather-once" FSDP2-Muon benaderingen.
Correctheid: Het trainingsverloop komt met hoge precisie overeen met een full-matrix DDP Muon referentie (geprint verschil in loss $\leq 0$ ), wat bevestigt dat de owner exact de gradiëntmatrix ontvangt die nodig is voor de update.

Betekenis en Claims

Het artikel claimt dat MatrixFSDP de "systeemmismatch" tussen matrix-optimizers en parameter-sharding oplost. De primaire betekenis ligt in het mogelijk maken van grootschalige training met matrix-optimizers (die een betere token-efficiëntie en convergentie bieden) zonder de prohibitieve communicatiekosten of geheugenoverheads van eerdere benaderingen te dragen.

De auteurs benadrukken dat MatrixFSDP bijzonder effectief is voor multi-node ZeRO-3 runs met grote shard spans, waarbij de kosten voor het reconstrueren van matrices bij elke stap anders de trainingstijd zouden domineren. Het artikel merkt bescheiden op dat de winsten kleiner zijn bij single-node runs of onder vaste model strong scaling waarbij het aantal owner-rollen schaarser wordt, wat potentieel de owner fanout een nieuwe bottleneck maakt. Het werk claimt niet de coördinaat-gewijze optimizers voor niet-matrix parameters te vervangen, maar biedt een verenigde, efficiënte weg voor het afhandelen van de 2D-matrix componenten van moderne modellen.