Más allá de la caché KV: entrenamiento y servicio del modelo de atención lineal Laplace 3B de i14 en AWS Trainium2 con PyTorch nativo
aws ·
Este artículo detalla cómo i14 entrenó y sirvió su modelo de atención lineal Laplace 3B en AWS Trainium2 utilizando PyTorch nativo a través de TorchNeuron Native. A diferencia de los transformadores tradicionales que dependen de una caché de clave-valor en crecimiento, Laplace 3B utiliza Atención Lineal con Puertas (GLA) para mantener un estado recurrente fijo de 29,5 MB por secuencia, independientemente de la longitud del contexto. La implementación requirió tres adaptaciones de ingeniería principales: reescribir el kernel GLA con operadores estándar de PyTorch, optimizar la comunicación distribuida para el optimizador Muon y desenrollar estáticamente el paso hacia atrás para el compilador. El entrenamiento en una sola instancia trn2.48xlarge alcanzó aproximadamente 120.000 tokens por segundo con una convergencia estable de la pérdida. El rendimiento de inferencia llegó a 11.544 tokens de salida por segundo en 1.280 secuencias concurrentes mediante la compilación del paso de decodificación para reducir la sobrecarga de despacho de operadores. Estos resultados demuestran que las arquitecturas personalizadas pueden lograr alta eficiencia en el hardware de AWS mientras ofrecen un uso de memoria predecible para aplicaciones de contexto largo.
aws