Treinar o Qwen3 no AWS Trainium com torchtitan
aws ·
Este artigo demonstra como treinar o modelo Qwen3 no hardware AWS Trainium usando o framework torchtitan. Os autores descrevem três métodos de configuração distintos: utilizar o Hugging Face Transformers v5 como backend, empregar a implementação nativa do Qwen3 no torchtitan e integrar kernels da Interface de Kernel Neuron (NKI) para modelos mistos de especialistas. Todas as abordagens são executadas em uma única instância trn2.48xlarge com 64 NeuronCores lógicos, aproveitando interfaces padrão do PyTorch por meio do TorchNeuron. O guia detalha pré-requisitos específicos, incluindo versões de software necessárias e configuração do ambiente usando a AWS Deep Learning AMI. São fornecidos exemplos de código para ajuste fino a partir de pesos publicados, personalização de mecanismos de atenção e anexação de kernels NKI definidos pelo usuário por meio de um registro de substituição. A explicação passo a passo enfatiza a verificação da estabilidade do treinamento comparando métricas de perda e garantindo a compatibilidade dos checkpoints com os formatos do Hugging Face. Por fim, explica estratégias de paralelismo, como sharding paralelo de dados e paralelismo de tensores, para otimizar a utilização de recursos nos aceleradores.
aws