[feat] Init true on policy with qwen_dense

[maocheng23]

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
• Miles: [feat] Init true on policy with qwen_dense miles#1052

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:

1. Miles (launcher): picks contract, validates parallel layout against typed model profile, renders SGLang/Megatron/env args from a typed kernel policy, drives PPO logprob comparison.
2. SGLang (rollout numerical truth): produces logprobs using deterministic kernels — TP-invariant matmul + fixed binary tree all-reduce + fp32 log-softmax.
3. Megatron (parity target, this PR): 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. Backward path is fully differentiable — SGLang* layer classes inherit from standard Megatron layers and use the gradient-aware Megatron primitives (linear_with_grad_accumulation_and_async_allreduce).

In this PR (Megatron)

• New megatron/core/true_on_policy/ package:
  • schema.py — vendored shared identity, byte-identical with SGLang and Miles copies
  • contracts.py — MegatronTrueOnPolicyContract with policy_for(config) factory; resolver is registry lookup, not branch
  • provider.py — SGLangSpecProvider extending BackendSpecProvider
  • Layer classes: SGLangNorm, SGLangColumnParallelLinear, SGLangRowParallelLinear, SGLangCoreAttention, SGLangFinalRMSNorm, SGLangQKRMSNorm
  • matmul.py — sglang_reference_matmul with K-block fixed-tree partial-sum ordering matching SGLang's row-linear contract, delegating to gradient-aware Megatron primitives in backward
  • attention_fa3.py, rope.py, cp_layout.py, bias_dropout.py, runtime.py, sglang_backend.py
• Forward-path leak collapse — zero if self.config.use_sglang: branches in forward paths:
  • transformer/attention.py — dtype boundaries via runtime policy
  • transformer/transformer_layer.py — residual contract via runtime policy
  • transformer/transformer_block.py — block-level true-on-policy behavior + Ulysses CP recompute fallback
  • transformer/linear_cross_entropy.py — LM-head input cast via runtime policy
  • tensor_parallel/layers.py — deterministic row-parallel reduction via runtime policy
  • distributed/distributed_data_parallel.py — Ulysses CP gradient scaling via runtime policy
  • models/gpt/gpt_layer_specs.py — spec provider selection from runtime policy
• Phase 4 — retire use_sglang boolean field on TransformerConfig; replaced by true_on_policy_contract: Optional[str]. Backward-compat warning removed.
• All 19 typed runtime-policy fields land — cast_attention_input_to_dense_math_dtype, use_sglang_residual_pair, use_sglang_final_norm, deterministic_row_parallel_reduce, defer_ulysses_cp_loss_scaling_to_grad_sum, etc.

Validation

• ✅ CPU unit tests green (test_sglang_extension.py, test_tree_all_reduce.py, test_true_on_policy_logits.py)
• 🔴 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.
• 🟡 Backward smoke test for every training-path kernel — listed in the contract checklist but not yet executed; gating item before training runs.

Out of scope

• Qwen3-MoE / Qwen3-Next contracts — additive after this stack lands; spec provider has grouped_mlp_modules slot reserved (currently falls back to standard GroupedMLP)
• Alignment harness (layer-dump comparator, first-divergence comparator) — separate follow-up PR
• BackendSpecProvider extension for hybrid attention (Qwen3-Next) — needed when that contract is added
• Model-family cross-check at validator — recommended before contract [1/8] fix: misc compatibility fixes for PyTorch and TE #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
• Backward smoke test for SGLang layer classes
• Training smoke run with PPO loss (validates Miles bug fix in companion PR)

🤖 Generated with Claude Code

[Zhichenzzz]

LGTM, thanks!