[Sampling Replay] Bounded sampling-support replay

[nanjiangwill]

Summary

Add end-to-end bounded sampling-support replay for top-p/top-k rollout training.

SGLang samples each token after temperature and support filters. A full-vocabulary actor softmax is therefore not the same policy as the bounded rollout distribution. Sampling-support replay records the realized support S_t and defines both behavior and actor scores on that fixed support:

• SGLang returns log mu(a_t | s_t, S_t).
• Miles computes log pi_theta(a_t | s_t, S_t) by masking and renormalizing current actor logits.
• S_t remains fixed during optimization; Miles does not recompute the current actor top-p boundary.

Sampling replay series

• [Sampling Replay] Add sampling-support log-prob primitives #2200 adds TP/CP-aware masked actor log-probability primitives.
• [Sampling Replay] Represent and transport rollout sampling support #2595 adds compact sampling-support representation and transport.
• [Sampling Replay] Bounded sampling-support replay #3354 activates replay end to end across rollout capture, session propagation, trainer scoring, validation, E2E, and documentation.

Runtime support uses SGLang v0.5.20 from #3321 and native-router chat forwarding from radixark/sgl-router-for-miles#20.

End-to-end behavior

• Replay activates when rollout top-p is below 1 or rollout top-k is positive.
• SGLang v0.5.20 supplies the realized support and its support-normalized selected-token log probability.
• Standard, legacy, multi-turn, tool-call, and session v1/v2 paths preserve the same contract.
• Both the Miles router and a compatible native SGLang router preserve the capture request.
• Environment and tool tokens receive singleton support.
• Megatron and FSDP apply support only to actor scores. Reference, teacher, value, evaluation, and entropy semantics remain full-vocabulary or unchanged.
• Evaluation disables training-only support capture.

Ownership and fail-closed boundaries

SGLang owns sampling-mask capacity, overflow, custom logit processors, and speculative-decoding compatibility. Miles does not duplicate those policies and does not add a speculative-decoding gate.

Miles rejects cases its trainer cannot currently replay faithfully:

• finite top-p without a positive top-k;
• request temperature different from rollout temperature;
• frequency, presence, or repetition penalties and logit bias;
• prefill log-probability recomputation;
• incomplete or inconsistent response metadata.

Router compatibility is a transport boundary, not part of the replay algorithm. The Miles router forwards raw request bodies. The native SGLang router uses typed request schemas, so its chat and generate schemas must declare return_sampling_mask. sgl-router-for-miles#20 adds the missing chat field and verifies forwarding to a worker.

Request-specific positive top-k may differ from the launch default. SGLang enforces the physical support limit.

Current KL and distillation limitation

The actor still produces full-vocabulary logits, so sampling replay is not mathematically incompatible with KL or distillation. The current loss interface exposes one actor score per token. Replay uses that slot for the support-normalized policy score, while KL and distillation require a separate full-policy actor score. Until both are carried independently, validation rejects nonzero KL coefficient, KL loss, and on-policy distillation with replay.

Validation

• Focused CPU coverage exercises trainer math, direct rollout, sessions, transport, FSDP, and argument validation.
• Repository pre-commit and all scheduled CPU CI shards pass on the current head.
• The included 2xH200 Qwen3-0.6B FSDP E2E exercises native support capture and replay across five GRPO updates under true-on-policy mode. --ci-test checks rollout/trainer log-probability equality on every rollout, with log-probability-difference and KL metric gates.
• The E2E passes on the current head with the native router.

Acknowledgements

Thank you to Zilin Zhu for the original top-p replay implementation in THUDM/slime#2102.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[guapisolo]

good. Seems like the rebase include args resolve logic on main now

[guapisolo]

[resolved] cc cmt

Follow-up on the sampling-support replay e2e, now that args.use_sampling_support_replay is derived once (thanks for folding that in). A true-on-policy run turns the log-prob check into an exact one, and it passes on the current head.

[P3] The current e2e only catches large normalization errors. tests/e2e/megatron/test_qwen3_0_6b_sampling_support_replay.py runs one rollout. Its log-prob check is --ci-test's abs_tol=0.03 on mean log_probs vs rollout_log_probs, and the register_ci_gate verdict is informational. I reproduced a regression where the actor scores over the full vocabulary instead of the replayed support. On this data it moves the mean log-prob by about 0.037, just above the tolerance, so a milder top-p or a support that is wrong on only some positions would keep the job green.

Proposal. Megatron still rejects --true-on-policy-mode (set_default_megatron_args raises NotImplementedError), so the exact check has to run on FSDP. With true-on-policy on, --ci-test asserts log_probs == rollout_log_probs on every rollout. The test below keeps Qwen3-0.6B, top-p 0.8 / top-k 32, 2 GPUs and 5 rollouts. It uses gsm8k with 1024 response tokens: with dapo at 128 tokens every response truncates, reward and advantage are 0, and the 5 steps barely move the weights.

Results on d7e28bdd92, H200, SGLang from the radixark/miles:dev image (5a8da8cc3f):

Run	Rollouts	Response tokens	train_rollout_logprob_abs_diff	Result
Proposed test, unchanged	5	~57k	0.0 at every step	pass, ~4 min, raw reward 0.31-1.0
Same config	10	~116k	0.0 at every step	pass
dapo, 128 tokens	5	~2.6k	0.0 at every step	pass
dapo, 128 tokens, FSDP forward-only mask disabled	2	-	-	fails at rollout 0: trainer -0.2095 vs rollout -0.1730

Caveat on "exactly 0". It holds through rounding, not by construction. SGLang's replayed log-prob is log(p_sel / sum_S p) in fp32 (Sampler._build_sampling_mask_output), the trainer's is a masked log_softmax in bf16, and both are compared in bf16. Positions with a single-token support give exactly 0 on both sides; the rest can differ by one bf16 ulp when the two fp32 values straddle a rounding boundary. Replaying both formulas offline on the same logits (122k sampled tokens, SGLang's own top_k_top_p_min_p_sampling_from_probs_torch, 72% single-token supports) gave 0 flips and an expected rate of about 4.6e-6 per sampled token, i.e. about 0.26 flips per run of this size.

• The --ci-test assertion compares bf16-reduced per-rollout means, which a single one-ulp flip almost never moves, so I don't expect the proposed test to flake.
• The per-token train_rollout_logprob_abs_diff can show a tiny nonzero value in roughly a quarter of runs, so it should stay an informational gate.
• Making it exact by construction needs SGLang's on-policy path to compute the replayed log-prob with the trainer's kernel (masked log_softmax); that is a separate SGLang change.

I'd keep the Megatron test as a smoke run until Megatron supports true-on-policy, since it is the only e2e over the Megatron replay path.

Optional naming nit: use_rollout_sampling_mask on forward_only and the FSDP local is the per-call decision (actor scoring with replay on), while the run switch is args.use_sampling_support_replay. A name like apply_rollout_sampling_mask would keep it from reading as a second switch.

tests/e2e/fsdp/test_qwen3_0_6b_fsdp_sampling_support_replay_true_on_policy.py

import os

from tests.ci.ci_register import register_cuda_ci
from tests.ci.metric_history import register_ci_gate

import miles.utils.external_utils.command_utils as U

register_cuda_ci(
    est_time=600,
    suite="stage-c-2-gpu-h200",
    labels=["fsdp", "sglang", "replay"],
    hardware=["hopper"],
)
register_ci_gate(metric_key="train/train_rollout_logprob_abs_diff")

MODEL_NAME = "Qwen3-0.6B"
NUM_GPUS = 2


def prepare() -> None:
    U.exec_command_cpu("mkdir -p /root/models /root/datasets")
    U.exec_command_cpu(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
    U.hf_download_dataset("zhuzilin/gsm8k")


def execute() -> None:
    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} "

    # Bounded top-p/top-k enables sampling-support replay. Under true-on-policy the actor's
    # support-normalized log-probs must equal the rollout's exactly; --ci-test asserts this on
    # every rollout step.
    rollout_args = (
        "--prompt-data /root/datasets/gsm8k/train.parquet "
        "--input-key messages "
        "--label-key label "
        "--apply-chat-template "
        "--rm-type math "
        "--num-rollout 5 "
        "--rollout-batch-size 4 "
        "--n-samples-per-prompt 4 "
        "--rollout-max-response-len 1024 "
        "--rollout-temperature 1 "
        "--rollout-top-p 0.8 "
        "--rollout-top-k 32 "
        "--global-batch-size 16 "
    )

    optimizer_args = (
        "--optimizer adam "
        "--lr 1e-6 "
        "--lr-decay-style constant "
        "--weight-decay 0.1 "
        "--adam-beta1 0.9 "
        "--adam-beta2 0.98 "
    )

    algorithm_args = "--advantage-estimator grpo " "--kl-loss-coef 0.00 " "--kl-coef 0.00 " "--entropy-coef 0.00 "

    fsdp_args = "--train-backend fsdp " "--update-weight-buffer-size 536870912 "

    sglang_args = (
        "--rollout-num-gpus-per-engine 1 " "--sglang-mem-fraction-static 0.4 " "--sglang-decode-log-interval 1000 "
    )

    true_on_policy_args = (
        "--sglang-enable-deterministic-inference "
        "--sglang-true-on-policy-contract qwen3_dense_true_on_policy_v1 "
        "--sglang-attention-backend fa3 "
        "--attn-implementation flash_attention_3 "
        "--deterministic-mode "
        "--true-on-policy-mode "
    )

    train_args = (
        ckpt_args
        + rollout_args
        + optimizer_args
        + algorithm_args
        + fsdp_args
        + sglang_args
        + true_on_policy_args
        + U.get_default_wandb_args(__file__)
        + "--ci-test --actor-num-nodes 1 --actor-num-gpus-per-node 2 --colocate "
    )

    U.execute_train(
        train_args=train_args,
        num_gpus_per_node=NUM_GPUS,
        megatron_model_type=None,
        extra_env_vars={
            "NCCL_ALGO": "allreduce:tree",
            "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
            "CUBLAS_WORKSPACE_CONFIG": ":4096:8",
            "CUDA_DEVICE_MAX_CONNECTIONS": "1",
        },
    )


if __name__ == "__main__":
    prepare()
    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
        os.environ.pop(proxy_var, None)
    execute()

[guapisolo]

@claude review always

[github-actions]

✅ tests/e2e/fsdp/test_qwen3_0_6b_fsdp_sampling_support_replay_true_on_policy.py passed on stage-c-2-gpu-h200 in 1h13m42s — workflow run

[claude]

I reviewed this PR and didn't find any bugs. Because it's a large, complex change to RL training correctness (rollout capture, session propagation, and dual-backend actor scoring across FSDP and Megatron), a human look would still be worthwhile even though this automated pass found nothing and a codeowner has already approved.

What was reviewed: the updated diff after the commits addressing the earlier inline feedback (bounded top-p/top-k contract validation, router forwarding fix, doc clarifications, e2e naming); the sampling-support capture/threading path through session state, config, and request-args; and the actor-side masked log-prob reconstruction in both backends. Checked whether replay-enabling flags are validated against router mode, and whether a custom Megatron loss calling get_rollout_sampling_masks would crash — both were investigated and ruled out as not blocking for this change.

Extended reasoning...

This PR adds a new sampling-support replay feature spanning rollout capture (SGLang mask extraction), session/wire serialization, and actor loss computation in both FSDP and Megatron backends across 46 files; no auth/crypto/injection surface, but it does touch core training-correctness logic owned by multiple CODEOWNERS. Since the prior review, 14 commits addressed a codeowner's (guapisolo) earlier inline objections and that same codeowner has since approved, and this run found no new bugs, so approval risk is lower; however the size and centrality of the change to RL correctness still argue for a human pass rather than auto-approval.