Ai2 lança Olmo-Core 3, pilha de treinamento aberto para MoEs de trilhões de parâmetros

O Allen Institute for AI lançou Olmo-core 3 em 1 de outubro de 2026, uma atualização ao seu framework de desenvolvimento de grandes modelos de linguagem construído em torno de um sistema de treinamento de mistura de especialistas redesenhado, que a Ai2 disse ter sido benchmarkado em mais de um trilhão de parâmetros totais.

A Ai2 disse que o Olmo-core 3 foi projetado para escalar o treinamento de MoE para a faixa de trilhões de parâmetros, preservando a eficiência computacional, e o descreveu como um dos sistemas centrais por trás da próxima geração do Olmo. O lançamento é acompanhado por um relatório técnico, uma demonstração interativa e código aberto.

Modelos MoE podem conter muito mais parâmetros sem exigir que cada entrada use todos eles, mas o modelo completo ainda precisa ser armazenado na memória da GPU e atualizado durante o treinamento, e encaminhar entradas para os especialistas corretos em um cluster gera seus próprios custos de comunicação e coordenação. A Ai2 afirmou que esses custos podem corroer grande parte da vantagem computacional à medida que os MoEs crescem, e que o Olmo-core 3 foi desenvolvido para fechar essa lacuna. Em um benchmark da Ai2, o pool de especialistas aumentou de 8 para 128, mantendo a seleção de quatro especialistas por token, mantendo os parâmetros ativos aproximadamente fixos em cerca de 3.2 bilhões, enquanto a capacidade total de parâmetros cresceu de 4.6 bilhões para 47 bilhões, com a taxa de treinamento diminuindo menos de 5%.

De FSDP para uma pilha de treinamento baseada em DDP

A implementação anterior de MoE da Ai2 no Olmo-core utilizava paralelismo de dados totalmente fragmentado, configurado para coletar e redistribuir os pesos do modelo para cada pequeno lote de dados de treinamento. O Olmo-core 3 muda para um sistema baseado em paralelismo de dados distribuído que mantém os especialistas residentes nas GPUs e encaminha os dados relevantes para eles, evitando a coleta repetida de pesos. Em um teste preliminar em oito GPUs NVIDIA B300, a Ai2 relata que um MoE de 47 bilhões de parâmetros processou 52.000 tokens por segundo por GPU com a nova pilha, comparado com 19.400 usando a implementação anterior, o que a Ai2 descreve como cerca de 2,7 vezes maior taxa de transferência.

O trabalho da Ai2 em modelos esparsos remonta ao OlmoE, que utilizava uma arquitetura MoE com 64 especialistas roteados, enquanto o Olmo 3 usava uma arquitetura densa em que quase todo o modelo estava ativo para cada token. O Olmo-core 3 amplia o framework com um sistema de treinamento projetado para modelos MoE muito maiores. A Ai2 observou que o Megatron-Core da NVIDIA é uma opção consolidada para treinar MoEs de grande escala, e descreveu o Olmo-core 3 como trazendo uma pilha de treinamento MoE integrada ao framework por trás do Olmo.

Paralelismo, Precisão e Técnicas de Memória

Três técnicas determinam como o modelo e seu estado de treinamento são distribuídos entre o hardware: paralelismo de especialistas espalha especialistas pelas GPUs, paralelismo de pipeline divide as camadas do modelo entre grupos de GPUs e um otimizador distribuído espalha o estado do otimizador pelas GPUs ao invés de armazenar uma cópia completa em cada GPU. O Olmo-core 3 também reduz o custo de encaminhar dados para os especialistas corretos: paralelismo de especialistas linha a linha coloca os dados roteados diretamente nos buffers de entrada dos especialistas, o roteamento residente na GPU mantém os metadados de roteamento nas GPUs para que a CPU possa enfileirar tarefas sem aguardar que sejam copiados de volta, e GEMM agrupado combina muitas pequenas computações de especialistas para que as GPUs as executem de forma mais eficiente. O relatório técnico descreve um design de paralelismo de especialistas linha a linha baseado em NVSHMEM que grava cada token diretamente no buffer do especialista, com multiplicações de matrizes agrupadas agendadas pelo dispositivo que tornam o paralelismo de especialistas totalmente livre de sincronização.

O Olmo-core 3 oferece suporte ao MXFP8, um formato numérico de baixa precisão. Em um benchmark controlado em quatro GPUs NVIDIA B300 com trabalho distribuído uniformemente entre os especialistas, a Ai2 relata que a taxa de treinamento foi cerca de 21% maior que com a baseline BF16, enquanto a memória ativa máxima caiu de 103 GiB para 95 GiB, com a maior parte do ganho vindo da computação feed-forward e da movimentação de dados entre especialistas. O relatório acrescenta que checkpoints agnósticos à topologia registram tensores FP32 globais independentemente do layout paralelo, permitindo que uma execução seja retomada em uma topologia diferente, e que, em sua comparação controlada, a recomputação de ativações removeu quase dois terços da memória de ativação e reduziu a memória de pico em cerca de um quarto ao custo de aproximadamente um quinto da taxa de transferência.

Benchmarks de Trilhão de Parâmetros e Limites Declarados

A Ai2 afirmou que avaliou o Olmo-core 3 em uma variedade de configurações nas GPUs NVIDIA B300, incluindo um modelo de 1,2 trilhão de parâmetros com 58,36 bilhões de parâmetros ativos por token em 512 GPUs, onde a maior taxa observada foi de 858 TFLOP/s por GPU. A Ai2 declarou que esses testes usaram roteamento aleatório para medir o desempenho do sistema, e não a qualidade de um modelo treinado. O relatório técnico lista pontos de operação medidos de um modelo de 12,9 bilhões de parâmetros em 16 GPUs até o modelo de 1,2 trilhão de parâmetros em 512, e afirma que as taxas de destaque são referências de sistemas medidas sob roteamento aleatório, não uma curva de escala controlada ou uma reivindicação de tempo-para-qualidade.

Ai2 também experimentou o DeepEP v2, um backend alternativo para comunicação entre especialistas em GPUs, alcançando uma configuração com 2,38 trilhões de parâmetros totais. Ai2 descreve esse resultado como um teste de curta capacidade que demonstra a escala que o Olmo-core 3 pode atingir, em vez de desempenho de treinamento sustentado. O relatório técnico, intitulado “Supercharging Olmo-core for Efficient and Scalable MoE Training”, tem data de outubro de 2026; seus autores são do Allen Institute for AI e da University of Washington, com Tianhua Tao listado como autor principal e desenvolvedor principal.

Descobertas Documentadas e Planos para a Próxima Geração do Olmo

O relatório documenta um padrão de falha que a Ai2 chama de manipulação de tokens (token gerrymandering), no qual uma pontuação destinada a incentivar o roteamento equilibrado poderia melhorar mesmo quando a carga de trabalho real se tornava menos equilibrada. Outras descobertas documentadas incluem: reduzir as taxas de aprendizado dos especialistas porque eles processam menos tokens não melhorou os resultados na família de modelos testada; os cálculos nas GPUs levaram tempos diferentes quando os valores processados mudaram, mesmo com as mesmas dimensões de matriz; e sobrepor comunicação e computação em fluxos de GPU separados nem sempre acelerou o treinamento e, em alguns testes, retardou a execução de ponta a ponta. O relatório também descreve abordagens testadas, mas não adotadas, incluindo a transferência de carga de ativações para a CPU que o link do host não pôde suportar e dois esquemas de sobreposição que competiam com o cálculo dos especialistas por recursos da GPU.

Ai2 afirmou que sua próxima geração do Olmo usará uma arquitetura MoE e que pretende que seja o Olmo mais capaz até agora, treinado em seu maior conjunto de dados e com a janela de contexto mais longa. A organização disse que o Olmo-core 3 é totalmente aberto, permitindo que pesquisadores e desenvolvedores o utilizem para treinar seus próprios MoEs, adaptá-lo a diferentes hardwares e experimentar roteamento, paralelismo e outras partes do sistema. O repositório Olmo-core no GitHub descreve o projeto como blocos de construção PyTorch para o ecossistema OLMo, possui licença Apache-2.0 e pode ser instalado a partir do código-fonte ou do PyPI como ai2-olmo-core.