Nvidia corta custo de treino determinístico para 2% em 2.432 GPUs

Treinar modelos nessa escala consome milhares de dias de GPU e envolve muitas dimensões de paralelismo, computação em baixa precisão e salvamento distribuído de checkpoints. Nesses ambientes, duas execuções iniciadas com dados e parâmetros idênticos costumam divergir porque operações paralelas acumulam pequenas variações numéricas. A equipe da Nvidia concentrou a solução no pull request Megatron-LM PR #7262, que registra impressões digitais de tensores ordenadas por rank para localizar onde o não-determinismo começa. O sistema detecta discrepâncias no primeiro passo divergente.

Mudanças em kernels preservam ordem de cálculo

A Nvidia alterou a execução de kernels para manter o paralelismo sem bagunçar as somas. Os desenvolvedores deram posições privadas de saída para escritores de grouped-GEMM e fixaram a ordem de redução matemática. Isso impede que a chegada aleatória de threads modifique os bits finais dos tensores. A validação cobre configurações em formatos de precisão reduzida como FP8 e FP4. A estabilidade numérica exige disciplina rígida de hardware e software.

O artigo adverte que valores arredondados de loss escondem divergências graves em gradientes. Para validar o comportamento, o fluxo gera impressões digitais de métricas como loss de treino, loss de modelagem de linguagem, loss de balanceamento de carga, loss de previsão multi-token (MTP) e norma de gradiente. No caso do modelo Nemotron de trilhão de parâmetros, o caminho de checkpoint salva uma fonte da verdade em precisão total e reconstrói o estado em tempo de execução pelo mesmo método de conversão adotado no treino. Um teste pareado encontrou desvio logo no passo 377 ao checar métricas completas.

Validação depende de configurações idênticas

A garantia bit a bit exige que o usuário fixe dados, arquitetura, configurações de runtime, paralelismo e software. Se uma máquina trocar de versão ou mudar a ordem dos lotes, o alinhamento quebra. O Megatron Core promete duas frentes: reprodutibilidade independente entre dois treinos do zero e continuidade exata entre um treino contínuo e outro restaurado de checkpoints salvos. Os autores recomendam testar o estado do otimizador, os geradores de números aleatórios (RNG) e os parâmetros restaurados contra uma execução ininterrupta.

Os números apresentados pela Nvidia vêm de três cargas de trabalho da família Nemotron avaliadas pela própria fabricante. A empresa não detalhou no texto o impacto do método fora de seu ecossistema de hardware nem forneceu testes em configurações com GPUs de outras arquiteturas. Para aplicar a rotina, quem opera o Megatron Core precisa consultar o guia do usuário e validar as receitas suportadas contra checkpoints próprios. O código depende de validação constante de receitas e módulos para evitar regressões.