Codebase Scaling als Open Source für nicht-kommerzielle Forschung

Codebase Scaling

Wir veröffentlichen unsere Modell-Trainings-Codebasis Scaling und stellen sie öffentlich unter der Open Aleph License bereit, die ausdrücklich nicht-kommerzielle Forschung und Bildungszwecke erlaubt. Begleitend zum öffentlichen GitHub-Repository beleuchtet dieser Blogpost Design- und Engineering-Entscheidungen, die wir bei der Entwicklung unserer Trainings-Codebasis getroffen haben. Scaling wurde zur Entwicklung unserer zeitgleich veröffentlichten neuen Modelle Pharia-1-LLM-control und Pharia-1-LLM-control-aligned eingesetzt.

Einführung

In den letzten Jahren ist das Rechenbudget für das Training modernster Deep-Learning-Modelle – insbesondere LLMs – drastisch gestiegen. Grund dafür ist das Wachstum sowohl der Modell- als auch der Datengröße: Modelle erreichen Hunderte Milliarden Parameter, Trainingssets wachsen auf mehrere Billionen Tokens. Trainingscode so zu engineeren, dass er auf diese rechenintensiven Szenarien skaliert, ist keine triviale Aufgabe. Die Speicher- und Compute-Anforderungen des LLM-Trainings verlangen Parallelisierung über viele GPUs hinweg, damit Training überhaupt machbar wird.

Parallelisierung bringt viel Komplexität in eine Standard-Trainingsschleife. Der resultierende Code ist bekanntermaßen schwer zu warten und zu erweitern. Viele Forschungslabs und Unternehmen veröffentlichen zwar Modellgewichte, ihren Trainingscode halten sie jedoch oft proprietär und stellen ihn nicht der Forschungsgemeinschaft zur Verfügung. Andererseits sind öffentliche Bibliotheken für Large-Scale-Training häufig schwer zu durchdringen oder anzupassen.

Wir freuen uns daher, Scaling zu veröffentlichen – Aleph Alphas parallele Trainings-Codebasis.

Scaling enthält die folgenden Features:

• Frameworks: Wir liefern allgemeine Bausteine für verteiltes Modelltraining (scaling.core) sowie eine vollwertige Suite für Large-Scale-Transformer-Training (scaling.transformer).
• Effizientes Training: Die Codebasis unterstützt Daten-, Pipeline- und Modell-(Tensor-)Parallelismus, Mixed-Precision-Training und moderne Performance-Optimierungen wie Optimizer-State-Sharding und Activation-Checkpointing.
• Codequalität: Wir setzen auf strenges Typing, Pydantic-Klassen und umfangreiche automatisierte Tests. Das sorgt für sauberen Code, sichere Entwicklung und weniger Fehlerpotenzial.

Scaling wird unter der Open Aleph License veröffentlicht und kann für nicht-kommerzielle Forschung und Bildungszwecke genutzt werden. Wir hoffen, dass die Bereitstellung unserer Codebasis für die Machine-Learning-Community unabhängige Forschung im Bereich des Large-Scale-Trainings unterstützt.

Scaling ist eine vollständige Eigenentwicklung, basierend auf unserer Forschung und Arbeit – nicht geforkt von bestehenden Open-Source-Referenzen. Scaling wird auch in Zukunft aktiv weiterentwickelt und verbessert, damit die Trainingsperformance modernen Standards entspricht. Außerdem werden wir Innovationen aus unserem Forschungsteam einfließen lassen, um unsere Arbeit transparent und reproduzierbar zu machen.

Im Folgenden beschreiben wir die beiden Module scaling.core und scaling.transformer etwas detaillierter. Scaling setzt – verglichen mit größeren Codebases wie Megatron-LM oder DeepSpeed – auf Prägnanz und Erweiterbarkeit. Diese Designentscheidung reduziert Komplexität, erleichtert Wartung und Optimierung und erhält gleichzeitig Performance und Skalierbarkeit.

Eine allgemeine Library für verteiltes Training: scaling.core

scaling.core ist der weitgehend modell-agnostische Maschinenraum der Scaling-Library, gebaut auf PyTorch, der das parallele Training von Machine-Learning-Modellen im großen Maßstab ermöglicht. Es unterstützt die folgenden Parallelisierungstechniken:

1. Datenparallelismus. Ein Batch Trainingsdaten wird auf mehrere Modellkopien aufgeteilt, die unabhängige Forward-/Backward-Pässe ausführen; anschließend werden die resultierenden Gewichtsgradienten aggregiert. So kann mit kleineren Batchgrößen pro GPU trainiert werden, was Activation- und Gradient-Speicher reduziert.
2. Pipeline-Parallelismus. Die Schichten eines Modells werden in sequentielle „Pipeline-Stufen“ gruppiert, wobei die Berechnung jeder Stufe auf einem anderen Gerät stattfindet. Aktivierungen und Gradienten werden anschließend an die nachfolgenden bzw. vorhergehenden Pipeline-Stufen kommuniziert. Jede Stufe muss nur einen Teil des Modells materialisieren, was den Speicherverbrauch drastisch senkt.
3. Modell-(Tensor-)Parallelismus. Große Matrixmultiplikationen machen einen Großteil der Compute-Anforderungen moderner Deep Neural Networks aus. Diese Kernoperationen können auf Geräte aufgeteilt werden, wobei jedes Gerät einen Teil der Gewichtsmatrix besitzt und die entsprechenden Teilberechnungen durchführt.

Alle drei Techniken können gleichzeitig eingesetzt werden – das ergibt 3D-Parallelismus. Darüber hinaus implementiert Scaling weitere wichtige Performance-Optimierungen, etwa Mixed-Precision-Training, geshardte Optimizer-Zustände (bekannt als ZeRO) sowie Activation-Checkpointing.

Wir haben Wert darauf gelegt, dass Nutzer:innen ihre eigenen 3D-Trainings-Anwendungen bauen können, ohne tiefgehendes Wissen über die zugrunde liegende Implementierung zu benötigen. Ein Beispiel: Damit eine eigene Modellarchitektur in Scaling läuft, musst du sie nur als Instanz unserer ParallelModule-Klasse implementieren. ParallelModule funktioniert sehr ähnlich wie torch.nn.Sequential und nimmt eine sequentielle Modellstruktur an, die sich in Layer zerlegen lässt. Die Layer selbst haben keine besonderen Anforderungen und sind ganz normale PyTorch-Module. In Kombination mit unseren Trainer- und Optimizer-Abstraktionen, die meist out-of-the-box nutzbar sind, ermöglicht ParallelModule dann eine daten- und pipeline-parallele Trainingsschleife. Um Modellparallelismus zu nutzen, ist ein zusätzlicher Schritt nötig: Ersetze alle Instanzen von torch.nn.Linear in deinem Modell durch die von uns bereitgestellten Klassen ColumnParallelLinear bzw. RowParallelLinear. Bei der allgemeinen Komplexität verteilten Trainings gibt es natürlich einige Feinheiten, die wir hier nicht abdecken. Wir laden dich ein, unser Repository und die umfangreichen Ressourcen darin zu erkunden, um loszulegen.

Eine Suite für LLM-Training: scaling.transformer

Da Transformer die Standardarchitektur für Sprachmodellierung und der Hauptanwendungsfall für Large-Scale-Distributed-Training sind, bieten wir in scaling.transformer eine Trainings-Suite für Large Language Models.

Mit den Bausteinen aus scaling.core implementieren wir eine modernste Transformer-Architektur und Trainingsschleife. Die Modellarchitektur lässt sich über eine prägnante, aber flexible Config-Klasse konfigurieren, mit der sich eine Vielzahl von Transformer-Varianten realisieren lässt. Einige Architekturoptionen, die wir unterstützen:

• Multi-Query- und Grouped-Query-Attention,
• verschiedene MLP-Typen (z.B. SwiGLU),
• Rotary-Positional-Embeddings,
• parameter-effiziente Fine-Tuning-Methoden,

und viele weitere Features moderner LLM-Architekturen. Zusätzlich zur Architektur stellen wir LLM-spezifische Funktionalität für Datenladen und Vorverarbeitung bereit.

Natürlich unterstützt scaling.transformer vollständig verteiltes Training mit allen Parallelisierungs- und Optimierungsoptionen, die die scaling.core-Komponenten bieten.

Ein Transformer-Training oder Fine-Tuning lässt sich über ein einheitliches Entry-Point-Skript train.py starten, das eine Konfigurationsdatei im YAML-Format entgegennimmt – ein einfacher und nachvollziehbarer Workflow für Experimente.

Schließlich liefert scaling.transformer eine schlanke Inferenz-Funktionalität, mit der trainierte Modelle getestet und evaluiert werden können. Sie unterstützt KV-Caching für bessere Inferenz-Performance.