feat(grpo): add M2PO for stale-rollout training

[primorLee]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Implements M2PO (Second-Moment Trust Policy Optimization) for stale-rollout GRPO, following the final ICLR 2026 Algorithm 1. Related RFC: #9954.

What changed

• adds loss_type=m2po and the paper default m2_threshold=0.04
• computes log(pi_current / pi_behavior) from the actual rollout log-probabilities for vLLM/Megatron paths
• permits old_per_token_logps fallback only for synchronous native HF generation; deployment-backed paths fail closed if behavior log-probabilities are missing on any rank
• masks only the active PPO trust-region quadrants and removes the largest squared log-ratio outliers until the remaining batch M2 is within budget
• preserves the paper's original valid-token denominator after masking
• selects once over the complete optimizer batch: HF currently requires gradient_accumulation_steps=1; Megatron concatenates all optimizer-step micro-batches before selection
• computes distributed selection across pure data-parallel groups so reconstructed TP/CP replicas are not counted twice
• keeps Transformers/HF and Megatron-SWIFT on the same shared core implementation
• logs m2_before, m2_after, masked_fraction, and trust_region_fraction
• rejects unsupported combinations that would silently change the objective: sequence-level IS, extra rollout IS correction without separately retained training-engine behavior log-probabilities, off-policy sequence masking, dual clipping, Liger loss, HF sequence parallel/dynamic loss chunking, and Megatron multi-step reuse
• requires deterministic Megatron mask prepasses; LoRA uses --lora_dropout 0, and non-zero model dropout/stochastic depth/BatchNorm/router jitter fail closed
• documents the objective and support boundary in English and Chinese

The reference experiments use beta=0, so M2PO defaults to zero KL coefficient when the user does not explicitly set beta; an explicit non-zero value remains supported with a warning.

The final paper describes TIS as complementary to M2PO. That composition needs both rollout-engine and training-engine behavior-policy log-probabilities. This implementation currently retains only the rollout behavior tensor, so it rejects an additional rollout_importance_sampling_mode instead of double-correcting the same ratio.

Validation

Validated head 4ad0ed6a464119cc2d8fd3165320e58513ad0cae:

• Windows targeted M2PO + adjacent rollout-IS regression: 26 passed
• Linux/Python 3.10 targeted M2PO + adjacent rollout-IS regression: 26 passed
• repository pre-commit hooks on every changed file: flake8, isort, YAPF, whitespace, EOF, quote, merge-conflict, and mixed-line-ending checks passed
• Python compilation checks for the shared core and both trainer integrations passed
• git diff --check passed
• a conflict-free merge tree against current origin/main (174f15e2a812c2669171c1006d8441348ecf35e5) passed the same 26 tests

The tests cover strict/fallback behavior-policy routing, active trust-region quadrants, final Algorithm 1 equivalence, one optimizer-batch selection across micro-batch splits, original-denominator gradients, masked padding non-finites, invalid inputs, variable distributed token counts, and exclusion of reconstructed CP replicas.

RTX 4090 CUDA oracle at head

• PyTorch 2.8.0+cu128, CUDA 12.8
• valid tokens: 3630; trust-region tokens: 1872; kept trust-region tokens: 938
• M2: 0.307053 -> 0.039919 at threshold 0.04
• masked fraction: 0.257300
• CUDA mask matched an independent CPU Algorithm 1 loop and the shared CPU implementation
• analytic-vs-autograd maximum gradient error: 1.16e-10
• missing deployment behavior log-probabilities and stochastic Megatron prepasses were both rejected as intended

Real two-node Megatron Context Parallel smoke

The topology/CP fix in parent 60da0d4048e25fa9d916e84ef666085aac737eac was validated with one NVIDIA A800 80GB (rank 0) and one NVIDIA GeForce RTX 4090 (rank 1):

• aligned runtime: PyTorch 2.8.0+cu128, CUDA 12.8, NCCL 2.27.3, Megatron-Core 0.16.1, Transformer Engine 2.16.0
• real two-node CUDA/NCCL world size 2; all_reduce(rank + 1) returned 3.0 on both GPUs; both ranks exited 0
• initialized context_parallel_size=2; confirmed CP group size 2, DP+CP group size 2, and pure-DP group size 1
• called production swift.megatron.trainers.utils.reconstruct_tensor_cp on load-balanced CP shards; both ranks reconstructed [0.1, 0.2, 0.3, 0.4]
• negative control with the old DP+CP group produced replica-divergent masks: [[T,T,F,F], [T,T,T,F]]
• pure-DP selection produced replica-identical masks: [[T,T,F,F], [T,T,F,F]]
• backward gradients matched the analytic shard-local oracle: [[-0.27629274, 0], [-0.30535069, 0]]
• metrics matched the oracle: M2 0.075 -> 0.025, masked fraction 0.5, valid/trust/kept counts 4/4/2

The two-node run used SSH-forwarded loopback sockets plus a process-scoped bind shim because the GPUs are on different private networks. The shim changed only socket port selection; tensor operations, process groups, collectives, reconstruction, masking, loss, and backward used the production PyTorch/NCCL/Megatron/ms-swift paths.

The A800 instance was offline while validating 4ad0ed6, so the two-node smoke above is intentionally attributed to its exact parent commit rather than presented as a head rerun. Head changes were revalidated through the optimizer-batch distributed tests and the 4090 CUDA oracle above; they do not change CP reconstruction.

Review status

The behavior-policy decision from #9954 is now resolved fail-closed for deployment-backed rollouts, and the optimizer-batch/parallelism boundaries are explicit. This is ready for maintainer review.