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

shoumikhin · 2026-08-22T22:52:50Z

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

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 already handles this for the C++ Module API. It
builds planned memory per method when that method needs device buffers, and
shares one set of arenas across methods only when every buffer is host memory.
This change makes pybindings do the same thing instead of inventing a second
approach.

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.

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

New unit test, test_program_loads_when_one_method_is_device_planned in
extension/pybindings/test/test_pybindings.py. It exports a two method program
where one method is planned onto CUDA and the other onto the host, then checks
that the program loads, that the host method runs, and that it still runs after
an attempt to load the device method. It first asserts that the exported program
really does carry a CUDA-tagged planned buffer, so the test cannot quietly
degrade into a plain multi-method test if planning stops tagging devices. It
needs no GPU and runs in the existing CPU-only pytest job, which is where this
class of regression was previously invisible.

Device execution was measured on Linux x86_64 with an NVIDIA A100 (compute
capability 8.0), CUDA 12.8, Python 3.12. The check was done by swapping the
built extension module in and out, so both states are measured on the same
machine with the same model files.

Case	Before	After
`_load_for_executorch`, the loader that already worked	pass, max diff 7.45e-08	pass, max diff 7.45e-08
`Runtime.load_program`, the reported failure	fails, error 0x12	pass, max diff 7.45e-08
CPU-only program, no device tags, no `.ptd`	pass, max diff 2.98e-08	pass, max diff 2.98e-08
Two methods sharing planned buffers	not run	pass, max diff 1.19e-07 and 2.38e-07

The fixed path produces the same numbers as the loader that already worked, so
the result is correct and not merely non-crashing.

Negative control, confirming the test can actually fail: reverting to the
unpatched module reproduces not backed by CUDA device memory and a non-zero
exit, and reapplying the change returns it to a pass.

Reproducer, a small Linear plus ReLU model exported with the CUDA backend at
default settings, weights written with write_tensor_data_to_file and passed
back as data_path:

et = to_edge_transform_and_lower(
    export(model, example), partitioner=[CudaPartitioner([spec])]
).to_executorch()
et.write_tensor_data_to_file(out_dir)

# fails before this change, passes after
Runtime.get().load_program(pte_path, data_path=ptd_path) \
    .load_method("forward").execute([*example])

Not covered

The A100 numbers above were measured before planned memory was split per
method. The split changes when a device-planned method allocates, not what it
allocates, and the two loaders still share the same allocation code, but the
table has not been re-measured against the current revision.
Device execution was verified on x86_64 with an A100 only. Not verified on
aarch64 or on Jetson devices.
The new unit test exercises the host side of a mixed program. It cannot check
that the device method loads, because CPU-only CI has no device allocator.

pytorch-bot · 2026-08-22T22:52:53Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

⏳ No Failures, 2 Pending

As of commit c7458ee with merge base cff6f4d ():
💚 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.

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-22T22:58:48Z

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.

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 exactly as before. This is the same split extension/module/module.cpp already makes. 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 with one method planned onto CUDA and the other onto the host, then checks that the program loads, that the host method runs, and that it still runs after an attempt to load the device method. It needs no GPU.

…am (#22058) # 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. ```python # 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. 1. `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. 2. 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. 3. 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. 4. The caller reads the device for each buffer through `MethodMeta::memory_planned_buffer_device`. 5. 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. --------- Co-authored-by: Anthony Shoumikhin <shoumikhin@users.noreply.github.com> Co-authored-by: r <r@e>

shoumikhin requested review from JacobSzwejbka and larryliu0820 as code owners August 22, 2026 22:52

Copilot AI lite review requested due to automatic review settings August 22, 2026 22:52

shoumikhin temporarily deployed to cadence August 22, 2026 22:52 — 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 22, 2026

Copilot AI reviewed Aug 22, 2026

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

shoumikhin force-pushed the fix-pybindings-planned-buffer-device branch from 2825ea8 to c7458ee Compare August 22, 2026 22:57

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

Copilot AI reviewed Aug 22, 2026

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

shoumikhin closed this Aug 23, 2026

shoumikhin force-pushed the fix-pybindings-planned-buffer-device branch from c7458ee to 2052786 Compare August 23, 2026 01:01

shoumikhin requested review from abhinaykukkadapu, digantdesai, kirklandsign, psiddh, rascani and robert-kalmar as code owners August 23, 2026 01:01

github-actions Bot added ciflow/trunk module: arm Issues related to arm backend labels Aug 23, 2026

shoumikhin mentioned this pull request Aug 23, 2026

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

Merged

Navigation Menu

Uh oh!

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

Fix device placement for memory-planned buffers in Runtime.load_program#22057
shoumikhin wants to merge 1 commit into
mainfrom
fix-pybindings-planned-buffer-device

shoumikhin commented Aug 22, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Aug 22, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

github-actions Bot commented Aug 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

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

Uh oh!

The problem

Why it happens

The fix

Existing programs are unaffected

Test plan

Not covered

Uh oh!

pytorch-bot Bot commented Aug 22, 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/22057

⏳ No Failures, 2 Pending

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!

github-actions Bot commented Aug 22, 2026

This PR needs a release notes: label

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

shoumikhin commented Aug 22, 2026 •

edited

Loading

pytorch-bot Bot commented Aug 22, 2026 •

edited

Loading

This PR needs a `release notes:` label