NVIDIA-NeMo / NVIDIA-NeMo/RL

DTensorPolicyWorkerV2 sets reference policy from resumed weights on restart

Open
#2,955 2 comments 0 reactions 0 assignees View on GitHub
bug community-request
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.

Image

**Current behaviour**

Plotting `train/kl_penalty` shows it reset to 0 after every resume.

Image

**Additional context**

Looks like this is a regression, it could have been introduced with #1709.

Contributor guide

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.