NVIDIA escala o determinismo bit a bit no Megatron Core para 2.432 GPUs com apenas 2% de overhead

O problema do treinamento em larga escala e a solução da NVIDIA

No treinamento de modelos com trilhões de parâmetros em milhares de GPUs, as abordagens tradicionais de depuração deixam de funcionar. Múltiplas dimensões de paralelismo, computação de baixa precisão (MXFP8, NVFP4) e checkpoints distribuídos tornam muito mais difícil reproduzir bugs e validar correções. A NVIDIA implementou no Megatron Core o determinismo bit a bit — a garantia de que duas execuções iniciadas a partir do mesmo estado seguirão uma trajetória numérica idêntica, desde que os dados, a arquitetura, a configuração de paralelismo e o software permaneçam inalterados.

O Megatron Core oferece duas garantias fundamentais: reprodutibilidade independente (duas execuções iniciadas a partir do mesmo estado permanecem idênticas bit a bit em cada passo) e retomada exata a partir de checkpoint (a recuperação corresponde exatamente a uma execução contínua). Isso é crítico para treinamentos de produção de alto custo — em escala de trilhão de parâmetros, mesmo uma pequena lentidão pode custar milhares de dias de GPU.

Quatro vantagens práticas

O determinismo bit a bit resolve quatro problemas de depuração e validação:

• Reproduzir bugs — capturar um pico de loss em um passo específico e, em seguida, alterar um componente por vez para isolar a causa • Continuidade de checkpoints — retomar um treinamento interrompido sem alterar a trajetória numérica • Detecção de defeitos — usar discrepâncias entre execuções que, de outro modo, seriam idênticas para identificar corrupção de checkpoints ou problemas de software e hardware • Comparação do impacto de mudanças — diferenciar o efeito numérico de uma alteração intencional da variação natural entre execuções

Metodologia para depurar o não determinismo

No PR #7262 do Megatron-LM, a NVIDIA propôs um sistema de rastreamento progressivo que registra impressões digitais ordenadas de tensores entre os ranks. O processo estreita a busca, das métricas de ponta a ponta até um kernel específico:

• Métricas de ponta a ponta — identificar a primeira iteração em que uma métrica serializada diverge (loss de treinamento, loss de modelagem de linguagem, norma do gradiente, loss de MTP) • Rastreamento semântico amplo — abranger coletivos, comunicação entre pipelines, recomputação, operações do otimizador e reduções de gradiente para determinar em que fase do treinamento ocorre a divergência • Rastreamento de módulos e camadas — localizar o problema em um componente específico do modelo (camada Transformer, bloco MoE, estágio do otimizador) • Rastreamento em nível de operação — impressões digitais de operações ATen e chamadas de extensão para encontrar a primeira operação com entradas iguais e saídas diferentes • Ferramentas em nível de kernel — identificar o kernel responsável e o mecanismo causador do não determinismo

É fundamental não limitar o rastreamento apenas à iteração em que a divergência de loss se torna visível; o primeiro bit divergente pode surgir muito antes. O rastreamento amplo localiza a iteração e os ranks onde ocorre a divergência mais precoce, e só então o rastreamento detalhado é ativado, restrito a essa janela estreita.

As impressões digitais são calculadas na GPU usando torch.hash_tensor, registrando formato (shape), tipo de dado (dtype) e número de elementos. Para MXFP8 e NVFP4, tanto os valores codificados quanto os buffers de escala são registrados. Para dados sensíveis à ordem (mapas de roteamento, saídas de dispatch do MoE), é usado um rastreamento por blocos (chunk-wise).

Otimização para um Nemotron de um trilhão de parâmetros

Uma receita totalmente determinística, com overhead elevado, é útil para depuração, mas custa caro demais para produção. A NVIDIA realizou uma co-otimização dos caminhos determinístico e não determinístico em um modelo híbrido Nemotron — que combina camadas SSM (modelo de espaço de estados) no estilo Mamba com camadas de atenção Transformer —, otimizando três cargas de trabalho:

• Nemotron 3 Ultra: redução de ~17% para 2% após uma otimização de preenchimento de buffer, e depois para 1,5% em 3.072 GPUs • Proxy híbrido do Triton: redução de ~60% para 37% após a restauração da fusão MoE-MLP, e depois para 2% • Caminho otimizado de gradiente de peso com CuTeDSL: ~3,6% de overhead em 256 GPUs, e depois ~2% em 2.432 GPUs

Em 2.432 GPUs, a receita de produção atingiu cerca de 2% de overhead de determinismo em regime estável, mantendo o determinismo bit a bit ao longo de 800 passos de treinamento. As otimizações em nível de kernel incluem a alocação de slots de saída privados para os escritores de grouped-GEMM e a aplicação de uma ordem fixa de redução para preservar o paralelismo sem deixar de recuperar resultados determinísticos. O Megatron Core valida essas receitas, testa kernels e verifica módulos em todas as configurações de paralelismo, incluindo suporte a formatos de baixa precisão FP8 e FP4.