[https://nvbugs/6248837][fix] Free worker GPU memory on executor shutdown#15099
Conversation
📝 WalkthroughWalkthroughThis PR updates the test waiver list to change which parameterization of the Kimi K25 NVFP4 test is exempted from running. The waiver for the ChangesTest Waiver Parameterization Update
🎯 1 (Trivial) | ⏱️ ~2 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/bot run --stage-list "GB200-8_GPUs-2_Nodes-PyTorch-1" |
|
PR_Github #52747 [ run ] triggered by Bot. Commit: |
|
PR_Github #52747 [ run ] completed with state
|
|
/bot run --stage-list "GB200-8_GPUs-2_Nodes-PyTorch-1" |
|
PR_Github #53037 [ run ] triggered by Bot. Commit: |
|
PR_Github #53037 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #53231 [ run ] triggered by Bot. Commit: |
|
/bot run --stage-list "GB200-8_GPUs-2_Nodes-PyTorch-1" |
|
PR_Github #53231 [ run ] completed with state
|
|
PR_Github #53291 [ run ] triggered by Bot. Commit: |
|
PR_Github #53291 [ run ] completed with state
|
|
/bot run --stage-list "GB200-8_GPUs-2_Nodes-PyTorch-1" |
|
PR_Github #53518 [ run ] triggered by Bot. Commit: |
|
PR_Github #53518 [ run ] completed with state
|
|
/bot run --stage-list "GB200-8_GPUs-2_Nodes-PyTorch-1" |
|
PR_Github #53552 [ run ] triggered by Bot. Commit: |
|
PR_Github #53552 [ run ] completed with state
|
|
/bot run --stage-list "GB200-8_GPUs-2_Nodes-PyTorch-1" |
|
PR_Github #53604 [ run ] triggered by Bot. Commit: |
|
PR_Github #53604 [ run ] completed with state |
tianyuxbear
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #56905 [ run ] triggered by Bot. Commit: |
|
PR_Github #56905 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56947 [ run ] triggered by Bot. Commit: |
|
PR_Github #56947 [ run ] completed with state
|
…down
GenerationExecutorWorker.shutdown() dropped the engine reference but
never returned the engine's GPU memory to the driver. Reference cycles
keep the model tensors alive until a later GC, and PyTorch's caching
allocator holds the freed blocks as "reserved" instead of releasing them.
In a persistent worker process that is reused to build another LLM --
successive multi-GPU tests in one CI shard, or any long-running process
that reloads models -- this leftover, roughly weights-sized reservation
can exhaust the device and OOM the next model load at model.to("cuda").
This is why TestKimiK25::test_nvfp4[tp8] OOMed on all ranks when it ran
after a large model (DeepSeek-R1) in the same shard, yet passed in
isolation.
Force gc.collect() + torch.cuda.empty_cache() in
GenerationExecutorWorker.shutdown(), after the engine and the
torch.distributed process group are torn down, so each worker returns
the previous model's memory to the driver before the next LLM is built.
Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
…le cases Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #57004 [ run ] triggered by Bot. Commit: |
|
PR_Github #57004 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #57312 [ run ] triggered by Bot. Commit: |
|
PR_Github #57312 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #57561 [ run ] triggered by Bot. Commit: |
|
PR_Github #57561 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57682 [ run ] triggered by Bot. Commit: |
|
PR_Github #57682 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #57696 [ run ] triggered by Bot. Commit: |
|
PR_Github #57696 [ run ] completed with state |
…down (NVIDIA#15099) Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
5 participants
Description
accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8]intermittently OOMed at executor init (model.to("cuda")) on theGB200-8_GPUs-2_Nodesstage, yet passed when run alone. Root cause: GPU memory was not released between tests that share a CI shard.A shard runs as one long-lived MPI job (8 ranks / 2 nodes) whose worker processes persist across all tests in the shard.
GenerationExecutorWorker.shutdown()dropped the engine reference but never returned its GPU memory to the driver — reference cycles keep the model tensors alive until a later GC, and PyTorch's caching allocator holds the freed blocks as "reserved". So a large predecessor (DeepSeek-R1, ~110 GiB/GPU) left its weights resident in the workers, and the next large model (e.g. Kimi-K2.5, ~76 GiB/GPU) then OOMed on all 8 ranks with only ~73 GiB free. This also explains the flakiness: it only failed when a heavy test ran immediately before in the same shard.Fix
Force
gc.collect()+torch.cuda.empty_cache()inGenerationExecutorWorker.shutdown()(after the engine and torch.distributed process group are torn down), so each worker returns the previous model's memory to the driver before the next LLM is built.Local vs CI debugging summary
A pre-
LLM()nvidia-smiprobe (diagnostic only, not in this PR) measured the device state at test start:python3process)After the fix, the worker-shutdown step returns that memory at the test→test handoff (
free 72.76 → 169.69 GiB) and the previously-failing tests pass.Same-root-cause OOM tests
The same worker-shutdown OOM affected several GB200 2-node tests. Status in this PR:
TestKimiK25::test_nvfp4[tp8]TestKimiK25::test_nvfp4[dep8]TestKimiK25::test_nvfp4[tp8_attn_dp]TestDeepSeekR1::test_fp8_blockscale[throughput_mtp]TestDeepSeekR1::test_fp8_blockscale[throughput_mtp_trtllm]TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_mtp]Test Coverage
TestKimiK25::test_nvfp4[tp8|dep8|tp8_attn_dp]andTestDeepSeekR1::test_fp8_blockscale[throughput_mtp|throughput_mtp_trtllm]on theGB200-8_GPUs-2_Nodes-PyTorch-*stages.