`nemo_rl.models.automodel.router_replay`#

Rollout Routing Replay (R3) for the Automodel policy backend.

vLLM records the top-k experts it routed each token to during rollout (routed_experts, shape [B, S, layer slots, topk]). Small numerical differences between vLLM and the training forward flip near-tied routing decisions, and the flips compound across layers and positions until a token’s logprob can differ by tens of nats. Replaying the rollout selection through Automodel’s per-gate RouterReplay hooks makes the policy score each token through the experts that generated it. Only the discrete selection is replayed; routing weights are still computed from the live router, so gradients keep flowing into it.

Module Contents#

Classes#

ReplayGate

A policy MoE gate whose expert selection nemo-rl can replay.

Functions#

`validate_automodel_router_replay_config`	Reject setups where R3 cannot get routes or gate tokens would not align with them.
`require_automodel_moe_model`	Reject models that load through HF, whose MoE gates have no replay hooks.
`enable_routing_replay_in_automodel_kwargs`	Build every Automodel MoE gate with a `RouterReplay` handle.
`_decoder_layer_index`	Index `i` of the `layers.<i>` block that owns `module_name`, if any.
`_replaying_apply`	Swap in the replay target; tokens marked keep-live retain the gate’s own selection.
`router_replay_gates`	The policy’s MoE gates in layer order, each set up for nemo-rl replay.
`_payload_layers`	Which `routed_experts` layer slot each gate replays.
`microbatch_routed_experts`	The microbatch’s `routed_experts` (token-aligned with `input_ids`), padding unreplayed.
`_validated_targets`	Per-gate replay targets `[tokens, gates, topk]` and keep-live mask `[tokens, gates, 1]`.
`replay_routes`	Replay `routed_experts` ([B, S, layer slots, topk], aligned with input_ids) on `gates`.
`microbatch_replay_context`	Replay the microbatch’s rollout routes on `gates`; a no-op when `gates` is None.

API#

class nemo_rl.models.automodel.router_replay.ReplayGate#

A policy MoE gate whose expert selection nemo-rl can replay.

name: str#: None

layer: Optional[int]#: None

num_experts: Optional[int]#: None

handle: nemo_automodel.components.moe.router_replay.RouterReplay#: None

nemo_rl.models.automodel.router_replay.validate_automodel_router_replay_config( config: nemo_rl.models.policy.PolicyConfig, ) → None#

Reject setups where R3 cannot get routes or gate tokens would not align with them.

Raises:

ValueError – If generation is not vLLM (the only backend that returns routes), or policy.automodel_cfg is disabled.
NotImplementedError – If sequence packing, context parallelism or sequence parallelism is on; the gate would then see tokens that do not line up with the [batch, seq] routes.

nemo_rl.models.automodel.router_replay.require_automodel_moe_model( automodel_kwargs: dict[str, Any], model_config: Any, ) → None#

Reject models that load through HF, whose MoE gates have no replay hooks.

Call after force_hf is settled; otherwise HF rejects the injected moe_overrides with an opaque TypeError.

Raises:: ValueError – If force_hf is set or the architecture has no Automodel implementation.

nemo_rl.models.automodel.router_replay.enable_routing_replay_in_automodel_kwargs( automodel_kwargs: dict[str, Any], ) → None#: Build every Automodel MoE gate with a RouterReplay handle.

nemo_rl.models.automodel.router_replay._decoder_layer_index(module_name: str) → Optional[int]#: Index i of the layers.<i> block that owns module_name, if any.

nemo_rl.models.automodel.router_replay._replaying_apply( handle: nemo_automodel.components.moe.router_replay.RouterReplay, indices: torch.Tensor, ) → torch.Tensor#: Swap in the replay target; tokens marked keep-live retain the gate’s own selection.

nemo_rl.models.automodel.router_replay.router_replay_gates( model: torch.nn.Module, ) → list[nemo_rl.models.automodel.router_replay.ReplayGate]#

The policy’s MoE gates in layer order, each set up for nemo-rl replay.

Replay is installed on these instances only (Automodel’s own RouterReplay API is untouched). MTP gates are excluded: vLLM never routes them.

nemo_rl.models.automodel.router_replay._payload_layers( gates: list[nemo_rl.models.automodel.router_replay.ReplayGate], num_payload_layers: int, ) → list[int]#

Which routed_experts layer slot each gate replays.

vLLM may emit one slot per MoE layer, or one per decoder layer (hybrid models such as NemotronH also count Mamba/attention layers); the latter is indexed by decoder layer.

nemo_rl.models.automodel.router_replay.microbatch_routed_experts( data_dict: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], input_ids: torch.Tensor, ) → torch.Tensor#

The microbatch’s routed_experts (token-aligned with input_ids), padding unreplayed.

Right padding of routed_experts is 0 on the driver path ([0, ..., 0] rows would repeat one expert) and -1 from the data plane; positions past each sequence’s input_lengths get the missing-route sentinel so they keep the live selection.

Raises:: ValueError – If the microbatch has no routed_experts or they are not [batch, seq, layer slots, topk] aligned with input_ids.

nemo_rl.models.automodel.router_replay._validated_targets( gates: list[nemo_rl.models.automodel.router_replay.ReplayGate], routed_experts: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#

Per-gate replay targets [tokens, gates, topk] and keep-live mask [tokens, gates, 1].

Only all-sentinel rows mean “no route captured” and keep the live selection. Any other invalid row (partially negative, out of range, or repeating an expert) is corruption.

nemo_rl.models.automodel.router_replay.replay_routes( gates: list[nemo_rl.models.automodel.router_replay.ReplayGate], routed_experts: torch.Tensor, ) → Iterator[None]#

Replay routed_experts ([B, S, layer slots, topk], aligned with input_ids) on gates.

Wrap both the forward and the backward of a microbatch so activation-checkpoint recomputation replays the same selection.