NVIDIA-NeMo / NVIDIA-NeMo/RL

Qwen2.5VL-7B-Instruct model loss starts with nan

Open
#874 1 comment 1 reaction 0 assignees View on GitHub
bug t-pytdensor
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

For GRPO training with `Qwen2.5VL-7B-Instruct`, the loss at step 0 is `nan` and training does not proceed. token logprob errors are also consistently less than 1.05. Setting the KL loss to 0 seems to work (there are no nans). This problem does not occur with `Qwen2.5VL-3B-Instruct` or `Qwen2VL-7B-Instruct` (with `kl` penalty factor set to 0.01).

The problem is that the KL loss in `calculate_kl_penalty_joschu2020` function overflows for high/low values of `r`

Steps to reproduce
`uv run examples/run_vlm_grpo.py policy.model_name=Qwen/Qwen2.5VL-7B-Instruct`

The simplest fix that works for me is to clamp the value of the loss before computing KL

Image

KL loss is generally very unstable (and high variance) which has led researchers to explore other forms of Bregman-divergences as well
- https://proceedings.mlr.press/v155/ni21a.html
- https://arxiv.org/abs/2106.12112

For now, the simple fix I proposed above works.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.