Return an error instead of aborting when device memory cannot be allocated - #22085
Merged
Conversation
shoumikhin
requested review from
JacobSzwejbka and
kirklandsign
as code owners
August 24, 2026 15:57
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22085
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 156 PendingAs of commit 820c4e2 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
|
JacobSzwejbka
approved these changes
Aug 24, 2026
…cated A model can ask for its memory-planned buffers to live on an accelerator instead of on the CPU. When Module loads such a model it asks the runtime for a block of memory on that device. That request can fail for two ordinary reasons that have nothing to do with the model file: no allocator has been registered for that device type in this build, or the device is out of memory. Both were checked with ET_CHECK_MSG, which terminates the whole process. A Python caller got a SIGABRT and a core dump with no traceback and no chance to fall back to the CPU. These are properties of the machine, not of the program, so load_method now returns the error to the caller and the process survives. The same function also ignored the error from two MethodMeta lookups, and the loop that decides whether a model uses device buffers at all discarded a failed device query, which then read as "this buffer is on the CPU". That last path is not reachable today, because the loop bounds keep the index in range and that is the only case the query rejects, but silently treating a failed query as CPU would hand a backend host memory, so it now reports the error too. Test plan: - Added ModuleDeviceMemoryTest.DeviceAllocationFailureIsReportedNotFatal. It makes the test allocator refuse the request, then checks that load_method returns MemoryAllocationFailed, that the method is not left half loaded, and that a later attempt with the device healthy allocates normally. - Verified the test catches the bug: with the fix reverted the test binary exits with signal 6 (abort) and prints no result. With the fix, all 6 tests in the suite pass. - module_device_memory_test.cpp was only registered in the internal build, so no open source job ran it. Added it to extension/module/test/CMakeLists.txt along with the model file it needs. The full extension_module_test binary now runs 57 tests and all pass. - clang-format 18.1.3, cmake-format 0.6.13 and cmake-lint all report no changes on the touched files.
The device allocator registry is a process wide static. Registering a
second allocator for the same device type aborts the process on
purpose. GoogleTest calls SetUpTestSuite again for every repeat
iteration, so running this suite with --gtest_repeat=2 aborted instead
of passing.
Check the registry first and register only when nothing is there yet.
This is what runtime/core/test/device_allocator_test.cpp and
runtime/core/test/device_memory_buffer_test.cpp already do.
Test plan:
ctest -R '^extension_module_test$'
1/1 Test #54: extension_module_test ... Passed
extension_module_test --gtest_repeat=3
[ PASSED ] 57 tests. three times, exit 0
With the check removed again, the same repeat command aborts with
exit 134, so the test really covers this.
The comment claimed the forced allocator failure also stands in for a
device whose allocator was never registered. It does not. Two different
errors are possible here. A missing allocator makes
DeviceMemoryBuffer::create return Error::NotFound, while a registered
allocator that refuses the request returns
Error::MemoryAllocationFailed. This test only reaches the second one.
Comment only, no behavior change.
Test plan:
ctest -R '^extension_module_test$'
1/1 Test #54: extension_module_test ... Passed
…gnal The callee's only error return is a range check that this loop's bound makes unreachable, and broken metadata yields a CPU device rather than an error. Describe what the handling is actually for.
shoumikhin
force-pushed
the
fix-module-device-alloc-abort
branch
from
August 24, 2026 22:58
1894ed7 to
820c4e2
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.
What is going on
A model can be exported so that some of its memory-planned buffers live on an
accelerator instead of on the CPU. When
Moduleloads such a model, it asks theruntime for a block of memory on that device.
That request can fail for two ordinary reasons that say nothing about the model
file being loaded:
Both were checked with
ET_CHECK_MSG, which terminates the whole process. Soloading a device-annotated model on a machine without the matching backend, or
on a busy GPU, killed the process outright. Through the Python bindings
(
_load_for_executorch_from_bufferand friends build one of theseModuleobjects) that shows up as an abort and a core dump: no traceback, no error to
catch, no chance to fall back to the CPU.
Neither condition means the program is broken. They describe the machine it is
running on, so
load_methodnow returns the error to the caller and theprocess stays alive.
Before this change, both reasons ended the same way:
After, the two reasons stay apart, because
DeviceMemoryBuffer::createalreadyreports them with different codes:
Also in this change
Two other
MethodMetalookups in the same function, one for the buffer sizeand one for the buffer device, also used
ET_CHECK_MSG. They return the errornow as well, so all three abort sites in that function are gone.
Separately, the loop that decides whether a model uses device buffers at all
threw away a failed device query, which then read as "this buffer is on the
CPU". It returns the error now instead of ignoring it.
That last path is not reachable today. The only case the query rejects is an
out-of-range buffer index, and the loop bounds already keep the index in range.
It is fixed for consistency, not because it can fire. Quietly treating a failed
query as CPU would hand a backend host memory, which is a far harder failure to
debug than an error return.
Test coverage
extension/module/test/module_device_memory_test.cppwas registered in theinternal build only, so no open source job ran any of it. This change adds it
to
extension/module/test/CMakeLists.txt, along with the model file it needs,so it now builds and runs as part of
extension_module_test.Test plan
ModuleDeviceMemoryTest.DeviceAllocationFailureIsReportedNotFatal.It makes the test allocator refuse the request, then checks that
load_methodreturnsMemoryAllocationFailed, that the method is not lefthalf loaded, and that a later attempt with the device healthy allocates
normally.
binary exits with signal 6 (abort) and prints no result at all. With the fix,
all 6 tests in the suite pass.
extension_module_testbinary: 57 tests, all pass. The 51pre-existing tests are unaffected by adding the device test file to the same
binary.
extension_module_test --gtest_repeat=3: 57 tests pass on everyiteration, exit 0. The suite registers its allocator only once, so process
reuse is safe. Put the unconditional registration back and the same command
aborts with exit 134.
clang-formatreports no changes needed on the touched source file.cmake-format0.6.13 andcmake-lintreport no changes needed on thetouched
CMakeLists.txt.Not covered
Nothing here runs against real accelerator hardware. The failure is driven
through the existing mock allocator, which stands in for a device that is out
of memory. The committed test therefore covers the
Error::MemoryAllocationFailedbranch only. TheError::NotFoundbranch, fora device type with no registered allocator, is not exercised by a committed
test.