R3 PR: Rollout Routing Replay

[erictang000]

Overview

This PR adds support for Rollout Routing Replay (R3) from (See Paper).

See #815 for tracking of future tasks to fully support routing replay in all settings.

We add the following flags to enable R3:

cfg.generator.inference_engine.enable_return_routed_experts=True
cfg.trainer.policy.megatron_config.moe_enable_routing_replay=True

cfg.generator.inference_engine.enable_return_routed_experts=True is a pass through argument to vLLM, which records expert router indices (returning a list of dimension (batch_size, seq_len, num_layers, top_k).

We then pass this list rollout_expert_indices list through to Megatron's native RouterReplay feature (link).

When cfg.trainer.policy.megatron_config.moe_enable_routing_replay is set to true, Megatron initializes an instance of RouterReplay on each training worker rank. RouterReplay.set_replay_data(per_layer_data) can be used to set router decisions, and RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD) and RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_BACKWARD) can be used to set the routing mode to be forward or backward.

Results

GSM8K Training on moonlight16b-a3b shows R3 improves training stability - this can be seen both in logprob diffs as well as in clip_ratio, grad_norm, and loss, which otherwise explode and collapse training.

Supported Settings

Router Replay is supported for the following settings:

Generator Settings

• use_conversation_multi_turn=True and use_conversation_multi_turn=False
• batched=False and batched=True
• async_engine=True and async_engine=False
• NOT retokenize_chat_history mode - i.e. self.use_conversation_multi_turn and self.custom_chat_template
• NOT self.generator_cfg.step_wise_trajectories - there are some question marks about how to support this when using step wise training and not strictly appending (what should the routing look like for per turn obs that the inference engine doesn't see? - do we need to disable routing overrides for those tokens?)
• fully_async training - technically should work but not tested in this PR. Tracking in [skyrl-train] Enable routing replay in SkyRL #815

Inference Engine Settings

• TP and EP and DP should be supported from the vLLM side
• NOTE: _SKYRL_USE_NEW_INFERENCE is not supported - this will be added in a follow up PR
• NOTE: cfg.generator.distributed_executor_backend must be set to mp - hanging related to a Ray Compiled Graph issue occurs when using the default ray vLLM distributed executor backend. (see [Bug]: Generation hangs until RAY_CGRAPH_get_timeout (300s) with Ray compiled DAG executor vllm-project/vllm#36237 for details on the error that comes up)
• NOTE: The above use of mp also means that serving must be single node per engine, until we add support for using the mp backend with multi-node serving - progress tracked here: [vllm] Fully enable mp distributed executor backend in vLLM #1309

Trainer Settings

• TP, EP, DP are all supported. CP is in progress in this PR but needs more testing. CP + PP will be added in a follow up PR.

Custom Generator support

• Custom generators using SkyRL's inference engine should just plumb through

Tests

Adds test_router_replay.py, which includes:

• test_logprobs - integration test that runs a training batch through vllm, and through megatron with and without R3, to verify that logprob diffs are lower with routing replay
• test_forward_backward - unit test for forward_backward that verifies that a training step can complete successfully when routing replay indices are passed in

Adds test_generator_multi_turn_gsm8k_router_replay to test_skyrl_gym_generator to verify that the SkyRLGymGenerator plumbs through the router indices in an expected format.

Rollout Routing Replay

Relevant resources:
vLLM PR: vllm-project/vllm#28284
Verl PR: verl-project/verl#4101
Mindlab blog: https://macaron.im/mindlab/research/router-replay-r3-why-it-failed-and-how-we-fixed-it
Megatron-LM API guide: https://github.com/NVIDIA/Megatron-LM/blob/main/docs/api-guide/router_replay.md

---

[erictang000]

Forward pass with router replay showing lower logprob diff!

[erictang000]

Current State:

For small scale tests, routing replay seems to be working as shown above - tested only for TP=8 serving and TP=4 and EP=8 training.

What's not working:

• Deepseek style models (like moonlight-16b-a3b) - the average logprobs for both with and without routing replay seem to be way off from what is expected from the inference engine (like ~8 average logprobs vs 0.2 average logprobs) - maybe this is specific to moonlight + some megatron config? not sure. Reproducible by running the test_logprobs tests
• Running a real training batch runs into hanging for all models - see [Bug]: Generation hangs until RAY_CGRAPH_get_timeout (300s) with Ray compiled DAG executor vllm-project/vllm#36237 for the relevant stack trace. This seems to happen pretty much only when router replay is enabled on the vllm end. We could try upgrading vllm to 0.17.0 to see if anything has been pushed that could help fix this. Reproducible by running a large enough training batch with router replay on the test_logprobs tests

TODOs:

• fix above bugs
• clean up code
• add checks for settings where r3 won't be supported for now (async RL, batched generator, use_conversation_multi_turn=false)
• run full end to end test showing r3 minimizes drift on DAPO (with at least one model family - others just need to pass tests.

[erictang000]

Deepseek style models (like moonlight-16b-a3b) - the average logprobs for both with and without routing replay seem to be way off from what is expected from the inference engine (like ~8 average logprobs vs 0.2 average logprobs) - maybe this is specific to moonlight + some megatron config? not sure. Reproducible by running the test_logprobs tests

Solved! the issue was that in our test we were setting cfg.trainer.use_sample_packing=False, and also setting NVTE_FUSED_ATTN=0. This was causing flash-attn to be used without sample packing for the moonlight forward pass - flash attn produces incorrect logprobs that differ greatly from vLLM:

vLLM logprobs     - mean: -2.564655
Megatron (replay) - mean: -9.231623
Megatron (no rep) - mean: -9.647593

after setting cfg.trainer.use_sample_packing=True and NVTE_FUSED_ATTN=1 (to allow setting transformer_config_kwargs["attention_backend"] to "fused" correctly

vLLM logprobs     - mean: -0.223607, std: 0.674102
Megatron (replay) - mean: -0.223626, std: 0.674850
Megatron (no rep) - mean: -0.224379, std: 0.677036
With replay    - logprob diff mean: 0.006648, std: 0.021737
Without replay - logprob diff mean: 0.011115, std: 0.035957

[erictang000]

Running a real training batch runs into hanging for all models - see vllm-project/vllm#36237 for the relevant stack trace. This seems to happen pretty much only when router replay is enabled on the vllm end. We could try upgrading vllm to 0.17.0 to see if anything has been pushed that could help fix this. Reproducible by running a large enough training batch with router replay on the test_logprobs tests

verified that cherry picking the changes from #1300 to use the mp backend allow us to work around the compiled graph timeout

[erictang000]

skipping these for now - need to test some smaller models out - tracking as a follow up task in #815, will do ASAP along with PP + CP support

[erictang000]

this is needed for "allenai/OLMoE-1B-7B-0924", which is used in the skyrl_gym_generator test - i plan to port over this model to the other router replay tests (and maybe megatron moe tests in general) since it's supported in Megatron-Bridge and is a 7B with 1B activated MoE that makes it nice for CI.

[erictang000]

additional curves for DAPO for posterity (green is DAPO on moonlight w/ TIS, blue is same but w/o TIS + router replay)

Time for fwd/train is slightly higher with R3: