Entraînement d'un modèle d'attention linéaire sur AWS Trainium2 avec PyTorch natif

Les chercheurs de i14 et du Centre d'innovation en IA générative d'AWS ont réussi à entraîner et à servir le modèle Laplace 3B à l'aide des accélérateurs AWS Trainium2. Publiés en octobre 2026, ces travaux démontrent comment les architectures d'attention linéaire peuvent être optimisées pour un matériel spécialisé sans dépendre de noyaux GPU personnalisés.

Cette collaboration met en évidence une évolution par rapport à l'attention softmax traditionnelle des transformeurs, qui souffre d'une mise à l'échelle quadratique lors de l'entraînement et de coûts mémoire croissants lors de l'inférence. En tirant parti de la prise en charge native de PyTorch via TorchNeuron Native, l'équipe a atteint des métriques de performance significatives tout en maintenant une empreinte mémoire fixe par séquence.

Ce qui s'est passé

Les grands modèles de langage reposent généralement sur des architectures de transformeurs utilisant l'attention softmax. Lors de l'inférence, ces modèles maintiennent un cache clé-valeur (KV) qui augmente linéairement avec la longueur de la conversation ou du document. Cette croissance limite la concurrence, augmente les coûts de service et force souvent les applications à abandonner le contexte historique. Pour remédier à ces limitations, i14 a développé Laplace 3B, un modèle basé sur les principes de l'attention linéaire.

Contrairement aux transformeurs standards, Laplace 3B utilise un état récurrent qui reste constant à 29,5 Mo par séquence, indépendamment de la longueur du contexte. L'équipe a collaboré avec AWS pour adapter cette architecture au matériel Trainium2. Ils ont utilisé TorchNeuron Native, une fonctionnalité bêta qui permet d'exécuter directement du code PyTorch natif sur les accélérateurs AWS Trainium. Cette approche a préservé la conception Gated Linear Attention (GLA) du modèle ainsi que son utilisation de l'optimiseur Muon.

Les résultats ont montré des gains d'efficacité substantiels. Sur une seule instance trn2.48xlarge, le modèle a été entraîné à environ 120 000 tokens par seconde. Lors de l'inférence, il a servi 11 544 tokens de sortie par seconde sur 1 280 séquences concurrentes. Ces chiffres démontrent que les modèles d'attention linéaire peuvent atteindre un débit élevé sur un matériel spécialisé lorsque la pile logicielle est correctement alignée avec l'architecture sous-jacente.

Comment cela fonctionne

Les modèles d'attention linéaire réorganisent le calcul autour d'un état récurrent plutôt que de comparaisons par paires de tokens. Dans les transformeurs traditionnels, les scores d'attention augmentent de manière quadratique avec la longueur de la séquence lors de l'entraînement. Laplace 3B remplace cela par un mécanisme qui met à jour un état matriciel de taille fixe pour chaque token. Cela change la complexité de l'entraînement de O(T^2) à O(T) et maintient la mémoire d'inférence constante. Le modèle compresse l'historique dans cet état, ce qui nécessite une évaluation minutieuse de la qualité de récupération pour des charges de travail spécifiques.

L'adaptation de ce modèle pour Trainium2 a nécessité trois modifications d'ingénierie spécifiques. Premièrement, l'équipe a remplacé les noyaux Triton optimisés pour GPU par des opérateurs PyTorch standard. Ils ont structuré le noyau GLA pour utiliser 16 blocs de 128 tokens, s'alignant sur les 128 partitions du tampon interne NeuronCore. Deuxièmement, ils ont repensé la communication pour l'optimiseur Muon. En utilisant deux collectifs all-to-all au lieu de méthodes de rassemblement et de diffusion, ils ont réduit le trafic par rang de 24 fois et augmenté le débit de 1,63 fois. Troisièmement, ils ont déroulé statiquement la passe arrière séquentielle pour respecter les limites de taille de graphe et de mémoire du compilateur.

L'architecture Trainium2 a joué un rôle crucial dans ces optimisations. Chaque NeuronCore dispose d'un State Buffer (SBUF) géré par logiciel de 28 MiB avec 128 partitions. Contrairement à un cache, le compilateur planifie chaque transfert de données à l'avance, garantissant qu'aucune éviction ou erreur ne se produit. En mappant chaque bloc de 128 tokens sur une partition, l'équipe a obtenu une augmentation de 14 % du débit mesuré. Cet alignement entre la taille des blocs du modèle et les partitions physiques du matériel était essentiel pour libérer les performances.

Détails clés

• Laplace 3B maintient un état récurrent fixe de 29,5 Mo par séquence, indépendant de la longueur du contexte.
• L'entraînement sur une instance trn2.48xlarge a atteint environ 120 000 tokens par seconde.
• L'inférence a servi 11 544 tokens de sortie par seconde sur 1 280 séquences concurrentes.
• La mise en œuvre a utilisé TorchNeuron Native pour exécuter des opérateurs PyTorch standard sur Trainium2.
• Le trafic de communication pour l'optimiseur Muon a été réduit de 24 fois grâce aux collectifs all-to-all.
• Le modèle utilise 16 blocs de 128 tokens pour s'aligner sur les 128 partitions SBUF des NeuronCores.

Pourquoi c'est important

Pour les ingénieurs développant des produits d'IA, le coût croissant des caches KV constitue un goulot d'étranglement majeur. À mesure que les contextes s'allongent, l'utilisation de la mémoire explose, limitant le nombre d'utilisateurs concurrents qu'un serveur peut gérer. Les modèles d'attention linéaire comme Laplace 3B offrent une voie vers des coûts de service prévisibles. Comme la taille de l'état n'augmente pas avec le contexte, les applications peuvent conserver des historiques longs sans augmentations proportionnelles de l'allocation mémoire. Ceci est particulièrement précieux pour les flux de travail agentiques et l'analyse de documents volumineux où une mémoire persistante est critique.

De plus, ces travaux démontrent que le matériel spécialisé comme Trainium2 peut prendre en charge de nouvelles architectures sans nécessiter de réécriture de bas niveau des noyaux. En utilisant PyTorch natif via TorchNeuron Native, les développeurs peuvent tirer parti des optimisations spécifiques au matériel tout en restant dans un cadre familier. Le succès de la stratégie de communication all-to-all pour l'optimiseur Muon fournit également un modèle pour l'entraînement distribué sur des systèmes fractionnés. Il montre que la réduction de la surcharge de communication peut générer des gains de débit significatifs, même lorsque les FLOPs de calcul restent inchangés.

Ce que vous pouvez faire

• Évaluez les modèles d'attention linéaire pour les cas d'utilisation nécessitant de longues fenêtres de contexte et une utilisation mémoire prévisible.
• Testez TorchNeuron Native si vous déployez des architectures personnalisées sur des instances AWS Trainium.
• Alignez la taille des blocs de données sur les partitions du tampon matériel pour maximiser le débit sur les accélérateurs spécialisés.
• Envisagez des motifs de communication all-to-all pour les optimiseurs distribués afin de réduire les goulots d'étranglement du trafic réseau.
• Surveillez la qualité de récupération lorsque vous passez de l'attention softmax aux mécanismes de compression d'état.
• Profilez votre moteur d'inférence pour identifier si la croissance du cache KV limite vos niveaux de concurrence.