Au-delà du cache KV : entraînement et déploiement du modèle d'attention linéaire Laplace 3B de i14 sur AWS Trainium2 avec PyTorch natif
aws ·
Cet article détaille comment i14 a entraîné et déployé son modèle d'attention linéaire Laplace 3B sur AWS Trainium2 en utilisant PyTorch natif via TorchNeuron Native. Contrairement aux transformateurs traditionnels qui reposent sur un cache clé-valeur croissant, Laplace 3B utilise une attention linéaire à portes (Gated Linear Attention) pour maintenir un état récurrent fixe de 29,5 Mo par séquence, indépendamment de la longueur du contexte. L'implémentation a nécessité trois adaptations techniques majeures : la réécriture du noyau GLA avec des opérateurs PyTorch standard, l'optimisation de la communication distribuée pour l'optimiseur Muon, et le déroulement statique du passage arrière pour le compilateur. L'entraînement sur une seule instance trn2.48xlarge a atteint environ 120 000 tokens par seconde avec une convergence stable de la perte. Les performances d'inférence ont atteint 11 544 tokens de sortie par seconde sur 1 280 séquences concurrentes grâce à la compilation de l'étape de décodage afin de réduire les frais généraux de dispatch des opérateurs. Ces résultats démontrent que les architectures personnalisées peuvent atteindre une grande efficacité sur le matériel AWS tout en offrant une utilisation mémoire prévisible pour les applications à long contexte.
aws