Skip to content

Preserve nested control-flow subgraphs during pt2e convert#4540

Open
atharvjairath wants to merge 1 commit into
pytorch:mainfrom
atharvjairath:fix-pt2e-nested-while-loop-subgraphs
Open

Preserve nested control-flow subgraphs during pt2e convert#4540
atharvjairath wants to merge 1 commit into
pytorch:mainfrom
atharvjairath:fix-pt2e-nested-while-loop-subgraphs

Conversation

@atharvjairath

Copy link
Copy Markdown

Fixes #4455.

GraphModule.delete_all_unused_submodules() only inspects the immediate graph, so convert() could delete submodules that are referenced only from nested control-flow subgraphs. In nested while_loop cases, this left dangling references and caused DuplicateDQPass graph linting to fail.

This adds a recursion-aware cleanup wrapper that preserves submodules referenced from nested GraphModule children, and adds a regression test covering nested while_loop conversion.

Test plan:

  • python -m ruff check torchao/quantization/pt2e/convert.py test/quantization/pt2e/test_quantize_pt2e.py
  • python -m pytest test/quantization/pt2e/test_quantize_pt2e.py::TestQuantizePT2E::test_nested_while_loop_convert_preserves_subgraphs -q
  • python -m pytest test/quantization/pt2e/test_duplicate_dq.py -q
  • git diff --check

@pytorch-bot

pytorch-bot Bot commented Jun 27, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4540

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

convert_pt2e crashes in DuplicateDQPass — GraphModule rebuild drops submodules only referenced by other submodules' internal graphs

2 participants