Fine-Tuning & Trainingpytorchhuggingfacetransformersaccelerate

Why your fine-tuning loss becomes NaN, and how to fix it

Error
loss: nan

Also appears as

  • RuntimeError: Function 'LogSoftmaxBackward0' returned nan values
  • grad_norm: nan

Short answer

NaN loss during training is most often caused by fp16 numeric overflow in gradients or activations, which bf16 avoids because of its wider exponent range. Other common causes are a learning rate spike (especially right after warmup), a small number of corrupt or malformed training samples, and unsafe division or log operations in a custom loss function. Switch to bf16 first if your hardware supports it, then check for corrupt samples and unstable LR.

Affects: Fine-tuning with fp16 mixed precision most commonly, also occurs with bf16 and full fp32 training

Fastest path to a stable loss

  1. 1If your GPU supports it (Ampere or newer), switch from fp16 to bf16 mixed precision, since bf16's wider exponent range eliminates most overflow-driven NaNs immediately.
  2. 2If you must use fp16, enable gradient clipping (max_grad_norm=1.0) and confirm the loss scaler is not stuck at a very high scale value.
  3. 3Lower the learning rate and/or lengthen the warmup period, since a rate spike right after warmup ends is a common trigger for the first NaN step.
  4. 4Scan your dataset for empty strings, all-padding sequences, or extremely long outlier sequences and filter them out before training.
  5. 5If using a custom loss function, add epsilon values to any division or log operation and clamp inputs to avoid log(0) or division by zero.

How to confirm this is your problem

  • Loss value logs as literally nan starting at a specific step, often early in training or right after a learning rate change
  • Gradient norm (grad_norm) also shows nan or inf at the same step the loss goes NaN
  • Training was stable for many steps and then suddenly diverges to NaN with no configuration change
  • Loss scaler value (if using fp16 with dynamic scaling) has grown very large just before the NaN appears
  • Model outputs become garbage or repetitive tokens immediately after the NaN step if training continues

Root causes and fixes

Most common

fp16 numeric overflow in activations or gradients

fp16 has a much smaller representable exponent range than fp32 or bf16. During backpropagation, gradient values or intermediate activations (especially in attention softmax or large logit values) can exceed fp16's maximum representable magnitude, overflowing to inf, and any subsequent arithmetic on an inf value produces NaN.

Fix: Switch to bf16 mixed precision on any Ampere-or-newer GPU (A100, H100, RTX 30-series and later), which has the same exponent range as fp32 and essentially eliminates this overflow class of NaN.

Commands
accelerate launch --mixed_precision bf16 train.py
Common

Learning rate spike, often right after warmup ends

If the learning rate schedule ramps up too aggressively or peaks higher than the model can tolerate at that point in training, a single large weight update can push activations into an unstable numeric range. This is especially common the first time the rate hits its full peak value right after a short warmup period.

Fix: Extend the warmup period (more steps at a lower rate before reaching peak LR) and lower the peak learning rate, particularly for full fine-tuning where sensitivity to LR spikes is higher than with LoRA.

Common

Corrupt or malformed training samples

A small number of pathological examples (empty strings after tokenization, sequences that are entirely padding with no real content, or encoding errors that produce out-of-vocabulary token ids) can produce degenerate loss values on the specific batch that contains them, and that single bad batch's NaN then propagates through the optimizer state to every subsequent step.

Fix: Filter the dataset for empty or near-empty sequences after tokenization, validate that all token ids fall within the tokenizer's actual vocabulary size, and log which batch index triggers the first NaN to identify the specific sample.

Commands
python -c "import datasets; ds=datasets.load_from_disk('train'); print(min(len(x['input_ids']) for x in ds))"
Occasional

Unsafe division or log operation in a custom loss function

Custom loss implementations (weighted cross-entropy, custom regularization terms, or manual perplexity calculations) sometimes divide by a count that can be zero for certain batches, or take log of a value that can reach exactly zero due to floating point underflow. Either operation produces NaN or inf that then contaminates the entire loss graph.

Fix: Add a small epsilon (like 1e-8) to any denominator or log argument in custom loss code, and use torch.clamp to keep values within a safe numeric range before the operation.

Rare

Gradient accumulation interacting badly with loss scaling

When using gradient accumulation with fp16 dynamic loss scaling, the scaler can grow the scale factor during a run of stable steps and then overflow on the very next step if the true gradient magnitude was borderline, particularly right at an accumulation boundary where accumulated gradients are finally unscaled and applied.

Fix: Reduce the initial loss scale or disable dynamic scale growth if this pattern recurs, or switch to bf16 which does not require loss scaling at all and removes this failure mode entirely.

Diagnostic commands

Check which precision mode is active

python -c "from transformers import TrainingArguments; print(TrainingArguments.fp16, TrainingArguments.bf16)"

Confirms whether fp16 or bf16 is active in your current TrainingArguments. If fp16 is True and your GPU supports bf16, switching removes the most common source of NaN outright.

Log gradient norm every step to find the exact failure point

python -c "print(trainer.state.log_history)"

Find the exact step where grad_norm first shows nan or inf. If it correlates with a learning rate schedule transition (end of warmup, restart), that points to an LR-driven spike rather than a data issue.

Isolate the batch that first produces NaN

python -c "for i,b in enumerate(dl):\n    out=model(**b)\n    if torch.isnan(out.loss): print(i); break"

If a specific batch index reliably reproduces the NaN, inspect that batch's raw text and token ids directly for empty sequences, encoding errors, or unexpectedly extreme lengths.

Check for inf or NaN in model weights after the failure

python -c "print(any(torch.isnan(p).any() for p in model.parameters()))"

If weights themselves already contain NaN, the run must be restarted from the last good checkpoint since NaN weights cannot recover on their own; if weights are still clean, the NaN is confined to the current batch's loss computation.

Stopping it from happening again

  • Default to bf16 mixed precision on any hardware that supports it instead of fp16.
  • Always enable gradient clipping (max_grad_norm around 1.0) as a cheap safeguard against rare gradient spikes.
  • Validate and clean datasets (no empty sequences, no out-of-vocabulary tokens, reasonable length bounds) before every training run, not just once at dataset creation.
  • Checkpoint frequently (every few hundred steps) so a NaN failure only costs a small amount of retraining time.
  • Add automated NaN detection that halts training and dumps the offending batch immediately rather than letting corrupted optimizer state silently continue.

When this becomes an architecture problem

If NaN losses recur across multiple datasets and configurations even after switching to bf16, clipping gradients, and cleaning data, the issue may be in a custom model modification, a specific base model's numerical sensitivity, or the training framework version itself, and is worth a deeper systematic review. If this is happening in a production fine-tuning pipeline that customers or internal teams depend on, building automated NaN detection and rollback into the pipeline is worth doing properly rather than firefighting each occurrence.

Frequently asked questions

Does switching from fp16 to bf16 fix NaN loss?

In most cases yes, because bf16 has the same exponent range as fp32 and does not overflow the way fp16 can during backpropagation. bf16 requires Ampere-generation or newer NVIDIA GPUs (A100, H100, RTX 30-series and later); older GPUs cannot use it and need other mitigations like gradient clipping and lower learning rates.

Can I recover a training run after the loss becomes NaN?

Not from the point of failure itself, since NaN values in the optimizer state and model weights do not resolve on their own. You need to resume from the last saved checkpoint before the NaN occurred and address the underlying cause (precision, learning rate, or data) before continuing.

How do I find which training example is causing NaN loss?

Log the batch index alongside the loss value at every step, then re-run inference on just that batch outside the training loop to isolate whether a specific sample (empty text, corrupt tokens, or extreme length) is responsible versus a general precision or learning rate issue.

Is gradient clipping enough to prevent NaN loss on its own?

Gradient clipping helps but does not fully prevent NaN, since it clips the gradient after it exists and cannot fix an already-overflowed fp16 value or a NaN produced inside a custom loss calculation. Combine clipping with bf16 precision and data validation for reliable protection.

Related problems

Training loss not decreasing during fine-tuning

Training loss that stays flat is most often caused by a LoRA adapter that targets the wrong modules (missing q_proj/k_proj/v_proj/o_proj), a learning rate that is too low for LoRA or too high and bouncing, or labels that were never masked so the model is trying to learn the prompt tokens as if they were random noise. Check target_modules first, then the label mask, then the learning rate.

Gradient checkpointing errors during fine-tuning

Gradient checkpointing errors during fine-tuning almost always come from three sources: the use_reentrant parameter left unset (it now must be explicit and False is usually correct for transformer models), an attention implementation that isn't fully compatible with checkpointing's re-computation approach, or leaving use_cache=True enabled while checkpointing is on, which conflicts because checkpointing recomputes the forward pass and a live KV cache assumes it won't be recomputed. Set use_reentrant=False and use_cache=False together.

Training dataset format errors during fine-tuning

Dataset format errors happen because the trainer expects a specific schema (either a messages list of role/content dicts, or a prompt/completion pair, or a single text field) and your JSONL doesn't match it, because samples are missing an EOS token so the model never learns to stop generating, or because a fixed max_length silently truncates long examples and cuts off labels partway through the intended response. Confirm your exact schema against what SFTTrainer or your data collator expects before training.

CUDA out of memory even though nvidia-smi shows free VRAM

This almost always means memory fragmentation: the allocator has enough total free memory but no single contiguous block large enough for the requested allocation, because the address space is broken into many small free-and-used segments from prior allocations of different sizes. The fix is enabling PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, reducing allocation size variability, or restarting the process to reset the address space.

Guide

Fine-Tuning Failure Modes: What Actually Goes Wrong

Fine-tuning failure modes that actually derail enterprise projects: catastrophic forgetting, eval overfitting, data leakage, and how to catch each one.

Guide

Building Fine-Tuning Datasets From Enterprise Data

Build fine-tuning datasets from enterprise data: instruction formats, deduplication methods, PII scrubbing, and quality filtering that actually works.

Guide

LoRA vs QLoRA: Choosing the Right Fine-Tuning Method

LoRA vs QLoRA for enterprise fine-tuning: rank and alpha choices, real VRAM math by model size, and when each method actually wins.

Guide

Multi-Node LLM Training Infrastructure: Networking and Storage

Multi-node LLM training infrastructure explained: InfiniBand vs RoCE tradeoffs, storage throughput needs, and cluster topology for enterprise fine-tuning.

Still stuck, or tired of fighting your own infrastructure?

Netray deploys and operates on-prem AI for regulated manufacturers and defense suppliers. We have debugged this stack in production, on air-gapped networks, at scale.