Entraîner Qwen3 sur AWS Trainium avec torchtitan

Utiliser torchtitan sur AWS Trainium signifie l'utiliser comme vous l'avez toujours fait. TorchNeuron enregistre Trainium comme un dispositif PyTorch natif, ce qui permet à torchtitan de l'utiliser via les interfaces PyTorch standard. Vous pouvez utiliser ces configurations pour pré-entraîner à partir d'une initialisation aléatoire ou affiner des poids publiés. Vous pouvez également personnaliser l'implémentation du modèle et les noyaux pour votre charge de travail. Dans cet article, vous entraînez Qwen3 de trois manières sur un seul trn2.48xlarge.

La première méthode utilise Transformers v5 comme backend de modélisation dans torchtitan. Elle exécute les modèles du Hub via AutoModelForCausalLM et lit et écrit directement les safetensors publiés. La deuxième méthode passe au Qwen3 natif de torchtitan en changeant config.model_spec. La troisième ajoute les noyaux Neuron Kernel Interface (NKI) de mixture-of-experts (MoE) provenant des recettes torchtitan-neuron, puis utilise le registre de surcharges de torchtitan pour attacher votre propre noyau. Les trois configurations utilisent le même torchrun modèle de lancement et la même structure de configuration.

Prérequis

Pour suivre ce guide, utilisez un seul trn2.48xlarge avec soit l'AWS Deep Learning AMI pour Neuron, soit un conteneur d'entraînement Neuron PyTorch. Les deux environnements fournissent le pilote, le runtime et le compilateur Neuron nécessaires pour l'entraînement sur Trainium.

Avant d'installer torchtitan-neuron, vérifiez si votre environnement l'inclut déjà. La commande suivante affiche la version installée et le chemin du paquet :

python3 -c "import torchtitan, torchtitan_neuron; print(torchtitan.__version__, torchtitan_neuron.__file__)"

Ce guide utilise les versions indiquées dans le tableau suivant. Pour reproduire les exemples, fixez ces versions. L'intégration de Transformers v5 et les modules de surcharge spécifiques à Trainium peuvent changer entre les versions de torchtitan.

Lorsque torch_neuronx est installé, TorchNeuron enregistre Trainium comme le neuron dispositif et torchtitan le détecte au démarrage. En mode eager, Neuron compile les opérations au fur et à mesure de leur exécution. Avec torch.compile, Neuron compile le graphe capturé. Ce même modèle de lancement fonctionne pour les trois configurations.

Pour les versions utilisées dans ce guide, installez torchtitan-neuron et les dépendances restantes de torchtitan avec --no-deps pour préserver la version Neuron de torch :

git clone https://github.com/aws-neuron/torchtitan-neuron
cd torchtitan-neuron
pip install -r requirements.txt

Avant de commencer l'entraînement, récupérez les artefacts du modèle Qwen3 depuis Hugging Face Hub. Stockez-les dans un répertoire local que vous pourrez passer à la configuration d'entraînement :

hf download Qwen/Qwen3-4B-Instruct-2507 \
  --local-dir "$PWD/assets/hf/Qwen3-4B-Instruct-2507"

Chaque configuration d'entraînement est une fonction de configuration Python. Les trois exemples partent d'une configuration de base que vous mettez à jour à chaque étape.

Entraîner avec Hugging Face Transformers v5

Dans la première configuration, Transformers v5 construit Qwen3 via AutoModelForCausalLM. La définition du modèle du Hub fournit l'implémentation utilisée pour les passes avant et arrière. Vous sélectionnez le modèle en définissant son identifiant Hub dans la configuration.

Commencez par transformers_modeling_backend_full(). Il laisse non définis les champs qui correspondent aux clés de configuration Hugging Face, de sorte que AutoConfig.from_pretrained() charge l'architecture du modèle depuis le Hub.

Créez neuron_qwen3.py à la racine du dépôt torchtitan-neuron, où vous exécutez la commande d'entraînement. Ajoutez les imports, le chemin des artefacts du modèle et la configuration de base :

import os
from torchtitan.experiments.transformers_modeling_backend.config_registry import (
    transformers_modeling_backend_full,
)
from torchtitan.experiments.transformers_modeling_backend.configs import (
    TransformersBackendConfig,
)

ASSETS = os.path.abspath("./assets/hf/Qwen3-4B-Instruct-2507")

def v5_qwen3_4b() -> TransformersBackendConfig:
    config = transformers_modeling_backend_full()
    config.hf_model = "Qwen/Qwen3-4B-Instruct-2507"
    config.hf_assets_path = ASSETS
    config.dataloader.dataset = "c4_test"
    config.training.steps = 10
    config.training.local_batch_size = 1
    config.parallelism.data_parallel_shard_degree = -1
    return config

torchrun --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0 \
  --local-ranks-filter 0 --role rank --tee 3 \
  -m torchtitan.train --module neuron_qwen3 --config v5_qwen3_4b

Cette commande lance 64 processus d'entraînement, un pour chaque NeuronCore logique. --module neuron_qwen3 importe le module de configuration local. --config v5_qwen3_4b sélectionne la fonction de configuration. Le réglage de parallélisme des données répartit le modèle sur les 64 ranks.

Charger et enregistrer des checkpoints Hugging Face

Pour effectuer un fine-tuning à partir de poids publiés, activez le chargement de checkpoints et faites pointer initial_load_path vers ASSETS. L'adaptateur de state-dictionary ajuste le préfixe initial model. et gère les embeddings liés. Il enregistre le checkpoint final sous forme de safetensors Hugging Face.

def v5_qwen3_4b_finetune() -> TransformersBackendConfig:
    config = v5_qwen3_4b()
    config.checkpoint.enable = True
    config.checkpoint.folder = "checkpoint_qwen3_4b_ft"
    config.checkpoint.initial_load_path = ASSETS # absolute path required
    config.checkpoint.initial_load_in_hf = True
    config.checkpoint.initial_load_model_only = True
    config.checkpoint.last_save_in_hf = True
    config.checkpoint.last_save_model_only = True
    return config

Utilisez la perte de la première étape d'entraînement comme vérification rapide que torchtitan a chargé les poids publiés. Pour Qwen/Qwen3-4B-Instruct-2507, l'exécution vérifiée a produit une perte de 6.54 à l'étape 1 avec les poids publiés et de 12.43 avec une initialisation aléatoire.

Après l'entraînement, chargez le checkpoint exporté avec AutoModelForCausalLM et load_state_dict. Un aller-retour réussi ne signale aucune clé manquante ou inattendue autre que lm_head.weight, que Qwen3-4B lie aux embeddings de tokens. Cela confirme que les fichiers restent un checkpoint Hugging Face standard.

Figure 1 : le checkpoint publié démarre avec une loss de 6,54, et les poids exportés se chargent sans clés manquantes ni inattendues

Entraîner avec les modèles natifs de torchtitan

Dans la seconde configuration, torchtitan construit son modèle natif Qwen3Model. Ajoutez les imports Qwen3 natifs et la fonction de config à neuron_qwen3.py:

Cet exemple part de qwen3_32b() car la version testée de torchtitan ne fournit pas de constructeur générique de qwen3_4b() config. La ligne suivante remplace la spécification du modèle 32B par la spécification du modèle 4B. Les champs restants reproduisent l'exécution testée.

from torchtitan.models.qwen3 import model_registry
from torchtitan.models.qwen3.config_registry import qwen3_32b
from torchtitan.trainer import Trainer

def native_qwen3_4b() -> Trainer.Config:
    config = qwen3_32b()
    config.model_spec = model_registry("4B", attn_backend="sdpa")
    config.hf_assets_path = ASSETS
    config.dataloader.dataset = "c4_test"
    config.training.steps = 10
    config.training.local_batch_size = 1
    config.training.global_batch_size = 64
    config.parallelism.data_parallel_shard_degree = -1
    return config

torchrun --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0 \
  --local-ranks-filter 0 --role rank --tee 3 \
  -m torchtitan.train --module neuron_qwen3 --config native_qwen3_4b

config.model_spec sélectionne la classe du modèle et définit comment torchtitan partitionne le modèle et mappe les clés de son checkpoint. Ici, model_registry("4B", attn_backend="sdpa") sélectionne l'implémentation Qwen3 native de torchtitan avec l'attention par produit scalaire mis à l'échelle (SDPA).

Ajouter des kernels NKI pour le mixture of experts

Dans la troisième configuration, torchtitan exécute des couches mixture-of-experts avec des kernels écrits en NKI, le langage de kernels Neuron. Ajoutez les imports requis et la fonction de config Qwen3 MoE à neuron_qwen3.py:

from torchtitan.config import CompileConfig
from torchtitan_neuron.models.qwen3.config_registry import qwen3_nki_moe_config

def nki_qwen3_30b_a3b() -> Trainer.Config:
    config = qwen3_nki_moe_config(
        "30B-A3B",
        compile_config=CompileConfig(enable=False, backend="neuron"),
        n_layers=4,
        num_experts=64,
        top_k=4,
    )
    config.hf_assets_path = ASSETS
    return config

torchrun --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0 \
  --local-ranks-filter 0 --role rank --tee 3 \
  -m torchtitan.train --module neuron_qwen3 --config nki_qwen3_30b_a3b

Cette configuration utilise le sharding data parallel sur 64 ranks avec un degré de tensor parallel de 1. Elle n'utilise pas l'expert parallelism. num_experts=64 définit le nombre d'experts dans le modèle, et non le degré d'expert parallel.

Définir CompileConfig(enable=False) exécute cet exemple en mode eager. TorchNeuron prend en charge les opérateurs NKI enregistrés à la fois en exécution eager et torch.compile.

Dans la version de torchtitan utilisée pour ce guide, qwen3_nki_moe_config() remplace le bloc MoE par l'implémentation NKI d'experts groupés. Elle sélectionne également SDPA et l'embedding rotatif sans gather. Lorsque vous activez torch.compile, ajoutez "grad_clip" à CompileConfig.components pour compiler le clipping de la norme des gradients.

Les kernels NKI entrent dans l'entraînement en tant qu'opérateurs PyTorch. wrap_nki rend un kernel NKI appelable depuis PyTorch, et nki_op enregistre l'opérateur. register_autograd attache sa passe backward à l'aide des métadonnées de routage sauvegardées lors de la passe forward.

L'exemple réduit Qwen3-30B-A3B à 4 couches et 64 experts afin qu'il tienne sur un seul trn2.48xlarge. Il route chaque token vers 4 experts avec top_k=4.

Attacher votre propre kernel NKI

Dans certains cas, vous pourriez vouloir optimiser une opération pour les formes de tenseurs de votre modèle. Un kernel spécialisé peut adapter son tiling et son déplacement de données à ces formes. Il peut aussi remplacer une séquence d'opérations PyTorch critique pour les performances par un seul opérateur optimisé. Vous pouvez écrire le kernel vous-même ou partir d'un kernel de nkilib. Cet exemple utilise rope_hf de nkilib.core.embeddings, qui attend la disposition des tenseurs de Hugging Face et prend en charge son propre mode backward. Cet exemple illustre l'enregistrement du kernel et l'intégration à autograd, et non une amélioration des performances.

Créer my_neuron_kernels.py dans la racine du dépôt, où vous lancez l'entraînement. torchtitan stocke la requête (query) et la clé (key) sous la forme [batch, seq, heads, head_dim], tandis que rope_hf s'attend à [batch, heads, seq, head_dim]. Convertissez la disposition à l'intérieur de l'opérateur et allouez les tenseurs de sortie avec torch.empty_like():

import torch
from nkilib.core.embeddings.rope_hf import rope_hf
from torch_neuronx import nki_op
from torch_neuronx.nki_hop import wrap_nki
from torch_neuronx.utils import get_logical_neuron_cores

LNC = (int(get_logical_neuron_cores()),)
wrapped_rope = wrap_nki(rope_hf)

@nki_op("my_kernels::rope_fwd")
def rope_fwd(query, key, rope_cache):
    q = query.permute(0, 2, 1, 3).contiguous()
    k = key.permute(0, 2, 1, 3).contiguous()
    cache = rope_cache.reshape(-1, rope_cache.shape[-1])
    q_out, k_out = wrapped_rope[LNC](
        q,
        k,
        torch.empty_like(q),
        torch.empty_like(k),
        rope_cache=cache,
    )
    return (
        q_out.permute(0, 2, 1, 3).contiguous(),
        k_out.permute(0, 2, 1, 3).contiguous(),
    )

rope_hf calcule les gradients lorsqu'on l'appelle avec backward=True. Enregistrez un second opérateur pour ce mode. Sauvegardez le cache rotatif pendant la passe avant afin que la passe arrière puisse le réutiliser. Renvoyez None pour le gradient du cache, car le cache n'est pas différentiable.

@nki_op("my_kernels::rope_bwd")
def rope_bwd(grad_query, grad_key, rope_cache):
    gq = grad_query.permute(0, 2, 1, 3).contiguous()
    gk = grad_key.permute(0, 2, 1, 3).contiguous()
    cache = rope_cache.reshape(-1, rope_cache.shape[-1])
    out_q, out_k = wrapped_rope[LNC](
        gq,
        gk,
        torch.empty_like(gq),
        torch.empty_like(gk),
        rope_cache=cache,
        backward=True,
    )
    return (
        out_q.permute(0, 2, 1, 3).contiguous(),
        out_k.permute(0, 2, 1, 3).contiguous(),
    )

def _setup_context(ctx, inputs, output):
    ctx.save_for_backward(inputs[2])

def _backward(ctx, grad_query, grad_key):
    (rope_cache,) = ctx.saved_tensors
    grad_q, grad_k = rope_bwd(grad_query, grad_key, rope_cache)
    return grad_q, grad_k, None

rope_fwd.register_autograd(_backward, setup_context=_setup_context)

Parce que rope_hf exige une longueur de séquence prise en charge, validez cette contrainte avant d'appeler le kernel. Dérivez CosSinRoPE et appelez rope_fwd depuis apply_rotary_emb:

from dataclasses import dataclass
from torchtitan.config import derive, override
from torchtitan.models.common.rope import CosSinRoPE

class NKIRoPE(CosSinRoPE):
    @dataclass(kw_only=True, slots=True)
    class Config(CosSinRoPE.Config):
        pass

    @staticmethod
    def apply_rotary_emb(query, key, rope_cache):
        stride = 128 * LNC[0]
        if query.shape[1] % stride:
            raise ValueError(
                f"rope_hf requires a sequence length that is a multiple of "
                f"128 * LNC={stride}, got {query.shape[1]}"
            )
        q_out, k_out = rope_fwd(query, key, rope_cache)
        return q_out.type_as(query), k_out.type_as(key)

Enregistrez NKIRoPE comme le remplacement de CosSinRoPE. derive() copie les champs de configuration partagés dans la config de remplacement :

@override(
    "nki_rope",
    target=CosSinRoPE.Config,
    exact=True,
    description="nkilib rope_hf kernel in place of the PyTorch rotation",
)
def _nki_rope(cfg: CosSinRoPE.Config) -> NKIRoPE.Config:
    return derive(cfg, NKIRoPE.Config)

Pour charger la surcharge personnalisée, passez son nom de module à --override.imports. torchtitan importe l'entrée du registre avant de construire le modèle :

torchrun --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0 \
  --local-ranks-filter 0 --role rank --tee 3 \
  -m torchtitan.train --module neuron_qwen3 --config native_qwen3_4b \
  --override.imports my_neuron_kernels

Avant le début de l'entraînement, torchtitan journalise chaque nœud de config que la surcharge remplace. Vérifiez que la sortie contient une entrée nki_rope pour chaque couche Qwen3 :

[Override] nki_rope: model_spec.model.layers.0.attention.rope CosSinRoPEConfig -> NKIRoPE.Config
[Override] nki_rope: model_spec.model.layers.1.attention.rope CosSinRoPEConfig -> NKIRoPE.Config

Définissez la longueur de séquence sur un multiple de 128 * LNC, soit 256 lorsque LNC vaut 2. Remodelez le cache rotatif en [seq, head_dim * 2], avec les valeurs de cosinus dans la première moitié et les valeurs de sinus dans la seconde. Le kernel valide ces contraintes avant le début de l'entraînement.

Pour un exemple torchtitan complet, consultez torchtitan_neuron/models/gpt_oss/nki_attention.py. Il connecte les kernels nkilib flash-attention via les mêmes mécanismes d'opérateur et de surcharge. Il gère également les entrées DTensor pour le parallélisme de tenseurs.

Avant d'entraîner avec un kernel personnalisé, comparez ses sorties avec l'implémentation qu'il remplace. Dans le test numérique avec des entrées float32, rope_hf et le rotary embedding PyTorch de torchtitan ont produit une différence absolue maximale de 0.0 pour la query et la key.

Lorsque vous enregistrez plusieurs surcharges, limitez chacune aux instances de module prévues avec fqns. Par exemple, fqns=["layers.1[0-9].attention.rope"] cible le rotary embedding des couches 10 à 19. torchtitan s'arrête lorsque deux surcharges revendiquent le même nœud et avertit lorsqu'une surcharge ne correspond à aucun nœud.

Examinez la sortie de l'entraînement

Au démarrage, vérifiez que torchtitan indique 64 rangs et le modèle attendu. La première étape inclut la compilation et prend plus de temps que les étapes suivantes. Chaque étape d'entraînement indique la loss, la norme du gradient, l'utilisation de la mémoire à large bande passante (HBM), le débit et l'utilisation des FLOPs du modèle (MFU). Pour l'exécution avec checkpoint, vérifiez que la loss de la première étape est proche de 6.54. Pour l'exécution avec noyau personnalisé, vérifiez que la nki_rope L'override apparaît une fois pour chaque couche de Qwen3.

Mise à l'échelle sur toute l'instance

Les mêmes paramètres de parallélisme s'appliquent aux configurations Transformers et aux configurations de modèles natifs. A trn2.48xlarge expose 64 NeuronCores logiques à LNC=2. Définir data_parallel_shard_degree = -1 pour sharder les paramètres du modèle, les gradients et l'état de l'optimiseur sur chaque rang qui torchrun commence.

Définir config.parallelism.tensor_parallel_degree pour partager (sharder) les poids de chaque couche entre un groupe de ranks. Chaque rank calcule une tranche de chaque multiplication matricielle. Avec le parallélisme de séquence activé par défaut, les couches de normalisation s'exécutent sur des activations partagées le long de la dimension de séquence. torchtitan combine le parallélisme de tenseurs avec le Fully Sharded Data Parallel (FSDP) sur un multidimensionnel DeviceMesh.

Observez la HBM par rank à mesure que vous augmentez la taille du modèle, la longueur de séquence ou la taille de lot locale. Ces valeurs déterminent la quantité de mémoire requise par chaque rank.

Nettoyage

Si vous souhaitez conserver vos checkpoints, copiez-les vers un stockage persistant avant de procéder au nettoyage. Si vous prévoyez de réutiliser l'instance, arrêtez-la. Lorsque vous n'avez plus besoin de l'instance, résiliez-la.

Conclusion

AWS Trainium fonctionne comme un appareil torchtitan pris en charge dans les trois configurations de ce guide pratique. Vous pouvez entraîner Qwen3 via Hugging Face Transformers v5 et charger ou enregistrer directement des safetensors publiés. Vous pouvez également entraîner avec l'implémentation native de Qwen3 de torchtitan en définissant config.model_spec. Pour les modèles mixture-of-experts, utilisez les recettes NKI de torchtitan-neuron. Pour les kernels personnalisés, utilisez le registre de remplacement (override registry) pour associer des opérateurs NKI via le même flux de lancement.

Pour obtenir des conseils de configuration, consultez la documentation AWS Neuron. Pour le code source actuel et les recettes d'entraînement, consultez le dépôt torchtitan.

À propos des auteurs