DTensorPolicyWorkerV2 sets reference policy from resumed weights on restart
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Summary**
When resuming a training run with DTensorPolicyWorkerV2 (`policy.dtensor_cfg._v2: true`), the reference policy is set *after* the resume checkpoint has been loaded into the model.
Lines https://github.com/NVIDIA-NeMo/RL/blob/9b331d2da51c43158fb5030b882d592a7fa7e543/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py#L306-L316 carry out the setup of the model, loading a checkpoint if `weights_path` is provided.
Then, the reference model is set from that model in lines https://github.com/NVIDIA-NeMo/RL/blob/9b331d2da51c43158fb5030b882d592a7fa7e543/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py#L332-L335
This impacts training runs in GRPO with `loss_fn.reference_policy_kl_penalty > 0`, as the penalty will be computed w.r.t. the resumed checkpoint rather than the original policy.
**Steps to reproduce**
1. GRPO training run with `policy.dtensor_cfg._v2: true` and `loss_fn.reference_policy_kl_penalty > 0`.
2. Train, checkpoint after a small number of steps (e.g. 10), resume from the checkopint
3. Plot `train/kl_penalty`
**Expected behaviour**
Reference model is set to original policy, not resumed checkpoint. Plotting `train/kl_penalty` should result in a continuous curve, like in v0.4.0.
**Current behaviour**
Plotting `train/kl_penalty` shows it reset to 0 after every resume.
**Additional context**
Looks like this is a regression, it could have been introduced with #1709.
Contributor guide
Research direction
Start in nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py, especially lines 306-316 and 332-335, and trace how weights_path and the reference model are initialized. Reproduce the GRPO scenario with policy.dtensor_cfg._v2 enabled and loss_fn.reference_policy_kl_penalty > 0, then resume a checkpoint. Done means the reference model remains the original policy and train/kl_penalty stays continuous across resume.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100