[Sampling Replay] Bounded sampling-support replay

docs/advanced/index.md

@@ -7,7 +7,7 @@ This section covers the Miles features that the Core-features section of the
 homepage points at: low-precision training (FP8 / MXFP8 / NVFP4 / INT4 QAT),
 Rollout Routing Replay for MoE, fast weight updates over P2P RDMA,
 disaggregated RL rollout through an external service, fault tolerance,
-speculative decoding, and LoRA training and serving.
+speculative decoding, sampling-support replay, and LoRA training and serving.
 
 <CardGroup cols={2}>
 
@@ -38,6 +38,13 @@ speculative decoding, and LoRA training and serving.
 
   </Card>
 
+  <Card title="Sampling-Support Replay" icon="filter" href="/advanced/sampling-support-replay">
+
+    Preserve top-p and top-k rollout probabilities when the actor recomputes
+    token log probabilities for training.
+
+  </Card>
+
   <Card title="On-Policy Distillation" icon="graduation-cap" href="/advanced/on-policy-distillation">
 
     Train a student on its own rollouts while matching teacher token

docs/advanced/sampling-support-replay.md

@@ -0,0 +1,85 @@
+---
+title: Sampling-Support Replay
+description: Train against the same bounded sampling distribution used during rollout.
+---
+
+Top-p and top-k change the distribution that generates a rollout token. They
+remove tokens from the vocabulary and renormalize the remaining probability
+mass. Recomputing a token's log probability over the full vocabulary therefore
+does not reproduce its behavior-policy probability, even when the rollout and
+trainer weights are identical.
+
+Sampling-support replay preserves that distribution. For every generated token,
+SGLang returns the realized set of token IDs that remained after sampling
+filters. Miles transports this ragged support with the sample and applies it to
+the actor logits before the softmax:
+
+$$
+q_\theta(a_t \mid S_t) =
+\frac{\exp(z_\theta(a_t) / T)}
+{\sum_{j \in S_t} \exp(z_\theta(j) / T)}.
+$$
+
+The support `S_t` is fixed rollout data; `z_θ` comes from the actor being
+trained. This gives PPO/GRPO a denominator and numerator in the same sampled
+probability space while retaining gradients through the current actor logits.
+
+## Enable replay
+
+```bash
+ROLLOUT_ARGS+=(
+  --rollout-temperature 1.0
+  --rollout-top-p 0.95
+  --rollout-top-k 64
+)
+```
+
+Replay is enabled whenever `--rollout-top-p` is below `1` or
+`--rollout-top-k` is positive. A finite top-p run must also set a positive
+top-k so that the captured support is bounded. Miles asks SGLang to capture its
+native sampling support with `return_sampling_mask` and uses the returned
+support-normalized log probability as the rollout log probability.
+
+`--rollout-top-k` is the default for rollout requests, not a global upper bound
+on request-specific top-k values. SGLang owns support capacity through
+`--sglang-sampling-mask-max-tokens` (or the matching server-group override) and
+rejects requests whose realized support cannot be represented.
+
+## Correctness requirements
+
+Miles rejects configurations that it cannot replay faithfully:
+
+- Every resolved training request must have top-p, top-k, and temperature. Agentic
+  sessions fill omitted values from the sampling defaults registered when the
+  session is created.
+- Request temperature must match the registered training temperature
+  (`--rollout-temperature` for the standard rollout paths).
+- Frequency, presence, and repetition penalties, `logit_bias`, and custom logit
+  processors are not supported because the trainer does not replay those logit
+  transformations.
+- `--recompute-logprobs-via-prefill` is incompatible because that path does not
+  preserve the per-token support.
+- The selected router must preserve `return_sampling_mask`. The Miles router
+  forwards raw request bodies. The native SGLang router requires typed chat and
+  generate request schemas that declare the field.
+
+The SGLang backend must return one complete support and its normalized log
+probability for every sampled token. Miles validates that response contract but
+does not add separate restrictions for speculative decoding.
+
+Tool and environment tokens are recorded with singleton support. Evaluation
+requests do not capture this training-only metadata.
+
+## Current objective limitation
+
+Replay currently exposes only the support-normalized actor score to the loss.
+Reference KL and on-policy distillation also need a full-vocabulary actor score,
+so Miles rejects `--use-kl-loss`, nonzero `--kl-coef`, and `--use-opd` with
+replay until the loss interface carries both scores.
+
+## Monitor replay
+
+On the first update from identical weights,
+`train/train_rollout_logprob_abs_diff` should be near numerical tolerance.
+Afterward, interpret it with version-lag and clipping metrics because it also
+reflects policy updates and staleness.

docs/docs.json

@@ -220,6 +220,7 @@
        "advanced/low-precision",
        "advanced/int4-qat",
        "advanced/speculative-decoding",
+       "advanced/sampling-support-replay",
        "advanced/lora",
        "advanced/on-policy-distillation"
       ]

docs/user-guide/cli-reference.md

@@ -80,7 +80,8 @@ then push up until you OOM.
 | Flag | Default | What |
 |---|---|---|
 | `--rollout-temperature` | `1.0` | Sampling temperature. |
-| `--rollout-top-p` | `1.0` | Top-p truncation. |
+| `--rollout-top-p` | `1.0` | Top-p truncation. Values below `1` enable [sampling-support replay](/advanced/sampling-support-replay) and require a positive top-k. |
+| `--rollout-top-k` | `-1` | Top-k truncation. Positive values enable [sampling-support replay](/advanced/sampling-support-replay). |
 | `--rollout-max-response-len` | `–` | Max tokens per response. |
 | `--rollout-stop-token-ids` | model default | Stop token IDs. Override when generations don't stop. |
 | `--apply-chat-template` | off | Apply the tokenizer's chat template. |
@@ -186,8 +187,8 @@ Sections mirror the launch-script argument groups.
 |---|---|---|---|
 | `--rollout-max-response-len` | int | – | Max tokens per response. |
 | `--rollout-temperature` | float | `1.0` | Sampling temperature. |
-| `--rollout-top-p` | float | `1.0` | Top-p truncation. |
-| `--rollout-top-k` | int | `-1` | Top-k truncation (-1 disables). |
+| `--rollout-top-p` | float | `1.0` | Top-p truncation. Values below `1` require bounded [sampling-support replay](/advanced/sampling-support-replay). |
+| `--rollout-top-k` | int | `-1` | Top-k truncation (`-1` disables). Positive values enable [sampling-support replay](/advanced/sampling-support-replay). |
 | `--rollout-stop` | str+ | – | Stop strings. |
 | `--rollout-stop-token-ids` | int+ | – | Stop token IDs. |
 

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
@@ -369,6 +370,7 @@ def _compute_log_prob(
             active_model.eval()
         else:
             active_model = self.model
+        use_rollout_sampling_mask = model_tag == "actor" and self.args.use_sampling_support_replay
 
         try:
             forward_data_store = []
@@ -392,13 +394,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 +420,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 +509,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 self.args.use_sampling_support_replay
+                else ()
+            )
 
             for step_id in range(num_steps_per_rollout):
                 self.optimizer.zero_grad(set_to_none=True)
@@ -523,6 +536,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

@@ -430,6 +430,7 @@ def compute_log_prob(
                 rollout_id=rollout_id,
                 store_prefix=store_prefix,
                 fp32_output=False,
+                use_rollout_sampling_mask=store_prefix == "" and self.args.use_sampling_support_replay,
             )
 
     @with_logs

miles/backends/megatron_utils/model.py

@@ -29,6 +29,7 @@
 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
@@ -268,6 +269,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 +285,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 +308,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 +331,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 +350,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 +418,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 args.use_sampling_support_replay 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 +464,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,6 +19,7 @@
     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.types import RolloutBatch
 
@@ -111,6 +112,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 args.use_sampling_support_replay else None
 
     log_probs_and_entropy = get_log_probs_and_entropy(
         logits,
@@ -121,6 +123,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/rollout/generate_hub/multi_turn.py

@@ -51,7 +51,12 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     for _turn in range(args.generate_max_turns):
         # ----------------------- Call inference endpoint -------------------------
 
-        payload, halt_status = compute_request_payload(args, sample.tokens, input.sampling_params)
+        payload, halt_status = compute_request_payload(
+            args,
+            sample.tokens,
+            input.sampling_params,
+            evaluation=input.evaluation,
+        )
         if payload is None:
             sample.status = halt_status
             break

miles/rollout/generate_hub/single_turn.py

@@ -35,7 +35,11 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         input_ids = prompt_ids
 
     payload, halt_status = compute_request_payload(
-        args, input_ids=input_ids, sampling_params=sampling_params, multimodal_inputs=sample.multimodal_inputs
+        args,
+        input_ids=input_ids,
+        sampling_params=sampling_params,
+        multimodal_inputs=sample.multimodal_inputs,
+        evaluation=input.evaluation,
     )
     if payload is None:
         sample.status = halt_status