Análise aprofundada · Cerca de 7 min de leitura
Qwen3 no Trainium: três caminhos com torchtitan para verificar pesos e kernels
O guia da AWS de 21 de outubro integra Transformers, Qwen3 nativo e kernels NKI em uma instância Trainium2, sem demonstrar vantagem de velocidade ou custo.
- Publicado aqui
- Data da fonte
- Fonte verificada
Um guia técnico da AWS de 21 de outubro explica como treinar Qwen3 com torchtitan no Trainium usando três configurações: um backend Transformers v5, o modelo nativo do torchtitan e um exemplo reduzido de mistura de especialistas com kernels NKI. O avanço útil é uma receita concreta de integração, incluindo verificações de checkpoints e de resultados numéricos. Não é o lançamento de um novo modelo Qwen nem um benchmark que comprove treinamento mais barato.
Ouça as notícias
Ouça este artigo ou continue ouvindo as notícias.
Preparando a leitura…
Voz do dispositivo
O navegador ou dispositivo processa a voz. O AISubHub não envia o texto a uma API de áudio. As vozes e a reprodução com a tela bloqueada variam conforme o dispositivo.
Ao continuar, uma frase curta pode se repetir.
Resumo
- O guia utiliza uma trn2.48xlarge: 16 aceleradores Trainium2, 128 NeuronCores físicos e 64 lógicos com LNC=2. A inicialização abre 64 processos de treinamento.
- Os exemplos densos usam os arquivos do Qwen3-4B-Instruct-2507 e dez passos com c4_test. O exemplo MoE é reduzido a quatro camadas e 64 especialistas, com quatro especialistas selecionados por token; não treina o 30B-A3B completo e sem alterações.
- A AWS informa perdas no primeiro passo de 6.54 com pesos publicados e 12.43 com inicialização aleatória. Seu teste de embeddings rotativos registra diferença absoluta máxima de 0.0 com entradas float32; essas verificações tratam de questões diferentes.
Uma integração de dispositivo, três pontos para mudar o modelo
O TorchNeuron registra Trainium como um dispositivo neuron nativo no PyTorch, permitindo que o torchtitan utilize interfaces habituais. O primeiro caminho deixa o Transformers v5 construir o modelo por AutoModelForCausalLM. O segundo troca essa implementação pelo Qwen3 nativo do torchtitan por meio de config.model_spec. O terceiro utiliza operadores NKI específicos do Trainium nas camadas de mistura de especialistas e demonstra uma substituição personalizada. São escolhas de implementação sob uma mesma entrada de treinamento, e não três modelos Qwen recém-lançados.
O guia fixa torch 2.14.0, torch_neuronx 2.14.3, Transformers 5.16.0 e torchtitan-neuron release-2.33, versão 0.2.2+neuron.1. Ele exige uma AMI de aprendizado profundo para Neuron ou um contêiner de treinamento com driver, runtime e compilador. Nossa leitura é que código PyTorch conhecido reduz a superfície de integração; isso não torna opcional a combinação de dependências testada. A AWS avisa expressamente que a integração do Transformers e os módulos de substituição podem mudar entre versões do torchtitan.
A receita 4B verifica inicialização e portabilidade do checkpoint
A primeira configuração especifica Qwen/Qwen3-4B-Instruct-2507, utiliza c4_test, define dez passos e lote local de um, e fragmenta o estado do paralelismo de dados entre os 64 ranks iniciados. Carregar a arquitetura é diferente de carregar pesos publicados. A configuração de ajuste fino ativa explicitamente o carregamento do checkpoint a partir dos arquivos baixados e exporta o modelo final em Hugging Face safetensors. A perda reportada no primeiro passo é 6.54 com pesos publicados, contra 12.43 com inicialização aleatória.
Essa diferença é uma verificação útil de inicialização, não evidência de que dez passos produzam um assistente melhor ou estabeleçam convergência final. A AWS também recarrega o checkpoint exportado com Transformers para verificar suas chaves, considerando o lm_head.weight vinculado aos embeddings. O caminho nativo começa no construtor qwen3_32b porque a versão testada não oferece um construtor genérico 4B; depois substitui a especificação do modelo por 4B e define lote global 64. Portanto, o nome do construtor não significa que o exemplo treinou um modelo de 32B.
A configuração MoE é deliberadamente menor que sua família
A receita Qwen3-30B-A3B é reduzida a quatro camadas e 64 especialistas para caber em uma instância. top_k=4 encaminha cada token a quatro especialistas. A fonte usa fragmentação do paralelismo de dados em 64 ranks, grau de paralelismo tensorial um e nenhum paralelismo de especialistas. O número de especialistas é uma configuração interna do modelo, não a quantidade de dispositivos de paralelismo de especialistas. Separar esses significados evita transformar uma pequena demonstração de integração em uma alegação de treinamento do modelo original completo.
O exemplo roda em modo eager, com compilação desativada na configuração. Operadores NKI registrados também funcionam com torch.compile, mas trocar o modo de execução é uma decisão separada. Na versão testada, a receita substitui o bloco MoE por kernels de especialistas agrupados e seleciona atenção por produto escalar escalado e embedding rotativo sem gather. A AWS explica como os operadores recebem autograd com informações salvas na passagem direta. Não apresenta uma comparação de vazão que demonstre que essa configuração específica é mais rápida.
Um kernel personalizado precisa de verificações de layout, gradiente e destino
O exemplo de embeddings rotativos mostra por que registrar vai além de chamar um kernel. O torchtitan armazena query e key como lote, sequência, cabeças e dimensão da cabeça; rope_hf espera lote, cabeças, sequência e dimensão da cabeça. O encapsulamento transforma o layout, fornece espaço de saída e desfaz a transformação. Um operador de retropropagação separado reutiliza o cache rotativo salvo, enquanto o próprio cache não tem gradiente. O comprimento de sequência deve ser múltiplo de 128 vezes LNC: 256 quando LNC é dois.
Um registro de substituições então troca a configuração do módulo pretendido. Os logs de inicialização devem mostrar uma substituição nki_rope para cada camada Qwen3; padrões de nome mais restritos podem selecionar determinadas camadas. A AWS informa diferenças absolutas máximas de 0.0 para query e key em uma comparação float32 com a implementação PyTorch. Isso valida aquele teste numérico, não todos os tipos de dados, formatos ou treinamentos. O artigo apresenta explicitamente o exemplo como integração do kernel com autograd, e não como melhoria de desempenho.
Transformar o guia em um piloto medido
O primeiro passo inclui compilação e demora mais que os seguintes. Os logs expõem perda, norma do gradiente, uso de memória de alta largura de banda, vazão e utilização de FLOPs do modelo, mas o guia não fornece comparação controlada de custo completo. Conforme crescem o modelo, o comprimento de sequência ou o lote local, a memória por rank continua sendo uma restrição prática. O paralelismo tensorial pode ser combinado com paralelismo de dados totalmente fragmentado em um DeviceMesh multidimensional; a configuração MoE mostrada não deve ser confundida com esse experimento maior.
O AISubHub sugere preservar as versões testadas, verificar o modelo selecionado e os 64 ranks, conferir carregamento e exportação do checkpoint e comparar saídas do kernel personalizado antes de iniciar uma execução longa. Depois, medir o aquecimento separadamente da execução sustentada e avaliar convergência nos dados pretendidos. Salve checkpoints necessários antes de parar ou encerrar o hardware alugado. São etapas de avaliação inferidas do guia, não resultados que reproduzimos, e não estabelecem direito de assinatura, franquia gratuita de treinamento ou vantagem universal do Trainium.
Por que importa
Análise do AISubHub: as equipes podem separar compatibilidade do modelo, portabilidade de checkpoints e correção dos kernels em tarefas menores antes de tentar um treinamento completo. A estrutura comum de inicialização reduz alterações entre receitas, mas não elimina requisitos de versões, memória, compilação ou comportamento numérico.
Limites e disponibilidade
Este é um guia técnico de Josh Longenecker e Alex Hoffman, da AWS, não um teste de hardware do AISubHub. Ele não fornece preço comparável de treinamento completo, resultado final de convergência nem aceleração comprovada do kernel personalizado. Disponibilidade e cobrança da instância, versões exatas e validação da carga continuam sob responsabilidade do operador. Não há anúncio de benefício de assinatura para consumidores.
Fontes e evidências
Data da fonte: · Fonte verificada:
Conteúdo editorial com apoio de IA, verificado em fontes primárias. Alegações dos fornecedores são atribuídas; inclusão não é recomendação.