Repository navigation
Fix device placement of constants, shared inputs and delegate outputs - #24088
Open
shoumikhin wants to merge 16 commits into
Open
shoumikhin wants to merge 16 commits into
shoumikhin wants to merge 16 commits into
Conversation
shoumikhin
requested review from
JacobSzwejbka and
larryliu0820
as code owners
October 22, 2026 19:09
shoumikhin
added
the
release notes: exir
meta-cla
Bot
added
the
CLA Signed
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24088
Note: Links to docs will display an error until the docs builds have been completed. ❌ 3 New FailuresAs of commit 0ac8098 with merge base 185668b ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
shoumikhin
force-pushed
the
device-pass-constants
branch
2 times, most recently
from
October 22, 2026 23:08
3aadad6 to
894b0fa
Compare
Gasoonjia
reviewed
Oct 23, 2026
Gasoonjia
left a comment
Gasoonjia
left a comment
Contributor
There was a problem hiding this comment.
Thanks for find and fixing the issue! I left some thoughts and we can discuss.
| new_spec = copy.copy(spec) | ||
| new_spec.init_mem_planning_fields() | ||
| # A copy's output is planned memory, even when it copies a constant. | ||
| new_spec.const = False |
Contributor
There was a problem hiding this comment.
i feel like it is a liitle bit dangerous if we overwrite the const attribute silently here. Maybe we can rename the function to highlight it, or create a new function. WDYT?
Contributor
Author
There was a problem hiding this comment.
Good point. I renamed the helper to _copy_output_spec. Its docstring now says the output is always planned memory, so it is never a constant, even when the tensor it copies is one. It is only used to build the spec of an _h2d_copy or _d2h_copy output, so no other spec gets its const flag changed.
| if ( | ||
| self.skip_h2d_for_method_inputs | ||
| and self._is_placeholder(arg) | ||
| and not arg_spec.const |
Contributor
There was a problem hiding this comment.
im ok with this change to fix the correctness isssue but from api design perspective it doesn't quite in line with skip_h2d_for_method_inputs: constant var is also pert of method inputs (although implictly), and if user turn it on but still need to do h2d copy for constant it is not elegant.
My understanding is we may need to have devcie-resident constant materialization so that we can treat contant as regular input, in the device support scenerio.
Contributor
Author
There was a problem hiding this comment.
I agree that a constant is a kind of method input too. The problem is that today the runtime always points a constant at its data in host memory (in the program or in the external data file), whatever device its spec says. So if we skip the copy, the delegate gets a host pointer that is marked as device memory. Until the runtime can load constants into device memory, we need the copy to stay correct. I added a line to the comment on skip_h2d_for_method_inputs that says constants are still copied. Device-resident constants sound like a good follow-up, and then this flag could skip their copy too.
shoumikhin
force-pushed
the
device-pass-constants
branch
from
October 23, 2026 02:25
894b0fa to
ea00fc9
Compare
shoumikhin
changed the title
shoumikhin
force-pushed
the
device-pass-constants
branch
from
October 23, 2026 04:13
ea00fc9 to
e12db30
Compare
shoumikhin
force-pushed
the
device-pass-constants
branch
from
October 23, 2026 09:44
e12db30 to
71dd63f
Compare
shoumikhin
changed the title
shoumikhin
force-pushed
the
device-pass-constants
branch
from
October 23, 2026 18:03
71dd63f to
e1ed5ce
Compare
shoumikhin
force-pushed
the
device-pass-constants
branch
from
October 23, 2026 19:15
e1ed5ce to
635d39e
Compare
Contributor
Author
|
Since the last push, I rebased on main and made these changes:
Tested: all device pass tests pass at every commit, and 12 programs with real CUDA delegates on an H100 match eager PyTorch. |
Contributor
|
sry can you give me a brief introduction about why we need to update the dim order parts for the device input changing? |
Contributor
Author
|
@Gasoonjia It is needed by the constants change in this PR. When a constant or a parameter feeds a CUDA delegate from outside the delegate, this PR now copies it to the device. That copy moves the bytes the way the program stores them. The program stores a tensor that is neither contiguous nor channels-last, such as a transposed weight, as contiguous. Without the dim order change, the delegate is compiled for the weight's original strides, so it would read that copy in the wrong order. That is why the constants change refuses such a constant with an error instead of copying it.
So without it, a model with a transposed weight that two CUDA delegates read fails to export. With it, the model exports and matches eager on CUDA. |
Contributor
|
Thanks for the explanation! The changes for the current target LGTM. Since this PR touches the dim order pass, which every backend runs in to_edge, I added more CI labels for better coverage. Some thoughts:
I will OOO next week i will defer to @JacobSzwejbka to review and stamp. |
The emitter stores a constant tensor that is neither contiguous nor channels-last as contiguous. convert_constant_dim_order_pass makes lifted constants and buffers contiguous in to_edge, before partitioning, so a delegate is lowered with the layout the program will store. Parameters were left out. So a non-contiguous parameter that a delegate reads as an input, such as a transposed weight left outside the delegate, was lowered with its original strides but stored contiguous. A delegate that compiles those strides in, like the CUDA backend, would read the wrong elements from a copy of the stored bytes. The pass now converts such parameters too, keeping them nn.Parameter.
When two delegates on the same device are connected, PropagateDevicePass removes the device -> host -> device copy pair between them, and the second delegate reads the first one's device output directly. That is wrong when a host op writes the host copy in place in between. A mutable custom op that stays on the host (like a KV cache update) does this: to_executorch removes auto_functionalized before this pass, so the op writes its argument in place. The second delegate then reads the device tensor from before the write and gets a stale value. Now the pass keeps the round-trip when any op that reads the host copy also writes it, according to the op's schema. The check reuses _schema_writes_input from the buffer write-back pass. It does not look at where the write happens, so it also keeps a round-trip whose second delegate runs before the write. The new test runs a mutable host op between two delegates on cuda:0 and checks that the second delegate reads a fresh copy of the written host tensor. It fails without the change.
…item A lowered module that returns one tensor and is called from another module has no getitem after its call. The device pass only placed copies after getitems, so such an output got no _d2h_copy. A host op that read it got device memory, and another delegate on the same device got an _h2d_copy with a device source, which the copy kernel rejects at run time. The device-to-host copy is now inserted after the call itself in that case. The existing elision then removes the round-trip between two calls on the same device.
…tputs is set With skip_d2h_for_method_outputs, a delegate output that the method returns should stay on the device. The pass skipped the device-to-host copy only when the graph output was the copy's only reader, and it decided this while it was still inserting copies. So a returned output that something else also read was returned on the host: - A delegate on the same device: its input copy was still a reader at that point. The round-trip was removed later, but the host copy stayed. - A host op, or a delegate on another device: these need the host copy, so the graph output kept reading it too. For example, `a = f(x, y); return a, torch.sin(a)` returned `a` on the host. The pass now inserts every device-to-host copy, removes the round-trips, and then points the graph output at the device tensor. A copy stays only for the readers that still need it. When a host op writes the host copy in place, the device tensor never sees the write, so the output reads a new device copy of the written host copy.
A constant that feeds a delegate on a device, such as a buffer the partitioner leaves outside the delegate, broke device memory planning. - Without skip_h2d_for_method_inputs, its _h2d_copy cloned the constant's spec, constant flag included, so memory planning skipped the copy's output and emission failed. - With it, the constant was treated as a method input: its placeholder was tagged with the device while its data stays in the program on the host, so the delegate got a host pointer labeled as device memory. A copy's output is now never a constant, and skip_h2d_for_method_inputs applies only to placeholders that are not constants. The config comment and the CUDA docs say that this copy runs on every call, and the pass logs a warning that names the constant and its size, so the cost is visible at export time. The copy moves the constant's bytes as they are stored. The program stores a constant that is neither contiguous nor channels-last as contiguous, so such a constant, which a pass that runs after to_edge can create, is now refused with a clear error instead of getting a copy with the wrong layout.
With skip_h2d_for_method_inputs, a method input had to have exactly one user, because its spec is changed in place to the delegate's device. A program that calls one lowered module more than once, or passes one input to two delegates, could not keep its inputs on the device. Changing the spec is safe when every user is a delegate on that same device, so that case is now allowed. An input that a host op or a delegate on another device also reads is still refused, and the error names that reader and the device the input was going to. When one of the readers is the write-back of a mutated input or mutable buffer, the error names it and says so, and for a mutable buffer it points to tag_mutated_buffer, which lets the partitioner keep the buffer inside the delegate. The CUDA docs now list these rules.
When several delegates on one device read the same tensor, the pass made a separate host-to-device copy for each of them. A constant read by three delegates on cuda:0, such as a lookup table, was copied three times on every call. Method inputs were copied the same way when skip_h2d_for_method_inputs is off. The pass visits delegates in graph order, so a copy made for an earlier delegate on the same device already comes before the later ones. It now reuses that copy. A tensor read on two devices still gets one copy for each device. The warning for a constant copied on every call is logged once per copy. The reused copy stays in device memory from the first delegate that reads it to the last, so memory planning cannot give that space to other tensors in between. That is the cost of saving the copy on every call. A tensor that any op writes in place, such as an argument of a mutable custom op or a mutable buffer with its write-back, is not reused, because an earlier copy can hold the value from before the write. The check does not look at where the write happens, so it also covers writes before the first reader or after the last one, where reuse would be safe. Those tensors get one copy for each delegate, as before. For example, a constant read by delegates on cuda:0, cuda:1 and cuda:0 now gets two copies instead of three.
SpecPropPass gives the result of an in-place op the TensorSpec of the tensor it writes, so memory planning puts both in one buffer. When a host op writes a device delegate output in place, for example relu_ or index_put_ made by reinplace_pass, PropagateDevicePass moves the op to the host copy of the output, but the result kept the spec of the device output. Memory planning then refused the program: In-place node aten_relu__default ... result 0 does not share the TensorSpec of its mutated input _d2h_copy_default. This happened both for an output read through getitem and for a single tensor returned by the call. When the pass moves readers to the host copy, it already pointed the graph output at the copy's spec. It now does the same for in-place results and their own readers, so they share the copy's spec. A returned in-place result is the result of a host op, so the method returns it from the host, also with skip_d2h_for_method_outputs. The new test runs relu_ and then mul_ on a delegate output, for getitem and for a single tensor, with skip_d2h_for_method_outputs off and on. It fails without this change.
Before PropagateDevicePass reuses a device copy of a tensor, removes a device -> host -> device round-trip, or returns a delegate output on the device instead of its host copy, it checks that no host op writes that tensor in place. The check only looked at ops that write the tensor itself. A mutable custom op on the host can also write a view of the tensor. When to_executorch removes auto_functionalized, it makes that view again, for example with aten.as_strided, and the op writes the view. The check missed this, so a later delegate, or the method output, could get the value from before the write. For example: t = x * y a = t + y # delegate on cuda:0 accumulate(t.view(-1), y.view(-1)) # host op, writes t return a, t + y # delegate on cuda:0 Here the second delegate read the device copy of t made before accumulate. The check now also counts writes to views of the tensor, and to views of those, as the op schemas declare them. The result of an in-place op counts as a view of its input. A view_copy does not, because it still makes its own tensor at this point, and ReplaceViewCopyWithViewPass turns it into a view later only where no write can tell the difference. The tests for round-trip removal, returned outputs and copy reuse each get a case that writes through a view. These cases fail without this change. aten.as_strided has no out variant yet, so in these cases to_executorch is told to keep it.
The device pass tests covered a delegate output read by delegates on two devices and a constant read on two devices. A method input read by delegates on different devices was only covered with skip_h2d_for_method_inputs, where it is refused. The new test has delegates on cuda:0, cuda:1 and cuda:0 that read the same method input. The second and third each read the output of the delegate before them, which is on the other device. The input gets one copy for each device, and both delegates on cuda:0 read the same one. Each output crosses to the other device through the host. The test fails if copy reuse ignores the device index.
The device pass fixes were only tested on the graph and on memory planning. This adds tests that lower small models to real CUDA delegates, run them, and compare the results with eager PyTorch. Each add becomes its own delegate, and parameters and buffers stay outside the delegates, so the program copies them to the device. - A transposed parameter. Without the parameter fix, to_executorch fails, because the copy would move the contiguous bytes the program stores to a delegate compiled for the transposed strides, and the pass refuses that. - A buffer read by two delegates. The tests need a CUDA GPU and are skipped otherwise.
With skip_h2d_for_method_inputs and skip_d2h_for_method_outputs, users expect no copies at the method boundary. Two cases still copy on every call: - a constant that feeds a device delegate from outside it, because its data is stored on the host. The pass already logs a warning with its name and size. - with skip_d2h_for_method_outputs, a returned delegate output that a custom host op with a mutable argument, such as a KV cache update, writes in place. The method returns a device copy of the written host copy, so the output goes to the host and back on every call. This was silent. The second case now logs a warning too. It gives the output's size, the device and the op that writes it, and suggests moving that op into the delegate. To name the op, the write check now has a variant, _in_place_writer, that returns the writer. The config comment for skip_d2h_for_method_outputs and the CUDA backend docs mention the copies. The test for the second case now checks the warning, and fails without it.
PropagateDevicePass keeps a device -> host -> device round-trip between two delegates on one device, and does not reuse a device copy of a tensor, when a host op writes the tensor in place. It looked for such a write anywhere in the graph. But a write that runs after the delegate has read the copy cannot change what the delegate reads. So some programs paid a copy on every call that main did not make, for example when reinplace_pass turns a host op into an in-place op, which it only does on the last use of a tensor, or when a mutable buffer is written back after its readers. The round-trip now stays only when the host copy is written before it is copied back to the device. A copy is reused only when no write runs between it and the delegate that would reuse it, and the newest copy on that device is the one considered. The graph order is the order the method runs in, so it decides what comes before. The new test writes two tensors on the host after the second of two delegates on one device has read them. The round-trip is removed and the tensor both delegates read gets one copy. The test fails without this change, and the tests that write between the delegates still pass.
With skip_d2h_for_method_outputs, when a host op writes the host copy of a returned delegate output in place, the method returns a device copy of the written tensor. If a later delegate on the same device also reads that tensor, it already reads a device copy of the written host copy, so the method made two identical copies on every call. The method output now reads the copy that the later delegate reads, when no op writes the tensor in place between that copy and the end of the method. Otherwise it still gets its own copy. The new test returns a delegate output that a host op writes and that a second delegate on the same device reads. There is one device copy, read by both. The test fails without this change.
…der_pass _update_placeholder_meta looks up the input spec of a lifted tensor and then its placeholder. The second check tested the input spec again, which was already known to exist, so a missing placeholder failed later with an AttributeError on None instead of the intended "Missing placeholder" error. Parameters now go through this function too. The new test points a parameter's input spec at a name with no placeholder and expects the error. It fails without this change.
backends/cuda/tests/test_device_placement.py runs programs whose host/device copies come from PropagateDevicePass and whose parameter layout comes from convert_constant_dim_order_pass. The CUDA workflow did not start for changes to those passes, so a change to them alone could merge without running the test. The workflow's pull_request paths and the unittest-cuda job now include those two passes and the device pass config, the same way they already include the pybindings paths for the pybindings device test. The other CUDA jobs keep their own path checks, so they do not start for these files.
shoumikhin
force-pushed
the
device-pass-constants
branch
from
October 24, 2026 13:06
aea6d31 to
0ac8098
Compare
Contributor
Author
|
@Gasoonjia Thanks for the review. I pushed these changes:
|
This branch was successfully deployed
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.
What is wrong today
When a delegate runs on a device (such as CUDA),
PropagateDevicePassadds a host-to-device copy (_h2d_copy) before each delegate input and a device-to-host copy (_d2h_copy) after each output.skip_h2d_for_method_inputsandskip_d2h_for_method_outputsskip those copies for method inputs and outputs.Six cases go wrong:
relu_,to_executorchalso fails in memory planning._d2h_copy, so a host op reads device memory.skip_d2h_for_method_outputs, it still comes back on the host.to_executorchfails, or, with the input skip flag, the delegate gets a host pointer labeled as device memory.What this change does
Roughly one commit per item:
to_edgenow makes a parameter that is neither contiguous nor channels-last contiguous before partitioning, as it already does for constants and buffers. This applies to every backend. With_skip_dim_order=True, only the parameter itself changes: ops computed from it keep the old strides in the graph, so a delegate that reads such an op's result can expect the wrong layout, as for buffers and constants today.to_executorchneedsToOutVarPass(ignore_to_out_var_failure=True), becauseaten.as_stridedhas no out variant yet)._d2h_copyright after the call.skip_d2h_for_method_outputs, a returned output now stays on the device even when another op reads it. If a host op writes its host copy in place, a device copy of the written tensor is returned, a later delegate on that device that reads the same copy shares it, and a warning says that this copies on every call.What was tested
exir/tests/test_propagate_device_pass.py: all tests pass at every commit.exir/tests/test_passes.py: the new parameter cases, one on the pass and one throughto_edgewith and without_skip_dim_order, fail without the fix. A small fix in the same pass makes a missing placeholder raise its intended error, with a test.backends/cuda/tests/test_device_placement.py(new, needs a CUDA GPU): runs a transposed parameter and a buffer read by two delegates on real CUDA delegates and compares with eager. Theunittest-cudajob incuda.ymlnow also runs when the device pass, its config, or the dim order pass changes, so this test runs on such changes. An earlier version of this test ran on a CUDA GPU: it passed, and without the parameter fixto_executorchrefused the transposed parameter. CI runs the current version.Follow-up issues are being filed for keeping constants in device memory and for the smaller follow-ups from review.