Skip to content

fix(engine): handle torch 2.9 allocator API deprecation - #245

Open
fangyuan-3149 wants to merge 1 commit into
FlashML-org:mainfrom
fangyuan-3149:fix/allocator-deprecation
Open

fix(engine): handle torch 2.9 allocator API deprecation#245
fangyuan-3149 wants to merge 1 commit into
FlashML-org:mainfrom
fangyuan-3149:fix/allocator-deprecation

Conversation

@fangyuan-3149

Copy link
Copy Markdown

Summary

Fix expandable_segments allocator setting broken by PyTorch 2.9 API migration.

Background & Impact

expandable_segments is a CUDA allocator switch that lets freed GPU memory blocks be split, merged, and reused at arbitrary sizes instead of being locked into fixed-size buckets.

Why it matters:
In MoE offload/hybrid paths, each step fetches expert weights from host RAM, dequantizes them to BF16 (7–9 MB), computes, and immediately frees the buffer. The block size varies per layer/step because the active expert count changes.

  • Without it: Freed blocks are locked into fixed-size buckets. A freed 7 MB block cannot satisfy a 9 MB request, so the allocator keeps asking the OS for fresh VRAM. Reserved balloons far above actual working set (e.g., 78 GB reserved for a 30 GB working set on an 8 GB card), causing OOM.
  • Enabled: Freed memory can be split, merged, and reused at any size. Reserved ≈ allocated.

Scope: Benefits every workload with variable-sized allocate/free cycles — MoE offload/hybrid dequant, KV cache growth, long-context serving. Memory-constrained GPUs benefit most, but all GPUs gain the same allocator efficiency.

Change

PyTorch 2.9 moved the internal API for expandable_segments from
torch.cuda.memory._set_allocator_settingstorch._C._cuda_setAllocatorSettings.
The old path still works but emits a FutureWarning, polluting CI logs.

Fix: Try the new private API torch._C._cuda_setAllocatorSettings first; on AttributeError fall back to the legacy torch.cuda.memory._set_allocator_settings. Behavior is identical; only the warning is removed.

Verification

  • python -m py_compile passes
  • pytest -m "not slow" — no new warnings, CI clean
  • Typical MoE offload/hybrid workloads continue to serve stably

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant