Jenseits des KV-Cache: Training und Serving von i14s Laplace 3B Linearaufmerksamkeitsmodell auf AWS Trainium2 mit nativem PyTorch
aws ·
Dieser Artikel beschreibt, wie i14 sein Laplace-3B-Modell mit linearer Aufmerksamkeit auf AWS Trainium2 trainiert und bereitgestellt hat, wobei natives PyTorch über TorchNeuron Native eingesetzt wurde. Im Gegensatz zu herkömmlichen Transformern, die auf einem wachsenden Schlüssel-Wert-Cache basieren, nutzt Laplace 3Gates Linear Attention, um einen festen rekurrenten Zustand von 29,5 MB pro Sequenz beizubehalten, unabhängig von der Kontextlänge. Die Implementierung erforderte drei wesentliche ingenieurtechnische Anpassungen: das Umschreiben des GLA-Kernels mit Standard-PyTorch-Operatoren, die Optimierung der verteilten Kommunikation für den Muon-Optimizer sowie das statische Abarbeiten der Rückwärtsausführung für den Compiler. Das Training auf einer einzelnen trn2.48xlarge-Instanz erreichte etwa 120.000 Tokens pro Sekunde bei stabiler Verlustkonvergenz. Die Inferenzleistung belief sich auf 11.544 Ausgabetokens pro Sekunde über 1.280 gleichzeitige Sequenzen hinweg, indem der Decode-Schritt kompiliert wurde, um den Overhead durch Operator-Dispatch zu reduzieren. Diese Ergebnisse zeigen, dass benutzerdefinierte Architekturen hohe Effizienz auf AWS-Hardware erreichen können und gleichzeitig eine vorhersehbare Speichernutzung für Anwendungen mit langem Kontext bieten.
aws