diff --git a/docs/source/Instruction/Command-line-parameters.md b/docs/source/Instruction/Command-line-parameters.md index 400f43b72b..a5fca0c828 100644 --- a/docs/source/Instruction/Command-line-parameters.md +++ b/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。 diff --git a/docs/source/Instruction/GRPO/DeveloperGuide/loss_types.md b/docs/source/Instruction/GRPO/DeveloperGuide/loss_types.md index fed681229a..cfa7183980 100644 --- a/docs/source/Instruction/GRPO/DeveloperGuide/loss_types.md +++ b/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` diff --git a/docs/source/Megatron-SWIFT/Command-line-parameters.md b/docs/source/Megatron-SWIFT/Command-line-parameters.md index ab4f47fe86..e0c284a6b6 100644 --- a/docs/source/Megatron-SWIFT/Command-line-parameters.md +++ b/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 参数 diff --git a/docs/source_en/Instruction/Command-line-parameters.md b/docs/source_en/Instruction/Command-line-parameters.md index 5387a6cd2a..4c431e3572 100644 --- a/docs/source_en/Instruction/Command-line-parameters.md +++ b/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. diff --git a/docs/source_en/Instruction/GRPO/DeveloperGuide/loss_types.md b/docs/source_en/Instruction/GRPO/DeveloperGuide/loss_types.md index 7429c0ee5c..c85aeb855f 100644 --- a/docs/source_en/Instruction/GRPO/DeveloperGuide/loss_types.md +++ b/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` diff --git a/docs/source_en/Megatron-SWIFT/Command-line-parameters.md b/docs/source_en/Megatron-SWIFT/Command-line-parameters.md index 3eec9866b7..3192e4ac6a 100644 --- a/docs/source_en/Megatron-SWIFT/Command-line-parameters.md +++ b/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: diff --git a/swift/arguments/rlhf_args.py b/swift/arguments/rlhf_args.py index 8b1097c353..05f3500d70 100644 --- a/swift/arguments/rlhf_args.py +++ b/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. diff --git a/swift/megatron/arguments/megatron_args.py b/swift/megatron/arguments/megatron_args.py index 5afaacd220..2db130c55b 100644 --- a/swift/megatron/arguments/megatron_args.py +++ b/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: diff --git a/swift/megatron/trainers/grpo_trainer.py b/swift/megatron/trainers/grpo_trainer.py index 6532546261..1d30d9880b 100644 --- a/swift/megatron/trainers/grpo_trainer.py +++ b/swift/megatron/trainers/grpo_trainer.py @@ -16,6 +16,7 @@ expand_advantage_to_per_token, get_local_rollout_values) from swift.rl_core.data import GRPOBatch, GRPOSample from swift.rl_core.grpo_algorithm import score_completions +from swift.rl_core.m2po import compute_m2po_log_ratio, compute_m2po_masks_for_batches, compute_m2po_token_loss_from_mask from swift.rl_core.resample import resample_encode_failed_inputs from swift.rlhf_trainers.gkd_helpers import (assemble_teacher_completion_logprobs, build_opsd_samples, build_teacher_requests, encode_teacher_view, @@ -63,6 +64,55 @@ def __init__(self, args: MegatronArguments, template: Template, **kwargs): def prepare_model(self): super().prepare_model() self._load_teacher_model() + if self.args.loss_type == 'm2po': + self._validate_m2po_prepass_determinism() + + def _validate_m2po_prepass_determinism(self) -> None: + """Reject train/eval stochasticity that would invalidate the mask prepass.""" + stochastic = [] + config_fields = [ + 'hidden_dropout', + 'attention_dropout', + 'embedding_dropout', + 'drop_path_rate', + 'stochastic_depth', + 'moe_router_jitter_eps', + 'moe_input_jitter_eps', + ] + for field_name in config_fields: + value = getattr(self.config, field_name, None) + try: + enabled = value is not None and float(value) > 0 + except (TypeError, ValueError): + enabled = False + if enabled: + stochastic.append(f'config.{field_name}={value}') + + for model_idx, model in enumerate(self.unwrapped_models): + for module_name, module in model.named_modules(): + if isinstance(module, torch.nn.modules.batchnorm._BatchNorm): + stochastic.append(f'model[{model_idx}].{module_name}={type(module).__name__}') + continue + module_type = type(module).__name__.lower() + if isinstance(module, torch.nn.modules.dropout._DropoutNd): + probability = module.p + elif 'droppath' in module_type or 'stochasticdepth' in module_type: + probability = getattr(module, 'drop_prob', getattr(module, 'p', None)) + else: + continue + try: + enabled = probability is not None and float(probability) > 0 + except (TypeError, ValueError): + enabled = False + if enabled: + stochastic.append(f'model[{model_idx}].{module_name}={type(module).__name__}(p={probability})') + + if stochastic: + details = ', '.join(stochastic[:8]) + if len(stochastic) > 8: + details += f', ... ({len(stochastic)} total)' + raise ValueError('Megatron loss_type=m2po requires deterministic train/eval policy forwards because ' + f'the optimizer-batch mask is precomputed from old log-probabilities. Found: {details}') def train(self, train_dataset, val_dataset): if self.dynamic_sample or self.truncation_strategy == 'delete': @@ -100,6 +150,9 @@ def _init_grpo_params(self): self.fipo_clip_high_only = args.fipo_clip_high_only self.fipo_safety_threshold = args.fipo_safety_threshold + # M2PO, https://arxiv.org/abs/2510.01161 + self.m2_threshold = args.m2_threshold + # DAPO, https://arxiv.org/abs/2503.14476 self.dynamic_sample = args.dynamic_sample self.max_resample_times = args.max_resample_times @@ -237,8 +290,75 @@ def _build_rollout_buffer(self, data_iterator): micro_batch_data[i:i + num_mini_batch] for i in range(0, len(micro_batch_data), num_mini_batch) ] assert len(mini_batch_data) == num_gen_steps + if self.loss_type == 'm2po': + for optimizer_batch_data in mini_batch_data: + self._prepare_m2po_optimizer_batch(optimizer_batch_data) return mini_batch_data + def _prepare_m2po_optimizer_batch(self, optimizer_batch_data: List[Dict[str, Any]]) -> None: + """Select one M2PO mask across every micro-batch in an optimizer step. + + The selection is non-linear: applying Algorithm 1 independently to + micro-batches is not equivalent to applying it to their concatenation. + ``old_per_token_logps`` is the current training-engine policy here + because Megatron M2PO permits one immediate step per generation only. + """ + if not mpu.is_pipeline_last_stage(): + return + + grpo_batches = [data['grpo_batch'] for data in optimizer_batch_data] + dp_with_cp_group = mpu.get_data_parallel_group(with_context_parallel=True) + ready_flags = torch.tensor([ + bool(grpo_batches), + bool(grpo_batches) and all(batch.rollout_per_token_logps is not None for batch in grpo_batches), + bool(grpo_batches) and all(batch.old_per_token_logps is not None for batch in grpo_batches), + bool(grpo_batches) and all(batch.advantages is not None for batch in grpo_batches), + ], + dtype=torch.int32, + device=self.device) + torch.distributed.all_reduce(ready_flags, op=torch.distributed.ReduceOp.MIN, group=dp_with_cp_group) + if not ready_flags[0].item(): + raise ValueError('M2PO requires at least one Megatron micro-batch per optimizer step.') + if not ready_flags[1].item(): + raise ValueError('Megatron M2PO requires rollout_per_token_logps from the behavior policy on every rank.') + if not ready_flags[2].item(): + raise ValueError('Megatron M2PO requires old_per_token_logps to select the optimizer-batch mask.') + if not ready_flags[3].item(): + raise ValueError('Megatron M2PO requires per-token advantages before mask selection.') + + log_ratios = [] + completion_masks = [] + advantages = [] + for grpo_batch in grpo_batches: + log_ratios.append( + compute_m2po_log_ratio( + grpo_batch.old_per_token_logps, + grpo_batch.old_per_token_logps, + grpo_batch.rollout_per_token_logps, + allow_old_policy_fallback=False, + )) + completion_mask = grpo_batch.completion_mask + if self.args.overlong_filter: + truncated_mask = grpo_batch.truncated_mask.unsqueeze(-1).expand_as(completion_mask) + completion_mask = completion_mask & (~truncated_mask) + completion_masks.append(completion_mask) + advantages.append(grpo_batch.advantages) + + # Context-parallel ranks hold replicas of the reconstructed sequence. + # Select across pure DP only; each CP replica independently obtains the + # same logical optimizer-batch mask. + pure_dp_group = mpu.get_data_parallel_group(with_context_parallel=False) + batch_masks, metrics = compute_m2po_masks_for_batches( + log_ratios, + completion_masks, + advantages, + m2_threshold=self.m2_threshold, + process_group=pure_dp_group, + ) + for grpo_batch, batch_mask in zip(grpo_batches, batch_masks): + grpo_batch.m2po_mask = batch_mask + grpo_batch.m2po_metrics = metrics + def _generate_and_score_completions(self, inputs: DataType): # Get or create the rollout group (TP×PP×CP) @@ -333,7 +453,7 @@ def _generate_and_score_completions(self, inputs: DataType): if any(m['grpo_batch'].teacher_per_token_logps is not None for m in mini_batch_data): self._log_teacher_kl_metric(mini_batch_data) - if self.loss_type in ['cispo', 'dapo', 'fipo']: + if self.loss_type in ['cispo', 'dapo', 'fipo', 'm2po']: # Calculate num_items_in_batch # Count completion tokens from all mini_batch_data (this includes gathered data from rollout_group) # Use completion_mask.sum() for both padding_free and non-padding_free modes @@ -971,13 +1091,15 @@ def loss_func(self, output_tensor: torch.Tensor, data: Dict[str, Any]): should_compute_rollout_metrics = (not self.disable_rollout_importance_sampling and (self.rollout_importance_sampling_mode is not None or self.log_rollout_offpolicy_metrics)) + all_have_rollout = rollout_per_token_logps is not None if should_compute_rollout_metrics: dp_group = mpu.get_data_parallel_group(with_context_parallel=True) has_flag = torch.tensor([1 if rollout_per_token_logps is not None else 0], dtype=torch.int32, device=per_token_logps.device) torch.distributed.all_reduce(has_flag, op=torch.distributed.ReduceOp.MIN, group=dp_group) - should_compute_rollout_metrics = has_flag.item() > 0 + all_have_rollout = has_flag.item() > 0 + should_compute_rollout_metrics = should_compute_rollout_metrics and all_have_rollout if should_compute_rollout_metrics: # Compute off-policy diagnostic metrics rollout_correction_metrics = self._compute_rollout_offpolicy_metrics(old_per_token_logps, @@ -1015,10 +1137,20 @@ def loss_func(self, output_tensor: torch.Tensor, data: Dict[str, Any]): per_token_kl = None # Compute log ratio for importance sampling - log_ratio = per_token_logps - old_per_token_logps + if self.loss_type == 'm2po': + log_ratio = compute_m2po_log_ratio( + per_token_logps, + old_per_token_logps, + rollout_per_token_logps, + allow_old_policy_fallback=False, + ) + else: + log_ratio = per_token_logps - old_per_token_logps # Compute importance weights based on level - if self.importance_sampling_level == 'token': + if self.loss_type == 'm2po': + log_importance_weights = log_ratio + elif self.importance_sampling_level == 'token': log_importance_weights = log_ratio elif self.importance_sampling_level in ['sequence', 'sequence_token']: # Sequence-level: compute mean log ratio per sequence @@ -1034,7 +1166,8 @@ def loss_func(self, output_tensor: torch.Tensor, data: Dict[str, Any]): f"Unknown importance sampling level: {self.importance_sampling_level}. Possible values are 'token' " ",'sequence' and 'sequence_token'.") - coef_1 = torch.exp(log_importance_weights) + coef_1 = torch.ones_like(log_importance_weights) \ + if self.loss_type == 'm2po' else torch.exp(log_importance_weights) # advantages is per-token [B, T] (expanded at batch construction so the OPD-RL signed teacher # log-ratio is added per token). Edge loss types that need a per-sequence advantage (real / fipo / @@ -1045,6 +1178,7 @@ def loss_func(self, output_tensor: torch.Tensor, data: Dict[str, Any]): 'off_policy_sequence_mask.') fipo_metrics = None + m2po_metrics = None if self.loss_type == 'cispo': clamped_ratios = torch.clamp(coef_1, max=self.epsilon_high).detach() per_token_loss = -clamped_ratios * advantages * per_token_logps @@ -1054,6 +1188,11 @@ def loss_func(self, output_tensor: torch.Tensor, data: Dict[str, Any]): is_positive = advantages > 0 soft_gate = torch.where(is_positive, gate_pos, gate_neg) per_token_loss = -soft_gate * advantages + elif self.loss_type == 'm2po': + if grpo_batch.m2po_mask is None: + raise RuntimeError('Megatron M2PO mask was not prepared for the complete optimizer batch.') + per_token_loss = compute_m2po_token_loss_from_mask(log_ratio, advantages, grpo_batch.m2po_mask) + m2po_metrics = grpo_batch.m2po_metrics elif self.loss_type in ['grpo', 'bnpo', 'dr_grpo', 'dapo', 'fipo']: if self.loss_type == 'fipo': fipo_weight, fipo_metrics = self._compute_fipo_influence(log_ratio, coef_1, advantages, completion_mask) @@ -1135,8 +1274,8 @@ def loss_func(self, output_tensor: torch.Tensor, data: Dict[str, Any]): loss = (per_token_loss * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) elif self.loss_type == 'dr_grpo': loss = (per_token_loss * completion_mask).sum() / (micro_batch_size * self.max_completion_length) - elif self.loss_type in ['cispo', 'dapo', 'fipo']: - # CISPO, DAPO, and FIPO: Normalize by total completion tokens across all processes + elif self.loss_type in ['cispo', 'dapo', 'fipo', 'm2po']: + # Token-level objectives: normalize by all valid completion tokens across processes. num_items_in_batch = grpo_batch.num_items_in_batch dp_size = mpu.get_data_parallel_world_size() normalizer = num_items_in_batch / dp_size @@ -1216,6 +1355,10 @@ def loss_func(self, output_tensor: torch.Tensor, data: Dict[str, Any]): avg_metric['fipo/safety_keep_ratio'] = ((fipo_metrics['safety_mask'].float() * completion_mask).sum() / completion_token_count).clone().detach() + if m2po_metrics is not None: + for key in ['m2_before', 'm2_after', 'masked_fraction', 'trust_region_fraction']: + avg_metric[f'm2po/{key}'] = m2po_metrics[key].clone().detach() + if self.loss_type == 'cispo': # CISPO: Only track upper bound clipping # coef_1 is [batch_size, max_seq_len] or [batch_size, 1] depending on importance_sampling_level @@ -1223,7 +1366,7 @@ def loss_func(self, output_tensor: torch.Tensor, data: Dict[str, Any]): cispo_clip_ratio = (is_cispo_clipped.float() * completion_mask).sum() / completion_token_count # Store local clip ratio, _all_reduce_metric will handle averaging across ranks self._metrics[mode]['cispo_clip_ratio'].append(cispo_clip_ratio) - elif self.loss_type in ['sapo', 'real']: + elif self.loss_type in ['sapo', 'real', 'm2po']: # SAPO / REAL: No hard clipping, skip clipping metrics pass elif self.loss_type in ['grpo', 'bnpo', 'dr_grpo', 'dapo', 'fipo']: @@ -1635,6 +1778,7 @@ def _collect_config_info(self) -> Dict[str, str]: 'offpolicy_sequence_mask': 'enable' if self.args.off_policy_sequence_mask_delta is not None else 'disable', 'rollout_importance_sampling': 'enable' if self.args.rollout_importance_sampling_mode is not None else 'disable', - 'loss_type': str(self.args.loss_type) + 'loss_type': str(self.args.loss_type), + 'm2_threshold': str(self.args.m2_threshold), } return config diff --git a/swift/rl_core/data.py b/swift/rl_core/data.py index f859cba1d4..32653ec922 100644 --- a/swift/rl_core/data.py +++ b/swift/rl_core/data.py @@ -312,6 +312,8 @@ class GRPOBatch: 3. ``advantages`` — computed from gathered rewards. 4. ``rollout_per_token_logps``, ``num_items_in_batch`` — optional, filled when rollout IS / DAPO is enabled. + 5. ``m2po_mask``, ``m2po_metrics`` — optional optimizer-batch M2PO + selection, prepared before Megatron splits the batch into forward passes. """ completion_mask: torch.Tensor # [B, T] truncated_mask: torch.Tensor # [B] @@ -324,6 +326,8 @@ class GRPOBatch: advantages: Optional[torch.Tensor] = None # [B, T] per-token (base broadcast minus per-token teacher KL) num_items_in_batch: Optional[torch.Tensor] = None # scalar logits_to_keep: Optional[int] = None + m2po_mask: Optional[torch.Tensor] = None # [B, T] + m2po_metrics: Optional[Dict[str, torch.Tensor]] = None def to_device(self, device) -> 'GRPOBatch': """Move all tensor fields to ``device`` in place (Ray: collated on the CPU @@ -332,6 +336,11 @@ def to_device(self, device) -> 'GRPOBatch': v = getattr(self, f.name) if isinstance(v, torch.Tensor): setattr(self, f.name, v.to(device)) + elif isinstance(v, dict): + setattr(self, f.name, { + k: item.to(device) if isinstance(item, torch.Tensor) else item + for k, item in v.items() + }) return self diff --git a/swift/rl_core/m2po.py b/swift/rl_core/m2po.py new file mode 100644 index 0000000000..eb2c965481 --- /dev/null +++ b/swift/rl_core/m2po.py @@ -0,0 +1,224 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import math +import torch +import torch.distributed as dist +from typing import Dict, List, Tuple + + +def compute_m2po_log_ratio( + per_token_logps: torch.Tensor, + old_per_token_logps: torch.Tensor, + rollout_per_token_logps: torch.Tensor = None, + allow_old_policy_fallback: bool = True, +) -> torch.Tensor: + """Compute ``log(pi_current / pi_behavior)``. + + ``old_per_token_logps`` is a valid behavior-policy fallback only for the + synchronous native-generation path, where generation and training use the + same model engine. Deployment-backed rollout paths must provide the actual + rollout log probabilities instead of silently substituting a freshly + recomputed training-engine policy. + """ + if rollout_per_token_logps is not None: + behavior_logps = rollout_per_token_logps + elif allow_old_policy_fallback: + behavior_logps = old_per_token_logps + else: + raise ValueError('M2PO requires rollout_per_token_logps from the behavior policy for this rollout path; ' + 'old-policy fallback is only valid for synchronous native generation.') + if behavior_logps is None: + raise ValueError('M2PO requires rollout_per_token_logps or old_per_token_logps.') + if per_token_logps.shape != behavior_logps.shape: + raise ValueError('Current and behavior-policy log probabilities must have identical shapes, got ' + f'{per_token_logps.shape} and {behavior_logps.shape}.') + return per_token_logps - behavior_logps + + +def _distributed_context(process_group=None) -> Tuple[int, int]: + if not dist.is_available() or not dist.is_initialized(): + return 1, 0 + return dist.get_world_size(process_group), dist.get_rank(process_group) + + +def _all_gather_variable(values: torch.Tensor, process_group=None) -> Tuple[torch.Tensor, int]: + """Gather one-dimensional tensors with different lengths in process-group rank order.""" + world_size, rank = _distributed_context(process_group) + if world_size == 1: + return values, 0 + + local_count = torch.tensor([values.numel()], dtype=torch.long, device=values.device) + gathered_counts = [torch.zeros_like(local_count) for _ in range(world_size)] + dist.all_gather(gathered_counts, local_count, group=process_group) + counts = torch.cat(gathered_counts) + max_count = int(counts.max().item()) + local_offset = int(counts[:rank].sum().item()) + if max_count == 0: + return values.new_empty(0), local_offset + + padded = values.new_zeros(max_count) + padded[:values.numel()] = values + gathered_values = [torch.empty_like(padded) for _ in range(world_size)] + dist.all_gather(gathered_values, padded, group=process_group) + return torch.cat([rank_values[:count] + for rank_values, count in zip(gathered_values, counts.tolist())]), local_offset + + +def compute_m2po_mask( + log_ratio: torch.Tensor, + completion_mask: torch.Tensor, + advantages: torch.Tensor, + m2_threshold: float = 0.04, + process_group=None, +) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + """Compute the final-paper M2PO token mask and batch diagnostics. + + M2PO constrains only the two quadrants where PPO clipping would be active: + ``advantage > 0, ratio > 1`` and ``advantage < 0, ratio < 1``. The largest + squared behavior-policy log-ratio outliers are removed until the second + moment of the remaining trust-region tokens is at most ``m2_threshold``. + + When distributed training is initialized, the selection is computed across + the supplied process group. Every rank therefore uses the same batch-level + threshold while receiving a mask aligned with its local tokens. + """ + if not math.isfinite(m2_threshold) or m2_threshold < 0: + raise ValueError(f'm2_threshold must be finite and non-negative, got {m2_threshold}.') + if log_ratio.shape != completion_mask.shape or log_ratio.shape != advantages.shape: + raise ValueError('log_ratio, completion_mask, and advantages must have identical shapes, got ' + f'{log_ratio.shape}, {completion_mask.shape}, and {advantages.shape}.') + + valid_mask = completion_mask.bool() + world_size, _ = _distributed_context(process_group) + + with torch.no_grad(): + detached_log_ratio = log_ratio.detach() + detached_advantages = advantages.detach() + invalid_flags = torch.stack([ + (~torch.isfinite(detached_log_ratio) & valid_mask).any(), + (~torch.isfinite(detached_advantages) & valid_mask).any(), + ]).to(dtype=torch.int32) + if world_size > 1: + dist.all_reduce(invalid_flags, op=dist.ReduceOp.MAX, group=process_group) + if invalid_flags[0].item(): + raise ValueError('M2PO received non-finite behavior-policy log-ratios on valid completion tokens.') + if invalid_flags[1].item(): + raise ValueError('M2PO received non-finite advantages on valid completion tokens.') + + trust_region_mask = valid_mask & (((detached_advantages > 0) & (detached_log_ratio > 0)) + | ((detached_advantages < 0) & (detached_log_ratio < 0))) + + flat_trust_indices = torch.nonzero(trust_region_mask.reshape(-1), as_tuple=False).squeeze(-1) + local_values = detached_log_ratio.float().square().reshape(-1)[flat_trust_indices] + global_values, local_offset = _all_gather_variable(local_values, process_group) + + global_valid_count = valid_mask.sum().to(dtype=torch.long) + if world_size > 1: + dist.all_reduce(global_valid_count, op=dist.ReduceOp.SUM, group=process_group) + + trust_count = global_values.numel() + keep_count = trust_count + keep_global = torch.ones(trust_count, dtype=torch.bool, device=global_values.device) + if trust_count: + sorted_values, order = torch.sort(global_values) + prefix_counts = torch.arange(1, trust_count + 1, dtype=sorted_values.dtype, device=sorted_values.device) + prefix_means = torch.cumsum(sorted_values, dim=0) / prefix_counts + keep_count = int((prefix_means <= m2_threshold).sum().item()) + keep_global.zero_() + keep_global[order[:keep_count]] = True + + local_keep = keep_global[local_offset:local_offset + local_values.numel()] + final_mask = valid_mask.clone() + final_mask.reshape(-1)[flat_trust_indices] = local_keep + + zero = global_values.new_zeros(()) + m2_before = global_values.mean() if trust_count else zero + m2_after = global_values[keep_global].mean() if keep_count else zero + valid_count = int(global_valid_count.item()) + masked_count = trust_count - keep_count + + metrics = { + 'm2_before': m2_before, + 'm2_after': m2_after, + 'masked_fraction': zero.new_tensor(masked_count / valid_count if valid_count else 0.0), + 'trust_region_fraction': zero.new_tensor(trust_count / valid_count if valid_count else 0.0), + 'valid_count': zero.new_tensor(valid_count), + 'trust_region_count': zero.new_tensor(trust_count), + 'kept_trust_region_count': zero.new_tensor(keep_count), + } + + return final_mask, metrics + + +def compute_m2po_masks_for_batches( + log_ratios: List[torch.Tensor], + completion_masks: List[torch.Tensor], + advantages: List[torch.Tensor], + m2_threshold: float = 0.04, + process_group=None, +) -> Tuple[List[torch.Tensor], Dict[str, torch.Tensor]]: + """Select M2PO tokens once across all micro-batches in one optimizer batch.""" + num_batches = len(log_ratios) + if num_batches == 0: + raise ValueError('M2PO requires at least one micro-batch.') + if len(completion_masks) != num_batches or len(advantages) != num_batches: + raise ValueError('log_ratios, completion_masks, and advantages must contain the same number of batches.') + + for batch_idx, (log_ratio, completion_mask, advantage) in enumerate(zip(log_ratios, completion_masks, advantages)): + if log_ratio.shape != completion_mask.shape or log_ratio.shape != advantage.shape: + raise ValueError(f'M2PO micro-batch {batch_idx} has mismatched shapes: {log_ratio.shape}, ' + f'{completion_mask.shape}, and {advantage.shape}.') + + numels = [value.numel() for value in log_ratios] + flat_mask, metrics = compute_m2po_mask( + log_ratio=torch.cat([value.reshape(-1) for value in log_ratios]), + completion_mask=torch.cat([value.reshape(-1) for value in completion_masks]), + advantages=torch.cat([value.reshape(-1) for value in advantages]), + m2_threshold=m2_threshold, + process_group=process_group, + ) + split_masks = [ + value.reshape_as(completion_mask) for value, completion_mask in zip(flat_mask.split(numels), completion_masks) + ] + return split_masks, metrics + + +def compute_m2po_token_loss_from_mask( + log_ratio: torch.Tensor, + advantages: torch.Tensor, + m2po_mask: torch.Tensor, +) -> torch.Tensor: + """Return the unreduced M2PO loss for an already selected optimizer-batch mask.""" + if log_ratio.shape != advantages.shape or log_ratio.shape != m2po_mask.shape: + raise ValueError('log_ratio, advantages, and m2po_mask must have identical shapes, got ' + f'{log_ratio.shape}, {advantages.shape}, and {m2po_mask.shape}.') + + m2po_mask = m2po_mask.bool() + # Inactive tokens must not evaluate an exponential (or multiply a NaN + # advantage), otherwise a masked non-finite padding value can still poison + # autograd through a zero-times-NaN derivative. + active_log_ratio = torch.where(m2po_mask, log_ratio, torch.zeros_like(log_ratio)) + active_advantages = torch.where(m2po_mask, advantages, torch.zeros_like(advantages)) + # Keep the exponent finite even for a non-trust-region outlier. The mask is + # still computed from the unclamped log-ratio, so this does not weaken M2. + clamped_log_ratio = torch.clamp(active_log_ratio, min=-20, max=20) + stable_log_ratio = active_log_ratio + (clamped_log_ratio - active_log_ratio).detach() + return -torch.exp(stable_log_ratio) * active_advantages + + +def compute_m2po_token_loss( + log_ratio: torch.Tensor, + completion_mask: torch.Tensor, + advantages: torch.Tensor, + m2_threshold: float = 0.04, + process_group=None, +) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]: + """Return the unreduced final-paper M2PO loss and its non-differentiable mask.""" + m2po_mask, metrics = compute_m2po_mask( + log_ratio=log_ratio, + completion_mask=completion_mask, + advantages=advantages, + m2_threshold=m2_threshold, + process_group=process_group, + ) + per_token_loss = compute_m2po_token_loss_from_mask(log_ratio, advantages, m2po_mask) + return per_token_loss, m2po_mask, metrics diff --git a/swift/rlhf_trainers/args_mixin.py b/swift/rlhf_trainers/args_mixin.py index eb6d5f1fe9..61ab057bb7 100644 --- a/swift/rlhf_trainers/args_mixin.py +++ b/swift/rlhf_trainers/args_mixin.py @@ -344,6 +344,7 @@ class GRPOArgumentsMixin(RolloutTrainerArgumentsMixin): fipo_safety_threshold (Optional[float]): Safety threshold for negative advantages. Tokens with `advantage < 0` and importance ratio above this value have their FIPO influence weight capped to `[0.8, 1.0]` to avoid over-penalization. Defaults to 4.0. + m2_threshold (float): Batch-level second-moment threshold used by M2PO. Defaults to 0.04. advantage_estimator (Literal['grpo', 'rloo', 'reinforce_plus_plus']): The advantage estimation function to use. 'grpo' calculates the relative advantage within a group. Options are 'grpo', 'rloo', 'reinforce_plus_plus'. Defaults to 'grpo'. @@ -468,6 +469,9 @@ class GRPOArgumentsMixin(RolloutTrainerArgumentsMixin): 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 + num_generations_eval: Optional[int] = None # dataset diff --git a/swift/rlhf_trainers/grpo_trainer.py b/swift/rlhf_trainers/grpo_trainer.py index e261ef0458..48f44d52ff 100644 --- a/swift/rlhf_trainers/grpo_trainer.py +++ b/swift/rlhf_trainers/grpo_trainer.py @@ -54,6 +54,7 @@ expand_advantage_to_per_token) from swift.rl_core.data import GRPOBatch, GRPOSample from swift.rl_core.grpo_algorithm import score_completions +from swift.rl_core.m2po import compute_m2po_log_ratio, compute_m2po_token_loss from swift.rlhf_trainers.gkd_helpers import (assemble_teacher_completion_logprobs, build_opsd_samples, build_teacher_requests, encode_teacher_view, fetch_teacher_parsed_by_routing, remap_teacher_logps_to_student_frame, @@ -887,6 +888,9 @@ def _compute_loss(self, model, model_inputs, grpo_batch, origin_data=None): if not should_chunk: return self._compute_loss_single(model, model_inputs, grpo_batch) else: + if self.loss_type == 'm2po': + raise ValueError('HF loss_type=m2po does not support dynamic loss chunking because selecting ' + 'independent masks for each chunk changes the optimizer-batch objective.') # maybe dynamic rollout num for multi-turn training return self._compute_loss_chunked(model, model_inputs, grpo_batch, origin_data) @@ -999,7 +1003,9 @@ def _compute_loss_and_metrics(self, model, model_inputs: Dict[str, Any], grpo_ba self.rollout_importance_sampling_mode is not None or self.log_rollout_offpolicy_metrics) local_has_rollout = grpo_batch.rollout_per_token_logps is not None - should_compute_rollout_metrics = should_compute_rollout_metrics and all(gather_object([local_has_rollout])) + all_have_rollout = all(gather_object([local_has_rollout])) \ + if should_compute_rollout_metrics or self.loss_type == 'm2po' else local_has_rollout + should_compute_rollout_metrics = should_compute_rollout_metrics and all_have_rollout rollout_is_weights = None if (not self.disable_rollout_importance_sampling and should_compute_rollout_metrics): rollout_per_token_logps = grpo_batch.rollout_per_token_logps @@ -1019,8 +1025,19 @@ def _compute_loss_and_metrics(self, model, model_inputs: Dict[str, Any], grpo_ba pass # rollout_is_weights is a local variable, initialized to None above - log_ratio = per_token_logps - old_per_token_logps - if self.importance_sampling_level == 'token': + if self.loss_type == 'm2po': + rollout_per_token_logps = grpo_batch.rollout_per_token_logps if all_have_rollout else None + log_ratio = compute_m2po_log_ratio( + per_token_logps, + old_per_token_logps, + rollout_per_token_logps, + allow_old_policy_fallback=not self.use_vllm, + ) + else: + log_ratio = per_token_logps - old_per_token_logps + if self.loss_type == 'm2po': + log_importance_weights = log_ratio + elif self.importance_sampling_level == 'token': log_importance_weights = log_ratio elif self.importance_sampling_level in ['sequence', 'sequence_token']: seq_level_log_weights = ((log_ratio * completion_mask).sum(-1) @@ -1036,7 +1053,8 @@ def _compute_loss_and_metrics(self, model, model_inputs: Dict[str, Any], grpo_ba f"Unknown importance sampling level: {self.importance_sampling_level}. Possible values are 'token' " "and 'sequence'.") - coef_1 = torch.exp(log_importance_weights) + coef_1 = torch.ones_like(log_importance_weights) \ + if self.loss_type == 'm2po' else torch.exp(log_importance_weights) # advantages is per-token [B, T] (expanded at batch construction so the OPD-RL signed # teacher log-ratio is added per token). Edge loss types that need a per-sequence @@ -1047,6 +1065,7 @@ def _compute_loss_and_metrics(self, model, model_inputs: Dict[str, Any], grpo_ba 'off_policy_sequence_mask.') fipo_metrics = None + m2po_metrics = None if self.loss_type == 'cispo': clamped_ratios = torch.clamp(coef_1, max=self.epsilon_high).detach() per_token_loss = -clamped_ratios * advantages * per_token_logps @@ -1057,6 +1076,13 @@ def _compute_loss_and_metrics(self, model, model_inputs: Dict[str, Any], grpo_ba soft_gate = torch.where(is_positive, gate_pos, gate_neg) per_token_loss = -soft_gate * advantages + elif self.loss_type == 'm2po': + per_token_loss, _, m2po_metrics = compute_m2po_token_loss( + log_ratio=log_ratio, + completion_mask=completion_mask, + advantages=advantages, + m2_threshold=self.m2_threshold, + ) elif self.loss_type == 'real': per_token_loss = torch.zeros_like(per_token_logps) elif self.loss_type in ['grpo', 'bnpo', 'dr_grpo', 'dapo', 'fipo']: @@ -1137,8 +1163,8 @@ def _compute_loss_and_metrics(self, model, model_inputs: Dict[str, Any], grpo_ba if self.beta != 0.0: kl_loss = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) loss = loss + kl_loss * self.beta - elif self.loss_type in ['cispo', 'dapo', 'fipo']: - # CISPO, DAPO, and FIPO: Normalize by total completion tokens across all processes + elif self.loss_type in ['cispo', 'dapo', 'fipo', 'm2po']: + # Token-level objectives: normalize by all valid completion tokens across processes. normalizer = grpo_batch.num_items_in_batch / self.accelerator.num_processes loss = (per_token_loss * completion_mask).sum() / normalizer else: @@ -1188,6 +1214,9 @@ def masked_batch_mean(x): 'safety_keep_ratio': self.accelerator.gather_for_metrics(fipo_safety_keep).nanmean().item(), } + if m2po_metrics is not None: + metrics_data['m2po'] = {key: value.item() for key, value in m2po_metrics.items()} + if per_token_kl is not None: mean_kl = masked_batch_mean(per_token_kl) metrics_data['kl'] = self.accelerator.gather_for_metrics(mean_kl).nanmean().item() @@ -1203,7 +1232,7 @@ def masked_batch_mean(x): cispo_clip_ratio = masked_batch_mean(is_cispo_clipped.float()) gathered_cispo_clip_ratio = self.accelerator.gather_for_metrics(cispo_clip_ratio) metrics_data['clipping'] = {'cispo_clip_ratio': gathered_cispo_clip_ratio.nanmean().item()} - elif self.loss_type in ['sapo', 'real']: + elif self.loss_type in ['sapo', 'real', 'm2po']: pass else: is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0) @@ -1260,6 +1289,10 @@ def _update_metrics(self, metrics_data): for key, value in metrics_data['fipo'].items(): self._metrics[mode][f'fipo/{key}'].append(value) + if 'm2po' in metrics_data: + for key in ['m2_before', 'm2_after', 'masked_fraction', 'trust_region_fraction']: + self._metrics[mode][f'm2po/{key}'].append(metrics_data['m2po'][key]) + # Update clipping metrics if 'clipping' in metrics_data: clipping = metrics_data['clipping'] @@ -1344,6 +1377,7 @@ def _aggregate_and_update_metrics(self, all_metrics_data, mode): cispo_clip_values = [] entropy_thresholds = [] fipo_values = {} + m2po_values = [] for chunk_metrics, chunk_weight in all_metrics_data: chunk_tokens = chunk_metrics['completion_token_count'] @@ -1371,6 +1405,9 @@ def _aggregate_and_update_metrics(self, all_metrics_data, mode): for key, value in chunk_metrics['fipo'].items(): fipo_values.setdefault(key, []).append((value, weight)) + if 'm2po' in chunk_metrics: + m2po_values.append(chunk_metrics['m2po']) + # Collect clipping metrics (weighted by tokens) if 'clipping' in chunk_metrics: clipping = chunk_metrics['clipping'] @@ -1423,6 +1460,29 @@ def weighted_avg(values): if fipo_values: aggregated_metrics['fipo'] = {key: weighted_avg(values) for key, values in fipo_values.items()} + if m2po_values: + valid_count = sum(value['valid_count'] for value in m2po_values) + trust_count = sum(value['trust_region_count'] for value in m2po_values) + kept_count = sum(value['kept_trust_region_count'] for value in m2po_values) + masked_count = trust_count - kept_count + aggregated_metrics['m2po'] = { + 'm2_before': + sum(value['m2_before'] * value['trust_region_count'] for value in m2po_values) / max(trust_count, 1.0), + 'm2_after': + sum(value['m2_after'] * value['kept_trust_region_count'] + for value in m2po_values) / max(kept_count, 1.0), + 'masked_fraction': + masked_count / max(valid_count, 1.0), + 'trust_region_fraction': + trust_count / max(valid_count, 1.0), + 'valid_count': + valid_count, + 'trust_region_count': + trust_count, + 'kept_trust_region_count': + kept_count, + } + # Update metrics self._update_metrics(aggregated_metrics) @@ -2104,6 +2164,7 @@ def _collect_config_info(self) -> Dict[str, str]: 'offpolicy_sequence_mask': 'enable' if self.off_policy_sequence_mask_delta is not None else 'disable', 'rollout_importance_sampling': 'enable' if self.rollout_importance_sampling_mode is not None else 'disable', 'loss_type': str(self.loss_type), + 'm2_threshold': str(self.m2_threshold), } return config @@ -2143,6 +2204,9 @@ def _prepare_algorithm_params(self): self.fipo_clip_high_only = args.fipo_clip_high_only self.fipo_safety_threshold = args.fipo_safety_threshold + # M2PO, https://arxiv.org/abs/2510.01161 + self.m2_threshold = args.m2_threshold + # RLOO, self.advantage_estimator = args.advantage_estimator self.kl_in_reward = args.kl_in_reward diff --git a/tests/train/test_m2po.py b/tests/train/test_m2po.py new file mode 100644 index 0000000000..78b9aef63c --- /dev/null +++ b/tests/train/test_m2po.py @@ -0,0 +1,239 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import pytest +import tempfile +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from pathlib import Path + +from swift.rl_core.m2po import (compute_m2po_log_ratio, compute_m2po_mask, compute_m2po_masks_for_batches, + compute_m2po_token_loss, compute_m2po_token_loss_from_mask) + + +def test_m2po_log_ratio_prefers_rollout_policy_and_falls_back_to_old_policy(): + current = torch.tensor([[-1.0, -2.0]]) + old = torch.tensor([[-1.1, -2.1]]) + rollout = torch.tensor([[-1.5, -2.5]]) + + torch.testing.assert_close(compute_m2po_log_ratio(current, old, rollout), current - rollout) + torch.testing.assert_close(compute_m2po_log_ratio(current, old), current - old) + + with pytest.raises(ValueError, match='behavior policy'): + compute_m2po_log_ratio(current, old, allow_old_policy_fallback=False) + + +def test_m2po_selects_once_across_all_optimizer_micro_batches(): + log_ratios = [torch.tensor([[0.05, 0.4]]), torch.tensor([[0.19, 0.21]])] + completion_masks = [torch.ones_like(value, dtype=torch.bool) for value in log_ratios] + advantages = [torch.ones_like(value) for value in log_ratios] + + batch_masks, metrics = compute_m2po_masks_for_batches(log_ratios, completion_masks, advantages, m2_threshold=0.04) + independently_selected = [ + compute_m2po_mask(log_ratio, completion_mask, advantage, m2_threshold=0.04)[0] + for log_ratio, completion_mask, advantage in zip(log_ratios, completion_masks, advantages) + ] + + assert [mask.tolist() for mask in batch_masks] == [[[True, False]], [[True, True]]] + assert [mask.tolist() for mask in independently_selected] == [[[True, False]], [[True, False]]] + assert metrics['masked_fraction'].item() == pytest.approx(0.25) + + +def test_m2po_masks_largest_trust_region_outlier(): + log_ratio = torch.tensor([[0.1, 0.2, 0.3]]) + completion_mask = torch.ones_like(log_ratio, dtype=torch.bool) + advantages = torch.ones_like(log_ratio) + + mask, metrics = compute_m2po_mask(log_ratio, completion_mask, advantages, m2_threshold=0.04) + + assert torch.equal(mask, torch.tensor([[True, True, False]])) + assert metrics['m2_before'].item() == pytest.approx((0.01 + 0.04 + 0.09) / 3) + assert metrics['m2_after'].item() == pytest.approx(0.025) + assert metrics['masked_fraction'].item() == pytest.approx(1 / 3) + assert metrics['trust_region_fraction'].item() == pytest.approx(1.0) + + +def test_m2po_only_constrains_active_ppo_quadrants_and_valid_tokens(): + log_ratio = torch.tensor([[1.0, -1.0, 1.0, -1.0, 2.0]]) + advantages = torch.tensor([[-1.0, 1.0, 1.0, -1.0, 1.0]]) + completion_mask = torch.tensor([[True, True, True, True, False]]) + + mask, metrics = compute_m2po_mask(log_ratio, completion_mask, advantages, m2_threshold=0.1) + + assert torch.equal(mask, torch.tensor([[True, True, False, False, False]])) + assert metrics['trust_region_fraction'].item() == pytest.approx(0.5) + + +def test_m2po_vectorized_selection_matches_algorithm_one(): + threshold = 0.04 + for seed in range(10): + generator = torch.Generator().manual_seed(seed) + log_ratio = 0.8 * torch.randn(4, 7, generator=generator) + advantages = torch.randn(4, 7, generator=generator) + completion_mask = torch.rand(4, 7, generator=generator) > 0.2 + + actual, _ = compute_m2po_mask(log_ratio, completion_mask, advantages, threshold) + + expected = completion_mask.clone() + trust_region = completion_mask & (((advantages > 0) & (log_ratio > 0)) | ((advantages < 0) & (log_ratio < 0))) + active = torch.nonzero(trust_region.reshape(-1), as_tuple=False).squeeze(-1) + second_moment = log_ratio.float().square().reshape(-1) + while active.numel() and second_moment[active].mean() > threshold: + largest = torch.argmax(second_moment[active]) + expected.reshape(-1)[active[largest]] = False + active = torch.cat((active[:largest], active[largest + 1:])) + + assert torch.equal(actual, expected) + + +def test_m2po_loss_keeps_original_denominator_and_masks_gradients(): + log_ratio = torch.tensor([[0.0, 1.0]], requires_grad=True) + completion_mask = torch.ones_like(log_ratio, dtype=torch.bool) + advantages = torch.ones_like(log_ratio) + + per_token_loss, mask, metrics = compute_m2po_token_loss(log_ratio, completion_mask, advantages, m2_threshold=0.1) + loss = (per_token_loss * completion_mask).sum() / completion_mask.sum() + loss.backward() + + assert torch.equal(mask, torch.tensor([[True, False]])) + assert loss.item() == pytest.approx(-0.5) + assert torch.allclose(log_ratio.grad, torch.tensor([[-0.5, 0.0]])) + assert metrics['masked_fraction'].item() == pytest.approx(0.5) + + +def test_m2po_empty_mask_and_extreme_non_trust_ratio_are_finite(): + empty_log_ratio = torch.tensor([[1.0, -1.0]], requires_grad=True) + empty_mask = torch.zeros_like(empty_log_ratio, dtype=torch.bool) + empty_loss, final_mask, metrics = compute_m2po_token_loss(empty_log_ratio, empty_mask, + torch.ones_like(empty_log_ratio)) + + assert not final_mask.any() + assert empty_loss.sum().item() == pytest.approx(0.0) + assert metrics['masked_fraction'].item() == pytest.approx(0.0) + + extreme_log_ratio = torch.tensor([[100.0]], requires_grad=True) + extreme_loss, _, _ = compute_m2po_token_loss(extreme_log_ratio, + torch.ones_like(extreme_log_ratio, + dtype=torch.bool), -torch.ones_like(extreme_log_ratio)) + extreme_loss.sum().backward() + + assert torch.isfinite(extreme_loss).all() + assert torch.isfinite(extreme_log_ratio.grad).all() + assert extreme_log_ratio.grad.item() > 0 + + padded_log_ratio = torch.tensor([[float('nan'), 0.0]], requires_grad=True) + padded_advantages = torch.tensor([[float('nan'), 1.0]]) + padded_loss = compute_m2po_token_loss_from_mask(padded_log_ratio, padded_advantages, torch.tensor([[False, True]])) + padded_loss.sum().backward() + + assert torch.isfinite(padded_loss).all() + assert torch.isfinite(padded_log_ratio.grad).all() + assert padded_log_ratio.grad[0, 0].item() == pytest.approx(0.0) + + +def test_m2po_rejects_invalid_inputs(): + values = torch.zeros(1, 2) + mask = torch.ones_like(values, dtype=torch.bool) + advantages = torch.ones_like(values) + + with pytest.raises(ValueError, match='non-negative'): + compute_m2po_mask(values, mask, advantages, m2_threshold=-0.01) + with pytest.raises(ValueError, match='finite'): + compute_m2po_mask(values, mask, advantages, m2_threshold=float('nan')) + with pytest.raises(ValueError, match='identical shapes'): + compute_m2po_mask(values, mask[:, :1], advantages) + with pytest.raises(ValueError, match='non-finite'): + compute_m2po_mask(torch.tensor([[float('nan'), 0.0]]), mask, advantages) + with pytest.raises(ValueError, match='advantages'): + compute_m2po_mask(values, mask, torch.tensor([[float('nan'), 0.0]])) + + +def _distributed_m2po_worker(rank, world_size, init_method, result_queue): + dist.init_process_group('gloo', rank=rank, world_size=world_size, init_method=init_method) + try: + local_ratios = [torch.tensor([[0.05, 0.4]]), torch.tensor([[0.2, 0.21]])][rank] + ratio_batches = [local_ratios[:, :1], local_ratios[:, 1:]] + local_masks, metrics = compute_m2po_masks_for_batches( + ratio_batches, + [torch.ones_like(value, dtype=torch.bool) for value in ratio_batches], + [torch.ones_like(value) for value in ratio_batches], + m2_threshold=0.04, + ) + local_mask = torch.cat(local_masks, dim=-1) + result_queue.put((rank, local_mask.tolist(), {key: value.item() for key, value in metrics.items()})) + finally: + dist.destroy_process_group() + + +def test_m2po_uses_one_threshold_across_distributed_ranks(): + world_size = 2 + spawn_context = mp.get_context('spawn') + result_queue = spawn_context.SimpleQueue() + with tempfile.TemporaryDirectory() as tmp_dir: + init_method = (Path(tmp_dir) / 'm2po_dist_init').as_uri() + mp.spawn( + _distributed_m2po_worker, + args=(world_size, init_method, result_queue), + nprocs=world_size, + join=True, + ) + + results = sorted([result_queue.get() for _ in range(world_size)]) + assert results[0][1] == [[True, False]] + assert results[1][1] == [[True, True]] + for _, _, metrics in results: + assert metrics['m2_after'] == pytest.approx((0.05**2 + 0.2**2 + 0.21**2) / 3) + assert metrics['masked_fraction'] == pytest.approx(0.25) + + +def _distributed_cp_replica_worker(rank, world_size, init_method, result_queue): + dist.init_process_group('gloo', rank=rank, world_size=world_size, init_method=init_method) + try: + # Context-parallel ranks reconstruct the same full sequence before the loss. + local_ratios = torch.tensor([[0.1, 0.3]]) + completion_mask = torch.ones_like(local_ratios, dtype=torch.bool) + advantages = torch.ones_like(local_ratios) + + replicated_mask, _ = compute_m2po_mask( + local_ratios, + completion_mask, + advantages, + m2_threshold=0.04, + ) + + # Simulate pure-DP groups for dp_size=1, cp_size=world_size. All ranks must + # create the groups in the same order even though each rank uses only its own. + pure_dp_groups = [dist.new_group(ranks=[group_rank]) for group_rank in range(world_size)] + pure_dp_mask, _ = compute_m2po_mask( + local_ratios, + completion_mask, + advantages, + m2_threshold=0.04, + process_group=pure_dp_groups[rank], + ) + result_queue.put((rank, replicated_mask.tolist(), pure_dp_mask.tolist())) + finally: + dist.destroy_process_group() + + +def test_m2po_excludes_context_parallel_replicas_from_threshold_selection(): + world_size = 2 + spawn_context = mp.get_context('spawn') + result_queue = spawn_context.SimpleQueue() + with tempfile.TemporaryDirectory() as tmp_dir: + init_method = (Path(tmp_dir) / 'm2po_cp_replica_init').as_uri() + mp.spawn( + _distributed_cp_replica_worker, + args=(world_size, init_method, result_queue), + nprocs=world_size, + join=True, + ) + + results = sorted([result_queue.get() for _ in range(world_size)]) + replicated_masks = [replicated_mask for _, replicated_mask, _ in results] + pure_dp_masks = [pure_dp_mask for _, _, pure_dp_mask in results] + + # Counting both CP replicas keeps three of four duplicated values, so only one + # replica masks the 0.3 token. Pure-DP selection gives both replicas the mask + # that Algorithm 1 produces for the single logical sequence. + assert sorted(replicated_masks) == sorted([[[True, True]], [[True, False]]]) + assert pure_dp_masks == [[[True, False]], [[True, False]]]