[Bug] Broadcasting full rollout_routed_experts data to every Megatron training rank will cause host OOM

Bug Description

When --use-rollout-routing-replay is enabled, the complete DP-local rollout_routed_experts payload is transferred to every Megatron training actor in the same DP group.

Each actor then locally performs padding, concatenation, CP slicing, and optionally TP slicing through prepare_routed_experts_for_routing_replay().

For large MoE training batches, this creates a large host-memory peak at the beginning of MegatronTrainRayActor.train. Under some configurations, the peak causes host OOM.

rollout_routed_experts has an approximate shape of:

[num_tokens, num_moe_layers, router_topk]

Its size therefore grows with all of the following:

• rollout batch size;
• sequence length;
• number of MoE layers;
• router top-k.

Unlike token-level scalar fields such as log probabilities, rollout_routed_experts stores multiple expert IDs for every token and every MoE layer.

example

In my situation, I train glm-4.7-flash (num_moe_layers=47, router_topk=4) with this config :

• --use-rollout-routing-replay
• rollout batch size: 96
• group size: 16
• max response len: 64k
• cp=2 tp=4

I was trainning using fully async on 2 node each with 8*H800 and 1.6TB Memory, 1node for rollout and 1 node for training. The average resp len is about 32k, so the total size of rollout data would be:

96 prompts
× 16 samples per prompt
× 32,768 tokens
× 47 MoE layers
× 4 experts per token
× 4 bytes per int32
≈ 35.25 GiB

If all 8 Megatron training actors receive and hold their own copy of this payload, the aggregate host-memory consumption can reach approximately: 35.25 GiB × 8 actors ≈ 282 GiB.

This is only the routed-experts preparation overhead. It is in addition to the Ray Object Store payload, model processes, training data, pinned routing replay buffers, and other runtime memory. As a result, this memory spike can easily lead to host OOM.

The following graph shows the corresponding USS spike at the beginning of MegatronTrainRayActor.train:

Steps to Reproduce

1. Train glm-4.7-flash with --use-rollout-routing-replay enabled.
2. Use a fully async setup with:
   • 1 node (8× H800) for rollout
   • 1 node (8× H800) for Megatron training
   • rollout batch size: 96
   • group size: 16
   • max response length: 64k
   • average response length: ~32k
   • CP: 2
   • TP: 4
3. Start training and monitor host memory usage at the beginning of MegatronTrainRayActor.train.

Expected Behavior

The training is heathy and will not encouter host OOM.

Actual Behavior

The training is heathy for a few steps, then it encouters an OOM.

Environment

• slime version: 0.3.2
• Python version: 3.12
• PyTorch version: 2.11.0+cu129
• CUDA/ROCm version: 13.0
• GPU type and count: 2*8 H800
• OS: Ubuntu
• SGLang version (if relevant):
• Megatron-LM version (if relevant):

Logs

Additional Context

No response

Pre-submission Checklist

• I have read the CONTRIBUTING.md and understand the collaboration scope.
• I have read the documentation and my issue is not addressed there.
• I have searched for existing issues and this is not a duplicate.
• I have provided a minimal, reproducible example.

Start by tracing MegatronTrainRayActor.train and prepare_routed_experts_for_routing_replay() to follow how rollout_routed_experts is transferred, padded, concatenated, and sliced. Reproduce the stated fully async configuration while monitoring host memory, then verify that training no longer creates a per-rank payload peak that causes OOM.