When a call resumes from a saved prefix, which expert routes does it train with?

NVIDIA-NeMo/RL #4612 · stacked on #4583 · permalinks pin head 2548c679. Review comments are in a pending review; their links work once it is published.

Before
Setup refused to run prefix recovery and router replay together (the check, at the base). MoE runs with router replay had to restart a cut call from scratch.
This PR
Saves each call's expert routes with every prefix it saves. On resume, it keeps the saved routes for the old tokens and takes the rest from the resumed call.
Status
The design is sound — implemented in this PR, with open review asks. The mix matches what an uninterrupted run produces with default settings. The review asks the PR to write that down, and to fix one crash.

Three pages on this PR — this one: what the feature does, compared with a run that never restarted.

Background. In an MoE model each token goes to a few experts. Router replay saves vLLM's choice for every position as one route row. Row t is the experts used when token t was the input (trainer). Prefix recovery saves a running call's tokens at each checkpoint, so after a restart it resumes instead of starting over. Dotted words have a short explanation: hover or tap them.

At each cut, freeze_for_checkpoint stores the route rows for the new tokens next to them in the saved row. On resume, _generation_cut_route_extras checks and joins those pieces. vLLM then re-reads the saved tokens with the new weights, so it computes new routes for them too. When the call ends, _merge_generation_cut_routed_experts keeps the saved rows for the old tokens and the new rows for everything after them. From there the row goes through the same final save as any call.

Same call, three ways

position input token 012345 101112131415 promptgenerated before the cutafter Uninterrupted run weights updated at the cut, KV cache kept (the default) 1 2 3 4 5 – Resumed call's own routes vLLM re-reads 10…14 with the new weights 90 91 92 93 5 – What this PR trains with saved rows for 0…3, resumed rows from 4 on 1 2 3 4 5 – same rows = uninterrupted run How lane 3 is built cut 1 [1,2,3,ph] cut 2 [4] join on resume [1,2,3,4,ph] merge at the end [1,2,3,4,5,–] ph = a placeholder row for the last position (it has no route yet); each new piece replaces it and adds a new one. route from the weights that sampled the token route from the new weights no route (last position)

The numbers come from the PR's own test, test_router_replay_survives_repeated_prefix_cuts_and_resume. It cuts after 12, 13, cuts again after 14, resumes, and generates 15. It asserts the final rows are [1, 2, 3, 4, 5, …].

Why keep the old rows instead of the resumed call's own

Using the resumed call's own routes (lane 2) would delete almost all of this PR. But tokens 12–14 were sampled with the old weights and routes, and their saved sampling probabilities come from that run. Training them with routes 92, 93 would make training's probabilities differ from the saved ones. Also, by default recompute_kv_cache_after_weight_updates is False. So a run that never restarts also keeps the old routes for tokens it already read, and uses new routes only after a weight update (lane 1). The PR reproduces that.

One thing does differ: after a restart, token 15 is sampled while reading 10…14 with the new weights, not from the old cache. Neither choice avoids that; it is the cost of resuming without a saved cache.

📝 review comment — vllm_worker_async.py:1420

So what. The design choice is right, but the PR never says it. The review asks for one sentence in the merge's docstring and in the prefix-recovery docs: a resumed row trains with the saved routes for its old tokens, and this matches a run without a restart when the KV cache is kept. With recompute_kv_cache_after_weight_updates: true it does not match, so say that too. The open problems are elsewhere: a crash on one resume path and how the routes are collected.