NVIDIA demonstra determinismo bit a bit no Megatron Core para pré-treinamento com milhares de GPUs

Treinar um modelo com trilhões de parâmetros em milhares de GPUs custa milhões de dólares e leva semanas. Se no meio da execução aparece um pico de loss ou um NaN, em geral não é possível reproduzi-lo: a ordem das operações muda de uma execução para outra, e o erro não se manifesta. A NVIDIA publicou no Developer Blog um material sobre como obter determinismo bit a bit no pré-treinamento com o Megatron Core. Duas execuções com os mesmos dados e a mesma configuração devem produzir resultados idênticos, até o último bit.

Por que o determinismo no pré-treinamento é tão raro

A raiz do problema está na aritmética de ponto flutuante: a adição nela não é associativa. Somar os mesmos números em outra ordem dá um resultado ligeiramente diferente. Em uma única GPU, isso quase não aparece. Em um cluster, a ordem resulta de vários fatores: operações atômicas nos kernels, a escolha do algoritmo de redução nas comunicações coletivas, o escalonamento de threads e a distribuição do trabalho entre os dispositivos.

A fonte lista os fatores que complicam o cenário:

• vários eixos de paralelismo ao mesmo tempo (dados, tensores, pipeline e outros), cada um com suas próprias reduções e trocas de dados;
• cálculos em baixa precisão, nos quais os erros de arredondamento são maiores e dependem mais da ordem das operações;
• checkpoints distribuídos, que precisam ser salvos e carregados de modo que a retomada do treinamento não altere a trajetória;
• escala: trilhões de parâmetros e milhares de GPUs, onde a probabilidade de uma falha rara ao longo da execução se aproxima de 1.

O trecho disponível não revela as técnicas específicas que o Megatron Core usa para garantir a ordem das operações, nem o custo do determinismo em velocidade. Esses dados devem ser consultados na versão completa do material.

O que a reprodutibilidade bit a bit oferece

O anúncio cita três benefícios práticos: depuração, validação e retomada reproduzível do treinamento.

O primeiro diz respeito à busca pelas causas de uma falha. Se a execução é determinística, é possível reiniciá-la a partir do último checkpoint anterior à anomalia e obter a mesma anomalia. Isso transforma um erro "fantasma" em um bug comum, sobre o qual dá para colocar verificações e fazer bisseção por camadas ou passos.

O segundo diz respeito à validação de correções. No treinamento não determinístico, duas execuções sempre divergem, e só é possível saber estatisticamente se uma correção influenciou o resultado. Com a reprodutibilidade bit a bit, qualquer divergência em relação à referência significa que a mudança afetou algo. A ausência de divergência prova que a refatoração ou a otimização não alterou o comportamento numérico.

O terceiro diz respeito às retomadas. Em clusters grandes, falhas de hardware acontecem com regularidade. Se, depois de carregar o checkpoint, o treinamento segue exatamente a mesma trajetória que teria seguido sem a interrupção, os engenheiros podem ter certeza de que o reinício não introduziu deriva oculta.

O preço e onde isso é útil

O determinismo costuma ser comprado com desempenho: kernels rápidos e não determinísticos e reduções atômicas muitas vezes são trocados por alternativas mais lentas, porém previsíveis. Por isso, a pergunta importante é quanto a vazão cai em configurações reais. O trecho analisado não traz a resposta, e sem medições é cedo para avaliar o ganho.

Mesmo assim, a lógica de uso é clara. Não é necessário manter o modo determinístico o tempo todo nas execuções mais caras. Mas ele é útil em alguns cenários:

• na reprodução e análise de falhas, quando é preciso uma repetição exata;
• na verificação de novos kernels, na troca de precisão (por exemplo, a migração para uma precisão mais baixa) ou na mudança do esquema de paralelismo;
• em testes de regressão do código de treinamento: uma execução de referência com poucos passos, com a qual cada build é comparado bit a bit.

Conclusão prática para as equipes

Equipes que treinam modelos próprios com o Megatron Core devem criar uma execução de controle curta em modo determinístico e incluí-la no CI. Qualquer mudança no código, nos drivers ou na configuração de paralelismo é comparada primeiro com a referência, por meio de hashes dos pesos após um número fixo de passos. Uma divergência mostra imediatamente que a mudança afetou o comportamento numérico.

Antes da adoção, é preciso medir o overhead no seu cluster e na sua precisão. Também é preciso verificar se a retomada a partir do checkpoint realmente produz uma trajetória idêntica na sua topologia. O determinismo completo deve ser ativado de forma pontual: na depuração, na validação e na auditoria de execuções. O treinamento principal pode rodar no modo rápido, se as medições mostrarem uma perda de velocidade perceptível.