Além do cache KV: treinamento e serviço do modelo de atenção linear Laplace 3B da i14 no AWS Trainium2 com PyTorch nativo
aws ·
Este artigo detalha como o i14 treinou e serviu seu modelo de atenção linear Laplace 3B na AWS Trainium2 utilizando PyTorch nativo por meio do TorchNeuron Native. Diferentemente dos transformers tradicionais, que dependem de um cache chave-valor em crescimento, o Laplace 3B utiliza Atenção Linear com Portas (Gated Linear Attention) para manter um estado recorrente fixo de 29,5 MB por sequência, independentemente do comprimento do contexto. A implementação exigiu três grandes adaptações de engenharia: reescrever o kernel GLA com operadores padrão do PyTorch, otimizar a comunicação distribuída para o otimizador Muon e desenrolar estaticamente o passo de retropropagação para o compilador. O treinamento em uma única instância trn2.48xlarge atingiu aproximadamente 120.000 tokens por segundo, com convergência estável da perda. O desempenho de inferência alcançou 11.544 tokens de saída por segundo em 1.280 sequências concorrentes, ao compilar a etapa de decodificação para reduzir a sobrecarga de despacho de operadores. Esses resultados demonstram que arquiteturas personalizadas podem alcançar alta eficiência em hardware da AWS, oferecendo uso de memória previsível para aplicações de contexto longo.
aws