Replay sampler expert routes and token supports in Miles training

[kailash109]

Sampling currently drops SGLang's expert routes and token sampling supports, so clients cannot replay either during training. This adds opt-in capture and a sampler-to-forward_backward handoff for the Miles LoRA backend.

Changes

• lilo.sample_with_replay(...) returns per-sequence replay metadata; sequence.replay.training_inputs(...) turns it into Tinker tensors. JSON and protobuf use the existing loss_fn_inputs transport.
• Forward expert IDs into Miles' native router replay queues, including CP/SP packing and cleanup on success or failure. Requests with and without router replay are scheduled separately.
• Replay captured top-k/top-p/min-p support and temperature in the Tinker loss. Use support-normalized behavior log probabilities for the importance ratio; preserve ordinary sequence.logprobs separately.
• Keep replay aligned through multi-node sequence padding and DP filler datums. Validate routes against the resolved model before GPU dispatch, rejecting invalid expert IDs, mixed padding, duplicate selections, and incompatible dimensions.
• Setup and sample/train examples are in docs/sampler-replay.md. Native Tinker router handoff and multi-LoRA replay lifecycle are being upstreamed in Miles #3706.

Router capture requires ROLLOUT_RETURN_ROUTED_EXPERTS = True on the pool definition and --use-rollout-routing-replay in Miles' extra_args. Sampling-mask capture requires no extra server flag. Both request flags default to false.

Validation

• Current revision: 594 CPU tests passed, 1 skipped, on Python 3.12 with Torch 2.10 CPU; Ruff and diff checks passed.
• Real Tinker 0.24.1 SDK over local HTTP: capture replay metadata, construct a datum, and submit it through forward_backward to a recording trainer. Tests cover JSON/protobuf round trips, alignment, CP/DP padding, masked gradients, queue cleanup, and rejecting invalid routes without dispatching or poisoning the runtime.
• Two 10-step DAPO-math MoE runs, Matplotlib plots, and route audit: replay reduced token-weighted mean absolute sampler–trainer logprob error by 39.2%. These runs used ebeeac6 plus the compatibility/configuration fixes and instrumentation detailed in that comment. This revision includes the Miles import-path fix; the review changes have not been rerun on GPUs.

The GPU experiment used attention LoRA with frozen experts, TP4/EP4/CP1 and activation recomputation. It does not establish reward convergence, capture overhead, CP>1, or multi-node GPU correctness. The helper handles a single continuation, not automatic multi-turn stitching. Full-parameter replay is not implemented here.

[kailash109]

Live MoE validation completed: two matched 10-step DAPO-math runs, using Qwen3-30B-A3B-Instruct-2507. Router replay reduced token-weighted mean absolute sampler–trainer logprob error by 39.2%. These measurements used PR revision ebeeac6 plus the local compatibility/configuration changes and audit instrumentation described below.

Training mode	Successful optimizer steps	Mean absolute logprob difference (nats)	Mean answer reward
Replay off	10	0.017256	0.28125
Replay on	10	0.010492	0.28125

Each batch was also evaluated with replay both on and off on identical tokens and weights, before its optimizer update. Replay improved all 20 paired comparisons, separating the routing effect from differences in the two training trajectories:

Expert-index correctness:

• All 320 rollout sequences passed shape, token-coverage, expert-range and distinct-top-8 checks. Normalized int32 SHA-256 hashes matched the received tensors on all four trainer ranks.
• All 1,284 enabled rank-level passes consumed all 48 layer streams in forward, and in backward recomputation when gradients were requested. Disabled passes left replay queues/cursors empty. No missing hashes or unconsumed input batches.
• Negative control on one sample: correct routes gave 0.008589 MAE; native routes 0.017059; deliberately shifted valid expert IDs 4.660218. This was forward-only and did not update weights, confirming the supplied indices affect the computation.

Changes needed to run this validation:

Change	Reason
miles_runtime/runtime.py: import miles.utils.lora.arguments instead of the removed miles.utils.multi_lora, with a regression test	Current pinned Miles failed to start with the old import.
New isolated Qwen3-30B-A3B recipe: 4×H200 trainer, TP4/EP4/CP1, attention LoRA rank 32; 2×H200 SGLang sampler, TP2; capture/replay flags and activation recomputation	Exercise a real MoE forward and backward path. Set --rotary-base 10000000 to match this checkpoint; the inherited 1,000,000 value failed model validation.
LILO_REPLAY_AUDIT=1 logging in miles_runtime/replay.py	Record input hashes and per-layer forward/backward queue cursors. The existing route handoff, packing and replay queue logic did not need changes.
DAPO-math controller, fixture, analyzer and history/checkpoint verifier	Matched initial model/optimizer states; same-token paired checks; negative control; per-step checkpoint/resume; Matplotlib outputs. Controller setup also needed Jinja2, explicit apply_chat_template(return_dict=False), immutable-datum reconstruction and unique checkpoint bookkeeping.
Experiment-only app registry and deployment environment forwarding	Restrict the isolated deployment to this recipe. Missing environment forwarding caused a callback to look up the default app; frontend reconciliation still served the run. The forwarding fix was deployed after measurements.

These implementation changes remain in the validation worktree; this comment does not claim the unmodified PR head passed the canary. The runtime compatibility fix, audit instrumentation and regression-test patch is available for review. The separate results branch contains documentation/assets; it does not change this PR's source branch.

Infrastructure and checks: 35 sampler queue-full HTTP 503 retries, all 35 recovered (queue limit 8, sixteen concurrent sample requests). No observed CUDA/NCCL/OOM failures during completed training. Sampled GPU hardware throttle flags and uncorrected ECC counters were clear. Setup included a CPU controller preemption before any update, controller compatibility fixes, and checkpointed restarts; an occupied-slot resume error was resolved by unloading retired sessions. Both persisted W&B histories and final checkpoints were independently verified at optimizer step 10. 106 targeted replay/Miles tests passed, plus Ruff and diff checks. Experiment GPUs have been released.

Scope: this is an attention-LoRA replay smoke test on DAPO math prompts, using group-normalized binary answer reward and asymmetric PPO clipping, with frozen experts. It is not a full DAPO reproduction, reward-convergence result or throughput benchmark. Both modes requested route metadata, while only the enabled mode supplied it to its gradient pass. At the 4,096-token response cap, about 71% of completions truncated. Equal-reward groups produced zero advantages, so 6/10 baseline and 7/10 replay updates had nonzero gradients; all ten optimizer steps succeeded. CP>1 and multi-node trainers remain untested.

Matplotlib answer-reward plot

Per-step metrics and route audit · Matplotlib source · Experiment notes

[micahtyong]

35 sampler queue-full HTTP 503 retries, all 35 recovered (queue limit 8, sixteen concurrent sample requests)

Not a blocker for this PR, but did this lead to any scheduling slow-down? Do we need to increase the queue limit?

This adds opt-in capture and a sampler-to-forward_backward handoff for the Miles LoRA backend.

Looking at the results, this seems like a net-positive change. Are there any downsides to having this an always-on capture that I'm not seeing?

Or are we just adopting whatever the default is from Miles (which I'm also okay with)

[micahtyong]

Any reason these hooks need to live in Lilo? Should we eventually upstream this into Miles?

[kailash109]

Yes. The native Tinker route handoff and multi-LoRA replay lifecycle are being upstreamed in radixark/miles#3706 (currently draft). Miles already owns the router selection and CP/SP packing; these hooks fill the missing Tinker-to-training handoff. We can remove the corresponding Lilo adapter once that PR lands and the deployed Miles revision includes it. The sampler response extension and sampling-support/temperature adapter are separate from that upstream router-replay PR.

[kailash109]

No way i didnt realize my agent responded to your review commments

The main problem seems to be that in their sample and train payloads, they seem to adhere to the Tinker spec very closely (ie. for the sample path, make sure that the request follows tinker_types.SampleRequest, which doesn't have any flexibliity for bringing in replay indices/top p mask)

I made the upstream pr just in case but i doubt it'll get merged, for now seems that we can't rely on their willingness to make api changes

[kailash109]

--use-rollout-routing-replay can be passed through to miles via extra_args in config, but the miles tinker build_train_data doesn't pass this into the train batch https://github.com/radixark/miles/blob/cc76e23915b2132ecfff65b97331fd83b02637e6/miles/tinker/runtime.py#L27-L43

the patch we have now is a messy fix -- todo try and upstream this into miles multi lora

[kailash109]

Upstream draft is now open: radixark/miles#3706. It covers the native Tinker training-data handoff and multi-LoRA router replay lifecycle. The local adapter remains necessary until that is merged and included in our deployed Miles revision; sampler flag/response transport and sampling-mask replay are outside that upstream PR's scope.

[kailash109]

Addressed the review in 4eb5d34: moved imports to module scope in the replay/runtime/cluster code and touched tests, replaced fixed-field getattr calls with direct access, and reject invalid expert routes before distributed execution. Also included the Miles lora.arguments import-path fix identified during the earlier GPU validation.

On the two questions above:

• Queue-full retries: the 35 recovered 503s incurred retry/backoff time. We did not run a matched test without queue pressure, so the measurements do not quantify their effect on end-to-end step time. The setup used sixteen concurrent sample requests and a queue limit of eight. Increasing the queue can absorb bursts and reduce rejection/retry overhead, but it does not increase GPU service capacity and can increase queued-request latency. I would tune concurrency and the queue together using queue occupancy, retry count, and sampling latency rather than raise the limit based only on the fact that all retries recovered. No queue setting changes are included here.
• Why capture stays opt-in: for the validated Qwen3-30B-A3B model, expert routes alone cost 48 layers × 8 experts × 4 bytes = 1,536 bytes per captured input position. At 65,536 positions that is 96 MiB raw / 128 MiB base64, before JSON and the training submission. Capturing also needs server buffers, copies, serialization and client memory. Sampling-support capture can be larger still when top-p retains many vocabulary entries. Both validation modes requested route metadata, so the 39.2% logprob-error improvement measures using replay, not the overhead of enabling capture. Opt-in is therefore intentional; it is not just inheriting a default. Replay also changes which experts/support the trainer uses, so clients should choose it explicitly.

The upstream route handoff/lifecycle work is tracked in Miles #3706.

After resolving the overlap with current main in 0338a76, validation is 594 passed, 1 skipped on Python 3.12 with CPU Torch 2.10; Ruff and diff checks passed. The earlier 10-step GPU comparison and Matplotlib plots remain the GPU evidence; these review fixes have not been rerun on GPUs.