[Bug] seq-mask-tis skips token-level truncation and can produce non-finite PPO updates

Author: ai-yangCreated Aug 8, 2026Updated Aug 8, 2026

Summary

The seq-mask-tis branch says that its sequence-level geometric mean is used only for filtering and that correction coefficients still use token-level truncated importance sampling (TIS). However, it exponentiates old_log_probs - rollout_log_probs without applying the configured low/high clamp.

This makes the mode differ from its documented semantics even for finite ratios. With a sufficiently stale rollout, exp() overflows; an accepted sequence then has an infinite loss/gradient, while a rejected sequence can evaluate 0 * inf and produce NaN.

Severity: medium-high. The mode is opt-in, but once selected a single extreme off-policy sequence can poison an expensive PPO optimizer step and subsequent training.

Affected code

At commit bc71bb19464aca306b33080b2d2bb45d154e2f49, openrlhf/models/loss.py contains:

python

elif self.vllm_is_correction_type == "seq-mask-tis":
    # seq-mask-tis: use sequence-level geometric mean only for filtering,
    # correction coefficients still use TIS (token-level clamp)
    seq_log_ratio = masked_mean(rollout_log_ratio, action_mask, dim=-1)
    seq_is = torch.exp(seq_log_ratio)
    seq_mask = (seq_is >= low_threshold) & (seq_is <= high_threshold)
    vllm_is = torch.exp(rollout_log_ratio).detach()
    loss = seq_mask.unsqueeze(-1) * vllm_is * loss

The regular tis branch immediately below correctly applies .clamp(min=low_threshold, max=high_threshold).

Reproduction

CPU-only, using PyTorch 2.13.0. All log-probabilities below are non-positive.

python

import math
import torch
from openrlhf.models.loss import PolicyLoss

loss_fn = PolicyLoss(
    policy_loss_type="ppo",
    enable_vllm_is_correction=True,
    vllm_is_truncated_threshold=[0.5, 5.0],
    vllm_is_correction_type="seq-mask-tis",
)

def run(rollout_log_probs):
    old_log_probs = torch.tensor([[0.0, -1000.0]])
    log_probs = old_log_probs.clone().requires_grad_(True)
    loss, *_ = loss_fn(
        log_probs,
        old_log_probs,
        torch.ones_like(log_probs),
        action_mask=torch.ones_like(log_probs),
        rollout_log_probs=torch.tensor([rollout_log_probs]),
    )
    loss.backward()
    print(loss, log_probs.grad, torch.isfinite(loss), torch.isfinite(log_probs.grad).all())

# Finite token ratios [10, 0.1] have geometric mean 1 and pass the sequence filter.
d = math.log(10.0)
old_log_probs = torch.tensor([[-0.1, -0.1 - d]])
log_probs = old_log_probs.clone().requires_grad_(True)
moderate_loss, *_ = loss_fn(
    log_probs,
    old_log_probs,
    torch.ones_like(log_probs),
    action_mask=torch.ones_like(log_probs),
    rollout_log_probs=torch.tensor([[-0.1 - d, -0.1]]),
)
print(moderate_loss)  # -5.05 on main; TIS-clamped result should be -2.75

# Token ratios are exp([1000, -1000]); geometric mean is 1, so the sequence is accepted.
run([-1000.0, 0.0])

# Token ratios are exp([1000, -996]); geometric mean is exp(2) > 5, so it is rejected.
run([-1000.0, -4.0])

Actual result on main:

moderate_loss=-5.05 (configured [0.5, 5.0] bounds were bypassed)
loss=-inf, gradient contains -inf, finite=False
loss=nan, gradient contains nan, finite=False

The rejected case becomes False * inf, which is NaN rather than zero.

Expected behavior

After sequence-level filtering, token coefficients should be clamped to [low_threshold, high_threshold], exactly as the comment and the tis implementation specify. With the thresholds above, the accepted case should use weights [5.0, 0.5] and produce a finite loss of -2.75; the rejected case should produce finite zero loss/gradients.

Suggested fix

Apply the same clamp as the tis branch before detaching:

python

vllm_is = torch.exp(rollout_log_ratio).clamp(min=low_threshold, max=high_threshold).detach()

Add regression coverage for both an accepted extreme sequence and a rejected extreme sequence, asserting finite loss and gradients.

Reachability

This is reachable through the normal CLI path:

--algo.advantage.is_correction_enable
--algo.advantage.is_correction_type seq-mask-tis
  -> vLLM rollout_log_probs
  -> Experience.rollout_log_probs
  -> PolicyLoss(..., policy_loss_type="ppo")

Duplicate check

I searched open and closed issues and all PRs for seq-mask-tis, token clamp, TIS overflow, and importance sampling NaN; no existing report or fix covers this branch.

PR #1293 is adjacent but distinct: it bounds the exponent in the gspo policy-ratio branch (loss.py:170-178). This bug is in PPO's seq-mask-tis correction branch (loss.py:207-214) and still reproduces with #1293 applied. PR #1240 addresses invalid values at masked positions; this reproduction uses an all-ones action mask.

Source: OpenRLHF/OpenRLHF

View original on GitHub View discussion on GitHub