Era el lote 47.000. Una variante de ResNet que había estado entrenando durante seis horas en un conjunto de datos de imágenes médicas personalizado. La pérdida estaba convergiendo limpiamente (1,4, 1,1, 0,87, 0,73) y luego, nada. No es un error. No es un accidente. Solo nana.
Agregué torch.autograd.set_detect_anomaly(True) y reinicié. El entrenamiento se ralentizó (aproximadamente entre 7 y 10 veces más por lote solo en la CPU) y después de tres horas finalmente obtuve un seguimiento de la pila que apuntaba a una capa que, francamente, se veía bien. El verdadero culpable fue un programador de tasa de aprendizaje que interactuaba mal con una capa de normalización personalizada dos capas arriba. set_detect_anomaly me había señalado el síntoma, no la fuente.
Esa sesión de depuración me costó la mayor parte de un día. Entonces construí algo mejor.
Los NaN no bloquean su modelo, lo corrompen silenciosamente. Cuando te das cuenta, ya estás depurando la capa equivocada.
Código completo: https://github.com/Emmimal/pytorch-nan-detector/
El problema con set_detect_anomaly
PyTorch viene con torch.autograd.set_detect_anomaly(True), que es la recomendación estándar para depurar problemas de NaN. Funciona reteniendo el gráfico de cálculo completo y verificando anomalías durante el pase hacia atrás. Esto es poderoso, pero conlleva serios costos que lo hacen inadecuado para cualquier cosa más allá de una rápida verificación de cordura local.
El problema principal es que obliga al motor de autograduación de PyTorch a entrar en un modo sincrónico donde guarda activaciones intermedias para cada operación. En GPU, esto significa romper el proceso de ejecución asincrónica: cada lanzamiento del kernel debe completarse antes de que comience el siguiente. El resultado, como se informa en la documentación de PyTorch y se observa ampliamente en la práctica, es una sobrecarga que oscila entre aproximadamente 10 y 15 veces en la CPU y 50 a 100 veces en la GPU para modelos más grandes.[1][2].
Hay un segundo problema: set_detect_anomaly indica dónde se propagó el NaN en el paso hacia atrás, no necesariamente dónde se originó. Si un NaN ingresa a su red en la capa 3 de un modelo de 50 capas, el paso hacia atrás mostrará un error en algún lugar del cálculo del gradiente para una capa posterior, y tendrá que trabajar hacia atrás desde allí.
Mi punto de referencia, ejecutado en una pequeña CPU MLP (64→256→256→10), midió:
En este modelo pequeño la diferencia absoluta es modesta. A escala (un transformador con cientos de millones de parámetros en múltiples GPU), la brecha es la diferencia entre una ejecución de entrenamiento que se completa y otra que no.
El enfoque: ganchos delanteros
La API Register_forward_hook de PyTorch le permite adjuntar una devolución de llamada a cualquier nn.Module que se activa cada vez que ese módulo completa un paso hacia adelante.[3]. La devolución de llamada recibe el módulo en sí, sus entradas y sus salidas. Esto significa que puede inspeccionar cada tensor que fluye a través de cada capa en tiempo real, sin impacto en el gráfico de cálculo, sin sincronización forzada y sin activaciones retenidas.
La idea clave es que solo necesita realizar la verificación de NaN, no reproducir el cálculo. Una verificación de torch.isnan() y torch.isinf() en un tensor de salida es una única invocación del kernel CUDA y se completa en microsegundos.
def gancho(módulo, entradas, salida): if torch.isnan(salida).any(): print(f"NaN detectado en {layer_name}")
Ese es el núcleo de la idea. Lo que sigue es la versión reforzada para producción.
La implementación
La fuente completa está disponible en: https://github.com/Emmimal/pytorch-nan-detector/
Analizaré los cuatro componentes que importan.
Componente 1: la clase de datos NaNEvent
Cuando se detecta un NaN, necesita más que una declaración impresa. Necesita un registro estructurado que pueda inspeccionar después del hecho, registrarlo en el disco o enviarlo a un sistema de alerta.
@dataclass clase NaNEvent: lote_idx: int nombre_capa: str tipo_módulo: cadena entrada_has_nan: bool salida_has_nan: bool entrada_has_inf: bool salida_has_inf: bool forma_salida: tupla estadísticas_salida: dict = campo(default_factory=dict) is_backward: bool = False
El campo output_stats contiene el mínimo, el máximo y la media de los valores finitos en el tensor de salida en el momento de la detección. Esto es sorprendentemente útil: una salida de capa donde 3 valores son NaN pero el resto son finitos cuenta una historia diferente a una que es completamente NaN.
El indicador is_backward distingue si el evento quedó atrapado en un gancho hacia adelante o hacia atrás, lo cual es importante para el análisis de la causa raíz.
Componente 2: Registro de gancho seguro para hilos
La consideración de producción más importante es la seguridad del hilo. DataLoader de PyTorch ejecuta procesos de trabajo que pueden activar enlaces directos desde subprocesos en segundo plano. Si muta triggered = True y self.event = ev sin un bloqueo, obtendrá condiciones de carrera en configuraciones de varios trabajadores.
self._lock = threading.Lock() def _make_fwd_hook(self, Layer_name: str): def hook(módulo, entradas, salida): with self._lock: if self.triggered y self.stop_on_first: return current_batch = self._batch_idx # … las comprobaciones del tensor ocurren fuera del bloqueo si out_nan o out_inf: self._record_event(…) # bloqueo recuperado dentro gancho de retorno
Las comprobaciones del tensor se realizan fuera del bloqueo porque torch.isnan() es de solo lectura y seguro para subprocesos. Sólo se bloquean las mutaciones de estado compartido.
Componente 3: memoria limitada
Un problema sutil con las ejecuciones de entrenamiento largas: si acumula tiempos generales en una lista ilimitada, eventualmente agotará la memoria en ejecuciones que duren millones de lotes. La solución es un límite simple:
_OVERHEAD_CAP = 1000 con self._lock: si len(self._overhead_ms) < self._OVERHEAD_CAP: self._overhead_ms.append(transcurrido)
La misma lógica se aplica a all_events cuando stop_on_first=False: un parámetro max_events (predeterminado 100) evita la acumulación ilimitada durante ejecuciones patológicas.
Componente 4: Guardia de norma de gradiente
El camino más común en el mundo real hacia un NaN no es un error que produce nan directamente; es una tasa de aprendizaje demasiado alta que hace que las normas de gradiente exploten hasta inf, que luego se propaga a los pesos y produce activaciones de NaN en el siguiente paso hacia adelante. Para cuando disparas tu gancho delantero, ya estás un paso tarde.
El método check_grad_norms() soluciona esto recorriendo todos los parámetros después de loss.backward() y registrando un GradEvent para cualquier parámetro cuya norma de gradiente exceda un umbral:
def check_grad_norms(self) -> bool: si self.grad_norm_warn es Ninguno: devuelve False para nombre, módulo en self.model.named_modules(): para pname, parámetro en module.named_parameters(recurse=False): si param.grad es Ninguno: continúa norma = param.grad.detach().float().norm().item() si no es math.isfinite(norm) o norma > self.grad_norm_warn: # registrar GradEvent
En la demostración siguiente, este método detecta la explosión de gradiente en el lote 1: un paso de entrenamiento completo antes de que el NaN apareciera en el pase hacia adelante.
Uso
Básico: administrador de contexto
desde nan_detector importe NaNDetector con NaNDetector(modelo) como det: para lote_idx, (x, y) en enumerar(cargador): det.set_batch(batch_idx) pérdida = criterio(modelo(x), y) pérdida.backward() det.check_grad_norms() optimizador.step() si det.triggered: print(det.event) break
Cuando se activa el detector, det.event contiene el NaNEvent completo con el nombre de la capa, el tipo de módulo, el índice de lote y las estadísticas de salida.
Producción: circuito de entrenamiento directo
de nan_detector importar pérdidas de train_with_nan_guard, evento = train_with_nan_guard( modelo, cargador, criterio, optimizador, dispositivo="cuda", grad_norm_warn=50.0,) si evento: print(f"NaN en el lote {event.batch_idx}, capa {event.layer_name}")
Avanzado: ganchos hacia atrás + nombres de capas legibles
Para detectar NaN de gradiente directamente (no solo advertencias de normas), habilite check_backward=True. Utilice OrderedDict al crear modelos secuenciales para obtener nombres legibles en todos los resultados de los registros:
de colecciones import OrderedDict model = nn.Sequential(OrderedDict([ ("fc1", nn.Linear(16, 32)), ("relu1", nn.ReLU()), ("fc2", nn.Linear(32, 1)), ])) con NaNDetector(model, check_backward=True, grad_norm_warn=10.0) como det: …
Sin OrderedDict, PyTorch nombra las capas por índice (0.peso, 2.bias). Con él, obtienes fc1.weight, fc2.bias, algo pequeño que ahorra tiempo real al depurar modelos profundos.
Saltar capas
Se espera que algunos tipos de capas produzcan resultados no finitos en condiciones normales: nn.Dropout durante la evaluación, ciertas capas de normalización durante el primer paso hacia adelante antes de que se establezcan las estadísticas de ejecución. Sáltelos con:
det = NaNDetector (modelo, skip_types = (nn.Dropout, nn.BatchNorm1d))
Salida de demostración
La ejecución de las tres demostraciones produce el siguiente resultado:
────────────────────────────── ────────────────────────────── Demostración 1: detección directa de NaN + trazado de curva de pérdida ────────────────────────────── ────────────────────────────── [NaNDetector] Se adjuntan 5 ganchos. ============================================================== NaN/Inf detectado! [PASO ADELANTE] Lote: 12 Capa: capa4 Tipo: Banderas lineales: NaN en ENTRADA, NaN en SALIDA Forma de salida: (8, 1) Estadísticas de salida: min=n/a (todas no finitas) max=n/a (todas no finitas) media=n/a (todas no finitas) ============================================================== [NaNDetector] Separado. Gastos generales promedio: 0,109 ms/paso directo ────────────────────────────── ────────────────────────────── Demostración 2: Detección hacia atrás/de normas de graduación + trazado de normas de graduación ────────────────────────────── ────────────────────────────── [NaNDetector] Se adjuntan 8 ganchos (+ hacia atrás). [ADVERTENCIA GradNorm] lote = 1 capa = fc1.weight norma = umbral inf = 10.0 [ADVERTENCIA GradNorm] lote = 1 capa = fc1.bias norma = umbral inf = 10.0 [ADVERTENCIA GradNorm] lote = 1 capa = fc2.weight norma = umbral inf = 10.0 [ADVERTENCIA GradNorm] lote = 1 capa = fc2.bias norma=4.37e+18 umbral=10.0 Capturado en el lote 1
La sobrecarga de gancho de 0,109 ms por pase hacia adelante en la Demostración 1 es el número real que puede citar. La cifra de referencia de ~3 ms refleja un modelo más grande con cinco devoluciones de llamadas de gancho registradas ejecutándose simultáneamente, que es el escenario de producción más realista.
Limitaciones conocidas
Los ganchos directos ven activaciones, no todos los cálculos. Si un NaN se origina dentro de un método forward() personalizado de torch.autograd.Function, o dentro de una extensión C++/CUDA que no aparece a través de submódulos nn.Module con nombre, el enlace directo no lo detectará. Utilice check_backward=True para cobertura del lado del gradiente y grad_norm_warn para alerta temprana.
Básculas aéreas con profundidad de modelo. El punto de referencia se ejecutó en un MLP de 5 capas. Un transformador de 200 capas tendrá 200 devoluciones de llamadas de gancho disparadas por pase hacia adelante. La sobrecarga sigue siendo inferior a un milisegundo por gancho, pero se acumula. Mitigue utilizando skip_types para excluir capas no paramétricas como ReLU, Dropout y LayerNorm si la sobrecarga se convierte en una preocupación.
Las relaciones de referencia de la CPU son ruidosas. La relación de gastos generales entre NaNDetector y set_detect_anomaly varió entre 5× y 6× en las ejecuciones de mis pruebas, porque los microbenchmarks de la CPU en una escala de submilisegundos son sensibles a la programación del sistema operativo y al estado de la caché. Los números absolutos de milisegundos son más estables. La cifra de 50 a 100 × citada para GPU se extrae de la documentación de PyTorch y de los puntos de referencia de la comunidad.[1][2], no mis propias medidas de GPU.
Lo que esto no reemplaza
Esta es una herramienta de depuración y seguimiento, no un sustituto de una buena higiene del entrenamiento. Las recomendaciones estándar aún se aplican: recorte de gradiente (torch.nn.utils.clip_grad_norm_), programación cuidadosa de la tasa de aprendizaje, normalización de entrada e inicialización de peso. NaNDetector le indica dónde y cuándo ocurrió el problema; no le dice por qué, y solucionar la causa raíz aún requiere criterio de ingeniería.
Si está utilizando NaN en el entrenamiento de precisión mixta (fp16/bf16), los culpables más comunes son el desbordamiento de escala de pérdida y la inestabilidad de la norma de capa, y vale la pena investigarlos directamente antes de recurrir a un gancho de depuración.
Metodología de referencia
Todos los puntos de referencia se ejecutaron en CPU (Windows 11, PyTorch 2.x) usando un MLP de 4 capas con dimensión de entrada 64, dos capas ocultas de 256 y dimensión de salida 10. El tamaño del lote fue 64. Cada método ejecutó 30 pases hacia adelante. La primera pasada se incluyó en la media: los efectos del arranque en frío son reales y deben contarse. Los tiempos se midieron con time.perf_counter() solo alrededor de la llamada de reenvío, sin incluir la carga de datos ni el cálculo de pérdida.
La función de evaluación comparativa completa está incluida en el código fuente y se puede ejecutar con evaluación comparativa (n_batches=30, lote_size=64).
Referencias
[1]Documentación de PyTorch. "Mecánica de Autograd: detección de anomalías". pytorch.org. Disponible en: https://pytorch.org/docs/stable/autograd.html#anomaly-detection
[2]Documentación de PyTorch. torch.autograd.set_detect_anomaly. pytorch.org. Disponible en: https://docs.pytorch.org/docs/stable/autograd.html
[3]Documentación de PyTorch. torch.nn.Module.register_forward_hook. pytorch.org. Disponible en: https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_forward_hook
[4]Documentación de PyTorch. torch.nn.Module.register_full_backward_hook. pytorch.org. Disponible en: https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_hook
[5]Documentación de PyTorch. "Recorte de gradiente – clip_grad_norm_". pytorch.org. Disponible en: https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
[6]Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G.,… y Chintala, S. (2019). PyTorch: una biblioteca de aprendizaje profundo de alto rendimiento y estilo imperativo. Preimpresión de arXiv arXiv:1912.01703. https://doi.org/10.48550/arXiv.1912.01703
[7]Fundación de software Python. threading: paralelismo basado en subprocesos. Documentación de Python 3. Disponible en: https://docs.python.org/3/library/threading.html
[8]Fundación de software Python. clases de datos: clases de datos. Documentación de Python 3. Disponible en: https://docs.python.org/3/library/dataclasses.html
[9]Cazador, JD (2007). Matplotlib: un entorno de gráficos 2D. Computación en ciencia e ingeniería, 9(3), 90–95. https://doi.org/10.1109/MCSE.2007.55
Divulgación
Yo mismo construí y escribí sobre esta herramienta. No existe patrocinio, afiliación con PyTorch o la Fundación PyTorch, ni relación financiera con ninguna de las empresas mencionadas en este artículo. Los puntos de referencia se ejecutaron en mi propio hardware y se pueden reproducir utilizando el código del repositorio vinculado anteriormente.
Todo el código de este artículo es original. La herramienta fue escrita desde cero; no se utilizó como base ninguna biblioteca de detección de NaN de código abierto existente. Si usa esto en su propio trabajo, se agradece la atribución, pero no es obligatoria: el código tiene licencia del MIT.
La comparación comparativa con set_detect_anomaly se basa en mis propias mediciones en una configuración de hardware específica. Los resultados variarán según la arquitectura del modelo, el hardware y la versión de PyTorch. La cifra de sobrecarga de GPU de 50 a 100 × se extrae de la documentación oficial de PyTorch.[1][2]y no es mi propia medida de GPU.
Código fuente completo, incluidas las tres demostraciones y la función de referencia: https://github.com/Emmimal/pytorch-nan-detector/