Conheça o TorchNeuronx: Pesquisa De Modelos No AWS Trainium

Nas postagens anteriores, apresentamos o TorchNeuron, o backend nativo em PyTorch para o AWS Trainium (o pacote torch-neuronx. Com ele, os modelos da família Transformers executam seu código de modelagem no Trainium, assim como fazem em qualquer outro dispositivo, permitindo que qualquer modelo do Hub rode.

Mas quais modelos de fato se beneficiam desta capacidade? O recurso Space, disponível na visão de Selecionados, responde a essa pergunta com exatidão. Ele indica para quais dos modelos mais populares do Hub o suporte foi implementado e fornece o script necessário para rodá-los.

Como Usar O Recurso Space

O Space abre na visão de Selecionados, onde disponibilizamos uma curadoria própria dos modelos mais baixados e com melhor classificação de tendência do Hub. Cada linha apresenta o resultado obtido, o número de unidades Trainium utilizadas, o tempo total da execução (incluindo compilação) e a memória máxima ocupada no dispositivo.

Para consultar um modelo específico, mude para a visão de Tendências ou Downloads e digite parte do nome. Essas visões seguem as classificações originais do Hub, executando uma tarefa por vez e indicando explicitamente se o modelo foi testado ou não. Quando um modelo foi deixado de fora por design, o Space justifica a ausência — a maior parte das vezes não foi por incompatibilidade técnica, mas sim pela falta de prioridade no curadoria.

Também é possível selecionar as versões específicas do compilador Neuron (neuronx-cc) e da biblioteca Transformers para comparar resultados entre versões diferentes.

Como Funciona O Processo De Teste

Ao selecionar um modelo, o Space exibe todos os componentes necessários para reproduzir a execução, incluindo a pilha completa de testes, os pacotes adicionais a instalar e o script run.py propriamente dito. O fluxo segue a receita publicada no anúncio do TorchNeuron: carrega o modelo com from_pretrained, particiona quando necessário e gera texto com cache estático e torch.compile(backend="neuron").

Na prática, não é preciso usar o torchrun: o script próprio já gerencia as faudas de torch.distributed automaticamente quando o modelo é particionado.

Uma primeira execução sempre começa com a compilação das grafias do modelo. Para pular essa etapa, é possível copiar as grafias construídas pelo mesmo compilador durante o teste, extraídas do bucket público NEFF. O Space fornece o comando exato, incluindo a versão correta do compilador.

Critérios De Inclusão

Cada rodada de teste carrega o modelo em uma instância trn2 particionado sobre o menor número possível de unidades Trainium capazes de armazenar suas weights. A maioria dos modelos usa particionamento tensorial; modelos com arquiteturas Mixture-of-Experts que não possuem plano de particionamento tensorial usam particionamento exclusivo (expert parallelism).

O teste então gera 32 tokens e os verifica. Em uma única unidade, os tokens são comparados com aqueles gerados por uma execução idêntica no CPU do mesmo checkpoint. Quando o modelo está particionado, uma pequena modela julgadora garante que o texto saído seja fluente em inglês.

Até o momento, apenas modelos que couberem em 4 dispositivos Trainium2 são testados, o que equivale a 16 núcleos no máximo.

Base De Dados Publica

Todas as informações exibidas pelo Space derivam do conjunto de dados público aws-neuron/torch-neuronx-model-support, estruturado como um arquivo Parquet para cada par de versões neuronx-cc e biblioteca. Cada linha contém o resultado do modelo, a mensagem de erro caso houvesse falhas, as versões completas da pilha e o próprio script run.py utilizado.

É possível consultar o banco de dados diretamente através do card publicado no repositório. O card fornece o significado exato de cada resultado obtido.

Próxima vez que questionar se seu modelo roda no Trainium, vale a pena consultar o Space. É simples, transparente e dá toda a confiança necessária para tomar a decisão certa.

Com informações de: Hugging Face Blog