modelscope / modelscope/ms-swift

[bug] grad explode to nan while training Qwen 3 Reranker with tuner unsloth

Open
#6,992 1 comment 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

stale
Dominant language
Python
Stars
15.7k
Forks
1.7k
Avg merge
1d 16h
Merged PRs (30d)
136

Description

Describe the bug
I encountered a gradient explosion issue resulting in NaN values when training Qwen/Qwen3-Reranker-8B (Listwise Generative Reranker) using ms-swift with the unsloth backend.

Symptoms:

  1. Normal Run: If I run the training normally, the training loop continues, but the logs report grad_norm = nan. The loss value still fluctuates/decreases, but the gradients are clearly broken.
  2. Debug Run: When enabling torch.autograd.set_detect_anomaly(True) and adding a debug hook to sdpa, the training crashes with RuntimeError: Function 'MmBackward0' returned nan values.
  3. Gradient Analysis: By hooking into _debug_sdpa, I observed the gradients flowing backward into the Attention layer growing exponentially from 1e-6 to 1e+36 within a single step, eventually hitting Infinity/NaN.

Reproduction Script:

#!/bin/bash
# Listwise training configuration
CACHE_DIR="./webq_cached_dataset"
export CUDA_VISIBLE_DEVICES=1
export PYTORCH_ALLOC_CONF=expandable_segments:True
export UNSLOTH_COMPILE_DISABLE=1

echo ">>> Start SFT..."

uv --preview-features extra-build-dependencies run swift sft \
    --model Qwen/Qwen3-Reranker-8B \
    --task_type generative_reranker \
    --loss_type listwise_generative_reranker \
    --train_type lora \
    --tuner_backend unsloth \
    --torch_dtype bfloat16 \
    --learning_rate 2e-4 \
    --gradient_accumulation_steps 16 \
    --lora_rank 64 \
    --lora_alpha 128 \
    --dataset 'MTEB/scidocs-reranking' \
    # ... (rest of parameters as per context)

Debug Log & Stack Trace:
Using torch.autograd.set_detect_anomaly(True) and with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]):, the gradients explode as follows:

# logs from `output.register_hook(check_grad_output)` for  output of F.sdpa
[SDPA Backward] grad Max: 8.046627044677734e-06, Min: -6.794929504394531e-06
...
[SDPA Backward] grad Max: 566935683072.0, Min: -807453851648.0
...
[SDPA Backward] grad Max: 2.648071397852762e+36, Min: -2.1599954931504883e+36
# error raised
RuntimeError: Function 'MmBackward0' returned nan values in its 0th output.
  File "/home/User/workspace/ir_train/.venv/lib/python3.12/site-packages/accelerate/accelerator.py", line 2740, in backward
    loss.backward(**kwargs)
  File "/home/User/workspace/ir_train/.venv/lib/python3.12/site-packages/unsloth_zoo/gradient_checkpointing.py", line 598, in backward
    torch.autograd.backward(outputs_with_grad, args_with_grad)

Your hardware and system info

  • GPU: NVIDIA A100 (40GB)
  • CUDA: 12.8
  • System: Linux
  • Python: 3.12
  • Dependencies (from uv config):
    • torch==2.9.0 (installed via index pytorch-cu128)
    • unsloth (latest git)
    • ms-swift<=3.10.3
    • accelerate==1.11
    • deepspeed>=0.18.2

Additional context

  • The issue seems specific to the combination of Unsloth backend and the Listwise Generative Reranker task.
  • The explosion happens during the backward pass after the attention mechanism (gradients coming from the upper layers into SDPA are already huge).
  • attn_impl was set to sdpa for debugging, but flash-attn exhibits the same behavior (NaN grad_norm).

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running the provided swift sft reproduction with the listed Qwen3-Reranker, PyTorch, Unsloth, and ms-swift configuration. Compare anomaly-detection output and _debug_sdpa gradient logs, then trace the backward path through accelerate and unsloth_zoo/gradient_checkpointing.py; done means identifying the source of the NaN gradients and confirming a stable training run.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.