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

docs/source/Instruction/Command-line-parameters.md

@@ -633,7 +633,8 @@ reward模型参数将在PPO、GRPO中使用；teacher模型参数在GKD与GRPO
 - reward_model_plugin: 奖励模型逻辑，默认为orm逻辑, 详细见[自定义奖励模型](./GRPO/DeveloperGuide/reward_model.md#自定义奖励模型)。
 - dataset_shuffle: 是否对dataset进行随机操作，默认为True。
 - truncation_strategy: 用于处理输入长度超过 max_length 的样本，支持 delete 和 left 两种策略，分别表示删除该样本和从左侧裁剪。默认值为 left。若使用 delete 策略，被删除的超长样本或编码失败的样本将在原数据集中通过重采样进行替换。
-- loss_type: loss 归一化的类型，可选项为['grpo', 'bnpo', 'dr_grpo', 'dapo', 'cispo', 'sapo', 'real', 'fipo'], 默认为'grpo', 具体参考[文档](./GRPO/DeveloperGuide/loss_types.md)
+- loss_type: 策略损失类型，可选项为['grpo', 'bnpo', 'dr_grpo', 'dapo', 'cispo', 'sapo', 'real', 'fipo', 'm2po']，默认为'grpo'，具体参考[文档](./GRPO/DeveloperGuide/loss_types.md)
+- m2_threshold: M2PO 的批次级二阶矩阈值，默认为0.04。
 - fipo_decay_rate: FIPO Future-KL 折扣半衰参数，实际折扣为`2 ** (-1 / fipo_decay_rate)`，默认值为32.0。
 - fipo_clip_range: FIPO influence weight 裁剪范围，默认值为0.2；设置为None或0时不裁剪。
 - fipo_clip_high_only: 是否只将FIPO influence weight裁剪到`[1.0, 1.0 + fipo_clip_range]`，默认值为True。

docs/source/Instruction/GRPO/DeveloperGuide/loss_types.md

@@ -122,6 +122,29 @@ FIPO 的 influence weight 默认不参与梯度计算，并使用与 DAPO 相同
 
 **归一化维度:** 全局 token 维度（所有进程的 completion token 总数）
 
+## M2PO
+
+`--loss_type m2po --m2_threshold 0.04 --beta 0`
+
+[M2PO](https://arxiv.org/abs/2510.01161) 使用行为策略与当前策略之间对数概率比的批次级二阶矩，替代 PPO
+的固定裁剪区间。算法只约束 PPO 实际触发裁剪的两个区域：`(A > 0, ratio > 1)` 和
+`(A < 0, ratio < 1)`。它按二阶矩从大到小屏蔽异常 token，直到剩余 trust-region token 的平均二阶矩
+不超过 `m2_threshold`。
+
+使用 vLLM 或 Megatron rollout 时，ratio 必须使用实际采样行为策略的 `rollout_per_token_logps`，任一 rank
+缺失都会直接报错。只有同步的原生 HF generation（生成与训练使用同一模型引擎）允许回退到
+`old_per_token_logps`。分布式训练会在数据并行组内统一选择阈值。按照论文定义，被屏蔽 token 的策略损失
+置零，但分母仍使用屏蔽前的全部有效 completion token。论文默认配置为 `m2_threshold=0.04`、`beta=0`。
+
+当前 HF 路径要求 `gradient_accumulation_steps=1`、`sequence_parallel_size=1`，且不支持动态 loss chunk；
+Megatron 路径要求 `steps_per_generation=1`，并在一个 optimizer batch 的所有 micro-batch 上只选择一次 mask。
+Context Parallel 重建产生的副本不会被重复计入。由于 Megatron 在 loss forward 前预计算全批次 mask，策略 forward
+必须是确定性的；LoRA 训练需设置 `--lora_dropout 0`，非零模型 dropout、stochastic depth、BatchNorm 或 router
+jitter 会直接报错。最终版论文中的 TIS 组合还需要单独保留训练引擎侧的行为策略 log-prob；当前路径没有该独立
+张量，因此拒绝再叠加 `rollout_importance_sampling_mode`，以避免重复校正。
+
+**归一化维度：** M2PO 屏蔽前的全局有效 token 维度。
+
 ## SAPO
 
 `--loss_type sapo`

docs/source/Megatron-SWIFT/Command-line-parameters.md

@@ -366,7 +366,8 @@ Megatron训练参数继承自Megatron参数和基本参数（**与ms-swift共用
 - reward_weights: 每个奖励函数的权重。必须与奖励函数和奖励模型的总数量匹配。默认为 None，即所有奖励的权重都相等，为`1.0`。
   - 提示：如果GRPO训练中包含`--reward_model`，则其加在奖励函数的最后位置。
 - truncation_strategy: 对输入长度超过 `max_length`的处理方式，支持`delete`和`left`，代表删除、左侧裁剪，默认为`left`。注意对于多模态模型，左裁剪可能会裁剪掉多模态token导致模型前向报错shape mismatch。使用`delete`方式，对于超长数据和编码失败的样例会在原数据集中重采样其他数据作为补充。
-- loss_type: loss 归一化的类型，可选项为['grpo', 'bnpo', 'dr_grpo'], 默认为'grpo', 具体查看该[pr](https://github.com/huggingface/trl/pull/3256#discussion_r2033213348)。
+- loss_type: 策略损失类型，可选项为['grpo', 'bnpo', 'dr_grpo', 'dapo', 'cispo', 'sapo', 'real', 'fipo', 'm2po']，默认为'grpo'。
+- m2_threshold: M2PO 的批次级二阶矩阈值，默认为0.04。
 - log_completions: 是否记录训练中的模型生成内容，默认为False。
 - vllm_mode: vLLM 集成模式，可选项为 `server` 和 `colocate`。server 模式使用 `swift rollout` 拉起的 vLLM 服务器进行采样，colocate 模式在程序内部署 vLLM。使用server端时，
 - vllm_mode server 参数

docs/source_en/Instruction/Command-line-parameters.md

@@ -647,7 +647,8 @@ The meanings of the following parameters can be referenced [here](https://huggin
 - reward_model_plugin: The logic for the reward model, which defaults to ORM logic. For more information, please refer to [Customized Reward Models](./GRPO/DeveloperGuide/reward_model.md#custom-reward-model).
 - dataset_shuffle: Whether to shuffle the dataset randomly. Default is True.
 - truncation_strategy: The method to handle inputs exceeding `max_length`. Supported values are `delete` and `left`, representing deletion and left-side truncation respectively. The default is `left`. With the delete strategy, over-long or encoding-failed samples are discarded, and new samples are resampled from the original dataset to maintain the intended batch size.
-- loss_type: The type of loss normalization. Options are ['grpo', 'bnpo', 'dr_grpo', 'dapo', 'cispo', 'sapo', 'real', 'fipo'], default is 'grpo'. For details, refer to this [doc](./GRPO/DeveloperGuide/loss_types.md)
+- loss_type: The policy loss type. Options are ['grpo', 'bnpo', 'dr_grpo', 'dapo', 'cispo', 'sapo', 'real', 'fipo', 'm2po'], default is 'grpo'. For details, refer to this [doc](./GRPO/DeveloperGuide/loss_types.md)
+- m2_threshold: The batch-level second-moment threshold for M2PO. Defaults to 0.04.
 - fipo_decay_rate: Half-life parameter for FIPO Future-KL. The actual discount is `2 ** (-1 / fipo_decay_rate)`. Default is 32.0.
 - fipo_clip_range: Clipping range for the FIPO influence weight. Default is 0.2; set to None or 0 to disable clipping.
 - fipo_clip_high_only: Whether to clip the FIPO influence weight to `[1.0, 1.0 + fipo_clip_range]` only. Default is True.

docs/source_en/Instruction/GRPO/DeveloperGuide/loss_types.md

@@ -122,6 +122,33 @@ The FIPO influence weight is detached by default and uses the same global token
 
 **Normalization Dimension:** Global token dimension (total completion tokens across all processes)
 
+## M2PO
+
+`--loss_type m2po --m2_threshold 0.04 --beta 0`
+
+[M2PO](https://arxiv.org/abs/2510.01161) replaces PPO's fixed clipping interval with a batch-level
+second-moment constraint on the behavior-policy log-ratio. It considers only tokens in the active PPO trust-region
+quadrants, `(A > 0, ratio > 1)` and `(A < 0, ratio < 1)`, then masks the largest squared log-ratio outliers until the
+mean second moment of the remaining trust-region tokens is at most `m2_threshold`.
+
+With vLLM or Megatron rollouts, the ratio must use `rollout_per_token_logps` from the actual sampling behavior policy;
+a missing tensor on any rank raises an error. Falling back to `old_per_token_logps` is allowed only for synchronous
+native HF generation, where generation and training use the same model engine. In distributed training, the threshold
+is selected jointly across the data-parallel group. Following the paper, masked policy-loss terms are zeroed while the
+denominator remains the number of all valid completion tokens. The reference experiments use `m2_threshold=0.04` and
+`beta=0`.
+
+The current HF path requires `gradient_accumulation_steps=1` and `sequence_parallel_size=1`, and does not support
+dynamic loss chunking. The Megatron path requires `steps_per_generation=1` and selects the mask once across every
+micro-batch in an optimizer batch; reconstructed Context Parallel replicas are not counted twice. Because Megatron
+precomputes the full-batch mask before the loss forward, policy forwards must be deterministic: LoRA training requires
+`--lora_dropout 0`, and non-zero model dropout, stochastic depth, BatchNorm, or router jitter raises an error. Composing
+M2PO with TIS as described in the final paper additionally requires separately retained training-engine behavior log
+probabilities. The current path does not retain that independent tensor, so it rejects
+`rollout_importance_sampling_mode` to avoid applying the correction twice.
+
+**Normalization Dimension:** Global valid-token dimension before M2PO masking.
+
 ## SAPO
 
 `--loss_type sapo`

docs/source_en/Megatron-SWIFT/Command-line-parameters.md

@@ -389,7 +389,8 @@ In addition to inheriting the training parameters, the following parameters are
 - reward_weights: Weights for each reward function. Must match the total number of reward functions and reward models. Default is None, meaning all rewards have equal weights of `1.0`.
   - Tip: If GRPO training includes `--reward_model`, it is added at the end of the reward functions.
 - truncation_strategy: The method to handle inputs exceeding `max_length`. Supported values are `delete` and `left`, representing deletion and left-side truncation respectively. The default is `left`. Note that for multi-modal models, left-side truncation may remove multi-modal tokens and cause a shape mismatch error during model forward. With the delete strategy, over-long or encoding-failed samples are discarded, and new samples are resampled from the original dataset to maintain the intended batch size.
-- loss_type: Loss normalization type. Options are `['grpo', 'bnpo', 'dr_grpo']`. Default is `'grpo'`. See this [PR](https://github.com/huggingface/trl/pull/3256#discussion_r2033213348) for details.
+- loss_type: Policy loss type. Options are `['grpo', 'bnpo', 'dr_grpo', 'dapo', 'cispo', 'sapo', 'real', 'fipo', 'm2po']`. Default is `'grpo'`.
+- m2_threshold: Batch-level second-moment threshold for M2PO. Defaults to 0.04.
 - log_completions: Whether to log model-generated content during training. Default is False.
 - vllm_mode: vLLM integration mode. Options are `server` and `colocate`. Server mode uses the vLLM server launched by `swift rollout` for sampling, while colocate mode deploys vLLM within the program. When using server mode:
 - vllm_mode server parameters:

swift/arguments/rlhf_args.py

@@ -1,4 +1,5 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+import math
 import os
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Literal, Optional
@@ -353,7 +354,8 @@ def _init_grpo(self):
             raise ValueError("GRPO requires `truncation_strategy 'left' or 'delete'`, "
                              f"Current value: `truncation_strategy='{self.truncation_strategy}'`.")
         if self.beta is None:
-            self.beta = 0.04  # https://arxiv.org/abs/2402.03300
+            # The M2PO reference setup uses no auxiliary KL loss; keep the existing GRPO default otherwise.
+            self.beta = 0.0 if self.loss_type == 'm2po' else 0.04
         if self.async_generate:
             logger.info('Using async mode. This is a approximate version which '
                         'will use the old weights to generate responses to accelerate. '
@@ -542,6 +544,8 @@ def _check_grpo(self):
             raise ValueError('GRPO with vLLM is not compatible with `device_map`. '
                              'Please set NPROC_PER_NODE equal to num_processes.')
         if self.use_liger_kernel:
+            if self.loss_type == 'm2po':
+                raise ValueError('loss_type=m2po is not supported with use_liger_kernel.')
             liger_kernel_version = version.parse(importlib.metadata.version('liger-kernel'))
             if liger_kernel_version < version.parse('0.7.0'):
                 raise ValueError('Please update liger-kernel to 0.7.0 or later: pip install -U liger-kernel')
@@ -568,9 +572,37 @@ def _check_grpo(self):
             raise NotImplementedError('Currently, async_generate is not supported with multi-turn functionality.')
 
         self._check_opd_rl()
+        self._check_m2po()
         self._check_rlsd()
         self._check_sdar()
 
+    def _check_m2po(self):
+        """Validate combinations that would change the final-paper M2PO objective."""
+        if self.loss_type != 'm2po':
+            return
+        if not math.isfinite(self.m2_threshold) or self.m2_threshold < 0:
+            raise ValueError(f'm2_threshold must be finite and non-negative, got {self.m2_threshold}.')
+        if self.gradient_accumulation_steps != 1:
+            raise ValueError('HF loss_type=m2po requires gradient_accumulation_steps=1 because the M2PO mask '
+                             'must be selected once over the complete optimizer batch.')
+        if self.sequence_parallel_size > 1:
+            raise ValueError('HF loss_type=m2po does not yet support sequence_parallel_size > 1 because '
+                             'reconstructed sequence-parallel replicas must be excluded from mask selection.')
+        if self.importance_sampling_level != 'token':
+            raise ValueError('loss_type=m2po requires importance_sampling_level=token.')
+        if self.rollout_importance_sampling_mode is not None:
+            raise ValueError('The current loss_type=m2po path directly uses rollout log-probabilities as the '
+                             'behavior policy and does not retain the separate training-engine behavior '
+                             'log-probabilities required to compose M2PO with rollout importance sampling.')
+        if self.off_policy_sequence_mask_delta is not None:
+            raise ValueError('loss_type=m2po cannot be combined with off_policy_sequence_mask_delta.')
+        if self.delta is not None:
+            raise ValueError('loss_type=m2po replaces PPO clipping and cannot be combined with delta.')
+        if self.use_liger_kernel:
+            raise ValueError('loss_type=m2po is not supported with use_liger_kernel.')
+        if self.beta != 0:
+            logger.warning(f'M2PO uses beta=0 in the reference experiments, but beta={self.beta} was requested.')
+
     def _check_rlsd(self):
         """Validate RLSD (Self-Distilled RLVR) advantage reweighting parameters.
 

swift/megatron/arguments/megatron_args.py

@@ -1,5 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import json
+import math
 import megatron.core
 import os
 import torch
@@ -93,6 +94,9 @@ class RLHFMegatronArgumentsMixin:
     fipo_clip_high_only: bool = True
     fipo_safety_threshold: Optional[float] = 4.0
 
+    # M2PO https://arxiv.org/abs/2510.01161
+    m2_threshold: float = 0.04
+
     epsilon: float = 0.2
     epsilon_high: Optional[float] = None
     delta: Optional[float] = None
@@ -214,10 +218,11 @@ def __post_init__(self):
             return
         default_loss_type = {'kto': 'kto', 'dpo': 'sigmoid', 'grpo': 'grpo'}
         default_beta = {'gkd': 0.5, 'grpo': 0.04}
-        if self.beta is None:
-            self.beta = default_beta.get(self.rlhf_type, 0.1)
         if self.loss_type is None:
             self.loss_type = default_loss_type.get(self.rlhf_type)
+        if self.beta is None:
+            self.beta = 0.0 if self.rlhf_type == 'grpo' and self.loss_type == 'm2po' else default_beta.get(
+                self.rlhf_type, 0.1)
         if self.rlhf_type == 'kto':
             self._init_kto()
         if self.rlhf_type == 'grpo':
@@ -405,10 +410,33 @@ def _check_not_supported():
                     (f'"REAL loss requires that the training micro_batch_size ({self.micro_batch_size}) '
                      f'is a multiple of num_generations ({self.num_generations}). Please adjust your batch parameters.')
 
+            if self.loss_type == 'm2po':
+                if not math.isfinite(self.m2_threshold) or self.m2_threshold < 0:
+                    raise ValueError(f'm2_threshold must be finite and non-negative, got {self.m2_threshold}.')
+                if self.importance_sampling_level != 'token':
+                    raise ValueError('loss_type=m2po requires importance_sampling_level=token.')
+                if self.rollout_importance_sampling_mode is not None:
+                    raise ValueError('The current loss_type=m2po path directly uses rollout log-probabilities as '
+                                     'the behavior policy and does not retain the separate training-engine behavior '
+                                     'log-probabilities required to compose M2PO with rollout importance sampling.')
+                if self.off_policy_sequence_mask_delta is not None:
+                    raise ValueError('loss_type=m2po cannot be combined with off_policy_sequence_mask_delta.')
+                if self.delta is not None:
+                    raise ValueError('loss_type=m2po replaces PPO clipping and cannot be combined with delta.')
+                if self.tuner_type in {'lora', 'lora_llm'} and self.lora_dropout != 0:
+                    raise ValueError('Megatron loss_type=m2po requires lora_dropout=0 because its optimizer-batch '
+                                     'mask prepass must match the subsequent training forward exactly.')
+                if self.beta != 0:
+                    logger.warning(
+                        f'M2PO uses beta=0 in the reference experiments, but beta={self.beta} was requested.')
+
         _check_not_supported()
         if self.dataset_shuffle is not None:
             self.train_dataloader_shuffle = self.dataset_shuffle
         self._init_generation_batch_params()
+        if self.loss_type == 'm2po' and self.steps_per_generation != 1:
+            raise ValueError('Megatron loss_type=m2po requires steps_per_generation=1 so the precomputed '
+                             'optimizer-batch mask uses the current policy.')
         self.remove_unused_columns = False
         logger.info(f'Setting args.remove_unused_columns: {self.remove_unused_columns}')
         if self.truncation_strategy is None: