Debug NaN Loss in QLoRA

Training a QLoRA model and suddenly your loss prints nan — or worse, it silently becomes nan after a few hundred steps. You're not alone. NaN (Not a Number) loss is one of the most frustrating issues in LLM finetuning because the error messages are cryptic and the causes are many. But unlike other debugging puzzles, debug nan loss during qlora training follows a predictable pattern. This lesson will give you a mental model and a step-by-step plan to find and fix NaN loss fast, so you can get back to training a model that actually converges.

The problem this lesson solves

NaN loss means your model's weights have become nan or inf somewhere during backpropagation. When that happens, the optimizer has nothing real to update, so the model is effectively ruined — you have to stop and restart. The pain is real: hours of GPU time wasted, a finetuning pipeline that breaks at 2 AM, and no clear idea where to start.

QLoRA adds a special twist. It quantizes the base model weights to 4-bit (using bitsandbytes), which saves memory but introduces new failure modes. Mixed precision training (fp16 or bf16) amplifies numerical issues. And LoRA itself can push the model into unstable territory if the learning rate is too high. So the problem is not just “one” bug — it's a family of issues, and you need a debug strategy that works for QLoRA specifically.

By the end of this lesson, you'll be able to spot the usual suspects, apply targeted fixes, and verify that your training run is stable — without spending days guessing.

Core concept / mental model

Think of NaN loss as a chain reaction gone wrong. The chain is:

1. Forward pass — you feed batches of tokens through the model.
2. Loss calculation — cross-entropy loss compares predictions to labels.
3. Backward pass — gradients are computed for every parameter.
4. Optimizer step — weights are updated using those gradients.

If any link in that chain produces a nan or inf, the whole chain collapses. The classic culprits are:

• Learning rate too high — gradients explode, weights become inf, then nan.
• Mixed precision issues — fp16 has a limited range; small gradients can underflow to 0, large ones overflow to inf.
• Bad data — sequences with all-special-token masks, or labels that are all -100 (ignored tokens), produce a loss of 0/0 which is nan.
• Quantization bugs — 4-bit quantized base model (nf4 or fp4) may behave strangely in some hardware or library versions.
• Long sequences — gradients can accumulate and overflow during backprop through many layers.

Here's a simple mental picture:

Imagine you're driving a car. NaN loss is the check-engine light. The light doesn't tell you which part failed, but it tells you something is wrong. Your job is to look at the engine (the training loop), the fuel (data), and the tires (precision) to find the root cause.

In QLoRA, the engine is mostly frozen (quantized base), and you're only tuning a small set of LoRA adapters. That's good news — the problem is often in the hyperparameters or data rather than the model architecture.

How it works step by step

When you hit NaN loss, follow this ordered checklist. It goes from the cheapest (check the data) to the most invasive (change the model).

1. Reproduce and capture the moment

Don't change anything yet. First, your debugging goal is to see when it happens. Is it at step 1? Step 500? Only after a device warm-up? Use logging_steps and checkpoints to capture the context.

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="./qlora-out",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    logging_steps=10,
    save_steps=500,
    fp16=True,
)

2. Look at the data

NaN loss often comes from the data more than the model. Check for: - Empty sequences after tokenization - Labels all -100 (ignore index) - Sequences longer than the model's max_length causing truncation to nothing - Very imbalanced classes if you're doing classification

3. Check the loss function

If labels are all -100, the mean loss over zero valid tokens is 0/0 = nan. This is a classic.

4. Verify precision settings

Mixed precision (fp16=True) can cause overflow if your loss is large. Try bf16=True if your GPU supports it — bf16 has a wider range and is much more forgiving.

5. Look at the gradient norms

If you see grad_norm spiking to inf just before NaN, that's an exploding gradient problem. Clip gradients with max_grad_norm (default is 1.0, but you can lower it).

6. Try a smaller learning rate

For QLoRA, the learning rate is often between 1e-4 and 5e-4. If you're using 1e-3 or higher, NaN is almost guaranteed.

7. Re-check the QLoRA setup

Sometimes the issue is in the quantization layer. Update bitsandbytes to the latest version, check that your GPU is compatible (e.g., no nf4 issues on older GPUs), and try fp4 instead of nf4.

Hands-on walkthrough

Let's put this into practice with a full QLoRA training script. We'll add instrumentation to catch the NaN early.

Setup and imports

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
)
from peft import LoraConfig, get_peft_model

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    "your-base-model",
    quantization_config=bnb_config,
    device_map="auto",
)

tokenizer = AutoTokenizer.from_pretrained("your-base-model")

Detect NaN in the training loop

Instead of waiting for the trainer to crash, add a simple check. Here's how to log a warning when loss becomes NaN:

import torch
from transformers import Trainer

class NaNCheckTrainer(Trainer):
    def training_step(self, model, inputs):
        loss = super().training_step(model, inputs)
        if torch.isnan(loss) or torch.isinf(loss):
            self.control.should_log = True  # force a log
            print(f"⚠️ NaN loss detected at step {self.state.global_step}")
            # Optional: raise an exception to stop training
            # raise ValueError("NaN loss!")
        return loss

trainer = NaNCheckTrainer(
    model=model,
    args=args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)
trainer.train()

Expected output (when things work):

{'loss': 1.7234, 'learning_rate': 0.0002, 'epoch': 0.02}
{'loss': 1.5121, 'learning_rate': 0.0002, 'epoch': 0.04}

If you see 'loss': nan with a warning, you know exactly when it happened.

Fix the data before training

Prevent the all--100 label trap:

from datasets import Dataset

def filter_bad_sequences(dataset):
    """Remove examples where all labels are -100, which cause NaN loss."""
    def has_valid_label(batch):
        # Labels are typically 1D arrays of ints, -100 means ignore
        labels = batch["labels"]
        return any(l != -100 for l in labels)

    return dataset.filter(has_valid_label)

filtered_ds = filter_bad_sequences(your_dataset)
print(f"Filtered {len(your_dataset) - len(filtered_ds)} bad examples")

Try bf16 instead of fp16

args = TrainingArguments(
    ...
    bf16=True,  # instead of fp16=True
)

On supported hardware (Ampere or newer), bf16 solves many overflow problems.

Compare options / when to choose what

Troubleshooting & edge cases

Here are the common failure modes and their fixes, straight from real QLoRA runs.

All labels are -100 (data issue)

Symptom: Loss is nan from step 1. Cause: Your tokenizer produces a label mask that has no non--100 tokens for that batch. Fix: Use a data collator that handles padding properly, e.g., DataCollatorForLanguageModeling or DataCollatorForSeq2Seq. Or filter the dataset.

Gradient overflow (training instability)

Symptom: Loss drops normally, then jumps to inf and turns to nan over a few steps. Cause: Learning rate too high for QLoRA; or long sequences accumulate gradients. Fix: Reduce learning_rate to 1e-4 or lower, and lower max_grad_norm to 0.5. Also consider increasing gradient_accumulation_steps instead of raising batch size.

bf16/fp16 mixed precision glitches

Symptom: NaN only when fp16=True, but works in fp32. Cause: Values above 65,504 in fp16 become inf. Fix: Switch to bf16=True, or reduce the loss scale (HF's default is okay, but you can try fp16_opt_level="O1").

Quantization kernel issues

Symptom: NaN after several steps, sometimes random. Cause: Old bitsandbytes version or incompatibility with your GPU. Fix: Update bitsandbytes, use quant_type="fp4", or switch to bnb_4bit_quant_type="nf4" if you're on fp4. Also ensure bnb_4bit_compute_dtype matches your compute dtype.

Sequence length is too long

Symptom: Loss becomes nan only on very long sequences in the batch. Cause: Activations/gradients overflow in fp16. Fix: Truncate tokenizer to max_length=2048 (or lower), and enable gradient_checkpointing=True to reduce memory pressure.

What you learned & what's next

You now understand the core idea behind debug nan loss during qlora training — it's a chain reaction from data to precision to optimizer. You've applied a practical exercise with a NaNCheckTrainer and a data-filtering function. You know how to compare options like learning rate, gradient clipping, bf16, and quantization tweaks. And you have a concrete troubleshooting checklist for the most common edge cases.

These skills directly map to the learning objectives: explaining the core idea and completing a practical exercise. You also practiced connecting the dots to the rest of the QLoRA workflow.

Next in the track, you'll learn how to evaluate your QLoRA model — checking whether the finetuned model actually improved on your target task, using metrics like perplexity and accuracy. That's the natural next step after you've stabilized training and saved your checkpoints.

Now go forth — and may your loss curves be smooth and your gradients finite!