Copy a tensor back from the device before ETDump writes it - #22095
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22095
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 16 PendingAs of commit 74995b7 with merge base 381a30c ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
shoumikhin
force-pushed
the
fix-etdump-device-tensor
branch
from
August 24, 2026 19:00
4dacc67 to
307b649
Compare
|
shoumikhin
force-pushed
the
fix-etdump-device-tensor
branch
from
August 24, 2026 21:59
61cce05 to
9d41702
Compare
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: 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, 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. Not covered: The ATen mode branch of the device type conversion only compiles. There is no ATen mode CMake build to run it in, and an ATen mode tensor in ExecuTorch carries no device metadata today, so the branch has no caller that can reach it.
The re2-guarded suite links neither re2 nor etdump_filter.cpp, so it cannot link where re2::re2 does exist. That predates this change, which only narrows when it is attempted; say so rather than implying the branch is merely rare. Also record why the two new tests do not call executorch_target_link_shared_runtime: the root CMakeLists adds devtools before it defines executorch_shared, so the target does not exist at that point. They run under Buck and the default static CMake build, which have a single registry.
shoumikhin
force-pushed
the
fix-etdump-device-tensor
branch
from
August 24, 2026 22:58
9d41702 to
80bda29
Compare
Gasoonjia
reviewed
Aug 25, 2026
Gasoonjia
reviewed
Aug 25, 2026
…copy helper Two changes from review feedback, no behavior change. Reuse the existing device mapping. to_runtime_device_type had its own CPU/CUDA switch in the USE_ATEN_LIB branch; that is exactly what extension/aten_util/aten_bridge.h's torch_to_executorch_device already does, so call it instead of duplicating the mapping. This pulls aten_bridge in as an aten-only dependency of etdump_flatcc; the non-ATen build and the CMake build, which never compile this branch, are unaffected. Stop exposing the device-copy path. write_device_tensor_or_return_error was a private member declared in the header. Nothing outside the class uses it, so move it into the anonymous namespace as a free function that takes the DataSinkBase it needs, and drop the declaration from etdump_flatcc.h. Only write_tensor_or_return_error remains, which is the one entry point callers use; it routes device tensors to the helper as before.
Gasoonjia
approved these changes
Aug 25, 2026
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
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 bufferto 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
NotFoundandlogs 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.cppregisters the existingMockCudaAllocator, which backs its device memory with host memory, and has twotests. 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.cppcovers the case wherenothing 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.txtwas not referenced by any parentCMakeLists.txt, so nothing in that directory was built by CMake. This addsadd_subdirectory(tests)todevtools/etdump/CMakeLists.txtunderBUILD_TESTING, so both new tests are picked up byctest, which is how the C++tests run.
The pre-existing
sdk_etdump_teststarget stays out of the CMake build. Itcompiles
etdump_test.cpp, which includesetdump_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:With
etdump_flatcc.cppreverted to the old code and everything rebuilt, bothfail:
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:Checked that
etdump_flatcc.cppstill compiles with-DUSE_ATEN_LIB.clang-formatreports 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::writewouldmemcpydevicememory 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.cppreads the serializeddevice 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.
LogTensorOnCpuDoesNotStageThroughTheAllocatorpasses with the production changereverted, 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.