OOM from intra-step memory fragmentation
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
We hit OOM as early as the 2nd GRPO backward step from a combination of inter-step tensor retention and intra-step fragmentation (additional context below). `expandable_segments` is the documented mitigation for the fragmentation but is unavailable in colocated mode due to vLLM's `CuMemAllocator` (pytorch#147851).
**Steps/Code to reproduce bug**
We're running a single-node variant of the reference recipe `grpo-qwen3.5-35ba3b-2n8g-megatron-ep16.yaml` on 1n8g with EP=4, TP=2 instead of 2n8g / EP=16, TP=2. Our setup:
- Reproduced on both GCP and AWS
- NeMo-RL: `main` @ `e4eb1ee0`, colocated mode
- Model: Qwen3.5-35B-A3B (MoE, 256 experts, 40 layers)
- Hardware: 8× H200 141 GiB
- Parallelism: TP=2, EP=4, sequence parallel, activation checkpointing, bf16
- vLLM 0.17.1, transformers 5.3.0, flashinfer 0.6.4, Python 3.13
- GRPO, `gpu_memory_utilization=0.45`, 32 rollouts/step (4 prompts × 8 generations)
**Expected behavior**
Training proceeds for 100+ steps without OOM.
**Additional context**
Based on what we've found so far, three factors appear to compound:
### 1. Inter-step retention in `MoEAlltoAllTokenDispatcher`
Over multiple runs, three dispatcher attributes remain alive across `offload_after_refit()`:
- `self.probs`
- `self.routing_map`
- `self.reversed_local_input_permutation_mapping`
Each tensor is small (~16 MB) but pins a multi-GB allocator segment. Pattern is 40 layers × 3 attributes × N steps. Nulling these at the end of `combine_preprocess()` eliminated the retention in our setup. We're flagging this in case it makes sense to fix upstream in Megatron-Core.
We observed this on `zhiyul/hemil/automodel-transformers-v5`. Does this dispatcher observation match what you've seen on the existing `zhiyul/oom_repro_w_cpu_profiler`, `geshen/work_oom_fix`, `guyueh/refit_cleanup_ipc_buffer` repros?
### 2. Intra-step fragmentation independent of (1)
Even with (1) addressed, we measured `max_memory_reserved` running 5–26 GiB above `max_memory_allocated` per rank per step during forward/backward. Real peak need is ~78–96 GiB; reserved inflates to 83–122 GiB. This is allocator behavior *within* the step, not carry-over (inter-step cleanup doesn't address it).
### 3. `expandable_segments` is the documented mitigation for (2), but unavailable in colocated mode
NeMo-RL already enables `expandable_segments:True` on Hopper+ (PR #594) for non-colocated. vLLM's `CuMemAllocator` (`cumem.py`, v0.17.0+) is incompatible with `expandable_segments`, so colocated runs can't opt in. Tracking issue: [pytorch#147851](https://github.com/pytorch/pytorch/issues/147851).
We can't test `expandable_segments` in colocated mode for this workload. Is there a planned path forward for colocated MoE training on H200 given the `expandable_segments` block?
Contributor guide
Assessment
This issue has not been assessed yet.