Flaky test_value_worker_train_decreases_loss: tiny Qwen2 value head trains with HF's default 0.1 dropout
- 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
Assessment
This issue has not been assessed yet.