Add On-Policy Self-Distillation (OPSD) via --opd-type self - #1
Open
HJSang wants to merge 6 commits into
Open
Conversation
Implements "Self-Distilled Reasoner" style OPSD in the Megatron backend. A single model is both student and teacher, differing only by context: the student sees the problem; the frozen teacher (initial-policy checkpoint) also sees privileged information (the ground-truth solution) and scores the student's response tokens in one forward pass. The loss is a direct, full-vocab token-level generalized JSD over the student's on-policy tokens, with per-token clipping and no task reward. - ppo_utils: vocab-parallel generalized JSD (compute_vocab_parallel_jsd) with a differentiable all-reduce; gradients flow only through the student. - loss: get_response_logits extractor, opsd_loss_function, and "opsd" dispatch. - actor: privileged teacher forward (compute_teacher_response_logits) and train_actor wiring that skips the RL advantage/ref/old-logprob passes. - rollout/data: build the teacher's [prompt + privileged_info + response] sequence; new Sample.privileged_info and --opsd-privileged-info-key. - args: --opd-type self, --opsd-beta, --opsd-jsd-clip; validation. - docs + example script (run-qwen3-8B-opsd.sh). - tests: CPU unit tests for the JSD core (value, gradient, non-negativity). MVP limitations: context-parallel-size must be 1; teacher response logits are held over the full vocab between forwards (memory ~ response_len x vocab).
…data - Match the OPSD/GOLD generalized-JSD mixture convention exactly: M = (1-beta)*student + beta*teacher; jsd = beta*KL(teacher||M) + (1-beta)*KL(student||M). Endpoints unchanged (beta=0 -> KL(teacher||student), beta=1 -> KL(student||teacher)). - tests/test_opsd_jsd.py: reference now uses F.kl_div, independently matching OPSD. - tests/test_qwen2.5_0.5B_opsd.py: end-to-end Megatron smoke test (--opd-type self), registered in the megatron CI matrix (regenerated pr-test.yml from the template). - Example + docs: use the reference dataset siyanzhao/Openthoughts_math_30k_opsd (fields: problem, solution).
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between 5c5ea38e00a71ad6930aec1a8f30c95fda84e78f and 5396f18. 📒 Files selected for processing (11)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThis PR implements On-Policy Self-Distillation (OPSD): adds data/CLI support for privileged info, generates teacher tokens in rollout, computes distributed vocab-parallel JSD loss, runs frozen teacher forward in actor, validates with CPU and distributed tests, documents self-mode workflow, provides Qwen3/Qwen2.5 examples, and updates CI test matrix. ChangesOn-Policy Self-Distillation (OPSD) Implementation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 080dd3fff1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| data_iterator, | ||
| num_microbatches, | ||
| store_prefix="teacher_", | ||
| rollout_data.update(self.compute_teacher_response_logits(rollout_data, num_microbatches)) |
There was a problem hiding this comment.
This stores teacher_response_logits in rollout_data before log_rollout_data(...) runs below. I checked slime/backends/megatron_utils/data.py: the logger only skips a fixed list of keys and otherwise torch.cat(...).clone()s tensor lists, so OPSD duplicates the full [response_tokens, vocab_shard] teacher cache on TP0 just to log a meaningless mean. For realistic OPSD responses/vocabs this can add multiple GB and OOM before training; skip this key from rollout logging or log only a scalar summary.
Useful? React with 👍 / 👎.
| response_tokens = list(sample.tokens[prompt_length:]) | ||
| privileged_info = sample.privileged_info or "" | ||
| privileged_tokens = self.tokenizer.encode(privileged_info, add_special_tokens=False) | ||
| teacher_tokens.append(prompt_tokens + privileged_tokens + response_tokens) |
There was a problem hiding this comment.
Appending privileged_tokens here makes the teacher sequence potentially much longer than the student sequence, but _split_train_data_by_dp still builds the DP/microbatch schedule from data["tokens"] and that schedule is reused for the teacher forward. With --use-dynamic-batch-size, several short student samples whose privileged solutions are long can be packed into one teacher forward, violating max_tokens_per_gpu and causing OOMs or sequence-length overruns. Include the teacher token lengths in the OPSD packing cap, or build a teacher-specific schedule.
Useful? React with 👍 / 👎.
| if "teacher" in self.weights_backuper.backup_tags: | ||
| if self.args.use_routing_replay: | ||
| os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough" | ||
| opsd = self.args.use_opd and self.args.opd_type == "self" |
There was a problem hiding this comment.
When users combine OPSD with --disable-compute-advantages-and-returns, this gate skips the only call that populates teacher_response_logits, while argument validation still forces loss_type to opsd; the training loss then raises ValueError. Since OPSD is documented as pure distillation with no advantages, this flag is a plausible configuration, so either disallow it for OPSD or run the teacher forward independently of advantage computation.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/test_opsd_jsd.py (1)
1-70:⚠️ Potential issue | 🟠 Major | ⚖️ Poor tradeoffOPSD JSD unit tests are missing from the CPU CI matrix.
These CPU-only tests (
test_opsd_jsd.py) are not listed in thecpu-unittestjob matrix in.github/workflows/pr-test.ymlor.github/workflows/pr-test.yml.j2, so they will never run in CI. This defeats the purpose of having unit tests for the JSD computation.Add the following entry to the
cpu-unittesttests list in.github/workflows/pr-test.yml.j2:{'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0}, + {'test_file': 'test_opsd_jsd.py', 'num_gpus': 0}, {'test_file': 'utils/test_hf_checkpoint_saver.py', 'num_gpus': 0},Then regenerate
.github/workflows/pr-test.ymlviagenerate_github_workflows.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_opsd_jsd.py` around lines 1 - 70, The CPU-only unit tests in tests/test_opsd_jsd.py are not included in the cpu-unittest job matrix, so add an entry for tests/test_opsd_jsd.py to the cpu-unittest tests list in .github/workflows/pr-test.yml.j2 (so the test is run by CI), then run the repository's workflow generator script generate_github_workflows.py to regenerate .github/workflows/pr-test.yml; ensure the added test path matches the test file name tests/test_opsd_jsd.py and commit both the updated .j2 template and the regenerated .yml.
🧹 Nitpick comments (3)
examples/on_policy_distillation/run-qwen3-8B-opsd.sh (1)
150-160: 💤 Low valueQuote array expansions to prevent word splitting.
Shellcheck flags unquoted
${ARRAY[@]}expansions (lines 150–160). While this works in the current script, unquoted expansions can cause issues if array elements contain spaces or glob characters. Bash best practice is"${ARRAY[@]}".🛠️ Recommended fix
-- python3 train.py \ --actor-num-nodes 1 \ --actor-num-gpus-per-node 2 \ --rollout-num-gpus 4 \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${OPSD_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${SGLANG_ARGS[@]} \ - ${MISC_ARGS[@]} \ - ${RM_ARGS[@]} + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${OPSD_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + "${RM_ARGS[@]}"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/on_policy_distillation/run-qwen3-8B-opsd.sh` around lines 150 - 160, The script is expanding several argument arrays without quotes which can lead to word splitting or globbing; update each array expansion in the final command invocation—MODEL_ARGS, CKPT_ARGS, ROLLOUT_ARGS, OPTIMIZER_ARGS, OPSD_ARGS, WANDB_ARGS, PERF_ARGS, EVAL_ARGS, SGLANG_ARGS, MISC_ARGS, and RM_ARGS—to use the safe form with quoted expansions (i.e., replace ${ARRAY[@]} with "${ARRAY[@]}" for each listed array) so elements containing spaces or special characters are preserved.docs/en/advanced/on-policy-distillation.md (1)
113-113: ⚡ Quick winClarify why
--context-parallel-size 1is required.The limitation states "OPSD requires
--context-parallel-size 1" but does not explain why. Users may wonder if this is a temporary implementation detail or a fundamental constraint. Adding one sentence explaining the reason (e.g., "due to the current teacher logit storage implementation") would help set expectations and guide future workarounds.📝 Suggested clarification
-> **Limitations (current MVP)**: OPSD requires `--context-parallel-size 1`. The teacher response logits are held over the full vocabulary between the teacher and student forwards, so memory scales with `response_length × vocab_size`; prefer it for smaller models / shorter responses for now. The example script is `examples/on_policy_distillation/run-qwen3-8B-opsd.sh`. +> **Limitations (current MVP)**: OPSD requires `--context-parallel-size 1` (because the current implementation stores full-vocab teacher logits without CP sharding). The teacher response logits are held over the full vocabulary between the teacher and student forwards, so memory scales with `response_length × vocab_size`; prefer it for smaller models / shorter responses for now. The example script is `examples/on_policy_distillation/run-qwen3-8B-opsd.sh`.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/en/advanced/on-policy-distillation.md` at line 113, Clarify why OPSD requires the flag `--context-parallel-size 1` by adding a short sentence after the limitation line: state that this is due to the current teacher logit storage implementation which retains full-vocabulary teacher response logits across teacher and student forwards (so memory scales with response_length × vocab_size), and mark it as an implementation limitation rather than a fundamental algorithmic constraint so readers know it may be relaxed in future work.slime/utils/ppo_utils.py (1)
240-297: 💤 Low valueConsider moving the import to the top of the file.
The
import mathat line 279 inside the function body is functional but unconventional. Moving it to the module-level imports improves readability and follows Python conventions.import torch import torch.distributed as dist import torch.nn.functional as F +import mathThen remove the inline import at line 279.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@slime/utils/ppo_utils.py` around lines 240 - 297, The inline "import math" inside compute_vocab_parallel_jsd should be moved to module-level imports to follow Python conventions; add "import math" at the top of the file alongside the other imports and remove the inline import statement within compute_vocab_parallel_jsd (the function defined as compute_vocab_parallel_jsd in slime/utils/ppo_utils.py), leaving the math usage (math.log) unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@slime/backends/megatron_utils/loss.py`:
- Around line 1102-1127: The zip between the student response generator
(get_responses(...) call) and teacher_response_logits uses strict=False which
can silently drop samples on count mismatch; change the zip call in the loop
that iterates over (student_logits_chunk, _), teacher_logits_chunk to use
strict=True so mismatched lengths raise an error (i.e., replace strict=False
with strict=True in the zip invocation that pairs get_responses(...) and
teacher_response_logits).
In `@slime/utils/arguments.py`:
- Line 1828: A long f-string containing the text ending with "field holding the
privileged information (ground-truth solution) shown to the teacher." is split
across lines and fails Black; fix it by placing the entire f-string on a single
line or wrapping it with parentheses and splitting into multiple adjacent
f-strings (or using a parentheses-wrapped f-string with implicit concatenation)
so it conforms to Black formatting rules—locate and update the offending
f-string in slime/utils/arguments.py (the string that starts/ends with the
quoted fragment shown in the diff) and ensure the result is syntactically the
same but formatted as a single-line f-string or properly parenthesized
multi-line f-string.
In `@slime/utils/data.py`:
- Line 262: When building the privileged_info field avoid letting a bare
KeyError surface: check if privileged_info_key is not None and not in data
before accessing data[privileged_info_key], and raise a clear error (e.g. raise
KeyError(f"Missing privileged_info_key '{privileged_info_key}' in dataset;
available keys: {list(data.keys())}")) or set privileged_info to None depending
on desired behavior; update the assignment that currently reads
privileged_info=data[privileged_info_key] if privileged_info_key is not None
else None to perform this presence check and raise the informative KeyError when
the key is absent.
In `@tests/test_qwen2.5_0.5B_opsd.py`:
- Around line 21-27: The file fails black formatting on the multi-line string
constructing ckpt_args (the tuple assigned to ckpt_args using MODEL_NAME and
torch_dist_ckpt); run the formatter and commit the change: run `black
tests/test_qwen2.5_0.5B_opsd.py` (or your project's pre-commit formatter) to
reformat the file so the ckpt_args assignment matches black style, then add and
commit the formatted file so CI pre-commit checks pass.
---
Outside diff comments:
In `@tests/test_opsd_jsd.py`:
- Around line 1-70: The CPU-only unit tests in tests/test_opsd_jsd.py are not
included in the cpu-unittest job matrix, so add an entry for
tests/test_opsd_jsd.py to the cpu-unittest tests list in
.github/workflows/pr-test.yml.j2 (so the test is run by CI), then run the
repository's workflow generator script generate_github_workflows.py to
regenerate .github/workflows/pr-test.yml; ensure the added test path matches the
test file name tests/test_opsd_jsd.py and commit both the updated .j2 template
and the regenerated .yml.
---
Nitpick comments:
In `@docs/en/advanced/on-policy-distillation.md`:
- Line 113: Clarify why OPSD requires the flag `--context-parallel-size 1` by
adding a short sentence after the limitation line: state that this is due to the
current teacher logit storage implementation which retains full-vocabulary
teacher response logits across teacher and student forwards (so memory scales
with response_length × vocab_size), and mark it as an implementation limitation
rather than a fundamental algorithmic constraint so readers know it may be
relaxed in future work.
In `@examples/on_policy_distillation/run-qwen3-8B-opsd.sh`:
- Around line 150-160: The script is expanding several argument arrays without
quotes which can lead to word splitting or globbing; update each array expansion
in the final command invocation—MODEL_ARGS, CKPT_ARGS, ROLLOUT_ARGS,
OPTIMIZER_ARGS, OPSD_ARGS, WANDB_ARGS, PERF_ARGS, EVAL_ARGS, SGLANG_ARGS,
MISC_ARGS, and RM_ARGS—to use the safe form with quoted expansions (i.e.,
replace ${ARRAY[@]} with "${ARRAY[@]}" for each listed array) so elements
containing spaces or special characters are preserved.
In `@slime/utils/ppo_utils.py`:
- Around line 240-297: The inline "import math" inside
compute_vocab_parallel_jsd should be moved to module-level imports to follow
Python conventions; add "import math" at the top of the file alongside the other
imports and remove the inline import statement within compute_vocab_parallel_jsd
(the function defined as compute_vocab_parallel_jsd in
slime/utils/ppo_utils.py), leaving the math usage (math.log) unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 21978481-9f23-4ad0-ab38-ef0ae7965207
📒 Files selected for processing (16)
.github/workflows/pr-test.yml.github/workflows/pr-test.yml.j2docs/en/advanced/on-policy-distillation.mdexamples/on_policy_distillation/run-qwen3-8B-opsd.shslime/backends/megatron_utils/actor.pyslime/backends/megatron_utils/loss.pyslime/backends/megatron_utils/model.pyslime/ray/placement_group.pyslime/ray/rollout.pyslime/rollout/data_source.pyslime/utils/arguments.pyslime/utils/data.pyslime/utils/ppo_utils.pyslime/utils/types.pytests/test_opsd_jsd.pytests/test_qwen2.5_0.5B_opsd.py
Comment on lines
+21
to
+27
|
|
||
| # Teacher (frozen init policy) and student both start from the converted checkpoint. | ||
| ckpt_args = ( | ||
| f"--hf-checkpoint /root/models/{MODEL_NAME}/ " | ||
| f"--ref-load {torch_dist_ckpt}/ " | ||
| f"--load {torch_dist_ckpt}/ " | ||
| ) |
There was a problem hiding this comment.
Fix black formatting issue (line 21).
The pre-commit hook failed because black reformatted line 21. You must apply the reformatting before the PR can pass CI.
Run black tests/test_qwen2.5_0.5B_opsd.py locally and commit the formatted result.
🧰 Tools
🪛 GitHub Actions: pre-commit / 0_Run pre-commit.txt
[error] 21-21: pre-commit hook 'black' failed because it reformatted files. Reformat needed in this test file (string concatenation in ckpt_args was reformatted).
🪛 GitHub Actions: pre-commit / Run pre-commit
[error] 21-21: black reformatted the file (Format code failed due to pending formatting changes).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_qwen2.5_0.5B_opsd.py` around lines 21 - 27, The file fails black
formatting on the multi-line string constructing ckpt_args (the tuple assigned
to ckpt_args using MODEL_NAME and torch_dist_ckpt); run the formatter and commit
the change: run `black tests/test_qwen2.5_0.5B_opsd.py` (or your project's
pre-commit formatter) to reformat the file so the ckpt_args assignment matches
black style, then add and commit the formatted file so CI pre-commit checks
pass.
Cross-checked against siyan-zhao/OPSD: - Add --opsd-temperature (default 1.0) applied to both logits before the JSD, matching the reference (was missing). compute_vocab_parallel_jsd gains a temperature arg; unit test covers it. - Skip teacher_tokens / teacher_total_lengths / teacher_response_logits in log_rollout_data (they are not token-level metrics; teacher_tokens would have crashed on .mean() over a long tensor). - opsd_loss_function: assert student/teacher sample counts match (strict zip). - JSD direction/mixture/clip-then-mask confirmed to match the reference exactly.
HJSang
force-pushed
the
claude/awesome-dirac-mlam4x
branch
from
June 14, 2026 00:08
5c5ea38 to
cc54f81
Compare
- #1 add --opsd-offload-teacher-logits to offload full-vocab teacher logits to CPU between forwards (moved back to device per micro-batch in the loss); keep the on-GPU view by default. Chunked JSD remains future work (THUDM#4, deferred). - THUDM#2 warn when privileged_info is empty/None (teacher==student context -> ~0 signal). - THUDM#3 repack teacher micro-batches by teacher lengths on the actor side, keeping the student's sample-to-rank assignment so response positions stay aligned (repack_micro_batches_by_length in dp_schedule.py; forward-only, per-rank). - THUDM#5 validate OPSD is not combined with --disable-compute-advantages-and-returns. - THUDM#6 skip loading the ref model under OPSD (never forwarded). - THUDM#7 descriptive error when a dataset row lacks the privileged-info field. Tests: add repack unit tests (coverage, token budget, oversized-sample-alone).
#1 (critical): vocab-parallel log-softmax normalizer was all-reduced with an identity backward, under-counting the student-logit gradient by ~1/TP when tensor-parallel size > 1 (the global normalizer couples all ranks' log-probs, so its cotangent must be all-reduced). Add _VocabParallelAllReduceSumGradAllReduce (all-reduce forward AND backward) for the normalizer; keep identity-backward _VocabParallelAllReduceSum for the final replicated jsd reduction. Add a distributed (TP=2, gloo) test that shards the vocab and checks the JSD value and student gradient against the dense single-process reference (would fail under the old identity backward; a TP=1 test cannot catch it). THUDM#2 restore the actor as the live model via try/finally around the teacher forward so an OOM there can't leave teacher weights live for backup("actor"). THUDM#3 clone the [R, V] teacher response slice instead of keeping a view into the full [1, T_padded, V] microbatch buffer (frees the padded buffer; cheap, strictly better than the view). --opsd-offload-teacher-logits still offloads to CPU. THUDM#4 reject --opd-type=self with kl_coef!=0 or --use-kl-loss (OPSD is pure distillation; the ref model is intentionally not loaded). THUDM#5 log when auto-setting --loss-type=opsd instead of overwriting silently.
- Extend the TP=2 vocab-parallel JSD test to beta in {0, 0.5, 1} (endpoints use
the forward/reverse-KL branches; mixture uses the general path) — all check the
student gradient against the dense reference.
- Assert no virtual pipeline parallelism when OPSD re-packs teacher micro-batches
under --use-dynamic-batch-size (the re-pack does not preserve VPP mb-group alignment).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements On-Policy Self-Distillation (OPSD) — the method from Self-Distilled Reasoner — in slime's Megatron backend, enabled with
--opd-type self.A single model is both student and teacher, differing only by context:
--opd-teacher-load), additionally conditioned on privileged information (the ground-truth solution). It scores the student's response tokens in a single forward pass — no generation.The loss is a direct, full-vocab token-level generalized JSD over the student's on-policy tokens, with per-token clipping and no task reward (pure distillation). This contrasts with the existing OPD modes, which fold a sampled-token reverse-KL into the advantage.
with mixture
M = (1-β)·student + β·teacher,JSD_β = β·KL(teacher‖M) + (1-β)·KL(student‖M)(β=0→forward KL, β=1→reverse KL, β=0.5→symmetric), matching OPSD's GOLD-trainer convention exactly.What's included
ppo_utils.py):compute_vocab_parallel_jsdwith a differentiable vocab-parallel all-reduce; gradients flow only through the student. Verified by CPU unit tests against an independentF.kl_divreference.loss.py):get_response_logitsextractor,opsd_loss_function, and the"opsd"dispatch case; per-token clamp via--opsd-jsd-clip.actor.py):compute_teacher_response_logitsruns the frozen teacher on the privileged sequence;train_actorskips the RL advantage/ref/old-logprob passes for OPSD.rollout.py,data.py,data_source.py,types.py): builds the teacher's[problem + privileged_info + response]sequence; newSample.privileged_infoand--opsd-privileged-info-key.arguments.py):--opd-type self,--opsd-beta,--opsd-jsd-clip,--opsd-privileged-info-key.docs/en/advanced/on-policy-distillation.md,examples/on_policy_distillation/run-qwen3-8B-opsd.sh(uses the reference datasetsiyanzhao/Openthoughts_math_30k_opsd, fieldsproblem/solution).tests/test_opsd_jsd.py(CPU unit tests for the JSD core) andtests/test_qwen2.5_0.5B_opsd.py(end-to-end Megatron smoke test, registered in the megatron CI matrix).Verification
F.kl_divreference; gradients flow only through the student; non-negativity; identical→0).MVP limitations (documented)
--context-parallel-size 1(asserted).response_length × vocab_size; prefer smaller models / shorter responses until a chunked / dual-resident-teacher follow-up.reason_first(teacher self-rationalization before scoring) and Qwen3 thinking-mode toggles from the reference repo are not yet implemented.https://claude.ai/code/session_013FyzAwBJEqE9mcoPJScD5u
Generated by Claude Code
Summary by CodeRabbit