Oltre la cache KV: addestramento e servizio del modello di attenzione lineare Laplace 3B di i14 su AWS Trainium2 con PyTorch nativo
aws ·
Questo articolo descrive come i14 abbia addestrato e servito il suo modello di attenzione lineare Laplace 3B su AWS Trainium2 utilizzando PyTorch nativo tramite TorchNeuron Native. A differenza dei transformer tradizionali che si basano su una cache chiave-valore in crescita, Laplace 3B utilizza l'Attenzione Lineare con Gate (Gated Linear Attention) per mantenere uno stato ricorrente fisso di 29,5 MB per sequenza, indipendentemente dalla lunghezza del contesto. L'implementazione ha richiesto tre principali adattamenti ingegneristici: la riscrittura del kernel GLA con operatori standard di PyTorch, l'ottimizzazione della comunicazione distribuita per l'ottimizzatore Muon e lo srotolamento statico del passaggio inverso per il compilatore. L'addestramento su un singolo istanza trn2.48xlarge ha raggiunto circa 120.000 token al secondo con una convergenza stabile della perdita. Le prestazioni di inferenza hanno raggiunto 11.544 token di output al secondo su 1.280 sequenze concorrenti, compilando il passo di decodifica per ridurre l'overhead di dispatch degli operatori. Questi risultati dimostrano che le architetture personalizzate possono raggiungere elevata efficienza sull'hardware AWS offrendo al contempo un utilizzo della memoria prevedibile per applicazioni a lungo contesto.
aws