MegaScale-Data: Scaling Dataloader for Multisource Large Foundation Model Training

Dit paper introduceert MegaScale-Data, een industriële architectuur voor gedistribueerde dataloaders die de trainingsefficiëntie van multimodale foundation modellen aanzienlijk verbetert door werklastonevenwichtigheid op te lossen en het geheugengebruik te verminderen via gedescentraliseerde voorverwerking en een gecentraliseerde data-orkestratie.

De kern

Stel je voor dat je een gigantisch kookteam leidt dat een enorme maaltijd moet bereiden voor duizenden gasten. Dit is wat er gebeurt wanneer computers "grote foundation modellen" (zoals slimme AI's) trainen. Ze moeten enorme hoeveelheden data "eten" om te leren.

Het probleem? De huidige systemen zijn als een chaotische keuken waar elke kok (de computer) zijn eigen ingrediënten moet halen, maar de recepten zijn zo complex en de ingrediënten zo verschillend, dat het team vastloopt.

Hier is MegaScale-Data, de oplossing van ByteDance en de Universiteit van Hongkong, uitgelegd als een slimme nieuwe keukenorganisatie.

1. Het Oude Probleem: De Chaotische Keuken

Stel je voor dat je een recept hebt dat tekst, foto's en video's combineert.

• Het onevenwicht: Sommige ingrediënten zijn makkelijk te snijden (tekst), andere zijn enorm zwaar om te verwerken (video's of hoge-res foto's). In de oude systemen kregen alle koks een willekeurige stapel ingrediënten. De kok met de zware video's werd overbelast en vertraagde het hele team, terwijl de kok met alleen tekst al klaar was en maar wachtte. Dit heet werklast-ongelijkheid.
• De verspilling: Elke kok had zijn eigen kopie van de ingrediëntenlijst, zijn eigen bestek en zijn eigen ruimte in de koelkast. Als je 100 verschillende soorten ingrediënten hebt, moet elke kok 100 keer zijn eigen lijst bijhouden. Dat neemt enorm veel koelkastruimte (geheugen) weg, terwijl ze eigenlijk allemaal naar dezelfde grote voorraadkast kijken.
• De dynamiek: Soms wil je eerst simpele recepten doen en later moeilijke (zoals een leerling die eerst leert lopen en dan rennen). De oude systemen waren te stijf om dit tijdens het koken aan te passen.

2. De Oplossing: MegaScale-Data als een Slimme Keukenmanager

MegaScale-Data is een nieuw systeem dat de keuken volledig herontwerpt. Het maakt drie grote veranderingen:

A. De "Gescheiden" Keuken (Disaggregation)

In plaats dat elke kok alles zelf doet (van het uit de voorraadkast halen tot het serveren), introduceert MegaScale-Data twee speciale rollen:

• De Ingredienten-Verzamelaars (Source Loaders): Dit zijn gespecialiseerde helpers die alleen naar de voorraadkast gaan. Ze halen de rauwe ingrediënten (bestanden) en maken ze klaar (bijv. foto's ontgrendelen). Omdat ze gescheiden zijn, hoeft geen enkele kok zijn eigen lijst bij te houden. Ze delen één grote, centrale voorraad.
• De Plaatser (Data Constructor): Dit is de chef die de ingrediënten van de verzamelaars ontvangt, ze in de juiste schalen verdeelt en op de borden zet voor de koks.

De analogie: In plaats dat elke kok zelf naar de markt moet, komen er speciale bezorgers die alles klaarzetten op een centraal eiland. De koks hoeven alleen maar te kiezen wat ze nodig hebben. Dit bespaart enorm veel ruimte in de koelkast (geheugen) en voorkomt dat iedereen dezelfde boodschappenlijst in zijn hand heeft.

B. De Centrale Chef (De "Orkestratie")

Vroeger keek elke kok alleen naar zijn eigen bord. Nu is er een Centrale Chef (de Planner).

• Deze chef ziet precies wat er op de planken ligt. Hij weet: "Ah, deze foto is zwaar om te verwerken, dus ik geef die aan de sterke kok." En: "Deze tekst is kort, die kan de snelle kok doen."
• Hij zorgt ervoor dat alle koks tegelijk klaar zijn. Niemand wacht, niemand loopt vast. Dit heet balancering.
• Hij kan ook tijdens het koken beslissen: "Laten we eerst 80% simpele gerechten doen en dan 20% moeilijke," zonder dat de koks hoeven te stoppen.

C. De Slimme Voorraadkast (Auto-Scaling)

Stel dat je ineens 10 keer meer video's moet verwerken dan normaal.

• In het oude systeem zouden de koks vastlopen.
• In MegaScale-Data ziet de Centrale Chef dit aankomen. Hij roept automatisch meer "Ingredienten-Verzamelaars" bij zich om de video's sneller te halen. Als het weer rustig wordt, sturen ze ze weer weg. Het systeem past zich automatisch aan aan de hoeveelheid werk.

3. Wat levert dit op?

Dankzij deze slimme organisatie hebben ze geweldige resultaten geboekt:

• 4,5x sneller: Het team kookt (traint de AI) 4,5 keer sneller omdat niemand meer hoeft te wachten.
• 13,5x minder ruimte: Ze gebruiken 13,5 keer minder koelkastruimte (computergeheugen), omdat ze niet alles dubbel hoeven op te slaan.

Samenvattend

MegaScale-Data is als het verschil tussen een groep mensen die elk hun eigen boodschappen doen, hun eigen recepten opschrijven en in een kleine keuken proppen, versus een professioneel restaurant met een centrale voorraad, gespecialiseerde bezorgers en een chef die precies regelt wie wat doet zodat niemand vastloopt.

Het zorgt ervoor dat de "grote hersenen" van de AI (de GPU's) altijd vol aan het werk zijn, zonder dat ze hoeven wachten op de "boodschappen" (de data).

Technische samenvatting

Probleemstelling

Het trainen van grote foundation-modellen (LFM's) op basis van data uit meerdere, diverse bronnen (tekst, afbeeldingen, video, etc.) stuit op fundamentele uitdagingen in de huidige data-ladingarchitecturen:

1. Werklastonevenwicht (Workload Imbalance): Door de kwadratische rekencomplexiteit van de attention-operator ($O(l^2)$) in transformer-modellen, leidt een niet-uniforme verdeling van steekproeflengten binnen een batch tot grote verschillen in rekentijd tussen verschillende data-parallelle ranks. Dit veroorzaakt "stragglers" (trage workers) die de gehele training vertragen, vooral bij hybride parallelisme (zoals Pipeline Parallelism en Context Parallelism).
2. Excessief Geheugengebruik: Bestaande frameworks repliceren per-loader bestands-toegangsstatussen (zoals socket-verbindingen, metadata en I/O-buffers) voor elke data-bron. Bij honderden data-bronnen en duizenden parallelle workers leidt dit tot lineaire geheugenoverhead die schaalt met het aantal bronnen en workers.
3. Redundantie bij Hybride Parallelisme: In configuraties zoals Pipeline Parallelism (PP) en Context Parallelism (CP) draait elke GPU vaak een identieke, onafhankelijke dataloader. Dit resulteert in massale redundantie: dezelfde data wordt meerdere keren opgehaald, verwerkt en in het geheugen opgeslagen, wat I/O-bandbreedte en geheugen verspillen.
4. Dynamische Data Mixing: Strategieën zoals curriculum learning vereisen dat de mix van data-bronnen tijdens het trainen dynamisch verandert. Bestaande systemen zijn niet flexibel genoeg om deze veranderingen efficiënt te schalen zonder resource-overprovisioning.

Methodologie: MegaScale-Data

MegaScale-Data is een industriële, gedistribueerde architectuur die de data-preprocessing volledig ontkoppelt van het trainingsproces en deze oplost via drie kerninnovaties:

1. Gedissocieerde Data-Preprocessing (Disaggregated Preprocessing):

   • Het systeem splitst de data-lading op in gespecialiseerde rollen binnen een actor-model:
     • Source Loaders: Specifieke actors per data-bron die verantwoordelijk zijn voor sample-level transformaties (bijv. JPEG-decodering, tokenisatie). Ze elimineren redundantie door bestandsstatussen te isoleren.
     • Data Constructors: Aggregators die output van Source Loaders ontvangen en batch-level operaties uitvoeren (bijv. padding, packing, parallelisme-transformaties). Ze zorgen voor gedeelde toegang tot data binnen parallelle groepen.
   • Dit elimineert redundantie op het niveau van data-bronnen en parallelisme.

2. Gecentraliseerde en Declaratieve Data Plane:

   • DGraph: Een state-tracking dataflow-graaf die de levenscyclus van steekproeven traceert en afhankelijkheden tussen bronnen en verwerkingsstappen in kaart brengt.
   • ClientPlaceTree: Een logische weergave van het apparaat-mesh (GPU-topologie) die het systeem in staat stelt hybride parallelisme (DP, PP, CP, TP) te begrijpen.
   • Primitieven: Gebruikers kunnen complexe data-mixingstrategieën declaratief definiëren (bijv. mix, distribute, balance, broadcast_at) zonder lage-niveau uitvoeringsdetails te hoeven programmeren. Dit stelt het systeem in staat om werklasten dynamisch te balanceren over modules (bijv. tussen visuele encoder en taal-backbone).

3. Multi-level Auto-partitioning en Schaling:

   • Offline Partitioning: Data-bronnen worden automatisch opgesplitst in meerdere Source Loaders gebaseerd op hun verwerkingskosten (heterogeniteit in kosten voor tekst vs. video vs. afbeeldingen).
   • Online Mixture-Driven Scaling: De "Planner" monitort de sampling-weights tijdens het trainen. Als de mix van data verandert, schalen Source Loaders dynamisch (resharding) om de doorvoer te optimaliseren en CPU-gebruik te minimaliseren.

Belangrijkste Bijdragen

• Architectuur: Een gedissocieerd actor-model dat redundantie in data-toegang en geheugengebruik elimineert voor multisource LFM-training.
• Orkestratie: Een declaratieve interface (DGraph + ClientPlaceTree) die hybride parallelisme-bewuste data-scheduling mogelijk maakt met minimale coderinspanning.
• Schalbaarheid: Adaptieve algoritmen voor het dynamisch optimaliseren van CPU-resources op basis van heterogene verwerkingskosten en veranderende data-mixen.
• Betrouwbaarheid: Implementatie van "Shadow Loaders" en differentieel checkpointing voor hoge beschikbaarheid en fouttolerantie zonder onderbreking van de data-stroom.

Resultaten

De prestaties zijn getest op clusters tot 4096 GPU's met visuele-taalmodellen (VLM's) en grote taalmodellen (LLM's):

• Doorvoersnelheid: MegaScale-Data bereikt tot 4,5x verbetering in de end-to-end trainingsdoorvoer vergeleken met state-of-the-art baselines (zoals PyTorch DataLoader, Ray Data, Pecan).
• Geheugenefficiëntie: Er is een reductie van 13,5x in CPU-geheugengebruik, voornamelijk door het elimineren van redundantie in bestandsstatussen en het delen van verwerkte data tussen parallelle ranks.
• Schalbaarheid: Het systeem behoudt hoge efficiëntie bij schaalvergroting (tot 1152+ GPU's), terwijl baselines vaak instorten door communicatiebottlenecks en geheugenoverhead.
• Fouttolerantie: Het gebruik van Shadow Loaders zorgt voor naadloze failover, wat de effectieve trainingsdoorvoer tijdens storingen met 8% verbetert.

Betekenis

MegaScale-Data lost een kritieke knelpunt op in de training van de volgende generatie multimodale foundation-modellen. Waar eerdere oplossingen zich richtten op het uitbesteden van CPU-taken of caching, adresseert MegaScale-Data de fundamentele inefficiënties van multisource data-orkestratie en geheugenredundantie in hybride parallelle omgevingen.

De oplossing maakt het mogelijk om:

1. Enorme, heterogene datasets (honderden bronnen) efficiënt te mixen zonder dat de data-pipeline de GPU's vertraagt.
2. Dynamische trainingsstrategieën (zoals curriculum learning) toe te passen zonder resource-overhead.
3. De complexiteit van hybride parallelisme (PP, CP, TP) te abstraheren voor data-engineers, waardoor ze zich kunnen focussen op data-strategie in plaats van infrastructuur.

Dit paper markeert een verschuiving van "naïeve" dataloading naar een geavanceerde, declaratieve en schaalbare data-plane die essentieel is voor de toekomstige schaalbaarheid van AI-training.