NVIDIA-NeMo / NVIDIA-NeMo/RL

Flaky test_value_worker_train_decreases_loss: tiny Qwen2 value head trains with HF's default 0.1 dropout

Open
#3,999 0 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
accuracy bug community-request waiting-on-maintainers
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Describe the bug**

`tests/unit/models/value/test_dtensor_value_worker.py::test_value_worker_train_decreases_loss` fails intermittently in the default unit-test suite, about 1 run in 20:

```
AssertionError: Value loss should not increase after 3 steps; got [0.018786393105983734, 0.019307831302285194, 0.020152460783720016]
assert 0.020152460783720016 <= (0.018786393105983734 + 0.001)
```

The regression value head is a `Qwen2ForTokenClassification` head. `GenericForTokenClassification` builds `nn.Dropout(classifier_dropout)` and falls back to 0.1 when the model config defines neither `classifier_dropout` nor `hidden_dropout`; `Qwen2Config` defines neither, so the tiny Qwen2 test checkpoint trains its value head under a 0.1 dropout nobody configured. The test trains three steps at `lr=5e-6` on a fixed batch and asserts `losses[-1] <= losses[0] + 1e-3`. The genuine change over two steps is about -2.6e-4, while the per-step dropout noise has a standard deviation of about 6.4e-4, so the tolerance sits inside the noise and the assertion fails in a ~5% tail.

**Steps/Code to reproduce bug**

On a 2-GPU machine, run the test repeatedly with the default suite's invocation:

```bash
for i in $(seq 1 100); do
uv run --group test bash tests/run_unit.sh unit/models/value/test_dtensor_value_worker.py::test_value_worker_train_decreases_loss || echo "run $i FAILED"
done
```

Observed: 5 failures in 100 runs (transformers 5.5.0, 2× H100, CI image). Printing `losses` in the test shows the loss rising between step 1 and step 3 in 38 of 100 runs; 33 of those stay inside the 1e-3 tolerance by chance.

Setting `classifier_dropout=0.0` on the `Qwen2Config` that `tests/unit/conftest.py::tiny_qwen2_model_path` saves removes the noise: 0 failures in 100 runs, monotone decrease in every run, per-step standard deviation about 3.7e-5.

**Expected behavior**

The test should be deterministic given the parameters: with a fixed batch and a fixed head, three optimiser steps at `lr=5e-6` decrease the loss every time, and the assertion should never depend on a dropout mask.

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.