[Bugfix] Export weight-cache IPC tensors separately for each client - #56472
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
This pull request has merge conflicts that must be resolved before it can be |
0a4a894 to
a2ed82c
Compare
|
Rebased onto the latest main ( Re-ran the CUDA IPC regression on A10: main gives 1 passed / 1 failed, and the rebased fix gives 2 passed with no GPU skips. Pre-commit and mypy 3.12 also passed. |
| @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA IPC requires a GPU") | ||
| @pytest.mark.parametrize("zero_copy_consumer", [False, True], ids=["copy", "mixed"]) | ||
| def test_ipc_cache_release_keeps_live_consumers(monkeypatch, zero_copy_consumer): |
There was a problem hiding this comment.
Can we use e2e test instead here?
There was a problem hiding this comment.
Added an E2E test using Qwen/Qwen3.5-0.8B and the existing daemon and runner helpers. It keeps a zero-copy engine alive while a copy engine loads, generates, and shuts down, then compares the zero-copy engine's output with a disk-loaded baseline. IPC fallback and prefix caching are disabled.
The E2E also passed without the fix on A10, so I kept the focused allocation-lifetime regression as well; replacing it would lose the test that reproduces the bug. The production fix is unchanged.
Validation on 2fc5add46001: 3 passed (the two ownership cases and the E2E), no skips. Without the fix, the ownership cases give 1 passed / 1 failed: the mixed-client case observes 0 allocated bytes where 4096 must remain alive. Pre-commit and mypy 3.12 also passed. The A10 runs use the target weight-cache and supporting Python files over an existing vLLM CUDA image built from d9fbe526c078.
There was a problem hiding this comment.
The E2E also passed without the fix on A10, so I kept the focused allocation-lifetime regression as well; replacing it would lose the test that reproduces the bug. The production fix is unchanged.
I see, if the E2E can't catch the regression, there's no need to keep it then. 🤔
There was a problem hiding this comment.
Removed the added model E2E and its helper changes in 5d49860235. Kept the copy-only and mixed-client allocation-lifetime regression. The production fix is unchanged.
Re-ran the retained tests on A10: 2 passed, no skips. Pre-commit and mypy 3.12 also passed.
a2ed82c to
2fc5add
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
2fc5add to
3c8cafa
Compare
|
Rebased onto main Re-ran the A10 checks: the new main gives 1 passed / 1 failed in the focused lifecycle tests, and the rebased fix gives 3 passed with no skips, including the Qwen model E2E. Pre-commit and mypy 3.12 passed. The A10 runs use the target weight-cache and supporting Python files over the existing CUDA image, as described in the PR. |
Generate reduction arguments for each get_state transfer instead of replaying a single export. This lets PyTorch retain the producer allocation when a copy client releases the daemon cache while a zero-copy client is still using its weights. Cover copy-only and mixed clients with a CUDA multiprocessing regression that checks live mappings and eventual allocation reclamation. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: 子华 <huaxi.shx@alibaba-inc.com>
3c8cafa to
5d49860
Compare
|
✅ @shaohuaxi, CI is now available for this PR.
|
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
|
/ci run |
|
✅ Triggered Buildkite CI #88786 for commit |
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #88786. |
…llm-project#56472) Signed-off-by: 子华 <huaxi.shx@alibaba-inc.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
2 participants
Purpose
Starting a copy-mode client can release the weight cache while another engine is still using zero-copy weights. The daemon currently calls
reduce_tensoronce and sends the same arguments to every client. Those arguments include a reference-count slot for one consumer, so releasing the copy client's temporary IPC tensors can bring the count to zero while the other engine still holds its mappings.Export the tensors for each
get_stateresponse instead. PyTorch then tracks each consumer separately and can keep the underlying allocation alive after the daemon drops its model. Exporting at startup is removed as well, so there is no unused initial reference. The request protocol stays the same.Related: #51259 fixed export reuse in packed weight transfers; this change covers the separate weight-cache daemon. I checked the related Fast Start PRs, including #56047 and #55468, and found no existing fix for this path.
This change and its tests were developed with AI assistance.
Test Plan
The CUDA multiprocessing regression covers copy-only and mixed zero-copy/copy clients. It checks that releasing the daemon cache keeps the producer allocation alive for a live zero-copy consumer, that consumers can still read their tensors, and that the allocation is reclaimed after the last consumer exits.
Test Result
The A10 runs use the target weight-cache and supporting Python files over an existing vLLM CUDA image built from
d9fbe526c078(PyTorch 2.13.0+cu130).5d49860235f7dbf49dad11e3)Earlier model validation of the unchanged production fix (
3c8cafa5ac) produced matching token IDs and text forQwen/Qwen3.5-0.8Bacross a copy client's load and shutdown, compared with a disk-loaded baseline.