Fix device placement for memory-planned buffers in Runtime.load_program - #22058
Open
shoumikhin wants to merge 13 commits into
Open
Fix device placement for memory-planned buffers in Runtime.load_program#22058shoumikhin wants to merge 13 commits into
shoumikhin wants to merge 13 commits into
Conversation
shoumikhin
requested review from
JacobSzwejbka and
larryliu0820
as code owners
August 23, 2026 01:08
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22058
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 31 Pending, 1 Unrelated FailureAs of commit 3862c42 with merge base 85ec16e ( NEW FAILURES - The following jobs have failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
meta-cla
Bot
added
the
CLA Signed
shoumikhin
force-pushed
the
fix-pybindings-planned-buffer-device
branch
from
August 23, 2026 03:08
a1dd903 to
386752e
Compare
A .pte file records where each memory-planned buffer has to live, on the host or on an accelerator, in the plan's non_const_buffer_device field. ProgramMemory never read that field and allocated every planned buffer as host memory, so a program asking for device memory received a host pointer. With the CUDA backend this failed at execute time with "not backed by CUDA device memory", while the same .pte loaded through _load_for_executorch worked. Allocate each planned buffer on the device it is tagged for. Buffer indices are plan local, so one set of arenas shared across methods by index cannot describe a file where one method plans onto the host and another onto an accelerator. A method with any device-tagged buffer therefore gets its own arenas, and every other method keeps using the shared host arenas. extension/module/module.cpp faces the same problem and answers it differently, because it has a share_memory_arenas flag and this loader has none. Module refuses to load a device-planned method when that flag is set, and otherwise builds per-method arenas for every method. Runtime.load_program has no such flag and shares arenas unconditionally today, so refusing would break loading files that already load. It keeps the shared host arenas for host methods and gives a device-planned method its own instead. Those arenas are built when the method is loaded rather than when the program is loaded, so a file containing one accelerator method still opens on a machine without that accelerator, still lists its methods, and its host methods still load and run. non_const_buffer_device is optional and MethodMeta reports CPU when it is absent, so CPU-only and older programs keep the shared arenas, the host allocation path, and the single argument HierarchicalAllocator that leaves MemoryManager::has_device_memory() false. Adds test_program_loads_when_one_method_is_device_planned, which exports a two method program where one method has a device tagged planned buffer and the other has none, then checks that the program loads, that the host method runs, that loading the device method reaches the device allocator and is refused there, and that the host method still runs afterwards. It needs no GPU, and it skips itself on a build that links the CUDA backend, since that registers a CUDA allocator which would satisfy the request. Verified that the test fails when the change to pybindings.cpp is reverted: the old loader put the device buffer in host memory, so the method loaded and no error was raised.
shoumikhin
force-pushed
the
fix-pybindings-planned-buffer-device
branch
from
August 23, 2026 04:05
386752e to
58cf1b1
Compare
Program.load_method does not treat every method the same way. A method whose memory planned buffers are all on the host reuses one set of arenas shared with the other host only methods of the same program, so running a second such method can overwrite the intermediates and outputs of the first. A method with a buffer placed on an accelerator gets its own arenas instead. Nothing said so. Say it in the docstring users read.
Device allocation is the only step here that can fail, for example when no allocator is registered for the requested device or when the device is out of memory. It ran last, so every host arena of that method was already allocated and zero filled before the failure was reported, and all of it was then discarded. Swap the two members so the step that can fail runs first. Member initialization follows declaration order, so the order is what makes this work and a reorder would silently undo it. Say that in a comment.
ProgramMemory hands one span per planned buffer and one device tag per planned buffer to HierarchicalAllocator, which checks that the two counts agree. That check aborts the process. In a Python extension that means the interpreter dies with no traceback and nothing the caller can catch. Both constructors keep the counts equal today, so this is not reachable, but the invariant lives only in the callers. Check it where the error can still become a Python exception.
load_method walked every planned buffer once through has_device_buffers to decide whether the method needs its own arenas, then walked them all again through make_method_memory to collect the sizes and devices. Each device lookup scans the program's sparse device list, so the second walk repeats work the first one already did. Have make_method_memory return nullptr when every buffer is on the host. One pass then answers both questions. has_device_buffers stays because the program constructor still needs only the answer, not the sizes.
pybindings.cpp includes runtime/core/device_memory_buffer.h but the pybindings targets did not depend on it. The build worked only because the header came in transitively through extension/module, and the rule opts out of automatic dependency checking, so nothing would report the omission.
Three comments describe the code inaccurately. The first says device buffers are allocated first because they hold the only allocation that can fail. Allocating a host arena can also fail, so the real reason is only that a device failure should happen before the host arenas are allocated and zero filled. The second implies the buffer and device count guard converts a reachable abort. Both vectors are filled in lockstep today, so the guard only fires if a future caller breaks that. The third says a device planned method inflates arenas nobody will read. Every host only method reads those arenas. What that method would inflate is their size.
Result::get() aborts the process if the Result holds an error, so a corrupt program file that reports a bad planned buffer size kills the interpreter instead of raising. This line was already rewritten by an earlier commit in this change, and the device buffer code added next to it already checks the same call, so the two now behave the same way.
The docstring said a second host only method may overwrite the outputs of the first. It cannot: outputs are copied out on every call, so what sharing affects is the intermediate values, not what execute returns. It also said a failed device allocation affects only that method. There is no unload, so whatever a device method does claim stays claimed for every method loaded afterwards. Both points are now stated as they are.
The test file imports DeviceType from executorch.exir.schema and executorch.exir.backend.test.device_util, and the three test targets picked both up only through make_test. A direct import needs a direct dependency, so the targets no longer break if make_test drops either one.
The existing test covers the case where no device allocator is registered and the load is refused. Nothing covered the case the change exists to fix: a build that does have a device allocator, where the arena has to come off the device. The new test lowers one of two methods to the CUDA backend, checks that free device memory drops by the planned buffer size when that method is loaded, checks that both methods return the right values, and checks that the memory comes back when the program is released. On the previous behavior the arena stayed in host memory and the backend rejected the pointer with an execute error, so this test fails without the fix. A real accelerator is required, so the test skips unless one is present. The CUDA build job now runs it, and the CUDA workflow now triggers on changes under extension/pybindings, which is where this code lives.
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.
Fix device placement for memory-planned buffers in Runtime.load_program
Replaces #22057. That pull request was force-pushed to a commit with no
history in common with main, which made GitHub close it permanently. Same
branch, same change, correct history.
The problem
ExecuTorch has two ways to load a model from Python. With the CUDA backend, one
of them works and the other crashes. Same
.ptefile, same.ptdweights file,same default export settings. Only the loader differs.
The crash looks like this:
Why it happens
A
.ptefile records where each memory-planned buffer has to live, on the hostor on an accelerator. That is the
non_const_buffer_devicefield in the plan.ProgramMemoryinextension/pybindings/pybindings.cppnever read that field.It allocated every planned buffer as host memory (
std::vector<uint8_t>). So aprogram that asked for device memory received a host pointer. The CUDA backend
checked the pointer, saw it was not device memory, and refused to run.
The fix
extension/module/module.cpphits the same problem in the C++ModuleAPI andanswers it differently, because it has a
share_memory_arenasflag and thisloader has none.
Modulerefuses to load a device-planned method when that flagis set, and otherwise builds per-method arenas for every method.
Runtime.load_programshares arenas unconditionally today and offers no way toturn that off, so refusing would stop files loading that load now. It keeps the
shared host arenas for host methods and gives a device-planned method its own.
ProgramMemorynow receives the per-buffer device list alongside the sizes,and allocates each buffer on the device it is tagged for. Host-tagged buffers
keep using
std::vector<uint8_t>, exactly as before.buffer 0 of another, so one set of arenas shared by index cannot describe a
file where one method plans onto the host and another onto an accelerator. A
method with any device-tagged buffer therefore gets its own arenas, and every
other method keeps using the shared host arenas.
load_method, not when the program is loaded. Afile containing one accelerator method still opens on a machine without that
accelerator, still lists its methods, and its host methods still load and
run. Only loading the accelerator method fails, and it fails naming the
buffer and the device it could not allocate.
MethodMeta::memory_planned_buffer_device.forwarders, so behavior is unchanged.
Program.load_methodinruntime/__init__.pydocuments all of this, includingthe part that is not new: two host-only methods of the same program share one
set of arenas and therefore overwrite each other's intermediate values.
Existing programs are unaffected
non_const_buffer_deviceis optional.MethodMeta::memory_planned_buffer_devicereturns
Device{CPU, 0}when the field is absent, which is the case forCPU-only programs and for
.ptefiles produced before the field existed. Such aprogram keeps the shared arenas, the host allocation path, and the single
argument
HierarchicalAllocator, soMemoryManager::has_device_memory()staysfalse for it as that constructor documents.
Test plan
Two tests in
extension/pybindings/test/test_pybindings.py.test_program_loads_when_one_method_is_device_plannedcovers the refusalpath and needs no GPU, so it runs in the existing CPU-only job. It exports a two
method program where one method has a device-tagged planned buffer and the other
has none, then checks that the program loads, that the host method runs, that
loading the device method reaches the device allocator and is refused there, and
that the host method still runs afterwards. It first asserts that the exported
program really does carry a CUDA-tagged planned buffer, so it cannot quietly
degrade into a plain multi-method test if planning stops tagging devices. It
skips itself on a build that links the CUDA backend, because that registers a
CUDA allocator at static init, the registry has no way to drop one, and the
request would then be satisfied with or without this change.
test_device_planned_method_allocates_on_the_devicecovers what the refusalis protecting: on a build that does have a device allocator, the arena has to
come off the device rather than out of host memory. It builds a two method
program where one method is lowered to the CUDA backend and the other is not,
then measures free device memory with
torch.cuda.mem_get_infoaround eachload_methodcall. It asserts that loading the host method takes no devicememory, that loading the device method takes at least 90 percent of the planned
device bytes, that both methods return correct numbers, that running the host
method in between does not disturb the device method, and that releasing the
program returns the memory. It skips unless the build links the CUDA backend and
a device is visible.
That second test can only run where a device allocator is registered, so this
pull request also wires it into the job that has one.
.ci/scripts/test-cuda-build.shruns it, and
.github/workflows/cuda.ymlnow triggers on changes underextension/pybindings/and toruntime/__init__.py. Before this, no job in therepository built the Python extension with a device allocator and then ran the
pybindings tests, which is exactly why this class of defect was invisible.
Measured
Linux x86_64, NVIDIA A100 80GB, compute capability 8.0, CUDA 13.0, Python 3.12,
torch 2.13.0. Two builds from one source tree, one CPU only and one with the
CUDA backend. The before column is the merge base with
main, produced byswapping only
pybindings.cppand rebuilding, so both columns are the samemachine, the same model files and the same everything else.
Two methods in one program,
forwardon the host andforward2lowered toCUDA, planned device bytes 50331648 (48 MiB):
load_programThe before column is not a generic failure. The CUDA backend names the defect
itself:
Suites, on the same two builds:
extension/pybindings/test/test_pybindings.py-k devicetest_device_planned_method_allocates_on_the_devicealone, run the way the CI script runs itThe two failures are
test_method_quantized_opsandtest_quantized_ops. Theyare pre-existing and unrelated: they need the quantized AOT library preloaded,
which the Buck target does and a bare
pytestinvocation does not. Theyreproduce identically at the merge base.
Linux aarch64, Jetson Orin Nano, Python 3.10, CPU-only build: identical counts to
the x86_64 CPU-only column above, including the device tests and the same two
pre-existing quantized-op failures. The CUDA success path is not reachable on
that board, because its GPU needs a PyTorch build pinned to a different version
than this repository requires, so only the host paths and the refusal path are
covered there.
Not covered
by any test. Delete that skip and both tests still pass, because nothing in
Python can observe the shared arena sizes. What it saves is host memory that
nothing reads, which grows with the model, so it is worth a C++ test later.
capability 8.0. Not measured on a Jetson board or on any non-CUDA
accelerator.
has_device_buffersandmake_method_memoryaskMethodMetafor one bufferat a time, and
MethodMeta::memory_planned_buffer_devicescans the sparsedevice list on each call, so the cost is the buffer count times the device
entry count. Real programs measured here have 2 or 3 buffers and 1 device
entry, and
extension/module/module.cppalready reads the same metadata thesame way, but both counts come from the file. Removing the concern properly
means a bulk accessor on
MethodMeta, which would fix both callers at onceand belongs in its own change.