[Bug] Routed-expert traces cause host OOM with offload enabled and CUDA OOM before SP split when disabled

[matrix72c]

Summary

In a multi-node MoE RLOO training run, offload_rollout_routed_experts only changes where a large set of rollout routing traces exhausts memory:

• With offload_rollout_routed_experts=True, learner host memory grows until Ray's 95% node-memory threshold is reached.
• With offload_rollout_routed_experts=False, the first learner step fails with CUDA OOM while moving rollout_routed_experts to the GPU, before sequence-parallel sharding is applied.

The current behavior suggests that rollout routing traces are materialized for too much of the learner batch at once, and that the non-offload path transfers the unsharded trace to every SP rank.

Environment

The relevant code path is still present on current upstream main at commit 575d7e0.

Observed training topology/configuration:

• 3 nodes, 8 H200 GPUs per node
• 16 training workers and 8 rollout workers
• Sequence parallel size: 8
• Train batch size: 32 rollout groups
• RLOO repeat K: 4
• Maximum packed sequence length: 262,144
• The first 32 rollout groups expanded to approximately 2,390 learner samples

The run used a nearby checkout with an unrelated rollout-session lifecycle patch. The routing-trace code discussed below matches upstream main.

Observed behavior

offload_rollout_routed_experts=True: host OOM

The rollout phase completes and the learner prepares approximately 2,390 training batches. Host memory then crosses Ray's 95% node-memory threshold. Eight training workers on the affected learner node each use approximately 100 GiB of memory. Ray kills workers and the training step fails before the first optimizer update.

Ray object-store usage is comparatively small, so the dominant pressure appears to be worker heap rather than the object store.

offload_rollout_routed_experts=False: CUDA OOM

The same workload reaches the first learner step, then fails in this path:

TrainingWorker.fit()
  -> seq_ctx = data["seq_ctx"].to(DEVICE)
  -> SequenceContext.to()
  -> self.rollout_routed_experts.to(device)
  -> CUDA out of memory

At failure, the H200 had less than 200 MiB free and PyTorch attempted to allocate another approximately 640 MiB. Reserved-but-unallocated memory was negligible, so this looked like capacity exhaustion rather than fragmentation.

Suspected cause

In TrainingWorker.fit(), device transfer happens before SP splitting:

seq_ctx.offload_rollout_routed_experts = self.config.offload_rollout_routed_experts
seq_ctx = data["seq_ctx"].to(DEVICE)
if self.sp_mesh.size() > 1:
    seq_ctx = seq_ctx.split(self.sp_mesh)

When offload is disabled, SequenceContext.to() moves the full routing trace to the device:

if (
    self.rollout_routed_experts is not None
    and not self.offload_rollout_routed_experts
    and hasattr(self.rollout_routed_experts, "to")
):
    self.rollout_routed_experts = self.rollout_routed_experts.to(device)

This means each SP rank may materialize the full routing trace on GPU before retaining only its local sequence shard.

With offload enabled, the GPU copy is avoided, but routing traces for a large number of prepared learner samples remain resident in CPU worker memory. The exact per-field host-memory breakdown has not yet been profiled, so this part is an inference from worker RSS, object-store usage, and the batch preparation/lifetime in the code.

Expected behavior

The routing-trace memory footprint should be bounded independently of the total number of learner samples in an RL step. With sequence parallelism, each rank should only transfer or retain its rank-local route-trace shard on GPU.

Possible fixes

1. Split rollout_routed_experts by SP rank on CPU before transferring it to the target GPU. A route-only rank-local/lazy-transfer path may be safer than reordering SequenceContext.to() for every field.
2. Stream learner samples in bounded chunks instead of retaining routing traces for the entire expanded training batch. Release each chunk after forward/backward while preserving effective batch size through gradient accumulation.
3. Load route traces only for the current batch/layer and release them immediately after use.
4. Add diagnostics for routing-trace bytes, active chunks, worker RSS/USS, and CUDA allocated/reserved memory.
5. Consider storing expert indices in a narrower integer dtype when the number of experts permits it, converting only at the point of use.

Suggested validation

• Add a unit test verifying that only the SP rank-local routing-trace shard is moved to each GPU.
• Add a test ensuring that the number of simultaneously retained route-trace chunks is bounded.
• Run a one-step multi-node smoke test and verify completion of forward/backward, optimizer update, and weight synchronization.
• Verify that RLOO grouping, token normalization, and gradient-accumulation semantics are unchanged.