Document and test the memory planning a device-activation export needs - #21703
Document and test the memory planning a device-activation export needs#21703shoumikhin wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21703
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 499bd94 with merge base 730b77a ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
The problem
ExecuTorch can export a model so its activations stay on the GPU. You hand over memory
that is already on the GPU, the model runs on it, and the output stays there. No copying
back and forth, which is the point of the feature.
Two settings control this, and they live in different places. The first says "do not
insert copies at the delegate boundary":
The second is memory planning, which by default reserves a buffer for every graph input
and output. That is the right default for an ordinary model, and wrong here: the runtime
reserves its own GPU buffer for an input you already allocated, and then fills it by
copying your memory into it. The copy the export asked to skip comes back at run time.
It is worse than a wasted copy. For device memory that copy is a plain host
memcpyinto a GPU pointer, which is undefined and crashes the process with no message.
Nothing said the two settings had to be paired, and no test exported such a program and
checked it. The result is a documented feature whose documented usage does not work.
The change
Documentation and a test. No behavior change.
The configuration now states the pairing, with a complete example:
The measurable difference, for a model with two small inputs:
The extra 512 bytes are a duplicate of the caller's two inputs.
Test plan
Added
test_skipping_copies_requires_unallocated_graph_io, which pins the pairing:with default planning the graph inputs are planned, and with the pairing none of the
inputs or outputs are. It needs no GPU, because it inspects the exported program rather
than running it.
Confirmed the test fails without the pairing:
Ran the full
test_propagate_device_passfile, 24 tests, all passing.Separately verified on hardware that a program exported with the pairing runs correctly
from C++ with GPU memory in and out, matching eager PyTorch exactly on two different GPU
architectures.