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

Das Paper stellt MegaScale-Data vor, eine industrielle verteilte Datenlade-Architektur für das Training multiquelliger großer Basismodelle, die durch disaggregierte Vorverarbeitung, eine zentrale Orchestrierungsebene und ein mehrstufiges Auto-Partitionierungsverfahren Workload-Ungleichgewichte und Speicherverbrauch reduziert und damit die Trainingsdurchsatzleistung um das 4,5-fache sowie den CPU-Speicherbedarf um das 13,5-fache verbessert.

Das Wesentliche

Stellen Sie sich vor, Sie sind der Chefkoch in einer riesigen, futuristischen Küche, die darauf spezialisiert ist, den ultimativen KI-Kochbuch-Algorithmus zu trainieren. Diese KI (das "Large Foundation Model") soll alles verstehen: Texte, Bilder, Videos und Audio.

Das Problem ist: Die Zutaten kommen aus tausenden verschiedenen Quellen. Manche sind frische Tomaten (Text), andere sind riesige, schwer zu schneidende Kürbisse (hochauflösende Bilder) und wieder andere sind winzige Gewürzpartikel (Audio).

Das alte Problem: Der chaotische Küchenchef
Früher hatte jede Station in der Küche (jeder "GPU-Server") ihren eigenen kleinen Koch, der sich selbstständig um die Zutaten kümmerte.

1. Das Chaos der Ungleichheit: Ein Koch bekam nur kleine, leicht zu schneidende Tomaten, während ein anderer einen riesigen, harten Kürbis bekam. Der mit dem Kürbis brauchte 10 Minuten, der mit den Tomaten nur 10 Sekunden. Der ganze Kochprozess musste warten, bis der langsamste fertig war. Das war extrem ineffizient.
2. Der Platzmangel: Jeder Koch hatte sein eigenes Regal, in dem er alle möglichen Zutatenlisten und Werkzeuge für alle Quellen lagerte. Wenn es 1000 Quellen gab, hatte jeder Koch 1000 Listen. Das Regal (der Arbeitsspeicher) war voll, obwohl jeder Koch eigentlich nur für seine eigene Station zuständig war.
3. Die Doppelarbeit: Wenn eine neue Rezeptur (ein "Batch") vorbereitet wurde, kopierte jeder Koch das gleiche Rezept und die gleichen Zutatenlisten. Statt einer Liste gab es hundert Kopien.

Die Lösung: MegaScale-Data (Der neue Küchenplan)
Das Paper beschreibt eine revolutionäre neue Küchenorganisation namens MegaScale-Data. Hier ist, wie es funktioniert, einfach erklärt:

1. Die Entkopplung: Spezialisten statt Generalisten

Statt dass jeder Koch alles selbst macht, teilen wir die Arbeit auf:

• Die "Zutaten-Entdecker" (Source Loaders): Das sind spezialisierte Helfer, die nur für eine bestimmte Quelle zuständig sind. Einer kümmert sich nur um Tomaten, einer nur um Kürbisse. Sie holen die rohen Zutaten und machen die erste grobe Bearbeitung (z. B. Schneiden). Sie müssen nicht wissen, wer am Ende isst.
• Die "Teller-Verteiler" (Data Constructors): Das sind die Köche an der Ausgabeküche. Sie nehmen die vorbereiteten Zutaten von den Entdeckern und stellen die Teller (die Datenpakete für die KI) zusammen.
• Der Vorteil: Wenn 1000 Quellen da sind, brauchen wir nicht 1000 Regale pro Koch. Wir haben nur 1000 Regale für die Entdecker, und die Verteiler teilen sich die Arbeit. Das spart enorm viel Platz im Regal (Speicher).

2. Der intelligente Tischaufseher (Der "Planner")

Früher saß jeder Koch am Tisch und entschied selbst, was er isst. Jetzt gibt es einen zentralen Tischaufseher (den "Planner").

• Die Waage (Orchestration): Der Aufseher sieht, wer wie viel Zeit braucht. Er weiß: "Der Kürbis-Koch braucht lange." Also gibt er ihm weniger Kürbisse, aber dafür mehr Tomaten, damit alle gleichzeitig fertig werden. Er balanciert die Arbeit so aus, dass kein Koch warten muss.
• Das dynamische Menü (Curriculum Learning): Manchmal will man am Anfang nur einfache Gerichte (einfache Daten) und später komplexe. Der Aufseher kann das Menü live ändern, ohne dass die Küche stillstehen muss. Er sagt einfach: "Heute mehr Tomaten, morgen mehr Kürbisse."

3. Die flexible Küche (Skalierbarkeit)

Stellen Sie sich vor, die Küche wird plötzlich doppelt so groß (mehr GPUs).

• Alte Methode: Jeder neue Koch müsste sich sofort sein eigenes riesiges Regal mit allen Listen kaufen. Das wäre teuer und langsam.
• MegaScale-Methode: Der Aufseher teilt die Arbeit neu auf. Wenn ein neuer Koch kommt, bekommt er einfach einen Teil der Liste von einem alten Koch. Niemand muss alles neu lernen. Die Küche wächst mühelos mit, egal ob 10 oder 4000 Köche da sind.

4. Der Sicherheitsplan (Fehlertoleranz)

Was passiert, wenn ein Koch krank wird oder ein Regal umfällt?

• Der Schatten-Koch (Shadow Loader): Es gibt immer einen "Zwilling" eines Kochs, der genau das Gleiche tut, aber im Hintergrund wartet. Wenn der Hauptkoch ausfällt, springt der Schatten sofort ein, ohne dass die Küche einen Moment lang anhalten muss. Es gibt keine Unterbrechung beim Servieren.

Das Ergebnis

Durch diese neue Organisation erreichen sie zwei Wunder:

1. Geschwindigkeit: Die Küche ist bis zu 4,5-mal schneller, weil niemand mehr warten muss und die Arbeit perfekt verteilt ist.
2. Platzersparnis: Der Platzbedarf für die Regale (Arbeitsspeicher) sinkt um das 13,5-fache, weil niemand mehr unnötige Kopien von Listen herumträgt.

Zusammenfassend:
MegaScale-Data verwandelt eine chaotische, ineffiziente Küche, in der jeder Koch alles selbst macht und wartet, in eine hochorganisierte Fabrik mit spezialisierten Stationen, einem klugen Manager und einem Sicherheitsnetz. So kann die KI lernen, schneller und mit weniger Ressourcen als je zuvor.

Technische Zusammenfassung

1. Problemstellung

Das Training von großen Fundamentmodellen (Large Foundation Models, LFMs), insbesondere multimodalen Modellen wie Vision-Language-Modellen (VLMs), steht vor zwei fundamentalen Herausforderungen im Bereich des Daten-Loading und der Datenorchestrierung, wenn Daten aus vielen verschiedenen Quellen stammen:

• Ungleichgewichtige Arbeitslast (Workload Imbalance): Aufgrund der quadratischen Komplexität des Attention-Mechanismus ($O(l^2)$) führt eine nicht-uniforme Verteilung von Sequenzlängen in den Daten zu erheblichen Ungleichgewichten zwischen den Daten-Parallel-Ranks. Ein Batch mit gemischten kurzen und langen Sequenzen verursacht, dass einige GPUs (Ranks) deutlich mehr Rechenzeit benötigen als andere, was zu „Stragglern" (langsamen Einheiten) und ineffizienter GPU-Nutzung führt. Dies wird durch die Heterogenität verschiedener Datenquellen (Text, Bild, Video) und deren unterschiedliche Verarbeitungskosten (z. B. Bilddekodierung vs. Text-Tokenisierung) weiter verschärft.
• Skalierbarkeits- und Speicherprobleme: Herkömmliche DataLoaders replizieren den Zugriff auf Dateistrukturen (Sockets, Metadaten, I/O-Puffer) für jeden Parallelismus-Rank und jede Datenquelle. Bei hunderten von Datenquellen und hybriden Parallelismus-Strategien (Data Parallelism, Pipeline Parallelism, Context Parallelism) führt dies zu einem linearen Anstieg des Speicherverbrauchs (CPU-RAM). Zudem wird bei hybriden Parallelismen (z. B. Pipeline Parallelism) oft dieselbe Datenmenge redundant von mehreren GPUs geladen und verarbeitet, was I/O-Bandbreite und Speicher verschwendet.
• Dynamische Datenmischung: Techniken wie Curriculum Learning erfordern eine dynamische Anpassung der Mischverhältnisse der Datenquellen während des Trainings. Bestehende Systeme bieten hierfür keine ausreichende API-Unterstützung für eine globale Koordination über heterogene Quellen hinweg.

2. Methodik: MegaScale-Data

MegaScale-Data ist eine industrietaugliche, verteilte Architektur für das Laden und Vorverarbeiten von Daten, die auf drei Kerninnovationen basiert:

A. Entkoppelte Vorverarbeitung (Disaggregated Preprocessing)

Statt den DataLoader direkt mit dem Trainingsprozess zu koppeln, wird das System in spezialisierte Rollen aufgeteilt:

• Source Loaders: Diese sind dedizierte Akteure (Actors) für einzelne Datenquellen. Sie übernehmen die quellenspezifische Vorverarbeitung auf Sample-Ebene (z. B. JPEG-Decodierung, Tokenisierung). Dies eliminiert redundante Datei-Zugriffszustände und Metadaten über die verschiedenen Ranks hinweg.
• Data Constructors: Diese Aggregatoren empfangen die vorverarbeiteten Samples von den Source Loaders. Sie führen Batch-Level-Operationen durch (z. B. Padding, Packing, Tensor-Transformationen) und verteilen die Daten basierend auf der Parallelismus-Konfiguration an die Trainer. Dies ermöglicht das Teilen von Daten zwischen Ranks (z. B. in Context Parallelism), ohne dass jeder Rank die Daten neu laden muss.

B. Zentrale deklarative Daten-Orchestrierung

Ein zentraler Planner steuert den gesamten Datenfluss über eine deklarative Schnittstelle:

• DGraph (Dataflow Graph): Ein zustandsbehafteter Graph, der den Lebenszyklus von Trainingsstichproben (Samples) verfolgt und Abhängigkeiten zwischen Quellen und Transformationsstufen modelliert.
• ClientPlaceTree: Ein logisches Topologie-Modell des GPU-Clusters, das die Parallelismus-Strategien (DP, PP, TP, CP) abbildet.
• Primitiven: Das System bietet hochlevelige APIs wie mix() (für dynamische Mischverhältnisse), distribute() (für Partitionierung über Parallelismus-Achsen), cost() (für Schätzung von Rechenkosten) und balance() (für Lastverteilung). Dies ermöglicht eine Lastverteilung, die die Heterogenität der Datenquellen und Parallelismus-Strukturen berücksichtigt, bevor die Daten in das Modell fließen.

C. Multi-Level Auto-Partitioning und Skalierung

• Source Auto-Partitioning: Das System partitioniert Datenquellen automatisch in mehrere Source-Loader-Akteure, basierend auf den unterschiedlichen Vorverarbeitungskosten der Quellen.
• Mixture-Driven Scaling: Der Planner überwacht die Mischverhältnisse der Datenquellen in Echtzeit. Wenn sich die Sampling-Gewichte ändern (z. B. durch Curriculum Learning), skaliert der AutoScaler dynamisch die Anzahl der Source-Loader und Worker, um Engpässe zu vermeiden und Ressourcen effizient zu nutzen.

D. Fehlertoleranz und Elastizität

Das System nutzt „Shadow Loaders" (Hot-Standby-Instanzen) und differenzielle Checkpointing-Mechanismen, um bei Ausfällen von Source Loaders eine sofortige Wiederherstellung ohne Datenverlust zu gewährleisten. Zudem unterstützt es elastisches Resharding bei Änderungen der Parallelismus-Konfiguration des Trainings-Clusters.

3. Wichtige Beiträge

1. Entkoppelte Architektur: Ein verteiltes Actor-Modell, das redundante Datenzugriffe und Speicherverbrauch bei multisource LFM-Training eliminiert.
2. Deklarative Orchestrierung: Einführung von DGraph und ClientPlaceTree, die eine hybride Parallelismus-orientierte Datenplanung mit minimalem Programmieraufwand ermöglichen.
3. Adaptive Skalierung: Algorithmen zur dynamischen Optimierung der CPU-Nutzung basierend auf heterogenen Vorverarbeitungskosten und sich ändernden Datenmischverhältnissen.
4. Industrie-Implementierung: Ein vollständiges System, das auf Ray aufbaut und in Produktionsumgebungen (ByteDance) eingesetzt wird.

4. Ergebnisse

Die Evaluation wurde auf Clustern mit bis zu 4096 GPUs durchgeführt (u. a. mit VLMs wie Llama + ViT):

• Durchsatzsteigerung: MegaScale-Data erreicht eine 4,5-fache Verbesserung des End-to-End-Trainingsdurchsatzes im Vergleich zu herkömmlichen Data-Parallel-Baselines.
• Speichereffizienz: Der CPU-Speicherbedarf (RAM) wurde um den Faktor 13,5 reduziert. Dies wird durch die Eliminierung redundanter Dateizugriffszustände und das Teilen von Daten zwischen Parallelismus-Ranks erreicht.
• Skalierbarkeit: Im Gegensatz zu Baselines, die bei 4000+ GPUs aufgrund von Kommunikationsengpässen zusammenbrechen, skaliert MegaScale-Data stabil.
• Lastverteilung: Durch die aktive Balance-Strategie werden Workload-Ungleichgewichte (Stragglers) signifikant reduziert, was zu einer gleichmäßigeren GPU-Auslastung führt.
• Trainingskonvergenz: Die Balancing-Strategie beeinflusst die Trainingsverlust-Konvergenz nicht negativ; die Ergebnisse bleiben vergleichbar mit nicht-balanceierten Baselines.

5. Bedeutung

MegaScale-Data adressiert eine kritische Lücke im Training moderner KI-Modelle. Während bisherige Fortschritte sich stark auf die Optimierung der Modell-Parallelität (TP, PP, DP) konzentrierten, vernachlässigten sie oft die Datenpipeline als Engpass.

Die Arbeit zeigt, dass für das Training von Multimodal-Modellen mit tausenden von Datenquellen eine zentrale, entkoppelte Datenorchestrierung essenziell ist. Sie beweist, dass durch die Trennung von Datenzugriff, Vorverarbeitung und Batch-Assembly sowie durch intelligente Lastverteilung nicht nur massive Speicherressourcen gespart werden, sondern auch die Trainingsgeschwindigkeit drastisch erhöht werden kann. Dies ist ein entscheidender Schritt hin zu effizienteren und kostengünstigeren Trainingsprozessen für die nächste Generation von Foundation Models.