Google adiciona treinamento com contexto de milhões de tokens ao MaxText em TPUs

O Google anunciou suporte a treinamento com contexto longo no MaxText, seu framework de treinamento de código aberto baseado em JAX, em 9 de novembro de 2026, permitindo que sequências de treinamento cheguem a 14 milhões de tokens em chips Cloud TPU 5p e TPU 7 (Ironwood), segundo a publicação Announcing Long-Context Training on TPUs with MaxText Context Parallelism no Google Developers Blog.

O que está mudando

• O MaxText, framework de treinamento de código aberto em JAX do Google, agora oferece suporte a treinamento com contexto longo em Cloud TPUs.
• As sequências de treinamento escalam de forma contínua até 14 milhões de tokens nos chips TPU 5p e TPU 7 (Ironwood).
• Quatro estratégias de paralelismo de contexto são suportadas: All-gather, A2A (DeepSpeed Ulysses), P2P (Ring Attention) e um híbrido USP que combina A2A e P2P.
• Um kernel Tokamax Splash Attention gerencia o cômputo e a memória de atenção sem materializar a sequência completa de chave e valor na HBM.
• O Llama 3.1 8B atinge 70% de utilização de FLOPs do modelo (MFU) na TPU 5p para sequências de até 4 milhões de tokens, enquanto a TPU 7 mantém 50% de MFU em todas as escalas de sequência abaixo de 6 milhões de tokens.

O que isso significa para quem constrói com IA

O paralelismo de contexto divide uma sequência longa entre vários chips, de modo que nenhum chip isolado precisa manter as ativações de toda a sequência, o que é justamente o que viabiliza o treinamento nesses comprimentos. Para equipes que pré-treinam ou ajustam seus próprios modelos de linguagem grandes, isso oferece um caminho de código aberto e nativo de TPU para janelas de contexto de vários milhões de tokens, bem além das janelas de 1 a 2 milhões de tokens comuns nos modelos de produção atuais, sem precisar construir infraestrutura própria de treinamento com contexto longo.

Disponibilidade

• A funcionalidade está disponível no MaxText v0.2.4 no PyPI.
• O Google fornece receitas de treinamento de exemplo para configurações de TPU 7 e TPU 5p.

Fontes

• Announcing Long-Context Training on TPUs with MaxText Context Parallelism (Google Developers Blog, 9 de novembro de 2026, oficial)

Escrito com ajuda de IA, apenas a partir das fontes listadas acima. Antes de ir ao ar, verificações automáticas conferiram as datas, as fontes e todas as versões de idioma.