Technische Samenvatting: FORGE (Fused On-Register Gradient Elimination)

1. Het Probleem: De Gradiëntvloer

Huidige geheugenefficiënte trainingsmethoden voor Large Language Models (LLM's) richten zich primair op drie van de vier geheugenpools op een GPU: optimizer states (via compressie), activaties (via recomputatie) en modelgewichten. Ze laten echter de vierde pool ongemoeid—de gematerialiseerde gewichtsgradiënt.

In de standaard reverse-mode differentiatie berekent de backward pass de volledige gradiënttensor ($\nabla W$) voor elke laag, schrijft deze naar het High-Bandwidth Memory (HBM), waarna de optimizer de gradiënt pas inleest om de gewichten bij te werken. Deze tweefasige planning creëert een "gradiëntvloer" waarbij de gradiënten voor alle lagen gelijktijdig in het geheugen moeten bestaan op de grens tussen de backward pass en de optimizer stap. Voor een 8B-parameter model consumeert deze transiënte gradiëntpool alleen al ongeveer 16 GiB, wat vaak bepaalt of een trainingsconfiguratie op een enkele GPU past. Bestaande methoden die de backward en optimizer stappen fuseren (bijv. LOMO, FlashOptim) opereren op het niveau van volledige parameters of lagen, wat betekent dat de volledige gradiënttensor voor een parameter nog steeds in het geheugen gematerialiseerd moet worden voordat de update plaatsvindt.

2. Methodologie: FORGE

FORGE (Fused On-Register Gradient Elimination) elimineert de gematerialiseerde gradiënttensor volledig door de optimizer stap direct te fuseren met de gewichtsgradiëntberekening op een tile-granulariteit (specifiek $128 \times 128$ tiles).

Kernmechanisme

In plaats van de volledige gradiëntmatrix te berekenen en op te slaan, verwerkt FORGE de gewichtsmatrix één tile tegelijk:

1. Tile-wise Gradiëntaccumulatie: Voor een specifieke gewichtstitle accumuleert de kernel de partiële producten van de upstream gradiënt ($\Delta Y$) en de input-activatie ($X$) direct in FP32 GPU-registers.
2. In-Register Optimalisatie: Zodra de gradiënt van de tile volledig in de registers is geaccumuleerd, wordt de optimizer update (bijv. AdamW) direct op die tile toegepast, eveneens binnen de registers.
3. Directe Verwijdering: De gradiënt-tile wordt direct na productie geconsumeerd. Deze wordt nooit als een tensor naar het HBM geschreven.
4. Write-Back: Alleen de bijgewerkte gewichtstitle en de (optioneel gekwantiseerde) optimizer state tiles (moments) worden teruggeschreven naar het globale geheugen.

Precisie en Getrouwheid

• Exactheid: Voor elke element-wise-separable optimizer (waarbij de update voor een gewicht alleen afhangt van de coördinaat van dat gewicht, bijv. AdamW, SGD, Lion), is de per-tile update wiskundig identiek aan een standaard volledige gradiëntstap. De enige afwijking is de volgorde van de floating-point sommatie over tokens, wat inherent is aan elke GEMM-implementatie en geen benadering van de optimizer-logica vormt.
• Low-Precision Getrouwheid: In BF16 of INT8 trainingsregimes behoudt FORGE de volledige precisie-getrouwheid op twee manieren:
  • De gradiënt wordt niet naar BF16 afgekapt voordat de optimalisatie plaatsvindt; deze wordt berekend en geconsumeerd in FP32 registers.
  • Om de "dead zone" van BF16 (waar kleine updates naar nul worden afgerond) te voorkomen, gebruikt FORGE stochastische rounding voor de gewicht write-backs, wat zorgt voor een onbevooroordeelde accumulatie van sub-ULP updates.
  • Voor INT8 optimizer states gebruikt FORGE kleinere blokgroottes (64 vs 2024 in standaard 8-bit optimizers) om kwantiseringsfouten te minimaliseren.

Compatibiliteit met Gedistribueerde Training

FORGE is bewezen exact onder Tensor Parallelism (TP) en Sequence Parallelism (SP) omdat de lokale gradiënt voor een shard compleet is op het moment van de backward pass. Het is echter incompatibel met Data Parallelism (DP) of Context Parallelism in de huidige vorm, aange omdat deze methoden een all-reduce van partiële gradiënten vereisen voordat de update kan worden toegepast, wat het materialiseren van de gradiënttensor noodzakelijk zou maken.

3. Belangrijkste Bijdragen

1. Structurele Eliminatie, Niet Compressie: FORGE identificeert de gematerialiseerde gradiënt als een artefact van de trainingsplanning in plaats van een wiskundige noodzaak. Het elimineert de $2P$ (parameter aantal) gradiëntgeheugenterm, waardoor de gradiëntpool wordt gereduceerd van $O(P)$ naar $O(T_R T_C)$ (tile grootte), onafhankelijk van de modelgrootte.
2. Bewijsbare Exactheid: Het artikel bewijst dat voor element-wise optimizers de tile-fused planning updates produceert die identiek zijn aan standaard AdamW op FP32 gradiënten. Deze exactheid blijft behouden onder tensor en sequence parallelism.
3. Full-Precision Getrouwheid bij Lage Precisie: Door gradiënten in FP32 registers te houden tot de optimalisatie, vermijdt FORGE de bias die wordt geïntroduceerd door het afkappen van gradiënten naar BF16 vóór de optimalisatie. Gecombineerd met stochastische rounding, komt dit overeen met FP32-master convergentie, zelfs bij gebruik van BF16 of INT8 states.
4. Generaliteit: De methode is architectuur-agnostisch voor lineaire lagen en ondersteunt acht volledig gefuseerde element-wise optimizers (AdamW, NAdam, RAdam, Lion, RMSprop, AdaGrad, SGD, SGD-momentum).

4. Experimentele Resultaten

Geëvalueerd op NVIDIA H200 en B200 GPU's over modellen variërend van 0.6B tot 14B parameters (Llama-3.1-8B, Qwen3 familie):

• Geheugenreductie: FORGE vermindert het piekgeheugen met 38–53% vergeleken met gefuseerde AdamW. Voor Llama-3.1-8B daalt het piekgeheugen van ~75 GiB naar 35.4 GiB (een reductie van 53%).
• Doorvoer: FORGE bereikt een 1.52× versnelling in staptijd (99.1 ms vs 150.9 ms) door de aparte optimizer kernel launch en de HBM round-trip voor gradiënten te elimineren.
• Schaalbaarheid:
  • Op een enkele H200 traint FORGE Qwen3-14B (wat met standaard gefuseerde AdamW tot Out-Of-Memory fouten leidt) op 63.2 GiB.
  • In een 4-GPU A100-40GB setup maakt FORGE een micro-batch grootte van 4 mogelijk (past in 24.6 GiB), terwijl standaard gefuseerde AdamW slechts micro-batch 1 toelaat. Dit stelt FORGE in staat om een 2.42× hogere doorvoer te halen dan de standaard baseline op dezelfde hardware.
• Convergentie: FORGE volgt de gefuseerde AdamW convergentie binnen 0.001 nats (gemiddelde afwijking) op continued pretraining taken, wat bevestigt dat de low-precision afwijkingen binnen de run-to-run ruis vallen.
• Composabiliteit: FORGE combineert probleemloos met FP8 kwantisatie en activation checkpointing, waarbij de dekwantiseringskosten die standaard optimizers doorgaans straffen, worden geabsorbeerd.

5. Betekenis en Claims

Het artikel claimt dat FORGE een structurele verschuiving vertegenwoordigt in plaats van een incrementele optimalisatie. Door de gradiëntvloer te verwijderen, zet het optimizer geheugenoverhead om in bruikbare ruimte voor grotere batches, langere sequenties of grotere modellen op bestaande hardware.

• Capabiliteitsgrens: De auteurs stellen dat de eliminatie van de gradiëntvloer niet slechts een constante-factor verbetering is, maar een capabiliteitsgrens. Zo is bijvoorbeeld de standaard training configuratie voor 14B modellen op enkele H200's onmogelijk door de gradiëntvloer, terwijl FORGE dit wel mogelijk maakt.
• Orthogonaliteit: De methode is orthogonaal aan andere geheugenbesparende technieken (zoals state kwantisatie of activation recomputation) en kan worden gecombineerd met deze om het geheugenvoetafdruk verder te verkleinen.
• Beperkingen: De methode is momenteel beperkt tot element-wise separable optimizers en ondersteunt geen data parallelism zonder aanpassingen. Het vereist ook dat de input gradiënt wordt gelezen voordat de gewicht update plaatsvindt, wat het gebruik van tied embeddings of MoE experts met meerdere gradiënt consumenten in een enkele stap uitsluit.

Samenvattend demonstreert FORGE dat de "gradiëntvloer" een onnodig artefact is van huidige trainingsloops. Door de optimizer op tile-niveau in de backward pass te fuseren, bereikt het significante geheugenbesparingen en snelheidswinsten terwijl de wiskundige exactheid behouden blijft voor een brede klasse van optimizers.