Google lança agentes para migrar PyTorch para JAX em TPUs

O que é o projeto

O repositório "Accelerator Agents", hospedado na organização AI-Hypercomputer no GitHub, reúne ferramentas com IA para acelerar o desenvolvimento de machine learning em TPUs do Google Cloud. O projeto disponibiliza agentes que auxiliam na migração de código, na otimização de kernels e no ajuste de desempenho, com o objetivo de permitir que desenvolvedores aproveitem as TPUs com mais velocidade.

O README traz um aviso importante: trata-se de um projeto não oficialmente suportado pelo Google, e ele não é elegível para o programa de recompensas por vulnerabilidades de código aberto da empresa (Google Open Source Software Vulnerability Rewards Program). O código é distribuído sob a licença Apache 2.0, com 67 estrelas e 17 forks registrados na página.

MaxCode: migração de PyTorch para JAX

O primeiro agente, chamado MaxCode, facilita a conversão de modelos e bases de código existentes em PyTorch para JAX, ajudando usuários a migrar suas cargas de trabalho para rodar de forma eficiente em TPUs, com uso de estruturas de alto desempenho como o MaxText. Suas funções incluem:

• Conversão automatizada de blocos de código funcionais e camadas de modelo de PyTorch para JAX;
• Integração com MaxText, gerando código JAX compatível com a estrutura para treinamento e inferência imediatos em TPUs;
• Abordagem human-in-the-loop: o agente redige implementações iniciais que os desenvolvedores podem revisar e refinar.

O projeto ressalta que o MaxCode está em desenvolvimento ativo, com trabalho contínuo para melhorar a qualidade da migração e ampliar a cobertura de modelos.

MaxKernel: desenvolvimento de kernels de alto desempenho

O segundo agente, MaxKernel, é uma ferramenta especializada para o desenvolvimento de kernels de alto desempenho em TPUs. Ele auxilia engenheiros a escrever, otimizar e depurar kernels personalizados, com foco específico no Pallas, a linguagem de kernels do JAX. As funcionalidades divulgadas são:

• Redação de kernels Pallas do zero ou a partir de implementações de referência em JAX;
• Conversão de kernels personalizados de CUDA/GPU para rodar de forma otimizada em TPUs;
• Otimização e profiling, com sugestões e diagnósticos para melhorar o desempenho dos kernels (MFU);
• Geração automática de harness de testes, criando código boilerplate para verificação de corretude e checagens de compilação.

Requisitos e contribuições

Para começar, o projeto exige uma máquina virtual no Google Cloud: uma VM apenas com CPU é suficiente para o MaxCode, enquanto uma TPU VM é recomendada para o MaxKernel. Também são necessários Python 3.11 ou superior e acesso à API do Gemini, usada para as capacidades de raciocínio dos agentes. A instalação começa com a clonagem do repositório, e instruções específicas de cada agente ficam em seus respectivos subdiretórios.

O README convida contribuições externas, com diretrizes detalhadas no arquivo CONTRIBUTING.md, e também referencia um código de conduta. Além dos dois agentes principais, a estrutura do repositório inclui um diretório chamado JAXBench, sem descrição detalhada no texto apresentado.