fix: capture GRPO reference policy from base weights, not resumed checkpoint

[CanadaApollo6]

What does this PR do ?

Fixes the v2 dtensor worker snapshotting the GRPO KL reference policy after the resume checkpoint has been applied, which silently resets the reference to the resumed policy on every restart and zeroes the KL penalty.

Issues

Fixes #2955

Root cause

DTensorPolicyWorkerV2.__init__ snapshots the reference after setup_model_and_optimizer() returns — but since the v2 init refactor in #1709 (commit 4cbf8d58), that call applies the resume checkpoint internally (automodel/setup.py:855). The v1 worker has the correct ordering: snapshot at dtensor_policy_worker.py:450 before load_checkpoint at :496.

The fix

One file, no API change. When init_reference_model and a weights_path are both requested, the worker calls setup_model_and_optimizer(..., weights_path=None, optimizer_path=None), snapshots the reference from the base weights, then issues the same checkpoint_manager.load_checkpoint(model=..., weights_path=..., optimizer=..., optimizer_path=..., scheduler=...) call setup would have made. This restores the v1 ordering. All other paths (no reference model, fresh start) are unchanged.

Alternative shape considered: capture the reference inside setup_model_and_optimizer and return it on ModelAndOptimizerState. Rejected for blast radius: that NamedTuple is consumed via strict positional unpacks by both the v2 policy worker and the v2 value worker (dtensor_value_worker_v2.py), and a second, unrelated ModelAndOptimizerState exists for megatron. Happy to switch shapes if maintainers prefer it.

Evidence

Regression test (red on main, green with this fix)

The included test trains a tiny llama, checkpoints, resumes with init_reference_model=True, and fingerprints the reference via reference-policy logprobs on a fixed batch.

On main:

AssertionError: reference policy tracks the resumed checkpoint (#2955):
max|ref-base|=6.998e-01, max|ref-resumed|=0.000e+00
======================= 1 failed, 25 warnings in 41.27s ========================

The reference logprobs are bit-identical to the resumed checkpoint's (max|ref-resumed| = 0.0).

With this fix:

PASSED tests/unit/models/policy/test_dtensor_worker_v2.py::test_dtensor_v2_reference_policy_ignores_resume_checkpoint
======================= 1 passed, 25 warnings in 48.32s ========================

The resumed worker's log also shows the corrected ordering directly:

(DTensorPolicyWorkerV2) No weights path provided. Loaded base HF weights via from_pretrained (default policy init)
(DTensorPolicyWorkerV2) Loading weights from /tmp/.../policy/weights   <- deferred load, after the reference snapshot

Weight-space provenance (diagnostic instrumentation, not part of this PR)

L2 distance from the reference snapshot to the base weights vs. the resumed checkpoint, measured at worker init on a resume (Qwen2.5-0.5B / tiny-llama, all 16,462,144 compared params):

Before (main):

[#2955][rank0] ||ref - base||            = 5.172137e+01
[#2955][rank0] ||ref - resumed_ckpt||    = 0.000000e+00
[#2955][rank0] VERDICT: reference policy tracks CHECKPOINT (BUG #2955)

After (this fix):

[#2955][rank0] ||ref - base||            = 0.000000e+00
[#2955][rank0] ||ref - resumed_ckpt||    = 1.103226e+02
[#2955][rank0] VERDICT: reference policy tracks BASE policy (correct)

Both distances are exact zeros in opposite directions: before the fix the reference is the resumed checkpoint; after, it is the base policy.

A training-level KL-curve comparison (continuous vs. interrupt-and-resume GRPO run) is in progress and can be attached as a follow-up comment if useful.

Local evidence environment: torch 2.11.0+cu130 (CUDA 13.0), 1x RTX 3080, Python 3.13, buggy baseline = main @ 5ae12621; re-validated after rebase onto 78a581bb.

Usage

uv run --extra automodel --group test pytest \
  "tests/unit/models/policy/test_dtensor_worker_v2.py::test_dtensor_v2_reference_policy_ignores_resume_checkpoint" \
  --automodel-only --hf-gated

Needs 1 GPU and an HF token with meta-llama/Llama-3.2-1B access (tokenizer only; the model under test is a tiny random-weight llama built by the existing fixture).

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? (the new regression test red→green; details above)
• Did you add or update any necessary documentation? (n/a — behavior fix restoring documented v1 semantics)

Additional Information

The commit is DCO signed off. The regression test reuses the existing test_dtensor_worker_v2.py fixtures and adds a single-GPU cluster fixture, so it runs on any 1-GPU machine.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[hemildesai]

/claude review

[CanadaApollo6]

Hey @hemildesai noticed the review command you sent didn't seem to fire. Pinging so you have visibility