[RFC] Support M2PO for stale-rollout GRPO

[primorLee]

Checklist

• I have searched existing issues and pull requests, and this is a new feature request.

Feature Request Description

I would like to contribute support for M2PO (Second-Moment Trust Policy Optimization) to ms-swift's GRPO training stack.

M2PO was published at ICLR 2026. It targets off-policy RL with stale rollout data, which is directly relevant to the fully async rollout item in the current ms-swift roadmap. Instead of PPO's fixed token-wise clipping interval, the final paper's Algorithm 1 constrains the batch-level second moment of the policy/behavior log-ratio and masks only the largest outliers in the active trust-region quadrants.

Paper and reference implementation:

• Paper: https://arxiv.org/abs/2510.01161
• Reference repository: https://github.com/Infini-AI-Lab/M2PO

Proposed scope

1. Add m2po as a GRPO loss_type, with the paper default m2_threshold=0.04.
2. Compute log(pi_current / pi_behavior) from rollout_per_token_logps when rollout log-probabilities are available.
3. Apply M2 masking only to tokens in the two PPO clipping quadrants: (A > 0, ratio > 1) and (A < 0, ratio < 1).
4. Remove the largest squared log-ratio outliers until the second moment of the remaining trust-region tokens is at most the configured threshold.
5. Preserve the paper's denominator: average the masked objective over all valid completion tokens, not only the unmasked tokens.
6. Keep Transformers/HF and Megatron-SWIFT behavior aligned.
7. Add diagnostics such as m2_before, m2_after, masked_fraction, and trust_region_fraction.
8. Add deterministic unit tests for masking quadrants, threshold boundaries, padding/empty masks, numerical stability, and loss gradients, plus a small stale-rollout comparison against GRPO.

I already have an independent implementation of the final Algorithm 1 under review in AgentScope/Trinity-RFT#616. I would adapt the implementation and its edge-case tests to ms-swift's current GRPOBatch and distributed training abstractions rather than copying framework-specific code.

Before implementation, I would appreciate maintainer guidance on three details:

1. Should the initial PR include both the Transformers/HF and Megatron-SWIFT trainers, or land one backend first?
2. Should loss_type=m2po require rollout_per_token_logps, or fall back to old_per_token_logps when rollout log-probabilities are unavailable?
3. Should the M2 threshold be computed globally across data-parallel ranks, or per local training batch/micro-batch?

Pull Request

Yes. If this proposal fits the roadmap, I plan to submit the implementation, CPU unit tests, configuration/docs, and a reproducible small-model stale-rollout validation. PRs will be split if maintainers prefer separate backend or validation changes.

[primorLee]

Implementation update from draft PR #9965:

• The shared final-Algorithm-1 implementation is now wired into both HF and Megatron and uses a global data-parallel threshold.
• A topology audit found and fixed a Megatron CP correctness issue: log-probs are reconstructed to the full logical sequence before loss, so including CP ranks in M2 selection duplicated every token and could give CP replicas different masks at the threshold boundary. Selection now uses the pure-DP group, with a two-process negative-control/parity test.
• The M2PO + adjacent rollout-IS suites pass (25 tests), and a forced-stale RTX 4090 loss/gradient oracle matches CPU for the mask, loss, gradients, and metrics (M2 0.057690 -> 0.039966 at threshold 0.04).

I am keeping the PR in draft. The remaining semantic decision is item 2 above: for async/vLLM/Megatron stale rollouts, silently falling back to freshly recomputed old_per_token_logps can stop representing the behavior policy. My proposed fail-closed boundary is to require rollout log-probs for those paths and allow fallback only for a provably synchronous native-HF path where behavior == old. A real multi-GPU Megatron CP smoke test also remains pending.

[hjh0119]

Once it is ready for review, just let us know and we will take a look promptly.

[primorLee]

Validation update for #9965: the real two-node Megatron Context Parallel gate is now closed at head 60da0d4048e25fa9d916e84ef666085aac737eac.

• rank 0: NVIDIA A800 80GB; rank 1: NVIDIA GeForce RTX 4090
• aligned PyTorch 2.8.0+cu128, CUDA 12.8, NCCL 2.27.3, Megatron-Core 0.16.1, and Transformer Engine 2.16.0
• the CUDA/NCCL preflight returned all_reduce(rank + 1) == 3.0 on both ranks, both exit code 0
• production reconstruct_tensor_cp restored the complete token sequence identically on both CP ranks
• the DP+CP negative control reproduced replica-divergent masks, while pure-DP selection produced identical masks
• backward gradients and M2/count metrics matched the analytic oracle; both smoke ranks exited with code 0

The PR body now includes the exact validation scope and transport disclosure. I am keeping the PR in Draft only for the remaining behavior-policy interface decision above: whether stale async/vLLM/Megatron paths should fail closed when rollout log-probabilities are absent instead of silently using freshly recomputed old_per_token_logps.

[primorLee]

PR #9965 is now ready for review at 4ad0ed6a464119cc2d8fd3165320e58513ad0cae.

The adversarial pass closed the remaining behavior-policy decision fail-closed for vLLM/Megatron rollouts and also fixed a separate optimizer-batch issue: M2PO's nonlinear mask is now selected once across all Megatron micro-batches instead of independently per micro-batch. Unsupported HF accumulation/sequence-parallel/dynamic-chunk paths and stochastic Megatron prepasses are explicitly rejected rather than silently changing the objective.

Validation at the new head includes 26 targeted/adjacent regressions on Windows and Linux, a clean merge-tree run against current main, full pre-commit, and an RTX 4090 CUDA mask/loss/gradient oracle (M2 0.307053 -> 0.039919, max gradient error 1.16e-10). The earlier real A800+4090 two-node CP/NCCL smoke remains attributed to its exact parent 60da0d4; the A800 instance was offline during the head rerun, which is disclosed in the PR body.

Current GitHub status: lint passed; CPU and NPU unittest jobs are running.