From 1c70f71648f95b91de80aa07b05aad8bedc1e457 Mon Sep 17 00:00:00 2001
From: Lu Ken <bluewish.ken.lu@live.cn>
Date: Tue, 14 Jul 2026 17:51:58 +0800
Subject: [PATCH] rollout: fix partial rollout loss mask alignment

---
 slime/rollout/sglang_rollout.py         |  2 +-
 slime/utils/types.py                    |  5 ++
 tests/test_partial_rollout_loss_mask.py | 90 +++++++++++++++++++++++++
 3 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_partial_rollout_loss_mask.py

diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py
index 9dcf5fce85..f4455d86db 100644
--- a/slime/rollout/sglang_rollout.py
+++ b/slime/rollout/sglang_rollout.py
@@ -229,7 +229,7 @@ async def generate_and_rm(
 ) -> Sample | list[Sample]:
     # mask previous off-policy generation for partial rollout
     if args.partial_rollout and args.mask_offpolicy_in_partial_rollout and sample.response_length > 0:
-        sample.loss_mask = [0] * sample.response_length
+        sample.mask_response_tokens(0)
 
     # For samples with existing response, check if they're complete
     if sample.status == Sample.Status.COMPLETED or sample.status == Sample.Status.TRUNCATED:
diff --git a/slime/utils/types.py b/slime/utils/types.py
index df4b10062c..fb65b66406 100644
--- a/slime/utils/types.py
+++ b/slime/utils/types.py
@@ -250,6 +250,11 @@ def get_reward_value(self, args) -> float:
     def effective_response_length(self):
         return sum(self.loss_mask) if self.loss_mask is not None else self.response_length
 
+    def mask_response_tokens(self, value: int = 0) -> None:
+        """Set the response-side loss mask while preserving length invariants."""
+        self.loss_mask = [int(value)] * self.response_length
+        self._validate_response_metadata_lengths()
+
     def append_response_tokens(
         self,
         args=None,
diff --git a/tests/test_partial_rollout_loss_mask.py b/tests/test_partial_rollout_loss_mask.py
new file mode 100644
index 0000000000..60dcdc22bb
--- /dev/null
+++ b/tests/test_partial_rollout_loss_mask.py
@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from argparse import Namespace
+
+import pytest
+
+from slime.utils.types import Sample
+
+NUM_GPUS = 0
+
+
+def _make_args() -> Namespace:
+    return Namespace(sglang_speculative_algorithm=False)
+
+
+def _resume_sample(
+    sample: Sample,
+    *,
+    partial_rollout: bool,
+    tokens: list[int],
+    log_probs: list[float],
+    mask_offpolicy_in_partial_rollout: bool = True,
+) -> Sample:
+    if partial_rollout and mask_offpolicy_in_partial_rollout and sample.response_length > 0:
+        sample.mask_response_tokens(0)
+
+    sample.append_response_tokens(
+        _make_args(),
+        tokens=tokens,
+        log_probs=log_probs,
+        trainable=True,
+        meta_info={"finish_reason": {"type": "stop"}},
+        text=" new",
+    )
+    sample.reward = 1.0
+    return sample
+
+
+@pytest.mark.unit
+def test_partial_resume_masks_old_tokens_and_appends_trainable_loss_mask():
+    sample = Sample(
+        tokens=[1, 2, 10, 11],
+        response="old",
+        response_length=2,
+        loss_mask=[1, 1],
+        rollout_log_probs=[-0.7, -0.8],
+        status=Sample.Status.ABORTED,
+    )
+
+    _resume_sample(
+        sample,
+        partial_rollout=True,
+        tokens=[30, 31, 32],
+        log_probs=[-0.3, -0.4, -0.5],
+    )
+
+    assert sample.response_length == 5
+    assert sample.loss_mask == [0, 0, 1, 1, 1]
+    assert len(sample.loss_mask) == sample.response_length
+    assert sample.rollout_log_probs == [-0.7, -0.8, -0.3, -0.4, -0.5]
+    assert sample.status is Sample.Status.COMPLETED
+
+
+@pytest.mark.unit
+def test_non_partial_resume_preserves_existing_loss_mask():
+    sample = Sample(
+        tokens=[1, 2, 10, 11],
+        response="old",
+        response_length=2,
+        loss_mask=[1, 0],
+        rollout_log_probs=[-0.7, -0.8],
+        status=Sample.Status.ABORTED,
+    )
+
+    _resume_sample(sample, partial_rollout=False, tokens=[30], log_probs=[-0.3])
+
+    assert sample.response_length == 3
+    assert sample.loss_mask == [1, 0, 1]
+    assert len(sample.loss_mask) == sample.response_length
+
+
+@pytest.mark.unit
+def test_fresh_full_rollout_still_gets_all_trainable_loss_mask():
+    sample = Sample(tokens=[1, 2], status=Sample.Status.PENDING)
+
+    _resume_sample(sample, partial_rollout=False, tokens=[30, 31], log_probs=[-0.3, -0.4])
+
+    assert sample.response_length == 2
+    assert sample.loss_mask == [1, 1]
+    assert len(sample.loss_mask) == sample.response_length