Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,510 @@
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
import zipfile
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Any
|
||||
from typing_extensions import deprecated
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch.fx.passes.infra.pass_base import PassResult
|
||||
from torch.types import FileLike
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AdditionalInputs",
|
||||
"Constraint",
|
||||
"CustomDecompTable",
|
||||
"default_decompositions",
|
||||
"Dim",
|
||||
"dims",
|
||||
"draft_export",
|
||||
"export",
|
||||
"ExportBackwardSignature",
|
||||
"ExportedProgram",
|
||||
"ExportGraphSignature",
|
||||
"FlatArgsAdapter",
|
||||
"load",
|
||||
"ModuleCallEntry",
|
||||
"ModuleCallSignature",
|
||||
"register_dataclass",
|
||||
"save",
|
||||
"ShapesCollection",
|
||||
"unflatten",
|
||||
"UnflattenedModule",
|
||||
]
|
||||
|
||||
# To make sure export specific custom ops are loaded
|
||||
import torch.export.custom_ops
|
||||
|
||||
from ._state_dict_utils import _restore_state_dict
|
||||
from .decomp_utils import CustomDecompTable
|
||||
from .dynamic_shapes import AdditionalInputs, Constraint, Dim, dims, ShapesCollection
|
||||
from .exported_program import (
|
||||
default_decompositions,
|
||||
ExportedProgram,
|
||||
ModuleCallEntry,
|
||||
ModuleCallSignature,
|
||||
)
|
||||
from .graph_signature import ExportBackwardSignature, ExportGraphSignature
|
||||
from .unflatten import FlatArgsAdapter, unflatten, UnflattenedModule
|
||||
|
||||
|
||||
PassType = Callable[[torch.fx.GraphModule], PassResult | None]
|
||||
|
||||
log: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def export(
|
||||
mod: torch.nn.Module,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: Mapping[str, Any] | None = None,
|
||||
*,
|
||||
dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None = None,
|
||||
strict: bool = False,
|
||||
preserve_module_call_signature: tuple[str, ...] = (),
|
||||
prefer_deferred_runtime_asserts_over_guards: bool = False,
|
||||
) -> ExportedProgram:
|
||||
"""
|
||||
:func:`export` takes any nn.Module along with example inputs, and produces a traced graph representing
|
||||
only the Tensor computation of the function in an Ahead-of-Time (AOT) fashion,
|
||||
which can subsequently be executed with different inputs or serialized. The
|
||||
traced graph (1) produces normalized operators in the functional ATen operator set
|
||||
(as well as any user-specified custom operators), (2) has eliminated all Python control
|
||||
flow and data structures (with certain exceptions), and (3) records the set of
|
||||
shape constraints needed to show that this normalization and control-flow elimination
|
||||
is sound for future inputs.
|
||||
|
||||
**Soundness Guarantee**
|
||||
|
||||
While tracing, :func:`export()` takes note of shape-related assumptions
|
||||
made by the user program and the underlying PyTorch operator kernels.
|
||||
The output :class:`ExportedProgram` is considered valid only when these
|
||||
assumptions hold true.
|
||||
|
||||
Tracing makes assumptions on the shapes (not values) of input tensors.
|
||||
Such assumptions must be validated at graph capture time for :func:`export`
|
||||
to succeed. Specifically:
|
||||
|
||||
- Assumptions on static shapes of input tensors are automatically validated without additional effort.
|
||||
- Assumptions on dynamic shape of input tensors require explicit specification
|
||||
by using the :func:`Dim` API to construct dynamic dimensions and by associating
|
||||
them with example inputs through the ``dynamic_shapes`` argument.
|
||||
|
||||
If any assumption can not be validated, a fatal error will be raised. When that happens,
|
||||
the error message will include suggested fixes to the specification that are needed
|
||||
to validate the assumptions. For example :func:`export` might suggest the
|
||||
following fix to the definition of a dynamic dimension ``dim0_x``, say appearing in the
|
||||
shape associated with input ``x``, that was previously defined as ``Dim("dim0_x")``::
|
||||
|
||||
dim = Dim("dim0_x", max=5)
|
||||
|
||||
This example means the generated code requires dimension 0 of input ``x`` to be less
|
||||
than or equal to 5 to be valid. You can inspect the suggested fixes to dynamic dimension
|
||||
definitions and then copy them verbatim into your code without needing to change the
|
||||
``dynamic_shapes`` argument to your :func:`export` call.
|
||||
|
||||
Args:
|
||||
mod: We will trace the forward method of this module.
|
||||
|
||||
args: Example positional inputs.
|
||||
|
||||
kwargs: Optional example keyword inputs.
|
||||
|
||||
dynamic_shapes:
|
||||
An optional argument where the type should either be:
|
||||
1) a dict from argument names of ``f`` to their dynamic shape specifications,
|
||||
2) a tuple that specifies dynamic shape specifications for each input in original order.
|
||||
If you are specifying dynamism on keyword args, you will need to pass them in the order that
|
||||
is defined in the original function signature.
|
||||
|
||||
The dynamic shape of a tensor argument can be specified as either
|
||||
(1) a dict from dynamic dimension indices to :func:`Dim` types, where it is
|
||||
not required to include static dimension indices in this dict, but when they are,
|
||||
they should be mapped to None; or (2) a tuple / list of :func:`Dim` types or None,
|
||||
where the :func:`Dim` types correspond to dynamic dimensions, and static dimensions
|
||||
are denoted by None. Arguments that are dicts or tuples / lists of tensors are
|
||||
recursively specified by using mappings or sequences of contained specifications.
|
||||
|
||||
strict: When disabled (default), the export function will trace the program through
|
||||
Python runtime, which by itself will not validate some of the implicit assumptions
|
||||
baked into the graph. It will still validate most critical assumptions like shape
|
||||
safety. When enabled (by setting ``strict=True``), the export function will trace
|
||||
the program through TorchDynamo which will ensure the soundness of the resulting
|
||||
graph. TorchDynamo has limited Python feature coverage, thus you may experience more
|
||||
errors. Note that toggling this argument does not affect the resulting IR spec to be
|
||||
different and the model will be serialized in the same way regardless of what value
|
||||
is passed here.
|
||||
|
||||
preserve_module_call_signature: A list of submodule paths for which the original
|
||||
calling conventions are preserved as metadata. The metadata will be used when calling
|
||||
torch.export.unflatten to preserve the original calling conventions of modules.
|
||||
|
||||
Returns:
|
||||
An :class:`ExportedProgram` containing the traced callable.
|
||||
|
||||
**Acceptable input/output types**
|
||||
|
||||
Acceptable types of inputs (for ``args`` and ``kwargs``) and outputs include:
|
||||
|
||||
- Primitive types, i.e. ``torch.Tensor``, ``int``, ``float``, ``bool`` and ``str``.
|
||||
- Dataclasses, but they must be registered by calling :func:`register_dataclass` first.
|
||||
- (Nested) Data structures comprising of ``dict``, ``list``, ``tuple``, ``namedtuple`` and
|
||||
``OrderedDict`` containing all above types.
|
||||
|
||||
"""
|
||||
from ._trace import _export
|
||||
|
||||
if not isinstance(mod, torch.nn.Module):
|
||||
raise ValueError(
|
||||
f"Expected `mod` to be an instance of `torch.nn.Module`, got {type(mod)}."
|
||||
)
|
||||
if isinstance(mod, torch.jit.ScriptModule):
|
||||
raise ValueError(
|
||||
"Exporting a ScriptModule is not supported. "
|
||||
"Maybe try converting your ScriptModule to an ExportedProgram "
|
||||
"using `TS2EPConverter(mod, args, kwargs).convert()` instead."
|
||||
)
|
||||
|
||||
try:
|
||||
return _export(
|
||||
mod,
|
||||
args,
|
||||
kwargs,
|
||||
dynamic_shapes,
|
||||
strict=strict,
|
||||
preserve_module_call_signature=preserve_module_call_signature,
|
||||
pre_dispatch=True,
|
||||
prefer_deferred_runtime_asserts_over_guards=prefer_deferred_runtime_asserts_over_guards,
|
||||
)
|
||||
except Exception as e:
|
||||
draft_export_msg = (
|
||||
"The error above occurred when calling torch.export.export. If you would "
|
||||
"like to view some more information about this error, and get a list "
|
||||
"of all other errors that may occur in your export call, you can "
|
||||
"replace your `export()` call with `draft_export()`."
|
||||
)
|
||||
|
||||
# For errors that we know can be caught by draft-export, add the message
|
||||
# to ask users to try out draft-export
|
||||
if isinstance(
|
||||
e,
|
||||
(
|
||||
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode,
|
||||
torch._subclasses.fake_tensor.UnsupportedOperatorException,
|
||||
torch._dynamo.exc.UserError,
|
||||
torch.fx.experimental.symbolic_shapes.ConstraintViolationError,
|
||||
),
|
||||
):
|
||||
new_msg = str(e) + "\n\n" + draft_export_msg
|
||||
e.args = (new_msg,)
|
||||
elif isinstance(e, RuntimeError) and "no fake impl registered" in str(e):
|
||||
new_msg = str(e) + "\n\n" + draft_export_msg
|
||||
e.args = (new_msg,)
|
||||
raise e
|
||||
|
||||
|
||||
DEFAULT_PICKLE_PROTOCOL = 2
|
||||
|
||||
|
||||
def save(
|
||||
ep: ExportedProgram,
|
||||
f: FileLike,
|
||||
*,
|
||||
extra_files: dict[str, Any] | None = None,
|
||||
opset_version: dict[str, int] | None = None,
|
||||
pickle_protocol: int = DEFAULT_PICKLE_PROTOCOL,
|
||||
) -> None:
|
||||
"""
|
||||
|
||||
.. warning::
|
||||
Under active development, saved files may not be usable in newer versions
|
||||
of PyTorch.
|
||||
|
||||
Saves an :class:`ExportedProgram` to a file-like object. It can then be
|
||||
loaded using the Python API :func:`torch.export.load <torch.export.load>`.
|
||||
|
||||
Args:
|
||||
ep (ExportedProgram): The exported program to save.
|
||||
|
||||
f (str | os.PathLike[str] | IO[bytes]) A file-like object (has to
|
||||
implement write and flush) or a string containing a file name.
|
||||
|
||||
extra_files (Optional[Dict[str, Any]]): Map from filename to contents
|
||||
which will be stored as part of f.
|
||||
|
||||
opset_version (Optional[Dict[str, int]]): A map of opset names
|
||||
to the version of this opset
|
||||
|
||||
pickle_protocol: can be specified to override the default protocol
|
||||
|
||||
Example::
|
||||
|
||||
import torch
|
||||
import io
|
||||
|
||||
|
||||
class MyModule(torch.nn.Module):
|
||||
def forward(self, x):
|
||||
return x + 10
|
||||
|
||||
|
||||
ep = torch.export.export(MyModule(), (torch.randn(5),))
|
||||
|
||||
# Save to file
|
||||
torch.export.save(ep, "exported_program.pt2")
|
||||
|
||||
# Save to io.BytesIO buffer
|
||||
buffer = io.BytesIO()
|
||||
torch.export.save(ep, buffer)
|
||||
|
||||
# Save with extra files
|
||||
extra_files = {"foo.txt": b"bar".decode("utf-8")}
|
||||
torch.export.save(ep, "exported_program.pt2", extra_files=extra_files)
|
||||
|
||||
"""
|
||||
if not isinstance(ep, ExportedProgram):
|
||||
raise TypeError(
|
||||
f"The 'ep' parameter must be an instance of 'ExportedProgram', got '{type(ep).__name__}' instead."
|
||||
)
|
||||
|
||||
from torch.export.pt2_archive._package import package_pt2
|
||||
|
||||
package_pt2(
|
||||
f,
|
||||
exported_programs={"model": ep},
|
||||
extra_files=extra_files,
|
||||
pickle_protocol=pickle_protocol,
|
||||
opset_version=opset_version,
|
||||
)
|
||||
|
||||
|
||||
def load(
|
||||
f: FileLike,
|
||||
*,
|
||||
extra_files: dict[str, Any] | None = None,
|
||||
expected_opset_version: dict[str, int] | None = None,
|
||||
) -> ExportedProgram:
|
||||
"""
|
||||
|
||||
.. warning::
|
||||
Under active development, saved files may not be usable in newer versions
|
||||
of PyTorch.
|
||||
|
||||
.. warning::
|
||||
:func:`torch.export.load()` uses pickle under the hood to load models. **Never load data from an untrusted source.**
|
||||
|
||||
Loads an :class:`ExportedProgram` previously saved with
|
||||
:func:`torch.export.save <torch.export.save>`.
|
||||
|
||||
Args:
|
||||
f (str | os.PathLike[str] | IO[bytes]): A file-like object (has to
|
||||
implement write and flush) or a string containing a file name.
|
||||
|
||||
extra_files (Optional[Dict[str, Any]]): The extra filenames given in
|
||||
this map would be loaded and their content would be stored in the
|
||||
provided map.
|
||||
|
||||
expected_opset_version (Optional[Dict[str, int]]): A map of opset names
|
||||
to expected opset versions
|
||||
|
||||
Returns:
|
||||
An :class:`ExportedProgram` object
|
||||
|
||||
Example::
|
||||
|
||||
import torch
|
||||
import io
|
||||
|
||||
# Load ExportedProgram from file
|
||||
ep = torch.export.load("exported_program.pt2")
|
||||
|
||||
# Load ExportedProgram from io.BytesIO object
|
||||
with open("exported_program.pt2", "rb") as f:
|
||||
buffer = io.BytesIO(f.read())
|
||||
buffer.seek(0)
|
||||
ep = torch.export.load(buffer)
|
||||
|
||||
# Load with extra files.
|
||||
extra_files = {"foo.txt": ""} # values will be replaced with data
|
||||
ep = torch.export.load("exported_program.pt2", extra_files=extra_files)
|
||||
print(extra_files["foo.txt"])
|
||||
print(ep(torch.randn(5)))
|
||||
"""
|
||||
if isinstance(f, (str, os.PathLike)):
|
||||
f = os.fspath(f)
|
||||
|
||||
extra_files = extra_files or {}
|
||||
|
||||
from torch.export.pt2_archive._package import load_pt2, PT2ArchiveContents
|
||||
|
||||
try:
|
||||
pt2_contents = load_pt2(
|
||||
f,
|
||||
expected_opset_version=expected_opset_version,
|
||||
)
|
||||
except RuntimeError:
|
||||
log.warning("Ran into the following error when deserializing", exc_info=True)
|
||||
pt2_contents = PT2ArchiveContents({}, {}, {})
|
||||
|
||||
if len(pt2_contents.exported_programs) > 0 or len(pt2_contents.extra_files) > 0:
|
||||
for k, v in pt2_contents.extra_files.items():
|
||||
extra_files[k] = v
|
||||
|
||||
return pt2_contents.exported_programs["model"]
|
||||
|
||||
# TODO: For backward compatibility, we support loading a zip file from 2.7. Delete this path in 2.9(?)
|
||||
with zipfile.ZipFile(f, "r") as zipf:
|
||||
if "version" not in zipf.namelist():
|
||||
raise RuntimeError(
|
||||
"We ran into an error when deserializing the saved file. "
|
||||
"Please check the warnings above for possible errors. "
|
||||
)
|
||||
|
||||
log.warning(
|
||||
"Trying to deserialize for the older format. This version of file is "
|
||||
"deprecated. Please generate a new pt2 saved file."
|
||||
)
|
||||
|
||||
# Check the version
|
||||
version = zipf.read("version").decode().split(".")
|
||||
from torch._export.serde.schema import (
|
||||
SCHEMA_VERSION, # todo change archive version to schema version
|
||||
)
|
||||
|
||||
if len(version) != len(SCHEMA_VERSION):
|
||||
raise AssertionError(
|
||||
"Version in the saved file has incorrect length, double check if the file is generated by torch.export.save()"
|
||||
)
|
||||
if version[0] != str(SCHEMA_VERSION[0]):
|
||||
raise RuntimeError(
|
||||
f"Serialized version {version} does not match our current "
|
||||
f"schema version {SCHEMA_VERSION}."
|
||||
)
|
||||
|
||||
from torch._export.serde.serialize import deserialize, SerializedArtifact
|
||||
|
||||
# Load serialized_ep and serialized_state_dict from the zip file
|
||||
|
||||
serialized_exported_program: bytes | None = None
|
||||
serialized_state_dict: bytes | None = None
|
||||
serialized_constants: bytes | None = None
|
||||
serialized_example_inputs: bytes | None = None
|
||||
|
||||
for file_info in zipf.infolist():
|
||||
file_content = zipf.read(file_info.filename)
|
||||
|
||||
if file_info.filename == "serialized_exported_program.json":
|
||||
serialized_exported_program = file_content
|
||||
elif file_info.filename == "serialized_state_dict.json":
|
||||
warnings.warn("This version of file is deprecated", stacklevel=2)
|
||||
serialized_state_dict = file_content
|
||||
elif file_info.filename == "serialized_constants.json":
|
||||
warnings.warn("This version of file is deprecated", stacklevel=2)
|
||||
serialized_constants = file_content
|
||||
elif file_info.filename == "serialized_state_dict.pt":
|
||||
serialized_state_dict = file_content
|
||||
elif file_info.filename == "serialized_constants.pt":
|
||||
serialized_constants = file_content
|
||||
elif file_info.filename == "serialized_example_inputs.pt":
|
||||
serialized_example_inputs = file_content
|
||||
elif file_info.filename.startswith("extra_files"):
|
||||
filename = file_info.filename.split("/", 1)[1]
|
||||
extra_files[filename] = file_content.decode("utf-8")
|
||||
|
||||
if serialized_exported_program is None:
|
||||
raise AssertionError("serialized_exported_program is None")
|
||||
if serialized_state_dict is None:
|
||||
raise AssertionError("serialized_state_dict is None")
|
||||
if serialized_constants is None:
|
||||
raise AssertionError("serialized_constants is None")
|
||||
if serialized_example_inputs is None:
|
||||
raise AssertionError("serialized_example_inputs is None")
|
||||
artifact: SerializedArtifact = SerializedArtifact(
|
||||
serialized_exported_program,
|
||||
serialized_state_dict,
|
||||
serialized_constants,
|
||||
serialized_example_inputs,
|
||||
)
|
||||
|
||||
# Deserialize ExportedProgram
|
||||
ep = deserialize(artifact, expected_opset_version)
|
||||
|
||||
return ep
|
||||
|
||||
|
||||
def draft_export(
|
||||
mod: torch.nn.Module,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: Mapping[str, Any] | None = None,
|
||||
*,
|
||||
dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None = None,
|
||||
preserve_module_call_signature: tuple[str, ...] = (),
|
||||
strict: bool = False,
|
||||
prefer_deferred_runtime_asserts_over_guards: bool = False,
|
||||
) -> ExportedProgram:
|
||||
"""
|
||||
A version of torch.export.export which is designed to consistently produce
|
||||
an ExportedProgram, even if there are potential soundness issues, and to
|
||||
generate a report listing the issues found.
|
||||
"""
|
||||
from ._draft_export import draft_export
|
||||
|
||||
return draft_export(
|
||||
mod=mod,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
dynamic_shapes=dynamic_shapes,
|
||||
preserve_module_call_signature=preserve_module_call_signature,
|
||||
strict=strict,
|
||||
prefer_deferred_runtime_asserts_over_guards=prefer_deferred_runtime_asserts_over_guards,
|
||||
)
|
||||
|
||||
|
||||
def register_dataclass(
|
||||
cls: type[Any],
|
||||
*,
|
||||
serialized_type_name: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Registers a dataclass as a valid input/output type for :func:`torch.export.export`.
|
||||
|
||||
Args:
|
||||
cls: the dataclass type to register
|
||||
serialized_type_name: The serialized name for the dataclass. This is
|
||||
required if you want to serialize the pytree TreeSpec containing this
|
||||
dataclass.
|
||||
|
||||
Example::
|
||||
|
||||
import torch
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class InputDataClass:
|
||||
feature: torch.Tensor
|
||||
bias: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class OutputDataClass:
|
||||
res: torch.Tensor
|
||||
|
||||
|
||||
torch.export.register_dataclass(InputDataClass)
|
||||
torch.export.register_dataclass(OutputDataClass)
|
||||
|
||||
|
||||
class Mod(torch.nn.Module):
|
||||
def forward(self, x: InputDataClass) -> OutputDataClass:
|
||||
res = x.feature + x.bias
|
||||
return OutputDataClass(res=res)
|
||||
|
||||
|
||||
ep = torch.export.export(Mod(), (InputDataClass(torch.ones(2, 2), 1),))
|
||||
print(ep)
|
||||
|
||||
"""
|
||||
pytree.register_dataclass(cls, serialized_type_name=serialized_type_name)
|
||||
@@ -0,0 +1,545 @@
|
||||
import getpass
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import IntEnum
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch._logging._internal
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._dynamo.exc import UserError, UserErrorType
|
||||
from torch._export.passes.insert_custom_op_guards import (
|
||||
get_op_profiles,
|
||||
insert_custom_op_guards,
|
||||
OpProfile,
|
||||
)
|
||||
from torch._utils_internal import log_draft_export_usage
|
||||
|
||||
from ._trace import _export, get_ep_stats
|
||||
from .dynamic_shapes import _DimHint, _DimHintType, Dim
|
||||
from .exported_program import ExportedProgram
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FailureType(IntEnum):
|
||||
MISSING_FAKE_KERNEL = 1
|
||||
DATA_DEPENDENT_ERROR = 2
|
||||
GUARD_ADDED = 3
|
||||
MISMATCHED_FAKE_KERNEL = 4
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.name
|
||||
|
||||
|
||||
def prettify_stack(stack: list[dict[str, str]], str_to_filename: dict[int, str]) -> str:
|
||||
res = ""
|
||||
for frame in stack:
|
||||
if frame["filename"] not in str_to_filename:
|
||||
continue
|
||||
|
||||
res += f"""
|
||||
File {str_to_filename[frame["filename"]]}, lineno {frame["line"]}, in {frame["name"]}""" # type: ignore[index]
|
||||
|
||||
res += f"\n {stack[-1]['loc']}"
|
||||
return res
|
||||
|
||||
|
||||
def prettify_frame_locals(
|
||||
loc: str, locals: dict[str, Any], symbols: dict[str, Any]
|
||||
) -> str:
|
||||
local_str = "\n".join(f" {k}: {v}" for k, v in locals.items())
|
||||
res = f"""
|
||||
Locals:
|
||||
{local_str}
|
||||
"""
|
||||
if any(v is not None for v in symbols.values()):
|
||||
symbol_str = "\n".join(
|
||||
f" {k}: {v}" for k, v in symbols.items() if v is not None
|
||||
)
|
||||
res += f"""
|
||||
Symbols:
|
||||
{symbol_str}
|
||||
"""
|
||||
return res
|
||||
|
||||
|
||||
def get_loc(filename: str, lineno: int) -> str | None:
|
||||
try:
|
||||
with open(filename) as f:
|
||||
for i, line in enumerate(f):
|
||||
if i == lineno - 1:
|
||||
return line.strip()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
class FailureReport:
|
||||
def __init__(
|
||||
self, failure_type: FailureType, data: dict[str, Any], xfail: bool = False
|
||||
) -> None:
|
||||
self.failure_type: FailureType = failure_type
|
||||
self.data: dict[str, Any] = data
|
||||
self.xfail: bool = xfail
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"FailureReport(failure_type={self.failure_type}, xfail={self.xfail}, data={self.data})"
|
||||
|
||||
def print(self, str_to_filename: dict[int, str]) -> str:
|
||||
if self.failure_type == FailureType.MISSING_FAKE_KERNEL:
|
||||
op = self.data["op"]
|
||||
|
||||
return f"""Missing fake kernel.
|
||||
torch.ops.{op} is missing a fake kernel implementation.
|
||||
|
||||
Please refer to https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ahugy69p2jmz for more detailed instructions on how to write a meta implementation.
|
||||
""" # noqa: B950
|
||||
|
||||
elif self.failure_type == FailureType.GUARD_ADDED:
|
||||
locals_info = (
|
||||
prettify_frame_locals(**self.data["frame_locals"])
|
||||
if self.data["frame_locals"]
|
||||
else ""
|
||||
)
|
||||
return f"""Guard Added.
|
||||
A guard was added during tracing, which might've resulted in some incorrect
|
||||
tracing or constraint violation error.
|
||||
Specifically, this guard was added: {self.data["expr"]}, where {self.data["symbol_to_sources"]}.
|
||||
This occurred at the following stacktrace: {prettify_stack(self.data["user_stack"], str_to_filename)}:
|
||||
{locals_info}
|
||||
And the following framework stacktrace: {prettify_stack(self.data["stack"], str_to_filename)}\n
|
||||
Because of this, we have modified the dynamic shapes structure to be the
|
||||
following. You can also use torch.export.Dim.AUTO instead to specify your
|
||||
dynamic shapes, and we will automatically infer the dynamism for you.
|
||||
```
|
||||
dynamic_shapes = {self.data["new_dynamic_shapes"]}
|
||||
```
|
||||
"""
|
||||
|
||||
elif self.failure_type == FailureType.DATA_DEPENDENT_ERROR:
|
||||
locals_info = (
|
||||
prettify_frame_locals(**self.data["frame_locals"])
|
||||
if self.data["frame_locals"]
|
||||
else ""
|
||||
)
|
||||
return f"""Data dependent error.
|
||||
When exporting, we were unable to evaluate the value of `{self.data["expr"]}`.
|
||||
This was encountered {self.data["occurrences"]} times.
|
||||
This occurred at the following user stacktrace: {prettify_stack(self.data["user_stack"], str_to_filename)}
|
||||
{locals_info}
|
||||
And the following framework stacktrace: {prettify_stack(self.data["stack"], str_to_filename)}\n
|
||||
As a result, it was specialized to a constant (e.g. `{self.data["result"]}` in the 1st occurrence), and asserts were inserted into the graph.
|
||||
|
||||
Please add `torch._check(...)` to the original code to assert this data-dependent assumption.
|
||||
Please refer to https://docs.google.com/document/d/1kZ_BbB3JnoLbUZleDT6635dHs88ZVYId8jT-yTFgf3A/edit#heading=h.boi2xurpqa0o for more details.
|
||||
""" # noqa: B950
|
||||
|
||||
elif self.failure_type == FailureType.MISMATCHED_FAKE_KERNEL:
|
||||
op = self.data["op"]
|
||||
reason = self.data["reason"]
|
||||
return f"""Mismatched fake kernel.
|
||||
torch.ops.{op} has a fake kernel implementation, but it has incorrect behavior, based on the real kernel.
|
||||
The reason for the mismatch is: {reason}.
|
||||
|
||||
Please refer to https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ahugy69p2jmz for more detailed instructions on how to write a fake implementation.
|
||||
""" # noqa: B950
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown failure type: {self.failure_type}")
|
||||
|
||||
|
||||
class DraftExportReport:
|
||||
def __init__(
|
||||
self,
|
||||
failures: list[FailureReport],
|
||||
str_to_filename: dict[int, str],
|
||||
expressions_created: dict[int, dict[str, Any]],
|
||||
op_profiles: dict[str, set[OpProfile]],
|
||||
):
|
||||
self.failures: list[FailureReport] = failures
|
||||
self.str_to_filename = str_to_filename
|
||||
self.expressions_created: dict[int, dict[str, Any]] = expressions_created
|
||||
self.op_profiles = op_profiles
|
||||
|
||||
def successful(self) -> bool:
|
||||
return len(self.failures) == 0 or all(
|
||||
failure.xfail for failure in self.failures
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"DraftExportReport({self.failures})"
|
||||
|
||||
def __str__(self) -> str:
|
||||
WARNING_COLOR = "\033[93m"
|
||||
GREEN_COLOR = "\033[92m"
|
||||
END_COLOR = "\033[0m"
|
||||
|
||||
if self.successful():
|
||||
return f"""{GREEN_COLOR}
|
||||
##############################################################################################
|
||||
Congratuations: No issues are found during export, and it was able to soundly produce a graph.
|
||||
You can now change back to torch.export.export()
|
||||
##############################################################################################
|
||||
{END_COLOR}"""
|
||||
|
||||
error = f"""{WARNING_COLOR}
|
||||
###################################################################################################
|
||||
WARNING: {len(self.failures)} issue(s) found during export, and it was not able to soundly produce a graph.
|
||||
Please follow the instructions to fix the errors.
|
||||
###################################################################################################
|
||||
|
||||
"""
|
||||
|
||||
for i, failure in enumerate(self.failures):
|
||||
error += f"{i + 1}. {failure.print(self.str_to_filename)}\n"
|
||||
error += END_COLOR
|
||||
return error
|
||||
|
||||
def apply_suggested_fixes(self) -> None:
|
||||
raise NotImplementedError("Not implemented yet")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExpressionCreatedNode:
|
||||
result_id: int
|
||||
argument_ids: list[int]
|
||||
record: dict[str, object]
|
||||
visited: bool = False
|
||||
|
||||
|
||||
class LogRecord:
|
||||
def __init__(self) -> None:
|
||||
self.log_count: dict[int, int] = {}
|
||||
self.logs: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
def _hash(self, element: tuple[str, dict[str, Any]]) -> int:
|
||||
key, data = element
|
||||
|
||||
if key == "missing_fake_kernel":
|
||||
return hash((key, data["op"]))
|
||||
elif key == "mismatched_fake_kernel":
|
||||
return hash((key, data["op"], data["reason"]))
|
||||
elif key == "propagate_real_tensors_provenance":
|
||||
return hash((key, json.dumps(data["user_stack"])))
|
||||
elif key == "guard_added":
|
||||
return hash((key, json.dumps(data["user_stack"])))
|
||||
elif key == "create_unbacked_symbol":
|
||||
return hash((key, json.dumps(data["user_stack"])))
|
||||
|
||||
return hash((key, json.dumps(data)))
|
||||
|
||||
def try_add(self, element: tuple[str, dict[str, str]]) -> bool:
|
||||
hash_value = self._hash(element)
|
||||
if hash_value in self.log_count:
|
||||
self.log_count[hash_value] += 1
|
||||
return False
|
||||
|
||||
self.log_count[hash_value] = 1
|
||||
self.logs.append(element)
|
||||
return True
|
||||
|
||||
def get_log_count(self, element: tuple[str, dict[str, Any]]) -> int:
|
||||
return self.log_count[self._hash(element)]
|
||||
|
||||
|
||||
class CaptureStructuredTrace(torch._logging._internal.LazyTraceHandler):
|
||||
def __init__(self) -> None:
|
||||
self.specific_log_keys = [
|
||||
"str",
|
||||
"exported_program",
|
||||
"propagate_real_tensors_provenance",
|
||||
"guard_added",
|
||||
"missing_fake_kernel",
|
||||
"mismatched_fake_kernel",
|
||||
"expression_created",
|
||||
"create_unbacked_symbol",
|
||||
]
|
||||
self.log_record: LogRecord = LogRecord()
|
||||
self.expression_created_logs: dict[int, ExpressionCreatedNode] = {}
|
||||
self.symbol_to_expressions: dict[str, list[dict[str, Any]]] = {}
|
||||
self.logger = logging.getLogger("torch.__trace")
|
||||
self.prev_get_dtrace = False
|
||||
|
||||
if root_dir := os.environ.get(torch._logging._internal.DTRACE_ENV_VAR):
|
||||
super().__init__(root_dir)
|
||||
else:
|
||||
sanitized_username = re.sub(r'[\\/:*?"<>|]', "_", getpass.getuser())
|
||||
root_dir = os.path.join(
|
||||
tempfile.gettempdir(),
|
||||
"export_" + sanitized_username,
|
||||
)
|
||||
super().__init__(root_dir)
|
||||
|
||||
self.setFormatter(torch._logging._internal.TorchLogsFormatter(trace=True))
|
||||
|
||||
def __enter__(self) -> "CaptureStructuredTrace":
|
||||
self.log_record = LogRecord()
|
||||
self.expression_created_logs = {}
|
||||
|
||||
# Remove the lazy trace handler if it exists
|
||||
possible_lazy_trace_handlers = [
|
||||
handler
|
||||
for handler in self.logger.handlers
|
||||
if isinstance(handler, torch._logging._internal.LazyTraceHandler)
|
||||
]
|
||||
for handler in possible_lazy_trace_handlers:
|
||||
self.logger.removeHandler(handler)
|
||||
|
||||
self.logger.addHandler(self)
|
||||
self.prev_get_dtrace = torch._logging._internal.GET_DTRACE_STRUCTURED
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
torch._logging._internal.GET_DTRACE_STRUCTURED = True
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback) -> None: # type: ignore[no-untyped-def]
|
||||
self.log_record = LogRecord()
|
||||
self.expression_created_logs = {}
|
||||
self.logger.removeHandler(self)
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
torch._logging._internal.GET_DTRACE_STRUCTURED = self.prev_get_dtrace
|
||||
self.prev_get_dtrace = False
|
||||
|
||||
def emit(self, record: Any) -> None:
|
||||
def _log_expression_created(
|
||||
emit_func: Callable[[Any], None], sym_node_id: int
|
||||
) -> None:
|
||||
# Log all the relevant expression_created logs
|
||||
if sym_node_id is None:
|
||||
return
|
||||
if res := self.expression_created_logs.get(sym_node_id, None):
|
||||
# Don't log the expression if we have already
|
||||
# printed it beforehand
|
||||
if not res.visited:
|
||||
res.visited = True
|
||||
for arg in res.argument_ids:
|
||||
_log_expression_created(emit_func, arg)
|
||||
|
||||
emit_func(res.record)
|
||||
|
||||
metadata = record.metadata
|
||||
for key in self.specific_log_keys:
|
||||
if key in metadata:
|
||||
if self.log_record.try_add((key, metadata[key])):
|
||||
if key == "expression_created":
|
||||
# We don't want to log all expression_created logs, only
|
||||
# the ones that are relevant to the
|
||||
# guards/propagate_real_tensor
|
||||
self.expression_created_logs[metadata[key]["result_id"]] = (
|
||||
ExpressionCreatedNode(
|
||||
metadata[key]["result_id"],
|
||||
metadata[key].get("argument_ids", []),
|
||||
record,
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
elif key == "propagate_real_tensors_provenance":
|
||||
_log_expression_created(
|
||||
super().emit, metadata[key].get("expr_node_id")
|
||||
)
|
||||
|
||||
elif key == "guard_added":
|
||||
if len(metadata[key]["symbol_to_sources"]) == 0:
|
||||
# We only want to include guards added that are relevant to
|
||||
# the symbolic shapes corresponding to the inputs which were
|
||||
# specified in the dynamic_shapes arg. These have a source.
|
||||
return
|
||||
elif metadata[key]["prefix"] == "runtime_assert":
|
||||
# This should've been captured by a
|
||||
# propagate_real_tensors log
|
||||
return
|
||||
|
||||
_log_expression_created(
|
||||
super().emit, metadata[key].get("expr_node_id")
|
||||
)
|
||||
|
||||
super().emit(record)
|
||||
|
||||
|
||||
def draft_export(
|
||||
mod: torch.nn.Module,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: Mapping[str, Any] | None = None,
|
||||
*,
|
||||
dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
|
||||
preserve_module_call_signature: tuple[str, ...] = (),
|
||||
strict: bool = False,
|
||||
pre_dispatch: bool = True,
|
||||
prefer_deferred_runtime_asserts_over_guards: bool = False,
|
||||
) -> ExportedProgram:
|
||||
start_time = time.time()
|
||||
kwargs = kwargs or {}
|
||||
dynamic_shapes = dynamic_shapes or {}
|
||||
|
||||
constraint_violation_msg = None
|
||||
capture_structured_log = CaptureStructuredTrace()
|
||||
|
||||
with (
|
||||
torch._functorch.config.patch(
|
||||
fake_tensor_propagate_real_tensors=True,
|
||||
generate_fake_kernels_from_real_mismatches=True,
|
||||
),
|
||||
capture_structured_log,
|
||||
):
|
||||
try:
|
||||
new_shapes = None
|
||||
ep = _export(
|
||||
mod,
|
||||
args,
|
||||
kwargs,
|
||||
dynamic_shapes=dynamic_shapes,
|
||||
strict=strict,
|
||||
pre_dispatch=pre_dispatch,
|
||||
preserve_module_call_signature=preserve_module_call_signature,
|
||||
prefer_deferred_runtime_asserts_over_guards=prefer_deferred_runtime_asserts_over_guards,
|
||||
)
|
||||
except Exception as exc:
|
||||
if (
|
||||
isinstance(exc, UserError)
|
||||
and exc.error_type == UserErrorType.CONSTRAINT_VIOLATION
|
||||
):
|
||||
constraint_violation_msg = exc.msg
|
||||
|
||||
def convert_dim_to_auto(dim: Any) -> Any:
|
||||
if isinstance(dim, Dim):
|
||||
return Dim.AUTO(min=dim.min, max=dim.max)
|
||||
elif isinstance(dim, _DimHint) and dim.type == _DimHintType.DYNAMIC:
|
||||
return Dim.AUTO(min=dim.min, max=dim.max)
|
||||
return dim
|
||||
|
||||
new_shapes = pytree.tree_map(convert_dim_to_auto, dynamic_shapes)
|
||||
ep = _export(
|
||||
mod,
|
||||
args,
|
||||
kwargs,
|
||||
dynamic_shapes=new_shapes,
|
||||
strict=strict,
|
||||
pre_dispatch=pre_dispatch,
|
||||
preserve_module_call_signature=preserve_module_call_signature,
|
||||
prefer_deferred_runtime_asserts_over_guards=prefer_deferred_runtime_asserts_over_guards,
|
||||
)
|
||||
else:
|
||||
log_draft_export_usage(
|
||||
error=True,
|
||||
export_time=time.time() - start_time,
|
||||
strict=strict,
|
||||
message=str(exc),
|
||||
type=f"{type(exc).__name__}.{type(exc).__qualname__}",
|
||||
)
|
||||
raise exc
|
||||
|
||||
torch._logging.dtrace_structured("exported_program", payload_fn=lambda: str(ep))
|
||||
|
||||
str_to_filename: dict[int, str] = {}
|
||||
failures: list[FailureReport] = []
|
||||
incorrect_custom_ops: set[str] = set()
|
||||
expressions_created: dict[int, dict[str, Any]] = {}
|
||||
|
||||
for log_name, log_contents in capture_structured_log.log_record.logs:
|
||||
failure_type = None
|
||||
|
||||
if log_name == "str":
|
||||
str_to_filename[log_contents[1]] = log_contents[0] # type: ignore[index]
|
||||
continue
|
||||
|
||||
elif log_name == "propagate_real_tensors_provenance":
|
||||
log_contents["occurrences"] = (
|
||||
capture_structured_log.log_record.get_log_count(
|
||||
(log_name, log_contents)
|
||||
)
|
||||
)
|
||||
|
||||
failure_type = FailureType.DATA_DEPENDENT_ERROR
|
||||
|
||||
elif log_name == "guard_added":
|
||||
if new_shapes is None:
|
||||
continue
|
||||
|
||||
failure_type = FailureType.GUARD_ADDED
|
||||
log_contents["new_dynamic_shapes"] = new_shapes
|
||||
elif log_name == "missing_fake_kernel":
|
||||
failure_type = FailureType.MISSING_FAKE_KERNEL
|
||||
incorrect_custom_ops.add(log_contents["op"])
|
||||
|
||||
elif log_name == "mismatched_fake_kernel":
|
||||
failure_type = FailureType.MISMATCHED_FAKE_KERNEL
|
||||
incorrect_custom_ops.add(log_contents["op"])
|
||||
|
||||
else:
|
||||
continue
|
||||
|
||||
if failure_type is None:
|
||||
raise AssertionError("failure_type cannot be None at this point")
|
||||
failures.append(
|
||||
FailureReport(
|
||||
failure_type,
|
||||
log_contents,
|
||||
)
|
||||
)
|
||||
|
||||
for k, v in capture_structured_log.expression_created_logs.items():
|
||||
if v.visited:
|
||||
expressions_created[k] = v.record
|
||||
|
||||
op_profiles = get_op_profiles(ep.graph_module, incorrect_custom_ops)
|
||||
report = DraftExportReport(
|
||||
failures, str_to_filename, expressions_created, op_profiles
|
||||
)
|
||||
|
||||
# Add asserts around custom ops
|
||||
insert_custom_op_guards(ep.graph_module, incorrect_custom_ops)
|
||||
|
||||
ep._report = report
|
||||
if not report.successful():
|
||||
log_filename = capture_structured_log.stream.name
|
||||
|
||||
warning_msg = f"""
|
||||
###################################################################################################
|
||||
WARNING: {len(report.failures)} issue(s) found during export, and it was not able to soundly produce a graph.
|
||||
To view the report of failures in an html page, please run the command:
|
||||
`tlparse {log_filename} --export`
|
||||
Or, you can view the errors in python by inspecting `print(ep._report)`.
|
||||
"""
|
||||
|
||||
if len(report.op_profiles) > 0:
|
||||
warning_msg += f"""
|
||||
While tracing we found {len(report.op_profiles)} operator(s) which do not have a fake kernel registered.
|
||||
If you intend to retrace the exported graph or run it with fake tensors, please run it under the
|
||||
following context manager, which will register a fake kernel for those operators.
|
||||
```
|
||||
with torch._library.fake_profile.unsafe_generate_fake_kernels(ep._report.op_profiles):
|
||||
# run with fake tensors
|
||||
```
|
||||
"""
|
||||
|
||||
warning_msg += """#################################################################################################"""
|
||||
|
||||
log.warning(warning_msg)
|
||||
|
||||
else:
|
||||
log.info(
|
||||
"""
|
||||
##############################################################################################
|
||||
Congratuations: No issues are found during export, and it was able to soundly produce a graph.
|
||||
You can now change back to torch.export.export()
|
||||
##############################################################################################
|
||||
"""
|
||||
)
|
||||
|
||||
log_draft_export_usage(
|
||||
error=False,
|
||||
export_time=time.time() - start_time,
|
||||
strict=strict,
|
||||
constraint_violations=constraint_violation_msg,
|
||||
report=ep._report,
|
||||
**get_ep_stats(ep),
|
||||
)
|
||||
return ep
|
||||
@@ -0,0 +1,113 @@
|
||||
import gc
|
||||
import types
|
||||
import typing
|
||||
import weakref
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
from torch.fx.experimental.symbolic_shapes import TrackedFake
|
||||
|
||||
|
||||
"""
|
||||
These functions are used to detect potential fake tensor leakage when using PT2 export.
|
||||
See NOTE [export non-strict fake tensor leak detection]
|
||||
|
||||
There are some complications that made this logic overly complicated:
|
||||
1) Python 3.10 and Python 3.12 have different ways of implementing referrer so
|
||||
we need to account for whether it is ref.__dict__ or the real ref object
|
||||
|
||||
2) There are some internal PT2 references to fake tensors like `TrackedFake`
|
||||
3) closures, generators, and bound methods can hold fake tensors.
|
||||
4) global object can hold onto a fake tensor
|
||||
|
||||
In general, these utils are our last resort to detect fake tensors. if the leak happens
|
||||
within the model attributes, we have a separate mechanism to detect. This tool relies a bit
|
||||
on garbage collector internal details, so I think it is unsafe to turn on by default, hence
|
||||
this tool should be used as debugging tool.
|
||||
"""
|
||||
|
||||
|
||||
# Things we never want to flag as leaks
|
||||
_SKIP_TYPES = (
|
||||
types.FrameType,
|
||||
types.ModuleType,
|
||||
)
|
||||
|
||||
|
||||
def _is_globals_or_locals(obj: typing.Any) -> bool:
|
||||
# These comparisons only make sense within this frame; still cheap to check.
|
||||
return obj is globals() or obj is locals()
|
||||
|
||||
|
||||
def _is_tracked_fake(obj: typing.Any) -> TypeIs[TrackedFake]:
|
||||
return isinstance(obj, TrackedFake)
|
||||
|
||||
|
||||
def _is_gm_meta_like_dict(d: dict, o: typing.Any) -> bool:
|
||||
# Hope gm.meta was a custom dict we can assert on
|
||||
return d.get("val") is o
|
||||
|
||||
|
||||
def _dict_is_attr_of_tracked_fake(d: dict) -> bool:
|
||||
"""
|
||||
Python 3.10 quirk: sometimes the referrer is obj.__dict__ instead of obj.
|
||||
Check if this dict is exactly the __dict__ of a TrackedFake.
|
||||
"""
|
||||
for parent in gc.get_referrers(d):
|
||||
if (
|
||||
hasattr(parent, "__dict__")
|
||||
and parent.__dict__ is d
|
||||
and _is_tracked_fake(parent)
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def find_legit_leaks_from_referrers(active_fakes: weakref.WeakSet) -> weakref.WeakSet:
|
||||
legit_leak: weakref.WeakSet = weakref.WeakSet()
|
||||
|
||||
# This is so that we don't falsely flag generator to be holding fake tensor
|
||||
fake_list = list(active_fakes)
|
||||
fake_list_id = id(fake_list)
|
||||
|
||||
for act in fake_list:
|
||||
# Track by id to avoid processing duplicate referrers
|
||||
seen = set()
|
||||
# Assume it's a leak unless we find only ignorable referrers
|
||||
flagged = False
|
||||
|
||||
for r in gc.get_referrers(act):
|
||||
rid = id(r)
|
||||
if rid in seen:
|
||||
continue
|
||||
seen.add(rid)
|
||||
|
||||
# Skip our own fake_list
|
||||
if rid == fake_list_id:
|
||||
continue
|
||||
|
||||
# Fast-path: skip obvious non-owners
|
||||
if _is_globals_or_locals(r):
|
||||
continue
|
||||
if isinstance(r, _SKIP_TYPES):
|
||||
continue
|
||||
if _is_tracked_fake(r):
|
||||
# TrackedFake should be ignored
|
||||
continue
|
||||
|
||||
# Handle dicts carefully (Python 3.10 sometimes shows __dict__)
|
||||
if isinstance(r, dict):
|
||||
if _is_gm_meta_like_dict(r, act):
|
||||
continue
|
||||
if _dict_is_attr_of_tracked_fake(r):
|
||||
continue
|
||||
flagged = True
|
||||
break
|
||||
|
||||
# Any other referrer we don't explicitly whitelist counts as a leak
|
||||
flagged = True
|
||||
break
|
||||
|
||||
if flagged:
|
||||
legit_leak.add(act)
|
||||
|
||||
return legit_leak
|
||||
@@ -0,0 +1,333 @@
|
||||
import contextlib
|
||||
from collections.abc import Generator
|
||||
|
||||
import torch
|
||||
from torch._decomp import global_decomposition_table
|
||||
from torch._decomp.decompositions import _rnn_helper, gather_params, gru_cell, lstm_cell
|
||||
from torch._higher_order_ops.while_loop import while_loop
|
||||
|
||||
|
||||
def one_layer_while_loop_lstm(inp, hidden, params, has_biases, reverse=False):
|
||||
"""
|
||||
1 layer fn for while loop LSTM
|
||||
|
||||
Args:
|
||||
inp: Input tensor of shape (seq_len, batch, input_size)
|
||||
hidden: Tuple of (hx, cx) hidden states
|
||||
params: List of weight and bias tensors
|
||||
has_biases: Whether biases are included
|
||||
reverse: Whether to process sequence in reverse
|
||||
|
||||
Returns:
|
||||
Tuple of (output, (final_hx, final_cx))
|
||||
"""
|
||||
ih_weight = params[0]
|
||||
hh_weight = params[1]
|
||||
ih_bias = params[2] if has_biases else None
|
||||
hh_bias = params[3] if has_biases else None
|
||||
hr_weight = (
|
||||
params[4] if len(params) == 5 else params[2] if len(params) == 3 else None
|
||||
)
|
||||
|
||||
hx = hidden[0].unsqueeze(0)
|
||||
cx = hidden[1].unsqueeze(0)
|
||||
|
||||
precomputed_input = torch.nn.functional.linear(inp, ih_weight, ih_bias)
|
||||
precomputed_input = precomputed_input.flip(0) if reverse else precomputed_input
|
||||
|
||||
# while loop rewrite
|
||||
step_output = torch.empty(
|
||||
precomputed_input.size(0),
|
||||
*tuple(hx.shape[1:]),
|
||||
dtype=hx.dtype,
|
||||
device=hx.device,
|
||||
)
|
||||
|
||||
def cond_fn(i, out, hx, cx):
|
||||
return i < precomputed_input.size(0)
|
||||
|
||||
def body_fn(idx, out, hx, cx):
|
||||
# Extract the integer value from idx and constrain it for data-dependent indexing
|
||||
i = idx.item()
|
||||
torch._check_is_size(i)
|
||||
torch._check_is_size(i, max=precomputed_input.size(0) - 1)
|
||||
hx, cx = lstm_cell(
|
||||
precomputed_input[i], hx, cx, hh_weight, hh_bias, hr_weight, chunk_dim=2
|
||||
)
|
||||
out = out.clone()
|
||||
# Squeeze the first dimension before storing (lstm_cell preserves the unsqueezed dim)
|
||||
out[i] = hx.squeeze(0)
|
||||
return idx + 1, out, hx, cx
|
||||
|
||||
cnt = torch.tensor(0, dtype=torch.int64)
|
||||
_, out, final_hx, final_cx = while_loop(
|
||||
cond_fn, body_fn, [cnt, step_output, hx, cx]
|
||||
)
|
||||
if reverse:
|
||||
out = out.flip(0)
|
||||
|
||||
# Use squeeze(1) to match original implementation
|
||||
return out, (final_hx.squeeze(1), final_cx.squeeze(1))
|
||||
|
||||
|
||||
def lstm_while_loop_impl(
|
||||
input,
|
||||
hx,
|
||||
params,
|
||||
has_biases,
|
||||
num_layers,
|
||||
dropout,
|
||||
train,
|
||||
bidirectional,
|
||||
batch_first,
|
||||
):
|
||||
"""
|
||||
LSTM implementation using while_loop for export compatibility.
|
||||
|
||||
This is a drop-in replacement for the default LSTM decomposition that uses
|
||||
while_loop instead of Python loops, making it more suitable for torch.export.
|
||||
|
||||
Args:
|
||||
input: Input tensor
|
||||
hx: Tuple of (h0, c0) hidden states
|
||||
params: List of weight and bias tensors
|
||||
has_biases: Whether biases are included
|
||||
num_layers: Number of LSTM layers
|
||||
dropout: Dropout probability
|
||||
train: Training mode
|
||||
bidirectional: Whether to use bidirectional LSTM
|
||||
batch_first: Whether batch dimension is first
|
||||
|
||||
Returns:
|
||||
Tuple of (output, h_n, c_n)
|
||||
"""
|
||||
if len(hx) != 2:
|
||||
raise AssertionError("lstm expects two hidden states")
|
||||
params = gather_params(params, has_biases, hx[0].size(2) != hx[1].size(2))
|
||||
hidden = list(zip(hx[0], hx[1]))
|
||||
layer_fn = one_layer_while_loop_lstm
|
||||
out, final_hiddens = _rnn_helper(
|
||||
input,
|
||||
hidden,
|
||||
params,
|
||||
has_biases,
|
||||
num_layers,
|
||||
dropout,
|
||||
train,
|
||||
bidirectional,
|
||||
batch_first,
|
||||
layer_fn,
|
||||
)
|
||||
final_hiddens = list(zip(*final_hiddens))
|
||||
return out, torch.stack(final_hiddens[0], 0), torch.stack(final_hiddens[1], 0)
|
||||
|
||||
|
||||
def one_layer_while_loop_gru(inp, hidden, params, has_biases, reverse=False):
|
||||
"""
|
||||
1 layer fn for while loop GRU
|
||||
|
||||
Args:
|
||||
inp: Input tensor of shape (seq_len, batch, input_size)
|
||||
hidden: Hidden state tensor
|
||||
params: List of weight and bias tensors
|
||||
has_biases: Whether biases are included
|
||||
reverse: Whether to process sequence in reverse
|
||||
|
||||
Returns:
|
||||
Tuple of (output, final_hidden)
|
||||
"""
|
||||
ih_weight = params[0]
|
||||
hh_weight = params[1]
|
||||
ih_bias = params[2] if has_biases else None
|
||||
hh_bias = params[3] if has_biases else None
|
||||
|
||||
precomputed_input = torch.nn.functional.linear(inp, ih_weight, ih_bias)
|
||||
precomputed_input = precomputed_input.flip(0) if reverse else precomputed_input
|
||||
cur_hidden = hidden.unsqueeze(0)
|
||||
|
||||
# while loop rewrite
|
||||
step_output = torch.empty(
|
||||
precomputed_input.size(0),
|
||||
*tuple(cur_hidden.shape[1:]),
|
||||
dtype=cur_hidden.dtype,
|
||||
device=cur_hidden.device,
|
||||
)
|
||||
|
||||
def cond_fn(i, out, cur_hidden):
|
||||
return i < precomputed_input.size(0)
|
||||
|
||||
def body_fn(idx, out, cur_hidden):
|
||||
# Extract the integer value from idx and constrain it for data-dependent indexing
|
||||
i = idx.item()
|
||||
torch._check_is_size(i)
|
||||
torch._check_is_size(i, max=precomputed_input.size(0) - 1)
|
||||
cur_hidden = gru_cell(
|
||||
precomputed_input[i], cur_hidden, ih_weight, ih_bias, hh_weight, hh_bias
|
||||
)
|
||||
out = out.clone()
|
||||
out[i] = cur_hidden.squeeze(0)
|
||||
return idx + 1, out, cur_hidden
|
||||
|
||||
cnt = torch.tensor(0, dtype=torch.int64)
|
||||
_, out, final_hidden = while_loop(cond_fn, body_fn, [cnt, step_output, cur_hidden])
|
||||
if reverse:
|
||||
out = out.flip(0)
|
||||
|
||||
return out, final_hidden.squeeze(0)
|
||||
|
||||
|
||||
def gru_while_loop_impl(
|
||||
input,
|
||||
hx,
|
||||
params,
|
||||
has_biases,
|
||||
num_layers,
|
||||
dropout,
|
||||
train,
|
||||
bidirectional,
|
||||
batch_first,
|
||||
):
|
||||
"""
|
||||
GRU implementation using while_loop for export compatibility.
|
||||
|
||||
This is a drop-in replacement for the default GRU decomposition that uses
|
||||
while_loop instead of Python loops, making it more suitable for torch.export.
|
||||
|
||||
Args:
|
||||
input: Input tensor
|
||||
hx: Hidden state tensor
|
||||
params: List of weight and bias tensors
|
||||
has_biases: Whether biases are included
|
||||
num_layers: Number of GRU layers
|
||||
dropout: Dropout probability
|
||||
train: Training mode
|
||||
bidirectional: Whether to use bidirectional GRU
|
||||
batch_first: Whether batch dimension is first
|
||||
|
||||
Returns:
|
||||
Tuple of (output, h_n)
|
||||
"""
|
||||
params = gather_params(params, has_biases, False)
|
||||
hidden = list(hx.unbind(0))
|
||||
layer_fn = one_layer_while_loop_gru
|
||||
out, final_hiddens = _rnn_helper(
|
||||
input,
|
||||
hidden,
|
||||
params,
|
||||
has_biases,
|
||||
num_layers,
|
||||
dropout,
|
||||
train,
|
||||
bidirectional,
|
||||
batch_first,
|
||||
layer_fn,
|
||||
)
|
||||
return out, torch.stack(final_hiddens, 0)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _register_rnn_while_loop_decomposition(
|
||||
rnn_op, rnn_impl
|
||||
) -> Generator[None, None, None]:
|
||||
"""
|
||||
Generic context manager for registering while_loop-based RNN decompositions.
|
||||
|
||||
Args:
|
||||
rnn_op: The aten operation to patch (e.g., torch.ops.aten.lstm.input)
|
||||
rnn_impl: The while_loop-based implementation function
|
||||
|
||||
Note:
|
||||
This is an internal helper. Use register_lstm_while_loop_decomposition()
|
||||
or register_gru_while_loop_decomposition() instead.
|
||||
"""
|
||||
registry = global_decomposition_table["post_autograd"]
|
||||
|
||||
# Save the original decomposition if it exists
|
||||
original_decomp = registry.get(rnn_op, None)
|
||||
|
||||
# Save the original py_kernel if it exists
|
||||
original_py_kernel = rnn_op.py_kernels.get(
|
||||
torch._C.DispatchKey.CompositeImplicitAutograd, None
|
||||
)
|
||||
|
||||
try:
|
||||
# Register our while_loop-based implementation
|
||||
registry[rnn_op] = rnn_impl
|
||||
rnn_op.py_kernels[torch._C.DispatchKey.CompositeImplicitAutograd] = rnn_impl
|
||||
yield
|
||||
finally:
|
||||
# Restore the original decomposition
|
||||
if original_decomp is not None:
|
||||
registry[rnn_op] = original_decomp
|
||||
else:
|
||||
# If there was no original, remove our registration
|
||||
registry.pop(rnn_op, None)
|
||||
|
||||
# Restore the original py_kernel
|
||||
if original_py_kernel is not None:
|
||||
rnn_op.py_kernels[torch._C.DispatchKey.CompositeImplicitAutograd] = (
|
||||
original_py_kernel
|
||||
)
|
||||
else:
|
||||
# If there was no original, remove our registration
|
||||
rnn_op.py_kernels.pop(torch._C.DispatchKey.CompositeImplicitAutograd, None)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def register_lstm_while_loop_decomposition() -> Generator[None, None, None]:
|
||||
"""
|
||||
Context manager that temporarily registers the while_loop-based LSTM decomposition.
|
||||
|
||||
The while_loop-based decomposition is more suitable for export and graph-based
|
||||
execution, as it avoids Python control flow that cannot be captured in the graph.
|
||||
This should support dynamic sequence lengths, however as while_loop does not
|
||||
support Autograd yet, an ExportedProgram created with this will not be trainable.
|
||||
|
||||
Usage::
|
||||
|
||||
from torch.export._patches import register_lstm_while_loop_decomposition
|
||||
from torch.export import export
|
||||
|
||||
with register_lstm_while_loop_decomposition():
|
||||
# Export your model with LSTM
|
||||
ep = export(model, (x, h0, c0))
|
||||
|
||||
Note:
|
||||
This context manager temporarily modifies the global decomposition table
|
||||
and py_kernels registration. The original registrations are restored when
|
||||
exiting the context.
|
||||
"""
|
||||
with _register_rnn_while_loop_decomposition(
|
||||
torch.ops.aten.lstm.input, lstm_while_loop_impl
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def register_gru_while_loop_decomposition() -> Generator[None, None, None]:
|
||||
"""
|
||||
Context manager that temporarily registers the while_loop-based GRU decomposition.
|
||||
|
||||
The while_loop-based decomposition is more suitable for export and graph-based
|
||||
execution, as it avoids Python control flow that cannot be captured in the graph.
|
||||
This should support dynamic sequence lengths, however as while_loop does not
|
||||
support Autograd yet, an ExportedProgram created with this will not be trainable.
|
||||
|
||||
Usage::
|
||||
|
||||
from torch.export._patches import register_gru_while_loop_decomposition
|
||||
from torch.export import export
|
||||
|
||||
with register_gru_while_loop_decomposition():
|
||||
# Export your model with GRU
|
||||
ep = export(model, (x, h0))
|
||||
|
||||
Note:
|
||||
This context manager temporarily modifies the global decomposition table
|
||||
and py_kernels registration. The original registrations are restored when
|
||||
exiting the context.
|
||||
"""
|
||||
with _register_rnn_while_loop_decomposition(
|
||||
torch.ops.aten.gru.input, gru_while_loop_impl
|
||||
):
|
||||
yield
|
||||
+55
@@ -0,0 +1,55 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD-style license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
|
||||
import torch
|
||||
from torch._higher_order_ops.auto_functionalize import (
|
||||
auto_functionalized,
|
||||
auto_functionalized_v2,
|
||||
)
|
||||
from torch._inductor.fx_passes.post_grad import decompose_auto_functionalized
|
||||
from torch.export import ExportedProgram
|
||||
from torch.fx import Graph
|
||||
|
||||
|
||||
def remove_self_clone(graph: Graph) -> None:
|
||||
for node in graph.nodes:
|
||||
if node.target is torch.ops.aten.copy_.default and node.args[0] == node.args[1]:
|
||||
node.replace_all_uses_with(node.args[0])
|
||||
graph.erase_node(node)
|
||||
|
||||
|
||||
def unsafe_remove_auto_functionalized_pass(
|
||||
ep: ExportedProgram,
|
||||
) -> ExportedProgram:
|
||||
"""
|
||||
This pass removes an instances of the higher order op 'auto_functionalized',
|
||||
and modifies the calling EP inplace to have the original mutator op.
|
||||
This pass doesn't perform safety checks to make sure that this inplace mutation is safe.
|
||||
"""
|
||||
|
||||
with ep.graph_module._set_replace_hook(ep.graph_signature.get_replace_hook()):
|
||||
for module in ep.graph_module.modules():
|
||||
if not isinstance(module, torch.fx.GraphModule):
|
||||
continue
|
||||
for node in ep.graph.nodes:
|
||||
if (
|
||||
node.op == "call_function" and node.target is auto_functionalized
|
||||
) or (
|
||||
node.op == "call_function" and node.target is auto_functionalized_v2
|
||||
):
|
||||
func = node.args[0]
|
||||
if not isinstance(func, torch._ops.OpOverload):
|
||||
raise AssertionError(
|
||||
f"Expected func to be an OpOverload, but got {type(func)}"
|
||||
)
|
||||
# re-inplace everything
|
||||
node.meta["only_clone_these_tensors"] = []
|
||||
decompose_auto_functionalized(ep.graph)
|
||||
remove_self_clone(ep.graph)
|
||||
ep.graph.eliminate_dead_code()
|
||||
|
||||
return ep
|
||||
@@ -0,0 +1,237 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import operator
|
||||
|
||||
import torch
|
||||
from torch._higher_order_ops.effects import _get_schema, with_effects
|
||||
|
||||
from .exported_program import ExportedProgram
|
||||
from .graph_signature import (
|
||||
CustomObjArgument,
|
||||
InputKind,
|
||||
InputSpec,
|
||||
OutputKind,
|
||||
OutputSpec,
|
||||
TokenArgument,
|
||||
)
|
||||
|
||||
|
||||
def _get_custom_obj_for_node(node, inputs_to_lifted_custom_objs, constants):
|
||||
"""Extract the custom object from a node's arguments."""
|
||||
custom_obj_node = node
|
||||
custom_obj_meta = custom_obj_node.meta["val"] # type: ignore[union-attr]
|
||||
if not isinstance(custom_obj_meta, CustomObjArgument):
|
||||
raise AssertionError(
|
||||
f"Expected custom_obj_meta to be a CustomObjArgument, but got {type(custom_obj_meta)}"
|
||||
)
|
||||
|
||||
if custom_obj_meta.fake_val:
|
||||
return custom_obj_meta.fake_val
|
||||
elif custom_obj_node.name in inputs_to_lifted_custom_objs: # type: ignore[union-attr]
|
||||
return constants[inputs_to_lifted_custom_objs[custom_obj_node.name]] # type: ignore[union-attr]
|
||||
else:
|
||||
raise RuntimeError(f"Unable to find custom obj for node {node}")
|
||||
|
||||
|
||||
def _replace_with_effects_node(
|
||||
node, ep, inputs_to_lifted_custom_objs, output_tokens, input_tokens, module
|
||||
):
|
||||
"""Replace a with_effects node with the underlying function call."""
|
||||
# Get the input nodes
|
||||
token_node, func, *node_args = node.args
|
||||
if token_node.op == "placeholder":
|
||||
input_tokens.append(token_node)
|
||||
|
||||
if not isinstance(func, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)):
|
||||
raise AssertionError(
|
||||
f"Expected func to be an OpOverload or HigherOrderOperator, but got {type(func)}"
|
||||
)
|
||||
|
||||
# Get the schema for the function
|
||||
if func is torch.ops.higher_order.call_torchbind:
|
||||
custom_obj = _get_custom_obj_for_node(
|
||||
node_args[0], inputs_to_lifted_custom_objs, ep.constants
|
||||
)
|
||||
schema = _get_schema(func, [custom_obj] + node_args[1:])
|
||||
else:
|
||||
schema = _get_schema(func, node_args)
|
||||
|
||||
# Create the replacement node
|
||||
with module.graph.inserting_before(node):
|
||||
new_node = module.graph.call_function(func, tuple(node_args), node.kwargs)
|
||||
|
||||
# Update getitem nodes that extract outputs from with_effects
|
||||
for user in list(node.users.keys()):
|
||||
if user.target is not operator.getitem:
|
||||
raise AssertionError(
|
||||
f"Expected user target to be operator.getitem, but got {user.target}"
|
||||
)
|
||||
# getitem(with_effects, 0) is the token node
|
||||
if user.args[1] == 0:
|
||||
for user_user in list(user.users.keys()):
|
||||
if user_user.op == "output":
|
||||
output_tokens.append(user)
|
||||
|
||||
# Copy metadata from old node to new node
|
||||
for k, v in node.meta.items():
|
||||
new_node.meta[k] = v
|
||||
if k == "unbacked_bindings":
|
||||
# Remove the extra layer for effect token
|
||||
old_bindings = new_node.meta[k]
|
||||
new_bindings = {
|
||||
k: path[1:] if path else path for k, path in old_bindings.items()
|
||||
}
|
||||
new_node.meta[k] = new_bindings
|
||||
|
||||
# Fix up the getitem nodes based on return count
|
||||
if len(schema.returns) == 1:
|
||||
# Single return: replace getitem(with_effects, 1) with the node itself
|
||||
for user in list(node.users.keys()):
|
||||
if user.args[1] == 1:
|
||||
user.replace_all_uses_with(new_node)
|
||||
new_node.meta["val"] = node.meta["val"][1]
|
||||
elif len(schema.returns) > 1:
|
||||
# Multiple returns: shift getitem indices down by 1
|
||||
for user in list(node.users.keys()):
|
||||
if user.args[1] >= 1:
|
||||
user.args = (new_node, user.args[1] - 1)
|
||||
new_node.meta["val"] = node.meta["val"][1:]
|
||||
else:
|
||||
# No returns
|
||||
if len(schema.returns) != 0:
|
||||
raise AssertionError(
|
||||
f"Expected schema.returns to be empty, but got {len(schema.returns)} returns"
|
||||
)
|
||||
if len(new_node.users) != 0:
|
||||
raise AssertionError(
|
||||
f"Expected new_node to have no users, but got {len(new_node.users)} users"
|
||||
)
|
||||
new_node.meta["val"] = None
|
||||
|
||||
|
||||
def _replace_invoke_subgraph_node(node, module, output_tokens, input_tokens):
|
||||
"""Replace an invoke_subgraph node to remove the token argument."""
|
||||
if node.args[0].op != "get_attr":
|
||||
raise AssertionError(
|
||||
f"Expected node.args[0].op to be 'get_attr', but got {node.args[0].op}"
|
||||
)
|
||||
submod = getattr(module, node.args[0].target)
|
||||
if not submod.meta.get("has_with_effects", False):
|
||||
return
|
||||
|
||||
# Remove token from inputs
|
||||
subgraph, identifier, token, *operands = node.args
|
||||
node.args = (subgraph, identifier, *operands)
|
||||
if token.op == "placeholder":
|
||||
input_tokens.append(token)
|
||||
|
||||
# Update getitem nodes to account for removed token output
|
||||
for user in list(node.users.keys()):
|
||||
if user.args[1] >= 1:
|
||||
user.args = (node, user.args[1] - 1)
|
||||
elif user.args[1] == 0:
|
||||
for user_user in list(user.users.keys()):
|
||||
if user_user.op == "output":
|
||||
output_tokens.append(user)
|
||||
|
||||
|
||||
def _remove_effect_tokens(ep: ExportedProgram) -> ExportedProgram:
|
||||
"""
|
||||
Removes the existence of tokens from the exported program, including:
|
||||
- Removes the input and output tokens
|
||||
- Replaces with_effects(token, func, args) with just func(args)
|
||||
|
||||
This function does an inplace modification on the given ExportedProgram.
|
||||
"""
|
||||
inputs_to_lifted_custom_objs = ep.graph_signature.inputs_to_lifted_custom_objs
|
||||
|
||||
# mark submodules with effects as having effects. This will be used in the following pass to remove effects from subgraphs
|
||||
for _, module in ep.graph_module.named_modules():
|
||||
if not isinstance(module, torch.fx.GraphModule):
|
||||
continue
|
||||
|
||||
with_effect_nodes = [
|
||||
node for node in module.graph.nodes if node.target is with_effects
|
||||
]
|
||||
if len(with_effect_nodes) > 0:
|
||||
module.meta["has_with_effects"] = True
|
||||
|
||||
# Process each module with the replace hook to ensure graph signature is updated
|
||||
with ep.graph_module._set_replace_hook(ep.graph_signature.get_replace_hook()):
|
||||
for _, module in ep.graph_module.named_modules():
|
||||
if not isinstance(module, torch.fx.GraphModule):
|
||||
continue
|
||||
|
||||
input_tokens = []
|
||||
output_tokens = []
|
||||
|
||||
# Process with_effects and invoke_subgraph nodes
|
||||
for node in module.graph.nodes:
|
||||
if node.target is with_effects:
|
||||
_replace_with_effects_node(
|
||||
node,
|
||||
ep,
|
||||
inputs_to_lifted_custom_objs,
|
||||
output_tokens,
|
||||
input_tokens,
|
||||
module,
|
||||
)
|
||||
elif node.target is torch.ops.higher_order.invoke_subgraph:
|
||||
_replace_invoke_subgraph_node(
|
||||
node, module, output_tokens, input_tokens
|
||||
)
|
||||
|
||||
# Remove tokens from the output node
|
||||
if len(output_tokens) > 0:
|
||||
output_node = next(reversed(module.graph.find_nodes(op="output")))
|
||||
output_args = output_node.args[0]
|
||||
if len(output_args) < len(output_tokens):
|
||||
raise AssertionError(
|
||||
f"{output_args} output arguments found\n"
|
||||
f"{output_tokens} output tokens found\n"
|
||||
f"{module.graph}"
|
||||
)
|
||||
output_node.args = (tuple(output_args[len(output_tokens) :]),)
|
||||
|
||||
module.graph.eliminate_dead_code()
|
||||
|
||||
# Remove tokens from the input placeholders
|
||||
for node in module.graph.nodes:
|
||||
if node.op == "placeholder" and node in input_tokens:
|
||||
module.graph.erase_node(node)
|
||||
|
||||
module.recompile()
|
||||
|
||||
num_tokens: int = 0
|
||||
input_token_names: list[str] = []
|
||||
new_input_specs: list[InputSpec] = []
|
||||
for inp in ep.graph_signature.input_specs:
|
||||
if inp.kind == InputKind.TOKEN:
|
||||
num_tokens += 1
|
||||
if not isinstance(inp.arg, TokenArgument):
|
||||
raise AssertionError(
|
||||
f"Expected inp.arg to be a TokenArgument, but got {type(inp.arg)}"
|
||||
)
|
||||
input_token_names.append(inp.arg.name)
|
||||
else:
|
||||
new_input_specs.append(inp)
|
||||
|
||||
num_out_tokens: int = 0
|
||||
new_output_specs: list[OutputSpec] = []
|
||||
output_token_names: list[OutputSpec] = []
|
||||
for out in ep.graph_signature.output_specs:
|
||||
if out.kind == OutputKind.TOKEN:
|
||||
num_out_tokens += 1
|
||||
output_token_names.append(out.arg.name)
|
||||
else:
|
||||
new_output_specs.append(out)
|
||||
|
||||
# Update graph signature
|
||||
ep.graph_signature.input_specs = new_input_specs
|
||||
ep.graph_signature.output_specs = new_output_specs
|
||||
|
||||
if num_tokens != num_out_tokens:
|
||||
raise AssertionError(
|
||||
f"Number of input tokens ({num_tokens}) does not match output tokens ({num_out_tokens})"
|
||||
)
|
||||
|
||||
return ep
|
||||
@@ -0,0 +1,47 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode
|
||||
from torch.overrides import TorchFunctionMode
|
||||
|
||||
|
||||
class AutogradStateOpsFailSafeguard(TorchFunctionMode):
|
||||
"""
|
||||
Detect grad state ops during exporting the graph and fail the process by
|
||||
raising an error, to avoid unexpected behavior. Those grad mode ops could be:
|
||||
`torch.no_grad`
|
||||
`torch.enable_grad`
|
||||
`torch.set_grad_enabled`
|
||||
|
||||
Export with predispatch mode is exempted.
|
||||
"""
|
||||
|
||||
def __torch_function__(self, func, types, args=(), kwargs=None):
|
||||
kwargs = kwargs or {}
|
||||
unsupported_grad_mode_ops = [
|
||||
torch._C._set_grad_enabled,
|
||||
]
|
||||
# It's only enabled while tracing, by confirming the torch dispatch mode is
|
||||
# any active PROXY. This is to allow the autograd ops out of tracing.
|
||||
current_state = torch._C.is_grad_enabled()
|
||||
if func in unsupported_grad_mode_ops:
|
||||
if len(args) != 1:
|
||||
raise AssertionError(
|
||||
f"Expected exactly 1 argument for grad mode op, but got {len(args)}"
|
||||
)
|
||||
changed_state = args[0]
|
||||
mode = torch._C._get_dispatch_mode(torch._C._TorchDispatchModeKey.PROXY)
|
||||
# Intend to check if it's not the pre_dispatch mode. It's allowed to use
|
||||
# autograd ops in pre_dispatch mode, e.g. `torch.no_grad`
|
||||
if (
|
||||
mode
|
||||
and isinstance(mode, ProxyTorchDispatchMode)
|
||||
and not mode.pre_dispatch
|
||||
and changed_state != current_state
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Encountered autograd state manager op {func} trying to change global autograd state "
|
||||
"while exporting. This is unsafe because we don't capture this op in torch.export "
|
||||
"today, hence we can't reflect the user intention soundly. You can fix this by "
|
||||
"adding a torch.no_grad() context around the export call."
|
||||
)
|
||||
return func(*args, **kwargs)
|
||||
@@ -0,0 +1,187 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
State dict utilities for torch.export.
|
||||
|
||||
This module provides utilities for restoring state dicts to traced modules,
|
||||
ensuring that FQNs (Fully Qualified Names) match the original module structure.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.fx
|
||||
|
||||
|
||||
def _get_underlying_module(
|
||||
module_or_method: torch.nn.Module | Callable[..., Any],
|
||||
) -> torch.nn.Module:
|
||||
"""Extract the underlying nn.Module from either a module or a bound method.
|
||||
|
||||
Args:
|
||||
module_or_method: Either an nn.Module or a bound method of an nn.Module.
|
||||
|
||||
Returns:
|
||||
The underlying nn.Module.
|
||||
|
||||
Raises:
|
||||
TypeError: If module_or_method is neither an nn.Module nor a bound method.
|
||||
"""
|
||||
if isinstance(module_or_method, torch.nn.Module):
|
||||
return module_or_method
|
||||
# Handle bound methods (e.g., module.method)
|
||||
if (
|
||||
mod_self := getattr(module_or_method, "__self__", None)
|
||||
) is not None and isinstance(mod_self, torch.nn.Module):
|
||||
return mod_self
|
||||
raise TypeError(
|
||||
f"Expected nn.Module or bound method of nn.Module, got {type(module_or_method)}"
|
||||
)
|
||||
|
||||
|
||||
def _clear_traced_params_buffers(
|
||||
traced_module: torch.fx.GraphModule, const_keys: Sequence[str]
|
||||
) -> None:
|
||||
"""Remove all parameters and buffers from traced module before restoring.
|
||||
|
||||
For constants (parameters/buffers that don't need FQN mapping), this function
|
||||
removes them from the _buffers dict and re-assigns them as direct attributes.
|
||||
This ensures constants don't show up as buffers in the state dict.
|
||||
|
||||
Args:
|
||||
traced_module: The traced GraphModule to clean up.
|
||||
const_keys: List of keys that represent constants to be cleared.
|
||||
"""
|
||||
for key in const_keys:
|
||||
if key not in traced_module._buffers:
|
||||
raise AssertionError(f"Key {key} not found in traced_module._buffers")
|
||||
# We don't want constants to show up as a buffer in the state dict.
|
||||
# Instead they should just be a direct attribute.
|
||||
buffer = traced_module._buffers[key]
|
||||
del traced_module._buffers[key]
|
||||
# Note: setattr will register the value per nn.Module rules:
|
||||
# - If it's a Tensor, it'll be re-registered as a buffer (ends up back in _buffers).
|
||||
# - Otherwise, it becomes a plain attribute (not part of state_dict).
|
||||
setattr(traced_module, key, buffer)
|
||||
|
||||
|
||||
def _restore_state_dict(
|
||||
original_module: torch.nn.Module | Callable[..., Any],
|
||||
traced_module: torch.fx.GraphModule,
|
||||
) -> None:
|
||||
"""
|
||||
Restores the state dict of the traced module to match the original module exactly.
|
||||
|
||||
This function ensures that:
|
||||
1. Parameters and buffers in the traced module use the same FQNs (Fully Qualified Names)
|
||||
as the original module.
|
||||
2. The ordering of parameters/buffers matches the original module.
|
||||
3. Graph nodes referencing the old names are updated to use the correct FQNs.
|
||||
|
||||
This is useful after using functional tracing APIs (like dynamo_graph_capture_for_export)
|
||||
that may flatten parameter/buffer names.
|
||||
|
||||
Args:
|
||||
original_module: The original nn.Module (or a bound method of one) that was traced.
|
||||
traced_module: The traced fx.GraphModule whose state dict needs to be restored.
|
||||
|
||||
Example::
|
||||
|
||||
import torch
|
||||
from torch._dynamo.functional_export import _dynamo_graph_capture_for_export
|
||||
from torch.export import _restore_state_dict
|
||||
|
||||
|
||||
class Model(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.layer = torch.nn.Linear(10, 10)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layer(x)
|
||||
|
||||
|
||||
model = Model()
|
||||
gm = _dynamo_graph_capture_for_export(model)(torch.randn(1, 10))
|
||||
|
||||
# Before: gm may have flattened names like "p_layer_weight"
|
||||
# After: gm will have proper FQNs like "layer.weight"
|
||||
_restore_state_dict(model, gm)
|
||||
"""
|
||||
# Extract the underlying module if a bound method was passed
|
||||
module = _get_underlying_module(original_module)
|
||||
|
||||
# Build ID-based lookups for traced module params/buffers
|
||||
# Collect all data first to avoid modifying during iteration
|
||||
traced_params: dict[int, tuple[str, torch.nn.Parameter]] = {}
|
||||
for name, param in traced_module.named_parameters(remove_duplicate=False):
|
||||
traced_params[id(param)] = (name, param)
|
||||
|
||||
traced_buffers: dict[int, tuple[str, torch.Tensor]] = {}
|
||||
for name, buffer in traced_module.named_buffers(remove_duplicate=False):
|
||||
traced_buffers[id(buffer)] = (name, buffer)
|
||||
|
||||
# Collect original module's parameters and buffers upfront to avoid
|
||||
# issues with shared tensor objects during iteration
|
||||
orig_params_list: list[tuple[str, torch.nn.Parameter]] = list(
|
||||
module.named_parameters(remove_duplicate=False)
|
||||
)
|
||||
orig_buffers_list: list[tuple[str, torch.Tensor]] = list(
|
||||
module.named_buffers(remove_duplicate=False)
|
||||
)
|
||||
|
||||
# Build mapping from old names to new names for graph node updates
|
||||
name_mapping: dict[str, str] = {}
|
||||
|
||||
# Track which traced names we've processed
|
||||
processed_traced_names: set[str] = set()
|
||||
|
||||
# Restore parameters in the order they appear in original module
|
||||
for orig_name, orig_param in orig_params_list:
|
||||
if id(orig_param) in traced_params:
|
||||
# This param exists in traced module - restore it with original FQN
|
||||
traced_name, traced_param = traced_params[id(orig_param)]
|
||||
processed_traced_names.add(traced_name)
|
||||
if traced_name != orig_name:
|
||||
# Only reassign if the name is different
|
||||
torch.fx.graph_module._assign_attr(
|
||||
traced_param, traced_module, orig_name
|
||||
)
|
||||
torch.fx.graph_module._del_attr(traced_module, traced_name)
|
||||
name_mapping[traced_name] = orig_name
|
||||
else:
|
||||
# This param doesn't exist in traced module - add it
|
||||
torch.fx.graph_module._assign_attr(orig_param, traced_module, orig_name)
|
||||
|
||||
# Restore buffers in the order they appear in original module
|
||||
for orig_name, orig_buffer in orig_buffers_list:
|
||||
if id(orig_buffer) in traced_buffers:
|
||||
# This buffer exists in traced module - restore it with original FQN
|
||||
traced_name, traced_buffer = traced_buffers[id(orig_buffer)]
|
||||
processed_traced_names.add(traced_name)
|
||||
if traced_name != orig_name:
|
||||
# Only reassign if the name is different
|
||||
torch.fx.graph_module._assign_attr(
|
||||
orig_buffer, traced_module, orig_name
|
||||
)
|
||||
torch.fx.graph_module._del_attr(traced_module, traced_name)
|
||||
name_mapping[traced_name] = orig_name
|
||||
else:
|
||||
# This buffer doesn't exist in traced module - add it
|
||||
torch.fx.graph_module._assign_attr(orig_buffer, traced_module, orig_name)
|
||||
|
||||
param_names = [v[0] for v in traced_params.values()]
|
||||
buffer_names = [v[0] for v in traced_buffers.values()]
|
||||
# Constants are traced params/buffers that weren't matched to any original param/buffer
|
||||
const_keys = list(
|
||||
set(param_names + buffer_names).difference(processed_traced_names)
|
||||
)
|
||||
|
||||
_clear_traced_params_buffers(traced_module, const_keys)
|
||||
|
||||
# Update get_attr nodes in the graph to use the correct FQNs
|
||||
for node in traced_module.graph.nodes:
|
||||
if node.op == "get_attr" and node.target in name_mapping:
|
||||
node.target = name_mapping[node.target]
|
||||
|
||||
traced_module.recompile()
|
||||
@@ -0,0 +1,472 @@
|
||||
import logging
|
||||
import operator
|
||||
import types
|
||||
from collections import defaultdict
|
||||
|
||||
import torch
|
||||
import torch.fx._pytree as fx_pytree
|
||||
import torch.utils._pytree as pytree
|
||||
from torch.export.exported_program import (
|
||||
ConstantArgument,
|
||||
ExportedProgram,
|
||||
ModuleCallSignature,
|
||||
)
|
||||
from torch.fx.passes.tools_common import legalize_graph, NodeList
|
||||
from torch.fx.passes.utils.fuser_utils import erase_nodes, fuse_as_graphmodule
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _get_getitem_users(node: torch.fx.Node) -> set[torch.fx.Node]:
|
||||
node_users = list(node.users.keys())
|
||||
getitem_users = set()
|
||||
for user in node_users:
|
||||
if user.op == "output":
|
||||
continue
|
||||
|
||||
if not (user.op == "call_function" and user.target is operator.getitem):
|
||||
raise AssertionError(
|
||||
f"Expected getitem node as user for {node}, instead got {user}"
|
||||
)
|
||||
getitem_users.update(list(user.users.keys()))
|
||||
return getitem_users
|
||||
|
||||
|
||||
def _try_remove_connecting_pytrees(curr_module_node: torch.fx.Node) -> None:
|
||||
"""
|
||||
We want to try to remove extraneous pytree flatten/unflatten calls between modules
|
||||
calls. Instead of having the following:
|
||||
graph():
|
||||
...
|
||||
%foo : [num_users=1] = call_module[target=foo](args = (%getitem_1, %getitem_2), kwargs = {})
|
||||
%tree_flatten_spec : [num_users=1] = call_function[target=torch.fx._pytree.tree_flatten_spec](args = (%foo, %_spec_1), kwargs = {})
|
||||
%getitem_4 : [num_users=1] = call_function[target=operator.getitem](args = (%tree_flatten_spec, 0), kwargs = {})
|
||||
%tree_unflatten_1 : [num_users=2] = call_function[target=torch.utils._pytree.tree_unflatten](args = ([%getitem_4], %_spec_2), kwargs = {})
|
||||
%getitem_5 : [num_users=1] = call_function[target=operator.getitem](args = (%tree_unflatten_1, 0), kwargs = {})
|
||||
%getitem_7 : [num_users=0] = call_function[target=operator.getitem](args = (%tree_unflatten_1, 1), kwargs = {})
|
||||
%getitem_6 : [num_users=1] = call_function[target=operator.getitem](args = (%getitem_5, 0), kwargs = {})
|
||||
%bar : [num_users=1] = call_module[target=bar](args = (%getitem_6,), kwargs = {})
|
||||
...
|
||||
|
||||
We could do the following, if we know that all the outputs of `foo` feed into `bar`:
|
||||
graph():
|
||||
...
|
||||
%foo : [num_users=1] = call_module[target=foo](args = (%getitem_1, %getitem_2), kwargs = {})
|
||||
%bar : [num_users=1] = call_module[target=bar](args = (%getitem_6,), kwargs = {})
|
||||
...
|
||||
|
||||
Currently this optimization only works for the case where all of the outputs
|
||||
of `foo` go directly into `bar`, and `bar` has no other inputs.
|
||||
""" # noqa: B950
|
||||
|
||||
log.debug("Trying to remove pytrees for module call %s", curr_module_node)
|
||||
|
||||
curr_module_users = list(curr_module_node.users.keys())
|
||||
if len(curr_module_users) != 1:
|
||||
raise AssertionError(
|
||||
f"Expected only one user for module node, instead got {list(curr_module_users)}"
|
||||
)
|
||||
flatten_node = curr_module_users[0]
|
||||
if not (
|
||||
flatten_node.op == "call_function"
|
||||
and flatten_node.target is fx_pytree.tree_flatten_spec
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Expected flatten_node to be a call_function with target tree_flatten_spec, "
|
||||
f"but got op={flatten_node.op}, target={flatten_node.target}"
|
||||
)
|
||||
|
||||
flatten_getitem_users = _get_getitem_users(flatten_node)
|
||||
if len(flatten_getitem_users) != 1:
|
||||
log.debug(
|
||||
"More than one user found for flatten node, %s: %s. "
|
||||
"Unable to fuse it with another unflatten call.",
|
||||
flatten_node,
|
||||
flatten_getitem_users,
|
||||
)
|
||||
return
|
||||
|
||||
unflatten_node = next(iter(flatten_getitem_users))
|
||||
if not (
|
||||
unflatten_node.op == "call_function"
|
||||
and unflatten_node.target is pytree.tree_unflatten
|
||||
):
|
||||
log.debug(
|
||||
"Flatten node %s's user is not a pytree.tree_unflatten. "
|
||||
"Instead it is: %s. Passing...",
|
||||
flatten_node,
|
||||
unflatten_node,
|
||||
)
|
||||
return
|
||||
|
||||
for i, arg in enumerate(unflatten_node.args[0]): # type: ignore[union-attr,arg-type]
|
||||
if arg not in flatten_node.users:
|
||||
log.debug(
|
||||
"Module %s's outputs are not all directly used as inputs to "
|
||||
"the subsequent module. Unable to fuse the connecting "
|
||||
"flatten/unflatten. The inputs to the subsequent module are: %s. ",
|
||||
curr_module_node,
|
||||
unflatten_node.args[0],
|
||||
)
|
||||
return
|
||||
|
||||
if not (
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
arg.op == "call_function"
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
and arg.target is operator.getitem
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
and arg.args[1] == i
|
||||
):
|
||||
log.debug(
|
||||
"Module %s's outputs are not all directly used in the same "
|
||||
"order as outputted. Unable to fuse the connecting "
|
||||
"flatten/unflatten. The inputs to the "
|
||||
"subsequent module are: %s. ",
|
||||
curr_module_node,
|
||||
unflatten_node.args[0],
|
||||
)
|
||||
return
|
||||
|
||||
# Unflatten has two levels of getitem, because it gets the args and kwargs
|
||||
unflatten_getitem_getitem_users = set()
|
||||
unflatten_getitem_users = _get_getitem_users(unflatten_node)
|
||||
for unflatten_getitem_user in unflatten_getitem_users:
|
||||
unflatten_getitem_getitem_users.update(
|
||||
list(unflatten_getitem_user.users.keys())
|
||||
)
|
||||
|
||||
if len(unflatten_getitem_getitem_users) != 1:
|
||||
log.debug(
|
||||
"More than one user found for unflatten node, %s: %s. "
|
||||
"Unable to fuse it with another flatten call.",
|
||||
unflatten_node,
|
||||
unflatten_getitem_getitem_users,
|
||||
)
|
||||
return
|
||||
|
||||
next_module_node = next(iter(unflatten_getitem_getitem_users))
|
||||
if next_module_node.op != "call_module":
|
||||
log.debug(
|
||||
"Unflatten node %s's user is not a call_module. "
|
||||
"Instead it is: %s. Passing...",
|
||||
unflatten_node,
|
||||
next_module_node,
|
||||
)
|
||||
return
|
||||
|
||||
# Directly put the outputs of the current module into the next module
|
||||
next_module_node.args = (curr_module_node,)
|
||||
|
||||
|
||||
def _remove_extraneous_pytrees(gm: torch.fx.GraphModule) -> None:
|
||||
"""
|
||||
Remove extraneous pytree flatten/unflatten calls.
|
||||
|
||||
We try a couple of optimizations here:
|
||||
1. Remove pytree flatten/unflatten calls between modules
|
||||
2. TODO: Remove module's in_spec + initial unflatten call
|
||||
3. TODO: Remove module's out_spec + final flatten call
|
||||
"""
|
||||
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "call_module" and node.target != "_guards_fn":
|
||||
_try_remove_connecting_pytrees(node)
|
||||
|
||||
gm.graph.eliminate_dead_code()
|
||||
|
||||
|
||||
def _construct_inputs(
|
||||
gm: torch.fx.GraphModule,
|
||||
signature: ModuleCallSignature,
|
||||
node_name_map: dict[str, torch.fx.Node],
|
||||
) -> tuple[list[torch.fx.Node], dict[str, torch.fx.Node]]:
|
||||
tree_unflatten_args: list[torch.fx.Node | None] = []
|
||||
for input_ in signature.inputs:
|
||||
if isinstance(input_, ConstantArgument) and input_.value is None:
|
||||
# Constants should be directly embedded into the graph and not used
|
||||
# as inputs
|
||||
tree_unflatten_args.append(None)
|
||||
elif input_.name not in node_name_map:
|
||||
# For unused inputs
|
||||
tree_unflatten_args.append(None)
|
||||
else:
|
||||
tree_unflatten_args.append(node_name_map[input_.name])
|
||||
|
||||
# Insert unflatten call
|
||||
from .unflatten import _generate_unflatten
|
||||
|
||||
unflatten_node = _generate_unflatten(gm, tree_unflatten_args, signature.in_spec)
|
||||
|
||||
if signature.in_spec.num_children != 2:
|
||||
raise AssertionError(
|
||||
f"Expected in_spec to have 2 children, but got {signature.in_spec.num_children}"
|
||||
)
|
||||
if signature.in_spec.type is not tuple:
|
||||
raise AssertionError(
|
||||
f"Expected in_spec type to be tuple, but got {signature.in_spec.type}"
|
||||
)
|
||||
args_spec, kwargs_spec = signature.in_spec.children()
|
||||
if args_spec.type is not tuple:
|
||||
raise AssertionError(
|
||||
f"Expected args_spec type to be tuple, but got {args_spec.type}"
|
||||
)
|
||||
if kwargs_spec.type is not dict:
|
||||
raise AssertionError(
|
||||
f"Expected kwargs_spec type to be dict, but got {kwargs_spec.type}"
|
||||
)
|
||||
|
||||
args_node = gm.graph.call_function(operator.getitem, (unflatten_node, 0))
|
||||
args_nodes = [
|
||||
gm.graph.call_function(operator.getitem, (args_node, i))
|
||||
for i in range(args_spec.num_children)
|
||||
]
|
||||
kwargs_node = gm.graph.call_function(operator.getitem, (unflatten_node, 1))
|
||||
kwargs_nodes = {
|
||||
k: gm.graph.call_function(operator.getitem, (kwargs_node, k))
|
||||
for k in kwargs_spec.context
|
||||
}
|
||||
return args_nodes, kwargs_nodes
|
||||
|
||||
|
||||
def _insert_call_module(
|
||||
gm: torch.fx.GraphModule,
|
||||
args_nodes: list[torch.fx.Node],
|
||||
kwargs_nodes: dict[str, torch.fx.Node],
|
||||
module_to_swap: torch.nn.Module,
|
||||
name: str,
|
||||
) -> torch.fx.Node:
|
||||
from .unflatten import _assign_attr, _AttrKind
|
||||
|
||||
_assign_attr(module_to_swap, gm, name, _AttrKind.MODULE)
|
||||
module_node = gm.graph.call_module(name, tuple(args_nodes), kwargs_nodes) # type: ignore[arg-type]
|
||||
return module_node
|
||||
|
||||
|
||||
def _deconstruct_outputs(
|
||||
gm: torch.fx.GraphModule,
|
||||
signature: ModuleCallSignature,
|
||||
module_node: torch.fx.Node,
|
||||
node_name_map: dict[str, torch.fx.Node],
|
||||
orig_outputs: tuple[torch.fx.Node, ...],
|
||||
) -> None:
|
||||
from .unflatten import _generate_flatten_spec
|
||||
|
||||
flatten_node = _generate_flatten_spec(gm, module_node, signature.out_spec)
|
||||
|
||||
for i, orig_output in enumerate(orig_outputs):
|
||||
# Use Proxy to record getitem access.
|
||||
proxy_out = torch.fx.Proxy(flatten_node)[i].node # type: ignore[index]
|
||||
orig_output.replace_all_uses_with(proxy_out, propagate_meta=True)
|
||||
|
||||
node_name_map[orig_output.name] = proxy_out
|
||||
|
||||
|
||||
def _swap_module_helper(
|
||||
gm: torch.fx.GraphModule,
|
||||
modules_to_swap: dict[str, torch.nn.Module],
|
||||
module_call_graph: dict[str, ModuleCallSignature],
|
||||
) -> torch.fx.GraphModule:
|
||||
log.debug("Starting graph:")
|
||||
log.debug(gm.graph)
|
||||
|
||||
legalize_graph(gm)
|
||||
|
||||
partitions: dict[str, NodeList] = defaultdict(list)
|
||||
|
||||
node_name_map: dict[str, torch.fx.Node] = {
|
||||
node.name: node for node in gm.graph.nodes
|
||||
}
|
||||
|
||||
# TODO: Handle the duplicate module case
|
||||
for node in gm.graph.nodes:
|
||||
if nn_module_stack := node.meta.get("nn_module_stack"):
|
||||
for path, _ in nn_module_stack.values():
|
||||
if path in modules_to_swap:
|
||||
partitions[path].append(node)
|
||||
break
|
||||
|
||||
for name, nodes in partitions.items():
|
||||
"""
|
||||
Given a graph like the following, and we want to swap out the submodule "foo":
|
||||
graph():
|
||||
%x : [num_users=1] = placeholder[target=x]
|
||||
%y : [num_users=2] = placeholder[target=y]
|
||||
%add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%y, %x), kwargs = {}), nn_module_stack = {"foo": ("foo", torch.nn.Module)}
|
||||
%sub : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%y, %add), kwargs = {}), nn_module_stack = {"bar": ("bar", torch.nn.Module)}
|
||||
return (sub,)
|
||||
|
||||
We will first partition out foo's subgraph:
|
||||
graph():
|
||||
%x : [num_users=1] = placeholder[target=x]
|
||||
%y : [num_users=2] = placeholder[target=y]
|
||||
%add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%y, %x), kwargs = {})
|
||||
return add
|
||||
|
||||
And then insert an unflatten + call_module + flatten to replace the subgraph:
|
||||
graph():
|
||||
%x : [num_users=1] = placeholder[target=x]
|
||||
%y : [num_users=1] = placeholder[target=y]
|
||||
|
||||
%_spec_0 : [num_users=1] = get_attr[target=_spec_0]
|
||||
%tree_unflatten : [num_users=2] = call_function[target=torch.utils._pytree.tree_unflatten](args = ([%x, %y], %_spec_0), kwargs = {})
|
||||
%getitem : [num_users=2] = call_function[target=operator.getitem](args = (%tree_unflatten, 0), kwargs = {})
|
||||
%getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%getitem, 0), kwargs = {})
|
||||
%getitem_2 : [num_users=1] = call_function[target=operator.getitem](args = (%getitem, 1), kwargs = {})
|
||||
%getitem_3 : [num_users=0] = call_function[target=operator.getitem](args = (%tree_unflatten, 1), kwargs = {})
|
||||
%foo : [num_users=0] = call_module[target=foo](args = (%getitem_1, %getitem_2), kwargs = {})
|
||||
%_spec_1 : [num_users=1] = get_attr[target=_spec_1]
|
||||
%tree_flatten_spec : [num_users=1] = call_function[target=torch.fx._pytree.tree_flatten_spec](args = (None, %_spec_1), kwargs = {})
|
||||
%getitem_4 : [num_users=1] = call_function[target=operator.getitem](args = (%tree_flatten_spec, 0), kwargs = {})
|
||||
|
||||
%sub : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%y, %getitem_4), kwargs = {})
|
||||
return (%sub,)
|
||||
|
||||
The `tree_unflatten` call will construct tensor inputs into the input
|
||||
format needed by the swapped eager module.
|
||||
The `call_module` node should now reference the swapped torch.nn.Module.
|
||||
The `tree_flatten_spec` call will deconstruct the eager outputs of the
|
||||
swapped module into tensors.
|
||||
""" # noqa: B950
|
||||
|
||||
submod_name = name.replace(".", "_")
|
||||
sub_gm, orig_inputs, orig_outputs = fuse_as_graphmodule(
|
||||
gm, nodes, f"fused_{submod_name}"
|
||||
)
|
||||
|
||||
log.debug("Fused subgraph nodes:")
|
||||
log.debug(sub_gm.graph)
|
||||
|
||||
signature: ModuleCallSignature = module_call_graph[name]
|
||||
|
||||
args_nodes, kwargs_nodes = _construct_inputs(gm, signature, node_name_map)
|
||||
module_node = _insert_call_module(
|
||||
gm, args_nodes, kwargs_nodes, modules_to_swap[name], name
|
||||
)
|
||||
_deconstruct_outputs(gm, signature, module_node, node_name_map, orig_outputs)
|
||||
|
||||
erase_nodes(gm, nodes)
|
||||
|
||||
log.debug("Swapped graph:")
|
||||
log.debug(gm.graph)
|
||||
|
||||
legalize_graph(gm)
|
||||
|
||||
log.debug("Before removing extraneous pytrees:")
|
||||
log.debug(gm.graph)
|
||||
|
||||
_remove_extraneous_pytrees(gm)
|
||||
log.debug("After removing extraneous pytrees:")
|
||||
log.debug(gm.graph)
|
||||
|
||||
gm.recompile()
|
||||
|
||||
return gm
|
||||
|
||||
|
||||
def _fix_input_output_signature(
|
||||
gm: torch.fx.GraphModule, signature: ModuleCallSignature
|
||||
) -> None:
|
||||
"""
|
||||
Given the unlifted module from calling ep.module(), we want to remove the
|
||||
pytree processing from the graph module's PyTreeCodeGen and instead make it
|
||||
nodes inside of the graph. This allows us to do some optimizations, like
|
||||
remove these pytree calls if it is unnecessary, and makes the PyTree part
|
||||
more obvious to graph passes.
|
||||
"""
|
||||
from torch.export.unflatten import _generate_flatten, _generate_unflatten
|
||||
|
||||
# Remove the registered pytree codegen because we will take care of it
|
||||
# through inserting pytree nodes into the graph
|
||||
gm.graph._codegen = torch.fx.graph.CodeGen()
|
||||
|
||||
old_placeholders = [node for node in gm.graph.nodes if node.op == "placeholder"]
|
||||
|
||||
new_placeholders = []
|
||||
forward_arg_names = signature.forward_arg_names
|
||||
if forward_arg_names is None:
|
||||
forward_arg_names = []
|
||||
if signature.in_spec.num_children != 2:
|
||||
raise AssertionError(
|
||||
f"Expected in_spec to have 2 children, but got {signature.in_spec.num_children}"
|
||||
)
|
||||
arg_spec = signature.in_spec.child(0)
|
||||
kwarg_spec = signature.in_spec.child(1)
|
||||
if arg_spec.type is not tuple:
|
||||
raise AssertionError(
|
||||
f"Expected arg_spec type to be tuple, but got {arg_spec.type}"
|
||||
)
|
||||
if kwarg_spec.type is not dict:
|
||||
raise AssertionError(
|
||||
f"Expected kwarg_spec type to be dict, but got {kwarg_spec.type}"
|
||||
)
|
||||
for i in range(arg_spec.num_children):
|
||||
forward_arg_names.append(f"arg_{i}")
|
||||
forward_arg_names.extend(kwarg_spec.context)
|
||||
|
||||
for arg in forward_arg_names:
|
||||
with gm.graph.inserting_before(old_placeholders[0]):
|
||||
new_placeholders.append(gm.graph.placeholder(arg))
|
||||
|
||||
# Insert flatten call for the inputs
|
||||
with gm.graph.inserting_before(old_placeholders[0]):
|
||||
flat_node = _generate_flatten(gm, tuple(new_placeholders))
|
||||
for i, old_placeholder in enumerate(old_placeholders):
|
||||
old_placeholder.op = "call_function"
|
||||
old_placeholder.target = operator.getitem
|
||||
old_placeholder.args = (flat_node, i)
|
||||
|
||||
# Insert unflatten call for the outputs
|
||||
output_node = next(node for node in gm.graph.nodes if node.op == "output")
|
||||
with gm.graph.inserting_before(output_node):
|
||||
unflat = _generate_unflatten(gm, output_node.args[0], signature.out_spec)
|
||||
output_node.args = (unflat,)
|
||||
|
||||
gm.recompile()
|
||||
|
||||
|
||||
def _swap_modules(
|
||||
ep: ExportedProgram, modules_to_swap: dict[str, torch.nn.Module]
|
||||
) -> torch.fx.GraphModule:
|
||||
"""
|
||||
Unlifts the given ExportedProgram into a fx.GraphModule, and then swaps
|
||||
previously traced modules with new eager modules specified. Returns a
|
||||
fx.GraphModule with a custom forward function.
|
||||
|
||||
Args:
|
||||
ep (ExportedProgram): Exported program to modify
|
||||
modules_to_swap (Dict[str, torch.nn.Module]): Mapping from module fqn to
|
||||
eager module to swap with. The specified module fqn should have also
|
||||
been specified in the `preserve_module_call_signature` argument to
|
||||
torch.export so that we know how to restore the calling convention
|
||||
to this argument.
|
||||
run_with_interpreter: Whether or not to run the graph using
|
||||
fx.Interpreter. Setting to true will help result in better error
|
||||
messages and easier debugging, but it has found to result in a QPS
|
||||
drop.
|
||||
"""
|
||||
module_call_graph = {
|
||||
entry.fqn: entry.signature for entry in ep.module_call_graph if entry.signature
|
||||
}
|
||||
|
||||
gm = ep.module()
|
||||
gm.validate_inputs = False # type: ignore[assignment]
|
||||
gm.graph.eliminate_dead_code() # type: ignore[operator, union-attr]
|
||||
if not isinstance(gm, torch.fx.GraphModule):
|
||||
raise AssertionError(
|
||||
f"Expected gm to be a torch.fx.GraphModule, but got {type(gm)}"
|
||||
)
|
||||
_fix_input_output_signature(gm, ep.module_call_graph[0].signature)
|
||||
|
||||
gm.module_call_graph = ep.module_call_graph
|
||||
gm.train = types.MethodType(type(gm).train, gm) # type: ignore[assignment]
|
||||
gm.eval = types.MethodType(type(gm).eval, gm) # type: ignore[assignment]
|
||||
|
||||
if not isinstance(gm, torch.fx.GraphModule):
|
||||
raise AssertionError(
|
||||
f"Expected gm to be a torch.fx.GraphModule, but got {type(gm)}"
|
||||
)
|
||||
gm = _swap_module_helper(gm, modules_to_swap, module_call_graph)
|
||||
|
||||
return gm
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,72 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from torch.utils._pytree import Context, TreeSpec
|
||||
|
||||
|
||||
def reorder_kwargs(user_kwargs: dict[str, Any], spec: TreeSpec) -> dict[str, Any]:
|
||||
"""Reorder user-provided kwargs to match the order in `spec`. `spec` is
|
||||
expected to be the in_spec of an exported program, i.e. the spec that
|
||||
results from flattening `(args, kwargs)`.
|
||||
|
||||
We need this to provide consistent input ordering, such so that users can
|
||||
pass in foo(a=a, b=b) OR foo(b=b, a=a) and receive the same result.
|
||||
"""
|
||||
# Make sure that the spec is actually shaped like (args, kwargs)
|
||||
if spec.type is not tuple:
|
||||
raise AssertionError(f"Expected spec type to be tuple, but got {spec.type}")
|
||||
if spec.num_children != 2:
|
||||
raise AssertionError(
|
||||
f"Expected spec to have 2 children, but got {spec.num_children}"
|
||||
)
|
||||
kwargs_spec = spec.child(1)
|
||||
if kwargs_spec.type is not dict:
|
||||
raise AssertionError(
|
||||
f"Expected kwargs_spec type to be dict, but got {kwargs_spec.type}"
|
||||
)
|
||||
|
||||
if set(user_kwargs) != set(kwargs_spec.context):
|
||||
raise ValueError(
|
||||
f"Ran into a kwarg keyword mismatch: "
|
||||
f"Got the following keywords {list(user_kwargs)} but expected {kwargs_spec.context}"
|
||||
)
|
||||
|
||||
reordered_kwargs = {}
|
||||
for kw in kwargs_spec.context:
|
||||
reordered_kwargs[kw] = user_kwargs[kw]
|
||||
|
||||
return reordered_kwargs
|
||||
|
||||
|
||||
def is_equivalent(
|
||||
spec1: TreeSpec,
|
||||
spec2: TreeSpec,
|
||||
equivalence_fn: Callable[[type | None, Context, type | None, Context], bool],
|
||||
) -> bool:
|
||||
"""Customizable equivalence check for two TreeSpecs.
|
||||
|
||||
Arguments:
|
||||
spec1: The first TreeSpec to compare
|
||||
spec2: The second TreeSpec to compare
|
||||
equivalence_fn: A function to determine the equivalence of two
|
||||
TreeSpecs by examining their types and contexts. It will be called like:
|
||||
|
||||
equivalence_fn(spec1.type, spec1.context, spec2.type, spec2.context)
|
||||
|
||||
This function will be applied recursively to all children.
|
||||
|
||||
Returns:
|
||||
True if the two TreeSpecs are equivalent, False otherwise.
|
||||
"""
|
||||
if not equivalence_fn(spec1.type, spec1.context, spec2.type, spec2.context):
|
||||
return False
|
||||
|
||||
# Recurse on children
|
||||
if spec1.num_children != spec2.num_children:
|
||||
return False
|
||||
|
||||
for child_spec1, child_spec2 in zip(spec1.children(), spec2.children()):
|
||||
if not is_equivalent(child_spec1, child_spec2, equivalence_fn):
|
||||
return False
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,890 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import copy
|
||||
import inspect
|
||||
import math
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
from itertools import chain
|
||||
from typing import Any
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._export.non_strict_utils import (
|
||||
_enter_enable_graph_inputs_of_type_nn_module,
|
||||
_exit_enable_graph_inputs_of_type_nn_module,
|
||||
_get_graph_inputs_of_type_nn_module,
|
||||
)
|
||||
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
|
||||
_convert_range_to_int,
|
||||
)
|
||||
from torch._export.utils import _check_input_constraints_for_graph
|
||||
from torch.export.unflatten import _assign_attr, _AttrKind
|
||||
from torch.fx.experimental.proxy_tensor import _pytree_subclasses_that_lose_info
|
||||
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
|
||||
from torch.fx.traceback import NodeSource, NodeSourceAction
|
||||
from torch.utils._sympy.solve import try_solve
|
||||
from torch.utils._sympy.value_ranges import ValueRanges
|
||||
|
||||
from ._remove_effect_tokens_pass import _remove_effect_tokens
|
||||
from ._tree_utils import reorder_kwargs
|
||||
from .exported_program import (
|
||||
ExportedProgram,
|
||||
ExportGraphSignature,
|
||||
InputKind,
|
||||
OutputKind,
|
||||
)
|
||||
|
||||
|
||||
def eq_spec(self: pytree.TreeSpec, other: pytree.TreeSpec) -> bool:
|
||||
"""
|
||||
Refinement of TreeSpec.__eq__ where, e.g., torch.Size(...) matches tuple(...).
|
||||
See _pytree_subclasses_that_lose_info in proxy_tensor.py for more details.
|
||||
"""
|
||||
|
||||
def _normalize_type(t):
|
||||
return str(_pytree_subclasses_that_lose_info.get(t, t))
|
||||
|
||||
def _match_normalized_structure(a, b):
|
||||
if a is b:
|
||||
return True
|
||||
if _normalize_type(a.type) != _normalize_type(b.type):
|
||||
return False
|
||||
if a.type is dict and b.type is dict:
|
||||
# in the case of dict, the context is list of keys and we allow the keys to be in any order
|
||||
if set(a.context) != set(b.context):
|
||||
return False
|
||||
elif a.context != b.context:
|
||||
return False
|
||||
if a.num_children != b.num_children:
|
||||
return False
|
||||
return all(
|
||||
_match_normalized_structure(a, b)
|
||||
for a, b in zip(a.children(), b.children())
|
||||
)
|
||||
|
||||
return _match_normalized_structure(self, other)
|
||||
|
||||
|
||||
def _check_inputs_match(args, kwargs, in_spec: pytree.TreeSpec) -> list:
|
||||
reordered_kwargs = reorder_kwargs(kwargs, in_spec)
|
||||
flat_args_with_path, received_spec = pytree.tree_flatten_with_path(
|
||||
(args, reordered_kwargs)
|
||||
)
|
||||
|
||||
if not eq_spec(received_spec, in_spec):
|
||||
raise ValueError( # noqa: B904
|
||||
"Trying to flatten user inputs with exported input tree spec: \n"
|
||||
f"{in_spec}\n"
|
||||
"but actually got inputs with tree spec of: \n"
|
||||
f"{received_spec}.\n"
|
||||
"Please check that the inputs have the same number and type of "
|
||||
"args and kwargs as the ones you used when tracing."
|
||||
)
|
||||
|
||||
return flat_args_with_path
|
||||
|
||||
|
||||
def _force_ep_signature_match(ep_guards_code: list[str], input_paths):
|
||||
# TODO (tmanlaibaatar)
|
||||
# This is band-aid solution to export new tracer replacing
|
||||
# shape env sources to flat_args. The real fix should be replacing
|
||||
# shape env sources to original user sources but this is quite
|
||||
# involved because you need to carefully construct new sources using
|
||||
# dynamo and replace all instances of it inside shape env. But it is
|
||||
# lot easier to manipulate after we turn them into strings and only
|
||||
# time we use these guards is during retracing or running exported program,
|
||||
# so it is probably ok to have "not useful" guards on ep for now.
|
||||
name_mapping = {}
|
||||
for idx, path in enumerate(input_paths):
|
||||
name_mapping[f"L['flat_args'][{idx}]"] = f"L{pytree.keystr(path)}"
|
||||
|
||||
new_guards_code = []
|
||||
for guard in ep_guards_code:
|
||||
for old_name, new_name in name_mapping.items():
|
||||
guard = guard.replace(old_name, new_name)
|
||||
new_guards_code.append(guard)
|
||||
|
||||
return new_guards_code
|
||||
|
||||
|
||||
def _force_gm_signature_match(ep_guards_code: list[str], signature):
|
||||
"""
|
||||
The signature of the originally exported module may not match
|
||||
the signature of the unlifted graph module extracted from the
|
||||
exported program. The guards code extracted from the exported
|
||||
program is based on the former, but the generated guards fn is
|
||||
based on the latter; thus we need to reconcile any such diff.
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
# Handle case where signatures may differ in var args.
|
||||
orig_arg_names = set()
|
||||
for g in ep_guards_code:
|
||||
# match substrings of the form L['<name>'][<number>]
|
||||
orig_arg_names.update(re.findall(r"L\[\'([^\']+)\'\]\[([0-9]+)\]", g))
|
||||
|
||||
sig_arg_names = set()
|
||||
for n in signature.parameters:
|
||||
# match substrings of the form <name>_<number>
|
||||
sig_arg_names.update(re.findall(r"(.+)_([0-9]+)", n))
|
||||
|
||||
# replace L['<name>'][<number>] with L['<name>_<number>']
|
||||
new_guards_code = ep_guards_code
|
||||
for match in orig_arg_names:
|
||||
if match in sig_arg_names:
|
||||
base, idx = match
|
||||
new_guards_code = [
|
||||
g.replace(f"L['{base}'][{idx}]", f"L['{base}_{idx}']")
|
||||
for g in new_guards_code
|
||||
]
|
||||
|
||||
return new_guards_code
|
||||
|
||||
|
||||
def _convert_guards_code_to_fn(
|
||||
guards_code: list[str],
|
||||
paths_of_placeholders: list[pytree.KeyPath],
|
||||
):
|
||||
"""
|
||||
Generates Python code given guards code and paths of placeholders.
|
||||
We assume that, based on source information,
|
||||
- the tracer generates the guards code
|
||||
- the input spec generates the paths of placeholders.
|
||||
|
||||
Example:
|
||||
|
||||
Suppose we are given the guards code "L['z']['k'].size()[1] == 3"
|
||||
and we are given that ['z']['k'] is the path of placeholder #2.
|
||||
Then we will generate:
|
||||
```
|
||||
torch._assert(
|
||||
args[2].size()[0] == 3,
|
||||
"Guard failed: z['k'].size()[0] == 3",
|
||||
)
|
||||
```
|
||||
|
||||
FAQ: Why do we generate code based on (flattened) args instead of
|
||||
the original (unflattened) inputs? Because this would require
|
||||
inserting an additional pytree.unflatten call in our graph.
|
||||
|
||||
FAQ: Why do we not emit RuntimeError on guard failure as we used to?
|
||||
Because it is inconvenient :/, get used to AssertionError instead.
|
||||
"""
|
||||
|
||||
import ast
|
||||
|
||||
from torch.fx.experimental.symbolic_shapes import SYMPY_INTERP
|
||||
|
||||
actual_guards_code = []
|
||||
shadow_guards_code = []
|
||||
for c in guards_code:
|
||||
a, s = c, c
|
||||
for idx, path in enumerate(paths_of_placeholders):
|
||||
# e.g., replace L['z']['k'] with args[2] for Python code (actual)
|
||||
a = a.replace("L" + pytree.keystr(path), f"args[{idx}]")
|
||||
# e.g., replace L['z']['k'] with z['k'] for error message (shadow)
|
||||
s = s.replace(
|
||||
"L" + pytree.keystr(path),
|
||||
path[0].key + pytree.keystr(path[1:]), # type: ignore[attr-defined]
|
||||
)
|
||||
actual_guards_code.append(a)
|
||||
shadow_guards_code.append(s.replace("\n", ""))
|
||||
|
||||
# generate function code as str
|
||||
code_str = "\ndef _(*args):\n"
|
||||
for actual, shadow in zip(actual_guards_code, shadow_guards_code):
|
||||
# printing guards code may potentially introduce redundant parens;
|
||||
# we can normalize them out for readability by parsing/unparsing
|
||||
# NOTE: this is not necessary for correctness, just deemed desirable
|
||||
_shadow = ast.unparse(ast.parse(shadow, mode="eval"))
|
||||
# actual code and shadow error message
|
||||
code_str += f' torch._assert({actual}, "Guard failed: {_shadow}")\n'
|
||||
code_str += " return\n"
|
||||
|
||||
# populate namespace with sympy globals, materialize function (named `_`)
|
||||
namespace = {**SYMPY_INTERP}
|
||||
exec(code_str, namespace)
|
||||
|
||||
# create and return a module whose forward is the materialized function
|
||||
# NOTE: we want Dynamo to trace through this module, to repopulate guards:
|
||||
# otherwise we would lose them when retracing
|
||||
# NOTE: calling this module will be a side effect (no users): so it must
|
||||
# be marked impure to avoid being not cleaned up by DCE
|
||||
guards_fn = GuardsFn()
|
||||
guards_fn.forward = torch._dynamo.dont_skip_tracing(namespace["_"]) # type: ignore[call-overload, method-assign]
|
||||
guards_fn._is_impure = True # type: ignore[assignment]
|
||||
return guards_fn
|
||||
|
||||
|
||||
@torch._dynamo.disable
|
||||
def _check_input_constraints_for_module(self, args, kwargs):
|
||||
flat_args_with_path = _check_inputs_match(args, kwargs, self._in_spec)
|
||||
_check_input_constraints_for_graph(
|
||||
self.graph.find_nodes(op="placeholder"),
|
||||
flat_args_with_path,
|
||||
self.range_constraints,
|
||||
)
|
||||
|
||||
|
||||
def _check_input_constraints_pre_hook(self, args, kwargs):
|
||||
# preserve current behavior for clients that do not want any validation
|
||||
if not self.validate_inputs:
|
||||
return
|
||||
|
||||
# when a guards function exists, assume that the graph does calls it!
|
||||
# so we do not need to check input constraints...but we still want
|
||||
# to check inputs match, otherwise we'd get obscure pytree errors
|
||||
if hasattr(self, "_guards_fn"):
|
||||
_check_inputs_match(args, kwargs, self._in_spec)
|
||||
return
|
||||
|
||||
# NOTE: for some reason, Dynamo is tracing into this, we should see why and
|
||||
# put compile at the right place. Until then, we can skip the input
|
||||
# constraint checks.
|
||||
if not torch.compiler.is_dynamo_compiling():
|
||||
_check_input_constraints_for_module(self, args, kwargs)
|
||||
|
||||
|
||||
def _unlift_inputs_as_getattr(
|
||||
gm: torch.fx.GraphModule,
|
||||
lifted_inputs: Sequence[str | None],
|
||||
) -> tuple[dict[str, torch.fx.Node], dict[str, torch.fx.Node]]:
|
||||
"""
|
||||
Unlift inputs referring to params/buffers/constants as getattr nodes in the
|
||||
graph
|
||||
"""
|
||||
unlifted_name_to_node = {}
|
||||
input_name_to_node = {}
|
||||
|
||||
placeholder_nodes = [node for node in gm.graph.nodes if node.op == "placeholder"]
|
||||
if len(lifted_inputs) != len(placeholder_nodes):
|
||||
raise AssertionError(
|
||||
f"Number of lifted inputs ({len(lifted_inputs)}) does not match "
|
||||
f"placeholder nodes ({len(placeholder_nodes)})"
|
||||
)
|
||||
for input_node, lifted_node in zip(placeholder_nodes, lifted_inputs):
|
||||
if lifted_node is None:
|
||||
input_name_to_node[input_node.name] = input_node
|
||||
|
||||
else:
|
||||
with gm.graph.inserting_after(input_node):
|
||||
# It is fine to ignore this warning because
|
||||
# it is guaranteed that we will populate this
|
||||
# attr later.
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore")
|
||||
getattr_node = gm.graph.get_attr(lifted_node)
|
||||
input_node.replace_all_uses_with(getattr_node)
|
||||
metadata = input_node.meta
|
||||
gm.graph.erase_node(input_node)
|
||||
getattr_node.meta = metadata
|
||||
getattr_node.meta["from_node"] = [
|
||||
NodeSource(
|
||||
input_node,
|
||||
"ExportedProgram.module().unlift()",
|
||||
[NodeSourceAction.CREATE, NodeSourceAction.REPLACE],
|
||||
)
|
||||
]
|
||||
unlifted_name_to_node[lifted_node] = getattr_node
|
||||
|
||||
return unlifted_name_to_node, input_name_to_node
|
||||
|
||||
|
||||
def _insert_copy_for_mutations(
|
||||
gm: torch.fx.GraphModule,
|
||||
mutated_outputs: Sequence[str | None],
|
||||
unlifted_name_to_node: dict[str, torch.fx.Node],
|
||||
input_name_to_node: dict[str, torch.fx.Node],
|
||||
) -> None:
|
||||
"""
|
||||
Find the all the buffers and inputs that were mutated and insert copy_
|
||||
operators to reflect mutations.
|
||||
"""
|
||||
output_node = gm.graph.output_node()
|
||||
outputs = pytree.tree_flatten(output_node.args)[0]
|
||||
if len(outputs) != len(mutated_outputs):
|
||||
raise AssertionError(
|
||||
f"Number of outputs ({len(outputs)}) does not match "
|
||||
f"mutated outputs ({len(mutated_outputs)})"
|
||||
)
|
||||
|
||||
user_output_nodes = []
|
||||
return_nodes_to_copy = {}
|
||||
for return_node, mutated_node_name in zip(outputs, mutated_outputs):
|
||||
if mutated_node_name is None:
|
||||
user_output_nodes.append(return_node)
|
||||
continue
|
||||
|
||||
if mutated_node_name in unlifted_name_to_node:
|
||||
mutated_node = unlifted_name_to_node[mutated_node_name]
|
||||
elif mutated_node_name in input_name_to_node:
|
||||
mutated_node = input_name_to_node[mutated_node_name]
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Could not find {mutated_node_name} in either buffer or input nodes"
|
||||
)
|
||||
|
||||
with gm.graph.inserting_before(output_node):
|
||||
copy_node = gm.graph.call_function(
|
||||
torch.ops.aten.copy_.default, (mutated_node, return_node)
|
||||
)
|
||||
return_nodes_to_copy[return_node] = copy_node
|
||||
|
||||
output_args = tuple(
|
||||
return_nodes_to_copy.get(node, node) for node in user_output_nodes
|
||||
)
|
||||
with gm.graph.inserting_before(output_node):
|
||||
# Only return user outputs
|
||||
new_output = gm.graph.output(output_args)
|
||||
output_node.replace_all_uses_with(new_output)
|
||||
gm.graph.erase_node(output_node)
|
||||
new_output.name = output_node.name
|
||||
new_output.meta.update(output_node.meta)
|
||||
new_output.meta["from_node"] = [
|
||||
NodeSource(
|
||||
output_node,
|
||||
"ExportedProgram.module().unlift()",
|
||||
[NodeSourceAction.CREATE, NodeSourceAction.REPLACE],
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _get_codegen(
|
||||
in_spec: pytree.TreeSpec,
|
||||
out_spec: pytree.TreeSpec | None,
|
||||
forward_arg_names: list[str] | None = None,
|
||||
) -> _PyTreeCodeGen:
|
||||
"""
|
||||
Create the codegen for the graph module based on the in/out specs
|
||||
"""
|
||||
if forward_arg_names:
|
||||
names = forward_arg_names
|
||||
elif (
|
||||
in_spec.type is tuple
|
||||
and in_spec.num_children == 2
|
||||
and in_spec.child(0).type is tuple
|
||||
and in_spec.child(1).type is dict
|
||||
):
|
||||
# if in_spec contains the args (tuple) and kwargs (dict)
|
||||
names = [f"arg_{i}" for i in range(in_spec.child(0).num_children)]
|
||||
# add kwarg names
|
||||
names.extend(in_spec.child(1).context)
|
||||
else:
|
||||
names = [f"arg_{i}" for i in range(in_spec.num_children)]
|
||||
|
||||
return _PyTreeCodeGen(
|
||||
_PyTreeInfo(
|
||||
names,
|
||||
in_spec,
|
||||
out_spec,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _unlift(
|
||||
gm: torch.fx.GraphModule,
|
||||
lifted_inputs: Sequence[str | None],
|
||||
mutated_outputs: Sequence[str | None],
|
||||
in_spec: pytree.TreeSpec,
|
||||
out_spec: pytree.TreeSpec | None,
|
||||
forward_arg_names: list[str] | None = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
lifted_inputs: A list matching the graph module's input nodes. For
|
||||
an input node that is referring to a lifted parameter/buffer, this
|
||||
list will contain the fqn the corresponding attribute. Otherwise, this
|
||||
list will contain None. This is used to unlift the lifted parameters as
|
||||
get_attr nodes.
|
||||
|
||||
mutated_outputs: A list matching the graph module's output nodes. For
|
||||
an output node that is referring to a mutated buffer or user input, this
|
||||
list will contain the name of the corresponding buffer or user input
|
||||
that needs to be mutated. Otherwise, this list will contain None. This
|
||||
is used to re-insert an inplace copy_ operator to copy the mutated
|
||||
values back to the original node.
|
||||
"""
|
||||
unlifted_name_to_node, input_name_to_node = _unlift_inputs_as_getattr(
|
||||
gm, lifted_inputs
|
||||
)
|
||||
_insert_copy_for_mutations(
|
||||
gm, mutated_outputs, unlifted_name_to_node, input_name_to_node
|
||||
)
|
||||
gm.graph._codegen = _get_codegen(in_spec, out_spec, forward_arg_names)
|
||||
gm.graph.lint()
|
||||
gm.recompile()
|
||||
return gm
|
||||
|
||||
|
||||
def _register_attrs_to_new_gm(
|
||||
new_gm: torch.fx.GraphModule,
|
||||
graph_signature: ExportGraphSignature,
|
||||
state_dict: dict[str, Any],
|
||||
constants: dict[str, Any],
|
||||
) -> None:
|
||||
non_persistent_buffers = set(graph_signature.non_persistent_buffers)
|
||||
for name in graph_signature.buffers:
|
||||
if name in non_persistent_buffers:
|
||||
persistent = False
|
||||
value = constants[name]
|
||||
else:
|
||||
persistent = True
|
||||
value = state_dict[name]
|
||||
_assign_attr(
|
||||
value, new_gm, name, attr_kind=_AttrKind.BUFFER, persistent=persistent
|
||||
)
|
||||
for name in graph_signature.parameters:
|
||||
value = state_dict[name]
|
||||
_assign_attr(
|
||||
value,
|
||||
new_gm,
|
||||
name,
|
||||
attr_kind=_AttrKind.PARAMETER,
|
||||
)
|
||||
|
||||
# Technically this doesn't account for the aliased multiple constants but
|
||||
# it is ok because we have a separate pass later in the stack that populates
|
||||
# the final gm.
|
||||
for name in chain(
|
||||
graph_signature.lifted_custom_objs, graph_signature.lifted_tensor_constants
|
||||
):
|
||||
value = constants[name]
|
||||
_assign_attr(
|
||||
value,
|
||||
new_gm,
|
||||
name,
|
||||
attr_kind=_AttrKind.CONSTANT,
|
||||
)
|
||||
|
||||
|
||||
class _StatefulGraphModuleFactory(type):
|
||||
"""
|
||||
Metaclass that ensures a private constructor for _StatefulGraphModule
|
||||
"""
|
||||
|
||||
def __call__(cls, *args, **kwargs):
|
||||
raise TypeError(
|
||||
f"{cls.__module__}.{cls.__qualname__} has no public constructor. "
|
||||
)
|
||||
|
||||
def _create(cls, root, graph, range_constraints=None):
|
||||
return super().__call__(
|
||||
root,
|
||||
graph,
|
||||
range_constraints=range_constraints,
|
||||
)
|
||||
|
||||
|
||||
class _StatefulGraphModule(torch.fx.GraphModule, metaclass=_StatefulGraphModuleFactory):
|
||||
def __init__(self, root, graph, range_constraints=None):
|
||||
super().__init__(root, graph)
|
||||
# Need to fix up non-persistent buffers.
|
||||
self.range_constraints = range_constraints or []
|
||||
self.validate_inputs = True
|
||||
|
||||
|
||||
def _create_stateful_graph_module(
|
||||
plain_graph_module: torch.fx.GraphModule,
|
||||
range_constraints,
|
||||
ep: ExportedProgram,
|
||||
) -> _StatefulGraphModule:
|
||||
stateful_gm = _StatefulGraphModule._create(
|
||||
plain_graph_module,
|
||||
plain_graph_module.graph,
|
||||
range_constraints=range_constraints,
|
||||
)
|
||||
|
||||
module_types = _get_graph_inputs_of_type_nn_module(ep.example_inputs)
|
||||
stateful_gm.register_forward_pre_hook(
|
||||
lambda *args, **kwargs: _enter_enable_graph_inputs_of_type_nn_module(
|
||||
module_types
|
||||
)
|
||||
)
|
||||
stateful_gm.register_forward_pre_hook(
|
||||
_check_input_constraints_pre_hook, with_kwargs=True
|
||||
)
|
||||
|
||||
stateful_gm.register_forward_hook(
|
||||
lambda *args, **kwargs: _exit_enable_graph_inputs_of_type_nn_module(
|
||||
module_types
|
||||
),
|
||||
always_call=True,
|
||||
)
|
||||
|
||||
# When we have a constant that has requires_grad=True, we need to detach it
|
||||
# when we unlift as the tensors that require gradients should be registered
|
||||
# via parameters. But this is problematic when we have aliasing two constants
|
||||
# because when we call detach, they will become different tensors. This dict
|
||||
# keeps track of this logic.
|
||||
original_tensor_to_detached_tensor = {}
|
||||
|
||||
# Fix up lifted tensor constants.
|
||||
# fx.GraphModule() constructor silently turns a constant attribute of plain_graph_module
|
||||
# into a buffer in stateful_gm and creates an inconsistency with graph_signature.
|
||||
# We fix this by de-registering these buffers in lifted_tensor_constants
|
||||
# and call _assign_attr(attr_kind=CONSTANT) to register them as constants.
|
||||
for constant_fqn in ep.graph_signature.lifted_tensor_constants:
|
||||
# Sometimes, the constant can require gradient, this is probably a bug in user code,
|
||||
# e.g. `self.const = torch.randn(2, 2, requires_grad=True)`.
|
||||
# We call detach on the constant_val since they're tensor constants and we don't need to
|
||||
# compute their gradients anyway.
|
||||
# Users should properly register it as parameter if they want it to require gradient.
|
||||
buffer = stateful_gm.get_buffer(constant_fqn)
|
||||
if buffer.requires_grad:
|
||||
warnings.warn(
|
||||
f"A model attribute `{constant_fqn}` requires gradient. "
|
||||
f"but it's not properly registered as a parameter. "
|
||||
f"torch.export will detach it and treat it as a constant tensor "
|
||||
f"but please register it as parameter instead.",
|
||||
stacklevel=2,
|
||||
)
|
||||
detached_buffer = buffer.detach()
|
||||
original_tensor_to_detached_tensor[buffer] = detached_buffer
|
||||
buffer = detached_buffer
|
||||
*prefix, field = constant_fqn.rsplit(".")
|
||||
submod = torch.fx.graph_module._get_attr_via_attr_list(stateful_gm, prefix)
|
||||
delattr(submod, field)
|
||||
_assign_attr(buffer, stateful_gm, constant_fqn, attr_kind=_AttrKind.CONSTANT)
|
||||
|
||||
# Constants are not preserved well when we create a new GraphModule unlike param/buffers
|
||||
for const_name, value in ep.constants.items():
|
||||
if not torch.fx.graph_module._has_attr(stateful_gm, const_name):
|
||||
if isinstance(value, torch.Tensor):
|
||||
if value.requires_grad:
|
||||
warnings.warn(
|
||||
f"A model attribute `{const_name}` requires gradient "
|
||||
f"but it's not properly registered as a parameter. "
|
||||
f"torch.export will detach it and treat it as a constant tensor "
|
||||
f"but please register it as parameter instead.",
|
||||
stacklevel=2,
|
||||
)
|
||||
if value in original_tensor_to_detached_tensor:
|
||||
value = original_tensor_to_detached_tensor[value]
|
||||
else:
|
||||
detached_value = value.detach()
|
||||
original_tensor_to_detached_tensor[value] = detached_value
|
||||
value = detached_value
|
||||
_assign_attr(
|
||||
value,
|
||||
stateful_gm,
|
||||
const_name,
|
||||
attr_kind=_AttrKind.CONSTANT,
|
||||
)
|
||||
|
||||
# Fix up non-persistent buffers. torch.fx does not distinguish between
|
||||
# persistent and non-persistent buffers, so we must restore that distinction
|
||||
# here.
|
||||
for buffer in ep.graph_signature.non_persistent_buffers:
|
||||
_assign_attr(
|
||||
plain_graph_module.get_buffer(buffer),
|
||||
stateful_gm,
|
||||
buffer,
|
||||
attr_kind=_AttrKind.BUFFER,
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
return stateful_gm
|
||||
|
||||
|
||||
def _get_input_paths(example_inputs, signature):
|
||||
"""
|
||||
Generate paths of placeholders, needed for generating the guards function.
|
||||
|
||||
NOTE: Here we make use of the example inputs used for export as well as
|
||||
the signature of the unlifted graph module (not preserved by export).
|
||||
"""
|
||||
|
||||
args, kwargs = example_inputs
|
||||
binded = signature.bind(*args, **kwargs)
|
||||
binded.apply_defaults()
|
||||
ctx = binded.arguments
|
||||
flat_example_inputs_with_paths = pytree.tree_leaves_with_path(ctx)
|
||||
return [path for path, _ in flat_example_inputs_with_paths]
|
||||
|
||||
|
||||
def _replace_sources(result_str: str, flat_input_paths: list[Any]):
|
||||
"""
|
||||
Given user specified input paths, maybe fix up the guard string
|
||||
to reflect user path instead of tracer path.
|
||||
"""
|
||||
name_mapping = {}
|
||||
for idx, path in enumerate(flat_input_paths):
|
||||
name_mapping[f"L['flat_args'][{idx}]"] = f"L{pytree.keystr(path)}"
|
||||
|
||||
replace = result_str
|
||||
for key, val in name_mapping.items():
|
||||
replace = replace.replace(key, val)
|
||||
return replace
|
||||
|
||||
|
||||
def _get_input_guards_for_graph(
|
||||
placeholders: list[torch.fx.Node],
|
||||
range_constraints: dict[sympy.Symbol, ValueRanges],
|
||||
paths_for_placeholders: list[pytree.KeyPath],
|
||||
):
|
||||
"""
|
||||
Guards generated by the tracer include conditions observed in code, but
|
||||
but do not include some additional checks we typically do in export.
|
||||
For example, when dynamic shapes get specialized, are specified to be
|
||||
within a range, or are specified to be in some equational relation,
|
||||
corresponding input invalidation is done within a pre_hook, specifically,
|
||||
`_check_input_constraints_for_graph`.
|
||||
|
||||
Here we generate guards corresponding to the checks that happen in
|
||||
`_check_input_constraints_for_graph`, and add them to the guards already
|
||||
generated by the tracer. In the future, it may be worthwhile to separate
|
||||
them so that we can allow clients to turn off one but not the other.
|
||||
(Looking at you, AOTI.)
|
||||
|
||||
NOTE: We should eventually reconcile this logic with `build_guards` that
|
||||
is used by AOT Precompile.
|
||||
"""
|
||||
|
||||
deferred_expressions = []
|
||||
new_guards_code = []
|
||||
sources: dict[sympy.Expr, str] = {}
|
||||
|
||||
def handle_symint(expr, src):
|
||||
if len(expr.free_symbols) == 1:
|
||||
# complex equations (e.g., involving derived dims) need to
|
||||
# handled later, since we may not have enough information
|
||||
# just as we are passing through the placeholders in order
|
||||
deferred_expressions.append((src, expr))
|
||||
if expr in sources:
|
||||
# expressions that appear in multiple sources should force
|
||||
# inputs corresponding to those sources to be equal
|
||||
# e.g., x.shape[0] == y.shape[1]
|
||||
orig_src = sources[expr]
|
||||
new_guards_code.append(f"{src} == {orig_src}")
|
||||
else:
|
||||
sources[expr] = src
|
||||
# process value ranges as elsewhere in export
|
||||
min_val, max_val = _convert_range_to_int(range_constraints[expr])
|
||||
if min_val > 2:
|
||||
new_guards_code.append(f"{src} >= {min_val}")
|
||||
if max_val < math.inf:
|
||||
new_guards_code.append(f"{src} <= {max_val}")
|
||||
|
||||
for placeholder, path in zip(placeholders, paths_for_placeholders):
|
||||
src = "L" + pytree.keystr(path)
|
||||
meta = placeholder.meta["val"]
|
||||
# specializations
|
||||
if isinstance(meta, int):
|
||||
new_guards_code.append(f"{src} == {meta}")
|
||||
if isinstance(meta, float):
|
||||
if meta == math.inf:
|
||||
new_guards_code.append(f"{src} == math.inf")
|
||||
elif meta == -math.inf:
|
||||
new_guards_code.append(f"{src} == -math.inf")
|
||||
else:
|
||||
new_guards_code.append(f"{src} == {meta}")
|
||||
elif isinstance(meta, str):
|
||||
new_guards_code.append(f"{src} == '{meta}'")
|
||||
# range constraints and equalities
|
||||
elif isinstance(meta, torch.SymInt) and meta.node.expr in range_constraints:
|
||||
handle_symint(meta.node.expr, src)
|
||||
elif isinstance(meta, torch.Tensor):
|
||||
for i, dim in enumerate(meta.shape):
|
||||
src = "L" + pytree.keystr(path) + f".size()[{i}]"
|
||||
if isinstance(dim, int):
|
||||
# specializations
|
||||
new_guards_code.append(f"{src} == {dim}")
|
||||
elif (
|
||||
isinstance(dim, torch.SymInt) and dim.node.expr in range_constraints
|
||||
):
|
||||
# range constraints and equalities
|
||||
handle_symint(dim.node.expr, src)
|
||||
|
||||
unification_map: dict[sympy.Symbol, sympy.Expr] = {}
|
||||
py_printer = torch.utils._sympy.printers.PythonPrinter()
|
||||
|
||||
# process complex equations (e.g., involving derived dims)
|
||||
for src, expr in deferred_expressions:
|
||||
# we know this is the only symbol in expr (see check above)
|
||||
symbol = next(iter(expr.free_symbols))
|
||||
if symbol in sources:
|
||||
# if s0 is already known to be directly sourced from inputs,
|
||||
# e.g., z.shape[2], we do not need to do anything further
|
||||
# (assume we have already processed constraints on s0 above)
|
||||
continue
|
||||
|
||||
# otherwise s0 has some "hidden" source like 'dim'
|
||||
# example: src = y.shape[1], expr = s0 + 1
|
||||
if symbol in unification_map:
|
||||
# suppose that we already know that s0 = x.shape[0] * 2
|
||||
# so we can emit the guard: x.shape[0] * 2 + 1 = y.shape[1]
|
||||
substitution = expr.subs(unification_map)
|
||||
new_guards_code.append(
|
||||
py_printer.doprint(sympy.Eq(substitution, sympy.Symbol(src)))
|
||||
)
|
||||
else:
|
||||
# we do not yet know what s0 is, but given s0 + 1 = y.shape[1],
|
||||
# we can solve for s0...now knowing that s0 = y.shape[1] - 1
|
||||
solution = try_solve(sympy.Eq(expr, sympy.Symbol(src)), symbol)
|
||||
if solution is not None:
|
||||
definition = solution[1]
|
||||
unification_map[symbol] = definition
|
||||
|
||||
return new_guards_code
|
||||
|
||||
|
||||
def _ok_to_generate_guards_fn():
|
||||
patterns = [
|
||||
"executorch",
|
||||
"modai",
|
||||
"on_device_ai",
|
||||
"torchao",
|
||||
]
|
||||
# force check_guards=False for files matching `patterns`
|
||||
# because they have too many calls to .module() and
|
||||
# do not like any call modules in the graph
|
||||
# TODO: fix these files to handle guard fns
|
||||
frame = inspect.currentframe()
|
||||
while frame is not None:
|
||||
if any(path in frame.f_code.co_filename for path in patterns):
|
||||
return False
|
||||
frame = frame.f_back
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _unlift_exported_program_lifted_states(
|
||||
ep: ExportedProgram, check_guards=True
|
||||
) -> torch.fx.GraphModule:
|
||||
check_guards = check_guards and _ok_to_generate_guards_fn()
|
||||
|
||||
source_node_dict = {
|
||||
node.name: node for node in ep.graph.nodes if node.op != "placeholder"
|
||||
}
|
||||
# placeholder node name might change after deepcopy
|
||||
placeholder_source_node_dict = {
|
||||
node.target: node for node in ep.graph.nodes if node.op == "placeholder"
|
||||
}
|
||||
|
||||
new_gm = torch.fx.GraphModule(ep.graph_module, copy.deepcopy(ep.graph))
|
||||
new_gm.meta.update(ep.graph_module.meta)
|
||||
ep = copy.copy(ep)
|
||||
ep._graph_signature = ExportGraphSignature(
|
||||
ep._graph_signature.input_specs, ep._graph_signature.output_specs
|
||||
)
|
||||
ep._graph_module = new_gm
|
||||
|
||||
# TODO T206340015
|
||||
if ep.verifiers[0].dialect != "TRAINING":
|
||||
ep = _remove_effect_tokens(ep)
|
||||
|
||||
_register_attrs_to_new_gm(new_gm, ep.graph_signature, ep.state_dict, ep.constants)
|
||||
forward_arg_names = (
|
||||
sig.forward_arg_names if (sig := ep.module_call_graph[0].signature) else None
|
||||
)
|
||||
lifted_inputs: list[str | None] = [
|
||||
(
|
||||
in_spec.target
|
||||
if in_spec.kind
|
||||
in (
|
||||
InputKind.BUFFER,
|
||||
InputKind.CONSTANT_TENSOR,
|
||||
InputKind.PARAMETER,
|
||||
InputKind.CUSTOM_OBJ,
|
||||
)
|
||||
else None
|
||||
)
|
||||
for in_spec in ep.graph_signature.input_specs
|
||||
]
|
||||
|
||||
mutated_outputs: list[str | None] = [
|
||||
(
|
||||
out_spec.target
|
||||
if out_spec.kind
|
||||
in (
|
||||
OutputKind.BUFFER_MUTATION,
|
||||
OutputKind.USER_INPUT_MUTATION,
|
||||
OutputKind.PARAMETER_MUTATION,
|
||||
)
|
||||
else None
|
||||
)
|
||||
for out_spec in ep.graph_signature.output_specs
|
||||
]
|
||||
|
||||
for node in new_gm.graph.nodes:
|
||||
source_node = None
|
||||
if node.op == "placeholder":
|
||||
source_node = placeholder_source_node_dict.get(node.target)
|
||||
else:
|
||||
if node.name in source_node_dict:
|
||||
source_node = source_node_dict.get(node.name)
|
||||
node.meta["from_node"] = [
|
||||
NodeSource(
|
||||
source_node,
|
||||
"ExportedProgram.module()",
|
||||
NodeSourceAction.CREATE,
|
||||
)
|
||||
]
|
||||
|
||||
if ep.call_spec.in_spec is None:
|
||||
raise AssertionError("ep.call_spec.in_spec cannot be None")
|
||||
new_gm = _unlift(
|
||||
new_gm,
|
||||
lifted_inputs,
|
||||
mutated_outputs,
|
||||
ep.call_spec.in_spec,
|
||||
ep.call_spec.out_spec,
|
||||
forward_arg_names=forward_arg_names,
|
||||
)
|
||||
unlift_gm = _create_stateful_graph_module(new_gm, ep.range_constraints, ep)
|
||||
unlift_gm.meta.update(ep.graph_module.meta)
|
||||
|
||||
# create a _guards_fn submodule and insert a call to it after placeholders
|
||||
graph = unlift_gm.graph
|
||||
placeholders = graph.find_nodes(op="placeholder")
|
||||
if check_guards and placeholders and ep.example_inputs:
|
||||
sig = inspect.signature(unlift_gm.forward)
|
||||
input_paths = _get_input_paths(
|
||||
ep.example_inputs,
|
||||
sig,
|
||||
)
|
||||
|
||||
# TODO (tmanlaibaatar)
|
||||
# This is band-aid solution to export new tracer replacing
|
||||
# shape env sources to flat_args. The real fix should be replacing
|
||||
# shape env sources to original user sources but this is quite
|
||||
# involved because you need to carefully construct new sources using
|
||||
# dynamo and replace all instances of it inside shape env. But it is
|
||||
# lot easier to manipulate after we turn them into strings and only
|
||||
# time we use these guards is during retracing or running exported program,
|
||||
# so it is probably ok to have "not useful" guards on ep for now.
|
||||
ep_guards = []
|
||||
for guard in ep._guards_code:
|
||||
ep_guards.append(_replace_sources(guard, input_paths))
|
||||
|
||||
guards_code = _get_input_guards_for_graph(
|
||||
placeholders, ep.range_constraints, input_paths
|
||||
)
|
||||
|
||||
ep_guards_code = _force_ep_signature_match(ep._guards_code, input_paths)
|
||||
ep_guards_code = _force_gm_signature_match(ep_guards_code, sig)
|
||||
guards_code.extend(ep_guards_code)
|
||||
unlift_gm._guards_fn = _convert_guards_code_to_fn(guards_code, input_paths)
|
||||
|
||||
root_nn_module_stack = torch.fx._utils.first_call_function_nn_module_stack(
|
||||
graph
|
||||
)
|
||||
with graph.inserting_after(placeholders[-1]):
|
||||
node = graph.call_module("_guards_fn", tuple(placeholders))
|
||||
node.meta["nn_module_stack"] = root_nn_module_stack
|
||||
|
||||
unlift_gm.recompile()
|
||||
|
||||
return unlift_gm
|
||||
|
||||
|
||||
class GuardsFn(torch.nn.Module):
|
||||
"""
|
||||
Module class for guard functions.
|
||||
"""
|
||||
|
||||
def forward(self, *args):
|
||||
pass
|
||||
@@ -0,0 +1,10 @@
|
||||
import torch
|
||||
|
||||
|
||||
class _WrapperModule(torch.nn.Module):
|
||||
def __init__(self, f): # type: ignore[no-untyped-def]
|
||||
super().__init__()
|
||||
self.f = f
|
||||
|
||||
def forward(self, *args, **kwargs): # type: ignore[no-untyped-def]
|
||||
return self.f(*args, **kwargs)
|
||||
@@ -0,0 +1,16 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
__all__ = ["ScriptObjectMeta"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScriptObjectMeta:
|
||||
"""
|
||||
Metadata which is stored on nodes representing ScriptObjects.
|
||||
"""
|
||||
|
||||
# Key into constants table to retrieve the real ScriptObject.
|
||||
constant_name: str
|
||||
|
||||
class_fqn: str
|
||||
@@ -0,0 +1,56 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import importlib
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
lib = torch.library.Library("export", "FRAGMENT") # noqa: TOR901
|
||||
|
||||
lib.define(
|
||||
"access_subclass_inner_tensor(Tensor src_subclass_tensor, str attr) -> Tensor"
|
||||
)
|
||||
|
||||
|
||||
@torch.library.impl(lib, "access_subclass_inner_tensor", "Autograd")
|
||||
# When running under torch.inference_mode(), we seem to skip AUtograd key
|
||||
# so we should desugar this op as soon as we start tracing to post-dispatch.
|
||||
@torch.library.impl(lib, "access_subclass_inner_tensor", "Python")
|
||||
def _access_subclass_inner_tensor(
|
||||
src_subclass_tensor: torch.Tensor, attr: str
|
||||
) -> torch.Tensor:
|
||||
from torch.utils._python_dispatch import is_traceable_wrapper_subclass
|
||||
|
||||
if not is_traceable_wrapper_subclass(src_subclass_tensor):
|
||||
raise AssertionError(
|
||||
f"Expected src_subclass_tensor to be a traceable wrapper subclass, "
|
||||
f"but got {type(src_subclass_tensor)}"
|
||||
)
|
||||
val = getattr(src_subclass_tensor, attr, None)
|
||||
if val is None or not isinstance(val, torch.Tensor):
|
||||
raise RuntimeError(
|
||||
f"Attribute {attr} is not a tensor or doesn't exist in {src_subclass_tensor}"
|
||||
)
|
||||
return val
|
||||
|
||||
|
||||
def _call_custom_autograd_function_in_pre_dispatch(function_cls_name, *args, **kwargs):
|
||||
"""
|
||||
Import a custom autograd function by string name and call it. This is pretty bad
|
||||
because:
|
||||
1) There is no schema
|
||||
|
||||
Ideally we should automatically wrap custom autograd functions with a custom op, but
|
||||
that is too much work because we need to schematize custom autograd functions. For now,
|
||||
we just hackily put it in the IR.
|
||||
"""
|
||||
# Parse module and class name
|
||||
module_name, class_name = function_cls_name.rsplit(".", 1)
|
||||
|
||||
# Import the module and get the class
|
||||
module = importlib.import_module(module_name)
|
||||
function_cls = getattr(module, class_name)
|
||||
if not hasattr(function_cls, "apply"):
|
||||
raise AssertionError(
|
||||
f"Expected function class {function_cls_name} to have 'apply' method"
|
||||
)
|
||||
return function_cls.apply(*args, **kwargs)
|
||||
@@ -0,0 +1,160 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from torch._export.utils import (
|
||||
_collect_all_valid_cia_ops,
|
||||
_collect_all_valid_cia_ops_for_aten_namespace,
|
||||
_get_decomp_for_cia,
|
||||
_is_aten_op,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["CustomDecompTable"]
|
||||
|
||||
|
||||
"""
|
||||
Core ATen ops with Composite Implicit Autograd dispatch that should be excluded from decomposition
|
||||
by default. The decomposition logic should eventually exclude all core-tagged CIA ops, but until all
|
||||
backends are ready, this list allows opt-in one at a time.
|
||||
"""
|
||||
PRESERVED_ATEN_CIA_OPS = {
|
||||
torch.ops.aten.upsample_bilinear2d.vec,
|
||||
torch.ops.aten.upsample_nearest2d.vec,
|
||||
# NB: don't use the C++ decomp, because it is not functional!
|
||||
torch.ops.aten.silu_backward.default,
|
||||
torch.ops.aten.mish_backward.default,
|
||||
torch.ops.aten._fused_rms_norm.default,
|
||||
}
|
||||
|
||||
|
||||
class CustomDecompTable(dict[torch._ops.OperatorBase, Callable]):
|
||||
"""
|
||||
This is a custom dictionary that is specifically used for handling decomp_table in export.
|
||||
The reason we need this is because in the new world, you can only *delete* an op from decomp
|
||||
table to preserve it. This is problematic for custom ops because we don't know when the custom
|
||||
op will actually be loaded to the dispatcher. As a result, we need to record the custom ops operations
|
||||
until we really need to materialize it (which is when we run decomposition pass.)
|
||||
|
||||
Invariants we hold are:
|
||||
1. All aten decomp is loaded at the init time
|
||||
2. We materialize ALL ops when user ever reads from the table to make it more likely
|
||||
that dispatcher picks up the custom op.
|
||||
3. If it is write operation, we don't necessarily materialize
|
||||
4. We load the final time during export, right before calling run_decompositions()
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
from torch._decomp import _core_aten_decompositions_post_autograd
|
||||
|
||||
# For aten ops, we load them up in the beginning
|
||||
self.decomp_table = _core_aten_decompositions_post_autograd()
|
||||
|
||||
for op in _collect_all_valid_cia_ops_for_aten_namespace():
|
||||
if op not in PRESERVED_ATEN_CIA_OPS and op not in self.decomp_table:
|
||||
self.decomp_table[op] = _get_decomp_for_cia(op)
|
||||
|
||||
# This is to track the *pending* deleted custom ops that haven't been materialized yet
|
||||
self.deleted_custom_ops = set()
|
||||
# When this is true, there shouldn't be any pending operations in the table.
|
||||
self.has_materialized = False
|
||||
|
||||
def __getitem__(self, key):
|
||||
self._materialize_if_needed()
|
||||
return self.decomp_table.__getitem__(key)
|
||||
|
||||
def __setitem__(self, key, value) -> None:
|
||||
self.decomp_table.__setitem__(key, value)
|
||||
|
||||
if key in self.deleted_custom_ops:
|
||||
self.deleted_custom_ops.remove(key)
|
||||
|
||||
def keys(self):
|
||||
self._materialize_if_needed()
|
||||
return self.decomp_table.keys()
|
||||
|
||||
def __delitem__(self, key) -> None:
|
||||
self.pop(key)
|
||||
|
||||
def update(self, other_dict): # type: ignore[override]
|
||||
for k, v in other_dict.items():
|
||||
self.decomp_table.__setitem__(k, v)
|
||||
|
||||
def __missing__(self, key) -> bool:
|
||||
return not self.__contains__(key)
|
||||
|
||||
def __contains__(self, key) -> bool:
|
||||
self._materialize_if_needed()
|
||||
return self.decomp_table.__contains__(key)
|
||||
|
||||
def __len__(self) -> int:
|
||||
self._materialize_if_needed()
|
||||
return self.decomp_table.__len__()
|
||||
|
||||
def __iter__(self):
|
||||
self._materialize_if_needed()
|
||||
return self.decomp_table.__iter__()
|
||||
|
||||
def __reversed__(self):
|
||||
self._materialize_if_needed()
|
||||
return self.decomp_table.__reversed__()
|
||||
|
||||
def copy(self) -> "CustomDecompTable":
|
||||
new_dict = CustomDecompTable()
|
||||
new_dict.decomp_table = self.decomp_table.copy()
|
||||
new_dict.deleted_custom_ops = self.deleted_custom_ops.copy()
|
||||
new_dict.has_materialized = self.has_materialized
|
||||
return new_dict
|
||||
|
||||
def pop(self, *args):
|
||||
def _pop_if_can(key):
|
||||
if _is_aten_op(key):
|
||||
return self.decomp_table.pop(key)
|
||||
|
||||
if key in self.decomp_table:
|
||||
# Even if we materialized it, we should add it to the deleted
|
||||
# custom ops list so that when we materialize next time,
|
||||
# we should respect user's intention.
|
||||
self.deleted_custom_ops.add(key)
|
||||
return self.decomp_table.pop(key)
|
||||
|
||||
if key in self.deleted_custom_ops:
|
||||
raise KeyError(f"{key} doesn't exist in the table")
|
||||
|
||||
self.deleted_custom_ops.add(key)
|
||||
# We would come here when user pops off something that is
|
||||
# not in the table. In this case, we just pretend that it
|
||||
# was in the table.
|
||||
return _get_decomp_for_cia(key)
|
||||
|
||||
if len(args) == 1:
|
||||
return _pop_if_can(args[0])
|
||||
|
||||
if len(args) == 2:
|
||||
try:
|
||||
return _pop_if_can(args[0])
|
||||
except KeyError:
|
||||
return args[1]
|
||||
|
||||
def items(self):
|
||||
self._materialize_if_needed()
|
||||
return self.decomp_table.items()
|
||||
|
||||
def materialize(self) -> dict[torch._ops.OperatorBase, Callable]:
|
||||
for op in _collect_all_valid_cia_ops():
|
||||
if _is_aten_op(op):
|
||||
continue
|
||||
elif op in self.decomp_table:
|
||||
continue
|
||||
elif op not in self.deleted_custom_ops:
|
||||
self.decomp_table[op] = _get_decomp_for_cia(op)
|
||||
|
||||
self.has_materialized = True
|
||||
self.deleted_custom_ops = set()
|
||||
return {**self.decomp_table}
|
||||
|
||||
def _materialize_if_needed(self) -> None:
|
||||
if not self.has_materialized:
|
||||
self.materialize()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,454 @@
|
||||
import copy
|
||||
import dataclasses
|
||||
import functools
|
||||
import os
|
||||
import types
|
||||
import typing
|
||||
import typing_extensions
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from torch.export.experimental._utils import _get_main_cpp_file, _get_make_file
|
||||
from torch.export.exported_program import _decompose_exported_program
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
|
||||
_InputT = typing_extensions.ParamSpec("_InputT")
|
||||
_RetT = typing.TypeVar("_RetT")
|
||||
|
||||
|
||||
__all__ = [] # type: ignore[var-annotated]
|
||||
|
||||
|
||||
def _copy_graph_module_and_signature(
|
||||
ep: torch.export.ExportedProgram,
|
||||
) -> tuple[torch.fx.GraphModule, torch.export.graph_signature.ExportGraphSignature]:
|
||||
# copy.deepcopy lets the objects override __deepcopy__ methods with graph_copy() and node_copy(),
|
||||
# and this can break placeholder names in some particular cases.
|
||||
# For example, node copying will avoid Python keywords like 'input', suffixing and renaming to 'input_1'.
|
||||
# So we manually overwrite placeholder names by reading the old graph.
|
||||
gm = copy.deepcopy(ep.graph_module)
|
||||
new_graph_signature = copy.deepcopy(ep.graph_signature)
|
||||
|
||||
# iterate over old/new graph modules
|
||||
for old_gm, new_gm in zip(ep.graph_module.modules(), gm.modules()): # type: ignore[union-attr]
|
||||
old_phs = [node for node in old_gm.graph.nodes if node.op == "placeholder"]
|
||||
new_phs = [node for node in new_gm.graph.nodes if node.op == "placeholder"]
|
||||
# iterate over placeholders
|
||||
if len(old_phs) != len(new_phs):
|
||||
raise AssertionError(
|
||||
f"Number of old placeholders ({len(old_phs)}) does not match "
|
||||
f"new placeholders ({len(new_phs)})"
|
||||
)
|
||||
for old_node, new_node in zip(old_phs, new_phs):
|
||||
new_node.name = old_node.name
|
||||
|
||||
return gm, new_graph_signature
|
||||
|
||||
|
||||
def _remove_detach_pass(
|
||||
gm: torch.fx.GraphModule, sig: torch.export.graph_signature.ExportGraphSignature
|
||||
) -> None:
|
||||
with gm._set_replace_hook(sig.get_replace_hook()):
|
||||
for node in list(reversed(gm.graph.nodes)):
|
||||
if node.op != "call_function":
|
||||
continue
|
||||
if (
|
||||
node.target is torch.ops.aten.detach.default
|
||||
and len(node.users) == 1
|
||||
and next(iter(node.users)).target is torch.ops.aten.detach.default
|
||||
):
|
||||
next(iter(node.users)).replace_all_uses_with(node)
|
||||
|
||||
gm.graph.eliminate_dead_code()
|
||||
gm.recompile()
|
||||
|
||||
|
||||
def _export_forward_backward(
|
||||
ep: torch.export.ExportedProgram, joint_loss_index: int = 0
|
||||
) -> torch.export.ExportedProgram:
|
||||
"""
|
||||
WARNING: This API is highly unstable and will be subject to change in the future.
|
||||
"""
|
||||
from torch._decomp import core_aten_decompositions
|
||||
|
||||
ep = _decompose_exported_program(
|
||||
ep,
|
||||
cia_to_decomp={},
|
||||
python_decomp_table=core_aten_decompositions(),
|
||||
joint_loss_index=joint_loss_index,
|
||||
# For serialization purpose, we don't want to decompose custom triton ops.
|
||||
# If users would like to decompose custom triton ops, they could do it
|
||||
# with run_decompositions() API.
|
||||
decompose_custom_triton_ops=False,
|
||||
)
|
||||
gm, new_graph_signature = _copy_graph_module_and_signature(ep)
|
||||
_remove_detach_pass(gm, new_graph_signature)
|
||||
|
||||
return ep._update(gm, new_graph_signature)
|
||||
|
||||
|
||||
def _sticky_export(
|
||||
forward_func: typing.Callable[_InputT, _RetT],
|
||||
dynamic_shapes_callback: typing.Callable[
|
||||
_InputT, list[typing.Any] | dict[str, typing.Any] | tuple[typing.Any, ...]
|
||||
]
|
||||
| None = None,
|
||||
) -> typing.Callable[_InputT, _RetT]:
|
||||
"""
|
||||
Lazily export the model on first forward call.
|
||||
Usage:
|
||||
model.forward = _sticky_export(model.forward, dynamic_shapes_callback=callback)
|
||||
"""
|
||||
model = forward_func.__self__ # type: ignore[attr-defined]
|
||||
original_forward = forward_func.__func__ # type: ignore[attr-defined]
|
||||
|
||||
@functools.wraps(forward_func)
|
||||
def wrapper(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
|
||||
# Unpatch forward to avoid recursion during export
|
||||
model.forward = types.MethodType(original_forward, model)
|
||||
|
||||
dynamic_shapes_spec = None
|
||||
if dynamic_shapes_callback:
|
||||
dynamic_shapes_spec = dynamic_shapes_callback(*args, **kwargs)
|
||||
|
||||
try:
|
||||
exported = torch.export.export(
|
||||
model,
|
||||
args,
|
||||
kwargs,
|
||||
dynamic_shapes=dynamic_shapes_spec,
|
||||
).module()
|
||||
wrapper._exported_artifact = exported # type: ignore[attr-defined]
|
||||
finally:
|
||||
# Restore the wrapper after export
|
||||
model.forward = wrapper
|
||||
|
||||
return exported(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _ExportMethod:
|
||||
overloads: dict[str, torch.export.ExportedProgram]
|
||||
fallbacks: list[torch.export.ExportedProgram]
|
||||
|
||||
|
||||
class _ExportPackage:
|
||||
"""
|
||||
An export package is a collection of torch.export()-ed PyTorch models consisting of
|
||||
a list of exported methods and their corresponding overloads. ExportPackage is introduced
|
||||
on top of torch.export() to support the following use cases:
|
||||
- Exporting a model with multiple methods if a model has multiple independent parts.
|
||||
- Exporting a function with multiple overloads based on tensor shapes or other metadata.
|
||||
|
||||
ExportPackage is designed to contain multiple methods (associated with method names) and for
|
||||
each method, it can have multiple overloads (associated with overload names).
|
||||
|
||||
Here is an example of the data structure for an ExportPackage:
|
||||
```
|
||||
ExportPackage(
|
||||
methods={
|
||||
"decoder": ExportMethod(
|
||||
overloads={
|
||||
"prefill": ExportedProgram(...),
|
||||
"decode": ExportedProgram(...),
|
||||
},
|
||||
fallbacks=[],
|
||||
),
|
||||
"encoder": ExportMethod(overloads={}, fallbacks=[ExportedProgram(...)]),
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
To export a model into an ExportPackage, users can use the exporter API provided by ExportPackage.
|
||||
Exporter is a decorator that takes a callable and returns a wrapper. The wrapper will export the
|
||||
function into an ExportPackage, when it's invoked with some sample inputs (similar to how
|
||||
torch.compile() works). For more details, please refer to the document on .exporter() method.
|
||||
|
||||
This design allows users to decouple the exported callables from the actual sample inputs which can
|
||||
be helpful for use cases where the exported callable is hidden behind helper functions or when sample
|
||||
inpusts are hard to get.
|
||||
|
||||
NOTE: This is an experimental API and anything can be changed in the future.
|
||||
|
||||
Example usage:
|
||||
```
|
||||
def fn(x):
|
||||
return x + 1
|
||||
|
||||
def main(f, x):
|
||||
x += 1
|
||||
ret = f(x)
|
||||
return ret + 1
|
||||
|
||||
package = ExportPackage()
|
||||
main(package.exporter(fn), torch.randn(3, 2))
|
||||
```
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.methods: dict[str, _ExportMethod] = {}
|
||||
|
||||
def _exporter(
|
||||
self,
|
||||
method: str,
|
||||
fn: typing.Callable[_InputT, _RetT],
|
||||
*,
|
||||
fallback: str = "once",
|
||||
) -> typing.Callable[_InputT, _RetT]:
|
||||
"""
|
||||
A function/module decorator that sets up a callable to be exported later invoked.
|
||||
By default the exporter will only trigger torch.export for once and error on
|
||||
later invocations. To customize this behavior, users have the following two options:
|
||||
1. Call .define_overload() method on the returned wrapper to define an overload.
|
||||
2. Adjust the fallback policy using `fallback` argument.
|
||||
|
||||
An "overload" is a named branch for an ExportMethod with a user defined precondition,
|
||||
typically based on input tensor shapes. It's up to a downstream backend implementation
|
||||
of ExportMethod to respect the precondition later in inference.
|
||||
|
||||
define_overload() takes arguments like the following:
|
||||
- A name, for indexing purposes in a backend.
|
||||
- A callable (spec) that:
|
||||
- Has the same model input signature as the original model code.
|
||||
- Returns an optional dynamic shape spec.
|
||||
|
||||
Exporter will only export an overload when the spec callable successfully returns
|
||||
a result without raising AssertionError.
|
||||
|
||||
For example:
|
||||
```
|
||||
package = ExportPackage()
|
||||
|
||||
|
||||
def prefill(x, xa, kv_cache):
|
||||
assert x.shape[1] == 3
|
||||
assert kv_cache == {}
|
||||
|
||||
|
||||
def decode(x, xa, kv_cache):
|
||||
assert x.shape[1] > 1
|
||||
assert len(kv_cache) > 0
|
||||
return {...} # dynamic shape specs here
|
||||
|
||||
|
||||
exporter = (
|
||||
package.exporter(decoder)
|
||||
.define_overload("prefill", prefill)
|
||||
.define_overload("decode", decode)
|
||||
)
|
||||
```
|
||||
|
||||
A "fallback" is exported when no overload precondition matches a given set of sample
|
||||
inputs. Overloads should
|
||||
Fallbacks don't have names and are ordered in a list. It's up to a backend to decide
|
||||
which fallback is used amony multiple ones.
|
||||
|
||||
A reference backend implementation of ExportMethod may look like the following:
|
||||
```
|
||||
def execute(method: ExportMethod, *args, **kwargs):
|
||||
for overload in method.overloads:
|
||||
if match_precondition(overload, *args, **kwargs):
|
||||
return execute_overload(overload, *args, **kwargs)
|
||||
for fallback in method.fallbacks:
|
||||
if match_precondition(fallback, *args, **kwargs):
|
||||
return execute_fallback(fallback, *args, **kwargs)
|
||||
```
|
||||
|
||||
Args:
|
||||
method(str): The method name for an exported part of PyTorch model. This
|
||||
will be saved together with the exported/compiled artifacts
|
||||
in any serialization format and can be used as the key to
|
||||
index ExportPackage methods later.
|
||||
fn(callable): A PyTorch function/module to be exported.
|
||||
fallback(str): The fallback policy to decide when to call torch.export
|
||||
- "once" is the default policy. Under this policy a PyTorch program is assumed
|
||||
to be only called once later and an error will be raised for subsequent
|
||||
runs.
|
||||
- "error" means the ExportMethod will never have any fallbacks, meaning
|
||||
users should define all the possible overloads ahead of time.
|
||||
|
||||
"""
|
||||
|
||||
fallbacks: list[torch.export.ExportedProgram] = []
|
||||
specs: dict[str, typing.Callable[_InputT, typing.Any]] = {}
|
||||
overloads: dict[str, torch.export.ExportedProgram] = {}
|
||||
self.methods[method] = _ExportMethod(fallbacks=fallbacks, overloads=overloads)
|
||||
|
||||
@functools.wraps(fn)
|
||||
def _exporter_context(*args, **kwargs): # type: ignore[no-untyped-def]
|
||||
import torch.export._wrapper_utils
|
||||
|
||||
model: torch.nn.Module
|
||||
if not isinstance(fn, torch.nn.Module):
|
||||
model = torch.export._wrapper_utils._WrapperModule(fn)
|
||||
else:
|
||||
model = fn
|
||||
|
||||
for k, v in specs.items():
|
||||
try:
|
||||
if isinstance(fn, torch.nn.Module):
|
||||
dynamic_shapes = v(fn, *args, **kwargs) # type: ignore[arg-type]
|
||||
else:
|
||||
# pyrefly: ignore [invalid-param-spec]
|
||||
dynamic_shapes = v(*args, **kwargs)
|
||||
except AssertionError:
|
||||
continue
|
||||
if k not in overloads:
|
||||
ep = torch.export.export(
|
||||
model, args, kwargs, dynamic_shapes=dynamic_shapes
|
||||
)
|
||||
overloads[k] = ep
|
||||
ep = overloads[k]
|
||||
return ep.module()(*args, **kwargs)
|
||||
|
||||
if fallback == "error":
|
||||
raise RuntimeError(
|
||||
f"Exporter: Cannot export fallback {fn} when fallback policy is set to 'error',"
|
||||
+ "please specify an overload or adjust the fallback policy."
|
||||
)
|
||||
elif fallback == "once":
|
||||
if len(fallbacks) > 0:
|
||||
raise RuntimeError(
|
||||
f"Exporter: Cannot export {fn} more than once, "
|
||||
+ "please specify an overload or adjust the fallback policy."
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(f"Unknown fallback policy: {fallback}")
|
||||
ep = torch.export.export(model, args, kwargs)
|
||||
|
||||
fallbacks.append(ep)
|
||||
return ep.module()(*args, **kwargs)
|
||||
|
||||
if isinstance(fn, torch.nn.Module):
|
||||
_exporter_context = torch._dynamo.eval_frame.OptimizedModule( # type: ignore[assignment] # noqa: F811
|
||||
fn,
|
||||
lambda _: _exporter_context, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
def _define_overload(
|
||||
overload: str, spec: typing.Callable[_InputT, typing.Any]
|
||||
) -> typing.Any:
|
||||
if overload in specs:
|
||||
raise AssertionError(f"Overload '{overload}' already exists in specs")
|
||||
if not callable(spec):
|
||||
raise AssertionError(f"spec must be callable, but got {type(spec)}")
|
||||
if not overload.isidentifier():
|
||||
raise AssertionError(
|
||||
f"Overload '{overload}' is not a valid Python identifier"
|
||||
)
|
||||
specs[overload] = spec
|
||||
return _exporter_context
|
||||
|
||||
if hasattr(fn, "_define_overload"):
|
||||
raise AssertionError("fn already has a '_define_overload' attribute")
|
||||
_exporter_context._define_overload = _define_overload # type: ignore[attr-defined]
|
||||
|
||||
# pyrefly: ignore [bad-return]
|
||||
return _exporter_context
|
||||
|
||||
@property
|
||||
def _method_overloads(
|
||||
self,
|
||||
) -> typing.Iterator[tuple[str, torch.export.ExportedProgram]]:
|
||||
for method, method_data in self.methods.items():
|
||||
for overload, ep in method_data.overloads.items():
|
||||
yield f"{method}:{overload}", ep
|
||||
|
||||
def _compiled_and_package(
|
||||
self,
|
||||
f: torch.types.FileLike,
|
||||
standalone: bool = False,
|
||||
package_example_inputs: bool = False,
|
||||
) -> None:
|
||||
options: dict[str, typing.Any] = {
|
||||
"aot_inductor.package": True,
|
||||
"aot_inductor.package_cpp_only": True,
|
||||
"always_keep_tensor_constants": True,
|
||||
# we'll change this back to False once we enable weight deduping for standalone mode
|
||||
"aot_inductor.package_constants_in_so": standalone,
|
||||
"aot_inductor_mode.compile_standalone": standalone,
|
||||
}
|
||||
aoti_files_map = {}
|
||||
model_names = []
|
||||
device_type = "cpu"
|
||||
for name, ep in self._method_overloads:
|
||||
name = name.replace(":", "__")
|
||||
model_names.append(name)
|
||||
options["aot_inductor.model_name_for_generated_files"] = name
|
||||
aoti_files = torch._inductor.aot_compile(
|
||||
ep.module(), # type: ignore[arg-type]
|
||||
ep.example_inputs[0],
|
||||
kwargs=ep.example_inputs[1],
|
||||
options=options,
|
||||
)
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
aoti_files_map[name] = aoti_files
|
||||
|
||||
from torch._inductor.package import package
|
||||
|
||||
pt2_path = package.package_aoti(
|
||||
f,
|
||||
aoti_files_map, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
if not standalone:
|
||||
return
|
||||
|
||||
if not isinstance(pt2_path, str):
|
||||
raise AssertionError(
|
||||
f"Expected pt2_path to be a string, but got {type(pt2_path)}"
|
||||
)
|
||||
base_directory = os.path.dirname(pt2_path)
|
||||
package_name = os.path.basename(pt2_path)[:-4]
|
||||
with zipfile.ZipFile(pt2_path, "r") as zip_ref:
|
||||
zip_ref.extractall(base_directory)
|
||||
|
||||
example_inputs_map: dict[str, int] | None = (
|
||||
{} if package_example_inputs else None
|
||||
)
|
||||
for name, ep in self._method_overloads:
|
||||
name = name.replace(":", "__")
|
||||
# TODO: also dump kwargs
|
||||
# TODO: currently only support list of Tensors and they need to be on the same device
|
||||
if not ep.example_inputs:
|
||||
continue
|
||||
|
||||
device_types: OrderedSet[str] = OrderedSet()
|
||||
|
||||
for inp in ep.example_inputs[0]:
|
||||
if isinstance(inp, torch.Tensor):
|
||||
device_types.add(inp.device.type)
|
||||
device_types.discard("cpu")
|
||||
if len(device_types) > 1:
|
||||
raise AssertionError(
|
||||
"Does not support mixing {}".format("+".join(list(device_types)))
|
||||
)
|
||||
device_type = "cpu" if len(device_types) == 0 else device_types.pop()
|
||||
|
||||
if package_example_inputs:
|
||||
if example_inputs_map is None:
|
||||
raise AssertionError(
|
||||
"example_inputs_map cannot be None when package_example_inputs is True"
|
||||
)
|
||||
example_inputs_map[name] = len(ep.example_inputs[0])
|
||||
for i, t in enumerate(ep.example_inputs[0]):
|
||||
path = Path(base_directory) / f"{name}_input_{i}.pt"
|
||||
torch.save(t, path)
|
||||
|
||||
cmake_file_str = _get_make_file(
|
||||
package_name, model_names, device_type=device_type
|
||||
)
|
||||
|
||||
with open(Path(base_directory) / "CMakeLists.txt", "w") as file:
|
||||
file.write(cmake_file_str)
|
||||
|
||||
main_file_str = _get_main_cpp_file(
|
||||
package_name, model_names, example_inputs_map, device_type=device_type
|
||||
)
|
||||
with open(Path(base_directory) / "main.cpp", "w") as file:
|
||||
file.write(main_file_str)
|
||||
@@ -0,0 +1,238 @@
|
||||
import logging
|
||||
|
||||
import torch
|
||||
from torch._inductor.utils import IndentedBuffer
|
||||
|
||||
|
||||
__all__ = [] # type: ignore[var-annotated]
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _get_main_cpp_file(
|
||||
package_name: str,
|
||||
model_names: list[str],
|
||||
example_inputs_map: dict[str, int] | None,
|
||||
device_type: str,
|
||||
) -> str:
|
||||
"""
|
||||
Generates a main.cpp file for AOTInductor standalone models in the specified package.
|
||||
|
||||
Args:
|
||||
package_name (str): Name of the package containing the models.
|
||||
model_names (List[str]): List of model names to include in the generated main.cpp.
|
||||
cuda (bool): Whether to generate code with CUDA support.
|
||||
example_inputs_map (Optional[Dict[str, List[Tensor]]]): A mapping from model name to
|
||||
its list of example input tensors. If provided, the generated main.cpp will
|
||||
load and run these inputs.
|
||||
|
||||
Returns:
|
||||
str: The contents of the generated main.cpp file as a string.
|
||||
"""
|
||||
|
||||
ib = IndentedBuffer()
|
||||
|
||||
ib.writelines(
|
||||
[
|
||||
"#include <dlfcn.h>",
|
||||
"#include <fstream>",
|
||||
"#include <iostream>",
|
||||
"#include <memory>",
|
||||
"#include <torch/torch.h>",
|
||||
"#include <vector>",
|
||||
"#include <torch/csrc/inductor/aoti_torch/tensor_converter.h>",
|
||||
]
|
||||
)
|
||||
if device_type == "cuda":
|
||||
if torch.version.hip:
|
||||
ib.writelines(
|
||||
[
|
||||
"#include <hip/hip_runtime.h>",
|
||||
]
|
||||
)
|
||||
|
||||
else:
|
||||
ib.writelines(
|
||||
[
|
||||
"#include <cuda.h>",
|
||||
"#include <cuda_runtime_api.h>",
|
||||
]
|
||||
)
|
||||
for model_name in model_names:
|
||||
ib.writeline(
|
||||
f'#include "{package_name}/data/aotinductor/{model_name}/{model_name}.h"'
|
||||
)
|
||||
|
||||
ib.newline()
|
||||
for model_name in model_names:
|
||||
ib.writeline(f"using torch::aot_inductor::AOTInductorModel{model_name};")
|
||||
|
||||
ib.writelines(
|
||||
[
|
||||
"using torch::aot_inductor::ConstantHandle;",
|
||||
"using torch::aot_inductor::ConstantMap;",
|
||||
"",
|
||||
"int main(int argc, char* argv[]) {",
|
||||
]
|
||||
)
|
||||
|
||||
with ib.indent():
|
||||
ib.writeline(f'std::string device_str = "{device_type}";')
|
||||
ib.writeline("try {")
|
||||
|
||||
with ib.indent():
|
||||
ib.writeline("c10::Device device(device_str);")
|
||||
|
||||
if example_inputs_map is not None:
|
||||
# TODO: add device
|
||||
for i, model_name in enumerate(model_names):
|
||||
num_inputs = example_inputs_map[model_name]
|
||||
|
||||
ib.writeline(f"// Load input tensors for model {model_name}")
|
||||
ib.writeline(f"std::vector<at::Tensor> input_tensors{i + 1};")
|
||||
ib.writeline(f"for (int j = 0; j < {num_inputs}; ++j) {{")
|
||||
with ib.indent():
|
||||
ib.writeline(
|
||||
f'std::string filename = "{model_name}_input_" + std::to_string(j) + ".pt";'
|
||||
)
|
||||
ib.writeline("std::ifstream in(filename, std::ios::binary);")
|
||||
ib.writeline("if (!in.is_open()) {")
|
||||
with ib.indent():
|
||||
ib.writeline(
|
||||
'std::cerr << "Failed to open file: " << filename << std::endl;'
|
||||
)
|
||||
ib.writeline("return 1;")
|
||||
ib.writeline("}")
|
||||
ib.writeline(
|
||||
"std::vector<char> buffer((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());"
|
||||
)
|
||||
ib.writeline(
|
||||
"torch::IValue ivalue = torch::pickle_load(buffer);"
|
||||
)
|
||||
ib.writeline(
|
||||
f"input_tensors{i + 1}.push_back(ivalue.toTensor().to(device));"
|
||||
)
|
||||
ib.writeline("}")
|
||||
ib.newline()
|
||||
|
||||
ib.newline()
|
||||
ib.writeline("\n// Create array of input handles")
|
||||
for i in range(len(model_names)):
|
||||
ib.writelines(
|
||||
[
|
||||
f"auto input_handles{i + 1} =",
|
||||
f" torch::aot_inductor::unsafe_alloc_new_handles_from_tensors(input_tensors{i + 1});",
|
||||
]
|
||||
)
|
||||
|
||||
ib.writeline("\n// Create array for output handles")
|
||||
for i in range(len(model_names)):
|
||||
ib.writeline(f"AtenTensorHandle output_handle{i + 1};")
|
||||
|
||||
ib.writeline("\n// Create and load models")
|
||||
for i, model_name in enumerate(model_names):
|
||||
ib.writelines(
|
||||
[
|
||||
f"auto constants_map{i + 1} = std::make_shared<ConstantMap>();",
|
||||
f"auto constants_array{i + 1} = std::make_shared<std::vector<ConstantHandle>>();",
|
||||
f"auto model{i + 1} = std::make_unique<AOTInductorModel{model_name}>(",
|
||||
f" std::move(constants_map{i + 1}),",
|
||||
f" std::move(constants_array{i + 1}),",
|
||||
" device_str,",
|
||||
f' "{package_name}/data/aotinductor/{model_name}/");',
|
||||
f"model{i + 1}->load_constants();",
|
||||
]
|
||||
)
|
||||
|
||||
if example_inputs_map is not None:
|
||||
ib.writeline("\n// Run the models")
|
||||
for i in range(len(model_names)):
|
||||
ib.writeline(
|
||||
f"torch::aot_inductor::DeviceStreamType stream{i + 1} = nullptr;"
|
||||
)
|
||||
ib.writeline(
|
||||
f"model{i + 1}->run(&input_handles{i + 1}[0], &output_handle{i + 1}, stream{i + 1}, nullptr);"
|
||||
)
|
||||
|
||||
ib.writeline("\n// Convert output handles to tensors")
|
||||
for i in range(len(model_names)):
|
||||
ib.writelines(
|
||||
[
|
||||
f"auto output_tensor{i + 1} =",
|
||||
f" torch::aot_inductor::alloc_tensors_by_stealing_from_handles(&output_handle{i + 1}, 1);",
|
||||
]
|
||||
)
|
||||
|
||||
ib.writeline("\n// Validate outputs")
|
||||
for i in range(len(model_names)):
|
||||
ib.writeline(
|
||||
f"""std::cout << "output_tensor{i + 1}\\n" << output_tensor{i + 1} << std::endl;"""
|
||||
)
|
||||
ib.writeline(
|
||||
f"""torch::save(output_tensor{i + 1}, "output_tensor{i + 1}.pt");"""
|
||||
)
|
||||
|
||||
ib.writeline("return 0;")
|
||||
|
||||
ib.writelines(
|
||||
[
|
||||
"} catch (const std::exception &e) {",
|
||||
]
|
||||
)
|
||||
with ib.indent():
|
||||
ib.writeline('std::cerr << "Error: " << e.what() << std::endl;')
|
||||
ib.writeline("return 1;")
|
||||
|
||||
ib.writeline("}")
|
||||
ib.writeline("}")
|
||||
|
||||
return ib.getvalue()
|
||||
|
||||
|
||||
def _get_make_file(package_name: str, model_names: list[str], device_type: str) -> str:
|
||||
ib = IndentedBuffer()
|
||||
|
||||
ib.writelines(
|
||||
[
|
||||
"cmake_minimum_required(VERSION 3.10)",
|
||||
"project(TestProject)",
|
||||
"",
|
||||
"set(CMAKE_CXX_STANDARD 20)",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
from torch._inductor.config import test_configs
|
||||
|
||||
if test_configs.use_libtorch:
|
||||
ib.writeline("find_package(Torch REQUIRED)")
|
||||
|
||||
if device_type == "cuda":
|
||||
if torch.version.hip:
|
||||
ib.writeline("find_package(hip REQUIRED)")
|
||||
else:
|
||||
ib.writeline("find_package(CUDA REQUIRED)")
|
||||
|
||||
ib.newline()
|
||||
for model_name in model_names:
|
||||
ib.writeline(f"add_subdirectory({package_name}/data/aotinductor/{model_name}/)")
|
||||
|
||||
ib.writeline("\nadd_executable(main main.cpp)")
|
||||
if device_type == "cuda":
|
||||
if torch.version.hip:
|
||||
ib.writeline("target_compile_definitions(main PRIVATE USE_HIP)")
|
||||
else:
|
||||
ib.writeline("target_compile_definitions(main PRIVATE USE_CUDA)")
|
||||
elif device_type == "xpu":
|
||||
ib.writeline("target_compile_definitions(main PRIVATE USE_XPU)")
|
||||
|
||||
model_libs = " ".join(model_names)
|
||||
ib.writeline(f"target_link_libraries(main PRIVATE torch {model_libs})")
|
||||
|
||||
if device_type == "cuda":
|
||||
if torch.version.hip:
|
||||
ib.writeline("target_link_libraries(main PRIVATE hip::host)")
|
||||
else:
|
||||
ib.writeline("target_link_libraries(main PRIVATE cuda ${CUDA_LIBRARIES})")
|
||||
elif device_type == "xpu":
|
||||
ib.writeline("target_link_libraries(main PRIVATE sycl ze_loader)")
|
||||
return ib.getvalue()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,759 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import dataclasses
|
||||
from collections.abc import Collection, Mapping
|
||||
from enum import auto, Enum
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from torch._library.fake_class_registry import FakeScriptObject
|
||||
from torch._library.opaque_object import get_opaque_type_name, is_opaque_type
|
||||
from torch._subclasses.fake_tensor import is_fake
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
from torch._functorch._aot_autograd.schemas import GraphSignature
|
||||
|
||||
__all__ = [
|
||||
"ConstantArgument",
|
||||
"CustomObjArgument",
|
||||
"ExportBackwardSignature",
|
||||
"ExportGraphSignature",
|
||||
"InputKind",
|
||||
"InputSpec",
|
||||
"OutputKind",
|
||||
"OutputSpec",
|
||||
"SymIntArgument",
|
||||
"SymFloatArgument",
|
||||
"SymBoolArgument",
|
||||
"TensorArgument",
|
||||
]
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TensorArgument:
|
||||
name: str
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TokenArgument:
|
||||
name: str
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class SymIntArgument:
|
||||
name: str
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class SymFloatArgument:
|
||||
name: str
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class SymBoolArgument:
|
||||
name: str
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class CustomObjArgument:
|
||||
name: str
|
||||
class_fqn: str
|
||||
fake_val: FakeScriptObject | None = None
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ConstantArgument:
|
||||
name: str
|
||||
value: int | float | bool | str | None
|
||||
|
||||
|
||||
ArgumentSpec = (
|
||||
TensorArgument
|
||||
| SymIntArgument
|
||||
| SymFloatArgument
|
||||
| SymBoolArgument
|
||||
| ConstantArgument
|
||||
| CustomObjArgument
|
||||
| TokenArgument
|
||||
)
|
||||
|
||||
|
||||
class InputKind(Enum):
|
||||
USER_INPUT = auto()
|
||||
PARAMETER = auto()
|
||||
BUFFER = auto()
|
||||
CONSTANT_TENSOR = auto()
|
||||
CUSTOM_OBJ = auto()
|
||||
TOKEN = auto()
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class InputSpec:
|
||||
kind: InputKind
|
||||
arg: ArgumentSpec
|
||||
target: str | None
|
||||
persistent: bool | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.kind == InputKind.BUFFER:
|
||||
if self.persistent is None:
|
||||
raise AssertionError("Failed to specify persistent flag on BUFFER.")
|
||||
if not isinstance(
|
||||
self.arg,
|
||||
(
|
||||
TensorArgument,
|
||||
SymIntArgument,
|
||||
SymFloatArgument,
|
||||
SymBoolArgument,
|
||||
ConstantArgument,
|
||||
CustomObjArgument,
|
||||
TokenArgument,
|
||||
),
|
||||
):
|
||||
raise AssertionError(f"expected valid arg type, got {type(self.arg)}")
|
||||
|
||||
def __str__(self):
|
||||
target = "" if self.target is None else f" target='{self.target}'"
|
||||
persistent = "" if self.persistent is None else f" persistent={self.persistent}"
|
||||
return f"{str(self.arg.name)}: {str(self.kind.name)}{target}{persistent}"
|
||||
|
||||
|
||||
class OutputKind(Enum):
|
||||
USER_OUTPUT = auto()
|
||||
LOSS_OUTPUT = auto()
|
||||
BUFFER_MUTATION = auto()
|
||||
PARAMETER_MUTATION = auto()
|
||||
GRADIENT_TO_PARAMETER = auto()
|
||||
GRADIENT_TO_USER_INPUT = auto()
|
||||
USER_INPUT_MUTATION = auto()
|
||||
TOKEN = auto()
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class OutputSpec:
|
||||
kind: OutputKind
|
||||
arg: ArgumentSpec
|
||||
target: str | None
|
||||
|
||||
def __post_init__(self):
|
||||
if not isinstance(
|
||||
self.arg,
|
||||
(
|
||||
TensorArgument,
|
||||
SymIntArgument,
|
||||
SymFloatArgument,
|
||||
SymBoolArgument,
|
||||
ConstantArgument,
|
||||
TokenArgument,
|
||||
CustomObjArgument,
|
||||
),
|
||||
):
|
||||
raise AssertionError(f"expected valid arg type, got {self.arg}")
|
||||
|
||||
def __str__(self):
|
||||
target = "" if self.target is None else f" target='{self.target}'"
|
||||
return f"{str(self.arg.name)}: {str(self.kind.name)}{target}"
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ExportBackwardSignature:
|
||||
gradients_to_parameters: dict[str, str]
|
||||
gradients_to_user_inputs: dict[str, str]
|
||||
loss_output: str
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ExportGraphSignature:
|
||||
"""
|
||||
:class:`ExportGraphSignature` models the input/output signature of Export Graph,
|
||||
which is a fx.Graph with stronger invariants guarantees.
|
||||
|
||||
Export Graph is functional and does not access "states" like parameters
|
||||
or buffers within the graph via ``getattr`` nodes. Instead, :func:`export`
|
||||
guarantees that parameters, buffers, and constant tensors are lifted out of
|
||||
the graph as inputs. Similarly, any mutations to buffers are not included
|
||||
in the graph either, instead the updated values of mutated buffers are
|
||||
modeled as additional outputs of Export Graph.
|
||||
|
||||
The ordering of all inputs and outputs are::
|
||||
|
||||
Inputs = [*parameters_buffers_constant_tensors, *flattened_user_inputs]
|
||||
Outputs = [*mutated_inputs, *flattened_user_outputs]
|
||||
|
||||
e.g. If following module is exported::
|
||||
|
||||
class CustomModule(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super(CustomModule, self).__init__()
|
||||
|
||||
# Define a parameter
|
||||
self.my_parameter = nn.Parameter(torch.tensor(2.0))
|
||||
|
||||
# Define two buffers
|
||||
self.register_buffer("my_buffer1", torch.tensor(3.0))
|
||||
self.register_buffer("my_buffer2", torch.tensor(4.0))
|
||||
|
||||
def forward(self, x1, x2):
|
||||
# Use the parameter, buffers, and both inputs in the forward method
|
||||
output = (
|
||||
x1 + self.my_parameter
|
||||
) * self.my_buffer1 + x2 * self.my_buffer2
|
||||
|
||||
# Mutate one of the buffers (e.g., increment it by 1)
|
||||
self.my_buffer2.add_(1.0) # In-place addition
|
||||
|
||||
return output
|
||||
|
||||
|
||||
mod = CustomModule()
|
||||
ep = torch.export.export(mod, (torch.tensor(1.0), torch.tensor(2.0)))
|
||||
|
||||
Resulting Graph is non-functional::
|
||||
|
||||
graph():
|
||||
%p_my_parameter : [num_users=1] = placeholder[target=p_my_parameter]
|
||||
%b_my_buffer1 : [num_users=1] = placeholder[target=b_my_buffer1]
|
||||
%b_my_buffer2 : [num_users=2] = placeholder[target=b_my_buffer2]
|
||||
%x1 : [num_users=1] = placeholder[target=x1]
|
||||
%x2 : [num_users=1] = placeholder[target=x2]
|
||||
%add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x1, %p_my_parameter), kwargs = {})
|
||||
%mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add, %b_my_buffer1), kwargs = {})
|
||||
%mul_1 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%x2, %b_my_buffer2), kwargs = {})
|
||||
%add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul, %mul_1), kwargs = {})
|
||||
%add_ : [num_users=0] = call_function[target=torch.ops.aten.add_.Tensor](args = (%b_my_buffer2, 1.0), kwargs = {})
|
||||
return (add_1,)
|
||||
|
||||
Resulting ExportGraphSignature of the non-functional Graph would be::
|
||||
|
||||
# inputs
|
||||
p_my_parameter: PARAMETER target='my_parameter'
|
||||
b_my_buffer1: BUFFER target='my_buffer1' persistent=True
|
||||
b_my_buffer2: BUFFER target='my_buffer2' persistent=True
|
||||
x1: USER_INPUT
|
||||
x2: USER_INPUT
|
||||
|
||||
# outputs
|
||||
add_1: USER_OUTPUT
|
||||
|
||||
To get a functional Graph, you can use :func:`run_decompositions`::
|
||||
|
||||
mod = CustomModule()
|
||||
ep = torch.export.export(mod, (torch.tensor(1.0), torch.tensor(2.0)))
|
||||
ep = ep.run_decompositions()
|
||||
|
||||
Resulting Graph is functional::
|
||||
|
||||
graph():
|
||||
%p_my_parameter : [num_users=1] = placeholder[target=p_my_parameter]
|
||||
%b_my_buffer1 : [num_users=1] = placeholder[target=b_my_buffer1]
|
||||
%b_my_buffer2 : [num_users=2] = placeholder[target=b_my_buffer2]
|
||||
%x1 : [num_users=1] = placeholder[target=x1]
|
||||
%x2 : [num_users=1] = placeholder[target=x2]
|
||||
%add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x1, %p_my_parameter), kwargs = {})
|
||||
%mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add, %b_my_buffer1), kwargs = {})
|
||||
%mul_1 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%x2, %b_my_buffer2), kwargs = {})
|
||||
%add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul, %mul_1), kwargs = {})
|
||||
%add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%b_my_buffer2, 1.0), kwargs = {})
|
||||
return (add_2, add_1)
|
||||
|
||||
Resulting ExportGraphSignature of the functional Graph would be::
|
||||
|
||||
# inputs
|
||||
p_my_parameter: PARAMETER target='my_parameter'
|
||||
b_my_buffer1: BUFFER target='my_buffer1' persistent=True
|
||||
b_my_buffer2: BUFFER target='my_buffer2' persistent=True
|
||||
x1: USER_INPUT
|
||||
x2: USER_INPUT
|
||||
|
||||
# outputs
|
||||
add_2: BUFFER_MUTATION target='my_buffer2'
|
||||
add_1: USER_OUTPUT
|
||||
|
||||
"""
|
||||
|
||||
input_specs: list[InputSpec]
|
||||
output_specs: list[OutputSpec]
|
||||
|
||||
# A list of parameters uniquely identified by mangled fully qualified name
|
||||
@property
|
||||
def parameters(self) -> Collection[str]:
|
||||
return tuple(
|
||||
s.target
|
||||
for s in self.input_specs
|
||||
if s.kind == InputKind.PARAMETER
|
||||
if isinstance(s.target, str)
|
||||
)
|
||||
|
||||
# A list of buffers uniquely identified by mangled fully qualified name
|
||||
@property
|
||||
def buffers(self) -> Collection[str]:
|
||||
return tuple(
|
||||
s.target
|
||||
for s in self.input_specs
|
||||
if s.kind == InputKind.BUFFER
|
||||
if isinstance(s.target, str)
|
||||
)
|
||||
|
||||
@property
|
||||
def non_persistent_buffers(self) -> Collection[str]:
|
||||
return tuple(
|
||||
s.target
|
||||
for s in self.input_specs
|
||||
if s.kind == InputKind.BUFFER
|
||||
if s.persistent is False
|
||||
if isinstance(s.target, str)
|
||||
)
|
||||
|
||||
# A list of lifted constant tensors
|
||||
@property
|
||||
def lifted_tensor_constants(self) -> Collection[str]:
|
||||
return tuple(
|
||||
s.target
|
||||
for s in self.input_specs
|
||||
if s.kind == InputKind.CONSTANT_TENSOR
|
||||
if isinstance(s.target, str)
|
||||
)
|
||||
|
||||
@property
|
||||
def lifted_custom_objs(self) -> Collection[str]:
|
||||
return tuple(
|
||||
s.target
|
||||
for s in self.input_specs
|
||||
if s.kind == InputKind.CUSTOM_OBJ
|
||||
if isinstance(s.target, str)
|
||||
)
|
||||
|
||||
# Graph node names of pytree-flattened inputs of original program
|
||||
@property
|
||||
def user_inputs(self) -> Collection[int | float | bool | str | None]:
|
||||
user_inputs: list[int | float | bool | str | None] = []
|
||||
for s in self.input_specs:
|
||||
if s.kind != InputKind.USER_INPUT:
|
||||
continue
|
||||
|
||||
if isinstance(
|
||||
s.arg,
|
||||
(
|
||||
TensorArgument,
|
||||
SymIntArgument,
|
||||
SymFloatArgument,
|
||||
SymBoolArgument,
|
||||
CustomObjArgument,
|
||||
),
|
||||
):
|
||||
user_inputs.append(s.arg.name)
|
||||
elif isinstance(s.arg, ConstantArgument):
|
||||
user_inputs.append(s.arg.value)
|
||||
else:
|
||||
raise RuntimeError(f"{s.arg} is not a valid user inputs")
|
||||
return tuple(user_inputs)
|
||||
|
||||
# Graph node names of pytree-flattened outputs of original program
|
||||
# For joint-graph purposes, will include the loss output.
|
||||
@property
|
||||
def user_outputs(self) -> Collection[int | float | bool | str | None]:
|
||||
user_outputs: list[int | float | bool | str | None] = []
|
||||
for s in self.output_specs:
|
||||
if s.kind not in [
|
||||
OutputKind.USER_OUTPUT,
|
||||
OutputKind.LOSS_OUTPUT,
|
||||
]:
|
||||
continue
|
||||
|
||||
if isinstance(
|
||||
s.arg,
|
||||
(TensorArgument, SymIntArgument, SymFloatArgument, SymBoolArgument),
|
||||
):
|
||||
user_outputs.append(s.arg.name)
|
||||
elif isinstance(s.arg, ConstantArgument):
|
||||
user_outputs.append(s.arg.value)
|
||||
elif isinstance(s.arg, CustomObjArgument):
|
||||
user_outputs.append(s.arg.name)
|
||||
else:
|
||||
raise RuntimeError(f"{s.arg} is not a valid user output")
|
||||
return tuple(user_outputs)
|
||||
|
||||
# A dictionary mapping graph input node names to parameters. If a graph input
|
||||
# name is found in this dictionary, it is guaranteed to be a lifted parameter.
|
||||
@property
|
||||
def inputs_to_parameters(self) -> Mapping[str, str]:
|
||||
return _immutable_dict(
|
||||
(s.arg.name, s.target)
|
||||
for s in self.input_specs
|
||||
if s.kind == InputKind.PARAMETER
|
||||
and isinstance(s.arg, TensorArgument)
|
||||
and isinstance(s.target, str)
|
||||
)
|
||||
|
||||
# A dictionary mapping graph input node names to buffers. If a graph input
|
||||
# name is found in this dictionary, it is guaranteed to be a lifted buffer.
|
||||
@property
|
||||
def inputs_to_buffers(self) -> Mapping[str, str]:
|
||||
return _immutable_dict(
|
||||
(s.arg.name, s.target) # type: ignore[union-attr, misc]
|
||||
for s in self.input_specs
|
||||
if s.kind == InputKind.BUFFER
|
||||
and isinstance(s.arg, TensorArgument)
|
||||
and isinstance(s.target, str)
|
||||
)
|
||||
|
||||
# A dictionary mapping graph output node names to buffers that are mutated in the
|
||||
# original program. Buffers that are not mutated will not be found in this dictionary.
|
||||
@property
|
||||
def buffers_to_mutate(self) -> Mapping[str, str]:
|
||||
return _immutable_dict(
|
||||
(s.arg.name, s.target)
|
||||
for s in self.output_specs
|
||||
if s.kind == OutputKind.BUFFER_MUTATION
|
||||
and isinstance(s.arg, TensorArgument)
|
||||
and isinstance(s.target, str)
|
||||
)
|
||||
|
||||
@property
|
||||
def parameters_to_mutate(self) -> Mapping[str, str]:
|
||||
return _immutable_dict(
|
||||
(s.arg.name, s.target)
|
||||
for s in self.output_specs
|
||||
if s.kind == OutputKind.PARAMETER_MUTATION
|
||||
and isinstance(s.arg, TensorArgument)
|
||||
and isinstance(s.target, str)
|
||||
)
|
||||
|
||||
@property
|
||||
def user_inputs_to_mutate(self) -> Mapping[str, str]:
|
||||
return _immutable_dict(
|
||||
(s.arg.name, s.target)
|
||||
for s in self.output_specs
|
||||
if s.kind == OutputKind.USER_INPUT_MUTATION
|
||||
and isinstance(s.arg, TensorArgument)
|
||||
and isinstance(s.target, str)
|
||||
)
|
||||
|
||||
# A dictionary mapping graph input node names to lifted tensor constants.
|
||||
@property
|
||||
def inputs_to_lifted_tensor_constants(self) -> Mapping[str, str]:
|
||||
return _immutable_dict(
|
||||
(s.arg.name, s.target)
|
||||
for s in self.input_specs
|
||||
if s.kind == InputKind.CONSTANT_TENSOR
|
||||
and isinstance(s.arg, TensorArgument)
|
||||
and isinstance(s.target, str)
|
||||
)
|
||||
|
||||
@property
|
||||
def inputs_to_lifted_custom_objs(self) -> Mapping[str, str]:
|
||||
return _immutable_dict(
|
||||
(s.arg.name, s.target)
|
||||
for s in self.input_specs
|
||||
if s.kind == InputKind.CUSTOM_OBJ
|
||||
and isinstance(s.arg, CustomObjArgument)
|
||||
and isinstance(s.target, str)
|
||||
)
|
||||
|
||||
@property
|
||||
def backward_signature(self) -> ExportBackwardSignature | None:
|
||||
loss_output = None
|
||||
gradients_to_parameters: dict[str, str] = {}
|
||||
gradients_to_user_inputs: dict[str, str] = {}
|
||||
for spec in self.output_specs:
|
||||
if spec.kind == OutputKind.LOSS_OUTPUT:
|
||||
if loss_output is not None:
|
||||
raise AssertionError("multiple LOSS_OUTPUT specs found")
|
||||
if not isinstance(spec.arg, TensorArgument):
|
||||
raise AssertionError(
|
||||
f"expected TensorArgument for LOSS_OUTPUT, got {type(spec.arg)}"
|
||||
)
|
||||
loss_output = spec.arg.name
|
||||
elif spec.kind == OutputKind.GRADIENT_TO_PARAMETER:
|
||||
if not isinstance(spec.target, str):
|
||||
raise AssertionError(
|
||||
f"expected str target for GRADIENT_TO_PARAMETER, got {type(spec.target)}"
|
||||
)
|
||||
if not isinstance(spec.arg, TensorArgument):
|
||||
raise AssertionError(
|
||||
f"expected TensorArgument for GRADIENT_TO_PARAMETER, got {type(spec.arg)}"
|
||||
)
|
||||
gradients_to_parameters[spec.arg.name] = spec.target
|
||||
elif spec.kind == OutputKind.GRADIENT_TO_USER_INPUT:
|
||||
if not isinstance(spec.target, str):
|
||||
raise AssertionError(
|
||||
f"expected str target for GRADIENT_TO_USER_INPUT, got {type(spec.target)}"
|
||||
)
|
||||
if not isinstance(spec.arg, TensorArgument):
|
||||
raise AssertionError(
|
||||
f"expected TensorArgument for GRADIENT_TO_USER_INPUT, got {type(spec.arg)}"
|
||||
)
|
||||
gradients_to_user_inputs[spec.arg.name] = spec.target
|
||||
|
||||
if loss_output is None:
|
||||
return None
|
||||
|
||||
return ExportBackwardSignature(
|
||||
loss_output=loss_output,
|
||||
gradients_to_parameters=gradients_to_parameters,
|
||||
gradients_to_user_inputs=gradients_to_user_inputs,
|
||||
)
|
||||
|
||||
# Map from assertion dependency token index to assertion dep token output
|
||||
# name in output. The shape of output after aot_autograd will be like:
|
||||
# (updated_inputs, user_outputs, dep_token).
|
||||
@property
|
||||
def assertion_dep_token(self) -> Mapping[int, str] | None:
|
||||
return None
|
||||
|
||||
@property
|
||||
def input_tokens(self) -> Collection[str]:
|
||||
input_tokens = []
|
||||
for s in self.input_specs:
|
||||
if s.kind == InputKind.TOKEN:
|
||||
if not isinstance(s.arg, TokenArgument):
|
||||
raise AssertionError(
|
||||
f"expected TokenArgument for TOKEN kind, got {type(s.arg)}"
|
||||
)
|
||||
input_tokens.append(s.arg.name)
|
||||
return tuple(input_tokens)
|
||||
|
||||
@property
|
||||
def output_tokens(self) -> Collection[str]:
|
||||
output_tokens = []
|
||||
for s in self.output_specs:
|
||||
if s.kind == OutputKind.TOKEN:
|
||||
if not isinstance(s.arg, TokenArgument):
|
||||
raise AssertionError(
|
||||
f"expected TokenArgument for TOKEN kind, got {type(s.arg)}"
|
||||
)
|
||||
output_tokens.append(s.arg.name)
|
||||
return tuple(output_tokens)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
assertion_dep_token = self.assertion_dep_token
|
||||
if assertion_dep_token is None:
|
||||
return
|
||||
if len(assertion_dep_token) != 1:
|
||||
raise AssertionError(
|
||||
f"expected exactly 1 assertion_dep_token, got {len(assertion_dep_token)}"
|
||||
)
|
||||
assertion_dep_token_index = next(iter(assertion_dep_token.keys()))
|
||||
expected_index = len(self.user_outputs) + len(self.buffers_to_mutate)
|
||||
if expected_index != assertion_dep_token_index:
|
||||
raise AssertionError(
|
||||
f"expected assertion_dep_token_index to be {expected_index}, got {assertion_dep_token_index}"
|
||||
)
|
||||
|
||||
def replace_all_uses(self, old: str, new: str):
|
||||
"""
|
||||
Replace all uses of the old name with new name in the signature.
|
||||
"""
|
||||
if not isinstance(old, str):
|
||||
raise AssertionError(f"expected old to be str, got {type(old)}")
|
||||
if not isinstance(new, str):
|
||||
raise AssertionError(f"expected new to be str, got {type(new)}")
|
||||
arg_types = (
|
||||
TensorArgument,
|
||||
SymIntArgument,
|
||||
SymFloatArgument,
|
||||
SymBoolArgument,
|
||||
CustomObjArgument,
|
||||
TokenArgument,
|
||||
)
|
||||
for o in self.output_specs:
|
||||
if isinstance(o.arg, arg_types):
|
||||
if o.arg.name == old:
|
||||
o.arg.name = new
|
||||
for i in self.input_specs:
|
||||
if isinstance(i.arg, arg_types):
|
||||
if i.arg.name == old:
|
||||
i.arg.name = new
|
||||
|
||||
def get_replace_hook(self, replace_inputs=False):
|
||||
def _(old, new, user):
|
||||
if user.op == "output":
|
||||
self.replace_all_uses(old.name, new)
|
||||
if replace_inputs and old.op == "placeholder":
|
||||
self.replace_all_uses(old.name, new)
|
||||
|
||||
return _
|
||||
|
||||
def __str__(self):
|
||||
input_specs = "\n".join(str(s) for s in self.input_specs)
|
||||
output_specs = "\n".join(str(s) for s in self.output_specs)
|
||||
return f"\n# inputs\n{input_specs}\n\n# outputs\n{output_specs}\n"
|
||||
|
||||
|
||||
def _immutable_dict(items):
|
||||
"""
|
||||
Creates a mapping where items cannot be added, deleted, or updated.
|
||||
NOTE: The immutability is shallow (like tuple is an immutable collection).
|
||||
"""
|
||||
from types import MappingProxyType
|
||||
|
||||
return MappingProxyType(dict(items))
|
||||
|
||||
|
||||
def _make_argument_spec(node, token_names) -> ArgumentSpec:
|
||||
from torch import ScriptObject, SymBool, SymFloat, SymInt
|
||||
from torch._library.fake_class_registry import FakeScriptObject
|
||||
|
||||
if isinstance(node, (int, bool, float, type(None), str)):
|
||||
# For const outputs we just directly return this
|
||||
return ConstantArgument(name="", value=node)
|
||||
|
||||
if "val" not in node.meta:
|
||||
raise AssertionError(
|
||||
f"{node} is not a constant or a node with a 'val' metadata field"
|
||||
)
|
||||
val = node.meta["val"]
|
||||
if node.name in token_names:
|
||||
return TokenArgument(name=node.name)
|
||||
elif is_fake(val):
|
||||
return TensorArgument(name=node.name)
|
||||
elif isinstance(val, SymInt):
|
||||
return SymIntArgument(name=node.name)
|
||||
elif isinstance(val, SymFloat):
|
||||
return SymFloatArgument(name=node.name)
|
||||
elif isinstance(val, SymBool):
|
||||
return SymBoolArgument(name=node.name)
|
||||
elif isinstance(val, ScriptObject):
|
||||
return CustomObjArgument(name=node.name, class_fqn=val._type().qualified_name()) # type: ignore[attr-defined]
|
||||
elif isinstance(val, FakeScriptObject):
|
||||
return CustomObjArgument(
|
||||
name=node.name, class_fqn=val.script_class_name, fake_val=val
|
||||
)
|
||||
elif is_opaque_type(type(val)):
|
||||
return CustomObjArgument(
|
||||
name=node.name, class_fqn=get_opaque_type_name(type(val)), fake_val=val
|
||||
)
|
||||
elif isinstance(val, (int, bool, str, float, type(None))):
|
||||
return ConstantArgument(name=node.name, value=val)
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"Encountered an unsupported object of type {type(val)} "
|
||||
f"while writing the metadata for exported program"
|
||||
)
|
||||
|
||||
|
||||
def _convert_to_export_graph_signature(
|
||||
graph_signature: "GraphSignature",
|
||||
gm: "torch.fx.GraphModule",
|
||||
non_persistent_buffers: set[str],
|
||||
) -> "ExportGraphSignature":
|
||||
from torch.utils import _pytree as pytree
|
||||
|
||||
is_joint = graph_signature.backward_signature is not None
|
||||
|
||||
# unpack objects
|
||||
user_inputs = set(graph_signature.user_inputs)
|
||||
inputs_to_parameters = graph_signature.inputs_to_parameters
|
||||
inputs_to_buffers = graph_signature.inputs_to_buffers
|
||||
user_outputs = set(graph_signature.user_outputs)
|
||||
buffer_mutations = graph_signature.buffers_to_mutate
|
||||
parameter_mutations = graph_signature.parameters_to_mutate
|
||||
user_input_mutations = graph_signature.user_inputs_to_mutate
|
||||
grad_params = (
|
||||
graph_signature.backward_signature.gradients_to_parameter # type: ignore[union-attr]
|
||||
if is_joint
|
||||
else {}
|
||||
)
|
||||
grad_user_inputs = (
|
||||
graph_signature.backward_signature.gradients_to_user_inputs # type: ignore[union-attr]
|
||||
if is_joint
|
||||
else {}
|
||||
)
|
||||
loss_output = (
|
||||
graph_signature.backward_signature.loss_output # type: ignore[union-attr]
|
||||
if is_joint
|
||||
else None
|
||||
)
|
||||
input_tokens = graph_signature.input_tokens
|
||||
output_tokens = graph_signature.output_tokens
|
||||
|
||||
inputs = [
|
||||
_make_argument_spec(node, input_tokens)
|
||||
for node in gm.graph.nodes
|
||||
if node.op == "placeholder"
|
||||
]
|
||||
outputs = [
|
||||
_make_argument_spec(node, output_tokens)
|
||||
for node in pytree.tree_leaves(next(iter(reversed(gm.graph.nodes))).args)
|
||||
]
|
||||
|
||||
def to_input_spec(inp: ArgumentSpec) -> InputSpec:
|
||||
if isinstance(inp, TokenArgument):
|
||||
return InputSpec(kind=InputKind.TOKEN, arg=inp, target=None)
|
||||
|
||||
if not isinstance(inp, TensorArgument):
|
||||
return InputSpec(kind=InputKind.USER_INPUT, arg=inp, target=None)
|
||||
name = inp.name
|
||||
if name in user_inputs:
|
||||
return InputSpec(kind=InputKind.USER_INPUT, arg=inp, target=None)
|
||||
elif name in inputs_to_parameters:
|
||||
return InputSpec(
|
||||
kind=InputKind.PARAMETER,
|
||||
arg=inp,
|
||||
target=inputs_to_parameters[name], # type: ignore[index]
|
||||
)
|
||||
elif name in inputs_to_buffers:
|
||||
return InputSpec(
|
||||
kind=InputKind.BUFFER,
|
||||
arg=inp,
|
||||
target=inputs_to_buffers[name], # type: ignore[index]
|
||||
persistent=(inputs_to_buffers[name] not in non_persistent_buffers), # type: ignore[index]
|
||||
)
|
||||
else:
|
||||
raise AssertionError(f"Unknown tensor input kind: {name}")
|
||||
|
||||
def to_output_spec(idx: int, o: ArgumentSpec) -> OutputSpec:
|
||||
if isinstance(o, TokenArgument):
|
||||
return OutputSpec(kind=OutputKind.TOKEN, arg=o, target=None)
|
||||
|
||||
if not isinstance(o, TensorArgument):
|
||||
return OutputSpec(kind=OutputKind.USER_OUTPUT, arg=o, target=None)
|
||||
name = o.name
|
||||
if idx < len(buffer_mutations) + len(parameter_mutations) + len(
|
||||
user_input_mutations
|
||||
) + len(output_tokens):
|
||||
if name in buffer_mutations:
|
||||
return OutputSpec(
|
||||
kind=OutputKind.BUFFER_MUTATION,
|
||||
arg=o,
|
||||
target=buffer_mutations[name], # type: ignore[index]
|
||||
)
|
||||
elif name in parameter_mutations:
|
||||
return OutputSpec(
|
||||
kind=OutputKind.PARAMETER_MUTATION,
|
||||
arg=o,
|
||||
target=parameter_mutations[name], # type: ignore[index]
|
||||
)
|
||||
elif name in user_input_mutations:
|
||||
return OutputSpec(
|
||||
kind=OutputKind.USER_INPUT_MUTATION,
|
||||
arg=o,
|
||||
target=user_input_mutations[name], # type: ignore[index]
|
||||
)
|
||||
else:
|
||||
raise AssertionError(f"Unknown tensor mutation kind: {name}")
|
||||
else:
|
||||
if name in user_outputs:
|
||||
return OutputSpec(kind=OutputKind.USER_OUTPUT, arg=o, target=None)
|
||||
|
||||
elif name in grad_params:
|
||||
return OutputSpec(
|
||||
kind=OutputKind.GRADIENT_TO_PARAMETER,
|
||||
arg=o,
|
||||
target=grad_params[name],
|
||||
)
|
||||
elif name in grad_user_inputs:
|
||||
return OutputSpec(
|
||||
kind=OutputKind.GRADIENT_TO_USER_INPUT,
|
||||
arg=o,
|
||||
target=grad_user_inputs[name],
|
||||
)
|
||||
elif name == loss_output:
|
||||
return OutputSpec(kind=OutputKind.LOSS_OUTPUT, arg=o, target=None)
|
||||
|
||||
else:
|
||||
raise AssertionError(f"Unknown tensor output kind: {name}")
|
||||
|
||||
input_specs = [to_input_spec(inp) for inp in inputs]
|
||||
output_specs = [to_output_spec(idx, o) for idx, o in enumerate(outputs)]
|
||||
return ExportGraphSignature(input_specs=input_specs, output_specs=output_specs)
|
||||
@@ -0,0 +1,97 @@
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch.export.exported_program import ExportedProgram
|
||||
|
||||
|
||||
__all__ = ["move_to_device_pass"]
|
||||
|
||||
|
||||
def move_to_device_pass(
|
||||
ep: ExportedProgram, location: torch.device | str | dict[str, str]
|
||||
) -> ExportedProgram:
|
||||
"""
|
||||
Move the exported program to the given device.
|
||||
|
||||
Args:
|
||||
ep (ExportedProgram): The exported program to move.
|
||||
location (Union[torch.device, str, Dict[str, str]]): The device to move the exported program to.
|
||||
If a string, it is interpreted as a device name.
|
||||
If a dict, it is interpreted as a mapping from
|
||||
the existing device to the intended one
|
||||
|
||||
Returns:
|
||||
ExportedProgram: The moved exported program.
|
||||
"""
|
||||
|
||||
def _get_new_device(
|
||||
curr_device: torch.device,
|
||||
location: torch.device | str | dict[str, str],
|
||||
) -> str:
|
||||
if isinstance(location, dict):
|
||||
if str(curr_device) in location:
|
||||
return location[str(curr_device)]
|
||||
else:
|
||||
return str(curr_device)
|
||||
else:
|
||||
return str(location)
|
||||
|
||||
# move all the state_dict
|
||||
for k, v in ep.state_dict.items():
|
||||
if isinstance(v, torch.nn.Parameter):
|
||||
ep._state_dict[k] = torch.nn.Parameter(
|
||||
v.to(_get_new_device(v.device, location)),
|
||||
v.requires_grad,
|
||||
)
|
||||
else:
|
||||
ep._state_dict[k] = v.to(_get_new_device(v.device, location))
|
||||
|
||||
# move all the constants
|
||||
for k, v in ep.constants.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
ep._constants[k] = v.to(_get_new_device(v.device, location))
|
||||
|
||||
# move example_inputs if they exist
|
||||
if ep.example_inputs is not None:
|
||||
args, kwargs = ep.example_inputs
|
||||
moved_args = pytree.tree_map_only(
|
||||
torch.Tensor,
|
||||
lambda tensor: tensor.to(_get_new_device(tensor.device, location)),
|
||||
args,
|
||||
)
|
||||
moved_kwargs = pytree.tree_map_only(
|
||||
torch.Tensor,
|
||||
lambda tensor: tensor.to(_get_new_device(tensor.device, location)),
|
||||
kwargs,
|
||||
)
|
||||
ep._example_inputs = (moved_args, moved_kwargs)
|
||||
|
||||
for m in ep.graph_module.modules():
|
||||
if isinstance(m, torch.fx.GraphModule):
|
||||
for node in m.graph.nodes:
|
||||
# move all the nodes kwargs with burnt-in device
|
||||
if "device" in node.kwargs:
|
||||
kwargs = node.kwargs.copy()
|
||||
kwargs["device"] = _get_new_device(kwargs["device"], location)
|
||||
node.kwargs = kwargs
|
||||
|
||||
if (
|
||||
node.op == "call_function"
|
||||
and node.target is torch.ops.aten.to.device
|
||||
):
|
||||
args = list(node.args)
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
args[1] = _get_new_device(args[1], location)
|
||||
node.args = tuple(args)
|
||||
|
||||
# move all the tensor metadata
|
||||
node.meta["val"] = pytree.tree_map(
|
||||
lambda v: v.to(_get_new_device(v.device, location))
|
||||
if isinstance(v, torch.Tensor)
|
||||
else v,
|
||||
node.meta.get("val"),
|
||||
)
|
||||
|
||||
ep.validate()
|
||||
return ep
|
||||
@@ -0,0 +1,4 @@
|
||||
from ._package import is_pt2_package, PT2ArchiveReader, PT2ArchiveWriter
|
||||
|
||||
|
||||
__all__ = ["PT2ArchiveWriter", "PT2ArchiveReader", "is_pt2_package"]
|
||||
File diff suppressed because it is too large
Load Diff
+213
@@ -0,0 +1,213 @@
|
||||
import collections
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensor
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
|
||||
def _end_ptr(tensor: torch.Tensor) -> int:
|
||||
if tensor.nelement():
|
||||
stop = tensor.view(-1)[-1].data_ptr() + tensor.element_size()
|
||||
else:
|
||||
stop = tensor.data_ptr()
|
||||
return stop
|
||||
|
||||
|
||||
class TensorProperties:
|
||||
def __init__(self, tensor: torch.Tensor):
|
||||
self.is_fake = isinstance(tensor, FakeTensor)
|
||||
self.is_contiguous = tensor.is_contiguous()
|
||||
self.storage_ptr = None
|
||||
self.storage_size = None
|
||||
self.start = None
|
||||
self.end = None
|
||||
|
||||
if not self.is_fake:
|
||||
# only get the storage pointer for real tensors
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.storage_ptr = tensor.untyped_storage().data_ptr()
|
||||
if self.is_contiguous:
|
||||
# only get storage size and start/end pointers for contiguous tensors
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.storage_size = tensor.untyped_storage().nbytes()
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.start = tensor.data_ptr()
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.end = _end_ptr(tensor)
|
||||
|
||||
# info to recover tensor
|
||||
self.shape = tensor.shape
|
||||
self.stride = tensor.stride()
|
||||
self.offset = tensor.storage_offset()
|
||||
|
||||
def is_complete(self) -> bool:
|
||||
"""
|
||||
Whether the tensor completely overlaps with its underlying storage
|
||||
"""
|
||||
if self.is_fake:
|
||||
# Theoretically, fake tensors should not appear in weights
|
||||
# But we handle this corner case to make it always complete
|
||||
return True
|
||||
if not self.is_contiguous:
|
||||
return False
|
||||
|
||||
if self.storage_ptr is None:
|
||||
raise AssertionError("storage_ptr cannot be None for complete check")
|
||||
if self.storage_size is None:
|
||||
raise AssertionError("storage_size cannot be None for complete check")
|
||||
if self.start is None:
|
||||
raise AssertionError("start cannot be None for complete check")
|
||||
if self.end is None:
|
||||
raise AssertionError("end cannot be None for complete check")
|
||||
return (
|
||||
self.start == self.storage_ptr
|
||||
and self.end == self.storage_ptr + self.storage_size
|
||||
)
|
||||
|
||||
|
||||
class Weights(dict):
|
||||
"""
|
||||
A dictionary mapping from weight name to a tuple of (tensor, TensorProperties).
|
||||
tensor represents the actual initial value of the weight.
|
||||
TensorProperties represents the properties of the weight that are needed to recover the weight.
|
||||
|
||||
We use two separate entries because `tensor` could be a clone of the original weight tensor,
|
||||
so it doesn't have the same property as the original weight (such as underlying storage pointer).
|
||||
"""
|
||||
|
||||
def __init__(self, weight_dict: dict[str, tuple[torch.Tensor, TensorProperties]]):
|
||||
super().__init__(weight_dict)
|
||||
|
||||
def get_weight(self, name: str) -> tuple[torch.Tensor, TensorProperties]:
|
||||
return self[name]
|
||||
|
||||
def get_weight_properties(self, name: str) -> TensorProperties:
|
||||
return self[name][1]
|
||||
|
||||
|
||||
def get_complete_tensor(
|
||||
group: OrderedSet[tuple[str, str]], models_weights: dict[str, Weights]
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Given a group of (model_name, weight_name) pairs whose tensors share the same
|
||||
underlying storage, return the complete (maximal) tensor covering that storage
|
||||
region.
|
||||
|
||||
This function handles two cases:
|
||||
|
||||
1. If any tensor in the group is already marked as complete, return it directly.
|
||||
2. Otherwise, all tensors in the group are assumed to be slices of a larger,
|
||||
contiguous tensor backed by the same storage. In this case, reconstruct
|
||||
the complete tensor by taking the union of their storage ranges. This assumes
|
||||
all tensors in the group have the same dtype.
|
||||
|
||||
Args:
|
||||
group: Set of (model_name, weight_name) tuples identifying tensors that share storage.
|
||||
models_weights: Dictionary mapping model names to their Weights objects.
|
||||
|
||||
Returns:
|
||||
The complete tensor (either found directly or reconstructed from slices).
|
||||
|
||||
Example:
|
||||
# Tensors a, b, c share storage:
|
||||
# a = full_tensor[0:5] -> start=addr_0, end=addr_5
|
||||
# b = full_tensor[3:8] -> start=addr_3, end=addr_8
|
||||
# c = full_tensor -> complete tensor
|
||||
|
||||
# Case 1: If c is in group -> return c
|
||||
# Case 2: If only a, b in group -> reconstruct from addr_0 to addr_8
|
||||
"""
|
||||
|
||||
if len(group) == 0:
|
||||
raise AssertionError("group cannot be empty")
|
||||
|
||||
start_addr = None
|
||||
end_addr = None
|
||||
for model_name, weight_name in group:
|
||||
tensor_property = models_weights[model_name].get_weight_properties(weight_name)
|
||||
|
||||
# Case 1: Found a complete tensor.
|
||||
if tensor_property.is_complete():
|
||||
return models_weights[model_name].get_weight(weight_name)[0]
|
||||
|
||||
# Case 2: Track the widest boundary across all slices.
|
||||
if tensor_property.start is not None:
|
||||
start_addr = (
|
||||
tensor_property.start
|
||||
if start_addr is None
|
||||
else min(start_addr, tensor_property.start)
|
||||
)
|
||||
if tensor_property.end is not None:
|
||||
end_addr = (
|
||||
tensor_property.end
|
||||
if end_addr is None
|
||||
else max(end_addr, tensor_property.end)
|
||||
)
|
||||
|
||||
# Case 2: Reconstruct complete tensor from slices.
|
||||
# Pick any tensor from the group as a reference (they all share the same storage).
|
||||
warnings.warn(
|
||||
"No complete tensor found in the group! Returning the first one. "
|
||||
"This may cause issues when your weights are not on CPU.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
model_name, weight_name = next(iter(group))
|
||||
reference_tensor = models_weights[model_name].get_weight(weight_name)[0]
|
||||
|
||||
# If no boundary information available (e.g., FakeTensor), return reference tensor as is.
|
||||
if start_addr is None and end_addr is None:
|
||||
return reference_tensor
|
||||
|
||||
# Validate that we have both boundaries.
|
||||
if start_addr is None or end_addr is None:
|
||||
raise AssertionError(
|
||||
f"Inconsistent boundary information: start={start_addr}, end={end_addr}. "
|
||||
"Unable to reconstruct complete tensor from group."
|
||||
)
|
||||
|
||||
# Reconstruct a view over the full contiguous storage range.
|
||||
storage = reference_tensor.untyped_storage()
|
||||
total_size_bytes = end_addr - storage.data_ptr()
|
||||
element_size = reference_tensor.element_size()
|
||||
# It assumes all tensors in the group have the same dtype.
|
||||
total_size = total_size_bytes // element_size
|
||||
|
||||
# Validate alignment: size must be multiples of element_size.
|
||||
if total_size_bytes % element_size != 0:
|
||||
raise AssertionError(
|
||||
f"Total size ({total_size_bytes} bytes) is not aligned with "
|
||||
f"element size ({element_size} bytes). Cannot reconstruct tensor safely. "
|
||||
f"Expected size to be a multiple of {element_size}."
|
||||
)
|
||||
|
||||
# Reconstruct a tensor that spans the needed storage range, the metadata will be handled separately.
|
||||
return torch.tensor(
|
||||
[], device=reference_tensor.device, dtype=reference_tensor.dtype
|
||||
).set_(
|
||||
storage,
|
||||
0,
|
||||
torch.Size([total_size]),
|
||||
(),
|
||||
)
|
||||
|
||||
|
||||
def group_weights(all_weights: dict[str, Weights]) -> list[OrderedSet[tuple[str, str]]]:
|
||||
"""
|
||||
Group weights that share the same underlying storage.
|
||||
|
||||
Returns a list of sets, each set contains a tuple of (model_name, weight_name).
|
||||
"""
|
||||
|
||||
weights_dict: dict[tuple[int, torch.dtype], OrderedSet[tuple[str, str]]] = (
|
||||
collections.defaultdict(OrderedSet)
|
||||
) # (storage_key, dtype) -> set(weight)
|
||||
|
||||
for model_name, weights in all_weights.items():
|
||||
for weight_name, (tensor, properties) in weights.items():
|
||||
weights_dict[(properties.storage_ptr, tensor.dtype)].add(
|
||||
(model_name, weight_name)
|
||||
)
|
||||
|
||||
return list(weights_dict.values())
|
||||
@@ -0,0 +1,35 @@
|
||||
# Defined in torch/csrc/export/pt2_archive_constants.h
|
||||
from torch._C._export import pt2_archive_constants
|
||||
|
||||
|
||||
AOTINDUCTOR_DIR: str = pt2_archive_constants.AOTINDUCTOR_DIR
|
||||
ARCHIVE_FORMAT_PATH: str = pt2_archive_constants.ARCHIVE_FORMAT_PATH
|
||||
ARCHIVE_FORMAT_VALUE: str = pt2_archive_constants.ARCHIVE_FORMAT_VALUE
|
||||
ARCHIVE_ROOT_NAME: str = pt2_archive_constants.ARCHIVE_ROOT_NAME
|
||||
ARCHIVE_VERSION_PATH: str = pt2_archive_constants.ARCHIVE_VERSION_PATH
|
||||
ARCHIVE_VERSION_VALUE: str = pt2_archive_constants.ARCHIVE_VERSION_VALUE
|
||||
CONSTANTS_DIR: str = pt2_archive_constants.CONSTANTS_DIR
|
||||
CONSTANTS_CONFIG_FILENAME_FORMAT: str = (
|
||||
pt2_archive_constants.CONSTANTS_CONFIG_FILENAME_FORMAT
|
||||
)
|
||||
CUSTOM_OBJ_FILENAME_PREFIX: str = pt2_archive_constants.CUSTOM_OBJ_FILENAME_PREFIX
|
||||
EXECUTORCH_DIR: str = pt2_archive_constants.EXECUTORCH_DIR
|
||||
EXTRA_DIR: str = pt2_archive_constants.EXTRA_DIR
|
||||
MODELS_DIR: str = pt2_archive_constants.MODELS_DIR
|
||||
MODELS_FILENAME_FORMAT: str = pt2_archive_constants.MODELS_FILENAME_FORMAT
|
||||
MODULE_INFO_PATH: str = pt2_archive_constants.MODULE_INFO_PATH
|
||||
MTIA_DIR: str = pt2_archive_constants.MTIA_DIR
|
||||
SAMPLE_INPUTS_DIR: str = pt2_archive_constants.SAMPLE_INPUTS_DIR
|
||||
SAMPLE_INPUTS_FILENAME_FORMAT: str = pt2_archive_constants.SAMPLE_INPUTS_FILENAME_FORMAT
|
||||
TENSOR_CONSTANT_FILENAME_PREFIX: str = (
|
||||
pt2_archive_constants.TENSOR_CONSTANT_FILENAME_PREFIX
|
||||
)
|
||||
WEIGHTS_CONFIG_FILENAME_FORMAT: str = (
|
||||
pt2_archive_constants.WEIGHTS_CONFIG_FILENAME_FORMAT
|
||||
)
|
||||
WEIGHT_FILENAME_PREFIX: str = pt2_archive_constants.WEIGHT_FILENAME_PREFIX
|
||||
WEIGHTS_DIR: str = pt2_archive_constants.WEIGHTS_DIR
|
||||
XL_MODEL_WEIGHTS_DIR: str = pt2_archive_constants.XL_MODEL_WEIGHTS_DIR
|
||||
XL_MODEL_WEIGHTS_PARAM_CONFIG_PATH: str = (
|
||||
pt2_archive_constants.XL_MODEL_WEIGHTS_PARAM_CONFIG_PATH
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user