NVIDIA demuestra determinismo bit a bit en Megatron Core para preentrenar con miles de GPU

Entrenar un modelo con billones de parámetros en miles de GPU cuesta millones de dólares y lleva semanas. Si a mitad de la ejecución aparece un pico de pérdida (loss) o un NaN, normalmente no es posible repetirlo: el orden de las operaciones cambia de una ejecución a otra y el error no se manifiesta. NVIDIA ha publicado en su Developer Blog un artículo sobre cómo lograr determinismo bit a bit en el preentrenamiento con Megatron Core. Dos ejecuciones con los mismos datos y la misma configuración deben dar resultados idénticos hasta el último bit.

Por qué el determinismo en el preentrenamiento es tan raro

La raíz del problema está en la aritmética de coma flotante: la suma no es asociativa. Si se suman los mismos números en otro orden, el resultado cambia ligeramente. En una sola GPU apenas se nota. En un clúster, el orden depende de muchos factores: operaciones atómicas en los kernels, la elección del algoritmo de reducción en las comunicaciones colectivas, la planificación de hilos y el reparto del trabajo entre dispositivos.

La fuente enumera varios factores que complican el problema:

• varias dimensiones de paralelismo a la vez (datos, tensores, pipeline y otras), cada una con sus propias reducciones e intercambios;
• cálculos en baja precisión, donde los errores de redondeo son mayores y dependen más del orden de las operaciones;
• checkpoints distribuidos, que hay que guardar y cargar de modo que reanudar el entrenamiento no altere la trayectoria;
• la escala: billones de parámetros y miles de GPU, donde la probabilidad de que ocurra un fallo poco frecuente durante una ejecución se acerca a uno.

En el fragmento disponible no se detallan las técnicas concretas con las que Megatron Core garantiza el orden de las operaciones ni el coste del determinismo en velocidad. Esos datos hay que consultarlos en la versión completa del material.

Qué aporta la reproducibilidad bit a bit

El anuncio señala tres beneficios prácticos: depuración, validación y reanudación reproducible del entrenamiento.

El primero afecta a la búsqueda de las causas de un fallo. Si la ejecución es determinista, se puede relanzar desde el último checkpoint previo a la anomalía y obtener la misma anomalía. Así, un error «fantasma» se convierte en un bug corriente, sobre el que se pueden poner comprobaciones y hacer bisección por capas o pasos.

El segundo afecta a la validación de correcciones. Con un entrenamiento no determinista, dos ejecuciones siempre divergen, y solo se puede saber de forma estadística si un fix ha influido en el resultado. Con reproducibilidad bit a bit, cualquier divergencia respecto a la referencia significa que el cambio ha afectado a algo. Y la ausencia de divergencia demuestra que un refactor o una optimización no han alterado el comportamiento numérico.

El tercero afecta a los reinicios. En clústeres grandes, los fallos de hardware ocurren con regularidad. Si, tras cargar un checkpoint, el entrenamiento sigue exactamente la misma trayectoria que habría seguido sin interrupción, los ingenieros pueden estar seguros de que el reinicio no ha introducido una deriva oculta.

El precio y dónde resulta útil

El determinismo suele pagarse con rendimiento: los kernels rápidos no deterministas y las reducciones atómicas se sustituyen a menudo por alternativas más lentas, pero predecibles. Por eso la pregunta clave es cuánto caerá el throughput en configuraciones reales. El fragmento citado no responde a esa pregunta, y sin mediciones es pronto para evaluar el balance.

Aun así, la lógica de uso es clara. No es necesario mantener el modo determinista de forma permanente en las ejecuciones más caras. Pero resulta útil en varios escenarios:

• al reproducir y analizar fallos, cuando hace falta una repetición exacta;
• al validar nuevos kernels, un cambio de precisión (por ejemplo, el paso a una más baja) o una modificación del esquema de paralelismo;
• en las pruebas de regresión del código de entrenamiento: una ejecución de referencia con pocos pasos, con la que se compara bit a bit cada build.

Conclusión práctica para los equipos

Los equipos que entrenan sus propios modelos con Megatron Core deberían montar una ejecución de control corta en modo determinista e incluirla en el CI. Cualquier cambio en el código, los drivers o la configuración de paralelismo se compara primero con la referencia mediante los hashes de los pesos tras un número fijo de pasos. Una divergencia indica de inmediato que el cambio ha afectado al comportamiento numérico.

Antes de implantarlo, conviene medir la sobrecarga en el propio clúster y con la precisión que se utilice. También hay que comprobar que el reinicio desde checkpoint da realmente una trayectoria idéntica en la topología de cada equipo. Lo razonable es activar el determinismo total de forma puntual: en depuración, validación y auditoría de ejecuciones. El entrenamiento principal puede seguir en modo rápido si las mediciones muestran una pérdida de velocidad apreciable.