Skip to content

[externalization] Add a two-step patch/sub-export API for module externalization #53

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gokulkrishna98 wants to merge 3 commits into
base: main
Choose a base branch
from
+733 −148
Open

[externalization] Add a two-step patch/sub-export API for module externalization #53

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
80a50c2
Add module externalization API
gokulkrishna98 Aug 3, 2026
a61e4a9
Address review: cross-reference no-call-site paths, test multi-entry
gokulkrishna98 Aug 3, 2026
67c610e
Unify phase/step numbering in externalize docstrings
gokulkrishna98 Aug 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion coreai_torch/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,11 @@
from ._decomp import get_decomp_table
from ._torch_metal_kernel import TorchMetalKernel
from .converter import TorchConverter
from .externalize import ExternalizeSpec
from .externalize import (
ExternalizeSpec,
_patch_model_for_externalization,
_subexport_and_restore,
)

__all__ = [
"__version__",
Expand All @@ -27,6 +31,8 @@
"TorchMetalKernel",
"get_decomp_table",
"generate_composite_decl",
"_patch_model_for_externalization",
"_subexport_and_restore",
]

_TORCH_MAX_VERSION = "2.13.0"
Expand Down
133 changes: 61 additions & 72 deletions coreai_torch/converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from __future__ import annotations

from collections import OrderedDict
from collections.abc import Iterator, Sequence
from collections.abc import Sequence
from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable, Optional, cast
Expand Down Expand Up @@ -63,13 +63,11 @@
from ._validate import validate_exported_program
from .externalize import (
ExternalizeSpec,
_ExportedModule,
_finalize_module_export,
_mark_externalize,
_prepare_externalized,
_PreparedModule,
_ExternalizedExportedProgram,
_find_marked_submodules,
_patch_model_for_externalization,
_restore_externalized,
_torch_export_module,
_subexport_and_restore,
)


Expand All @@ -86,6 +84,7 @@ class _StagedEntry:
module: torch.nn.Module | None = None
export_fn: Callable[..., Any] | None = None
externalize_modules: list[type | ExternalizeSpec] | None = None
_externalized_exported_programs: list[_ExternalizedExportedProgram] | None = None


class Context(_CoreAIAuthoringContext):
Expand Down Expand Up @@ -182,7 +181,7 @@ def _init_conversion_state(self) -> None:
self._externalized_lowerings: dict[str, Callable[..., Any]] = {}

# externalized modules
self._externalized_modules: list[_ExportedModule] = []
self._externalized_exported_programs: list[_ExternalizedExportedProgram] = []

self.user_input_names: Sequence[str] = []
self.user_output_names: Sequence[str] = []
Expand All @@ -201,19 +200,23 @@ def add_exported_program(
output_names: Sequence[str] | None = None,
state_names: Sequence[str] | None = None,
entrypoint_name: str = "main",
_externalized_exported_programs: list[_ExternalizedExportedProgram]
| None = None,
) -> Self:
"""Stage a pre-exported ExportedProgram for conversion.
"""Stage a pre-exported ``ExportedProgram`` for conversion.

The caller is responsible for calling ``torch.export.export()`` and
``run_decompositions()`` before passing the program.

Args:
input_names: Non-stateful forward() arg names only.
output_names: Return value names only (not mutation outputs).
state_names: One name per state, applied to both input and
mutation output. Order: buffers (registration order), then
mutated user inputs (signature order). Defaults to FX
placeholder names when not provided.
input_names: Non-stateful ``forward()`` arg names.
output_names: Return value names (not mutation outputs).
state_names: One name per state (buffers then mutated inputs).
Defaults to FX placeholder names.
_externalized_exported_programs: Result of
:func:`coreai_torch.externalize._subexport_and_restore`.
When set, emits composite ``coreai.graph``s for the patched
call sites in ``exported_program``.

Returns ``self`` for chaining.
"""
Expand All @@ -233,6 +236,7 @@ def add_exported_program(
output_names=output_names or [],
state_names=state_names or [],
entrypoint_name=entrypoint_name,
_externalized_exported_programs=_externalized_exported_programs,
)
)
return self
Expand Down Expand Up @@ -303,89 +307,67 @@ def add_pytorch_module(
)
return self

def _run_externalize_pipeline(self) -> None:
"""Mark, re-export, and export each externalized submodule.
def _run_externalize_pipeline_from_module(self) -> None:
"""Externalize from a live ``nn.Module`` + ``export_fn`` (Phases 1-3).

The model is patched temporarily and always restored, even on error.

Steps:
1. Mark matching submodules with custom ops via ``_mark_externalize``.
2. Re-export the whole model (now containing custom op calls).
3. For each marked submodule, export it standalone and decompose.
4. Store results in ``self._externalized_modules`` and update
``self.exported_program`` to the re-exported whole-model program.
Used by :meth:`add_pytorch_module`. Marks matching submodules with
custom ops, re-exports the whole model, then delegates Phases 2-3
(sub-export) to :func:`_subexport_and_restore`, which stashes the
results in ``self._externalized_exported_programs`` for
:meth:`_perform_externalization` to emit and restores the model
patch (``try/finally``) internally.
"""
assert self._module is not None
assert self._export_fn is not None
assert self._externalize_modules is not None

_mark_externalize(self._module, self._externalize_modules)
_patch_model_for_externalization(self._module, self._externalize_modules)

try:
# Re-export after marking — failure here is our bug, not the user's
try:
with self._progress_bar.status("Re-exporting for externalization..."):
whole_ep: ExportedProgram = self._export_fn(self._module)
inject_subbyte_tensors(whole_ep)
except Exception as e:
raise RuntimeError(
f"Internal error: re-export after externalization failed: {e}\n"
f"This is a coreai-torch bug. Please report it."
) from e

# Prepare and export each submodule
preps: Iterator[_PreparedModule] = _prepare_externalized(
self._module, whole_ep
)
exts: list[_ExportedModule] = []
with self._progress_bar.stream("Externalizing submodules") as advance:
for prep in preps:
try:
inner_ep = _torch_export_module(prep)
except Exception as e:
raise RuntimeError(
f"Internal error: failed to export submodule "
f"'{prep.name}': {e}\n"
f"This is a coreai-torch bug. Please report it."
) from e

# Use the standard decomposition table for inner modules.
# The user's export_fn may preserve composite ops like
# aten.scaled_dot_product_attention so they survive in the
# *whole-model* graph for externalization detection.
# Inside the externalized body those ops must be decomposed.
inner_ep = inner_ep.run_decompositions()
exts.append(_finalize_module_export(prep, inner_ep))
advance()

self._externalized_modules = exts
self.exported_program = whole_ep
finally:
_restore_externalized(self._module)
with self._progress_bar.status("Re-exporting for externalization..."):
whole_ep: ExportedProgram = self._export_fn(self._module)
inject_subbyte_tensors(whole_ep)
except Exception as e:
# Re-export after marking failed — our bug, not the user's. Restore
# here since _subexport_and_restore never gets a chance to.
_restore_externalized(_find_marked_submodules(self._module))
raise RuntimeError(
f"Internal error: re-export after externalization failed: {e}\n"
f"This is a coreai-torch bug. Please report it."
) from e

# The user's export_fn may preserve composite ops like
# aten.scaled_dot_product_attention so they survive in the
# *whole-model* graph for externalization detection; decompositions
# inside each submodule's own export happen in _subexport_and_restore.
self._externalized_exported_programs = _subexport_and_restore(
self._module, whole_ep, progress_bar=self._progress_bar
)
self.exported_program = whole_ep

def _perform_externalization(self, context) -> None:
"""Emit externalized submodule graphs and register their lowerings.

Modules are processed deepest-first so that inner lowerings are
available when parent graphs are built (nested externalization).

For each :class:`_ExportedModule`, this method:
For each :class:`_ExternalizedExportedProgram`, this method:

1. Temporarily swaps ``self.exported_program`` to the submodule's export.
2. Builds a ``coreai.graph noinline`` via :meth:`_get_graph_op`
(with ``composite_decl`` and ``private`` for composite op modules).
3. Registers a lowering that maps the custom op to ``coreai.invoke``.
4. Restores the whole-model program.
"""
if not self._externalized_modules:
if not self._externalized_exported_programs:
return

whole_program: ExportedProgram = self.exported_program

# Process deepest modules first so their lowerings exist
# when parent modules are built.
sorted_exts: list[_ExportedModule] = sorted(
self._externalized_modules,
sorted_exts: list[_ExternalizedExportedProgram] = sorted(
self._externalized_exported_programs,
key=lambda ext: ext.name.count("."),
reverse=True,
)
Expand Down Expand Up @@ -941,7 +923,14 @@ def to_coreai(self, *, entrypoints: Sequence[str] | None = None) -> AIProgram:
self._module = entry.module
self._export_fn = entry.export_fn
self._externalize_modules = entry.externalize_modules
self._run_externalize_pipeline()
self._run_externalize_pipeline_from_module()
self._perform_externalization(module.context)

# Handle externalization for pre-marked ExportedProgram path
elif entry._externalized_exported_programs is not None:
self._externalized_exported_programs = (
entry._externalized_exported_programs
)
self._perform_externalization(module.context)

self._get_graph_op(
Expand Down Expand Up @@ -971,7 +960,7 @@ def __repr__(self) -> str:
parts = [f" {entry.entrypoint_name}: {kind} {inputs} -> {outputs}"]
if entry.externalize_modules:
names = [
s.module_type.__name__
s.target_class.__name__
if isinstance(s, ExternalizeSpec)
else s.__name__
for s in entry.externalize_modules
Expand Down
Loading