NVIDIA escala el determinismo bit a bit en Megatron Core a 2.432 GPU con solo un 2% de sobrecarga

El problema del entrenamiento a gran escala y la solución de NVIDIA

Al entrenar modelos con billones de parámetros en miles de GPU, los enfoques tradicionales de depuración se rompen. Las múltiples dimensiones de paralelismo, los cálculos de baja precisión (MXFP8, NVFP4) y los checkpoints distribuidos dificultan enormemente reproducir errores y validar correcciones. NVIDIA ha incorporado a Megatron Core el determinismo bit a bit: la garantía de que dos ejecuciones iniciadas desde el mismo estado seguirán una trayectoria numérica idéntica, siempre que los datos, la arquitectura, el esquema de paralelismo y el software se mantengan sin cambios.

Megatron Core ofrece dos garantías clave: reproducibilidad independiente (dos ejecuciones lanzadas desde el mismo estado permanecen idénticas bit a bit en cada paso) y reanudación exacta desde un checkpoint (la recuperación coincide exactamente con una ejecución continua). Esto resulta crítico para los costosos entrenamientos en producción: incluso una pequeña ralentización, a escala de un billón de parámetros, puede costar miles de días-GPU.

Cuatro beneficios prácticos

El determinismo bit a bit resuelve cuatro problemas de depuración y validación:

• Reproducir errores: capturar un pico de loss en un paso concreto y luego modificar un componente a la vez para aislar la causa • Continuidad de checkpoints: reanudar un entrenamiento interrumpido sin alterar la trayectoria numérica • Detectar defectos: usar discrepancias entre ejecuciones que de otro modo serían idénticas para localizar corrupción de checkpoints o problemas de software y hardware • Comparar el impacto de los cambios: distinguir el efecto numérico de un cambio intencional de la variación natural entre ejecuciones

Metodología para depurar el no determinismo

NVIDIA propuso en el PR #7262 de Megatron-LM un sistema de trazado progresivo que registra huellas ordenadas de tensores entre rangos. El proceso reduce el área de búsqueda, desde métricas de extremo a extremo hasta un kernel concreto:

• Métricas de extremo a extremo: identificar la primera iteración en la que difiere una métrica serializada (loss de entrenamiento, loss de modelado de lenguaje, norma del gradiente, loss de MTP) • Trazado semántico amplio: cubrir colectivos, comunicación entre pipelines, recomputación, operaciones del optimizador y reducciones de gradiente para determinar en qué fase del entrenamiento se produce la divergencia • Trazado de módulos y capas: localizar el problema en un componente concreto del modelo (capa Transformer, bloque MoE, fase del optimizador) • Trazado a nivel de operación: huellas de operaciones ATen y llamadas a extensiones para encontrar la primera operación con entradas iguales y salidas distintas • Herramientas a nivel de kernel: identificar el kernel responsable y el mecanismo del no determinismo

Resulta fundamental no limitar el trazado a la iteración donde la divergencia del loss se hace visible; el primer bit que diverge puede aparecer mucho antes. El trazado amplio localiza la iteración y los rangos donde surge la divergencia más temprana, y solo entonces se activa el trazado detallado, limitado a esa ventana reducida.

Las huellas se calculan en la GPU mediante torch.hash_tensor, registrando forma, tipo de dato y número de elementos. En el caso de MXFP8 y NVFP4, se calculan huellas tanto de los valores codificados como de los búferes de escala. Para datos sensibles al orden (mapas de routing, salidas de dispatch de MoE) se utiliza un trazado por fragmentos (chunk-wise).

Optimización para un Nemotron de un billón de parámetros

Una receta totalmente determinista, con una sobrecarga elevada, resulta útil para depurar, pero es demasiado costosa para producción. NVIDIA llevó a cabo una co-optimización de las rutas determinista y no determinista en un modelo híbrido Nemotron, que combina capas SSM (modelo de espacio de estados) al estilo Mamba con capas de atención Transformer, optimizando tres cargas de trabajo:

• Nemotron 3 Ultra: reducción del ~17% al 2% tras una optimización de relleno de búfer, y después al 1,5% en 3.072 GPU • Proxy híbrido de Triton: reducción del ~60% al 37% tras restaurar la fusión MoE-MLP, y después al 2% • Ruta optimizada de gradiente de pesos con CuTeDSL: ~3,6% de sobrecarga en 256 GPU, y después ~2% en 2.432 GPU

En 2.432 GPU, la receta de producción alcanzó aproximadamente un 2% de sobrecarga en régimen estable, manteniendo el determinismo bit a bit durante 800 pasos de entrenamiento. Las optimizaciones a nivel de kernel incluyen la asignación de slots de salida privados para los escritores de grouped-GEMM y la aplicación de un orden fijo de reducción, que preserva el paralelismo sin dejar de recuperar resultados deterministas. Megatron Core valida estas recetas, prueba los kernels y verifica los módulos en todas las configuraciones de paralelismo, incluido el soporte para formatos de baja precisión FP8 y FP4.