训推不一致(Training-Inference Discrepancy)

关键概念

IS 权重分解

Importance Sampling 权重分解

• $\pi$：训练引擎的策略；$\mu$：推理引擎的策略。

• ${\theta }_{old}$：rollout时的参数；$\theta$：当前更新后的参数。

  $$\frac{{\pi }_{\theta }(y_t\mid x,y_{<t})}{{\mu }_{{\theta }_{old}}(y_t\mid x,y_{<t})}=\underset{\text{训推不一致}}{\underbrace{\frac{{\pi }_{{\theta }_{old}}(y_t\mid x,y_{<t})}{{\mu }_{{\theta }_{old}}(y_t\mid x,y_{<t})}}}\times \underset{\mathrm{策}\mathrm{略}\mathrm{陈}\mathrm{旧}\mathrm{度}}{\underbrace{\frac{{\pi }_{\theta }(y_t\mid x,y_{<t})}{{\pi }_{{\theta }_{old}}(y_t\mid x,y_{<t})}}}$$

策略陈旧度

• 训练引擎内，新策略${\pi }_{\theta }$和采样策略${\pi }_{{\theta }_{\text{old}}}$的差异，off-policy更新。

训推不一致

• 同一权重同一输入，训练引擎${\pi }_{{\theta }_{\text{old}}}$和推理引擎${\mu }_{{\theta }_{\text{old}}}$的数值差异。

问题

Policy Mistmatch 现象

训练和推理实现导致Policy不一致

• 训推2个引擎，由于计算精度、算子实现、并行方式不同等原因，

• 导致权重和输入相同，但token概率却不相同，产生微小且持续累积的数值偏差

  • 甚至可能完全相反，一个0 一个1。
  $$\text{rollout}\mu (a\mid s)\ne p(a\mid s)\text{train}$$

• Rollout：SGlang / VLLM BF16，TokenA = $0.2$， TokenB = ${10}^{-3}$

• Train：Megatron / FSDP FP32，TokenA = $0.25$，TokenB = ${10}^{-6}$

不一致带来的问题：不同token梯度贡献偏轻或偏重

从权重不一致来理解

• 把 on-policy 悄悄变成了 off-policy 训练。
• 比如TokenB
  • 训推绝对值差别小，但相对差达1000倍。Rollout 1000次出1，训练 100万次出现1次。
  • 训练端会看到自己分布下，极少出现的动作，被频繁采样送进来。
• 比如 TokenA
  • 按 0.2 概率采样，但 Train 希望估计自身策略对应的更新，应当是0.25，
  • 在实际采样端数据相对偏少。
• 直接等权训练，会让部分 token 梯度贡献偏轻或偏重，差异较大时，影响方向稳定性。
  • IS 采样的缺点

概率差异影响训练

例子

动作	Rollout u	Train p
A	0.8	0.6
B	0.2	0.4

策略梯度

• 优势：抑制还是鼓励；动作梯度：更新方向；采样概率：决定动作在这批数据出现多少次。

  $$\begin{array}{l}\mathrm{\nabla }J(\theta )=\frac{1}{N}\sum _{n=1}^N\sum _{t=0}^{T_n}\underset{\mathrm{权}\mathrm{重},\mathrm{优}\mathrm{势}\mathrm{值}}{\underbrace{A_{\pi }(s_t^n,a_t^n)}}\cdot \underset{\mathrm{动}\mathrm{作}a\mathrm{的}\mathrm{梯}\mathrm{度}}{\underbrace{\mathrm{\nabla }\log p_{\theta }(a_t^n\mid s_t^n)}}\end{array}$$$$g=E_{a\sim p(\cdot |s)}[A(s,a){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a\mid s)]$$

• 单token 梯度

  $$f(a)=A(s,a){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a\mid s)$$

问题分析

• 100个采样样本：80个A，20个B

• 样本：直接平均梯度

  $$0.8f(A)+0.2f(B)$$

• 训练：目标平均梯度

  $$0.6f(A)+0.4f(B)$$

• 问题：

  • 若直接采用 Rollout平均梯度，则A的更新影响偏大，B的更新影响偏小。
  • 比例变化，可能改变整体更新方向。

• 为何一定得对应当前分布，而不能是采样梯度，不能像SFT那样？

  • 由策略梯度定义，估计当前策略的期望回报，期望分布就是当前策略，而非采样策略。
  • SFT 则可以把外部分布 作为自己训练目标。

解法

• 下文，每个token 计算并乘以 rollout_is权重

  $$w(A)=\frac{0.6}{0.8}=0.75,w(A)=\frac{0.4}{0.2}=2$$

• 训练梯度

  $$w(A)\times 0.8f(A)+w(B)\times 0.2f(B)=0.6f(A)+0.4f(B)$$

原因

数值精度差异

训推不一致原因

1. 数值精度差异

• 浮点加法不满足结合律：$(a+b)+c\ne a+(b+c)$
• 训练引擎和推理引擎，可能有一下差异，导致浮点数运算顺序、数值舍入不同
  • 并行策略不同：Megatron TP=2，SGlang TP=8，Allreduce存在差异。
  • Kernal 实现方式不同：如Flash Attention V2 vs Triton attention等。
  • 精度格式不同或精度本身过低：Rollout FP8，训练 BF16

MoE 路由不一致

:::

2. MoE Router 路由不一致

• MoE 仅激活部分专家，训练推框架差异，可能导致相同输入，专家路由不同，选择不同专家。
• 不一致性引发激活参数子空间 突变，导致优化过程不稳定。

:::

工程配置未对齐

工程配置未对齐

3. 工程配置不对齐

• 异步时间滞后：rollout 使用过旧的权重
• 采样参数不同步：temperature等未对齐

假设模型预测下一个 token 的概率：

token	原始概率	累计概率
A	0.50	0.50
B	0.30	0.80
C	0.15	0.95
D	0.05	1.00

top-p 采样回放

采样

• top_p=0.8，集合S = {A, B}
• 在集合S内归一化，P(A|S)=0.5/0.8=0.625，P(B|S)=0.3/0.8=0.375
• 模型B原始概率：0.3；但 Rollout 实际概率 0.375

训练侧

• 如果直接用原始概率，计算IS权重

  $$\frac{0.3}{0.375}=0.8$$

• 参数明明没变，概率比却不是1，这就是采样差异。

采样候选集回放

• 训练时，也从S进行选择， 则B概率也是0.375，比率为1。

算法解决方案

Loss乘以rollout_is权重

token 更新权重

token 乘以权重 或 Mask

• token 权重：改变采样样本${\pi }_{\text{rollout}}$在平均梯度中的权重，对应目标分布${\pi }_{old}$。
  • 也就是设个权重，控制其更新贡献不能太大或直接mask掉。

$$w=\frac{{\pi }_{\text{old}}}{{\pi }_{\text{rollout}}},\mathrm{截}\mathrm{断}{w}_{\text{TIS}}=min(w,C)$$

• 为何需要截断？
  • 权重太大，会放大样本梯度贡献，太大了会带来波动。
  • 截断后的权重，近似修正分布差异，会引入一定偏差。
  • C 一般可设为2
• 计算代码 见下文
  • log_ratio = old_log_prob - rollout_log_prob,
  • rollout_is_weights = torch.exp(log_ratio)
• pg_loss 乘以权重：
  • pg_losses = pg_losses * rollout_is_weights

核心公式

文章博客

• Your Efficient RL Framework Secretly Brings You Off-Policy RL Training

Reinforce 算法为例

• 原始

$${\mathbb{E}}_{a\sim {\pi }_{\text{sampler}}(\theta )}[R(a)\cdot {\mathrm{\nabla }}_{\theta }\log {\pi }_{\text{learner}}(a,\theta )],$$

• 乘以权重后

  $${\mathbb{E}}_{a\sim {\pi }_{\text{sampler}}(\theta )}[\frac{{\pi }_{\text{learner}}(a,\theta )}{{\pi }_{\text{sampler}}(a,\theta )}\cdot R(a)\cdot {\mathrm{\nabla }}_{\theta }\log {\pi }_{\text{learner}}(a,\theta )].$$

• 截断IS，保持稳定

  $${\mathbb{E}}_{a\sim {\pi }_{\text{sampler}}(\theta )}[\underset{\text{truncated importance ratio}}{\underbrace{min(\frac{{\pi }_{\text{learner}}(a,\theta )}{{\pi }_{\text{sampler}}(a,\theta )},C)}}\cdot R(a)\cdot {\mathrm{\nabla }}_{\theta }\log {\pi }_{\text{learner}}(a,\theta )],$$

PPO 公式

• 原始

  $${\mathbb{E}}_{a\sim {\pi }_{{\theta }_{\mathrm{old}}}}[{\mathrm{\nabla }}_{\theta }min(\frac{{\pi }_{\theta }(a)}{{\pi }_{{\theta }_{\mathrm{old}}}(a)}\hat{A},\mathrm{clip}(\frac{{\pi }_{\theta }(a)}{{\pi }_{{\theta }_{\mathrm{old}}}(a)},1-ϵ,1+ϵ)\hat{A})].$$

• 修正后

  $${\mathbb{E}}_{a\sim {\pi }_{\text{sampler}}({\theta }_{\mathrm{old}})}[{\mathrm{\nabla }}_{\theta }min(\frac{{\pi }_{\text{learner}}(a,\theta )}{{\pi }_{\text{learner}}(a,{\theta }_{\mathrm{old}})}\hat{A},\mathrm{clip}(\frac{{\pi }_{\text{learner}}(a,\theta )}{{\pi }_{\text{learner}}(a,{\theta }_{\mathrm{old}})},1-ϵ,1+ϵ)\hat{A})].$$

• 截断IS

  $${\mathbb{E}}_{a\sim {\pi }_{\text{sampler}}({\theta }_{\mathrm{old}})}[\underset{\text{truncated importance ratio}}{\underbrace{min(\frac{{\pi }_{\text{learner}}(a,\theta )}{{\pi }_{\text{sampler}}(a,\theta )},C)}}{\mathrm{\nabla }}_{\theta }min(\frac{{\pi }_{\text{learner}}(a,\theta )}{{\pi }_{\text{learner}}(a,{\theta }_{\mathrm{old}})}\hat{A},\mathrm{clip}(\frac{{\pi }_{\text{learner}}(a,\theta )}{{\pi }_{\text{learner}}(a,{\theta }_{\mathrm{old}})},1-ϵ,1+ϵ)\hat{A})].$$

方案	做法	作者实验中的问题或效果
Vanilla-IS	直接乘原始 rollout IS 权重	极大权重会放大梯度波动
PPO-IS	将 rollout 概率直接作为 PPO ratio 的分母	训推差异可能让样本在更新前就触发 clipping
TIS	保留原 PPO ratio，额外乘上截断的 rollout 权重	在他们测试的较大训推差异场景中表现更稳定

TIS 截断策略

TIS Mask 策略

普通截断

• rollout_is_weights = raw_rollout_is_weights.clamp(max=rollout_is_threshold_upper)

IcePop (蚂蚁)

• 仅保留 is 留在特定区间的token，如果超过，权重则为0，直接全部Mask掉。

• 比如 [0.5, 2.0]，排除概率过于极端的样本，比如大于2 全 Mask。

  $$w=\frac{{\pi }_{\text{old}}}{{\pi }_{\text{rollout}}}>C\to \text{Mask}$$

• GLM5 就是使用的这个。

IPO

• 按照概率差绝对值来Mask token，不考虑rollout_is 大小，默认为0.3

  $$|{\pi }_{\text{old}}-{\pi }_{\text{rollout}}|>ϵ\to Mask$$

• Prime-RL IPO，DPPO

• 比如：

  • 从0.01变成0.06，6倍变化，直接被Mask，其实概率没啥变化。
  • 从0.9下降到0.2，0.22倍变化，直接下降了0.7的概率，策略发生很大变化，却被保留了。

Routing Replay (解决MoE问题)

提示

相关笔记

• Compute log_prob 源代码

提示

方法1：训练时强制采用采样路由路径

• 确保优化和推理的专家一致。
• 使用者：DeepSeek V3.2

方法2：R2 - Vanilla Routing Replay (策略陈旧度)

• 在训练引擎上固定路由。
• 梯度更新时，重放训练引擎 old policy 的路由 $e_{\text{old},t}^{\pi }$
• 减少策略陈旧度中 expert routing的影响。
• R2：训练重放 old policy 的路由
  • 计算 old_log_probs → 记录 [2, 9]
  • 更新 actor → 回放 [2, 9]

方法3：R3 - Rollout Routing Replay (训推不一致)

• 在训练引擎上固定路由。
• 梯度更新时，统一重放推理引擎 rollout 的路由 $e_{\text{old},t}^{\mu }$
• 同时减小 训推不一致和策略陈旧度中的影响。
• R3：训练重放 rollout时的路由
  • rollout → 记录 [2, 7]
  • 计算 old_log_probs → 回放 [2, 7]
  • 更新 actor → 回放 [2, 7]

R3 是否一定好 (Mimo 辩证看待 R3）

• R3虽然减少两种差异，
• 但也隐式偏置了 target policy，强制使用旧路由而非自然路由，引入 bias
• 但最新Mimo V2.6 都是使用的 R3

工程方案

减小数值差异

提示

• 切换到 FP16：训练更稳定、更快收敛、性能好，FP16 10位尾数，精度比 BF16 7位尾数 高。
  • 浮点数笔记
• 训推统一 FP8：统一 FP8 比 BF16训练 + FP8 rollout要好。
• 关键层使用FP32： MiniMax-M1 把精度最敏感层 (LM head、MoE router) 用 FP32，其余保持低精度

统一计算 Kernal

信息

• 使用 Bitwise 一致的推理 kernel
  • 替换对 TP Size 敏感的层 为 TP 不变的替代
• 对齐 Attention 实现 / Temperature 处理 / 权重同步
  • 确保训推使用相同的attention kernel，temp 缩放一致，rollout 引擎及时接收最新权重

统一工程配置/兜底逻辑

信息

• 对齐训练推理引擎配置：保持没有人工不对齐的配置误差
• 动态 LR Scheduler：基于 response length 动态触发 LR 衰减，把训推不一致保持在可控水平。

Policy Lag 模型滞后问题

提示

• off-policy：采样策略太旧，和学习策略不同

数据采集和训练之间的一般性分布漂移

实现代码

小心

代码

• ppo/rollout_corr_helper.py

Ray_trainer 调用 rollout_corr_helper

非bypass mode下，计算优势的时候，会调用计算is，代码见 _fit_compute_advantage

def _fit_compute_advantage(self, batch) -> DataProto:
    metrics = self.metrics
    timing_raw = self.timing_raw
    reward_tensor = self.reward_tensor
    reward_extra_infos_dict = self.reward_extra_infos_dict

    with marked_timer("adv", timing_raw, color="brown"):
        # we combine with rule-based rm
        reward_extra_infos_dict: dict[str, list]
        batch.batch["token_level_scores"] = reward_tensor

        if reward_extra_infos_dict:
            batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})

        # compute rewards. apply_kl_penalty if available
        if self.config.algorithm.use_kl_in_reward:
            batch, kl_metrics = apply_kl_penalty(
                batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty
            )
            metrics.update(kl_metrics)
        else:
            batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]

        # Compute rollout correction: IS weights, rejection sampling, and metrics
        # Only runs in decoupled mode (computes once per batch using stable π_old)
        # In bypass mode, this is skipped - actor computes metrics from evolving π_θ vs π_rollout
        rollout_corr_config = self.config.algorithm.get("rollout_correction", None)
        bypass_recomputing_logprobs = rollout_corr_config and rollout_corr_config.get("bypass_mode", False)
        if (
            rollout_corr_config is not None
            and "rollout_log_probs" in batch.batch
            and not bypass_recomputing_logprobs  # Only in decoupled mode
        ):
            from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_add_to_batch

            # Compute IS weights, apply rejection sampling, compute metrics
            batch, is_metrics = compute_rollout_correction_and_add_to_batch(batch, rollout_corr_config)  
            # IS and off-policy metrics already have rollout_corr/ prefix
            metrics.update(is_metrics)

        # compute advantages, executed on the driver process
        norm_adv_by_std_in_grpo = self.config.algorithm.get(
            "norm_adv_by_std_in_grpo", True
        )  # GRPO adv normalization factor

        batch = compute_advantage(
            batch,
            adv_estimator=self.config.algorithm.adv_estimator,
            gamma=self.config.algorithm.gamma,
            lam=self.config.algorithm.lam,
            num_repeat=self.config.actor_rollout_ref.rollout.n,
            norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
            config=self.config.algorithm,
        )
    return batch

rollout_corr 主函数

笔记

• 计算 log_ratio
• 计算 rollout_is_weights
• 计算 modified_response_mask
• 计算 offpolicy_metrics (KL, PPL, χ²)

def compute_rollout_correction_and_rejection_mask(
    old_log_prob: torch.Tensor,
    rollout_log_prob: torch.Tensor,
    response_mask: torch.Tensor,
    rollout_is: Optional[str] = None,
    rollout_is_threshold: Optional[str | float] = 2.0,
    rollout_is_batch_normalize: bool = False,
    rollout_rs: Optional[str] = None,
    rollout_rs_threshold: Optional[str | float] = None,
) -> tuple[Optional[DataProto], torch.Tensor, dict[str, float]]:
    """Unified interface for computing IS weights and rejection masks.

    This function combines IS weight calculation (truncated) and rejection sampling (masked)
    into a single pipeline.

    Key design:
    - Separation of IS weights (for variance reduction) and rejection masks (for sample filtering)
    - Comprehensive metrics tracking for mismatch diagnosis

    Args:
        old_log_prob: Log probabilities from the training policy (e.g., FSDP FP32),
            shape (batch_size, seq_length).
        rollout_log_prob: Log probabilities from the rollout policy (e.g., vLLM BF16),
            shape (batch_size, seq_length).
        response_mask: Binary mask for valid tokens (1=valid, 0=padding),
            shape (batch_size, seq_length).
        rollout_is: IS weight aggregation level (see compute_rollout_correction_weights for options).
            Set to None to disable IS weight computation.
        rollout_is_threshold: Threshold specification for IS weights.
            Single float implies TIS; "lower_upper" implies IcePop.
        rollout_rs: Rejection sampling aggregation modes as a comma separated string
            (see compute_rollout_rejection_mask for the full list). Set to None to disable
            rejection sampling.
        rollout_rs_threshold: Threshold specification string (see compute_rollout_rejection_mask for details).
            Provide one threshold per option (comma separated). For K1-style options, specify
            ``lower_upper`` to denote the lower/upper ratio bounds.
        rollout_is_batch_normalize: Whether to normalize IS weights to have mean=1.0 per batch.
            Default: False.

    Returns:
        Tuple containing:
            rollout_is_weights_proto: DataProto with IS weights (None if rollout_is is None),
                key "rollout_is_weights", shape (batch_size, seq_length).
            modified_response_mask: Response mask with rejection sampling applied,
                shape (batch_size, seq_length).
            metrics: Dictionary of all metrics (prefixed with "rollout_corr/"), including:
                - IS weight statistics
                - Rejection sampling rates
                - Policy mismatch metrics (KL, PPL, etc.)
    """
    # Validate input masks
    if not response_mask.any():
        raise ValueError("response_mask must contain at least one valid token (1).")
    if old_log_prob.shape != rollout_log_prob.shape:
        raise ValueError(
            f"old_log_prob shape {old_log_prob.shape} does not match rollout_log_prob shape {rollout_log_prob.shape}."
        )
    if old_log_prob.shape != response_mask.shape:
        raise ValueError(
            f"log_prob shape {old_log_prob.shape} does not match response_mask shape {response_mask.shape}."
        )

    # Step 1: Compute log ratio (log(π_train / π_rollout))
    log_ratio: torch.Tensor = old_log_prob - rollout_log_prob  
    metrics: dict[str, float] = {}

    # Step 2: Compute IS weights (if enabled)
    rollout_is_weights: Optional[torch.Tensor] = None
    if rollout_is is not None and rollout_is_threshold is not None:
        rollout_is_weights, is_metrics = compute_rollout_correction_weights(  
            log_ratio=log_ratio,
            response_mask=response_mask,
            rollout_is=rollout_is,
            rollout_is_threshold=rollout_is_threshold,
            rollout_is_batch_normalize=rollout_is_batch_normalize,
        )
        metrics.update(is_metrics)

    # Step 3: Compute rejection mask (if enabled)
    modified_response_mask: torch.Tensor = response_mask.clone()
    if rollout_rs is not None:
        if rollout_rs_threshold is None:
            raise ValueError(
                "rollout_rs_threshold must be explicitly provided when rollout_rs is enabled. "
                "Set rollout_rs_threshold to the desired threshold value."
            )
        modified_response_mask, rs_metrics = compute_rollout_rejection_mask(  
            log_ratio=log_ratio,
            response_mask=response_mask,
            rollout_rs=rollout_rs,
            rollout_rs_threshold=rollout_rs_threshold,
        )
        metrics.update(rs_metrics)

    # Step 4: Compute off-policy metrics (KL, PPL, χ², etc.)
    offpolicy_metrics: dict[str, float] = compute_offpolicy_metrics(  
        old_log_prob=old_log_prob,
        rollout_log_prob=rollout_log_prob,
        response_mask=response_mask,
    )
    metrics.update(offpolicy_metrics)

    # Step 6: Add "rollout_corr/" prefix to all metrics for logging consistency
    metrics_scalar: dict[str, float] = {}
    for key, value in metrics.items():
        if isinstance(value, torch.Tensor):
            metrics_scalar[f"rollout_corr/{key}"] = value.item()
        else:
            metrics_scalar[f"rollout_corr/{key}"] = value

    # Step 7: Wrap IS weights in DataProto for consistency with API
    rollout_is_weights_proto: Optional[DataProto] = None
    if rollout_is_weights is not None:
        rollout_is_weights_proto = DataProto.from_dict(tensors={"rollout_is_weights": rollout_is_weights})

    return rollout_is_weights_proto, modified_response_mask, metrics_scalar

计算 rollot_is_weights

rollout_is_weights 计算方式

token-level

• log_ratio = old_log_prob - rollout_log_prob，记得 clip 一下。
• rollout_is_weights = torch.exp(log_ratio_safe)

sequence-level

• log_ratio = old_log_prob - rollout_log_prob
• log_ratio_sum = verl_F.masked_sum(log_ratio, response_mask, axis=-1).unsqueeze(-1)
  • shape=[bs, 1]，记得 clip 一下。
• rollout_is_weights = torch.exp(log_ratio_sum).expand_as(log_ratio)

TIS 方式

• ICEPop：超过阈值的，都权重设为0，直接mask掉
• 普通clip：只做截断到最大权重。

IS 指标

指标	当前实现
rollout_is_mean	(\frac1N\sum_t w_t)
rollout_is_std	处理后权重的总体标准差
rollout_is_ratio_fraction_high	原始比值超过上界的比例
rollout_is_ratio_fraction_low	原始比值低于下界的比例
rollout_is_oob_ratio	IcePop 原始比值落在区间外的比例

def compute_rollout_correction_weights(
    log_ratio: torch.Tensor,
    response_mask: torch.Tensor,
    rollout_is: str = "token",
    rollout_is_threshold: str | float = 2.0,
    rollout_is_batch_normalize: bool = False,
) -> tuple[torch.Tensor, dict[str, float]]:
    """Compute importance sampling weights to correct for off-policy distribution shifts.

    This function calculates IS weights (π_train / π_rollout) using log ratios for numerical stability.
    It supports multiple aggregation levels and truncates extreme weights to prevent training instability.

    Key design:
    - Log-space computations to avoid overflow
    - Truncation of extreme weights (TIS: Truncated Importance Sampling)
    - Optional batch normalization (normalize to mean=1.0)
    - Metrics tracking for weight distribution analysis

    Args:
        log_ratio: Log ratio of training policy probability to rollout policy probability,
            shape (batch_size, seq_length).
        response_mask: Binary mask for valid tokens (1=valid, 0=padding),
            shape (batch_size, seq_length).
        rollout_is: IS weight aggregation level, must be one of:
            - "token": Per-token weights (biased, low variance)
            - "sequence": Per-sequence weight (product of tokens; unbiased, high variance)
        rollout_is_threshold: Threshold specification for IS weights.
            - Single float or float-like string: TIS, clamp weights to the upper bound
            - "lower_upper" string: IcePop, zero weights outside [lower, upper]
        rollout_is_batch_normalize: Whether to normalize IS weights to have mean=1.0 per batch,
            default False.

    Returns:
        Tuple containing:
            rollout_is_weights: Truncated IS weights (masked to zero for padding tokens),
                shape (batch_size, seq_length). If batch_normalize=True, normalized to mean=1.0.
            metrics: Dictionary of IS weight metrics (all scalars), including:
                - rollout_is_mean/max/min: Statistic of weights (before batch normalization)
                - rollout_is_eff_sample_size: Effective sample size (ESS)
                - rollout_is_seq_*: Sequence-level weight statistics
                - rollout_is_batch_norm_factor: Normalization factor (only if batch_normalize=True)
    """
    # Validate input parameters
    valid_is_levels = {"token", "sequence"}
    if rollout_is not in valid_is_levels:
        raise ValueError(f"Invalid rollout_is: {rollout_is}. Must be one of {valid_is_levels}.")
    rollout_is_threshold_upper, rollout_is_threshold_lower = _parse_rollout_is_threshold(rollout_is_threshold)
    use_icepop = rollout_is_threshold_lower is not None

    # Compute IS weights from log ratio (handles different aggregation levels)
    if rollout_is == "token":
        # Per-token IS weight: exp(log(π_train/π_rollout)) with safety clamp
        log_ratio_for_metrics: torch.Tensor = log_ratio
        log_ratio_safe: torch.Tensor = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
        raw_rollout_is_weights: torch.Tensor = torch.exp(log_ratio_safe) 

    elif rollout_is == "sequence":
        # Sequence-level IS weight: product of token ratios (exp(sum(log ratios)))
        log_ratio_sum: torch.Tensor = verl_F.masked_sum(log_ratio, response_mask, axis=-1).unsqueeze(
            -1
        )  # Shape: (batch_size, 1)
        log_ratio_for_metrics = log_ratio_sum

        log_ratio_sum_safe: torch.Tensor = torch.clamp(log_ratio_sum, min=-SAFETY_BOUND, max=SAFETY_BOUND) 
        raw_rollout_is_weights = torch.exp(log_ratio_sum_safe).expand_as(log_ratio)  # Broadcast to sequence length #

    else:
        raise ValueError(f"Unsupported rollout_is: {rollout_is}")

    # Zero out weights for padding tokens using response mask
    raw_rollout_is_weights = raw_rollout_is_weights * response_mask

    # Apply TIS for a single upper bound and IcePop for a lower_upper string.
    if not use_icepop:
        rollout_is_weights = raw_rollout_is_weights.clamp(max=rollout_is_threshold_upper)
    else:
        assert rollout_is_threshold_lower is not None
        token_kept_mask = (raw_rollout_is_weights >= rollout_is_threshold_lower) & (
            raw_rollout_is_weights <= rollout_is_threshold_upper
        )
        rollout_is_weights = torch.where(
            token_kept_mask, raw_rollout_is_weights, torch.zeros_like(raw_rollout_is_weights)
        )

    # Compute IS weight metrics.
    metrics: dict[str, float] = compute_is_metrics(
        rollout_is_weights=rollout_is_weights,
        raw_rollout_is_weights=raw_rollout_is_weights,
        log_ratio_for_metrics=log_ratio_for_metrics,
        response_mask=response_mask,
        rollout_is=rollout_is,
        rollout_is_threshold=rollout_is_threshold_upper,
        rollout_is_threshold_lower=rollout_is_threshold_lower,
    )
    if use_icepop:
        assert rollout_is_threshold_lower is not None
        oob_mask = (raw_rollout_is_weights < rollout_is_threshold_lower) | (
            raw_rollout_is_weights > rollout_is_threshold_upper
        )
        metrics["rollout_is_oob_ratio"] = verl_F.masked_mean(oob_mask.float(), response_mask).item()

    # Detach weights to prevent gradient flow (mathematically required by IS theory)
    # IS weights change the measure, not the objective. See §3.2.2 in docs/algo/rollout_corr_math.md
    rollout_is_weights = rollout_is_weights.detach()

    # Apply batch normalization if requested
    if rollout_is_batch_normalize:
        # Compute mean based on aggregation level
        mask_float = response_mask.to(dtype=rollout_is_weights.dtype)
        if rollout_is == "token":
            # Token-level: normalize over all token weights
            if torch.distributed.is_available() and torch.distributed.is_initialized():
                weights_mean = verl_F.distributed_masked_mean(rollout_is_weights, mask_float)
            else:
                weights_mean = verl_F.masked_mean(rollout_is_weights, response_mask)
        elif rollout_is == "sequence":
            # Sequence-level: normalize over sequence weights (one weight per sequence)
            # For each sequence, compute mean over valid tokens (they all have the same weight)
            # then average across sequences
            seq_weights = verl_F.masked_mean(rollout_is_weights, response_mask, axis=-1)  # (batch_size,)
            seq_mask = (response_mask.sum(dim=-1) > 0).to(dtype=rollout_is_weights.dtype)
            if torch.distributed.is_available() and torch.distributed.is_initialized():
                weights_mean = verl_F.distributed_masked_mean(seq_weights, seq_mask)
            else:
                weights_mean = (seq_weights * seq_mask).sum() / seq_mask.sum().clamp_min(1e-8)
        else:
            raise ValueError(f"Unsupported rollout_is: {rollout_is}")

        # Normalize to mean=1.0 (avoid division by zero)
        if weights_mean > 1e-8:
            rollout_is_weights = rollout_is_weights / weights_mean
            metrics["rollout_is_batch_norm_factor"] = weights_mean.item()
        else:
            metrics["rollout_is_batch_norm_factor"] = 1.0

    return rollout_is_weights, metrics

Rejection Sampling 计算 modified_response_mask

提示

核心

• 和IcePop 机制有点像，直接更改 response_mask，把相关token给mask掉。
• 根据训推概率差异，把超出阈值的位置，从训练mask中给除掉。

区别

• token IS 调整各 token 的权重
• sequence RS 根据整条序列的差异决定是否保留
• IcePop：只改weights，不改response_mask，分母不变，token-mean 聚合。
  • weights = [1, 1, 1, 0], response_mask = [1, 1, 1, 1]
• RS：改response_mask，分母直接减少。
  • response_mask = [1, 1, 1, 0]，有效token直接减少。

def compute_rollout_rejection_mask(
    log_ratio: torch.Tensor,
    response_mask: torch.Tensor,
    rollout_rs: str = "token_k1",
    rollout_rs_threshold: Optional[str | float] = None,
) -> tuple[torch.Tensor, dict[str, float]]:
    """Compute hard trust region mask using divergence estimators.

    This function enforces a hard trust region constraint by masking tokens/sequences
    where the estimated divergence (between training and rollout policies) exceeds
    a threshold. Unlike PPO's soft clipping, this provides a hard boundary.

    Multiple rejection criteria can be supplied via a comma separated `rollout_rs` string.
    All requested options must pass for a token/sequence to remain valid.

    Supported KL divergence-based modes (ideal = 0.0 unless noted):
    - "token_k{1,2,3}": Token-level divergences.
    - "seq_sum_k{1,2,3}": Sum of token divergences per sequence.
    - "seq_mean_k{1,2,3}": Mean of token divergences per sequence.
    - "seq_max_k{2,3}": Maximum token divergence per sequence.

    Args:
        log_ratio: Log ratio of training policy probability to rollout policy probability,
            shape (batch_size, seq_length).
        response_mask: Binary mask for valid tokens (1=valid, 0=padding),
            shape (batch_size, seq_length).
        rollout_rs: Comma separated rejection sampling options (e.g. "token_k1,seq_sum_k3").
        rollout_rs_threshold: Threshold specification string (required). Provide one entry per
            rollout_rs option separated by commas. Each entry must be a positive number.
            For K1-style options (``*k1``), specify ``lower_upper`` (e.g. ``"0.1_1.2"``)
            to denote lower/upper ratio bounds; other options accept a single upper bound.

    Returns:
        Tuple containing:
            modified_response_mask: Response mask with trust region violations masked (0=rejected),
                shape (batch_size, seq_length).
            metrics: Dictionary of trust region metrics (all scalars).
    """
    if rollout_rs is None or not isinstance(rollout_rs, str):
        raise ValueError("rollout_rs must be a non-empty string (comma separated for multiple options).")
    if rollout_rs_threshold is None:
        raise ValueError("rollout_rs_threshold must be provided for rejection sampling.")

    if log_ratio.shape[0] == 0:
        return response_mask, {}

    # rollout_rs supports chained criteria via comma separation (e.g. "token_k1,seq_mean_k3").
    #            Every listed option must pass; combined_mask aggregates them via logical AND.
    option_modes = [opt.strip() for opt in rollout_rs.split(",") if opt.strip()]
    if not option_modes:
        raise ValueError("rollout_rs must contain at least one valid option.")

    normalized_options: list[str] = []
    seen: set[str] = set()
    for opt in option_modes:
        if opt not in SUPPORTED_ROLLOUT_RS_OPTIONS:
            raise ValueError(
                f"Invalid rollout_rs option: {opt}. Must be one of {sorted(SUPPORTED_ROLLOUT_RS_OPTIONS)}."
            )
        if opt not in seen:
            normalized_options.append(opt)
            seen.add(opt)

    threshold_specs = _parse_rollout_rs_thresholds(normalized_options, rollout_rs_threshold)

    log_ratio_safe: torch.Tensor = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
    token_k1: torch.Tensor = -log_ratio_safe
    token_k2: torch.Tensor = 0.5 * log_ratio_safe**2
    token_k3: torch.Tensor = torch.exp(log_ratio_safe) - 1.0 - log_ratio_safe

    response_mask_bool: torch.Tensor = response_mask.bool()
    seq_valid_mask: torch.Tensor = response_mask.sum(dim=-1) > 0
    # combined_mask accumulates per-option passes; any failure flips tokens to 0.
    combined_mask: torch.Tensor = torch.ones_like(response_mask, dtype=log_ratio.dtype)
    metrics: dict[str, float] = {}

    def _sequence_sum(values: torch.Tensor) -> torch.Tensor:
        return verl_F.masked_sum(values, response_mask, axis=-1)

    def _sequence_mean(values: torch.Tensor) -> torch.Tensor:
        return verl_F.masked_mean(values, response_mask, axis=-1)

    def _sequence_max(values: torch.Tensor) -> torch.Tensor:
        mask_bool = response_mask.bool()
        neg_inf = torch.tensor(float("-inf"), device=values.device, dtype=values.dtype)
        masked_values = values.masked_fill(~mask_bool, neg_inf)
        max_values = masked_values.max(dim=-1).values
        return torch.where(max_values == neg_inf, torch.zeros_like(max_values), max_values)

    for option_name in normalized_options:
        thresholds_info = threshold_specs[option_name]
        is_k1_option = option_name.endswith("k1")
        upper_value = thresholds_info["upper"]
        lower_value = thresholds_info["lower"]
        apply_lower_threshold = is_k1_option
        lower_log: Optional[float] = None
        upper_log: Optional[float] = None

        if is_k1_option:
            if lower_value is None or upper_value is None:
                raise ValueError(
                    f"rollout_rs_threshold for option '{option_name}' must specify both lower and upper bounds."
                )
            lower_log = math.log(lower_value)
            upper_log = math.log(upper_value)
        else:
            if upper_value is None:
                raise ValueError(f"rollout_rs_threshold for option '{option_name}' must specify an upper bound.")

        level = "sequence" if option_name not in TOKEN_LEVEL_ROLLOUT_RS_OPTIONS else "token"

        per_token_stat: torch.Tensor
        per_sequence_stat: Optional[torch.Tensor] = None
        token_keep_bool: torch.Tensor

        if option_name == "token_k1":
            if lower_log is None:
                raise ValueError("Threshold specification for token_k1 must include lower and upper bounds.")
            per_token_stat = token_k1
            token_keep_bool = (per_token_stat >= lower_log) & (per_token_stat <= upper_log)
        elif option_name == "token_k2":
            per_token_stat = token_k2
            token_keep_bool = per_token_stat <= upper_value
        elif option_name == "token_k3":
            per_token_stat = token_k3
            token_keep_bool = per_token_stat <= upper_value
        elif option_name.startswith("seq_sum"):
            if option_name.endswith("k1"):
                if lower_log is None:
                    raise ValueError(
                        f"Threshold specification for option '{option_name}' must include lower and upper bounds."
                    )
                seq_stat = _sequence_sum(token_k1)
                seq_keep_bool_direct = (seq_stat >= lower_log) & (seq_stat <= upper_log)
            elif option_name.endswith("k2"):
                seq_stat = _sequence_sum(token_k2)
                seq_keep_bool_direct = seq_stat <= upper_value
            elif option_name.endswith("k3"):
                seq_stat = _sequence_sum(token_k3)
                seq_keep_bool_direct = seq_stat <= upper_value
            else:
                raise ValueError(f"Unsupported rollout_rs option: {option_name}.")
            per_sequence_stat = seq_stat
            token_keep_bool = seq_keep_bool_direct.unsqueeze(-1).expand_as(response_mask_bool)
            per_token_stat = seq_stat.unsqueeze(-1).expand_as(response_mask)
        elif option_name.startswith("seq_mean"):
            if option_name.endswith("k1"):
                if lower_log is None:
                    raise ValueError(
                        f"Threshold specification for option '{option_name}' must include lower and upper bounds."
                    )
                seq_stat = _sequence_mean(token_k1)
                seq_keep_bool_direct = (seq_stat >= lower_log) & (seq_stat <= upper_log)
            elif option_name.endswith("k2"):
                seq_stat = _sequence_mean(token_k2)
                seq_keep_bool_direct = seq_stat <= upper_value
            elif option_name.endswith("k3"):
                seq_stat = _sequence_mean(token_k3)
                seq_keep_bool_direct = seq_stat <= upper_value
            else:
                raise ValueError(f"Unsupported rollout_rs option: {option_name}.")
            per_sequence_stat = seq_stat
            token_keep_bool = seq_keep_bool_direct.unsqueeze(-1).expand_as(response_mask_bool)
            per_token_stat = seq_stat.unsqueeze(-1).expand_as(response_mask)
        elif option_name.startswith("seq_max"):
            if option_name.endswith("k2"):
                seq_stat = _sequence_max(token_k2)
                seq_keep_bool_direct = seq_stat <= upper_value
            elif option_name.endswith("k3"):
                seq_stat = _sequence_max(token_k3)
                seq_keep_bool_direct = seq_stat <= upper_value
            else:
                raise ValueError(f"Unsupported rollout_rs option: {option_name}.")
            per_sequence_stat = seq_stat
            token_keep_bool = seq_keep_bool_direct.unsqueeze(-1).expand_as(response_mask_bool)
            per_token_stat = seq_stat.unsqueeze(-1).expand_as(response_mask)
        else:
            raise ValueError(f"Unsupported rollout_rs option: {option_name}.")

        metrics_upper_threshold = upper_log if is_k1_option else upper_value
        metrics_lower_threshold = lower_log if (is_k1_option and lower_log is not None) else 0.0

        token_keep_mask = token_keep_bool.to(dtype=log_ratio.dtype)
        combined_mask = combined_mask * token_keep_mask
        seq_keep_bool_tensor = (~((~token_keep_bool) & response_mask_bool)).all(dim=-1)

        option_metrics = compute_rs_metrics(
            option_name=option_name,
            rs_statistic=per_token_stat,
            response_mask=response_mask,
            seq_valid_mask=seq_valid_mask,
            level=level,
            per_sequence_values=per_sequence_stat,
            rollout_rs_threshold=metrics_upper_threshold,
            rollout_rs_threshold_lower=metrics_lower_threshold,
            apply_lower_threshold=apply_lower_threshold,
        )
        metrics.update(option_metrics)

        token_masked_fraction = verl_F.masked_mean(1 - token_keep_mask, response_mask).item()
        seq_valid_float = seq_valid_mask.float()
        if seq_valid_float.sum() > 0:
            seq_keep_float = seq_keep_bool_tensor.to(dtype=log_ratio.dtype)
            seq_masked_fraction = (((1.0 - seq_keep_float) * seq_valid_float).sum() / seq_valid_float.sum()).item()
        else:
            seq_masked_fraction = 0.0
        metrics[f"rollout_rs_{option_name}_masked_fraction"] = token_masked_fraction
        metrics[f"rollout_rs_{option_name}_seq_masked_fraction"] = seq_masked_fraction

    final_mask = combined_mask
    metrics["rollout_rs_masked_fraction"] = verl_F.masked_mean(1 - final_mask, response_mask).item()
    final_keep_bool = (final_mask > 0.5) & response_mask_bool
    seq_has_masked: torch.Tensor = (~final_keep_bool & response_mask_bool).any(dim=-1)
    metrics["rollout_rs_seq_masked_fraction"] = seq_has_masked.float().mean().item()

    modified_response_mask: torch.Tensor = (response_mask * final_mask).to(dtype=response_mask.dtype)
    return modified_response_mask, metrics

Off policy 指标计算

提示

KL 和 K3L

$$\mathtt{\text{kl}}={\mathrm{mean}}_{\text{有效 token}}[\log \mu -\log p]$$$$\mathtt{\text{k3\_kl}}={\mathrm{mean}}_{\text{有效 token}}[e^{\mathrm{\Delta }}-1-\mathrm{\Delta }]$$

PPL 和 log PPL

χ² 与 IS 方差

def compute_offpolicy_metrics(
    old_log_prob: torch.Tensor,
    rollout_log_prob: Optional[torch.Tensor],
    response_mask: torch.Tensor,
) -> dict[str, Any]:
    """Compute off-policy diagnostic metrics (helper function).

    This helper function operates on raw tensors and is used internally by:
    - compute_rollout_correction_and_rejection_mask() in this module (automatically included)
    - Tests (test_rollout_corr.py, test_rollout_corr_integration.py)

    These metrics help diagnose the off-policy gap between rollout and training policies,
    which can arise from:
    - Policy mismatch (e.g., vLLM BF16 vs FSDP FP32)
    - Model staleness (training on trajectories from older checkpoints)
    - General distribution shifts

    Key metrics:
    - kl: Direct KL divergence estimator KL(π_rollout || π_training)
    - k3_kl: K3 KL estimator for stability (more stable for small KL)
    - training_ppl: Perplexity of training policy
    - rollout_ppl: Perplexity of rollout policy
    - log_ppl_diff: Difference in log perplexities
    - ppl_ratio: Ratio of training PPL to rollout PPL
    - chi2_token: Token-level χ² divergence E[ρ²] - 1
    - chi2_seq: Sequence-level χ² divergence E[(∏ρ_t)²] - 1

    Args:
        old_log_prob: Log probabilities from training policy, shape (batch_size, seq_length)
        rollout_log_prob: Log probabilities from rollout policy, shape (batch_size, seq_length)
        response_mask: Mask for valid tokens, shape (batch_size, seq_length)

    Returns:
        Dictionary of off-policy metrics (without prefix)
    """
    # Validate that we have at least one valid token
    assert response_mask.any(), "Expected at least one valid token in response_mask"

    metrics = {}

    # 1. Training policy perplexity (always available)
    # Formula: exp(-1/|T| * Σ log π_training(y_t|y_<t))
    # where |T| is the number of tokens generated by the model
    mean_log_prob_training = verl_F.masked_mean(old_log_prob, response_mask, axis=-1)  # (batch_size,)
    training_ppl = torch.exp(-mean_log_prob_training).mean()  # Batch mean of per-sequence PPL
    metrics["training_ppl"] = training_ppl.detach().item()

    # Also log log-ppl for easier analysis (avoids exponential scale)
    metrics["training_log_ppl"] = (-mean_log_prob_training).mean().detach().item()

    # 2. Compute rollout off-policy metrics (only if rollout_log_probs available)
    if rollout_log_prob is not None:
        # 2a. kl: Direct estimator for KL(π_rollout || π_training)
        # This is the standard KL divergence: E[log(π_rollout) - log(π_training)]
        # Positive value means rollout policy is more confident than training policy
        metrics["kl"] = verl_F.masked_mean(rollout_log_prob - old_log_prob, response_mask).detach().item()

        # 2b. k3_kl: K3 estimator for KL(π_rollout || π_training)
        # More stable for small KL values using: E[exp(log_ratio) - log_ratio - 1]
        # Formula: KL ≈ E[r - log(r) - 1] where r = π_training/π_rollout
        log_ratio = old_log_prob - rollout_log_prob
        k3_kl_matrix = torch.exp(log_ratio) - log_ratio - 1
        metrics["k3_kl"] = verl_F.masked_mean(k3_kl_matrix, response_mask).detach().item()

        # 2c. Rollout policy perplexity
        mean_log_prob_rollout = verl_F.masked_mean(rollout_log_prob, response_mask, axis=-1)  # (batch_size,)
        rollout_ppl = torch.exp(-mean_log_prob_rollout).mean()  # Batch mean of per-sequence PPL
        metrics["rollout_ppl"] = rollout_ppl.detach().item()
        metrics["rollout_log_ppl"] = (-mean_log_prob_rollout).mean().detach().item()

        # 2d. Log PPL difference (sequence-level perplexity difference)
        # log_ppl_diff = mean_log_prob_rollout - mean_log_prob_training
        # Since ppl = exp(-log_prob), we have:
        #   log(ppl_ratio) = log(training_ppl/rollout_ppl) = log_ppl_diff
        # Positive value means training assigns lower probability (higher PPL) than rollout
        log_ppl_diff = mean_log_prob_rollout - mean_log_prob_training
        metrics["log_ppl_diff"] = log_ppl_diff.mean().detach().item()
        metrics["log_ppl_abs_diff"] = log_ppl_diff.abs().mean().detach().item()
        metrics["log_ppl_diff_max"] = log_ppl_diff.max().detach().item()
        metrics["log_ppl_diff_min"] = log_ppl_diff.min().detach().item()

        # 2e. PPL ratio (how much higher is training PPL vs rollout PPL)
        # Compute per-sequence ratio first, then average
        # For numerical stability, compute in log space using log_ppl_diff
        # Note: log_ppl_diff = log(ppl_ratio), so ppl_ratio = exp(log_ppl_diff)
        # This is the inverse of geometric IS: ppl_ratio_i = 1 / geometric_is_i for each sequence
        ppl_ratio = torch.exp(log_ppl_diff).mean()  # mean(exp(log_ppl_diff)) = mean(ppl_ratio_i)
        metrics["ppl_ratio"] = ppl_ratio.detach().item()

        # 2f. Chi-squared divergence: χ²(π_training || π_rollout) = E_μ[ρ²] - 1
        # where ρ = π_training / π_rollout and μ = π_rollout (rollout distribution)
        # This measures the variance of importance sampling weights
        # Token-level: E_token[ρ²] - 1 (averaged over all tokens)
        log_ratio_safe = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
        rho_token = torch.exp(log_ratio_safe)  # ρ = π_training / π_rollout (token-level)
        rho_squared_token = rho_token.square()
        chi2_token = verl_F.masked_mean(rho_squared_token, response_mask) - 1.0
        metrics["chi2_token"] = chi2_token.detach().item()

        # Sequence-level: E_seq[(Π ρ_t)²] - 1 = E_seq[exp(2 * Σ log ρ_t)] - 1
        log_ratio_sum = verl_F.masked_sum(log_ratio, response_mask, axis=-1)  # Σ log ρ_t per sequence
        log_ratio_sum_safe = torch.clamp(log_ratio_sum, min=-SAFETY_BOUND, max=SAFETY_BOUND)
        rho_squared_seq = torch.exp(2.0 * log_ratio_sum_safe)  # (Π ρ_t)²
        chi2_seq = rho_squared_seq.mean() - 1.0
        metrics["chi2_seq"] = chi2_seq.detach().item()

    return metrics