Fix the crash when a program keeps its activations on the GPU - #21699
Closed
shoumikhin wants to merge 2 commits into
Closed
Fix the crash when a program keeps its activations on the GPU#21699shoumikhin wants to merge 2 commits into
shoumikhin wants to merge 2 commits into
Conversation
Contributor
Author
|
Stack from ghstack (oldest at bottom): |
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21699
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 124 PendingAs of commit 4fa7f46 with merge base 730b77a ( NEW FAILURE - The following job has failed:
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
added a commit
that referenced
this pull request
Aug 9, 2026
## The problem ExecuTorch can export a model two ways for a GPU. The default copies your input from main memory (the host) to the GPU for you. The other way skips those copies, so you hand over memory that is already on the GPU and the output stays there. That second way saves a round trip, which matters for a model in a loop. The second way crashes. It crashes from Python and from C++, immediately, before your data is touched: ``` [cuda_backend.cpp:429] Created new CUDA stream 0x10ae920 for method Segmentation fault (core dumped) ``` The reason is one line. When the runtime puts your input into the buffer the memory plan reserved, it always used a plain host copy: ```cpp std::memcpy(t_dst.mutable_data_ptr(), t_src.const_data_ptr(), t_src.nbytes()); ``` If that destination is GPU memory, a host `memcpy` into it is undefined and the process dies. The memory itself was allocated correctly on the GPU; only the copy assumed the host. ## The change Route the copy through the device that owns the memory instead of assuming the host. ExecuTorch already has an interface for this, and a backend registers an implementation when it is linked in, so the core runtime still names no GPU code and still builds for bare metal: ``` core (no GPU code) CUDA backend copy_between_devices() -> get_device_allocator(CUDA) ──────── register_device_allocator(...) -> allocator->copy_host_to_device() ──► cudaMemcpy ``` A host to host copy keeps using `memcpy` and does not need any device registered. One primitive was missing. The interface had host to device and device to host, but not device to device, which is exactly what this feature needs: your input is already on the GPU and the planned buffer is too. It is added with a default that returns `NotSupported`, so an allocator that cannot do it declines instead of being forced to implement it, and the CUDA one implements it. ## Test plan Added `device_copy_test.cpp`, which needs no GPU: the repository already has a mock allocator that copies with `memcpy` and counts calls, so the routing is observable. - a host to host copy still works and does not touch the allocator - host to device, device to host, and device to device each reach the matching call - zero bytes is accepted, which a tensor with an empty dimension relies on Before this change no test executed a program with device activations at all. The closest one loads the method and tolerates failure without running it, so the crash had no coverage. Also built the runtime and ran the existing runtime and extension tests on Linux x86_64. ## Known gaps, not fixed here Two problems on the same feature remain, and each deserves its own change: - The Python bindings build the source tensor without reading its device, so a GPU input is labelled as host memory. A caller passing GPU tensors from Python is still not correctly described to the runtime. - For a program whose output stays on the GPU, the Python return path reads that memory as if it were host memory. ghstack-source-id: 3d35054 ghstack-comment-id: 5233812152 Pull-Request: #21699
shoumikhin
added a commit
that referenced
this pull request
Aug 9, 2026
## The problem ExecuTorch can export a model two ways for a GPU. The default copies your input from main memory (the host) to the GPU for you. The other way skips those copies, so you hand over memory that is already on the GPU and the output stays there. That second way saves a round trip, which matters for a model in a loop. The second way crashes. It crashes from Python and from C++, immediately, before your data is touched: ``` [cuda_backend.cpp:429] Created new CUDA stream 0x10ae920 for method Segmentation fault (core dumped) ``` The reason is one line. When the runtime puts your input into the buffer the memory plan reserved, it always used a plain host copy: ```cpp std::memcpy(t_dst.mutable_data_ptr(), t_src.const_data_ptr(), t_src.nbytes()); ``` If that destination is GPU memory, a host `memcpy` into it is undefined and the process dies. The memory itself was allocated correctly on the GPU; only the copy assumed the host. ## The change Route the copy through the device that owns the memory instead of assuming the host. ExecuTorch already has an interface for this, and a backend registers an implementation when it is linked in, so the core runtime still names no GPU code and still builds for bare metal: ``` core (no GPU code) CUDA backend copy_between_devices() -> get_device_allocator(CUDA) ──────── register_device_allocator(...) -> allocator->copy_host_to_device() ──► cudaMemcpy ``` A host to host copy keeps using `memcpy` and does not need any device registered. One primitive was missing. The interface had host to device and device to host, but not device to device, which is exactly what this feature needs: your input is already on the GPU and the planned buffer is too. It is added with a default that returns `NotSupported`, so an allocator that cannot do it declines instead of being forced to implement it, and the CUDA one implements it. ## Test plan Added `device_copy_test.cpp`, which needs no GPU: the repository already has a mock allocator that copies with `memcpy` and counts calls, so the routing is observable. - a host to host copy still works and does not touch the allocator - host to device, device to host, and device to device each reach the matching call - zero bytes is accepted, which a tensor with an empty dimension relies on Before this change no test executed a program with device activations at all. The closest one loads the method and tolerates failure without running it, so the crash had no coverage. Also built the runtime and ran the existing runtime and extension tests on Linux x86_64. ## Known gaps, not fixed here Two problems on the same feature remain, and each deserves its own change: - The Python bindings build the source tensor without reading its device, so a GPU input is labelled as host memory. A caller passing GPU tensors from Python is still not correctly described to the runtime. - For a program whose output stays on the GPU, the Python return path reads that memory as if it were host memory. ghstack-source-id: 3d35054 ghstack-comment-id: 5233812152 Pull-Request: #21699
Contributor
Author
|
Closing this: the premise is wrong. This change made the input copy device aware, on the assumption that a program exported to keep its activations on the GPU has to copy the caller's input into a buffer the runtime reserved. It does not, and it should not. An ExecuTorch tensor refers to memory rather than owning it, and the runtime already has a path that points at the caller's buffer with no copy at all. The crash I was chasing came from my own export configuration. Memory planning allocates graph inputs by default, so the runtime reserved a GPU buffer for an input the caller already owned, and then filled it with a host memcpy. Exporting with Verified from C++ with GPU memory in and out, matching eager PyTorch exactly on two GPU architectures. Copying there would defeat the purpose of the feature, which its own documentation states: the caller provides tensors already on the target device. Two smaller changes replace this, each doing one thing:
|
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.
The problem
ExecuTorch can export a model two ways for a GPU. The default copies your input from
main memory (the host) to the GPU for you. The other way skips those copies, so you
hand over memory that is already on the GPU and the output stays there. That second
way saves a round trip, which matters for a model in a loop.
The second way crashes. It crashes from Python and from C++, immediately, before your
data is touched:
The reason is one line. When the runtime puts your input into the buffer the memory
plan reserved, it always used a plain host copy:
std::memcpy(t_dst.mutable_data_ptr(), t_src.const_data_ptr(), t_src.nbytes());If that destination is GPU memory, a host
memcpyinto it is undefined and the processdies. The memory itself was allocated correctly on the GPU; only the copy assumed the
host.
The change
Route the copy through the device that owns the memory instead of assuming the host.
ExecuTorch already has an interface for this, and a backend registers an implementation
when it is linked in, so the core runtime still names no GPU code and still builds for
bare metal:
A host to host copy keeps using
memcpyand does not need any device registered.One primitive was missing. The interface had host to device and device to host, but not
device to device, which is exactly what this feature needs: your input is already on the
GPU and the planned buffer is too. It is added with a default that returns
NotSupported, so an allocator that cannot do it declines instead of being forced toimplement it, and the CUDA one implements it.
Test plan
Added
device_copy_test.cpp, which needs no GPU: the repository already has a mockallocator that copies with
memcpyand counts calls, so the routing is observable.Before this change no test executed a program with device activations at all. The
closest one loads the method and tolerates failure without running it, so the crash had
no coverage.
Also built the runtime and ran the existing runtime and extension tests on Linux
x86_64.
Known gaps, not fixed here
Two problems on the same feature remain, and each deserves its own change:
is labelled as host memory. A caller passing GPU tensors from Python is still not
correctly described to the runtime.
as if it were host memory.