NVIDIA apresenta técnicas para tornar mais eficiente o treinamento de modelos biológicos baseados em arquitetura MoE
A NVIDIA publicou em seu portal para desenvolvedores um tutorial voltado ao treinamento eficiente de modelos de base biológicos que utilizam a arquitetura Mixture of Experts (MoE), conhecida em português como mistura de especialistas. O material mostra como aplicar técnicas de otimização combinadas no recipe MoE do framework BioNeMo e na biblioteca Transformer Engine (TE).
O ponto de partida do tutorial é o desafio de escalar modelos densos. Em um transformer denso, cada token — isto é, cada unidade elementar de informação processada pelo modelo — passa por todas as camadas do modelo, de modo que adicionar novas capacidades aumenta o custo computacional tanto no treinamento quanto na inferência, que é a etapa em que o modelo já treinado gera respostas a partir de novos dados. Já as arquiteturas MoE adotam uma estratégia diferente: elas utilizam várias sub-redes, chamadas de especialistas, mas ativam apenas um pequeno subconjunto delas para cada token processado.
O recipe MoE do BioNeMo reúne três otimizações principais. A primeira é o GroupedLinear, que melhora o cálculo dos especialistas ao agrupar operações que, de outra forma, seriam executadas separadamente. A segunda é o formato MXFP8, um esquema de quantização em precisão de 8 bits que reduz o uso de memória durante o treinamento, em contraste com o BF16, formato que utiliza 16 bits para representar cada peso e ativação do modelo. A terceira é o kernel GroupedMLP, que funde em um único passo três operações: a quantização, a função de ativação SwiGLU — um tipo de função matemática que ajuda o modelo a decidir quais informações reforçar — e o escalonamento dos pesos de roteamento, responsáveis por direcionar cada token ao especialista mais adequado.
Segundo o tutorial, modelos MoE aumentam a capacidade total de parâmetros, enquanto cargas de trabalho em genômica costumam trabalhar com sequências longas, o que pressiona a memória de ativações durante o treinamento. A memória de ativações é o espaço necessário para armazenar os valores intermediários produzidos à medida que os dados atravessam as camadas do modelo. O uso de precisão reduzida com MXFP8 ajuda a aliviar esse gargalo.
O material traz ainda trechos de código que ilustram a aplicação prática das técnicas. Um dos exemplos utiliza a classe MXFP8BlockScaling para configurar a receita de quantização e instancia um modelo do tipo TEMixtralMXFP8ForCausalLM, uma variante otimizada da arquitetura Mixtral — biblioteca otimizada da arquitetura Mixtral — para tarefas de modelagem causal de linguagem, na qual o modelo aprende a prever o próximo elemento de uma sequência. Outro exemplo mostra o uso da API TE autocast, que ativa automaticamente a precisão MXFP8 nos passos de propagação para frente e para trás, justamente as fases em que o modelo calcula as previsões e em seguida atualiza seus parâmetros com base nos erros cometidos.
O tutorial indica que os leitores podem consultar o recipe do BioNeMo para obter o código completo e os detalhes de implementação. O conteúdo se posiciona como uma referência prática para pesquisadores e engenheiros que desejam treinar modelos de fundação biológicos baseados em MoE de forma mais eficiente, combinando maior capacidade de parâmetros com menor consumo de memória e de poder computacional.
Dessa forma, a publicação reforça o papel das otimizações de precisão e de agrupamento de operações como caminhos viáveis para lidar com o aumento de escala dos modelos aplicados à biologia computacional. Ao reunir GroupedLinear, MXFP8 e GroupedMLP em uma única cadeia de ferramentas, o recipe do BioNeMo com o Transformer Engine oferece um ponto de partida concreto para quem trabalha com sequências longas de proteínas ou genomas e precisa contornar os limites de memória impostos pelos formatos tradicionais de 16 bits.
