Resolve externalized call sites by op target, not by recorded node name - #71
Open
metascroy wants to merge 2 commits into
Open
Resolve externalized call sites by op target, not by recorded node name#71metascroy wants to merge 2 commits into
metascroy wants to merge 2 commits into
Conversation
Author
|
cc @cymbalrush |
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.
Problem
add_exported_program(_externalized_exported_programs=...)is documented to emitcomposite graphs "for the patched call sites in
exported_program", but itregisters each lowering under
_ExternalizedExportedProgram.source_nodes— FXnode names captured back when
_subexport_and_restoreran.The two-step API introduced in #53 exists precisely so a caller can do work
between the phases;
_patch_model_for_externalization's own docstring showsep = my_export_or_quantize_pipeline(model). Any pass in that window thatrebuilds or renames nodes leaves the recorded names matching nothing. The
lowerings are then registered under dead keys, and conversion fails later with
an opaque error that names neither the submodule nor the cause:
The call sites are still there — only their names changed. The custom op target
survives any such transform.
Fix
TorchConverter._resolve_source_nodeslocates each call site by op target andpairs it with its
_ExternalizedExportedProgram, falling back to the recordednames when it cannot.
_perform_externalizationuses the resolved names.graph size rather than one graph walk per op name.
call site lives in its parent's program.
order. This assumes a transform preserves the relative order of an op's call
sites, as renaming and the usual lowering passes do.
graph changed shape rather than just its names, so pairing by position would
be meaningless. That op keeps its recorded names and a
UserWarningexplainswhy, instead of silently degrading to the error above.
_utils._externalized_op_nameis the inverse of the existing_find_custom_op_node/_find_all_custom_op_nodeslookups and shares_EXTERNALIZE_NAMESPACEwith them.Tests
test_call_sites_resolved_after_a_renaming_transform— renames everycall_functionnode between_subexport_and_restoreandadd_exported_program, then asserts both call sites still lower. Two callsites, so it covers ordered pairing as well as resolution. Fails without the
change.
test_mismatched_call_site_count_warns_and_falls_back— asserts theUserWarningon a count mismatch.Notes
No behaviour change for
add_pytorch_module: it converts the same program thesubmodules were prepared from, so resolution returns the recorded names.