Lift each constant once when exporting a partitioned graph - #4642
Open
shoumikhin wants to merge 1 commit into
Open
Lift each constant once when exporting a partitioned graph#4642shoumikhin wants to merge 1 commit into
shoumikhin wants to merge 1 commit into
Conversation
shoumikhin
force-pushed
the
fix/lift-duplicate-constant-placeholder
branch
2 times, most recently
from
August 30, 2026 20:10
9a9bfbf to
6935e4c
Compare
This was referenced Aug 30, 2026
shoumikhin
force-pushed
the
fix/lift-duplicate-constant-placeholder
branch
from
August 31, 2026 04:01
6935e4c to
1fd08cd
Compare
Inlining the partitions left in PyTorch copies each one's own get_attr back into the graph, so one constant can arrive as several get_attr nodes sharing a target. The lifting pass turned each of those into its own placeholder. fx makes a placeholder's name unique but not its target, and placeholder codegen emits the target, so the generated forward ended up with the same argument twice and failed to compile with "SyntaxError: duplicate argument". The signature also gained one input spec per read, all pointing at the same state dict key. Lifting each target once and pointing the later readers at that placeholder fixes both. Two things that were already broken here and are fixed in the same place, since the new branch is where they now belong. A lifted read that is also a graph output needs the output rename that only ran for the first read, or the signature names a node that was just erased and building the program fails. And the name that has to be unique in generated code is the target with dots replaced, not the target itself, so two attributes differing only by a dot still collided into one argument; taking the placeholder's own node name avoids that, which is what upstream's lifting pass does. Test plan: Two tests in tests/py/dynamo/lowering/test_exporter_lift_dedup.py. The first builds the graph the inliner produces, two constants each read twice and interleaved, and checks the lifted graph has three placeholders rather than five, that their targets are unique, that the signature carries two buffer specs and that each spec names the placeholder in the same position, that the module recompiles, and that running it gives the right value. It fails on the previous code with duplicate placeholders, and it also fails against a version that reuses whichever placeholder was lifted first rather than the one for this target: that version passes when only one constant is present, which is why there are two. The second drives the partitioner and the inliner and asserts they still produce duplicate get_attr targets. It passes either way, which is the point: it does not test the fix, it pins the shape the fix exists for, so if inlining stops producing duplicates and the dedup becomes dead code, that test says so rather than the first one quietly continuing to pass on a graph nothing produces. It builds no engine and saves nothing. Verified separately, since neither has a test yet: a deduplicated read that is also a graph output now exports where it raised before, and two attributes named a.b and a_b now produce distinct arguments and run to the right value. Both tests live here rather than with the exporter's other tests because this directory is collected by the lane that runs on every pull request. The exporter inlining tests still pass.
shoumikhin
force-pushed
the
fix/lift-duplicate-constant-placeholder
branch
from
August 31, 2026 06:15
1fd08cd to
efadd13
Compare
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.
Inlining a partitioned graph leaves one get_attr node per partition that reads a
constant, all carrying the same target. The lifting pass turned each of those into
its own placeholder. fx makes a placeholder's name unique but not its target, and
placeholder codegen emits the target, so the generated forward ended up with the
same argument twice and failed to compile with "SyntaxError: duplicate argument".
The signature also gained one input spec per read, all pointing at the same state
dict key.
Lifting each target once and pointing the later readers at that placeholder fixes
both. One get_attr node with several users was always fine; what breaks is several
get_attr nodes with the same target, which is what the partitions left in PyTorch
produce once the inliner copies them back into one graph.
Test plan:
Added a test that builds that graph directly: three get_attr nodes for one buffer.
It checks the lifted graph has two placeholders rather than four, that their
targets are unique, that the signature carries one buffer spec, that the module
recompiles, and that running it gives the right value.
Without the change all five of those fail, including the recompile with the
duplicate argument error. The value check is there because a version of the fix
that rewires the later readers to the wrong placeholder passes every structural
check while computing the wrong answer, and the buffer holds distinct values so
that a wrong read is visible.
It sits beside the other tests for lift rather than with the exporter's inlining
tests, because it needs no GPU and no TensorRT engine, so it runs in the lane that
every pull request runs. The tests it was written next to are only collected by
suites that run on a schedule, so it would not otherwise have run before merge.
The rest of the exporter inlining file still passes.