[skyrl-train] Enable routing replay in SkyRL

[erictang000]

This issue tracks the status of Rollout Routing Replay (R3) support being merged into SkyRL

• Initial PR - support for TP on inference, TP + CP + EP on Trainer: R3 PR: Rollout Routing Replay #1273
• Migrate tests to run on CI (use OlMoe since it's the smallest moe model supported on megatron bridge): [CI] Migrate MoE Tests to OlMoE #1384
• Plumb through router replay feature on new inference stack
• Test router replay plumbing for fully async training and step wise training
• Support for CP + PP on megatron Trainer (should just require slicing on megatron workers to get the right indices): [train][2/N] Support for Megatron PP + CP for R3 #1335
• Considering support for R2 - replaying the megatron forward pass expert routing - lower priority

[erictang000]

router replay data transfer:

Router replay can quickly make the size of the data batch to transfer between trainer and inference quite large. See example below where 32GB gets transferred per node for GLM-4.7-Flash, with batch size=128, n_samples_per_prompt=8, num_layers=47, top_k_experts=4, and max_context_length=20000.

This quickly results in ~16 GiB of expert indices, and since we are doing tp=4, we have 2 dp ranks per node materializing the data, resulting in 32GB per node

This can be a bottleneck for scaling to larger models with more layers and longer context length. We should take the following steps to improve this

• use uint8/int16 for expert indices (GLM-5 uses 256 routed experts, Kimi-2.5 uses 384, which needs more than uint8...) [router replay] downcast expert router indices to uint8/int16 to reduce space #1378
• remove unnecessary full batch data transfer to each node: [trainer] refactor dispatch_from_staged to individually serialize DP chunks to avoid materializing whole batch on all workers #1376
• remove padding tokens for expert indices
• ensure that all trainer ranks block on retrieving data to avoid hanging NCCL communication if different DP ranks take different amounts of time to retrieve from ray object store.
• Work on gpu based transfer for routed expert indices
• Time object serialization/deserialization from driver -> train workers