Отлавливаем NaN в PyTorch за 3 мс: forward hook для точной локализации первого NaN в слое

Тихий убийца обучения

Вы запускаете обучение, через 12 часов loss падает, и вдруг — NaN. Всё. Веса рассыпались, модель молчит, а где именно произошла катастрофа — загадка. Можно перезапустить с чекпоинта, но если NaN приходит регулярно, обучение превращается в русскую рулетку. Классические методы — torch.isnan после каждого слоя или torch.autograd.set_detect_anomaly(True) — либо громоздкие, либо убивают производительность. Я покажу, как за 3 миллисекунды зарегистрировать forward hook, который укажет точный слой и номер итерации, где впервые появился NaN.

Почему detect_anomaly — это костыль

Встроенная опция torch.autograd.set_detect_anomaly(True) полезна, но плата за нее — замедление в 2-5 раз. Она перехватывает исключения при обратном распространении, но не говорит, где именно в forward произошел NaN. А если модель глубиной 200 слоев, искать виновника ручным перебором — ад. Кроме того, detect_anomaly не сработает, если NaN появился только в forward без backward (например, при инференсе или в loss). Альтернатива — вставить проверки после каждого слоя вручную: 20 строк кода на слой, и через час вы уже ненавидите нейросети.

Forward hook решает все эти проблемы. Он не трогает граф вычислений, не замедляет backward, а лишь проверяет тензоры на выходе каждого модуля. Overhead — копеечный (единицы миллисекунд на батч). И главное — он говорит: «NaN появился в LayerNorm.3 на batch_idx=541, feature_map[0, 12, 45]». Это level up в отладке.

Как работает forward hook — за 30 секунд

В PyTorch любой nn.Module поддерживает три типа хуков: forward pre-hook (до выполнения), forward hook (после), backward hook (на градиенты). Нас интересует forward hook: функция, которая вызывается сразу после forward() слоя, получает на вход модуль, входной тензор и выходной. Мы внутри проверяем, содержит ли выход NaN или inf. Если да — логируем имя модуля, значения и останавливаем обучение (или просто пишем в лог).

Критичный момент: чтобы локализовать первый NaN, нужно остановить выполнение forward при обнаружении. Иначе NaN распространится на следующие слои, и вы увидите только финальный коллапс. Hook должен поднять флаг или бросить исключение, прерывающее forward. Мы сделаем это через raise RuntimeError.

Код: детектор NaN за 15 строк

Создаем класс-декоратор, который регистрирует hook на все модули модели. Для PyTorch 2.x используем model.modules() и register_forward_hook. Первый вариант — простой логгер на момент появления NaN. Второй — с детальным выводом координат NaN внутри тензора.

import torch
import torch.nn as nn

def nan_detector_hook(module, input, output, raise_on_nan=True):
    if isinstance(output, torch.Tensor):
        if torch.isnan(output).any() or torch.isinf(output).any():
            msg = f"NaN/Inf detected in {module.__class__.__name__} at {id(module)}"
            if raise_on_nan:
                raise RuntimeError(msg)
            else:
                print(msg)
    elif isinstance(output, (tuple, list)):
        for i, out in enumerate(output):
            if isinstance(out, torch.Tensor) and (torch.isnan(out).any() or torch.isinf(out).any()):
                msg = f"NaN/Inf detected in {module.__class__.__name__}[{i}]"
                if raise_on_nan:
                    raise RuntimeError(msg)
                else:
                    print(msg)

def register_nan_detector(model, raise_on_nan=True):
    handles = []
    for name, module in model.named_modules():
        handle = module.register_forward_hook(
            lambda mod, inp, out, n=name: nan_detector_hook(mod, inp, out, raise_on_nan)
        )
        handles.append(handle)
    return handles

Используем так:

model = MyBigModel()
handles = register_nan_detector(model)
try:
    output = model(data)
except RuntimeError as e:
    print("First NaN caught:", e)
finally:
    for h in handles:
        h.remove()

Точная координата NaN — расширенная версия

Просто знать, в каком слое возник NaN, иногда недостаточно. Если выход слоя — тензор 1024x512, нужно понять, какой именно элемент испортился. Допишем логирование индекса:

def nan_detector_hook_detail(module, input, output, raise_on_nan=True):
    if isinstance(output, torch.Tensor):
        nan_mask = torch.isnan(output) | torch.isinf(output)
        if nan_mask.any():
            first_nan_idx = nan_mask.flatten().nonzero(as_tuple=False)[0].item()
            # convert flat index to multi-dim
            shape = output.shape
            idx_nd = np.unravel_index(first_nan_idx, shape)
            val = output.flatten()[first_nan_idx].item()
            msg = (f"NaN/Inf in {module.__class__.__name__} at index {idx_nd}, "
                   f"shape={shape}, value={val}")
            if raise_on_nan:
                raise RuntimeError(msg)
            else:
                print(msg)

Этот код также выводит значение NaN (обычно 0/0 или inf*0). Часто причина — деление на ноль в LayerNorm или Softmax с огромными логитами. Увидев значение, вы сразу поймете проблему.

Кому это нужно и когда не сработает

Техника полезна любому, кто обучает модели глубже 10 слоев, особенно языковые модели с нуля или архитектуры с residual connections (ResNet, Transformer). NaN часто возникает из-за градиентного взрыва или нестабильных численных операций. В YOLOv2, например, BatchNorm может дать NaN при малом batch size — hook моментально укажет на этот слой.

Ограничения: если NaN появляется внутри custom autograd Function (например, свой triCubic interpolation), forward hook не сработает, потому что выход модуля может быть целым, а NaN возникает в промежуточных вычислениях. В таких случаях нужен backward hook или torch.autograd.Function с проверкой внутри. Но 90% современных моделей построены из стандартных nn.Module — для них решение идеально.

Сравнение с альтернативами

Как видите, forward hook — единственный метод, который не жертвует скоростью и дает максимальную информацию. Если вам нужно еще и отследить градиенты — комбинируйте с backward hook (но это уже другая история).

Продакшн-режим и профилирование

Можно встроить детектор в пайплайн обучения как callable, который активируется только при превышении порога loss или на определенных эпохах. Либо используйте torch.cuda.profiler в связке с hook — как в разборе NVIDIA Nsight. Overhead настолько мал, что детектор можно держать включенным постоянно во время экспериментов.

Итог: ловить NaN — не больно

Forward hook — это швейцарский нож для отладки нейросетей. Он дешевый, точный и легко встраивается. Вместо того чтобы гадать, где произошел взрыв, вы получаете четкий ответ за 3 мс. Добавьте этот инструмент в свой набор — и обучение перестанет быть черным ящиком. А если хотите копнуть глубже — вот как оптимизировать инференс LLM, чтобы на этапе продакшена NaN вообще не возникали.

Подписаться на канал