MoE-Training für biologische Grundmodelle optimieren
NVIDIA hat mit dem BioNeMo MoE Recipe und der Transformer Engine (TE) neue Optimierungsverfahren für das Training biologischer Foundational Models vorgestellt. Mit der wachsenden Parameteranzahl und Sequenzlänge von Genommodellen stießen dichte Transformer-Architekturen an ihre Grenzen. Der zunehmende Rechenaufwand für Training und Inferenz machte die Mixture-of-Experts-Architektur zur effektiven Alternative. Statt jedes Token durch alle Schichten zu leiten, aktiviert MoE pro Token nur eine kleine Auswahl spezialisierter Subnetzwerke. Diese Sparsity bringt Skaleneffizienz mit sich, erfordert jedoch präzise Implementierungen, um Engpässe bei GPU-Auslastung, Routing-Kommunikation und Speicherbedarf zu vermeiden. NVIDIA adressiert diese Herausforderungen mit gezielten Kernel-Optimierungen in der Transformer Engine. Ein zentrales Problem naiver MoE-Implementierungen ist die fragmentierte Expert-Exekution. Python-basierte Schleifen lösen für jeden Experten einzelne Kernel auf, was zu erheblichem Scheduling-Overhead führt. Der TE löst dies via GroupedLinear, das mehrere lineare Transformationen in einem einzigen Aufruf zusammenfasst. Durch die Übergabe von Token-Anzahlen pro Experten werden die Expert-Matrizen gesammelt und über die optimierte GEMM-Pipeline der TE ausgeführt. Dies reduziert die Kernel-Startkosten drastisch. Zusätzlich stellen die wachsende Modellspeicherkapazität zusammen mit langen biologischen Sequenzen enorme Anforderungen an den Aktivierungsspeicher. Während BF16 für Gewichte und Aktivierungen 16 Bit pro Wert benötigt, unterstützt die BioNeMo-Lösung FP8 und MXFP8. MXFP8 reduziert den Speicherbedarf auf 8 Bit pro Wert und weist jedem Block von 32 Werten einen individuellen Skalierungsfaktor zu, um numerische Genauigkeit und dynamischen Bereich zu erhalten. Auf NVIDIA Blackwell-GPUs werden MXFP8-Berechnungen durch spezialisierte Tensor-Core-Anweisungen hardwarebeschleunigt. Der dritte Engpass betrifft die Quantisierung und Dequantisierung zwischen 16-Bit-Mastergewichten und 8-Bit-Inferenzwegen. Naive Implementierungen führen diese Konvertierungen als separate Operationen aus. Die Transformer Engine eliminiert diesen Overhead durch einen GroupedMLP-Fusionskernel. Mithilfe der TE-Sequential-API werden Gate-Up-Projektionen, ScaledSwiGLU-Aktivierungen und Routing-Skalierungsvorgänge zu einem einzigen vorwärts- und rückwärtsgerichteten Kernel verschmolzen. Zwischenergebnisse werden nicht explizit im Speicher gehalten, was Framework-Overhead minimiert und den Durchsatz maximiert. Leistungstests auf acht NVIDIA B200 Tensor Core GPUs belegen die Wirksamkeit des Ansatzes. Im Vergleich zum Hugging-Face-Baseline-Transformers-Modell steigerte der BioNeMo MoE Recipe die Trainingsdurchsatzrate um bis zu 221 Prozent. Die Optimierungen ermöglichen skalierbare Expert-Parallelism-Konfigurationen bei gleichzeitiger Anpassung der Datenparallelität an die verfügbare Hardware. Entwickler können das Framework mit einer Zweier-GPU-Konfiguration validieren und anschließend auf Mixtral-8x7B-ähnliche Architekturen mit EP=8 und MXFP8-Präzision skalieren. Die vollständigen Launch-Skripte, Checkpoint-Logik und Benchmark-Befehle sind im BioNeMo Recipes-Repository sowie in der Transformer Engine Dokumentation verfügbar. Die Initiative etabliert einen offenen Standard für ressourceneffizientes Training hochkomplexer genomischer Modelle und untermauert den Einsatz hocheffizienter GPU-Workflows im Biotech-Sektor.
