Fix device placement for memory-planned buffers in Runtime.load_program by shoumikhin · Pull Request #22058 · pytorch/executorch

shoumikhin · 2026-08-23T01:07:59Z

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 .pte file, same .ptd weights file,
same default export settings. Only the loader differs.

# works
_load_for_executorch(pte_path, ptd_path).forward([x])

# crashes
Runtime.get().load_program(pte_path, data_path=ptd_path) \
    .load_method("forward").execute([x])

The crash looks like this:

[cuda_backend.cpp:548] Tensor 0 has device_type=CUDA but its data pointer
0x... is not backed by CUDA device memory
(cudaPointerGetAttributes err=0, cudaMemoryType=0).
RuntimeError: method->execute() failed with error 0x12

Why it happens

A .pte file records where each memory-planned buffer has to live, on the host
or on an accelerator. That is the non_const_buffer_device field in the plan.

ProgramMemory in extension/pybindings/pybindings.cpp never read that field.
It allocated every planned buffer as host memory (std::vector<uint8_t>). So a
program 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.

.pte says:  buffer 0 -> CUDA:0
before:     buffer 0 -> std::vector<uint8_t>        (host)   -> backend rejects
after:      buffer 0 -> DeviceMemoryBuffer::create  (device) -> backend accepts

The fix

extension/module/module.cpp hits the same problem in the C++ Module API 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 shares arenas unconditionally today and offers no way to
turn 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.

ProgramMemory now 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 indices are plan local. Buffer 0 of one method has nothing to do with
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.
Those arenas are built in load_method, not when the program is loaded. 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. Only loading the accelerator method fails, and it fails naming the
buffer and the device it could not allocate.
The caller reads the device for each buffer through
MethodMeta::memory_planned_buffer_device.
Two deprecated calls were replaced with their current names. Both are plain
forwarders, so behavior is unchanged.

Program.load_method in runtime/__init__.py documents all of this, including
the 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_device is optional. MethodMeta::memory_planned_buffer_device
returns Device{CPU, 0} when the field is absent, which is the case for
CPU-only programs and for .pte files produced before the field existed. Such a
program keeps the shared arenas, the host allocation path, and the single
argument HierarchicalAllocator, so MemoryManager::has_device_memory() stays
false 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_planned covers the refusal
path 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_device covers what the refusal
is 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_info around each
load_method call. It asserts that loading the host method takes no device
memory, 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 the memory is
returned once the methods and the program are all dropped. The test deletes both
methods before releasing the program: each method owns its own memory, so a method
kept alive after the program is released still holds its allocation. 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. The unittest-cuda job in
.github/workflows/cuda.yml runs it, right after the install that job already
does and before its builds, since the test needs nothing they produce. That
workflow now triggers on changes under extension/pybindings/ and to
runtime/__init__.py, and the job condition lists the same two paths so the job
actually fires on them. Before this, no job in the repository 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 by
swapping only pybindings.cpp and rebuilding, so both columns are the same
machine, the same model files and the same everything else.

Two methods in one program, forward on the host and forward2 lowered to
CUDA, planned device bytes 50331648 (48 MiB):

Measurement	Before	After
Device memory taken by `load_program`	0 MiB	0 MiB
Device memory taken by loading the host method	0 MiB	0 MiB
Device memory taken by loading the CUDA method	0 MiB	48 MiB
CUDA method produces correct numbers	no, error 0x12	yes
Host method produces correct numbers	yes	yes
CUDA method still correct after running the host method	not reached	yes
Device memory returned once the methods and program are dropped	nothing to return	48 MiB

The before column is not a generic failure. The CUDA backend names the defect
itself:

[cuda_backend.cpp:548] Tensor 0 has device_type=CUDA but its data pointer
0x7f5842fff010 is not backed by CUDA device memory
(cudaPointerGetAttributes err=0, cudaMemoryType=0).
[method.cpp:1528] CALL_DELEGATE execute failed at instruction 2: 0x12

Suites, on the same two builds:

Suite	CUDA build	CPU-only build
`extension/pybindings/test/test_pybindings.py`	39 passed, 1 skipped, 2 failed	39 passed, 1 skipped, 2 failed
the same file filtered to `-k device`	4 passed, refusal test skipped	4 passed, success test skipped
`test_device_planned_method_allocates_on_the_device` alone, run the way the CI script runs it	passed	skipped

The two failures are test_method_quantized_ops and test_quantized_ops. They
are pre-existing and unrelated: they need the quantized AOT library preloaded,
which the Buck target does and a bare pytest invocation does not. They
reproduce 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.

Landing order

This should land after or together with #22095. The device arenas added here are
allocated through the device allocator, and ETDump can be handed pointers into
them when a delegate logs its arguments. BufferDataSink::write does a plain host
memcpy, so recording a CUDA tensor would read device memory from the host.
#22095 is the fix for that path: it routes non-CPU tensors through a device copy
before writing. Landing this one first leaves that combination reachable whenever
event tracing is on.

Not covered

Leaving a device-planned method out of the shared host arenas is 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.
The device path is measured on one accelerator, an A100 with compute
capability 8.0. Not measured on a Jetson board or on any non-CUDA
accelerator.
has_device_buffers and make_method_memory ask MethodMeta for one buffer
at a time, and MethodMeta::memory_planned_buffer_device scans the sparse
device 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.cpp already reads the same metadata the
same way, but both counts come from the file. Removing the concern properly
means a bulk accessor on MethodMeta, which would fix both callers at once
and belongs in its own change.

pytorch-bot · 2026-08-23T01:08:03Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22058

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 175 Pending

As of commit 7b6ce88 with merge base 7d163c9 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

github-actions · 2026-08-23T01:08:42Z

This PR needs a `release notes:` label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

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.

The new test loads a method that the CUDA backend runs, so the process has to open the shared library the backend produces. That library needs a newer libstdc++ than the CI image keeps on its default search path. Every CUDA job in this workflow that runs a model already installs libstdcxx-ng and puts the conda library directory first on LD_LIBRARY_PATH. The one job that does not is test-cuda-builds, because until now it only built ExecuTorch and ran a plain torch matmul on the GPU, never a delegated method. Putting the test there made it fail on both CUDA versions with: libstdc++.so.6: version `GLIBCXX_3.4.29' not found Init failed for backend CudaBackend: 0x22 Move the test to unittest-cuda, the job for CUDA Python unit tests, which already has that setup. Run it right after the install and before the long builds, since it needs nothing they produce and failing early saves GPU time. Also add the two pybindings paths to that job's condition. They were added to the workflow path filter, but every job carries its own condition, so without this a change to pybindings alone would start the workflow and then skip every job in it. test-cuda-build.sh now matches main again, so this change no longer touches it. Test plan: checked the workflow file parses and the job count is unchanged, checked the shell script with bash -n and shellcheck and confirmed it is byte-identical to main, confirmed the test class and method names resolve, and evaluated the new job condition against sample file lists. The test itself needs a CUDA build of the pybindings that only CI produces, so CI is the check.

The loop counters added by this change used postfix increment. Postfix has to keep the old value around to return it, prefix does not, and for a plain size_t the two are identical after optimization. Prefix is the convention this file should follow, so use it consistently. Six loops in extension/pybindings/pybindings.cpp, all of them lines this change introduced. The three postfix loops that were already in the file before this change are left alone. Test plan: no behavior change. The counter type is size_t and the result of the increment expression is discarded in every one of the six loops, so prefix and postfix compute the same thing. Verified the file still compiles with the same command the build uses, and verified with a word level diff that the only tokens that moved are the six increments.

The docstring said the private arenas of a device-planned method are held until the program is released. That is the wrong owner. PyMethod keeps a shared_ptr to the ProgramMemory, so a caller who keeps the returned method alive and drops the program keeps the accelerator memory too. Releasing the program alone does not give it back. Say that the method owns the arenas, and add the reason the two usually look the same: this Program caches every method it loads, so the methods and the program normally die together. Documentation only, no behavior change. Test plan: black --check runtime/__init__.py -> unchanged python -m py_compile runtime/__init__.py -> ok

…he arena split unittest exits 0 when every test skips, so a missing CUDA backend or GPU would have read as a pass. Recover the status through tee with PIPESTATUS and treat a skip as an error. The pattern is anchored so a test whose name contains "skipped" does not trip it. Also document why a device-planned method gets its own arenas rather than joining the program-wide ones.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

### Summary ETDump records an intermediate tensor by handing the tensor's data pointer to a data sink, and every sink reads those bytes with a plain host read: ```cpp memcpy(cur_data_begin, ptr, length); ``` When the tensor lives on an accelerator that pointer is not host memory, so the read segfaults. A program placed on CUDA crashes as soon as tracing is turned on, which is exactly when someone is trying to debug it. Returning an error instead would not help. All four callers wrap the result in `ET_CHECK_MSG`, so an error aborts the process rather than skipping the tensor. This change brings the data back to host memory first. When the tensor is not on CPU, ETDump looks up the allocator registered for that device type, stages the bytes into a temporary host buffer with `copy_device_to_host`, writes that buffer to the sink and frees it. A tensor on CPU keeps the old path and copies nothing extra. If no allocator is registered for the device, ETDump now reports `NotFound` and logs the device type instead of reading the pointer anyway. ### Test plan Two new test files, each with a CMake target and a Buck target. `devtools/etdump/tests/etdump_device_test.cpp` registers the existing `MockCudaAllocator`, which backs its device memory with host memory, and has two tests. One logs a tensor tagged as CUDA and checks both that ETDump went through the allocator and that the bytes reached the debug buffer. The other logs a tensor on CPU and checks that the allocator was not used at all, so the CPU path is unchanged. `devtools/etdump/tests/etdump_device_no_allocator_test.cpp` covers the case where nothing is registered for the device. The registry is a process wide static with no way to remove an entry, so that case needs a binary that never registers anything, which is why it is a second file. `devtools/etdump/tests/CMakeLists.txt` was not referenced by any parent `CMakeLists.txt`, so nothing in that directory was built by CMake. This adds `add_subdirectory(tests)` to `devtools/etdump/CMakeLists.txt` under `BUILD_TESTING`, so both new tests are picked up by `ctest`, which is how the C++ tests run. The pre-existing `sdk_etdump_tests` target stays out of the CMake build. It compiles `etdump_test.cpp`, which includes `etdump_filter.h`, which needs re2, and the devtools build does not pull re2 in. It is now guarded on re2 being available rather than being silently unreachable. With this change both new binaries pass under `ctest`: ``` 1/2 Test #1: etdump_device_test ................ Passed 2/2 Test #2: etdump_device_no_allocator_test ... Passed ``` With `etdump_flatcc.cpp` reverted to the old code and everything rebuilt, both fail: ``` Expected equality of these values: g_mock_cuda.d2h_count_ Which is: 0 1 Death test: etdump_gen.log_evalue(EValue(tensor)) Result: failed to die. ``` Also reproduced the real crash on one NVIDIA H100, with a small program that allocates through `cudaMalloc`, tags a tensor as CUDA and logs it: ``` before: Segmentation fault (core dumped) after: debug buffer holds 1.5 2.5 3.5 4.5 ``` Checked that `etdump_flatcc.cpp` still compiles with `-DUSE_ATEN_LIB`. `clang-format` reports no changes needed on the four touched C++ files. ### Landing order #22058 adds device-planned arenas to the Python bindings that ETDump can be handed pointers into, and without this change `BufferDataSink::write` would `memcpy` device memory from the host. This should land before or together with it. ### Not covered The ATen mode branch of the device type conversion only compiles. There is no ATen mode CMake build to run it in, so nothing here executes it. It is reachable in principle: `runtime/executor/tensor_parser_aten.cpp` reads the serialized device type and index, including CUDA, and builds the ATen tensor on that device. An earlier version of this description said ATen tensors carry no device metadata, which is wrong. `LogTensorOnCpuDoesNotStageThroughTheAllocator` passes with the production change reverted, since CPU tensors already went straight to the data sink. It documents the CPU path rather than locking the fix; the other two tests are the ones that require the new branch. --------- Co-authored-by: r <r@e>

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

shoumikhin requested review from JacobSzwejbka and larryliu0820 as code owners August 23, 2026 01:08

Copilot AI lite review requested due to automatic review settings August 23, 2026 01:08

shoumikhin temporarily deployed to cadence August 23, 2026 01:08 — with GitHub Actions Inactive

Copilot AI reviewed Aug 23, 2026

shoumikhin temporarily deployed to cadence August 23, 2026 01:34 — with GitHub Actions Inactive

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 23, 2026

Copilot AI review requested due to automatic review settings August 23, 2026 03:08

shoumikhin force-pushed the fix-pybindings-planned-buffer-device branch from a1dd903 to 386752e Compare August 23, 2026 03:08

shoumikhin temporarily deployed to cadence August 23, 2026 03:08 — with GitHub Actions Inactive

Copilot AI reviewed Aug 23, 2026

shoumikhin temporarily deployed to cadence August 23, 2026 03:32 — with GitHub Actions Inactive

shoumikhin force-pushed the fix-pybindings-planned-buffer-device branch from 386752e to 58cf1b1 Compare August 23, 2026 04:05

Copilot AI review requested due to automatic review settings August 23, 2026 04:05

shoumikhin temporarily deployed to cadence August 23, 2026 04:05 — with GitHub Actions Inactive

Copilot AI reviewed Aug 23, 2026

shoumikhin temporarily deployed to cadence August 23, 2026 04:32 — with GitHub Actions Inactive

Copilot AI review requested due to automatic review settings August 24, 2026 02:50

Copilot AI reviewed Aug 24, 2026

shoumikhin temporarily deployed to cadence August 24, 2026 02:50 — with GitHub Actions Inactive

shoumikhin temporarily deployed to cadence August 24, 2026 03:19 — with GitHub Actions Inactive

Copilot AI review requested due to automatic review settings August 24, 2026 05:10

shoumikhin temporarily deployed to cadence August 24, 2026 05:10 — with GitHub Actions Inactive

shoumikhin and others added 9 commits August 24, 2026 15:57

Copilot AI review requested due to automatic review settings August 24, 2026 22:58

shoumikhin force-pushed the fix-pybindings-planned-buffer-device branch from 59c9394 to a3f584c Compare August 24, 2026 22:58

shoumikhin temporarily deployed to cadence August 24, 2026 22:58 — with GitHub Actions Inactive

Copilot AI reviewed Aug 24, 2026

shoumikhin temporarily deployed to cadence August 24, 2026 23:23 — with GitHub Actions Inactive

Merge branch 'main' into fix-pybindings-planned-buffer-device

be35a3d

Copilot AI review requested due to automatic review settings August 24, 2026 23:27

shoumikhin temporarily deployed to cadence August 24, 2026 23:27 — with GitHub Actions Inactive

Copilot AI reviewed Aug 24, 2026

shoumikhin temporarily deployed to cadence August 24, 2026 23:58 — with GitHub Actions Inactive

Merge branch 'main' into fix-pybindings-planned-buffer-device

06624e2

Copilot AI review requested due to automatic review settings August 25, 2026 01:03

Copilot AI reviewed Aug 25, 2026

shoumikhin temporarily deployed to cadence August 25, 2026 01:03 — with GitHub Actions Inactive

Merge branch 'main' into fix-pybindings-planned-buffer-device

7b6ce88

Copilot AI reviewed Aug 25, 2026

shoumikhin mentioned this pull request Aug 25, 2026

Enable fast ROCm Voxtral decode #22005

Merged

Navigation Menu

Uh oh!

Fix device placement for memory-planned buffers in Runtime.load_program - #22058

Fix device placement for memory-planned buffers in Runtime.load_program#22058
shoumikhin merged 18 commits into
mainfrom
fix-pybindings-planned-buffer-device

shoumikhin commented Aug 23, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Aug 23, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

github-actions Bot commented Aug 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

shoumikhin commented Aug 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Fix device placement for memory-planned buffers in Runtime.load_program

The problem

Why it happens

The fix

Existing programs are unaffected

Test plan

Measured

Landing order

Not covered

Uh oh!

pytorch-bot Bot commented Aug 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22058

⏳ No Failures, 175 Pending

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 23, 2026

This PR needs a release notes: label

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

shoumikhin commented Aug 23, 2026 •

edited

Loading

pytorch-bot Bot commented Aug 23, 2026 •

edited

Loading

This PR needs a `release notes:` label