Uh oh!
There was an error while loading. Please reload this page.
Fix device placement for memory-planned buffers in Runtime.load_program - #22057
Closed
shoumikhin wants to merge 1 commit into
Closed
Fix device placement for memory-planned buffers in Runtime.load_program#22057shoumikhin wants to merge 1 commit into
shoumikhin wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22057
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 2 PendingAs of commit c7458ee with merge base cff6f4d ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
shoumikhinforce-pushed
the
fix-pybindings-planned-buffer-device
branch
from
August 22, 2026 22:57
2825ea8 to
c7458eeCompareThis PR needs a |
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.
shoumikhinforce-pushed
the
fix-pybindings-planned-buffer-device
branch
from
August 23, 2026 01:01
c7458ee to
2052786Compareshoumikhin
requested review from
abhinaykukkadapu, digantdesai, kirklandsign, psiddh, rascani and robert-kalmar
as code ownersAugust 23, 2026 01:01
shoumikhin added a commit
that referenced
this pull request
Aug 25, 2026
…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>
Gasoonjia pushed a commit
that referenced
this pull request
Aug 25, 2026
…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>
Gasoonjia added a commit
that referenced
this pull request
Aug 25, 2026
…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>
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 freeto 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
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.cppalready handles this for the C++ModuleAPI. Itbuilds 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.
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.
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
New unit test,
test_program_loads_when_one_method_is_device_plannedinextension/pybindings/test/test_pybindings.py. It exports a two method programwhere 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.
_load_for_executorch, the loader that already workedRuntime.load_program, the reported failure.ptdThe 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 memoryand a non-zeroexit, and reapplying the change returns it to a pass.
Reproducer, a small
LinearplusReLUmodel exported with the CUDA backend atdefault settings, weights written with
write_tensor_data_to_fileand passedback as
data_path:Not covered
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.
aarch64 or on Jetson devices.
that the device method loads, because CPU-only CI has no device allocator.