Entrene Qwen3 en AWS Trainium con torchtitan

Usar torchtitan en AWS Trainium significa usarlo como siempre lo ha hecho. TorchNeuron registra Trainium como un dispositivo nativo de PyTorch, lo que permite que torchtitan lo utilice a través de las interfaces estándar de PyTorch. Puede usar estas configuraciones para preentrenar desde una inicialización aleatoria o para ajustar pesos publicados. También puede personalizar la implementación del modelo y los kernels para su carga de trabajo. En esta publicación, entrena Qwen3 de tres maneras en una sola trn2.48xlarge.

El primer método usa Transformers v5 como backend de modelado en torchtitan. Ejecuta modelos del Hub a través de AutoModelForCausalLM y lee y escribe safetensors publicados directamente. El segundo método cambia al Qwen3 nativo de torchtitan modificando config.model_spec. El tercero agrega kernels de mezcla de expertos (mixture-of-experts, MoE) de Neuron Kernel Interface (NKI) de las recetas de torchtitan-neuron y luego usa el registro de overrides de torchtitan para adjuntar su propio kernel. Las tres configuraciones usan el mismo torchrun patrón de lanzamiento y estructura de configuración.

Requisitos previos

Para seguir esta guía, use una sola trn2.48xlarge con la AWS Deep Learning AMI para Neuron o un contenedor de entrenamiento de Neuron PyTorch. Ambos entornos proporcionan el driver, el runtime y el compilador de Neuron necesarios para entrenar en Trainium.

Antes de instalar torchtitan-neuron, verifique si su entorno ya lo incluye. El siguiente comando imprime la versión instalada y la ruta del paquete:

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

Esta guía utiliza las versiones de la siguiente tabla. Para reproducir los ejemplos, fije estas versiones. La integración de Transformers v5 y los módulos de override específicos de Trainium pueden cambiar entre versiones de torchtitan.

Cuando torch_neuronx está instalado, TorchNeuron registra Trainium como el neuron dispositivo y torchtitan lo detecta al inicio. En modo eager, Neuron compila las operaciones a medida que se ejecutan. Con torch.compile, Neuron compila el grafo capturado. Este mismo patrón de lanzamiento funciona en las tres configuraciones.

Para las versiones usadas en esta guía, instale torchtitan-neuron y las dependencias restantes de torchtitan con --no-deps para preservar el build de Neuron de torch:

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

Antes de comenzar el entrenamiento, obtén los artefactos del modelo Qwen3 desde Hugging Face Hub. Guárdalos en un directorio local que puedas pasar a la configuración de entrenamiento:

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

Cada configuración de entrenamiento es una función de configuración de Python. Los tres ejemplos parten de una configuración base que actualizas en cada paso.

Entrena con Hugging Face Transformers v5

En la primera configuración, Transformers v5 construye Qwen3 mediante AutoModelForCausalLM. La definición del modelo del Hub proporciona la implementación utilizada para los pases hacia adelante y hacia atrás. Seleccionas el modelo estableciendo su ID del Hub en la configuración.

Comienza con transformers_modeling_backend_full(). Deja sin establecer los campos que se corresponden con las claves de configuración de Hugging Face, de modo que AutoConfig.from_pretrained() cargue la arquitectura del modelo desde el Hub.

Crea neuron_qwen3.py en la raíz del repositorio torchtitan-neuron, desde donde ejecutas el comando de entrenamiento. Añade los imports, la ruta de los artefactos del modelo y la configuración 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

Este comando lanza 64 procesos de entrenamiento, uno por cada NeuronCore lógico. --module neuron_qwen3 importa el módulo de configuración local. --config v5_qwen3_4b selecciona la función de configuración. La configuración de paralelismo de datos distribuye el modelo entre los 64 ranks.

Cargar y guardar checkpoints de Hugging Face

Para hacer fine-tuning a partir de pesos publicados, habilita la carga de checkpoints y apunta initial_load_path a ASSETS. El adaptador del diccionario de estados ajusta el prefijo inicial model. y gestiona los embeddings atados. Guarda el checkpoint final como safetensors de 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

Usa la pérdida del primer paso de entrenamiento como una comprobación rápida de que torchtitan cargó los pesos publicados. Para Qwen/Qwen3-4B-Instruct-2507, la ejecución verificada produjo una pérdida de 6.54 en el paso 1 a partir de los pesos publicados y de 12.43 con inicialización aleatoria.

Después del entrenamiento, carga el checkpoint exportado con AutoModelForCausalLM y load_state_dict. Un ciclo de ida y vuelta exitoso no reporta claves faltantes ni inesperadas aparte de lm_head.weight, que Qwen3-4B ata a los embeddings de tokens. Esto confirma que los archivos siguen siendo un checkpoint estándar de Hugging Face.

Figura 1: el checkpoint publicado comienza con una pérdida de 6.54, y los pesos exportados se cargan sin claves faltantes ni inesperadas

Entrenar con los modelos nativos de torchtitan

En la segunda configuración, torchtitan construye su Qwen3Modelnativa. Añada las importaciones nativas de Qwen3 y la función de configuración a neuron_qwen3.py:

Este ejemplo parte de qwen3_32b() porque la compilación de torchtitan probada no proporciona un constructor genérico de qwen3_4b() config. La siguiente línea reemplaza la especificación del modelo de 32B por la especificación de 4B. Los campos restantes reproducen la ejecución probada.

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 selecciona la clase del modelo y define cómo torchtitan fragmenta el modelo y mapea las claves de su checkpoint. Aquí, model_registry("4B", attn_backend="sdpa") selecciona la implementación nativa de Qwen3 de torchtitan con atención de producto punto escalado (SDPA).

Añadir kernels NKI para mezcla de expertos

En la tercera configuración, torchtitan ejecuta capas de mezcla de expertos con kernels escritos en NKI, el lenguaje de kernels de Neuron. Añada las importaciones necesarias y la función de configuración de Qwen3 MoE a 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

Esta configuración usa fragmentación de paralelismo de datos a través de 64 ranks con un grado de paralelismo de tensores de 1. No usa paralelismo de expertos. num_experts=64 establece el número de expertos en el modelo, no el grado de paralelismo de expertos.

Establecer CompileConfig(enable=False) ejecuta este ejemplo en modo eager. TorchNeuron admite operadores NKI registrados tanto en ejecución eager como en torch.compile.

En la compilación de torchtitan usada en esta guía, qwen3_nki_moe_config() reemplaza el bloque MoE con la implementación NKI de expertos agrupados. También selecciona SDPA y la incrustación rotatoria sin gather. Cuando habilite torch.compile, añada "grad_clip" a CompileConfig.components para compilar el recorte de norma de gradiente.

Los kernels NKI entran al entrenamiento como operadores de PyTorch. wrap_nki hace que un kernel NKI sea invocable desde PyTorch, y nki_op registra el operador. register_autograd adjunta su pasada hacia atrás usando metadatos de enrutamiento guardados durante la pasada hacia adelante.

El ejemplo reduce Qwen3-30B-A3B a 4 capas y 64 expertos para que quepa en una trn2.48xlarge. Enruta cada token a 4 expertos con top_k=4.

Adjuntar su propio kernel NKI

En algunos casos, puede querer optimizar una operación para las formas de tensores de su modelo. Un kernel especializado puede adaptar su tiling y movimiento de datos a esas formas. También puede reemplazar una secuencia crítica para el rendimiento de operaciones de PyTorch con un solo operador optimizado. Puede escribir el kernel usted mismo o comenzar con uno de nkilib. Este ejemplo usa rope_hf de nkilib.core.embeddings, que espera el diseño de tensores de Hugging Face y admite su propio modo hacia atrás. Este ejemplo demuestra el registro de kernels y la integración con autograd, no una mejora de rendimiento.

Crear my_neuron_kernels.py en la raíz del repositorio, donde lanzas el entrenamiento. torchtitan almacena query y key como [batch, seq, heads, head_dim], mientras que rope_hf espera [batch, heads, seq, head_dim]. Convierte la disposición dentro del operador y asigna los tensores de salida con 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 calcula los gradientes cuando lo llamas con backward=True. Registra un segundo operador para ese modo. Guarda la caché rotatoria durante la pasada hacia adelante para que la pasada hacia atrás pueda reutilizarla. Devuelve None para el gradiente de la caché porque la caché no es diferenciable.

@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)

Debido a que rope_hf requiere una longitud de secuencia soportada, valida esa restricción antes de llamar al kernel. Crea una subclase de CosSinRoPE y llama a rope_fwd desde 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)

Registra NKIRoPE como el reemplazo de CosSinRoPE. derive() copia los campos de configuración compartidos en la configuración de reemplazo:

@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)

Para cargar la anulación personalizada, pasa su nombre de módulo a --override.imports. torchtitan importa la entrada del registro antes de construir el modelo:

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

Antes de que comience el entrenamiento, torchtitan registra en el log cada nodo de configuración que la anulación reemplaza. Confirma que la salida contiene una entrada nki_rope por cada capa 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

Establece la longitud de la secuencia en un múltiplo de 128 * LNC, que es 256 cuando LNC es 2. Reformatea la caché rotatoria a [seq, head_dim * 2], con los valores del coseno en la primera mitad y los valores del seno en la segunda. El kernel valida estas restricciones antes de que comience el entrenamiento.

Para un ejemplo completo de torchtitan, consulta torchtitan_neuron/models/gpt_oss/nki_attention.py. Conecta los kernels nkilib flash-attention a través de los mismos mecanismos de operador y anulación. También gestiona entradas DTensor para el paralelismo de tensores.

Antes de entrenar con un kernel personalizado, compara sus salidas con la implementación que reemplaza. En la prueba numérica con entradas float32, rope_hf y el rotary embedding de PyTorch de torchtitan produjeron una diferencia absoluta máxima de 0.0 tanto para query como para key.

Cuando registras múltiples anulaciones, limita cada una a las instancias de módulo previstas con fqns. Por ejemplo, fqns=["layers.1[0-9].attention.rope"] apunta al rotary embedding en las capas 10 a 19. torchtitan se detiene cuando dos anulaciones reclaman el mismo nodo y advierte cuando una anulación no coincide con ningún nodo.

Revisa la salida del entrenamiento

Al iniciar, confirme que torchtitan informa 64 ranks y el modelo esperado. El primer paso incluye la compilación y tarda más que los pasos posteriores. Cada paso de entrenamiento informa la pérdida, la norma del gradiente, el uso de la memoria de alto ancho de banda (HBM), el rendimiento y la utilización de FLOPs del modelo (MFU). Para la ejecución del checkpoint, confirme que la pérdida del primer paso esté cerca de 6.54. Para la ejecución del kernel personalizado, confirme que la nki_rope override aparece una vez por cada capa de Qwen3.

Escalar en toda la instancia

Las mismas opciones de paralelismo se aplican a las configuraciones de Transformers y de modelos nativos. A trn2.48xlarge expone 64 NeuronCores lógicos en LNC=2. Establezca data_parallel_shard_degree = -1 particionar los parámetros del modelo, los gradientes y el estado del optimizador entre todos los rangos que torchrun comienza.

Conjunto config.parallelism.tensor_parallel_degree para fragmentar los pesos de cada capa en un grupo de rangos. Cada rango calcula una porción de cada multiplicación de matrices. Con el paralelismo de secuencia habilitado por defecto, las capas de normalización se ejecutan sobre activaciones fragmentadas a lo largo de la dimensión de secuencia. torchtitan combina el paralelismo de tensores con Fully Sharded Data Parallel (FSDP) en un multidimensional DeviceMesh.

Observe la HBM por rango a medida que aumenta el tamaño del modelo, la longitud de la secuencia o el tamaño del lote local. Estos valores determinan cuánta memoria requiere cada rango.

Limpiar

Si desea conservar sus checkpoints, cópielos en un almacenamiento persistente antes de realizar la limpieza. Si planea reutilizar la instancia, detente. Cuando ya no necesite la instancia, termínela.

Conclusión

AWS Trainium funciona como un dispositivo compatible con torchtitan en las tres configuraciones de este tutorial. Puede entrenar Qwen3 a través de Hugging Face Transformers v5 y cargar o guardar safetensors publicados directamente. También puede entrenar con la implementación nativa de Qwen3 de torchtitan configurando config.model_spec. Para modelos de mezcla de expertos (mixture-of-experts), use las recetas NKI de torchtitan-neuron. Para kernels personalizados, use el registro de overrides para adjuntar operadores NKI a través del mismo flujo de lanzamiento.

Para obtener orientación sobre la configuración, consulte la documentación de AWS Neuron. Para el código fuente actual y las recetas de entrenamiento, consulte el repositorio de torchtitan.

Sobre los autores