[Sampling Replay 3/8] Replay sampling support in actor scoring

miles/backends/fsdp_utils/actor.py

@@ -21,6 +21,7 @@
 )
 from miles.backends.training_utils.loss import compute_advantages_and_returns, get_log_probs_and_entropy, loss_function
 from miles.backends.training_utils.parallel import get_parallel_state, set_parallel_state
+from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks
 from miles.ray.train_actor import TrainRayActor
 from miles.utils import async_utils, train_dump_utils, train_metric_utils
 from miles.utils.context_utils import with_defer
@@ -32,6 +33,7 @@
 from miles.utils.processing_utils import load_processor, load_tokenizer
 from miles.utils.profile_utils import TrainProfiler
 from miles.utils.ray_utils import Box
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.timer import Timer, inverse_timer, timer
 from miles.utils.tracking_utils.tracking import init_tracking
 
@@ -369,6 +371,7 @@ def _compute_log_prob(
             active_model.eval()
         else:
             active_model = self.model
+        use_rollout_sampling_mask = model_tag == "actor" and sampling_support_replay_enabled(self.args)
 
         try:
             forward_data_store = []
@@ -392,13 +395,18 @@ def _compute_log_prob(
                             "response_lengths",
                             "max_seq_lens",
                         ]
+                        if use_rollout_sampling_mask:
+                            forward_only_keys.extend(["rollout_sampling_mask_ids", "rollout_sampling_mask_offsets"])
                         batch = get_batch(
                             data_iterator,
                             forward_only_keys,
                             self.args.data_pad_size_multiplier,
                             self.args.qkv_format,
                             get_position_ids=True,
                         )
+                        rollout_sampling_mask = None
+                        if use_rollout_sampling_mask:
+                            rollout_sampling_mask = get_rollout_sampling_masks(batch)
 
                         model_args = self._get_model_inputs_args(batch)
                         # keep logits in native bf16 (chunks upcast to fp32 downstream); avoids a full-vocab fp32 tensor (~5GB)
@@ -413,6 +421,7 @@ def _compute_log_prob(
                             response_lengths=batch["response_lengths"],
                             with_entropy=(store_prefix == ""),
                             max_seq_lens=batch.get("max_seq_lens", None),
+                            rollout_sampling_mask=rollout_sampling_mask,
                         )
 
                         batch_result = {
@@ -501,6 +510,11 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
         with routing_replay.stage(routing_replay.REPLAY_BACKWARD), timer("actor_train"):
             data_iterator.reset()
             num_steps_per_rollout = len(num_microbatches)
+            sampling_mask_keys = (
+                ("rollout_sampling_mask_ids", "rollout_sampling_mask_offsets")
+                if sampling_support_replay_enabled(self.args)
+                else ()
+            )
 
             for step_id in range(num_steps_per_rollout):
                 self.optimizer.zero_grad(set_to_none=True)
@@ -523,6 +537,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
                             "returns",
                             "ref_log_probs",
                             "rollout_log_probs",
+                            *sampling_mask_keys,
                         ],
                         self.args.data_pad_size_multiplier,
                         self.args.qkv_format,

miles/backends/megatron_utils/actor.py

@@ -33,6 +33,7 @@
 from miles.utils.ray_utils import Box
 from miles.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups
 from miles.utils.replay_base import all_replay_managers, routing_replay_manager
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor
 from miles.utils.timer import Timer, inverse_timer, timer
 from miles.utils.tracking_utils.structured_log import with_logs
@@ -430,6 +431,7 @@ def compute_log_prob(
                 rollout_id=rollout_id,
                 store_prefix=store_prefix,
                 fp32_output=False,
+                use_rollout_sampling_mask=store_prefix == "" and sampling_support_replay_enabled(self.args),
             )
 
     @with_logs

miles/backends/megatron_utils/model.py

@@ -29,11 +29,13 @@
 from miles.backends.megatron_utils.ft.types import TrainStepOutcome
 from miles.backends.megatron_utils.local_weight_checksum import dump_local_weight_checksums
 from miles.backends.megatron_utils.optimizer_state_reset import reset_optimizer_states
+from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks
 from miles.utils.audit_utils.witness.allocator import WitnessInfo
 from miles.utils.audit_utils.witness.module import witness_dump_and_clear_stale
 from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase
 from miles.utils.memory_utils import clear_memory
 from miles.utils.multi_lora import is_multi_lora_enabled
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor
 from miles.utils.tracking_utils.structured_log import log_structured
 
@@ -268,6 +270,7 @@ def forward_only(
     rollout_id: int,
     store_prefix: str = "",
     fp32_output: bool = True,
+    use_rollout_sampling_mask: bool = False,
 ) -> dict[str, list[torch.Tensor]]:
     """Run forward passes only and collect non-loss outputs (e.g., logprobs).
 
@@ -283,6 +286,8 @@ def forward_only(
         rollout_id: Rollout identifier (selects the per-rollout dump subdirectory).
         store_prefix: Prefix to prepend to stored output keys.
         fp32_output: Whether Megatron should upcast the complete model output to FP32.
+        use_rollout_sampling_mask: Whether to score over each rollout token's
+            captured sampling support.
 
     Returns:
         Aggregated outputs keyed by ``store_prefix + key``.
@@ -304,17 +309,20 @@ def forward_step(
     ) -> tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
         """Return the model output and its batch-bound loss callback."""
 
+        forward_only_keys = [
+            "tokens",
+            "loss_masks",
+            "multimodal_train_inputs",
+            "total_lengths",
+            "response_lengths",
+            "max_seq_lens",
+            "witness_ids",
+        ]
+        if use_rollout_sampling_mask:
+            forward_only_keys.extend(["rollout_sampling_mask_ids", "rollout_sampling_mask_offsets"])
         batch = get_batch(
             data_iterator,
-            [
-                "tokens",
-                "loss_masks",
-                "multimodal_train_inputs",
-                "total_lengths",
-                "response_lengths",
-                "max_seq_lens",
-                "witness_ids",
-            ],
+            forward_only_keys,
             args.data_pad_size_multiplier,
             args.qkv_format,
             allgather_cp=args.allgather_cp,
@@ -324,6 +332,7 @@ def forward_step(
         packed_seq_params = get_packed_seq_params(batch, args)
         total_lengths = batch["total_lengths"]
         response_lengths = batch["response_lengths"]
+        rollout_sampling_mask = get_rollout_sampling_masks(batch) if use_rollout_sampling_mask else None
 
         if "adapter_token_counts" in batch:
             from megatron.bridge.peft.multi_lora_layers import set_tokens_per_adapter_slot
@@ -342,15 +351,18 @@ def forward_step(
             fp32_output=fp32_output,
         )
 
-        return output_tensor, partial(
-            f,
+        callback_kwargs = dict(
             args=args,
             unconcat_tokens=unconcat_tokens,
             total_lengths=total_lengths,
             response_lengths=response_lengths,
             with_entropy=args.use_rollout_entropy,
             max_seq_lens=batch.get("max_seq_lens", None),
         )
+        if use_rollout_sampling_mask:
+            callback_kwargs["rollout_sampling_mask"] = rollout_sampling_mask
+
+        return output_tensor, partial(f, **callback_kwargs)
 
     # Turn on evaluation mode which disables dropout.
     for model_module in model:
@@ -407,6 +419,12 @@ def run_forward_backward_pass(
 ):
     """One pipeline forward/backward pass over the microbatches; no optimizer interaction."""
 
+    sampling_mask_keys = (
+        ("rollout_sampling_mask_ids", "rollout_sampling_mask_offsets")
+        if sampling_support_replay_enabled(args) and args.loss_type == "policy_loss"
+        else ()
+    )
+
     @dumper_phase_util.wrap_forward_step
     def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False) -> tuple[
         torch.Tensor,
@@ -447,6 +465,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
                 "loss_weights",
                 "target_tokens",
                 "sample_indices",
+                *sampling_mask_keys,
             ],
             args.data_pad_size_multiplier,
             args.qkv_format,

miles/backends/training_utils/loss_hub/losses.py

@@ -19,7 +19,9 @@
     compute_policy_loss,
 )
 from miles.backends.training_utils.parallel import get_parallel_state
+from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks
 from miles.utils.function_registry import load_function
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.types import RolloutBatch
 
 
@@ -111,6 +113,7 @@ def policy_loss_function(
     total_lengths = batch["total_lengths"]
     max_seq_lens = batch.get("max_seq_lens", None)
     calculate_entropy = args.entropy_coef != 0 or args.observe_training_entropy
+    rollout_sampling_mask = get_rollout_sampling_masks(batch) if sampling_support_replay_enabled(args) else None
 
     log_probs_and_entropy = get_log_probs_and_entropy(
         logits,
@@ -121,6 +124,7 @@ def policy_loss_function(
         with_entropy=calculate_entropy,
         entropy_requires_grad=args.entropy_coef != 0,
         max_seq_lens=max_seq_lens,
+        rollout_sampling_mask=rollout_sampling_mask,
     )
 
     log_probs = log_probs_and_entropy["log_probs"]

miles/backends/training_utils/sampling_mask.py

@@ -1,10 +1,26 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 
 import torch
 
 from miles.utils.sampling_mask import RolloutSamplingMask
 
 
+def get_rollout_sampling_masks(batch: Mapping[str, object]) -> list[RolloutSamplingMask]:
+    """Reconstruct the per-sample masks carried by the training-data wire."""
+    ids_batch = batch.get("rollout_sampling_mask_ids")
+    offsets_batch = batch.get("rollout_sampling_mask_offsets")
+    if ids_batch is None or offsets_batch is None:
+        raise ValueError("sampling-support actor scoring requires both sampling-mask wire fields")
+    if not isinstance(ids_batch, Sequence) or not isinstance(offsets_batch, Sequence):
+        raise TypeError("rollout sampling-mask ids and offsets must be sequences with one entry per sample")
+    if len(ids_batch) != len(offsets_batch):
+        raise ValueError(f"sampling-mask ids batch size {len(ids_batch)} != offsets batch size {len(offsets_batch)}")
+    return [
+        RolloutSamplingMask(ids=torch.as_tensor(ids), offsets=torch.as_tensor(offsets))
+        for ids, offsets in zip(ids_batch, offsets_batch, strict=True)
+    ]
+
+
 def build_local_sampling_mask(
     logits: torch.Tensor,
     sampling_mask: RolloutSamplingMask,

miles/utils/sampling_mask.py

@@ -5,6 +5,11 @@
 import torch
 
 
+def sampling_support_replay_enabled(args) -> bool:
+    """Whether rollout sampling uses a bounded support that must be replayed."""
+    return float(getattr(args, "rollout_top_p", 1.0)) < 1.0 or int(getattr(args, "rollout_top_k", -1)) > 0
+
+
 @dataclass(frozen=True, eq=False)
 class RolloutSamplingMask:
     """One sample's sampling mask: for each response position, the token ids

tests/fast/backends/megatron_utils/test_model.py

@@ -2,6 +2,7 @@
 from argparse import Namespace
 from dataclasses import dataclass, field
 from typing import Any
+from unittest.mock import Mock
 
 import pytest
 
@@ -120,3 +121,57 @@ def test_train_one_step_emits_the_train_tag_in_its_structured_event(
             )
 
         assert "train op=train_step rollout=7 step=3 attempt=2 outcome=NORMAL valid_step=true" in caplog.messages
+
+
+def test_forward_only_omits_sampling_mask_for_callbacks_that_do_not_replay_sampling_support(monkeypatch):
+    from miles.backends.megatron_utils import model as model_module
+
+    def forward_backward(**kwargs):
+        output, callback = kwargs["forward_step_func"](kwargs["data_iterator"][0], kwargs["model"][0])
+        return [callback(output)]
+
+    callback_kwargs = {}
+
+    def critic_callback(_output, **kwargs):
+        callback_kwargs.update(kwargs)
+        return {}
+
+    args = Namespace(
+        data_pad_size_multiplier=1,
+        qkv_format="thd",
+        allgather_cp=False,
+        enable_witness=False,
+        use_rollout_entropy=False,
+        custom_megatron_before_log_prob_hook_path=None,
+        seq_length=8,
+        micro_batch_size=1,
+    )
+    batch = {
+        "tokens": None,
+        "input_loss_masks": None,
+        "multimodal_train_inputs": None,
+        "unconcat_tokens": [],
+        "total_lengths": [],
+        "response_lengths": [],
+        "max_seq_lens": None,
+    }
+    dumper = Mock()
+    dumper.wrap_forward_step.side_effect = lambda fn: fn
+
+    monkeypatch.setattr(model_module, "DumperMegatronUtil", lambda *_args, **_kwargs: dumper)
+    monkeypatch.setattr(model_module, "get_model_config", lambda _model: Mock(timers=None))
+    monkeypatch.setattr(model_module, "get_batch", lambda *_args, **_kwargs: batch)
+    monkeypatch.setattr(model_module, "get_packed_seq_params", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(model_module, "get_forward_backward_func", lambda: forward_backward)
+    monkeypatch.setattr(model_module, "mpu", FakeMpu())
+
+    model_module.forward_only(
+        critic_callback,
+        args,
+        [Mock()],
+        [Mock()],
+        [1],
+        rollout_id=0,
+    )
+
+    assert "rollout_sampling_mask" not in callback_kwargs

tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py

@@ -136,18 +136,25 @@ def test_train_keeps_model_resident(actor_module, monkeypatch):
     worker.sleep.assert_not_called()
 
 
-def test_compute_log_prob_keeps_logits_in_model_precision(actor_module, monkeypatch):
+@pytest.mark.parametrize(
+    ("store_prefix", "expected"),
+    [("", True), ("ref_", False), ("teacher_", False)],
+)
+def test_compute_log_prob_replays_sampling_support_only_for_actor_scores(
+    actor_module, monkeypatch, store_prefix, expected
+):
     worker = object.__new__(actor_module.MegatronTrainRayActor)
-    worker.args = Namespace()
+    worker.args = Namespace(rollout_top_p=0.95, rollout_top_k=32)
     worker.model = [object()]
     forward_only = Mock(return_value={"log_probs": []})
     monkeypatch.setattr(actor_module, "forward_only", forward_only)
     monkeypatch.setattr(actor_module, "timer", lambda _name: nullcontext())
 
-    result = worker.compute_log_prob([], [], rollout_id=3)
+    result = worker.compute_log_prob([], [], rollout_id=3, store_prefix=store_prefix)
 
     assert result == {"log_probs": []}
     assert forward_only.call_args.kwargs["fp32_output"] is False
+    assert forward_only.call_args.kwargs["use_rollout_sampling_mask"] is expected
 
 
 def test_save_model_does_not_manage_lifecycle(actor_module, monkeypatch):