DTensorPolicyWorkerV2 sets reference policy from resumed weights on restart

[tommasobonomo]

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

RL/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Lines 306 to 316 in 9b331d2

	model_and_optimizer_state = setup_model_and_optimizer(
	config=config,
	tokenizer=self.tokenizer,
	runtime_config=runtime_config,
	distributed_context=distributed_context,
	checkpoint_manager=self.checkpoint_manager,
	is_vlm=self.is_vlm,
	init_optimizer=init_optimizer,
	weights_path=weights_path,
	optimizer_path=optimizer_path,
	)

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

RL/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Lines 332 to 335 in 9b331d2

	# Initialize reference model if requested
	self.reference_model_state_dict = None
	if init_reference_model:
	self.reference_model_state_dict = setup_reference_model_state(self.model)

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.

[CanadaApollo6]

I spent some time tracing where this regression originates, and I think it's an ordering inversion between the v1 and v2 workers rather than a logic error in any single function.

In v1 (dtensor_policy_worker.py) the reference is snapshotted before the resume checkpoint is loaded, so the get_cpu_state_dict(...) reference snapshot runs first, then self.load_checkpoint(weights_path, ...).

v2 refactored checkpoint loading into setup_model_and_optimizer (which applies weights_path internally via checkpoint_manager.load_checkpoint(...)), but the reference is still snapshotted after that call returns, in __init__ via setup_reference_model_state(self.model). So on resume, self.model already holds the resumed checkpoint, and the reference ends up equal to the resumed policy instead of the initial (model_name) policy. On a fresh run both orderings are equivalent, which is why this only surfaces on restart.

Net effect matches the report: with reference_policy_kl_penalty > 0, KL is computed against the resumed weights, so it collapses toward zero at the resume boundary and the anchor silently ratchets forward across restarts.

This is consistent with the #1709 hypothesis. If that's the PR that moved checkpoint loading into setup_model_and_optimizer, it's the point where the snapshot-before-load ordering got inverted. I'll confirm the exact commit.

I have a fix in progress, plus a minimal single-GPU repro (continuous vs checkpoint --> resume KL curves) and a regression test asserting the reference matches base rather than the resumed checkpoint. Before I open the PR, is there a preferred fix shape? My default is to capture the reference inside setup_model_and_optimizer from the base weights, before the weights_path overlay (restoring v1's ordering). The alternative is persisting the reference in the checkpoint and reloading on resume; more robust if the initial policy isn't reproducible from model_name, but a heavier change. Happy to go either direction.

[CanadaApollo6]

Opened #3048 with a fix and a regression test.

Confirmed the mechanism from my comment above: the ordering was inverted by the v2 init refactor in #1709 (4cbf8d58), which moved checkpoint loading inside setup_model_and_optimizer while the worker still snapshots the reference after that call returns. On a resume the reference comes out bit-identical to the resumed checkpoint (||ref - resumed_ckpt|| = 0.0 over all params, verdict details in the PR).

The fix restores the v1 ordering worker-side (snapshot from base weights, then load the checkpoint) with no API change; the alternative shape (returning the reference on ModelAndOptimizerState) is discussed in the PR description if maintainers prefer it.