LoRDO: Distributed Low-Rank Optimization with Infrequent Communication

Technische Samenvatting: LoRDO – Distributed Low-Rank Optimization met Infrequente Communicatie

1. Probleemstelling

De gedistribueerde training van fundatiemodellen via Distributed Data Parallel (DDP) wordt steeds vaker beperkt door de bandbreedte van de interconnect en de geheugenoverhead die nodig is voor het opslaan van optimizer-states (bijv. Adam-momentum). Hoewel strategieën die gebruikmaken van infrequente communicatie (lokale updates) de synchronisatiefrequentie verminderen, blijven ze een bottleneck vormen door de geheugen- en communicatiekosten van full-rank optimizer-states.

Low-rank adaptieve optimizers (bijv. GaLore, LDAdam) bieden een weg om deze beperkingen te verlichten door gradiënten te projecteren in een low-rank subspace. Het generaliseren van deze methoden naar het regime van infrequente synchronisatie vormt echter een kritieke optimalisatieuitdaging:

1. Lokale Projecties: Het lokaal berekenen van low-rank projecties op individuele workers met behulp van stochastische gradiënten introduceert aanzienlijke ruis door kleine effectieve batchgroottes per worker. Dit leidt tot suboptimale projectie-subspaces en prestatievermindering.
2. Globale Projecties: Het gebruik van een globale projectiematrix afgeleid van geaggregeerde pseudo-gradiënten (zoals in DDP) biedt een stabielere basis. Echter, in het regime van lokale updates beperkt deze aanpak de optimalisatiet traject permanent tot een vaste low-rank subspace. Zonder full-rank signalen stagneert de optimalisatie omdat de projectiebasis niet in staat is om te "verversen" of nieuwe richtingen te verkennen buiten de initiële subspace.

De kernvraag die wordt geadresseerd is: Hoe kunnen we hoogwaardige low-rank optimizers ontwerpen voor communicatie-efficiënte training die subspace-stagnatie voorkomen, terwijl we de low-rank geheugen- en communicatieoverhead behouden?

2. Methodologie: LoRDO

De auteurs stellen LoRDO (Low-Rank Distributed Optimization) voor, een framework dat low-rank optimalisatie verenigt met infrequente synchronisatie. De methodologie adresseert de spanning tussen projectiestabiliteit en subspace-exploratie via drie belangrijke mechanismen:

A. Globale Projectiematrices

In plaats van dat elke worker zijn eigen projectiematrix $Q_m$ bepaalt op basis van lokale stochastische gradiënten (wat lijdt onder hoge variantie), gebruikt LoRDO een globale projectiematrix $Q_t$. Deze matrix wordt berekend bij de synchronisatiegrens met behulp van de geaggregeerde pseudo-gradiënt ($\Delta_t$), die de totale verandering in modelparameters over $K$ lokale stappen vertegenwoordigt.

• Voordeel: De pseudo-gradiënt heeft een effectieve batchgrootte van $M \times B$ (waarbij $M$ het aantal workers is), wat resulteert in een stabieler signaal met verminderde variantie vergeleken met lokale gradiënten. Het incorporeert ook kromminginformatie die in de pseudo-gradiënt besloten ligt, wat leidt tot een kwalitatief betere basis-inschatting.

B. Full-Rank Quasi-Hyperbolic Momentum (QHM)

Om het probleem van de "stagnante subspace" dat inherent is aan globale projecties op te lossen, injecteert LoRDO een full-rank quasi-hyperbolic momentum term in de lokale update.

• Mechanisme: De update regel combineert de low-rank momentum met een full-rank gradiëntsignaal geschaald door de tweede momentum.
  $$\text{Update} \propto (1-\omega)\frac{\hat{G}_t}{\mu(\sqrt{\hat{v}_t} + \epsilon)} + \omega Q_t \left[ \frac{u_t}{\sqrt{v_t} + \epsilon} \right]$$
• Effect: Deze injectie voorkomt dat de geaggregeerde pseudo-gradiënt gevangen blijft binnen een vaste rank-$r$ subspace. Het stelt de optimalisatieprocedure in staat om rank-$r$ subspaces van verschillende workers temporeel te aggregeren, wat uiteindelijk volledige subspace-exploratie mogelijk maakt. Dit herstelt het vermogen om de volledige oplossingsruimte te verkennen zonder de communicatieoverhead te verhogen.

C. Error Feedback en Momentum Rotatie

LoRDO incorporeert error feedback lokaal om informatie te compenseren die verloren gaat tijdens de low-rank compressie. Daarnaast worden, na de berekening van een nieuwe projectiematrix, de optimizer-states (momenta) geroteerd om ervoor te zorgen dat ze geaccumuleerd blijven op dezelfde subspace, een techniek aangepast van LDAdam.

3. Belangrijkste Bijdragen

1. Analyse van Projectie-foutmodi: De auteurs demonstreren dat lokale projecties de prestaties schaden door hoge variantie vanuit kleine worker batchgroottes, terwijl globale projecties zonder full-rank signalen leiden tot subspace-stagnatie.
2. Herstel van Volledige Subspace-exploratie: LoRDO introduceert een low-rank optimizer die een full-rank gradiëntsignaal injecteert (via QHM) in lokale updates, terwijl een globale projectie behouden blijft die afgeleid is van geaggregeerde pseudo-gradiënten. Dit voorkomt stagnatie zonder de communicatie- of geheugenoverhead te verhogen.
3. DDP Pariteit en Efficiëntie: Het framework bereikt bijna-pariteit met synchrone low-rank DDP (perplexity gap < 1%) terwijl het de optimizer geheugen en communicatie met wel $8\times$–$12\times$ vermindert. Onder zware geheugenbeperkingen die lage ranks noodzakelijk maken, overtreft LoRDO zelfs DDP in perplexity.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op decoder-only transformers geschaald naar 16M, 125M en 720M parameters met de SmolLM2 data-mix.

• Stagnatie vs. Exploratie: In 16M modellen faalden globale projecties zonder QHM om te leren (perplexity 46.8) vergeleken met DDP (36.6), wat subspace-stagnatie bevestigt. Het toevoegen van full-rank QHM herstelde de prestaties, waardoor LoRDO-Global de lokale varianten kon overtreffen.
• Projectiekwaliteit: LoRDO-Global demonstreerde superieure stabiliteit en prestaties, vooral bij lagere ranks ($r=8, 16$), dankzij de hogere kwaliteit van de projectiebasis afgeleid van de geaggregeerde pseudo-gradiënt. Het vertoonde minder gevoeligheid voor veranderingen in de lokale batchgrootte vergeleken met lokale projectiemethoden.
• Opschaling naar 720M: Bij de 720M schaal bereikte LoRDO-Global ($r=256$) bijna-pariteit met low-rank DDP ($r=256$) met een perplexity gap van minder dan 1% (10.41 vs. 10.34). Het mat evenals de downstream taak nauwkeurigheid (ARC, HellaSwag, MMLU, PIQA) met DDP.
• Efficiëntiewinsten:
  • Communicatie: LoRDO vermindert de communicatieoverhead met ongeveer 10$\times$ vergeleken met low-rank DDP en 25$\times$ vergeleken met full-rank DDP.
  • Geheugen: Het vermindert de optimizer state geheugenoverhead met 8$\times$ vergeleken met full-rank tegenhangers.
  • Wall-Clock Tijd: In omgevingen met lage bandbreedte (gemodelleerd op 10 Gbps) presteert LoRDO aanzienlijk beter dan DDP in wall-clock tijd vanwege de verminderde synchronisatiefrequentie.

5. Betekenis en Claims

Het artikel claimt dat LoRDO de fundamentele spanning tussen projectiestabiliteit en subspace-exploratie in communicatie-efficiënte low-rank optimalisatie oplost. Door gebruik te maken van geaggregeerde pseudo-gradiënten voor projectiestabiliteit en full-rank quasi-hyperbolic momentum voor exploratie, maakt LoRDO het volgende mogelijk:

• Schaalbare Training: Competitieve prestaties op schaal (tot 720M parameters) met aanzienlijk lagere communicatie- en geheugenkosten.
• Resource Efficiëntie: Mogelijkheid tot pre-training op hardware met beperkte middelen (laag geheugen/bandbreedte) waar full-rank optimizers onhaalbaar zijn.
• Gedecentraliseerd Leren: Het leggen van de basis voor efficiënt leren in gedecentraliseerde omgevingen door modeltraining uit te breiden buiten enkele datacentra.

De auteurs merken beperkingen op, waaronder het gebrek aan een formeel convergentiebewijs en de huidige focus op adaptieve optimizers (Adam) in plaats van matrix-gebaseerde optimizers (bijv. Muon), wat zij overlaten aan toekomstig werk. Zij stellen echter dat de geobserveerde trends in kleinere modellen standhouden bij grotere schalen.