Routing Replay for MoEs

[qgallouedec]

Feature request

RecentRL approaches for training MoE models increasingly rely on Routing Replay, as described in the following papers:

• https://huggingface.co/papers/2507.18071
• https://huggingface.co/papers/2510.11370
• https://huggingface.co/papers/2512.01374

Without going into the training details, Routing Replay requires the ability to override the router during the forward pass, that is, to force the model to use a predefined set of router logits rather than computing new ones. This enables deterministic reproduction of expert selection.

AFAICT, Transformers currently does not expose a way to override router logits or manually control expert selection at inference/training time.

I imagine something along the following lines (minimal example):

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B-Instruct-2507", device_map="auto", dtype="auto")

input_ids = torch.tensor([[1, 2, 3, 4]], device="cuda")

# Standard forward pass, retrieving router logits
outputs = model(input_ids, output_router_logits=True)

# Forward pass with router logits injected (enabling Routing Replay)
model(input_ids, router_logits=outputs.router_logits)

Alternative

If we decide not to implement this feature, it would be nice to provide an example showing how to patch a MoE to enable this.

Motivation

See above.

Your contribution

I think I can do it.

[edbeeching]

@qgallouedec Excuse my naivity, but why not just provide the expert indices directly rather than the logits?

[qgallouedec]

For Qwen3-MOE (at least) the model requires both the expert indices and their corresponding weights for computing the hidden states

transformers/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py

Line 248 in fccb049

current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]

The indices and weights are derived directly from the gating logits:

transformers/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py

Line 267 in fccb049

router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) # (seq_len, top_k)

So having directly the logits is probably the easiest.