A NVIDIA publicou um tutorial sobre como treinar modelos de fundação biológica com arquitetura de especialistas mistos, ou MoE, usando a receita BioNeMo e a biblioteca Transformer Engine (TE). O foco é reduzir gargalos típicos desse tipo de modelo, como excesso de chamadas separadas e uso alto de memória.
O texto explica que, em vez de fazer cada token passar por todos os blocos, um modelo MoE ativa só alguns especialistas por token. Para tornar isso mais eficiente, a TE usa GroupedLinear para processar vários especialistas em uma única chamada e um kernel GroupedMLP que junta quantização, ativação SwiGLU e a escala aplicada às saídas do roteamento.
O tutorial também mostra o uso de MXFP8, formato de 8 bits com fator de escala por bloco de 32 valores. Segundo o material, em GPUs NVIDIA Blackwell essa representação é acelerada por hardware. Mesmo assim, os pesos principais continuam em 16 bits, com etapas de quantização e dequantização antes e depois do cálculo em baixa precisão.
Em um benchmark interno com oito GPUs NVIDIA B200 Tensor Core, a receita BioNeMo chegou a até 2,21x da vazão da implementação de referência da Hugging Face. A publicação ainda indica um teste inicial com duas GPUs e depois a escala para a configuração Mixtral-8x7B.