Novo buscador mostra quais modelos de IA rodam nos chips da Amazon
Hugging Face criou uma ferramenta que diz se um modelo de inteligência artificial funciona nos processadores Trainium da AWS, e entrega o código pronto para testá-lo.
A Hugging Face criou uma ferramenta de busca que mostra se um modelo de inteligência artificial roda nos chips Trainium da Amazon Web Services sem precisar de ajustes. O buscador lista os modelos mais populares testados, informa quantos núcleos do processador são necessários e entrega o script pronto para você rodar o modelo você mesmo em um servidor da AWS.
A Hugging Face (uma plataforma onde desenvolvedores compartilham e baixam modelos de IA já treinados) lançou uma ferramenta que responde a uma pergunta simples: o modelo de inteligência artificial que você quer usar funciona nos chips Trainium da Amazon? O novo buscador, chamado torch-neuronx model support Space, mostra quais dos modelos mais populares da plataforma rodam sem modificações nesses processadores — e entrega o script (pequeno programa de comandos) necessário para executá-los.
O buscador oferece três visualizações. A primeira mostra uma seleção própria da Hugging Face, com modelos escolhidos entre os mais baixados e os que estão em alta na plataforma. As outras duas seguem os rankings oficiais do site — por downloads ou por tendência — e deixam você buscar pelo nome do modelo. Cada linha da tabela informa se o modelo passou no teste, em quantos núcleos do Trainium ele foi distribuído, quanto tempo levou para rodar (incluindo a compilação, ou seja, a preparação do código para o processador) e quanta memória consumiu. Quando um modelo não foi testado, a ferramenta explica o motivo — geralmente porque não foi priorizado, às vezes porque é genuinamente incompatível.
Por enquanto, só modelos da biblioteca Transformers (um pacote de código aberto com modelos prontos para usar) foram testados, e apenas os voltados para geração de texto e para tarefas que misturam texto e imagem. A Hugging Face planeja expandir para outras bibliotecas e tipos de modelo. Cada teste carrega o modelo em uma instância trn2 (servidor com chip Trainium2, a versão mais recente do processador da Amazon), distribui os pesos (os parâmetros internos do modelo) pelo menor número possível de núcleos que consigam segurá-los e gera 32 tokens (pedaços de texto). Depois, confere se a saída faz sentido — comparando com uma execução em CPU (processador comum) ou usando um modelo auxiliar para verificar se o texto está fluente.
A ferramenta não só mostra o resultado: ela lista os pacotes que você precisa instalar e entrega o script run.py que foi usado no teste. Segundo informou a Hugging Face em seu blog, basta rodar python3 run.py em uma instância Trainium com o ambiente adequado — não é preciso usar torchrun (comando para rodar código distribuído), porque o próprio script cuida de distribuir o trabalho entre os núcleos quando necessário. Se quiser pular a etapa de compilação, você pode baixar os gráficos pré-compilados de um repositório público; o buscador dá o comando pronto.
Todos os dados do buscador vêm de um dataset (conjunto de dados) público, o aws-neuron/torch-neuronx-model-support, que você pode consultar diretamente com Python. Cada linha traz o resultado de um modelo, a mensagem de erro quando falhou, as versões de todos os componentes do sistema e o script que rodou. A documentação explica o significado exato de cada resultado — um PASS, por exemplo, só garante que o modelo carregou e gerou texto coerente, não diz nada sobre desempenho ou textos mais longos.
