Entraînement de Qwen3 sur AWS Trainium avec torchtitan et des kernels personnalisés

AWS a publié un tutoriel technique détaillé démontrant comment entraîner le grand modèle de langage Qwen3 sur ses accélérateurs Trainium2 en utilisant le framework torchtitan. Publié en octobre 2026, ce guide présente trois chemins de configuration distincts pour les ingénieurs souhaitant tirer parti des interfaces natives PyTorch sur du matériel spécialisé. Cette documentation sert de référence pratique pour les équipes de machine learning visant à optimiser les workflows de pré-entraînement et de fine-tuning sur l'infrastructure Amazon Web Services.

Ce qui s'est passé

L'article fournit un guide d'implémentation étape par étape pour exécuter les modèles Qwen3 sur une instance unique trn2.48xlarge, qui contient 16 accélérateurs Trainium2 et 128 NeuronCores physiques. Les auteurs démontrent que les développeurs peuvent utiliser torchtitan sur Trainium tout comme ils le feraient sur des clusters GPU standards, grâce à TorchNeuron qui enregistre Trainium comme un périphérique PyTorch natif. Cela permet au framework de détecter le matériel au démarrage et de compiler les opérations soit de manière immédiate (eager), soit via torch.compile.

Le tutoriel couvre trois configurations d'entraînement spécifiques. La première méthode utilise Hugging Face Transformers v5 comme backend de modélisation, permettant aux utilisateurs de charger directement des modèles depuis le Hub via AutoModelForCausalLM. La deuxième approche bascule vers l'implémentation native Qwen3 de torchtitan, offrant un meilleur contrôle sur l'architecture du modèle et les stratégies de sharding. La troisième méthode introduit une personnalisation avancée en intégrant des kernels mixture-of-experts de l'interface Neuron Kernel Interface (NKI), permettant aux développeurs d'ajouter des opérateurs personnalisés pour les couches critiques en termes de performance.

Chaque configuration utilise le même schéma de lancement torchrun, garantissant la cohérence entre les différents environnements expérimentaux. Le guide met l'accent sur la reproductibilité en spécifiant des versions logicielles exactes, notamment PyTorch 2.14.0, Transformers 5.16.0 et la release-2.33 de torchtitan-neuron. Il détaille également comment gérer les checkpoints, en vérifiant que les poids sauvegardés au format safetensors de Hugging Face peuvent être rechargés sans clés manquantes, confirmant ainsi l'intégrité du pipeline d'entraînement.

Comment cela fonctionne

Le mécanisme central repose sur l'intégration entre torchtitan et la stack AWS Neuron. Lorsque torch_neuronx est installé, il expose Trainium comme un périphérique neuron. Torchtitan utilise ensuite les primitives standard de distribution d'entraînement PyTorch, telles que Fully Sharded Data Parallel (FSDP) et le parallélisme tensoriel, pour répartir la charge de travail sur les 64 NeuronCores logiques disponibles sur l'instance. Le framework gère le sharding des paramètres du modèle, des gradients et des états de l'optimiseur, tandis que le parallélisme séquentiel gère le sharding des activations le long de la dimension de séquence.

Pour l'intégration de kernels personnalisés, le guide démontre comment utiliser le registre de surcharge (override registry) de torchtitan. Les développeurs peuvent écrire des kernels en NKI, le langage de kernel Neuron, et les envelopper comme des opérateurs PyTorch en utilisant wrap_nki et nki_op. Ces opérateurs sont ensuite enregistrés pour remplacer des modules spécifiques dans le graphe du modèle. Par exemple, un kernel personnalisé pour l'embedding rotatif peut remplacer l'implémentation PyTorch par défaut. Le système valide ces surcharges au démarrage, en s'assurant que chaque remplacement cible les noms qualifiés complets corrects au sein de la structure du modèle.

La boucle d'entraînement rapporte des métriques clés telles que la perte (loss), la norme du gradient, l'utilisation de la mémoire à large bande passante et l'utilisation des FLOPs du modèle. Cette visibilité permet aux ingénieurs de surveiller la consommation des ressources et d'ajuster dynamiquement les paramètres de parallélisme. Le processus de compilation se déroule lors de la première étape d'entraînement, ce qui prend plus de temps que les étapes suivantes, mais permet une exécution optimisée pour le reste de la session.

Détails clés

• La configuration nécessite une instance trn2.48xlarge avec 64 NeuronCores logiques configurés avec LNC=2.
• Les dépendances logicielles incluent PyTorch 2.14.0, Transformers 5.16.0 et torchtitan-neuron 0.2.2+neuron.1.
• Le chargement des checkpoints Hugging Face vérifie la correction en contrôlant qu'une perte à l'étape 1 proche de 6,54 est obtenue pour le modèle Qwen3-4B-Instruct-2507.
• Les kernels NKI personnalisés doivent respecter des contraintes spécifiques de longueur de séquence, telles que devoir être un multiple de 128 * LNC.
• Le parallélisme tensoriel shard les poids des couches entre les rangs, tandis que le parallélisme de données shard les paramètres et les états de l'optimiseur.
• Le registre de surcharge empêche les conflits en s'arrêtant si deux surcharges revendiquent le même nœud et en émettant un avertissement si aucun nœud ne correspond.

Pourquoi c'est important

Pour les ingénieurs logiciels et les praticiens du ML, ce guide abaisse la barrière à l'entrée pour l'entraînement de grands modèles sur du matériel non-NVIDIA. En maintenant la compatibilité avec les interfaces PyTorch standard et les écosystèmes Hugging Face, il réduit le besoin de refactoring de code extensif. Les équipes peuvent exploiter les dépôts de modèles et les formats de checkpoints existants, minimisant ainsi la charge opérationnelle associée à l'adoption de nouvelles architectures d'accélérateurs.

La capacité d'injecter des kernels NKI personnalisés est particulièrement significative pour l'optimisation des performances. À mesure que les modèles gagnent en complexité, les implémentations génériques peuvent ne pas utiliser pleinement les capacités matérielles. Permettre aux développeurs de remplacer des opérations spécifiques, telles que les mécanismes d'attention ou les embeddings rotatifs, par des kernels ajustés manuellement offre un contrôle plus fin sur l'efficacité computationnelle. Cette flexibilité est cruciale pour atteindre un débit compétitif et une rentabilité dans les jobs d'entraînement à grande échelle.

De plus, les étapes de vérification détaillées offrent une confiance accrue dans la fiabilité du processus d'entraînement. Confirmer que les checkpoints exportés se chargent correctement garantit que les modèles entraînés sur Trainium peuvent être déployés sans heurt dans des environnements de production reposant sur des stacks d'inférence standard. Cette validation de bout en bout soutient des pratiques MLOps robustes et facilite la collaboration entre les équipes utilisant différents backends matériels.

Ce que vous pouvez faire

• Configurez une instance trn2.48xlarge avec l'AMI Deep Learning AWS pour Neuron ou un conteneur PyTorch Neuron.
• Installez les versions spécifiées de torchtitan-neuron et des dépendances en utilisant le fichier requirements.txt fourni.
• Téléchargez les artefacts du modèle Qwen3 depuis Hugging Face Hub vers un répertoire local pour l'entraînement.
• Configurez le script d'entraînement pour utiliser le sharding de parallélisme de données sur tous les 64 rangs en définissant data_parallel_shard_degree = -1.
• Vérifiez les sorties des kernels personnalisés par rapport aux implémentations de base en utilisant des entrées float32 pour assurer la précision numérique.
• Surveillez l'utilisation de la mémoire à large bande passante par rang afin d'ajuster correctement les tailles de batch et les longueurs de séquence.