Visão geral
A AWS publicou um guia técnico mostrando como treinar o modelo Qwen3 usando o framework torchtitan sobre o hardware AWS Trainium. O ponto central da abordagem é que o TorchNeuron registra o Trainium como um dispositivo PyTorch nativo, o que permite ao torchtitan utilizá-lo pelas interfaces padrão do PyTorch — sem adaptações especiais no fluxo de trabalho.
O walkthrough cobre três configurações distintas, todas executadas em uma única instância trn2.48xlarge: treinamento via Hugging Face Transformers v5, treinamento com a implementação nativa Qwen3 do torchtitan, e uso de kernels NKI (Interface de Kernel Neuron) para camadas de Mistura de Especialistas (MoE). As três abordagens compartilham o mesmo padrão de lançamento com torchrun e a mesma estrutura de configuração.
Pré-requisitos
Para reproduzir os exemplos, a AWS recomenda usar uma instância trn2.48xlarge com a AWS Deep Learning AMI para Neuron ou um container de treinamento PyTorch Neuron. Ambos os ambientes já incluem o driver, runtime e compilador Neuron necessários.
Antes de instalar o torchtitan-neuron, vale verificar se o ambiente já o contém:
python3 -c "import torchtitan, torchtitan_neuron; print(torchtitan.__version__, torchtitan_neuron.__file__)"As versões utilizadas no guia são as seguintes:
- Instância: trn2.48xlarge — 16 aceleradores Trainium2, 128 NeuronCores físicos
- NeuronCores lógicos: 64 com LNC=2
- torch / torch_neuronx: 2.14.0 / 2.14.3
- transformers: 5.16.0
- torchtitan-neuron: release-2.33 (0.2.2+neuron.1)
Quando o torch_neuronx está instalado, o TorchNeuron registra o Trainium como dispositivo neuron e o torchtitan o detecta na inicialização. No modo eager, o Neuron compila as operações conforme são executadas. Com torch.compile, o Neuron compila o grafo capturado.
Para instalar o torchtitan-neuron preservando o build Neuron do torch:
git clone https://github.com/aws-neuron/torchtitan-neuron
cd torchtitan-neuron
pip install -r requirements.txtAntes de iniciar o treinamento, é necessário baixar os artefatos do modelo Qwen3 do Hugging Face Hub:
hf download Qwen/Qwen3-4B-Instruct-2507 \
--local-dir "$PWD/assets/hf/Qwen3-4B-Instruct-2507"Configuração 1: Treinamento com Hugging Face Transformers v5
Na primeira configuração, o Transformers v5 constrói o Qwen3 via AutoModelForCausalLM. A definição do modelo vinda do Hub é utilizada nos passos de forward e backward. O modelo é selecionado pelo seu ID no Hub.
Crie o arquivo neuron_qwen3.py na raiz do repositório torchtitan-neuron e adicione as importações, o caminho dos artefatos e a configuração 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 configtorchrun --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_4bEsse comando lança 64 processos de treinamento — um por NeuronCore lógico. O parâmetro --module neuron_qwen3 importa o módulo de configuração local, e --config v5_qwen3_4b seleciona a função de configuração.
Carregamento e salvamento de checkpoints Hugging Face
Para fazer fine-tuning a partir de pesos publicados, basta habilitar o carregamento de checkpoint e apontar initial_load_path para o diretório de artefatos. O adaptador de dicionário de estado ajusta o prefixo model. e lida com embeddings compartilhados. O checkpoint final é salvo no formato safetensors do 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 configA loss no primeiro passo serve como verificação rápida de que os pesos publicados foram carregados corretamente. No run validado pela AWS, o Qwen/Qwen3-4B-Instruct-2507 produziu loss de 6,54 a partir dos pesos publicados e de 12,43 a partir de inicialização aleatória. Após o treinamento, o checkpoint exportado pode ser carregado com AutoModelForCausalLM e load_state_dict sem chaves faltando ou inesperadas — exceto lm_head.weight, que o Qwen3-4B vincula aos embeddings de token.
Configuração 2: Treinamento com os modelos nativos do torchtitan
Na segunda configuração, o torchtitan constrói seu próprio Qwen3Model nativo. Adicione as importações e a função de configuração ao neuron_qwen3.py:
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 configtorchrun --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_4bO campo config.model_spec seleciona a classe do modelo e define como o torchtitan fragmenta o modelo e mapeia as chaves do checkpoint. Aqui, model_registry("4B", attn_backend="sdpa") seleciona a implementação nativa Qwen3 com atenção por produto escalar (SDPA — Scaled Dot Product Attention).
Configuração 3: Kernels NKI para Mistura de Especialistas (MoE)
Na terceira configuração, as camadas de Mistura de Especialistas (MoE) são executadas com kernels escritos em NKI, a linguagem de kernels do Neuron. Adicione as importações e a função de configuração:
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 configtorchrun --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_a3bEsta configuração usa fragmentação de dados paralelos (data parallel sharding) em 64 ranks com grau de paralelismo tensorial igual a 1. Não utiliza paralelismo de especialistas. O parâmetro num_experts=64 define o número de especialistas no modelo — não o grau de paralelismo entre eles. Definir CompileConfig(enable=False) executa o exemplo em modo eager. O TorchNeuron suporta operadores NKI registrados tanto em execução eager quanto com torch.compile.
O exemplo reduz o Qwen3-30B-A3B para 4 camadas e 64 especialistas para que caiba em uma única instância trn2.48xlarge. Cada token é roteado para 4 especialistas com top_k=4.
Adicionando seu próprio kernel NKI
Em alguns casos, pode ser vantajoso otimizar uma operação para os formatos de tensor específicos do seu modelo. Um kernel especializado pode ajustar seu tiling e movimentação de dados a esses formatos, ou substituir uma sequência crítica de operações PyTorch por um único operador otimizado.
O guia demonstra esse mecanismo usando rope_hf da nkilib.core.embeddings, que espera o layout de tensor do Hugging Face e suporta seu próprio modo backward. Crie o arquivo my_neuron_kernels.py na raiz do repositório:
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(),
)O torchtitan armazena query e key como [batch, seq, heads, head_dim], enquanto o rope_hf espera [batch, heads, seq, head_dim]. A conversão de layout acontece dentro do operador. Registre também um operador separado para o modo backward, salvando o cache rotacional durante o forward para reutilização:
@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)Em seguida, crie a subclasse NKIRoPE estendendo CosSinRoPE e chame rope_fwd a partir de 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)Registre o NKIRoPE como substituto do CosSinRoPE usando o decorator @override:
@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 carregar o override customizado, passe o nome do módulo via --override.imports:
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_kernelsAntes do início do treinamento, o torchtitan registra cada nó de configuração substituído pelo override. Confirme que a saída contém uma entrada nki_rope para cada camada do 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.ConfigDefina o comprimento de sequência como múltiplo de 128 * LNC — ou seja, 256 quando LNC é 2. Reformate o cache rotacional para [seq, head_dim * 2], com valores cosseno na primeira metade e seno na segunda. O kernel valida essas restrições antes do início do treinamento.
Paralelismo e escalabilidade
As mesmas configurações de paralelismo se aplicam tanto às configurações Transformers quanto às nativas. Uma instância trn2.48xlarge expõe 64 NeuronCores lógicos com LNC=2. Definir data_parallel_shard_degree = -1 fragmenta os parâmetros do modelo, gradientes e estado do otimizador entre todos os ranks iniciados pelo torchrun. Definir config.parallelism.tensor_parallel_degree fragmenta os pesos de cada camada entre um grupo de ranks — cada rank computa uma fatia de cada multiplicação matricial. Com o paralelismo de sequência habilitado por padrão, as camadas de normalização operam em ativações fragmentadas ao longo da dimensão de sequência.
O torchtitan combina paralelismo tensorial com FSDP (Paralelismo de Dados Totalmente Fragmentado) em uma DeviceMesh multidimensional. Acompanhe o uso de HBM (Memória de Alta Largura de Banda) por rank ao aumentar o tamanho do modelo, comprimento de sequência ou tamanho de batch local — esses valores determinam a memória necessária por rank.
Limpeza do ambiente
Se quiser manter os checkpoints gerados, copie-os para armazenamento persistente antes de encerrar o ambiente. Se planeja reutilizar a instância, pare-a. Quando não precisar mais dela, encerre-a.
Conclusão
O AWS Trainium funciona como um dispositivo torchtitan suportado nas três configurações apresentadas. É possível treinar o Qwen3 via Hugging Face Transformers v5, carregando e salvando safetensors publicados diretamente. Também é possível treinar com a implementação nativa Qwen3 do torchtitan definindo config.model_spec. Para modelos com Mistura de Especialistas (MoE), as receitas NKI do torchtitan-neuron estão disponíveis. Para kernels customizados, o registro de overrides permite anexar operadores NKI ao mesmo fluxo de lançamento.
Para orientações de configuração, consulte a documentação do AWS Neuron. Para o código-fonte atual e receitas de treinamento, consulte o repositório torchtitan.
Fonte
Train Qwen3 on AWS Trainium with torchtitan (https://aws.amazon.com/blogs/machine-learning/train-qwen3-on-aws-trainium-with-torchtitan/)
Leave a Reply