kempnerforge.resilience.health

GPU health monitoring and NaN detection.

Provides utilities for detecting common failures during training:

  • NaN/Inf in loss, gradients or parameters

  • GPU availability and basic health

  • NCCL liveness via lightweight collectives

Functions

check_gpu_health([device])

Run basic GPU health checks.

check_nccl_health([timeout_sec])

Check NCCL communication health via a lightweight all-reduce.

params_are_finite(model, device)

Check that no rank holds a NaN/Inf parameter of model.

Classes

NaNDetector

Detects and tracks NaN/Inf values in loss and gradients.

NaNState

Tracks NaN/Inf occurrences across training steps.
class kempnerforge.resilience.health.NaNState[source]

Bases: object

Tracks NaN/Inf occurrences across training steps.

consecutive_nans: int = 0
total_nans: int = 0
last_good_loss: float = inf
last_good_step: int = 0
nan_steps: list[int]
__init__(consecutive_nans=0, total_nans=0, last_good_loss=inf, last_good_step=0, nan_steps=<factory>)
Parameters:

  • consecutive_nans (int)

  • total_nans (int)

  • last_good_loss (float)

  • last_good_step (int)

  • nan_steps (list[int])

Return type:

None

class kempnerforge.resilience.health.NaNDetector[source]

Bases: object

Detects and tracks NaN/Inf values in loss and gradients.

Supports three responses to NaN:

  • "warn": Log a warning and continue.

  • "skip": Skip the optimizer step (zero gradients).

  • "raise": Raise a RuntimeError.

Once the consecutive NaN count reaches max_consecutive, the detector signals that a checkpoint rollback is recommended.

Parameters:

  • action – What to do when NaN is detected.

  • max_consecutive – Consecutive NaN steps before recommending rollback.

  • max_history – Number of NaN step indices to retain.

__init__(action='warn', max_consecutive=5, max_history=100)[source]
Parameters:

  • action (str)

  • max_consecutive (int)

  • max_history (int)

Return type:

None

check_loss(loss, step, grad_norm=None)[source]

Check a step’s loss, and its gradient norm when given, for NaN/Inf.

When running distributed, all-reduces a NaN flag so ALL ranks agree on whether to skip. Prevents rank desync where one rank sees NaN and skips its optimizer step while others proceed normally.

A finite loss does not guarantee finite gradients: a micro-batch whose labels are all masked has an exact zero cross-entropy loss whatever its activations, so the caller passes the gradient norm as well. Each rank’s norm goes into the same reduced flag as the loss, so every rank skips the step together even when a non-finite gradient reaches only some ranks’ norms.

Parameters:

  • loss (float) – The scalar loss value to check.

  • step (int) – Current training step.

  • grad_norm (float | None) – The step’s gradient norm, or None to check the loss alone.

Returns:

True if the loss and the gradient norm are finite on ALL ranks, False if any rank has NaN/Inf in either.

Raises:

RuntimeError – If action is “raise” and NaN is detected.

Return type:

bool

check_gradients(model, step)[source]

Check model gradients for NaN/Inf before optimizer step.

Parameters:
Returns:

True if all gradients are finite.

Return type:

bool

property should_rollback: bool

Whether consecutive NaN count suggests a checkpoint rollback.

reset()[source]

Reset NaN tracking state (e.g., after a rollback).

Return type:

None

kempnerforge.resilience.health.params_are_finite(model, device)[source]

Check that no rank holds a NaN/Inf parameter of model.

Each rank checks the parameters it holds (the local shard of a DTensor). When running distributed, one all-reduce of a flag over the default process group gives every rank the same answer, so this is a collective that every rank must call.

Parameters:

  • model (torch.nn.Module) – The model, or this rank’s pipeline stage of it.

  • device (torch.device | str) – Device for the all-reduced flag; it must suit the default process group’s backend (CUDA under NCCL).

Returns:

True if every parameter is finite on every rank.

Return type:

bool

kempnerforge.resilience.health.check_gpu_health(device=0)[source]

Run basic GPU health checks.

Performs:

  1. CUDA availability check

  2. Small test computation on the device

  3. Memory allocation test

Returns:

Dict with health check results.

Parameters:

device (int)

Return type:

dict[str, bool | str]

kempnerforge.resilience.health.check_nccl_health(timeout_sec=10.0)[source]

Check NCCL communication health via a lightweight all-reduce.

The all-reduce runs with async_op=True so work.wait(timeout=...) enforces the caller’s bound rather than falling back to the process-group default timeout (nccl_timeout_sec, 1800s). Without that, this function would sit for 30 minutes on a single stuck peer regardless of the timeout_sec argument.

Parameters:

timeout_sec (float) – Per-operation timeout for the collective. Returns False if the all-reduce does not complete within this budget.

Returns:

True on success, False on timeout, error, or world-size mismatch.

Return type:

bool