💬 NLP

Out of the Memory Barrier: A Highly Memory Efficient Training System for LLMs with Million-Token Contexts

Dit paper introduceert OOMB, een uiterst geheugenefficiënt trainingssysteem dat door middel van een chunk-recurrente architectuur met on-the-fly herberekening en geoptimaliseerd KV-cache-beheer het trainen van LLMs met contexten van miljoenen tokens mogelijk maakt op één enkele GPU.

Oorspronkelijke auteurs: Wenhao Li, Daohai Yu, Gen Luo, Yuxin Zhang, Fei Chao, Rongrong Ji, Yifan Wu, Jiaxin Liu, Ziyang Gong, Zimu Liao

Gepubliceerd 2026-03-03

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Wenhao Li, Daohai Yu, Gen Luo, Yuxin Zhang, Fei Chao, Rongrong Ji, Yifan Wu, Jiaxin Liu, Ziyang Gong, Zimu Liao

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 of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een enorme bibliotheek wilt bouwen, maar je hebt maar één heel klein bureau (je computerchip). Je wilt alle boeken van de wereld op dat ene bureau leggen om ze te bestuderen. Het probleem? De boeken zijn zo dik en talrijk dat ze je bureau volledig vullen, nog voordat je zelfs maar de eerste zin hebt gelezen.

Dit is precies het probleem waar kunstmatige intelligentie (AI) tegen aanloopt als ze probeert om "lange contexten" te leren. Dat betekent dat de AI teksten moet kunnen lezen die zo lang zijn als een heel boek, of zelfs een hele bibliotheek. Normaal gesproken zou je daar een hele serverruimte met duizenden dure computers voor nodig hebben.

Deze paper introduceert OOMB (Out Of the Memory Barrier), een slimme nieuwe manier om AI te trainen die dit probleem oplost alsof het niets is. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Geheugen-Explosie"

Wanneer een AI een lange tekst leest, moet hij onthouden wat hij eerder heeft gelezen om de context te begrijpen. In de computerwereld noemen we dit de KV-cache (Key-Value cache).

Hoe het nu werkt: Stel je voor dat je een verhaal leest en je schrijft elke zin op een los velletje papier. Als het verhaal 100.000 woorden lang is, heb je 100.000 velletjes papier nodig. Je bureau (het geheugen van je computer) is vol. Je kunt niet verder.
Het oude idee: Om dit op te lossen, bouwden mensen enorme kantoren met duizenden bureaus (duizenden computers) om die papieren te verspreiden. Dat is duur en complex.

2. De Oplossing van OOMB: De "Slimme Werkplek"

OOMB gebruikt een paar slimme trucs om die enorme stapel papier te verminderen tot een paar velletjes, zelfs als het verhaal miljoenen woorden lang is.

Truc 1: De "Vergeet-En-Herhaal" Methode (Chunk-Recurrent)

In plaats van het hele verhaal in één keer op je bureau te leggen, verdeel je het in kleine stukjes (chunks).

Hoe het werkt: Je leest een stukje, schrijft je gedachten op, en gooit het papier daarna direct weg. Je onthoudt alleen de kern van wat je hebt gelezen (de "geheugenspoor").
De magische stap: Als je later terug moet kijken om te leren van je fouten (tijdens het trainen), herbereken je die gedachten direct opnieuw. Het kost een beetje extra tijd om ze te herrekenen, maar het bespaart enorm veel ruimte op je bureau.
Analogie: Het is alsof je een film kijkt. In plaats van elke scène op te slaan op een harde schijf, kijk je naar een scène, onthoud je de plot, en gooi je de scène weg. Als je later moet uitleggen wat er gebeurde, "kijk" je de scène in je hoofd nog één keer na. Je hebt geen opslagruimte nodig voor de hele film, alleen voor het moment dat je nu kijkt.

Truc 2: De "Paged" Opbergkast (Paged Memory)

Zelfs met de "vergeet-en-herhaal" methode moet je de "geheugensporen" van alle eerdere stukjes bewaren. Normaal gesproken groeit deze stapel chaotisch, waardoor er veel ruimte wordt verspild (net als een kast waar je kleding niet goed past en er gaten tussen zitten).

De oplossing: OOMB gebruikt een paginabesturing. Stel je voor dat je boeken niet in een grote stapel legt, maar in een systeem van gelijke vakken (pagina's). Je vult vakje 1, dan vakje 2, dan vakje 3.
Het voordeel: Er is geen verspilde ruimte meer. Je kunt precies weten hoeveel ruimte je nodig hebt, ongeacht hoe lang het verhaal is.

Truc 3: De "Nachtdienst" (Asynchronous Offloading)

Als het verhaal extreem lang wordt (bijvoorbeeld 4 miljoen woorden), is zelfs die georganiseerde kast te groot voor je bureau.

De oplossing: OOMB heeft een nachtdienst (een CPU-geheugen in de achtergrond). Zodra een stukje van het verhaal niet direct nodig is, wordt het direct naar de achtergrond verplaatst.
De magie: Dit gebeurt terwijl de computer bezig is met iets anders. Het is alsof een assistent de boeken naar de kelder brengt terwijl jij nog steeds aan het lezen bent. Je merkt er niets van, maar je bureau blijft leeg.

Truc 4: De "Slimme Zoeker" (Sparse Attention)

Soms hoef je niet elk woord uit het verleden te onthouden om de context te begrijpen.

De oplossing: OOMB is slim genoeg om te weten welke delen van het verleden belangrijk zijn en welke niet. Het zoekt alleen naar de belangrijkste "pagina's" in de kast.
Analogie: Als je een detectiveverhaal leest, hoef je niet elke zin uit hoofdstuk 1 te onthouden om te weten wie de moordenaar is. Je onthoudt alleen de belangrijke aanwijzingen. OOMB doet precies dat: het negeert het ruis en focust op de signalen.

Het Resultaat: Een Wonder op één Chip

Door al deze trucs samen te voegen, kan OOMB een model trainen dat 4 miljoen woorden lang is (dat is een hele bibliotheek!) op één enkele moderne computerchip (een H200 GPU).

Vroeger had je daar een heel datacenter met honderden computers voor nodig. Nu past het op één apparaat.

De kosten: Elke extra 10.000 woorden in het verhaal kosten slechts een beetje extra geheugen (ongeveer 10 MB). Dat is alsof je voor elke extra pagina in een boek slechts een paar cent extra betaalt, in plaats van een heel nieuw gebouw.

Waarom is dit belangrijk?

Dit maakt het mogelijk voor onderzoekers en kleinere bedrijven om superkrachtige AI's te bouwen die hele boeken, juridische dossiers of historische archieven kunnen begrijpen, zonder dat ze miljarden moeten investeren in dure hardware. Het democratiseert de toegang tot de toekomst van AI.

Kortom: OOMB is de slimme verhuizer die ervoor zorgt dat je je hele bibliotheek op één klein bureau kunt organiseren, zodat je er met gemak doorheen kunt werken.

Probleemstelling

Het trainen van Large Language Models (LLMs) op zeer lange contexten (miljoenen tokens) wordt momenteel ernstig beperkt door de GPU-geheugenvraag, niet door de rekentijd. De twee hoofdoorzaken zijn:

Activaties: De geheugenfootprint van de tussenliggende activaties (activations) schaalt lineair met de sequentielengte.
KV-cache: De Key-Value cache, noodzakelijk voor de attention-mechanismen, moet tijdens de hele trainingsstap bewaard blijven en schaalt eveneens lineair met de contextlengte.

Bestaande technieken zoals ZeRO3 en tensor parallelism zijn vaak onvoldoende om deze geheugendruk te hanteren op enkele GPU's. Traditionele methoden vereisen vaak enorme clusters (bijv. 256 GPU's) om contexten van 4 miljoen tokens te trainen, wat de toegankelijkheid en kostenefficiëntie drastisch verlaagt.

Methodologie: Het OOMB-systeem

De auteurs introduceren OOMB (Out Of the Memory Barrier), een trainingsframework dat de geheugenbarrière direct aanpakt door een synergie van vier kerncomponenten:

Chunk-recurrent Training met Activeringsherberekening:
- In plaats van de volledige sequentie parallel te verwerken, wordt de input opgesplitst in kleinere "chunks".
- Tijdens de forward pass worden de activaties van een chunk berekend en direct verwijderd.
- Tijdens de backward pass worden deze activaties on-the-fly herberekend (recomputation).
- Resultaat: De geheugenfootprint voor activaties wordt constant ( $O(1)$ ), ongeacht de totale sequentielengte. Dit verplaatst de bottleneck echter volledig naar de KV-cache.
Gepagineerd Geheugenbeheer (Paged Memory Management):
- Om de groeiende KV-cache en zijn gradiënten efficiënt te beheren, introduceert OOMB een gepagineerd geheugenbeheersysteem (geïnspireerd op vLLM voor inferentie, maar aangepast voor training).
- Dit elimineert dure geheugenherallocaties en fragmentatie bij het toevoegen van nieuwe key-value paren.
- Er worden aangepaste Triton-kernels ontwikkeld die opereren op de KV-cache en gradiënten direct, waardoor de cache niet als een standaard PyTorch-activatie hoeft te worden opgeslagen. Dit zorgt voor in-place gradiëntaccumulatie en een slankere computationele graaf.
Asynchrone CPU Offloading:
- Omdat de KV-cache nog steeds lineair groeit met de contextlengte, wordt een mechanisme ontwikkeld om de cache en zijn gradiënten asynchroon naar het CPU-geheugen te verplaatsen.
- Door gebruik te maken van gepinnd CPU-geheugen en CUDA-streams, wordt de datatransferlatentie volledig gemaskeerd door de berekeningen (prefetching voor dichte attention, overlap voor sparse attention).
Pagina-niveau Sparse Attention:
- Om de kwadratische complexiteit van attention te verminderen, wordt pagina-niveau sparse attention geïntegreerd.
- Voor modellen met dichte attention (zoals Qwen2.5) wordt een Top-K pagina-retrieval methode gebruikt: er wordt een representatief vector per pagina berekend om de meest relevante key-pagina's te selecteren.
- Voor native sparse modellen (zoals GPT-OSS) worden lokale attention-lagen gebruikt.
- Dit vermindert zowel de rekencomplexiteit als de communicatie-overhead bij het offloaden.

Belangrijkste Bijdragen

O(1) Activatie-geheugen: Door chunk-wise training met recomputation wordt het geheugengebruik voor activaties constant, wat een fundamentele doorbraak is voor lange contexten.
Geoptimaliseerde KV-cache Management: De combinatie van gepagineerd geheugen, aangepaste kernels en asynchrone offloading maakt het mogelijk om de lineaire groei van de KV-cache te beheersen op één enkele GPU.
Synergie: Het paper demonstreert dat deze technieken niet alleen werken, maar elkaars effectiviteit versterken (bijvoorbeeld: sparse attention maakt offloading efficiënter door minder data te verplaatsen).

Resultaten

De auteurs hebben OOMB getest op het Qwen2.5-7B model op een enkele H200 GPU:

Geheugenefficiëntie: Voor elke extra 10.000 tokens aan context neemt de totale trainingsgeheugenvraag slechts toe met 10 MB.
Schalbaarheid: Het is mogelijk om een 4 miljoen tokens lange context te trainen op één enkele H200 GPU. Zonder deze technologie zou dit een groot cluster vereisen.
Prestaties:
- OOMB presteert beter dan bestaande parallelle training met FlashAttention en Ring Flash Attention (RFA) op het gebied van geheugengebruik.
- Bij sparse attention wordt een snelheidswinst van 3,7x (bij 64K tokens) tot 13,5x (bij 256K tokens) behaald ten opzichte van dichte attention baselines.
- De training van een 1M-token context op een 4x H200 systeem (met tensor parallelism) kostte ongeveer 10 uur voor 256 stappen, met succesvolle convergentie naar lage loss-waarden.
Nauwkeurigheid: Experimenten tonen aan dat de sparse attention benadering de modelprestaties niet significant degradeert, zelfs niet bij zeer lange contexten.

Betekenis

Dit werk vertegenwoordigt een aanzienlijke vooruitgang in de resource-efficiëntie voor het trainen van lange-context LLMs.

Democratisering: Het maakt het mogelijk voor onderzoekers en instellingen met beperkte hardware (bijv. één high-end GPU) om modellen te trainen op contexten die voorheen alleen met supercomputers mogelijk waren.
Duurzaamheid: Door de hardware-eisen drastisch te verlagen, wordt ook het totale energieverbruik en de koolstofvoetafdruk van het trainen van deze modellen verkleind.
Toekomstperspectief: Het biedt een nieuw paradigma voor lang-context training dat minder afhankelijk is van massale parallelisme en meer focust op slimme geheugentechnieken, hoewel er nog uitdagingen blijven qua latentie voor doorvoergevoelige scenario's en de nauwkeurigheid van sparse attention voor taken die een volledig dichte context vereisen.