[Bugfix] Don't unmap cuMem allocations freed while asleep

[aoshen02]

Purpose

my_free() in csrc/cumem_allocator.cpp always calls unmap_and_release() on the handle returned by the Python free callback. For an allocation that is asleep, sleep() has already unmapped it and released the physical handle, so freeing a cuMem-pool tensor while asleep double-unmaps it. The driver error is left in the global error_code, so the next create_and_map (e.g. wake_up()) or cuMem allocation fails with that stale error. This happens on engine teardown while asleep, or when a pool tensor is freed during sleep. It becomes more frequent once CUDA graph pools are cuMem-backed (#59160).

Fix (CUDA path only; the ROCm path already skips an asleep allocation via an empty chunk list):

• unmap_and_release() sets *p_memHandle = 0 after a successful release;
• my_free() skips unmap_and_release() for a zero handle and only frees the address reservation.

No Python change is needed: _python_free_callback already pops the entry, drops the CPU backup and returns the handle. The Python-only alternative, remapping physical memory before the free so the C++ unmap succeeds, can OOM when a colocated trainer holds the VRAM freed by sleep.

Related:

• Fixes item 2 of cumem allocator: double-free and stale error codes during sleep/wake cycles #36651: my_free on allocations already released by sleep().
• [Bugfix][ROCm] Fix cumem sleep and teardown #46203 fixed the same double unmap/release on the ROCm path by tracking asleep allocations. This is the CUDA counterpart.
• [ROCm] Fix cumem OOM double-free of allocation handles #46054 proposed the same zero-handle idea and was closed because the ROCm fix had landed; the CUDA path stayed unfixed.

Test Plan

New test_free_while_asleep in tests/basic_correctness/memory/cumem/test_cumem.py (discard and offload variants): sleep, free a pool tensor while asleep by releasing the pools with release_pools() (as engine shutdown does), then check that the pointer is untracked, the CPU backup is dropped, no memory is remapped, and wake_up() succeeds with the other tensors intact.

pytest tests/basic_correctness/memory/cumem/test_cumem.py -k free_while_asleep

Test Result

GB300, torch 2.13.0+cu130, extension rebuilt from this branch:

• test_cumem.py: 14/14 pass, including test_free_while_asleep[discard] and [offload].
• After the release_pools() test change: the 7 non-LLM tests pass on this branch, and the 9 non-LLM tests pass with this PR applied on top of [Feature] Release WorkspaceManager scratch on sleep #59156, which stores a list of pools per tag. With the unpatched extension both new cases still fail with the same error. The test no longer depends on the layout of allocator_and_pools.
• Unpatched extension: both new cases fail at wake_up() with CUDA Error: invalid argument at csrc/cumem_allocator.cpp:258, from the stale error_code left by the double unmap.
• End to end: DeepSeek-V4-Flash TP4 with cuMem graph pools ([Feature] Release the CUDA graph pool on sleep #59160).
  • Setup: put the engine to sleep, fill all but 511 MiB of VRAM from another process, then shut the engine down (--shutdown-timeout 60).
  • Unpatched: 4× CUDA Error: invalid argument.
  • Remapping in Python before free: 4× out of memory.
  • With this fix: a clean teardown.

---

AI assistance was used to prepare this change; I reviewed the diff.

Re-check on current main (bdd31c31d), GB300 (rebased; test-only conflict in test_cumem.py resolved by keeping both sides)

Re-checked on main bdd31c3 (rebased: the only conflict was the insertion point in
test_cumem.py next to the new #59156/#59158/#59160 tests; both kept), 4x GB300,
torch 2.13.0+cu130, extension rebuilt from this branch:

• test_cumem.py: 23/23 pass, 3 runs. With main's extension (shipped .so and a rebuild
  from main's csrc) both test_free_while_asleep cases fail at wake_up() with
  "CUDA Error: invalid argument at csrc/cumem_allocator.cpp:258".
• DeepSeek-V4-Flash TP4 + MTP, --enable-sleep-mode, --shutdown-timeout 60: SIGTERM while
  asleep gives 4x "CUDA Error: invalid argument at cumem_allocator.cpp:258" (one per
  worker) on main in all 3 launches and 0 with this PR. Asleep memory (11321 MiB/GPU),
  wake time (1.26-1.34 s) and outputs (3 sleep/wake cycles, 28/28 greedy and 28/28
  seeded top-k/top-p identical to the awake reference) are unchanged.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[AlanFokCo]

Independent H20 verification of the current patch: we applied the diff from 85eac2a2102b42520017dd0f1f06719e2916c826 to pinned main 6f75cc7d5d1b5adfb399cbf79a8231ae66844fcf, rebuilt its native allocator extension, and reran the focused checks on driver 580.126.09 / PyTorch 2.13.0+cu130.

• A native probe that frees one asleep allocation, wakes another and checks its contents completed all 20 rounds.
• python -m pytest tests/basic_correctness/memory/cumem/test_cumem.py -k test_free_while_asleep: both discard/offload cases passed, using the current patch's test source and the rebuilt extension with our matching pinned-main wheel.
• Independently recompiling that extension reproduced the tested binary's SHA256.

For comparison, our prior run on the same pinned main without the fix failed in the first native-probe round with invalid argument, and both pytest cases failed. The current patch produces the same C++ source on this base as the earlier tested fix; the current test file was refreshed rather than reusing the older copy.

This verifies the current patch applied to the pinned base, not the entire current PR-head checkout or complete engine shutdown. The scope is the allocator free path; we did not need to rerun the large-model matrix to establish this result.