Skip to content

Lift each constant once when exporting a partitioned graph - #4642

Open
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:fix/lift-duplicate-constant-placeholder
Open

Lift each constant once when exporting a partitioned graph#4642
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:fix/lift-duplicate-constant-placeholder

Conversation

@shoumikhin

@shoumikhin shoumikhin commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

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.

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
shoumikhin force-pushed the fix/lift-duplicate-constant-placeholder branch from 1fd08cd to efadd13 Compare August 31, 2026 06:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant