Olá, é um prazer inenarrável estar mais uma vez escrevendo este post. A Allen Institute for AI anunciou o Olmo-core 3, uma nova versão de sua infraestrutura aberta para treinar modelos de linguagem com arquitetura mixture-of-experts (MoE). A proposta é ambiciosa: permitir que pesquisadores trabalhem com modelos de até trilhões de parâmetros sem que o custo de computação cresça na mesma proporção.
O desafio é que modelos MoE não ficam eficientes automaticamente. Embora cada token seja encaminhado apenas para alguns especialistas, todos eles precisam ser distribuídos, armazenados e atualizados em uma rede de GPUs. É justamente nessa parte menos visível — a comunicação entre máquinas — que boa parte do ganho pode desaparecer. O Olmo-core 3 foi desenvolvido para reduzir esse desperdício.
Olmo-core 3 aumenta a capacidade dos modelos sem elevar tanto o custo
Em um dos testes divulgados, a equipe aumentou o conjunto de especialistas de 8 para 128, mantendo apenas quatro especialistas ativos para cada token. Com isso, a capacidade total do modelo passou de 4,6 bilhões para 47 bilhões de parâmetros, enquanto a quantidade de parâmetros ativos por token permaneceu próxima de 3,2 bilhões.
Mesmo com esse salto, a queda no throughput de treinamento foi inferior a 5% no benchmark apresentado. Em termos simples, o modelo passou a ter muito mais capacidade total sem exigir que cada entrada ativasse toda essa capacidade. É como ampliar uma oficina com dezenas de especialistas, mas chamar apenas os profissionais necessários para cada tarefa.
A AllenAI também afirma ter testado a infraestrutura em configurações acima de um trilhão de parâmetros. Esses números demonstram a escala que o sistema pretende alcançar, mas não devem ser confundidos automaticamente com a qualidade de um modelo treinado em produção.
Olmo-core 3 troca a forma de distribuir os especialistas entre as GPUs
Uma das principais mudanças está na arquitetura usada para o treinamento. A implementação anterior do Olmo-core utilizava fully sharded data parallelism (FSDP), reunindo e redistribuindo os pesos do modelo a cada pequeno lote de dados. O Olmo-core 3 passa a usar uma abordagem baseada em distributed data parallelism (DDP), mantendo os especialistas residentes nas GPUs e encaminhando os dados diretamente para eles.
Essa mudança evita parte do trabalho repetitivo de reunir os pesos antes de cada lote. Em um teste preliminar com oito GPUs NVIDIA B300, um modelo MoE de 47 bilhões de parâmetros processou 52 mil tokens por segundo por GPU com a nova infraestrutura. A implementação anterior alcançou 19,4 mil tokens por segundo por GPU, uma diferença de aproximadamente 2,7 vezes.
O resultado é uma comparação específica entre duas implementações e não uma promessa universal de desempenho. A velocidade pode variar conforme o hardware, a configuração da rede, o modelo e a distribuição dos dados.
Olmo-core 3 combina paralelismo e otimizações para o treinamento MoE
Para distribuir modelos grandes, o Olmo-core 3 combina diferentes formas de paralelismo. O paralelismo de especialistas divide os especialistas entre as GPUs; o paralelismo de pipeline distribui as camadas do modelo entre grupos de máquinas; e um otimizador distribuído divide o estado usado durante as atualizações.
Essa combinação reduz a quantidade de dados que cada GPU precisa manter localmente. Ao mesmo tempo, a infraestrutura tenta diminuir o custo de encaminhar cada token para o especialista correto. O sistema usa roteamento residente na GPU, organização direta dos dados nos buffers de entrada e operações agrupadas para que cálculos menores sejam executados de forma mais eficiente.
O framework também oferece suporte ao formato numérico MXFP8. Em um benchmark controlado com quatro GPUs NVIDIA B300, a equipe registrou throughput aproximadamente 21% maior do que usando BF16, enquanto a memória ativa de pico caiu de 103 GiB para 95 GiB. O ganho veio principalmente do cálculo nas camadas feed-forward e da movimentação de dados entre especialistas.
Há, porém, uma ressalva importante: usar menos bits não garante vantagem em qualquer cenário. A conversão entre formatos também tem custo, e a economia só aparece quando ela é menor do que o benefício obtido no cálculo e na comunicação.
Olmo-core 3 foi testado em configurações de até 2,38 trilhões de parâmetros
Nos testes descritos pela AllenAI, uma configuração com 1,2 trilhão de parâmetros foi executada em 512 GPUs NVIDIA B300, com 58,36 bilhões de parâmetros ativos por token. O maior throughput observado foi de 858 TFLOP/s por GPU, uma medida de computação útil por segundo em cada unidade.
A equipe também experimentou o DeepEP v2, uma alternativa para lidar com a comunicação entre especialistas, chegando a uma configuração de 2,38 trilhões de parâmetros totais. Esse teste foi de capacidade curta, e não uma execução completa de treinamento. Portanto, ele indica a escala que o sistema consegue atingir, mas não representa desempenho sustentado durante todo o treinamento de um modelo.
O relatório técnico ainda registra alguns resultados menos intuitivos. Uma métrica criada para incentivar o balanceamento do roteamento podia melhorar enquanto a carga real entre especialistas ficava mais desigual. Os pesquisadores chamaram esse problema de token gerrymandering. Além disso, aumentar a sobreposição entre comunicação e computação nem sempre acelerou o treinamento; em alguns testes, o efeito foi negativo.
Olmo-core 3 será a base aberta da próxima geração do Olmo
A AllenAI afirma que a próxima geração do Olmo usará uma arquitetura MoE, com um conjunto de dados maior e uma janela de contexto mais longa. Nesse cenário, o Olmo-core 3 funciona como a camada de infraestrutura que permitirá experimentar diferentes estratégias de roteamento, paralelismo e adaptação a novos hardwares.
O código está disponível no GitHub, e a organização também publicou um relatório técnico com os experimentos, as ablações e as decisões de engenharia. Há ainda um guia interativo que mostra como o paralelismo de dados, de especialistas e de pipeline pode ser combinado.
A importância do anúncio está menos em um novo modelo pronto para uso e mais na abertura da infraestrutura necessária para construí-lo. Treinar modelos gigantes continua exigindo hardware especializado e recursos consideráveis, mas disponibilizar as ferramentas pode reduzir a distância entre grandes laboratórios e grupos de pesquisa menores.
O que o Olmo-core 3 muda para a pesquisa em IA
O Olmo-core 3 mostra que escalar modelos MoE depende de muito mais do que adicionar parâmetros. É preciso coordenar memória, roteamento, comunicação, precisão numérica e execução distribuída como partes de um mesmo sistema. Uma melhoria isolada pode simplesmente transferir o gargalo para outra etapa.
Os resultados divulgados apontam para uma infraestrutura capaz de lidar com modelos de escala trilionária e de manter mais especialistas disponíveis sem aumentar proporcionalmente o trabalho realizado por token. Ainda assim, os benchmarks devem ser lidos com cuidado: alguns usam roteamento aleatório ou testes de capacidade, e não medem diretamente a qualidade de um modelo final.
No fim, a resposta para a pergunta inicial é simples, mas importante: tornar modelos MoE maiores não significa apenas ativar menos parâmetros em cada token. Significa construir uma rede de treinamento que consiga mover os dados e manter os especialistas ocupados sem desperdiçar o ganho de eficiência. É esse problema de engenharia que o Olmo-core 3 tenta resolver ao abrir sua infraestrutura para pesquisadores e desenvolvedores.
