Analyse approfondie · Environ 7 min de lecture
Qwen3 sur Trainium : trois voies torchtitan pour vérifier poids et noyaux
Le guide AWS du 21 octobre relie Transformers, Qwen3 natif et les noyaux NKI sur une instance Trainium2, sans établir un avantage de vitesse ou de coût.
- Publié ici
- Date de la source
- Source vérifiée
Un guide technique AWS publié le 21 octobre explique comment entraîner Qwen3 avec torchtitan sur Trainium selon trois configurations : un backend Transformers v5, le modèle natif de torchtitan et un exemple réduit de mélange d'experts utilisant des noyaux NKI. L'apport concret est une recette d'intégration avec vérification des checkpoints et des résultats numériques. Il ne s'agit ni d'un nouveau modèle Qwen ni d'un benchmark démontrant un entraînement moins cher.
Écouter les actualités
Écoutez cet article ou poursuivez la lecture des actualités.
Préparation de la lecture…
Voix de l’appareil
La voix est traitée par votre navigateur ou appareil. AISubHub n’envoie pas le texte à une API audio. Les voix et la lecture écran verrouillé dépendent de l’appareil.
La reprise peut répéter une courte phrase.
En bref
- Le guide utilise une trn2.48xlarge : 16 accélérateurs Trainium2, 128 NeuronCores physiques et 64 logiques avec LNC=2. Le lancement démarre 64 processus d'entraînement.
- Les exemples denses utilisent les fichiers Qwen3-4B-Instruct-2507 et dix étapes sur c4_test. L'exemple MoE est réduit à quatre couches et 64 experts, dont quatre sélectionnés par token ; il n'entraîne pas le modèle 30B-A3B complet et inchangé.
- AWS rapporte une perte à la première étape de 6.54 avec les poids publiés et de 12.43 avec une initialisation aléatoire. Son test de positionnement rotatif indique un écart absolu maximal de 0.0 avec des entrées float32 ; ces nombres vérifient des choses différentes.
Une intégration matérielle, trois façons de changer le modèle
TorchNeuron enregistre Trainium comme périphérique neuron natif dans PyTorch, ce qui permet à torchtitan d'employer les interfaces habituelles. La première voie confie la construction du modèle à AutoModelForCausalLM dans Transformers v5. La deuxième passe à l'implémentation Qwen3 native de torchtitan via config.model_spec. La troisième utilise des opérateurs NKI propres à Trainium pour les couches de mélange d'experts et montre une substitution personnalisée. Ce sont trois choix d'implémentation derrière un même point d'entrée d'entraînement, pas trois nouveaux modèles Qwen.
Le guide fixe torch 2.14.0, torch_neuronx 2.14.3, Transformers 5.16.0 et torchtitan-neuron release-2.33, version 0.2.2+neuron.1. Il exige une AMI de deep learning Neuron ou un conteneur d'entraînement comprenant pilote, runtime et compilateur. Notre lecture est que du code PyTorch familier réduit la surface d'intégration, sans rendre facultative la combinaison de dépendances testée. AWS prévient explicitement que l'intégration Transformers et les modules de substitution peuvent changer d'une version de torchtitan à l'autre.
La recette 4B vérifie initialisation et portabilité du checkpoint
La première configuration désigne Qwen/Qwen3-4B-Instruct-2507, utilise c4_test, fixe dix étapes et une taille de lot locale de un, puis fragmente l'état du parallélisme de données sur les 64 rangs lancés. Charger l'architecture est distinct de charger les poids publiés. La configuration de fine-tuning active explicitement la lecture du checkpoint depuis les fichiers téléchargés et exporte le modèle final au format Hugging Face safetensors. La perte rapportée à la première étape est de 6.54 avec les poids publiés, contre 12.43 avec une initialisation aléatoire.
Cet écart aide à vérifier l'initialisation, mais ne prouve pas que dix étapes améliorent un assistant ou établissent une convergence finale. AWS recharge aussi le checkpoint exporté avec Transformers pour vérifier ses clés, en tenant compte de lm_head.weight lié aux embeddings. La voie native part du constructeur qwen3_32b parce que la version testée n'offre pas de constructeur 4B générique ; elle remplace ensuite la spécification par 4B et fixe le lot global à 64. Le nom du constructeur ne signifie donc pas qu'un modèle 32B a été entraîné.
L'exemple MoE est volontairement réduit par rapport à sa famille
La recette Qwen3-30B-A3B est ramenée à quatre couches et 64 experts pour tenir sur une seule instance. top_k=4 dirige chaque token vers quatre experts. La source emploie une fragmentation du parallélisme de données sur 64 rangs, un degré de parallélisme tensoriel de un et aucun parallélisme d'experts. Le nombre d'experts est un réglage du modèle, pas le nombre de périphériques participant au parallélisme d'experts. Cette distinction évite de transformer une petite démonstration d'intégration en affirmation sur l'entraînement du modèle original complet.
L'exemple fonctionne en mode eager, avec la compilation désactivée dans sa configuration. Les opérateurs NKI enregistrés sont également compatibles avec torch.compile, mais changer de mode d'exécution constitue une décision distincte. Dans la version testée, la recette remplace le bloc MoE par des noyaux d'experts groupés et choisit l'attention par produit scalaire mis à l'échelle ainsi qu'un positionnement rotatif sans gather. AWS explique comment les opérateurs reçoivent leur implémentation autograd à partir d'informations conservées pendant la passe avant. Il n'y a pas de comparaison de débit démontrant que cette configuration précise est plus rapide.
Un noyau personnalisé exige de vérifier disposition, gradients et ciblage
L'exemple de positionnement rotatif montre qu'enregistrer un noyau ne se résume pas à l'appeler. torchtitan range query et key selon lot, séquence, têtes et dimension de tête ; rope_hf attend lot, têtes, séquence et dimension de tête. L'enveloppe change la disposition, alloue la sortie et inverse la transformation. Un opérateur backward séparé réutilise le cache rotatif sauvegardé, tandis que ce cache n'a lui-même pas de gradient. La longueur de séquence doit être un multiple de 128 fois LNC : 256 lorsque LNC vaut deux.
Un registre de substitutions remplace ensuite la configuration du module voulu. Les journaux de démarrage doivent montrer une substitution nki_rope par couche Qwen3 ; des motifs de nom plus précis peuvent cibler certaines couches. AWS rapporte des écarts absolus maximaux de 0.0 pour query et key lors d'une comparaison float32 avec l'implémentation PyTorch. Cela valide ce test numérique, pas tous les types, formes ou entraînements. L'article présente explicitement cet exemple comme une intégration du noyau et d'autograd, non comme un gain de performances.
Transformer le guide en pilote mesuré
La première étape comprend la compilation et dure plus longtemps que les suivantes. Les journaux exposent perte, norme du gradient, consommation de mémoire à large bande passante, débit et utilisation des FLOPs du modèle, mais le guide ne fournit aucune comparaison contrôlée du coût de bout en bout. Quand modèle, séquence ou lot local grandissent, la mémoire par rang reste une limite pratique. Le parallélisme tensoriel peut se combiner au parallélisme de données entièrement fragmenté sur un DeviceMesh multidimensionnel ; il ne faut pas confondre cet essai plus vaste avec le réglage MoE présenté.
AISubHub suggère de conserver les versions testées, vérifier le modèle choisi et les 64 rangs, contrôler chargement et export du checkpoint, puis comparer les sorties du noyau personnalisé avant une longue exécution. Il faut ensuite mesurer séparément l'échauffement et l'exécution soutenue, puis évaluer la convergence sur les données prévues. Les checkpoints utiles doivent être sauvegardés avant l'arrêt ou la résiliation du matériel loué. Ces étapes d'évaluation sont déduites du guide, sans reproduction par notre équipe, et n'établissent ni droit d'abonnement, ni quota gratuit d'entraînement, ni avantage universel de Trainium.
Pourquoi cela compte
Analyse AISubHub : les équipes peuvent vérifier séparément la compatibilité du modèle, la portabilité des checkpoints et la correction des noyaux avant un entraînement complet. Une structure de lancement commune réduit les changements de recette, mais ne supprime pas les contraintes de versions, mémoire, compilation et calcul numérique.
Limites et disponibilité
Il s'agit d'un guide technique de Josh Longenecker et Alex Hoffman, d'AWS, et non d'un test matériel AISubHub. Il ne donne ni prix comparable pour l'entraînement complet, ni résultat final de convergence, ni accélération démontrée du noyau personnalisé. Disponibilité et facturation de l'instance, versions exactes et validation de la charge restent à la charge de l'opérateur. Aucun avantage d'abonnement grand public n'est annoncé.
Sources et preuves
Date de la source: · Source vérifiée:
Articles assistés par IA, vérifiés avec les sources primaires citées. Les affirmations des fournisseurs sont attribuées ; la présence ne vaut pas recommandation.