LLM Training NaN/Inf Production Guide: Non-finite Gates, GradScaler Evidence, and Layer Bisect to Isolate the First Bad Step

Background: NaN is Often Not the Starting Point, but the Last Domino

When large model training encounters NaN, +Inf, or -Inf, the first thing you usually see in the logs is loss=nan, but the real fault may have occurred much earlier: anomalous input creates extreme logits, a layer’s activation overflows, gradients exceed the dynamic range under low precision, and then non-finite values enter gradient communication, optimizer states, and model parameters.

Therefore, “restarting the job after detecting NaN” is not a solution. A production system needs to answer five questions:

1. In which Step, Rank, Layer, and Tensor did the first non-finite value appear?
2. Have the bad values already entered the Collective, Optimizer State, or Checkpoint?
3. Did all ranks make the same decision about skipping the update?
4. Can the same batch, random state, and configuration be used for reproduction?
5. After recovery, how do we prevent a bad checkpoint from re-entering the training pipeline?

These five questions determine whether a numerical anomaly is a diagnosable event or an incident that requires a full restart.

Core Principles: Decomposing Numerical Anomalies into Four Boundaries

1. Input Boundary: First, Check if the Data is Already Non-finite

Before the training loop enters the model, perform lightweight checks on high-risk fields:

• Whether floating-point inputs contain NaN/Inf;
• Whether Token, Position ID, and Mask are out of bounds;
• Whether Labels contain disallowed values;
• Whether sample length, valid token count, and loss weights are anomalous;
• Whether the input domain for custom normalization, division, logarithm, exponent, and probability calculations is valid.

Performing a full scan on every large tensor permanently is not recommended. A more practical approach is: under normal operation, only check input summaries and key scalars; when a risk signal appears, enable deep scanning for the target batch.

2. Backward Boundary: Set Up Non-finite Gates Before Communication and Parameter Updates

A finite loss does not mean the gradients are finite. Anomalies can arise only during the backward pass, so the gate should be placed before gradient communication or the Optimizer Step.

PyTorch’s clip_grad_norm_ supports error_if_nonfinite=True, which raises an error if the total gradient norm is NaN/Inf. Megatron Core also provides configurations to check for NaN/Inf before gradient collectives. The goal of such gates is not to “clip bad gradients into good ones,” but to prevent bad values from spreading further.

3. Optimizer Boundary: Distinguish Normal Loss Scale Probing from Real Faults

In FP16 mixed-precision training, GradScaler amplifies the loss and then scales down the gradients after the backward pass. If non-finite gradients are found, scaler.step(optimizer) skips the parameter update and adjusts the Scale.

A single skip might just be the dynamic Scale finding a suitable range, but the following situations should trigger an alert:

• Multiple consecutive Steps are skipped;
• The skip ratio increases significantly within a short window;
• The Scale keeps decreasing but cannot recover;
• The same data shard or Layer triggers repeatedly;
• Anomalies persist even with BF16/FP32 replay.

Therefore, you must include the current Scale, previous Scale, whether a skip occurred, total gradient norm, learning rate, and batch fingerprint in the evidence package.

4. Distributed Boundary: Skipping Must Be a Global Decision

In data-parallel training, if one rank detects a bad gradient, it cannot just continue locally. It is recommended that each rank computes a local bad_flag, and then uses a small collective to aggregate them. If any rank reports an anomaly, all ranks execute the same action:

• Do not perform the parameter update;
• Clear the gradients for this step;
• Save the event metadata;
• Update or freeze the Loss Scale according to a unified policy;
• Enter a replay, degradation, or termination branch.

This rule is especially important because inconsistent skipping will cause the parameters and optimizer states of different ranks to diverge.

Engineering Implementation: Building a “Detect-Stop-Locate-Recover” Loop

Step 1: Record the First Bad Step, Not the Last Error Report

Maintain a minimal event ledger for each training step:

Field	Description
run_id / step / micro_step / global_batch_id	Training location
rank / dp_rank / tp_rank / pp_rank	Distributed location
dataset_version / shard_id / sample_ids	Data provenance
model_commit / tokenizer_hash / config_hash	Code and config fingerprint
precision / loss_scale / learning_rate	Precision and hyperparameters
total_grad_norm / skipped_step / first_bad_stage	Anomaly flags
rng_state_hash / checkpoint_parent	Reproducibility and lineage

Once a non-finite value is detected, immediately freeze the current event record. Subsequent logs can be appended, but the “first bad step” must not be overwritten.

Step 2: Save Evidence Hierarchically, Don’t Dump the Entire Model

Dumping all activations and gradients completely will quickly exhaust storage. A more sensible approach is a three-level evidence strategy:

Level	Content	Trigger
L1 Normal Evidence	Loss, Scale, Gradient Norm, Input Summary, Config Fingerprint	Every step
L2 Anomaly Evidence	First anomalous Rank, Parameter Group, Gradient Bucket, Module Name, Tensor Statistics	Non-finite value detected
L3 Localization Evidence	Input, Output, Gradient, and Operator Stack of the target Layer	Bad batch replay

Megatron FSDP provides debugging capabilities to report NaN gradients per weight, but the official documentation also notes its significant performance cost. Therefore, fine-grained detection should be limited to the replay window, not enabled permanently.

Step 3: Use Layer Bisect to Narrow Down the Anomaly Range

For models with many Transformer Blocks, use a binary search approach:

1. Replay with a fixed bad batch and fixed RNG state;
2. First, check if the activations and gradients at the boundary of the middle layer are finite;
3. Determine if the anomaly is in the first half or the second half;
4. Iteratively narrow it down to a specific Block;
5. Finally, add hooks to the Attention, Normalization, MLP, and custom operators within the target Block.

torch.autograd.detect_anomaly(check_nan=True) can provide the forward call trace that led to the failing backward function during debugging, but it significantly reduces performance, so it’s suitable for short-window replay of a bad batch.

Step 4: Establish Checkpoint Isolation Rules

When an anomaly occurs at a certain step, you cannot assume all previously saved checkpoints are safe. It is recommended to set three states:

State	Meaning	Action
committed	Passed forward, backward, gradient, and parameter finiteness checks	Allowed for recovery
suspect	Near the first bad step, not yet fully verified	Requires finiteness scan before release
quarantined	Explicitly contains non-finite parameters, optimizer states, or irreproducible anomalies	Forbidden for recovery

Recovery tasks should only read committed checkpoints. For suspect checkpoints, perform a finiteness scan of parameters, optimizer states, and key buffers before deciding whether to release them.

PyTorch Reference Implementation: Unified Stop-Loss Before Update

The following code shows the core sequence. A real distributed system should also include inter-rank bad_flag aggregation and evidence persistence.

import torch

def train_step(model, optimizer, scaler, batch, max_grad_norm: float):
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(**batch)
        loss = outputs.loss

    if not torch.isfinite(loss):
        raise FloatingPointError("non-finite loss before backward")

    scaler.scale(loss).backward()
    # Must unscale before reading, checking, or clipping the real gradients
    scaler.unscale_(optimizer)

    try:
        grad_norm = torch.nn.utils.clip_grad_norm_(
            model.parameters(),
            max_norm=max_grad_norm,
            error_if_nonfinite=True,
        )
    except RuntimeError as exc:
        optimizer.zero_grad(set_to_none=True)
        # Record batch fingerprint, loss scale, rank, and config fingerprint here
        raise FloatingPointError("non-finite gradients") from exc

    scaler.step(optimizer)
    scaler.update()

    return {
        "loss": float(loss.detach()),
        "grad_norm": float(grad_norm.detach()),
        "loss_scale": float(scaler.get_scale()),
    }

The critical order is: Backward → Unscale → Finiteness Check/Clip → Step → Scale Update. Do not use the gradient norm to judge training stability while the gradients are still scaled.

Applicable Scenarios

This approach is suitable for:

• FP16, BF16, or FP8 mixed-precision pre-training and fine-tuning;
• Distributed training using DDP, FSDP, Tensor Parallel, or Pipeline Parallel;
• Models with custom CUDA/Triton operators, complex loss functions, or multi-modal inputs;
• Large-scale training tasks that require long runtimes and cannot tolerate full restarts;
• Pipelines where training data is continuously added and may contain extreme samples.

Common Misconceptions

Misconception 1: Just Reduce the Learning Rate When NaN is Found

A high learning rate can indeed cause gradient explosion, but NaN can also come from illegal input domains, incorrect masks, zero denominators in normalization, low-precision overflow, corrupted samples, or custom kernels. Simply reducing the learning rate might temporarily mask the problem without eliminating the root cause.

Misconception 2: Gradient Clipping Can Fix NaN

Clipping only scales finite gradients proportionally. When gradients are already NaN/Inf, you should reject the update and locate the source, not continue clipping.

Misconception 3: No Need to Log After GradScaler Skips a Step

Occasional skips can be normal, but without evidence, you cannot distinguish between normal Scale probing and a systemic numerical fault. At a minimum, record the Scale, Step, Rank, batch, and consecutive skip count.

Misconception 4: Only Check on Rank 0

Anomalies might only appear on a specific data-parallel rank or pipeline stage. Only checking Rank 0 will miss local bad values and may allow the communication phase to propagate the anomaly to other devices.

Misconception 5: Keep Running Full Training with Anomaly Detection Enabled

Anomaly Detection is a diagnostic tool, not a normal monitoring solution. A better approach is to first lock down the bad batch with lightweight gates, and then use it for short-window replay.

Can BF16 Eliminate the Need for Non-finite Gates?

No. BF16 is less prone to overflow due to insufficient dynamic range compared to FP16, but illegal mathematical operations, anomalous inputs, gradient explosions, and custom kernel errors can still produce NaN/Inf.

Do I Need to Scan All Parameters to Confirm Checkpoint Safety?

Under normal conditions, you can use shard-level summaries, gradient bucket checks, and pre-save gates. After an anomaly occurs, a full finiteness scan of nearby checkpoints is safer. The frequency of per-parameter scans should be controlled to avoid significantly impacting training throughput.

Pre-Production Checklist

• Clear finiteness check boundaries are in place before input, loss, gradient, and parameter updates.
• GradScaler’s Scale, skip results, and consecutive skip count are included in metrics.
• All ranks are consistent in their skip, termination, or recovery actions for anomalous steps.
• The first bad step can be linked to data version, sample ID, configuration, code, and parent checkpoint.
• Replay is supported using a fixed Batch, RNG State, and environment fingerprint.
• Layer Bisect and target module hooks have switches and do not permanently slow down training.
• Suspect and corrupted checkpoints are not selected by the automatic recovery process.
• Scenarios with input NaN, gradient Inf, consecutive skips, and single-rank anomalies have been tested.

References

1. PyTorch, Automatic Mixed Precision examples
2. PyTorch, Autograd anomaly detection
3. PyTorch, clip_grad_norm_
4. NVIDIA Megatron Core, DistributedDataParallelConfig
5. NVIDIA Megatron FSDP, report_nan_in_param_grad
6. PyTorch, torch.isfinite