[Bug]: Distributed Tests (4 GPUs) failing in main branch CI

[njhill]

This is now consistently failing with CUDA OOM: https://buildkite.com/vllm/ci/builds/22221#01977f3a-71ea-41cb-bbeb-a43340a10124

I narrowed this down to #19572 which appears to have introduced the issue.

[mgoin]

@yewentao256 can you please help take a look?

[yewentao256]

Certainly, taking a look now

[mgoin]

It looks like the failing command within the test is

python3 ../examples/offline_inference/data_parallel.py

which does run ibm-research/PowerMoE-3b

[yewentao256]

#20153
Could we use this pr to fix temporally first? And I will rent a L4 cluster to reproduce and figure out why later

[yewentao256]

Using a L4 *4 cluster to verify, here are my results:

• This pr is the original issue: [Core] Enable CUDA graphs for DP + All2All kernels  #18724, but I do think it is reasonable to adopt this.
• And also verified that this pr's commit could solve this issue: [Bugfix] default set cuda_graph_sizes to max_num_seqs for v1 engine #20062
• However, main + this pr's update still OOM (quite interesting, something changed after that pr).

So I think update config to reduce the memory usage could solve this CI issue. CC @mgoin @njhill

[njhill]

Thanks @yewentao256 I think that's a good fix to unblock the CI.

But in theory we should not OOM with defaults ... the memory profiling should take cuda graphs into account. So I feel like there may be a more general issue to investigate here. cc @varun-sundar-rabindranath

[varun-sundar-rabindranath]

Digging further into this :-
PR #18724 enables cudagraphs for DP workloads.

Debug commit : The PR 18724
Debug machine : L4 x 4
Total GPU Memory : 22.4 GB
Memory used for KV Cache + Weights : 0.90 * 22.4 = 20.16 GB
Memory left for cuda-graph capture := 2.24 GB (logging using pynvml reports 2.6GB)

Experiments:

• data_parallel.py --dp-size 2 --tp-size 2 + enforce-eager : Pass

Testing data-parallel + cuda-graphs

• data_parallel.py --dp-size 2 --tp-size 2 + cuda-graphs for batch-sizes in {1, 2, 4, 8, 16, 24 ... 512} : OOM
• data_parallel.py --dp-size 2 --tp-size 2 + cuda-graphs for batch-sizes in {1, 2, 4, 8, 16, 24 ... 496} : Pass
• data_parallel.py --dp-size 2 --tp-size 1 + cuda-graphs for batch-sizes in {1, 2, 4, 8, 16, 24 ... 512}: Pass
• data_parallel.py --dp-size 4 --tp-size 1 + cuda-graphs for batch-sizes in {1, 2, 4, 8, 16, 24 ... 512}:OOM
  Comment: I see that for each batch-size cuda-graphs consumes an additional 41MB

Testing no data-parallel + cuda-graphs

• data_parallel.py --dp-size 1 --tp-size 1 + cuda-graphs for batch-sizes in {1, 2, 4, 8, 16, 24 ... 512} : Pass
• data_parallel.py --dp-size 1 --tp-size 1 + cuda-graphs for batch-sizes in {1, 2, 4, 8, 16, 24 ... 600} : OOM [edit]
• data_parallel.py --dp-size 1 --tp-size 1 + cuda-graphs for batch-sizes in {1, 2, 4, 8, 16, 24 ... 600} : Pass
  Comment : For the non DP work-load eventhough each cudagraph capture consumed memory - it is not as much as the DP case.

I believe this is generally a CUDA Graph capture issue and should be investigated further at the cuda-graph capture level .
Will continue investigating why DP graph-capture needs so much additional memory.