Stop ignoring a failed cudaSetDevice in CudaAllocator - #22093
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22093
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 182 PendingAs of commit afa82ae with merge base 368a849 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
|
Gasoonjia
approved these changes
Aug 24, 2026
CudaAllocator switches the current CUDA device before it allocates, frees or copies, and switches back afterwards. The result of that switch was thrown away in all three places. When the switch fails, for example because the requested device index does not exist on this machine, the work still went ahead on whatever device happened to be current. allocate() then returned success with a pointer that lives on a different device than the caller asked for. The caller stores the requested index next to that pointer, so the pointer and its recorded device disagree from then on, and the mistake only shows up much later as a wrong result or an unrelated CUDA error. Now allocate() and the copy helpers log the CUDA error and return Error::MemoryAllocationFailed and Error::Internal instead of going ahead. deallocate() returns void and cudaFree works on a pointer from any device under unified addressing, so it logs the error and still frees rather than leaking. Example, on a machine with one GPU: ``` auto r = CudaAllocator::instance().allocate(1024, /*index=*/1); // before: r.ok() == true, pointer actually on device 0 // after: r.ok() == false, r.error() == Error::MemoryAllocationFailed ``` Test plan: Added three tests to test_cuda_allocator.cpp that ask for device index device_count, which is one past the last valid ordinal on any machine, so the switch always fails: AllocateOnMissingDeviceFails CopyHostToDeviceOnMissingDeviceFails CopyDeviceToHostOnMissingDeviceFails Built and ran the suite on one NVIDIA H100. With this change: [ PASSED ] 9 tests. With the allocator reverted to the old code and the same three tests: [ FAILED ] CudaAllocatorTest.AllocateOnMissingDeviceFails [ FAILED ] CudaAllocatorTest.CopyHostToDeviceOnMissingDeviceFails [ FAILED ] CudaAllocatorTest.CopyDeviceToHostOnMissingDeviceFails The allocate failure showed the old behavior directly: the call reported success for a device that does not exist. The logged message on the fixed build: CudaAllocator::allocate: cudaSetDevice(1) failed: invalid device ordinal clang-format reports no changes needed on either touched file. Not covered: the deallocate path only logs, so no test asserts on it.
…e tests cudaSetDevice and cudaGetDevice failures were reported as MemoryAllocationFailed and InvalidArgument. Neither fits: a device that cannot be selected is not out of memory, and a well-formed ordinal that does not exist is an operational failure rather than a bad argument. Use Internal, which is what Exception.h maps every CUDA error to and what cuda_mutable_state.cpp already returns for this same call. Also fail instead of continuing when cudaGetDevice fails, since without the current device there is no way to switch or restore, and the allocation would land on whatever device happened to be current. The three tests that need a valid-but-absent ordinal now check that one exists in DeviceIndex (int8_t) rather than gating in SetUp, which skipped the six device-0 tests too. The bound is <= max, so 127 runs and 128 skips.
shoumikhin
force-pushed
the
fix-cuda-setdevice-unchecked
branch
from
August 24, 2026 22:58
b2fd544 to
afa82ae
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
CudaAllocatorswitches the current CUDA device before it allocates, frees orcopies, and switches back afterwards. The result of that switch was thrown away
in all three places:
(void)cudaSetDevice(index);When the switch fails, for example because the requested device index does not
exist on this machine, the work still went ahead on whatever device happened to
be current.
allocatethen returned success with a pointer that lives on adifferent device than the caller asked for. The caller stores the requested
index next to that pointer, so the pointer and its recorded device disagree
from then on, and the mistake surfaces much later as a wrong result or an
unrelated CUDA error somewhere else.
On a machine with one GPU:
What changes:
allocatelogs the CUDA error and returnsError::Internalinstead ofallocating on the wrong device. A device that cannot be selected is not an
out-of-memory condition, and
Error::Internalis whatException.hmapsevery CUDA error to and what
cuda_mutable_state.cppalready returns forthis same call.
return
Error::Internalinstead of copying against the wrong device.cudaGetDevicefails, instead of carrying onwithout switching. Without the current device there is no way to switch to the
requested one or restore afterwards, so the work would run against whatever
device happened to be current. This is the branch the tests cannot reach: on a
working CUDA host
cudaGetDevicedoes not fail, so it is verified byinspection rather than by the new tests.
deallocatereturnsvoid, andcudaFreeaccepts a pointer from any deviceunder unified addressing, so it logs the error and still frees the pointer
rather than leaking it. The failure is no longer silent.
Only the switch to the requested device is checked. The switch back to the
previous device stays best effort, because there is nothing useful to do if
restoring fails.
Test plan
Three new tests in
backends/cuda/runtime/test/test_cuda_allocator.cppask fordevice index
device_count, which is one past the last valid ordinal on anymachine, so the switch always fails:
AllocateOnMissingDeviceFailsCopyHostToDeviceOnMissingDeviceFailsCopyDeviceToHostOnMissingDeviceFailsDeviceIndexisint8_t, so each of the three first checks that one-past-the-lastordinal still fits in it: at 128 or more visible GPUs the count wraps to a negative
index, or at 256 back onto real device 0, and neither is a valid absent ordinal. The
check is per-test rather than in
SetUpso it does not also skip the six device-0tests. 127 devices still runs; 128 skips.
Built and ran the suite on one NVIDIA H100.
With this change:
Confirmed the tests really catch the bug. With the allocator reverted to the
old code and the same three tests:
The allocate failure shows the old behavior directly: the call reported success
for a device that does not exist.
The message the fixed build logs:
The six tests already in that file still pass, so the change does not disturb
the normal single device path.
clang-formatreports no changes needed on either touched file.Not covered
The
cudaGetDevicefailure branch is not covered by these tests. On a working CUDAhost that call does not fail, and there is no injection point in the suite, so
reverting just that branch leaves all three tests passing. Reaching it would need a
function-pointer seam or an injected fake CUDA runtime, both larger than the guard
they would protect.
deallocateonly logs, so no test asserts on it. Everything here ran on asingle GPU machine, so the switch always failed with an invalid ordinal. A
device that exists but is unavailable, for example one in exclusive compute
mode, takes the same code path but was not exercised.