[yellowbee686]

What does this PR do?

Fix policy-gradient distillation when rollout importance weights arrive as a NestedTensor.

Before calling the PPO policy loss, distillation_loss() converts log_prob, old_log_prob, advantages and response_mask from NestedTensor to padded dense tensors, but passes rollout_is_weights through unchanged. On the no-padding path the weights stay jagged and the policy loss fails. In the forward pass it raises ValueError: NestedTensor where_self(...): expected condition to be a jagged layout NestedTensor, and in backward it raises NotImplementedError: aten._nested_sum_backward.default.

This PR pads nested weights with 0.0 before the policy loss, matching the other dense inputs. Padded positions are already excluded by response_mask. Dense or absent weights are unchanged.

Checklist Before Starting

• Search for similar PRs: distillation nested, distillation importance weights. I also checked every open PR that touches verl/trainer/distillation/losses.py. [trainer, megatron, fsdp] fix: forward-KL topk distillation shape mismatch with pad_bshd_to_minibatch_max #8061 passes forced_max_seqlen, [algo, worker] fix: handle padding-only top-k distillation microbatches #7887 handles empty top-k micro-batches, and [trainer, fsdp, worker, algo, cfg, recipe] feat: add reverse-KL on-policy distillation with FSDP teacher #6920, [worker, perf] feat: defer scalar metric materialization #7096, [algo, fsdp, megatron, veomni] feat: add tail-aware top-k KL for on-policy distillation #7855 and others add losses or metrics. None converts rollout_is_weights, so this is not a duplicate.
• Format the PR title as [{modules}] {type}: {description}.

Test

Run from the repository root on CPU (Python 3.12, PyTorch 2.13):

PYTHONPATH=. python -m pytest tests/workers/test_distillation_policy_gradient_on_cpu.py -q

• With this PR: 12 passed.
• The same test file against current main: 4 failed, 8 passed, with the two errors above.

The tests call the real distillation_loss(), padding conversion and PPO policy loss. They compare loss values and student gradients against independent dense inputs, and confirm that the teacher receives no gradient. Coverage includes unequal response lengths, non-unit and zero weights, internal masks, dense and absent weights, and both token-mean and seq-mean-token-mean.

API and Usage Example

No API change.

Design & Code Changes

• verl/trainer/distillation/losses.py: two lines that pad nested rollout_is_weights before policy_loss_fn.
• tests/workers/test_distillation_policy_gradient_on_cpu.py: numerical and gradient regressions.

AI assistance (OpenAI Codex and Cursor) was used to write the code, tests and this description. This is a draft: the submitter is reviewing every changed line and will mark it ready after that review.

Checklist Before Submitting

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit run --files <changed files> passed all hooks.
• Add / Update the documentation: no user-facing change.
• Add unit tests: tests/workers/test_distillation_policy_gradient_on_cpu.py is collected by cpu_unit_tests.yml.
• Once your PR is ready for CI, send a message in the ci-request channel.
• Not related to the recipe submodule.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}

[cla-assistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

huangjunyi.0 seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}