IBM Research Data Loader verbessert das KI-Modelltraining für die Open-Source-Community

Wie überwinden Sie Engpässe, wenn Sie KI-Modelle mit riesigen Datenmengen trainieren? Auf der diesjährigen PyTorch-Konferenz präsentierte IBM Research einen bahnbrechenden Datenloader für groß angelegtes LLM-Training. Das Tool, jetzt für PyTorch-Nutzer verfügbar, zielt darauf ab, das Training in großem Umfang für eine möglichst breite Zielgruppe zu vereinfachen.

Die Idee für den Hochdurchsatz-Datenlader entstand aus praktischen Problemen der Forschung, die Forscher während des Modelltrainings beobachteten, da ihre Arbeit ein Werkzeug erforderte, das große Datenmengen über mehrere Geräte hinweg verarbeiten konnte – und das alles bei zunehmend effizienten GPUs. Wie IBM Research in seinem Blog zur Veröffentlichung anmerkt, „ist das alles einem Team von Forschern zu verdanken, die einfach die Werkzeuge gebaut haben, die sie brauchten, um eine Aufgabe zu erledigen.“

Davis Wertheimer von IBM Research erläutert einige der Herausforderungen, die bei groß angelegtem Training auftreten können: „Bei groß angelegtem Training gilt so etwas wie die 80/20-Regel. Achtzig Prozent der gesamten veröffentlichten Literatur befassen sich mit algorithmischen Kompromissen zwischen GPU-Speicher und Kommunikation sowie Berechnung. Aber wenn Sie tatsächlich versuchen, etwas zu entwickeln, können Sie sich in 80 % der Fälle auf einen sehr langen Abschnitt all dieser anderen praktischen Probleme verlassen, weil die Pipeline mit der Geschwindigkeit des engsten Engpasses läuft.“

Während das IBM-Team seine Trainingsplattform entwickelte, stieß es weiterhin auf Engpässe. „Je besser wir unsere GPUs nutzen, desto häufiger ist der Datenlader der Flaschenhals“, bemerkt Wertheimer.

Diese Erkenntnis führte zu einem dualen Entwicklungsprozess. „Es gab eine parallele Entwicklung einerseits unserer Trainingsplattform und andererseits unseres Datenladers, um mit den Geschwindigkeitsanforderungen unserer Trainingsplattform Schritt zu halten und einen Engpass zu vermeiden“, erklärt er.

Linsong Chu von IBM Research erläutert die wesentlichen Funktionen des Datenladers:

Zustandsabhängig und überprüfbar: „Immer wenn Sie ein Modell speichern, wird auch der Status des Datenladers gespeichert, und immer wenn Sie von einem Checkpoint aus wiederherstellen, müssen sowohl der Modellstatus als auch der Zustand des Datenladers gleichzeitig wiederhergestellt werden“, sagt Chu.

Automatische Neuskalierung von Checkpoints: Der Data Loader passt sich automatisch an Änderungen des Workload während ausgedehnter Trainingseinheiten an. „Ein Training kann leicht Wochen oder Monate dauern, und es gibt unzählige Gründe, warum man sein Workload zwischendurch anpassen muss“, bemerkt Chu.

Effizientes Datenstreaming: Das System unterstützt Datenstreaming ohne jeglichen Build-Overhead für das Mischen.

Asynchroner verteilter Betrieb: „Wir wollen, dass der Datenlader nicht blockierend ist“, erklärt Chu. „Beim Speichern des Datenladezustands möchten wir, dass die Speicherung in einer Form erfolgt, die keinerlei Kommunikation erfordert.“

Dynamische Datenmischung: Der Datenlader kann sich an unterschiedliche Datenmischungsverhältnisse anpassen, was für sich ändernde Trainingsanforderungen nützlich ist.

Effizientes globales Mischen: Das Tool adressiert Speicherengpässe beim Umgang mit großen Datensätzen und macht das Mischen auch bei Datenwachstum effizient.

PyTorch nativ, modular und umfangreich: Der Datenloader ist auf Anpassungsfähigkeit und Skalierbarkeit ausgelegt und auf zukünftiges Wachstum vorbereitet. „Was ist, wenn wir nächstes Jahr mit 30 Billionen, 50 Billionen oder 100 Billionen Token zu tun haben?“ fragt Chu. „Die Welt verändert sich schnell, deshalb müssen wir den Datenlader so entwickeln, dass er nicht nur heute, sondern auch morgen noch bestehen kann.“

Das IBM Research Team testete seinen Data Loader über mehrere Monate hinweg rigoros und führte Hunderte kleiner und großer Aufträge aus. Es beobachtete stabile und gleichmäßige Codezahlen. Außerdem arbeitet der gesamte Datenlader asynchron und ist nicht blockierend.

„Wir haben viele integrierte PyTorch-Funktionen genutzt, um all dies zu realisieren“, sagt Wertheimer. „Deshalb tragen wir bei und geben etwas zurück.“