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

@shoumikhinshoumikhin 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.

@github-actionsgithub-actionsBot added component: tests Issues re: Tests component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 29, 2026
@shoumikhin
shoumikhinforce-pushed the fix/lift-duplicate-constant-placeholder branch from ff312d1 to 9a9bfbfCompareAugust 30, 2026 18:04
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.
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signedcomponent: api [Python]Issues re: Python APIcomponent: coreIssues re: The core compilercomponent: dynamoIssues relating to the `torch.compile` or `torch._dynamo.export` pathscomponent: testsIssues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@shoumikhin