Skip to content

Resolve externalized call sites by op target, not by recorded node name - #71

Open
metascroy wants to merge 2 commits into
apple:mainfrom
metascroy:externalize-refactor
Open

Resolve externalized call sites by op target, not by recorded node name#71
metascroy wants to merge 2 commits into
apple:mainfrom
metascroy:externalize-refactor

Conversation

@metascroy

Copy link
Copy Markdown

Problem

add_exported_program(_externalized_exported_programs=...) is documented to emit
composite graphs "for the patched call sites in exported_program", but it
registers each lowering under _ExternalizedExportedProgram.source_nodes — FX
node names captured back when _subexport_and_restore ran.

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 shows
ep = my_export_or_quantize_pipeline(model). Any pass in that window that
rebuilds 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:

ValueError: unable to handle call function op:
  target: norm_rmsnorm_impl_7b5f5a2a.default, namespace: coreai_torch_ext

The call sites are still there — only their names changed. The custom op target
survives any such transform.

Fix

TorchConverter._resolve_source_nodes locates each call site by op target and
pairs it with its _ExternalizedExportedProgram, falling back to the recorded
names when it cannot. _perform_externalization uses the resolved names.

  • One pass per program, bucketing nodes by op, so cost stays linear in total
    graph size rather than one graph walk per op name.
  • Every program is searched, not just the whole-model one: a nested submodule's
    call site lives in its parent's program.
  • Multiple call sites of one op are paired in graph order, which is preparation
    order. This assumes a transform preserves the relative order of an op's call
    sites, as renaming and the usual lowering passes do.
  • If an op's call-site count no longer matches its prepared submodules, the
    graph changed shape rather than just its names, so pairing by position would
    be meaningless. That op keeps its recorded names and a UserWarning explains
    why, instead of silently degrading to the error above.

_utils._externalized_op_name is the inverse of the existing
_find_custom_op_node / _find_all_custom_op_nodes lookups and shares
_EXTERNALIZE_NAMESPACE with them.

Tests

  • test_call_sites_resolved_after_a_renaming_transform — renames every
    call_function node between _subexport_and_restore and
    add_exported_program, then asserts both call sites still lower. Two call
    sites, so it covers ordered pairing as well as resolution. Fails without the
    change.
  • test_mismatched_call_site_count_warns_and_falls_back — asserts the
    UserWarning on a count mismatch.

Notes

No behaviour change for add_pytorch_module: it converts the same program the
submodules were prepared from, so resolution returns the recorded names.

@metascroy

Copy link
Copy Markdown
Author

cc @cymbalrush

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.

1 participant