Why colocated MInf router replay records no routes when a reference model is loaded

Explainer for lauradang/RL #2, a PR into NVIDIA-NeMo/RL #4129. All permalinks pin the PR head b311ce9 and Megatron-LM bde6af2.

Before
The SingleController setup rejected router replay with Megatron generation (setup.py:1196, before), so this path could not run.
This PR
Turns it on, and requires every MInf model call to carry expert routes (tq_token_sink.py:694).
Status
Open — not fixed in this PR. Raised in a review comment on tq_token_sink.py:694.

Two pages on this PR, one question each — this one: the colocated bug.

Megatron keeps one list of routers per process. Every MoE router with replay turned on adds itself to RouterReplay.global_router_replay_instances when it is built (router_replay.py:112). Megatron Inference (MInf) finds routers only through this list: it tells each one to record (router_replay.py:61-64) and reads the recorded routes back from it (routing_metadata.py:75). When the KL penalty is above 0, the policy worker also builds a reference model, and that setup ends by emptying the whole list, the policy's routers included (setup.py:2666-2667). When generation shares the training GPUs (colocated) and has no separate inference model, MInf serves the training model itself (_gen_model), whose routers are no longer in the list.

The router list over time — a model with 3 MoE layers

RouterReplay.global_router_replay_instances — one list per process, shared by every model in it 1 policy model built router_replay.py:112 append(self) P1 P2 P3 length 3 the model MInf will serve 2 reference model built (KL > 0) megatron_policy_worker.py:742 P1 P2 P3 R1 R2 R3 length 6 R1…R3 never generate 3 finally: clear the whole list megatron/setup.py:2667 [ ] empty — P1…P3 removed too length 0 nothing adds P1…P3 back 4 MInf generates a token text_generation_controller.py:2944 RECORD for each router in [ ] → none records get_recorded_data() → routes = None every request, every step P1…P3 still route, unrecorded 5 route check fails the call tq_token_sink.py:694 routing_indices is None → ValueError rollout_reassembler.py:551 → group dropped no group is ever trained stalls at step 0, no crash policy router — the model MInf serves reference router — never generates red = the problem Training is not affected: each model keeps its own routers. Only MInf goes through the list.

The strict check in step 5 is correct — the PR does not want rollouts with missing routes. It is simply what turns the empty list into a 100% failure. Each failed call is marked failed (fail_call), the reassembler drops the group because no rollout has routes yet, and the SingleController never commits a dropped group (single_controller.py:1640-1655). The GPUs keep generating; the first training step never runs.

Which runs hit it — colocated only (a separate generation worker builds just one model):

MInf servesKL penaltylist when MInf runsresult
the training model> 0[ ]every call fails — this page
the training model0[P1, P2, P3]works
a separate inference model (reshard)> 0[I1, I2, I3]works, by luck: it is built after the clear (:823)
a separate inference model (reshard)0[P1, P2, P3, I1, I2, I3]likely fails: 6 layers of routes where 3 are expected (route_assembly.py:168); not run on a GPU

What would fix it — proposed, not in the PR

Nothing below is implemented. Rebuild the list from the model MInf actually serves, at the top of _initialize_inference_engine, before DynamicInferenceEngine(...) is built (megatron_worker.py:627). Same example: 3 MoE layers, colocated, KL 0.01.

# step 1: before the engine exists. list = [ ] (after the clear) if router_replay_enabled: served = self._gen_model() # the training model: routers P1, P2, P3 RouterReplay.global_router_replay_instances[:] = [ inst for inst, _ in _router_replay_instances_for_model(served) ] # list = [P1, P2, P3] (reshard + KL 0: [P1..P3, I1..I3] -> [I1, I2, I3]) if not RouterReplay.global_router_replay_instances: raise RuntimeError("router replay is on but the served model has no routers") # step 2: build the engine. It sizes the route buffer from the list: L = len(list) = 3 self.dynamic_inference_engine = DynamicInferenceEngine(...) # step 3: each step, RECORD reaches P1, P2, P3 -> routes for 3 layers -> the check at :694 passes

The order matters: the engine constructor captures CUDA graphs (dynamic_engine.py:481) and sizes its route buffer from the list (routing_metadata.py:44), so fixing the list later is too late. _router_replay_instances_for_model already exists (router_replay.py:159) and skips the extra multi-token-prediction (MTP) routers by default, which matches the trainer's layer count.

Two fixes that look right but are not: deleting the clear leaves R1…R3 in the list (length 6, not 3); saving and restoring the list around the reference build fixes the main case but not reshard with KL 0.

So what: a colocated MInf run with router replay and a KL penalty above 0 never trains — it generates forever at step 0 with no error that stops the job. The default grpo_math_1B.yaml (colocated on, KL 0.01) plus the MInf config the PR adds to docs/guides/router-replay.md hits it. No shipped recipe turns on router replay with MInf generation, so no recipe run shows it. Nothing here was run on a GPU; it follows from the code above.