Skip to content

externalize: skip submodules not invoked in the exported graph#18

Merged
gokulkrishna98 merged 3 commits into
apple:mainfrom
gokulkrishna98:dev/gokul/fix_extern_bug_v1
Jun 17, 2026
Merged

externalize: skip submodules not invoked in the exported graph#18
gokulkrishna98 merged 3 commits into
apple:mainfrom
gokulkrishna98:dev/gokul/fix_extern_bug_v1

Conversation

@gokulkrishna98

@gokulkrishna98 gokulkrishna98 commented Jun 16, 2026

Copy link
Copy Markdown
Contributor

Description:

  • When a model registers a submodule whose class matches an externalize_modules target, but the model's forward never calls it, the re-exported FX graph contains no custom-op node for that submodule. The pipeline previously raised:
  ValueError: Custom op for '<name>' not found in any ancestor program

from _find_program_for and aborting the entire conversion.

Fix:

  • Skip such submodules with a UserWarning instead of raising. The unused module's forward is still restored by
    _restore_externalized.

  • Also silences pre-existing F821 lint warnings on transform_with_custom_compression_ops references in skipped tests so pre-commit ruff-check passes.

Testing

  • python unit test
  • Added new unit test to test it.

When a model registers a submodule whose class matches an
externalize_modules target but the model's forward never calls it, the
re-exported FX graph contains no custom-op node for that submodule.
The pipeline previously raised
``ValueError: Custom op for '<name>' not found in any ancestor program``
from _find_program_for, aborting the whole conversion.

Skip these submodules (with a UserWarning) instead. The unused module's
forward is still restored by _restore_externalized.

Also silence pre-existing F821 lint warnings on transform_with_custom_compression_ops
references in skipped tests so pre-commit ruff-check passes.
@gokulkrishna98 gokulkrishna98 changed the title Skip externalize for submodules not invoked in the exported graph externalize: skip submodules not invoked in the exported graph Jun 16, 2026
@gokulkrishna98 gokulkrishna98 marked this pull request as ready for review June 16, 2026 21:46
@gokulkrishna98 gokulkrishna98 merged commit ea728d6 into apple:main Jun 17, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants