Allen Institute libera infraestrutura aberta para treinar modelos de IA com trilhões de parâmetros
Olmo-core 3 é um sistema que permite treinar modelos gigantes de inteligência artificial do tipo mixture-of-experts mantendo a eficiência computacional, e está disponível gratuitamente para qualquer pesquisador usar
O Allen Institute liberou o Olmo-core 3, uma infraestrutura aberta que permite treinar modelos de inteligência artificial do tipo mixture-of-experts (arquitetura em que cada entrada ativa só alguns dos muitos componentes especializados do modelo) em escalas de até trilhões de parâmetros mantendo a eficiência computacional, disponível gratuitamente para pesquisadores.
O Allen Institute for AI (AI2), instituto de pesquisa sem fins lucrativos baseado em Seattle, liberou o Olmo-core 3, uma atualização importante da infraestrutura que usa para desenvolver seus próprios modelos de linguagem. O sistema foi redesenhado para treinar modelos do tipo mixture-of-experts (MoE, ou mistura de especialistas — uma arquitetura em que o modelo contém muitos componentes especializados, mas cada entrada de texto ativa apenas alguns deles) em escalas que chegam à casa dos trilhões de parâmetros (os valores ajustáveis que um modelo aprende durante o treinamento). Segundo o instituto, em um teste comparativo, aumentar o número de especialistas disponíveis de 8 para 128 — mantendo fixo em quatro o número de especialistas ativos por token (as pequenas unidades de texto que um modelo processa) — fez a capacidade total do modelo saltar de 4,6 bilhões para 47 bilhões de parâmetros, enquanto a velocidade de treinamento caiu menos de 5%.
Treinar modelos grandes de IA exige muito poder computacional, o que encarece o processo e coloca o desenvolvimento de modelos avançados fora do alcance de muitos laboratórios acadêmicos e equipes menores. Modelos MoE são, em tese, mais eficientes — podem conter muito mais parâmetros sem que cada entrada precise ativar todos eles. Mas o modelo completo ainda precisa ficar armazenado na memória das GPUs (os processadores gráficos usados para treinar IA) e ser atualizado durante o treinamento, e distribuir entradas pelos especialistas certos em um cluster de GPUs cria custos próprios de comunicação e coordenação. Conforme os MoEs crescem, esses custos podem anular boa parte da vantagem computacional de usar apenas uma fração do modelo por entrada.
A versão anterior do sistema, usada no modelo OlmoE, dependia de uma técnica chamada FSDP (fully sharded data parallelism, ou paralelismo de dados totalmente fragmentado), que reunia e redistribuía os pesos do modelo (os parâmetros aprendidos) a cada pequeno lote de dados de treinamento. O Olmo-core 3 mudou para um sistema baseado em DDP (distributed data parallelism, ou paralelismo de dados distribuído), que mantém os especialistas residentes nas GPUs e roteia os dados relevantes até eles, evitando essa reunião repetida de pesos. Em um teste preliminar com oito GPUs NVIDIA B300, um MoE de 47 bilhões de parâmetros processou 52 mil tokens por segundo por GPU com a nova infraestrutura, contra 19,4 mil na implementação anterior — cerca de 2,7 vezes mais rápido, informou o instituto em seu blog oficial.
O sistema combina várias técnicas para distribuir modelos grandes entre GPUs: paralelismo de especialistas (que espalha os especialistas por diferentes GPUs), paralelismo de pipeline (que divide as camadas do modelo — os estágios sucessivos que transformam uma entrada — entre grupos de GPUs) e um otimizador distribuído (que espalha os dados adicionais usados para calcular atualizações durante o treinamento, em vez de manter uma cópia completa em cada GPU). Além disso, usa otimizações como roteamento residente em GPU (que mantém os metadados de roteamento nas próprias GPUs, permitindo que a CPU (o processador central) organize o trabalho sem esperar que essas informações sejam copiadas de volta) e suporte a MXFP8 (um formato numérico de menor precisão que representa alguns valores com menos bits, reduzindo computação e movimento de dados). Em testes controlados, habilitar MXFP8 nas partes do sistema onde mais ajudava aumentou a velocidade de treinamento em cerca de 21% em relação ao formato BF16 de maior precisão, enquanto o pico de memória ativa caiu de 103 para 95 gigabytes.
O Allen Institute testou o Olmo-core 3 em diversas configurações, incluindo um modelo de 1,2 trilhão de parâmetros com 58,36 bilhões de parâmetros ativos por token rodando em 512 GPUs, que atingiu 858 TFLOP/s/GPU (uma medida de computação útil do modelo por segundo em cada GPU). Também fizeram um teste de capacidade curto — não um treinamento completo — com 2,38 trilhões de parâmetros totais. A infraestrutura será usada na próxima geração do Olmo, que adotará arquitetura MoE, e está disponível abertamente no repositório do instituto no GitHub (uma plataforma onde desenvolvedores compartilham código) junto com um relatório técnico detalhado e uma demonstração interativa que mostra como as técnicas de paralelismo funcionam juntas.
