NaN-y w PyTorchu niszczą trening w ciszy — oto jak je złapać w 3ms

Twój model trenuje się normalnie, a po 47 epokach loss nagle skacze do NaN-a. Spędzasz dwie godziny szukając przyczyny w kodzie, a problem siedział w warstwie głębokim, którą nigdy nie sprawdzałeś. PyTorch nie zatrzymuje treningu — po prostu pozwala mu się toczyć w nieskończoność z wartościami, które nie mają sensu. Zamiast czekać na katastrofę, możesz złapać NaN dokładnie tam, gdzie się pojawia, z narzutem zaledwie 3ms na forward pass.

NaN-y niszczą Twój model bez ostrzeżenia — oto co się dzieje pod spodem

NaN-y w PyTorchu to ciche zabójcy. Model trenuje się, loss się zmienia, ale wartości są śmieci. Dopóki nie sprawdzisz metryki na walidacji, nie wiesz, że coś poszło nie tak. A nawet wtedy — jeśli nie masz automatycznego monitoringu — możesz przegapić moment, w którym sieć zaczęła się psować.

Dlaczego PyTorch nie zatrzymuje treningu? Bo NaN to legalny float. Operacja matematyczna, która powinna zwrócić liczbę, zwraca NaN — i framework nie ma powodu, żeby to blokować. Gradient descent pracuje dalej, parametry się aktualizują, ale w kierunku, który nie ma sensu. Po kilku iteracjach cała sieć może być skażona.

Główne źródła NaN-ów to trzy rzeczy. Pierwsza: zbyt duży learning rate. Gdy gradient jest duży, a krok aktualizacji jeszcze większy, parametry mogą przeskoczyć w region, gdzie funkcja strata zwraca nieskończoność, a potem NaN [2]. Druga: mixed precision bez GradScaler. W FP16 wartości są mniejsze, przepełnienia są łatwiejsze, a underflow do zera w FP16 prowadzi do dzielenia przez zero w backward pass [2]. Trzecia: operacje matematyczne na wartościach granicznych — log(0), sqrt(-1), log wartości ujemnej — wszystko to generuje NaN [2].

Koszt ciszy to godziny tracone na debugowanie. Typowy workflow: zauważysz NaN, wrócisz do ostatniego checkpointu, zaczniesz zmniejszać learning rate, testujesz, czekasz 30 minut na epokę, znowu NaN. Powtarzasz to 4–5 razy. Dwa dni pracy zamiast dwóch godzin iteracji.

Standardowe narzędzia PyTorcha są zbyt wolne do produkcji

PyTorch ma wbudowane narzędzia do detekcji NaN-ów. Żadne z nich nie sprawdza się w praktyce.

torch.autograd.detect_anomaly(True) to klasyczne rozwiązanie. Włączasz go, PyTorch dodaje asercje między krokami backward, i jeśli gdzieś pojawi się NaN, wiesz gdzie [3]. Problem: spowalnia trening 5–10 razy [3]. Dla modelu, który trenuje się 8 godzin, to 40–80 godzin debugowania. Nie do zaakceptowania w produkcji.

torch.isnan() i torch.isinf() to proste funkcje — zwracają tensor bool, gdzie True oznacza NaN lub Inf [4]. Możesz napisać pętlę, która sprawdza każdy tensor po forward pass. Działa szybko, ale nie mówi ci, gdzie dokładnie pojawił się NaN. Jeśli model ma 50 warstw, musisz ręcznie sprawdzić każdą [4].

Gradient clipping — na przykład torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) — to obrona, nie diagnostyka [2]. Zmniejsza skalę gradientów, żeby nie eksplodowały, ale nie mówi ci, że coś jest nie tak. Możesz clipować gradienty przez całą epokę, a potem odkryć, że model się nie nauczył nic sensownego.

Żadne z tych narzędzi nie daje ci odpowiedzi na pytanie: w której warstwie pojawił się NaN, w którym batchu, i dlaczego?

Forward hooki: jak złapać NaN dokładnie w warstwie, gdzie się pojawia

Rozwiązanie to forward hooki — funkcje, które rejestrują się na warstwach i uruchamiają po każdym forward pass. Każdy hook sprawdza, czy wyjście warstwy zawiera NaN lub Inf, i loguje to z kontekstem [1].

Architektura detektora to pętla po wszystkich warstwach modelu. Na każdej rejestrujesz hook, który:

1. Sprawdza, czy tensor wyjściowy zawiera NaN lub Inf
2. Jeśli tak — zapisuje nazwę warstwy, indeks batcha, rozmiar tensora i wartości diagnostyczne
3. Jeśli nie — nic nie robi

Narzut obliczeniowy? Sprawdzenie NaN na tensorze to pojedynczy CUDA kernel trwający kilka mikrosekund [1]. Dla całej sieci MLP (kilka warstw, tysiące parametrów) to zaledwie 3ms na forward pass [1]. Dla porównania: jeden forward pass na GPU trwa 10–50ms. Narzut to 6–30% czasu forward pass, ale w zamian masz pełną diagnostykę.

Co rejestrujesz? Nazwa warstwy (np. layer_5), indeks batcha (np. batch_42), rozmiar tensora (np. torch.Size([32, 512])), i pierwsze kilka wartości tensora przed i po operacji [1]. To wystarczy, żeby w 2 minuty prześledzić, co poszło nie tak.

Praktycznie: rejestrujesz hooki na istniejącym modelu w 5 linijkach kodu. Nie musisz refaktoryzować architektury, nie musisz zmieniać pętli treningowej. Hooki pracują w tle.

Gradient norm guard: druga linia obrony po backward()

Forward hooki łapią NaN-y w wyjściu warstwy. Ale problem może pojawić się wcześniej — w gradientach. Eksplodujące gradienty to wczesny sygnał, że coś się sypie.

Gradient norm guard to druga linia obrony. Po każdym backward pass mierzysz normę gradientów dla każdego parametru. Jeśli norma jest większa niż threshold (na przykład 10.0), zapisujesz stack trace i logowanie. To daje ci wczesne ostrzeżenie, zanim NaN się rozprzestrzeni [1].

Integracja z logowaniem: kiedy norma gradientu przekroczy threshold, zapisujesz:

• Nazwę parametru (np. layer_5.weight)
• Wartość normy (np. 15.3)
• Indeks kroku treningowego
• Stack trace z informacją, gdzie w kodzie doszło do backward pass

Połączenie forward hooków + gradient norm guard = pełna diagnostyka. Forward hooki mówią ci, gdzie pojawił się NaN w forward pass. Gradient norm guard mówi ci, że coś się sypie, zanim NaN się pojawi.

Praktyka: jak wdrożyć detektor w 30 minut bez refaktoryzacji kodu

Wdrożenie to proste. Załóżmy, że masz model model i pętlę treningową.

Krok 1: Rejestracja hooków na istniejącym modelu — 5 linii kodu.

def register_nan_hooks(model):
    for name, module in model.named_modules():
        module.register_forward_hook(
            lambda m, inp, out, n=name: check_nan(out, n)
        )

Krok 2: Funkcja check_nan() — sprawdza tensor, loguje jeśli NaN.

def check_nan(tensor, layer_name):
    if isinstance(tensor, tuple):
        tensor = tensor[0]
    if torch.isnan(tensor).any() or torch.isinf(tensor).any():
        print(f"NaN/Inf detected in {layer_name}")
        print(f"Tensor shape: {tensor.shape}")
        print(f"First values: {tensor.flatten()[:10]}")

Krok 3: Integracja z pętlą treningową — dodajesz check_grad_norms() po backward().

for epoch in range(num_epochs):
    for batch in dataloader:
        output = model(batch)
        loss = criterion(output, labels)
        optimizer.zero_grad()
        loss.backward()
        check_grad_norms(model)  # Nowa linia
        optimizer.step()

Funkcja check_grad_norms():

def check_grad_norms(model, threshold=10.0):
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norm = param.grad.norm()
            if grad_norm > threshold:
                print(f"High gradient norm in {name}: {grad_norm.item():.2f}")

Czytanie logów: gdy widzisz NaN/Inf detected in layer_5, wiesz, że problem jest w warstwie 5. Sprawdzasz, co się tam dzieje — czy to operacja matematyczna, czy przepełnienie w FP16, czy coś innego. Od NaN w layer_5 batch_42 do przyczyny to 2 minuty, zamiast 2 godzin.

Kiedy wdrażać: mixed precision, duże learning rate, nowe architektury — to są Twoje flagi

Detektor NaN-ów nie jest obowiązkowy dla każdego treningu. Ale są scenariusze, gdzie powinien być standardem.

Pierwszy: mixed precision bez GradScaler. Gdy używasz FP16 do przyspieszenia treningu, ale zapominasz o GradScaler, underflow w FP16 prowadzi do dzielenia przez zero i NaN [2]. Jeśli trenujesz z mixed precision, detektor powinien być włączony od pierwszego batcha.

Drugi: duży learning rate. Gdy eksperymentujesz z learning rate > 0.1, ryzyko eksplozji gradientów rośnie [2]. Detektor da ci wczesne ostrzeżenie.

Trzeci: nowe architektury. Gdy piszesz custom loss function lub nową warstwę, nie wiesz, czy będzie generować NaN-y. Detektor pracuje w tle i mówi ci, jeśli coś pójdzie nie tak.

Czwarty: duże modele. Gdy trenujesz model z milionami parametrów, eksplozja gradientów jest bardziej prawdopodobna. Detektor to ubezpieczenie.

Koszt vs benefit: 3ms narzutu na forward pass vs 2–4 godziny debugowania. ROI jest oczywisty. Jeśli trenujesz model przez 24 godziny, narzut to 7–14 minut całkowitego czasu. W zamian masz pewność, że jeśli pojawi się NaN, wiesz o tym natychmiast.

Ograniczenie: detektor nie zapobiega NaN-om, tylko je wykrywa. Jeśli problem to zbyt duży learning rate, musisz go zmniejszyć. Jeśli to custom loss function, musisz go poprawić. Detektor to narzędzie diagnostyczne, nie magiczna wróżka.

Werdykt: detektor powinien być standardem w każdym pipeline'u treningowym

NaN-y w PyTorchu to nie rzadkość — to systemowy problem, który pojawia się w każdym projekcie deep learning'u. Standardowe narzędzia PyTorch'a są albo zbyt wolne (detect_anomaly), albo zbyt proste (torch.isnan). Forward hooki + gradient norm guard to praktyczne rozwiązanie, które zajmuje 30 minut do wdrożenia i daje ci pełną diagnostykę.

Jeśli trenujesz modele w produkcji — mixed precision, duże learning rate, nowe architektury — detektor powinien być obowiązkowy. Koszt to 3ms na forward pass. Benefit to pewność, że jeśli coś pójdzie nie tak, wiesz o tym w 2 minuty, zamiast szukać przez 2 godziny.

Zacznij od rejestracji hooków na swoim modelu. Dodaj check_grad_norms() do pętli treningowej. Trenuj. Jeśli pojawi się NaN, będziesz wiedzieć dokładnie gdzie i dlaczego.

Źródła

[1] PyTorch NaNs Are Silent Killers — So I Built a 3ms Hook to Catch Them at the Exact Layer

https://towardsdatascience.com/pytorch-nans-are-silent-killers-i-built-a-3ms-hook-to-catch-them-at-the-exact-layer/

[2] NaN loss — causes and fixes - PyTorch - The Neural Base

https://theneuralbase.com/pytorch/learn/intermediate/nan-loss-causes-and-fixes/

[3] Nan Loss coming after some time - PyTorch Discuss

https://discuss.pytorch.org/t/nan-loss-coming-after-some-time/11568

[4] Pytorch Operation to detect NaNs - Stack Overflow

https://stackoverflow.com/questions/48158017/pytorch-operation-to-detect-nans