[feat] Init true on policy with qwen_dense by maocheng23 · Pull Request #1052 · radixark/miles

maocheng23 · 2026-04-28T08:32:11Z

Summary

Initial framework introducing true-on-policy for Qwen3-dense across SGLang, Megatron, and Miles. This is one of three tightly-coupled PRs that must land together — they share a single contract identifier qwen3_dense_true_on_policy_v1 defined by a vendored schema in each repo.

Companion PRs (must land in lockstep):

SGLang: [feat] Init true on policy with qwen_dense sgl-project/sglang#23961
Megatron-LM: [feat] Init true on policy with qwen_dense Megatron-LM#29

Target

Bit-identical (exact-zero) logprob parity between the SGLang rollout engine and the Megatron trainer for every scored response token at TP=1, TP>1, PP>1, and Ulysses CP for Qwen3-4B (dense).

Design

Three-layer contract architecture:

Miles (launcher, this PR): picks contract, validates parallel layout against typed model profile, renders SGLang/Megatron/env args from a typed kernel policy, drives PPO logprob comparison.
SGLang (rollout numerical truth): produces logprobs using deterministic kernels.
Megatron (parity target): reproduces SGLang's numerics in a differentiable training forward via SGLangSpecProvider layer classes.

The contract object owns its own runtime policy — adding a new architecture (e.g. Qwen3-MoE) is one new contract object + one new model profile entry, not edits across three repos.

In this PR (Miles)

New miles/true_on_policy/ package:
- schema.py — vendored shared identity, byte-identical with SGLang and Megatron copies
- contracts.py — TrueOnPolicyContract definition + lookup registry
- model_profiles.py — QWEN3_DENSE_PROFILE (sole registered family)
- config.py — TrueOnPolicyConfig → TrueOnPolicyKernelPolicy → TrueOnPolicyLaunchPlan; renders SGLang / Megatron / env args from typed contract
apply_true_on_policy_script_defaults(args) and build_true_on_policy_launch_plan(args) entry points
scripts/run_qwen3_4b.py launcher with --true-on-policy=true switch and Megatron model_provider routing through config.true_on_policy_contract
miles/utils/ppo_utils.py:
- compute_log_probs(..., true_on_policy_mode=True) — full-vocab gather → real-vocab truncate → fp32 log-softmax → gather
- NEW: _ReplicatedLossAllGatherLastDim autograd Function (the grad_norm bug fix). Megatron's standard all_gather_last_dim_from_tensor_parallel_region has a reduce-scatter backward, which is correct when each rank contributes a distinct output gradient. In the true-on-policy logprob path every TP rank computes the same scalar loss from the gathered full vocabulary, so reduce-scatter sums TP_size identical gradients into the local logits and scales them by TP_size. Replaced with a typed autograd Function whose backward narrows (splits) the gradient — the correct inverse of the forward all-gather when the downstream loss is replicated.
- This was visible as a ~1.8x grad_norm gap vs the off-policy baseline at clipfrac=0.
miles/backends/training_utils/loss.py — true_on_policy_mode handling (dtype management, response-only logits slicing across CP modes)
Phase 4 — stop emitting legacy --use-sglang and --sglang-rl-on-policy-target flags. Only emits --true-on-policy-contract.

Validation

✅ CPU unit tests green (tests/fast/true_on_policy/, tests/fast/utils/test_true_on_policy_logprobs.py)
✅ Test added for _ReplicatedLossAllGatherLastDim backward (verifies narrow vs reduce-scatter)
🔴 GPU exact-zero E2E gate not yet run at TP=1, TP>1, PP>1, CP. This is the next task on the stack before adopting Qwen3-MoE / Qwen3-Next.

Out of scope

Qwen3-MoE / Qwen3-Next model profiles — additive after this stack lands; would be one new TrueOnPolicyModelProfile entry per family
Alignment harness (layer-dump comparator, first-divergence comparator) — separate follow-up PR
Model-family cross-check at config validation — recommended before contract Add end-to-end accuracy tests to CI #2

Test plan

CPU unit tests pass in CI
GPU exact-zero E2E gate at TP=1
GPU exact-zero E2E gate at TP=2
grad_norm parity check (before/after _ReplicatedLossAllGatherLastDim fix)
PPO training smoke at planned configs

🤖 Generated with Claude Code

gemini-code-assist

Code Review

This pull request introduces a 'true-on-policy' launch contract for Qwen3-dense models, centralizing configuration for SGLang and Megatron backends. It includes updates to rollout log-probability computation, enabling recomputation via SGLang prefill, and adds robust handling for tensor-parallel vocab gathering and loss masking. My feedback focuses on improving the safety of list iteration in cp_utils.py by recommending strict=True for zip to ensure data alignment.

gemini-code-assist · 2026-04-28T08:37:19Z

+
+    local_masks = []
+    for i, (total_length, response_length, loss_mask) in enumerate(
+        zip(total_lengths, response_lengths, loss_masks, strict=False)


Using zip with strict=False can hide potential bugs if the input lists (total_lengths, response_lengths, loss_masks) have mismatched lengths. It's safer to use strict=True to ensure that all per-sample lists are correctly aligned. If they have different lengths, it would be better to raise an error and fail fast.

Suggested change

zip(total_lengths, response_lengths, loss_masks, strict=False)

zip(total_lengths, response_lengths, loss_masks, strict=True)

yueming-yuan · 2026-05-01T05:41:42Z

 def get_base_gpu_id(args, rank):
    num_gpus = min(args.num_gpus_per_node, args.rollout_num_gpus_per_engine)
-    if args.colocate:
+    if getattr(args, "colocate", False):


can we keep the original implementation here (just use args.xxx rather than getattr) to avoid silent errors?

yueming-yuan · 2026-05-01T05:45:17Z

    return pg_loss, loss_masks, metrics


+_POLICY_LOSS_DUMP_COUNTER = 0


nit: can we make this part more clean? (e.g., if the debug utils are not expected to be frequently used, we can remove them; if they are expected to be used in the future, can we move to a separate file, and just plug several lines in core logics?)

yueming-yuan · 2026-05-01T05:50:10Z

    train_fp8: bool = False
    enable_megatron_bridge: bool = False
    enable_mis: bool = False
+    tensor_model_parallel_size: int | None = None


nit: to keep the readability and reproducibility of these scripts, I might suggest not adding too many arguments and changing too much here, maybe we can just provide one verified config and let the user to edit the script according to their task (e.g. just one --true-on-policy, and any other necessary args; and we directly code the recommended settings in this file). This will also make it easier for us to maintain

yueming-yuan · 2026-05-01T05:54:37Z

    return eval_datasets


+def _maybe_enable_true_on_policy_sglang_cp_lm_head(args) -> None:


qq: why is this function empty?

Zhichenzzz

Most backends LGTM! Thank you for the great work! Just a quick reminder, there is a small issue raised by ci test, https://github.com/radixark/miles/actions/runs/25076368286/job/73469945178?pr=1052#step:9:347

Zhichenzzz · 2026-05-02T22:34:44Z

+    for i, (total_length, response_length, loss_mask) in enumerate(
+        zip(total_lengths, response_lengths, loss_masks, strict=False)
+    ):
+        max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None
+        prompt_length = total_length - response_length
+        _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(
+            total_length, response_length, qkv_format, max_seq_len
+        )
+        loss_mask_0 = loss_mask[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length]
+        loss_mask_1 = loss_mask[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length]
+        local_masks.append(torch.cat([loss_mask_0, loss_mask_1], dim=0))


Could we extract L147–157 and the duplicate slicing at L98–109 into a shared _slice_loss_mask_for_local_cp helper, so both call sites share a single source of truth?

Implements PR 11 from miles_migration.md on top of origin/main. Wires the Miles-side half of true-on-policy exact-zero alignment between SGLang rollout and Megatron scoring, all gated on --true-on-policy-mode so the default off-policy path is unchanged. Contract A (deterministic runtime mode): - Add --recompute-logprobs-via-prefill (requires --true-on-policy-mode). - When --true-on-policy-mode is set, auto-select sglang rl_on_policy_target=fsdp_tp (tp>1) or fsdp (tp=1) and enable deterministic inference. - --recompute-logprobs-via-prefill also enables SGLang prefill-only-deterministic-inference. Contract D (logprob scoring): - compute_log_probs / _calculate_log_probs_and_entropy_true_on_policy now: gather full padded TP vocab -> truncate to real vocab_size -> FP32 log_softmax -> gather target tokens. - get_log_probs_and_entropy threads vocab_size through so the truncation happens after gather, not before log_softmax. - policy_loss_function computes train_rollout_logprob_abs_diff from the train-side recomputed log-probs (not old_log_probs), which is what the exact-zero acceptance bar measures. Tests: - tests/fast/utils/test_true_on_policy_logprobs.py: - TP=1 truncate-after-gather parity. - Fake TP-sharded vocab gather/truncate/log-softmax parity. - tests/fast/backends/training_utils/test_true_on_policy_loss_metrics.py: - train_rollout_logprob_abs_diff uses recomputed train log-probs; zero when rollout matches, non-zero when perturbed. - tests/fast/utils/test_arguments.py: - --recompute-logprobs-via-prefill parsing. - --true-on-policy-mode propagation to SGLang server args for TP=1 fsdp and TP>1 fsdp_tp, with and without prefill recompute. Compatibility: - All new behavior is dormant unless --true-on-policy-mode is passed. - The fused vocab-parallel CE path is still the default for compute_log_probs when true_on_policy_mode is False. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

pre-commit isort fix for the new test file. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

The previous _gather_true_on_policy_full_logits used Megatron's all_gather_last_dim_from_tensor_parallel_region whose backward is reduce-scatter. When each TP rank computes the same replicated loss from the gathered full vocabulary, reduce-scatter sums TP_size identical gradients into the local logits, scaling the gradient by TP_size (visible as a ~1.8x grad_norm vs the off-policy baseline at clipfrac=0). Replace with a typed _ReplicatedLossAllGatherLastDim autograd Function whose backward narrows (splits) the gradient along the gathered last dim — the correct inverse of the forward all-gather when the downstream loss is replicated on every rank. Also includes pre-commit auto-fixes (isort, black) for the true-on-policy stack files. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Per review feedback, drop the in-progress journal files that should not ship in this PR. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

maocheng23 requested review from Zhichenzzz, fzyzcjy, guapisolo, yueming-yuan and yushengsu-thu as code owners April 28, 2026 08:32

This was referenced Apr 28, 2026

[feat] Init true on policy with qwen_dense maocheng23/sglang#2

Open

[feat] Init true on policy with qwen_dense radixark/Megatron-LM#29

Merged

gemini-code-assist Bot reviewed Apr 28, 2026

View reviewed changes

maocheng23 mentioned this pull request Apr 28, 2026

[feat] Init true on policy with qwen_dense sgl-project/sglang#23961

Merged

4 tasks

yueming-yuan reviewed May 1, 2026

View reviewed changes

This was referenced May 1, 2026

[feat] Init true on policy with qwen_moe maocheng23/sglang#3

Closed

[feat] Add Qwen3 MoE true-on-policy parity radixark/Megatron-LM#30

Open

[feat] Add Qwen3 MoE true-on-policy parity #1059

Open

Zhichenzzz reviewed May 2, 2026

View reviewed changes

yueming-yuan approved these changes May 14, 2026

View reviewed changes

maocheng23 and others added 12 commits May 17, 2026 16:06

Wire Miles true-on-policy Qwen3 dense flow

06ea7c6

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Add Miles true-on-policy launch contract

41e3f76

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Refine true-on-policy launch contract

8350e04

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Add Miles true-on-policy contract selection

6c47ed5

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Fix true-on-policy prefill logprob parity

3ab8605

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Add true-on-policy contract schema adapter

858bc34

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Make true-on-policy launch contract-only

992e381

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

Stop emitting SGLang legacy on-policy target

a354481

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

maocheng23 and others added 2 commits May 17, 2026 16:09

address cmts

e59efc0

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

chore: apply pre-commit auto-fixes

798b791

Co-authored-by: zju-stu-lizheng <lizheng.cs@zju.edu.cn> Co-authored-by: zyxiyy02 <282300612+zyxiyy02@users.noreply.github.com> Co-authored-by: Yi Zhang <1109276519@qq.com>

maocheng23 force-pushed the feat/true_on_policy_qwen_dense branch from 6fac9bb to 2047f7d Compare May 17, 2026 23:14

maocheng23 requested a review from jybsuper as a code owner May 17, 2026 23:14

maocheng23 force-pushed the feat/true_on_policy_qwen_dense branch from 2047f7d to 798b791 Compare May 18, 2026 06:15

maocheng23 merged commit b8649e6 into main May 18, 2026
18 of 19 checks passed

maocheng23 deleted the feat/true_on_policy_qwen_dense branch May 18, 2026 06:21

Navigation Menu

Uh oh!

[feat] Init true on policy with qwen_dense - #1052

[feat] Init true on policy with qwen_dense#1052
maocheng23 merged 14 commits into
mainfrom
feat/true_on_policy_qwen_dense

maocheng23 commented Apr 28, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Apr 28, 2026

Uh oh!

yueming-yuan May 1, 2026

Uh oh!

yueming-yuan May 1, 2026

Uh oh!

yueming-yuan May 1, 2026

Uh oh!

yueming-yuan May 1, 2026

Uh oh!

Zhichenzzz left a comment

Uh oh!

Zhichenzzz May 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	zip(total_lengths, response_lengths, loss_masks, strict=False)
	zip(total_lengths, response_lengths, loss_masks, strict=True)

		return pg_loss, loss_masks, metrics


		_POLICY_LOSS_DUMP_COUNTER = 0

		return eval_datasets


		def _maybe_enable_true_on_policy_sglang_cp_lm_head(args) -> None:

Uh oh!

Conversation

maocheng23 commented Apr 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Target

Design

In this PR (Miles)

Validation

Out of scope

Test plan

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

yueming-yuan May 1, 2026

Choose a reason for hiding this comment

Uh oh!

yueming-yuan May 1, 2026

Choose a reason for hiding this comment

Uh oh!

yueming-yuan May 1, 2026

Choose a reason for hiding this comment

Uh oh!

yueming-yuan May 1, 2026

Choose a reason for hiding this comment

Uh oh!

Zhichenzzz left a comment

Choose a reason for hiding this comment

Uh oh!

Zhichenzzz May 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

maocheng23 commented Apr 28, 2026 •

edited

Loading