[RL/DeepEP] Generator OOMs in vLLM cudagraph capture on memory-tight MoE (expandable_segments incompatible with DeepEP; gpu_memory_utilization is a no-op)

[yichuan-w]

Summary

Serving a large MoE (Qwen3-235B-A22B) generator through the DeepEP EXPAND inference path
(torchtitan/distributed/deepep/inference_override.py deepep_inference override, cudagraph
enabled) in the RL Search-R1 recipe OOMs during vLLM engine init -- inside cudagraph memory
profiling, before the KV cache is allocated. The OOM is CUDA caching-allocator fragmentation
that the usual fix (PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True) cannot address, because
DeepEP requires expandable_segments:False (NVSHMEM/RDMA register fixed GPU buffers that
expandable/VMM segments would remap and break).

Environment

• torchtitan main
• Model: Qwen3-235B-A22B, moe_comm_backend="deepep"
• Generator: vLLM, override torchtitan.distributed.deepep.inference_override, cudagraph FULL_DECODE_ONLY
• Generator parallelism DP=4 / TP=4 / EP=16 (also observed, marginally, at DP=2 / TP=4 / EP=8)
• GPU: H100 80GB
• PYTORCH_[CUDA_]ALLOC_CONF=expandable_segments:False (required by DeepEP)

Traceback (abridged)

EngineCore.__init__ -> _initialize_kv_caches -> determine_available_memory
  -> profile_cudagraph_memory -> _warmup_and_capture -> _dummy_run
  -> model.forward -> MoE
  -> torchtitan/distributed/deepep/deepep.py combine_tokens (hidden_states * per_row_score)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB.
  GPU 79.18 GiB total, ~1.9 GiB free; ~40 GiB allocated, ~2.3 GiB CUDA-graph pools,
  ~32.7 GiB reserved-but-unallocated (fragmentation).

Key findings

1. OOM is during cudagraph capture, before KV allocation. gpu_memory_utilization
   (gpu_memory_limit) has no effect -- 0.9 / 0.7 / 0.6 give byte-identical OOMs
   (~32.7 GiB reserved-but-unallocated each time), since it only sizes the not-yet-allocated
   KV pool.
2. expandable_segments:True (the fix the PyTorch OOM message itself suggests) is unusable --
   it breaks DeepEP's NVSHMEM/RDMA fixed-buffer registration.
3. The only knob that removes the OOM is enforce_eager (generator cudagraph.enable=False),
   which skips profile_cudagraph_memory / _warmup_and_capture entirely -- at the cost of the
   decode cudagraph speedup.
4. Fragmentation appears to grow with EP degree / num_max_tokens_per_rank: ~12.6 GiB (marginal)
   at EP=8 / max_num_batched_tokens=1024, ~32.7 GiB (fatal) at EP=16 / 2048. This is not cleanly
   isolated (EP, token budget, and single- vs multi-host all differed), so "scales with EP" is a
   hypothesis.

Impact

A DeepEP EXPAND generator with cudagraph cannot be brought up on a memory-tight large MoE; one is
forced to enforce_eager and lose the decode cudagraph speedup that FULL_DECODE_ONLY is meant
to provide.

Possible directions (for maintainers)

• Reuse a single DeepEP dispatch/combine scratch buffer across all captured cudagraph sizes so
  capture warmup does not churn / fragment the allocator.
• Give DeepEP capture-time temporaries a dedicated pool separate from the (expandable-segments-
  disabled) default pool.
• Document a DeepEP-safe PYTORCH_CUDA_ALLOC_CONF (e.g. max_split_size_mb /
  garbage_collection_threshold) that mitigates the fragmentation without expandable segments.
• Bound / expose the cudagraph capture sizes to cap the fragmentation.

Current workaround

Generator cudagraph.enable=False (enforce_eager). Keeps the DeepEP EXPAND dispatch; drops
graph capture.

[wwwjn]

@yichuan-w From my previous run, 235B model in generator (16GPUs) with DeepEP is easy to OOM, during the cudagraph capture phase. So it's highly likely because of the saved cudagraph itself takes a lot of memory.

The config I run is capture_size =[8] (only capture 1 cudagraph, because I set num_requests=8, so it only covers decode, FULL_DECODE_ONLY) 16GPU, 235B model, TP4 (4 kv_heads) DP4

But worth kick off a agent to dig into the size of the cudagraph - what's the estimated size of a single cudagraph for 235B model?