agentic RL top-p replay

miles/backends/megatron_utils/actor.py

@@ -781,7 +781,7 @@ def compute_log_prob(
         num_microbatches: list[int],
         store_prefix: str = "",
     ) -> dict[str, list[torch.Tensor]]:
-
+        use_rollout_sampling_masks = "rollout_sampling_masks" in data_iterator[0].rollout_data
         with timer(f"{store_prefix}log_probs"):
             return forward_only(
                 get_log_probs_and_entropy,
@@ -790,6 +790,7 @@ def compute_log_prob(
                 data_iterator,
                 num_microbatches,
                 store_prefix=store_prefix,
+                use_rollout_sampling_masks=use_rollout_sampling_masks,
             )
 
     def train(self, rollout_id: int, rollout_data_ref: Box) -> None:

miles/backends/megatron_utils/model.py

@@ -194,6 +194,7 @@ def forward_only(
     data_iterator: Sequence[DataIterator],
     num_microbatches: Sequence[int],
     store_prefix: str = "",
+    use_rollout_sampling_masks: bool = False,
 ) -> dict[str, list[torch.Tensor]]:
     """Run forward passes only and collect non-loss outputs (e.g., logprobs).
 
@@ -239,16 +240,18 @@ def forward_step(
         assert not return_schedule_plan, "forward_only step should never return schedule plan"
 
         # Get the batch.
+        keys = [
+            "tokens",
+            "loss_masks",
+            "multimodal_train_inputs",
+            "total_lengths",
+            "response_lengths",
+            "max_seq_lens",
+            "rollout_sampling_masks",
+        ]
         batch = get_batch(
             data_iterator,
-            [
-                "tokens",
-                "loss_masks",
-                "multimodal_train_inputs",
-                "total_lengths",
-                "response_lengths",
-                "max_seq_lens",
-            ],
+            keys,
             args.data_pad_size_multiplier,
             args.qkv_format,
             allgather_cp=args.allgather_cp,
@@ -268,15 +271,17 @@ def forward_step(
             **(batch["multimodal_train_inputs"] if batch["multimodal_train_inputs"] is not None else {}),
         )
 
-        return output_tensor, partial(
-            f,
-            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),
-        )
+        callback_kwargs = {
+            "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_masks:
+            callback_kwargs["rollout_sampling_masks"] = batch["rollout_sampling_masks"]
+        return output_tensor, partial(f, **callback_kwargs)
 
     # Turn on evaluation mode which disables dropout.
     for model_module in model:
@@ -381,23 +386,25 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
         """
 
         # Get the batch.
+        keys = [
+            "tokens",
+            "multimodal_train_inputs",
+            "packed_seq_params",
+            "total_lengths",
+            "response_lengths",
+            "loss_masks",
+            "log_probs",
+            "ref_log_probs",
+            "values",
+            "advantages",
+            "returns",
+            "rollout_log_probs",
+            "max_seq_lens",
+            "rollout_sampling_masks",
+        ]
         batch = get_batch(
             data_iterator,
-            [
-                "tokens",
-                "multimodal_train_inputs",
-                "packed_seq_params",
-                "total_lengths",
-                "response_lengths",
-                "loss_masks",
-                "log_probs",
-                "ref_log_probs",
-                "values",
-                "advantages",
-                "returns",
-                "rollout_log_probs",
-                "max_seq_lens",
-            ],
+            keys,
             args.data_pad_size_multiplier,
             args.qkv_format,
             allgather_cp=args.allgather_cp,

miles/backends/sglang_utils/sglang_engine.py

@@ -680,6 +680,9 @@ def _compute_server_args(
         for key in unused_keys:
             kwargs.pop(key)
 
+    if args.num_rollout != 0 and args.rollout_top_p < 1 and kwargs.get("disaggregation_mode", "null") != "null":
+        raise ValueError("sampling-mask replay is incompatible with prefill/decode disaggregation")
+
     return kwargs, external_engine_need_check_fields
 
 

miles/backends/training_utils/loss.py

@@ -21,7 +21,7 @@
     get_reinforce_plus_plus_baseline_advantages,
     get_reinforce_plus_plus_returns,
 )
-from miles.utils.types import RolloutBatch
+from miles.utils.types import RolloutBatch, RolloutSamplingMask
 
 from .cp_utils import (
     _allgather_cp_redistribute,
@@ -129,8 +129,8 @@ def get_responses(
     total_lengths: list[int],
     response_lengths: list[int],
     max_seq_lens: list[int] | None = None,
-) -> Iterator[tuple[torch.Tensor, torch.Tensor]]:
-    """Yield response-aligned `(logits_chunk, tokens_chunk)` pairs per sample.
+) -> Iterator[tuple[torch.Tensor, torch.Tensor, list[tuple[int, int]]]]:
+    """Yield response logits, targets, and response-row spans per sample.
 
     After squeezing batch dimension and applying temperature scaling, this
     function extracts the logits and tokens corresponding to response segments
@@ -149,9 +149,10 @@ def get_responses(
         response_lengths: Response segment lengths per sample.
 
     Yields:
-        Tuple of `(logits_chunk, tokens_chunk)` where `logits_chunk` is shape
+        Tuple of `(logits_chunk, tokens_chunk, response_spans)` where `logits_chunk` is shape
         `[R, V]` (policy) or `[R, 1]` (value) and `tokens_chunk` is shape `[R]`
         (1D int64), both aligned to response tokens for one sample.
+        `response_spans` contains ordered half-open row ranges in the full response.
     """
     parallel_state = get_parallel_state()
     qkv_format = args.qkv_format
@@ -183,7 +184,8 @@ def get_responses(
                 end += total_length
                 start = end - response_length
             logits_chunk = logits[start - 1 : end - 1]
-            tokens_chunk = tokens[-response_length:]
+            tokens_chunk = tokens[total_length - response_length : total_length]
+            response_spans = [(0, response_length)] if response_length else []
         elif args.allgather_cp:
             # DSA: global concat then contiguous CP split. Each rank owns logits for
             # global positions [chunk_start, chunk_end).
@@ -203,16 +205,23 @@ def get_responses(
             if e <= s:
                 logits_chunk = logits[0:0]
                 tokens_chunk = tokens[0:0]
+                response_spans = []
             else:
                 logits_chunk = logits[s - chunk_start : e - chunk_start]
                 tokens_chunk = tokens[(s + 1) - seq_start : (e + 1) - seq_start]
+                response_spans = [(s - logit_global_start, e - logit_global_start)]
             assert logits_chunk.size(0) == tokens_chunk.size(0), f"{logits_chunk.size(0)} vs {tokens_chunk.size(0)}"
         else:
             # TODO: this is super ugly... do better abstraction.
             chunk_size, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp(
                 total_length, response_length, qkv_format, max_seq_len
             )
 
+            prompt_length = total_length - response_length
+            response_spans = [
+                (start - prompt_length, stop - prompt_length) for start, stop in tokens_offset if start < stop
+            ]
+
             logits_0, logits_1 = logits[end : end + chunk_size], logits[end + chunk_size : end + 2 * chunk_size]
             end += 2 * chunk_size
 
@@ -230,7 +239,36 @@ def get_responses(
 
         seq_start += total_length
 
-        yield logits_chunk, tokens_chunk
+        yield logits_chunk, tokens_chunk, response_spans
+
+
+def _build_tp_sampling_mask(
+    logits: torch.Tensor,
+    sampling_mask: RolloutSamplingMask,
+    response_spans: list[tuple[int, int]],
+    tp_rank: int,
+) -> torch.Tensor:
+    mask = torch.zeros(logits.shape, dtype=torch.bool, device=logits.device)
+    if not response_spans:
+        return mask
+
+    offsets = sampling_mask.offsets
+    sizes = torch.cat([offsets[start + 1 : stop + 1] - offsets[start:stop] for start, stop in response_spans])
+    parts = [sampling_mask.token_ids[offsets[start].item() : offsets[stop].item()] for start, stop in response_spans]
+    token_ids = parts[0] if len(parts) == 1 else torch.cat(parts)
+
+    vocab_size = logits.size(-1)
+    vocab_start = tp_rank * vocab_size
+    owned_entries = ((token_ids >= vocab_start) & (token_ids < vocab_start + vocab_size)).nonzero(as_tuple=True)[0]
+    if owned_entries.numel() == 0:
+        return mask
+
+    rows = torch.repeat_interleave(sizes, output_size=token_ids.numel())[owned_entries]
+    cols = token_ids[owned_entries].long() - vocab_start
+    rows = rows.to(logits.device, non_blocking=True)
+    cols = cols.to(logits.device, non_blocking=True)
+    mask[rows, cols] = True
+    return mask
 
 
 def get_log_probs_and_entropy(
@@ -243,6 +281,7 @@ def get_log_probs_and_entropy(
     with_entropy: bool = False,
     non_loss_data: bool = True,
     max_seq_lens: list[int] | None = None,
+    rollout_sampling_masks: list[RolloutSamplingMask] | None = None,
 ) -> dict[str, list[torch.Tensor]]:
     """Compute per-token log-probabilities (and optionally entropy) on responses.
 
@@ -260,6 +299,7 @@ def get_log_probs_and_entropy(
         response_lengths: Response segment lengths per sample.
         with_entropy: If True, include "entropy" key in result.
         non_loss_data: Unused; kept for API compatibility.
+        rollout_sampling_masks: Fixed rollout support per full response; entropy remains unrestricted.
 
     Returns:
         Dict with key "log_probs" mapping to a list of `[R]` tensors per
@@ -268,26 +308,35 @@ def get_log_probs_and_entropy(
     """
     parallel_state = get_parallel_state()
     assert non_loss_data
+
     log_probs_list = []
     entropy_list = []
-    for logits_chunk, tokens_chunk in get_responses(
-        logits,
-        args=args,
-        unconcat_tokens=unconcat_tokens,
-        total_lengths=total_lengths,
-        response_lengths=response_lengths,
-        max_seq_lens=max_seq_lens,
+    for i, (logits_chunk, tokens_chunk, response_spans) in enumerate(
+        get_responses(
+            logits,
+            args=args,
+            unconcat_tokens=unconcat_tokens,
+            total_lengths=total_lengths,
+            response_lengths=response_lengths,
+            max_seq_lens=max_seq_lens,
+        )
     ):
+        sampling_mask = None
+        if rollout_sampling_masks is not None:
+            sampling_mask = _build_tp_sampling_mask(
+                logits_chunk, rollout_sampling_masks[i], response_spans, parallel_state.tp.rank
+            )
         log_prob, entropy = calculate_log_probs_and_entropy(
             logits_chunk,
             tokens_chunk,
             parallel_state.tp.group,
             with_entropy=with_entropy,
             chunk_size=args.log_probs_chunk_size,
             true_on_policy=args.true_on_policy_mode,
+            sampling_mask=sampling_mask,
         )
 
-        log_probs_list.append(log_prob.squeeze(-1))
+        log_probs_list.append(log_prob.reshape(-1))
         entropy_list.append(entropy)
 
     res = {
@@ -341,7 +390,7 @@ def get_values(
         per sample.
     """
     value_list = []
-    for logits_chunk, _ in get_responses(
+    for logits_chunk, _, _ in get_responses(
         logits,
         args=args,
         unconcat_tokens=unconcat_tokens,
@@ -649,6 +698,7 @@ def policy_loss_function(
         response_lengths=response_lengths,
         with_entropy=True,
         max_seq_lens=max_seq_lens,
+        rollout_sampling_masks=batch.get("rollout_sampling_masks"),
     )
 
     log_probs = log_probs_and_entropy["log_probs"]

miles/ray/rollout.py

@@ -824,6 +824,9 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
         if samples[0].rollout_log_probs is not None:
             train_data["rollout_log_probs"] = [sample.rollout_log_probs for sample in samples]
 
+        if self.args.rollout_top_p < 1:
+            train_data["rollout_sampling_masks"] = [sample.rollout_sampling_mask for sample in samples]
+
         if samples[0].rollout_routed_experts is not None:
             train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples]
 
@@ -953,6 +956,7 @@ def _stat(xs):
                 "round_number",
                 "sample_indices",
                 "rollout_log_probs",
+                "rollout_sampling_masks",
                 "rollout_routed_experts",
                 "prompt",
                 "teacher_log_probs",

miles/rollout/generate_hub/agentic_tool_call.py

@@ -49,7 +49,8 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         "agentic_tool_call.generate requires session_server_ip/session_server_port. "
         "Pass --use-session-server to start the session server."
     )
-    tracer = await OpenAIEndpointTracer.create(input.args)
+    capture = not input.evaluation and input.sampling_params["top_p"] < 1
+    tracer = await OpenAIEndpointTracer.create(input.args, capture_sampling_mask=capture)
 
     custom_agent_function: Callable = load_function(input.args.custom_agent_function_path)
     assert (
@@ -114,6 +115,8 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         sample.status = Sample.Status.ABORTED
         return GenerateFnOutput(samples=sample)
 
+    if capture and merged_sample.rollout_sampling_mask is None:
+        raise ValueError("a replay trajectory must carry a sampling mask")
     sample = apply_merged_session_sample(input.args, input.sample, merged_sample)
     sample.metadata.update(agent_metadata or {})
     sample.metadata.update(session_metadata)