feat: ProRLv2 - add seq-mask-tis truncated importance sampling type

[hijkzzz]

Add a new IS filtering mechanism "seq-mask-tis" that masks entire sequences based on the geometric mean of per-token IS ratios, while keeping non-truncated token-level IS weights for gradient correction. Also adds shared is_filter_drop_frac metric for both icepop and seq-mask-tis modes, and documents the new option in prorlv2.md.

Summary by CodeRabbit

• New Features

  • Introduced seq-mask-tis, a sequence-level alternative to ICE-POP for importance sampling that operates at the sequence level rather than per-token.
  • Added is_filter_drop_frac metric for monitoring filtering behavior.

• Documentation

  • Comprehensive guide covering seq-mask-tis with rationale, configuration, and comparison with ICE-POP.
  • Updated terminology and expanded examples with new feature references.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR introduces seq-mask-tis, a sequence-level importance sampling approach as an alternative to ICE-POP. Documentation is updated with terminology clarification and detailed comparison between methods. Implementation adds seq-mask-tis support to the loss function with validation rules and metric tracking.

Changes

Cohort / File(s)	Summary
Documentation Updates docs/guides/prorlv2.md	Updated terminology (CE-POP to ICE-POP), introduced seq-mask-tis feature with detailed comparison table, configuration examples, and rationale documentation.
Loss Function Implementation nemo_rl/algorithms/loss_functions.py	Extended truncated importance sampling to support three types (tis, icepop, seq-mask-tis). Added sequence-level masking logic using geometric mean computation, validation checks preventing seq-mask-tis with per-sequence IS, and filter-drop metric tracking for icepop and seq-mask-tis paths.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• feat: support truncated importance sampling #1348: Both PRs modify nemo_rl/algorithms/loss_functions.py to add and validate truncated importance-sampling behavior with token/sequence-level clipping logic.

Suggested labels

CI:L1, r0.4.0

Suggested reviewers

• terrykong

🚥 Pre-merge checks | ✅ 2 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Test Results For Major Changes	⚠️ Warning	PR description lacks test results documentation despite major feature affecting numerics; critical review comments highlight unresolved validation and metric aggregation issues.	Add test results to PR description showing seq-mask-tis produces correct loss values and metrics without convergence regressions; address unresolved review comments on validation and metric aggregation.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main change: adding a new seq-mask-tis truncated importance sampling type to ProRLv2, which aligns with the primary objective of the PR.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In `@nemo_rl/algorithms/loss_functions.py`:
- Around line 171-183: In the __init__ (or initializer) validation block where
truncated_importance_sampling_type is checked (the assert on
self.truncated_importance_sampling_type and the subsequent check for
"seq-mask-tis"), add a guard that when self.truncated_importance_sampling_type
is "icepop" or "seq-mask-tis", then self.truncated_importance_sampling_ratio_min
is not None; raise an assertion with a clear message referencing
truncated_importance_sampling_ratio_min and the allowed sampling types to
prevent later TypeError when comparing tensors to None in methods that use this
attribute.
- Line 579: The metric is_filter_drop_frac is a fraction and must not be summed
across packed sequences; instead implement the same special aggregation used for
min/max metrics in SequencePackingLossWrapper: when encountering key
"is_filter_drop_frac" in metrics_accum, accumulate a weighted sum (e.g.,
metrics_accum_sum["is_filter_drop_frac"] += val * weight) and a corresponding
total weight (metrics_accum_weight["is_filter_drop_frac"] += weight) where
weight is the number of examples/sequence length for that packed segment, then
compute the final fraction as metrics_accum_sum / metrics_accum_weight before
reporting; update the code paths that currently do metrics_accum[k] += val to
detect "is_filter_drop_frac" and use this weighted accumulation and final
division (follow the same pattern used for the existing min/max handling in
SequencePackingLossWrapper and use the same helper variables/keys to keep
aggregation consistent).

🧹 Nitpick comments (3)

docs/guides/prorlv2.md (2)

158-184: Documentation for seq-mask-tis looks good overall.

The new section clearly explains the mechanism, provides a comparison table with ICE-POP, and includes configuration snippets.

One observation: the reference bounds in the table (min=0.002, max=0.003) represent a very narrow band for the geometric-mean IS ratio. Consider adding a brief note explaining why these bounds are so far from 1.0 (and so tight), or pointing users to the referenced blog for tuning guidance. Users unfamiliar with the method may assume bounds closer to 1.0 are expected.

---

182-182: Clarify metric semantics difference.

Line 182 notes that is_filter_drop_frac represents "fraction of tokens (ICE-POP) or sequences (seq-mask-tis)" filtered out. Since the same metric name measures different granularities depending on the mode, this could cause confusion when comparing runs. Consider noting this caveat more prominently or suggesting users check which mode is active when interpreting this metric.

nemo_rl/algorithms/loss_functions.py (1)

418-424: Replace EN DASH (–) with HYPHEN-MINUS (-) in comments.

Ruff (RUF003) flags ambiguous Unicode EN DASH characters in these comment lines. While visually similar, they can cause issues with some tools and are not idiomatic in source code.

Proposed fix

-        # "tis"          – clamp IS weights to [0, max]
-        # "icepop"       – zero out tokens whose IS weight ∉ [min, max]   (ref bounds: 0.5–5)
-        # "seq-mask-tis" – zero out entire sequences whose geometric-mean
-        #                  IS ratio ∉ [min, max]; retained sequences keep
-        #                  raw (non-truncated) token-level IS weights      (ref bounds: 0.002–0.003)
+        # "tis"          - clamp IS weights to [0, max]
+        # "icepop"       - zero out tokens whose IS weight not in [min, max]   (ref bounds: 0.5-5)
+        # "seq-mask-tis" - zero out entire sequences whose geometric-mean
+        #                  IS ratio not in [min, max]; retained sequences keep
+        #                  raw (non-truncated) token-level IS weights      (ref bounds: 0.002-0.003)

The same applies to the comment block at lines 48-51:

-    #   "tis"          – clamp IS weights to max
-    #   "icepop"       – zero out tokens with IS weight outside [min, max]
-    #   "seq-mask-tis" – zero out sequences by geometric-mean IS ratio, non-truncated token IS correction
+    #   "tis"          - clamp IS weights to max
+    #   "icepop"       - zero out tokens with IS weight outside [min, max]
+    #   "seq-mask-tis" - zero out sequences by geometric-mean IS ratio, non-truncated token IS correction

[hijkzzz]

@terrykong @yfw the l1/l0 tests passed.

[terrykong]

thanks for contributing @hijkzzz. to help others who come across this PR, is there any experimental results you can share showing how this helps stability in your experiments?

In the Demystifying blog they show "Token-Level MIS < Sequence-Level MIS" for stability. Any reason why you implemented the token level one instead of the sequence level one first?

Also, i think it would be good to have unit tests for all these importance sampling techniques for correctness

[hijkzzz]

thanks for contributing @hijkzzz. to help others who come across this PR, is there any experimental results you can share showing how this helps stability in your experiments?

In the Demystifying blog they show "Token-Level MIS < Sequence-Level MIS" for stability. Any reason why you implemented the token level one instead of the sequence level one first?

Also, i think it would be good to have unit tests for all these importance sampling techniques for correctness

We found seq-based filtering to be more stable for MoE models.

[hijkzzz]

@terrykong all tests passed please merge it

[hijkzzz]

@terrykong fixed please merge it

[terrykong]

lgtm. @hijkzzz thanks for making nemo-rl better!