Skip to main content
Router replay, also called rollout routing replay (R3), makes the trainer send each sampled token through the same experts the sampler chose for it when you train a mixture-of-experts (MoE) model. To use it, record your samples and attach each training token’s source to its Datum.

​
Record samples and attach provenance

To record routes and attach them to a training example, create the sampling client with record_stability_info=True, sample with include_routed_experts=True, and describe the datum’s tokens with with_provenance():
  • record_stability_info=True records every sample from the client. Each returned SampledSequence carries a sequence_id that refers to its record. The record and its routes stay on Baseten, and the datum carries only the references. Create the client from the training client, as in the snippet, so the run’s own sampler records the routes that its trainer replays.
  • include_routed_experts=True captures the expert routes with each record. It requires a client created with record_stability_info=True.
  • with_provenance() returns a copy of the datum with two span lists. The model_input spans describe the input tokens in order, and the loss_fn_inputs spans describe target_tokens in order. Each span names a sample’s sequence_id, whether its tokens come from the sample’s prompt (PromptProvenanceSpan) or its completion (SampledProvenanceSpan), the position of its first token in that part (offset, 0 by default), and the number of tokens (length).
With the Tinker compatibility package, import PromptProvenanceSpan and SampledProvenanceSpan from tinker.types.

​
Requirements and limitations

  • Router replay works in both availability modes on MoE models whose sampler captures routes. On other models, a sample() call with include_routed_experts=True fails.
  • Recording takes tokenized text prompts. A recording client rejects prompts that contain images.
  • dpo doesn’t support router replay. A forward_backward() call with loss_fn="dpo" fails when its datums carry provenance.

​
Routing mismatch between the sampler and trainer

Each MoE layer has a router that sends every token to a few of the layer’s experts. The sampler and the trainer both run the router, but with different kernels and numerics, so for some tokens they choose different experts. The trainer then computes those tokens’ log-probabilities and gradients through experts that didn’t generate them. These routing differences widen the KL divergence between the sampler and trainer even when both run the same weights. The RL losses read the gap as a change in the policy, and in MoE models the mismatch can grow until training becomes unstable. R3 removes this source of mismatch: the sampler records the experts it chose for each token, and the trainer reuses them in its forward and backward passes. The trainer still computes the router’s gating weights from its own logits, so the router keeps training.

​
Next steps