i14 bringt Linear-Attention-Modell Laplace 3B mit nativem PyTorch auf AWS Trainium2
Das KI-Forschungsunternehmen i14 hat sein Linear-Attention-Modell Laplace 3B mit nativem PyTorch auf AWS Trainium2 trainiert und betrieben, wie aus einem AWS-Blogbeitrag hervorgeht. Der rekurrente Zustand des Modells bleibt unabhängig von der Kontextlänge bei 29,5 MB pro Sequenz, der wachsende KV-Cache klassischer Transformer entfällt.
- Laplace 3B hält einen festen rekurrenten Zustand von 29,5 MB pro Sequenz, unabhängig von der Kontextlänge.
- Auf einer einzelnen trn2.48xlarge-Instanz erreichte das Training rund 120.000 Token pro Sekunde.
- Im Betrieb waren es 11.544 Ausgabetoken pro Sekunde über 1.280 parallele Sequenzen.
- Das Team nutzte das Beta-Backend TorchNeuron Native und behielt dabei das GLA-Design und den Muon-Optimizer bei.
- Zwei all_to_all-Collectives in MuonFSDP2 senkten den Datenverkehr pro Rank um das 24-Fache und steigerten den Durchsatz um den Faktor 1,63.
Warum i14 auf den KV-Cache verzichtet
Die meisten großen Sprachmodelle nutzen Transformer-Architekturen mit Softmax-Attention. Bei der Inferenz wächst ihr Key-Value-Cache (KV-Cache) mit dem Kontext des Gesprächs oder Dokuments. Laut dem AWS-Beitrag kann das die Parallelität verringern, die Betriebskosten erhöhen und Agenten-Anwendungen dazu zwingen, nützlichen Verlauf zu verwerfen.
i14 hat Laplace 3B als Linear-Attention-Alternative entworfen. Das Modell nutzt Gated Linear Attention (GLA), um den bisherigen Kontext in einem Matrixzustand fester Größe zu komprimieren. Dadurch sinkt der Aufwand der Attention im Training von O(T^2) auf O(T), und der Inferenzzustand bleibt bei 29,5 MB pro Sequenz. i14 beschreibt sich als Unternehmen, das Modellarchitekturen, Komponenten und Laufzeitumgebungen entwickelt und dabei auf KI, kognitive Computational Science sowie Forschung dazu zurückgreift, wie biologische Systeme im Lauf der Zeit lernen.
Der Beitrag benennt den Zielkonflikt klar: Linear Attention komprimiert den Verlauf, statt jeden Key und Value zu speichern. Nutzer sollten daher die Retrieval-Qualität für ihre eigenen Workloads prüfen. Der feste Zustand kann langlaufenden Agenten, der Analyse großer Dokumente und Anwendungen mit persistentem Gedächtnis und planbaren Betriebsanforderungen zugutekommen.
Das Modell und seine Ergebnisse auf Trainium2
Laplace 3B ist ein Decoder mit 24 Schichten, in dem jede Schicht die Softmax-Attention durch Simple GLA ersetzt. Das Modell hat insgesamt 3,028 Mrd. Parameter, davon 2,718 Mrd. mit dem Muon-Optimizer und 309,6 Mio. mit AdamW trainiert. Es verwendet eine Hidden-Dimension von 3.072, 16 Heads, eine Sequenzlänge von 2.048 und ein Vokabular von 50.304. Trainiert wurde es mit 100 Mrd. Token aus FineWeb-Edu bei einer Batchgröße von 524.288 Token pro Schritt.
Jeder Head hält einen Matrixzustand von 96 x 192, der als abklingendes Protokoll vergangener Key-Value-Paare dient. Die Zerfallsrate pro Head wird bei der Initialisierung festgelegt und folgt einem geometrischen Schema, sodass die Heads unterschiedlich lange Gedächtnisspannen abdecken. Head 0 hat eine Halbwertszeit von einem Token, Head 15 von 2.048 Token.
Auf einer einzelnen trn2.48xlarge-Instanz erreichte Laplace 3B im Training laut Beitrag rund 120.000 Token pro Sekunde. Im Betrieb lieferte das Modell 11.544 Ausgabetoken pro Sekunde über 1.280 parallele Sequenzen. Eine trn2.48xlarge kombiniert 16 Trainium2-Chips, 64 logische NeuronCores, Speicher mit hoher Bandbreite und NeuronLink-v3-Konnektivität.
Drei Anpassungen für Trainium2
Die ursprüngliche Implementierung von Laplace 3B stützte sich auf GPU-optimierte Triton-Kernel. Das Team des AWS Generative AI Innovation Center (GenAIIC) schrieb zunächst den GLA-Kernel mit Standard-PyTorch-Operatoren neu. Nach einem Vergleich der Ansätze entschied es sich für die Beta-Version von TorchNeuron Native, dem nativen PyTorch-Backend für AWS Trainium, weil diese besser abschnitt. So blieben das GLA-Design und der Muon-Optimizer erhalten.
- GLA-Kernel in nativem PyTorch: Der GPU-spezifische Triton-Kernel fused_chunk_simple_gla wurde durch Standard-PyTorch-Operatoren ersetzt, mit 16 Chunks zu je 128 Token, passend zu den 128 Partitionen des On-Chip-Puffers der NeuronCore.
- Muon-Kommunikation: Zwei all_to_all-Collectives in MuonFSDP2 senkten den Datenverkehr pro Rank um das 24-Fache und steigerten den Durchsatz um den Faktor 1,63.
- Kompilierter Backward-Pass: Das Team rollte die 16 Chunk-Schritte der GLA-Rekurrenz statisch aus, damit sie in die Grenzen des Compilers für Graphgröße und On-Chip-Speicher passen.
Drei Eigenschaften von Trainium2 prägten die Arbeit: softwareverwalteter On-Chip-Speicher, eine schnelle Chip-zu-Chip-Verbindung und ein statischer Compiler, der die GLA-Rekurrenz als festen Graphen nachzeichnet. Jede NeuronCore verfügt über einen State Buffer von 28 MiB mit 128 Partitionen, und der Compiler plant jede Übertragung im Voraus.
Was das für Entwickler bedeutet
Der Beitrag zeigt, dass eine eigene Linear-Attention-Architektur auf Trainium2 laufen kann, ohne natives PyTorch zu verlassen. Er räumt aber ein, dass solche Modelle architekturspezifische Kernel und Compiler-Tuning erfordern, anders als Transformer mit breit optimierten Attention-Kerneln. Für Teams, die Anwendungen mit langem Kontext oder Agenten-Systeme bauen, bietet der feste Zustand pro Sequenz einen planbaren Speicherbedarf. Ob der komprimierte Zustand genügend Informationen behält, muss weiterhin für jeden Workload getestet werden.
