InternLM / InternLM/xtuner

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

Open
#2,025 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
5.2k
Forks
448
Avg merge
3d 15h
Merged PRs (30d)
26

Description

## 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`](https://github.com/InternLM/xtuner/commit/575d7e058040baa7f609b3d5d3f397653877bc25).

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:

```text
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()`](https://github.com/InternLM/xtuner/blob/575d7e058040baa7f609b3d5d3f397653877bc25/xtuner/v1/rl/trainer/worker.py#L616-L619), device transfer happens before SP splitting:

```python
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()`](https://github.com/InternLM/xtuner/blob/575d7e058040baa7f609b3d5d3f397653877bc25/xtuner/v1/data_proto/sequence_context.py#L605-L610) moves the full routing trace to the device:

```python
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.

Contributor guide

Open the contributing guide

Research direction

Start with TrainingWorker.fit() in xtuner/v1/rl/trainer/worker.py and SequenceContext.to() in xtuner/v1/data_proto/sequence_context.py, then trace the existing split and batch-preparation lifetimes. Add focused tests for rank-local transfer and bounded retained chunks, followed by the suggested one-step multi-node smoke test. Done means the step completes without host or CUDA OOM while RLOO grouping, normalization, accumulation, and synchronization remain unchanged.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.