Kernekoncepter
FAX ist eine JAX-basierte Bibliothek, die große verteilte und föderierte Berechnungen in Rechenzentren und geräteübergreifenden Anwendungen unterstützt. FAX integriert Bausteine für föderierte Berechnungen als Primitive in JAX, was drei Hauptvorteile bietet: Übersetzung in XLA HLO, Implementierung der Föderierten Automatischen Differenzierung und Interpretation in bestehende Produktions-Föderations-Computesysteme.
Resumé
Die Arbeit präsentiert FAX, eine JAX-basierte Bibliothek, die darauf ausgelegt ist, große verteilte und föderierte Berechnungen in Rechenzentren und geräteübergreifenden Anwendungen zu unterstützen.
FAX nutzt JAX's Sharding-Mechanismen, um native Unterstützung für TPUs und modernste JAX-Laufzeiten wie Pathways zu ermöglichen. FAX baut föderierte Berechnungen aus Primitiven auf, die in JAX eingebettet sind. Dies hat drei Hauptvorteile:
- FAX-Berechnungen können in XLA HLO übersetzt werden.
- FAX bietet eine vollständige Implementierung der Föderierten Automatischen Differenzierung, was den Ausdruck föderierter Berechnungen erheblich vereinfacht.
- FAX-Berechnungen können in bestehende Produktions-Föderations-Computesysteme interpretiert werden.
Die Arbeit zeigt, dass FAX ein leicht programmierbares, leistungsfähiges und skalierbares Framework für föderierte Berechnungen in Rechenzentren bietet. FAX kann effizient und skalierbar für das föderierte Training von Sprachmodellen eingesetzt werden.
Statistik
Die größte Kohorte umfasst 2048 Clients, die insgesamt 3,355 × 10^7 Token verarbeiten und 2,293 × 10^13 FLOPs pro Runde ausführen.
Für das 1B-Modell wird eine Kohorte von 512 Clients verwendet, die 8,389 × 10^6 Token verarbeiten und 1,638 × 10^13 FLOPs pro Runde ausführen.
Für das 8B-Modell wird eine Kohorte von 128 Clients verwendet, die 2,097 × 10^6 Token verarbeiten und 3,277 × 10^13 FLOPs pro Runde ausführen.
Citater
"FAX kann skalierbare verteilte und föderierte Berechnungen in Rechenzentren ermöglichen."
"FAX bietet eine vollständige Implementierung der Föderierten Automatischen Differenzierung, was den Ausdruck föderierter Berechnungen erheblich vereinfacht."
"FAX-Berechnungen können in bestehende Produktions-Föderations-Computesysteme interpretiert werden."