1061 lines
39 KiB
Python
1061 lines
39 KiB
Python
|
import dataclasses
|
||
|
|
import functools
|
||
|
|
import inspect
|
||
|
|
import logging
|
||
|
|
import re
|
||
|
|
import time
|
||
|
|
import warnings
|
||
|
|
from contextlib import contextmanager, nullcontext
|
||
|
|
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch._dynamo
|
||
|
|
import torch.fx
|
||
|
|
|
||
|
|
import torch.utils._pytree as pytree
|
||
|
|
from torch._dynamo.exc import UserError, UserErrorType
|
||
|
|
from torch._export.non_strict_utils import (
|
||
|
|
make_constraints,
|
||
|
|
make_fake_inputs,
|
||
|
|
make_fake_params_buffers,
|
||
|
|
)
|
||
|
|
from torch._export.passes.add_runtime_assertions_for_constraints_pass import (
|
||
|
|
_AddRuntimeAssertionsForInlineConstraintsPass,
|
||
|
|
)
|
||
|
|
from torch._export.passes.collect_tracepoints_pass import CollectTracepointsPass
|
||
|
|
from torch._export.passes.lift_constants_pass import (
|
||
|
|
ConstantAttrMap,
|
||
|
|
lift_constants_pass,
|
||
|
|
rewrite_script_object_meta,
|
||
|
|
)
|
||
|
|
from torch._export.wrappers import _wrap_submodules
|
||
|
|
from torch._functorch.aot_autograd import aot_export_module
|
||
|
|
from torch._guards import detect_fake_mode
|
||
|
|
from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
|
||
|
|
from torch._utils_internal import log_export_usage
|
||
|
|
from torch.export.exported_program import OutputKind
|
||
|
|
from torch.fx.experimental.symbolic_shapes import (
|
||
|
|
ConstraintViolationError,
|
||
|
|
free_unbacked_symbols,
|
||
|
|
GuardOnDataDependentSymNode,
|
||
|
|
ShapeEnv,
|
||
|
|
)
|
||
|
|
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
|
||
|
|
from torch.utils._sympy.value_ranges import ValueRangeError
|
||
|
|
|
||
|
|
from ._safeguard import AutogradStateOpsFailSafeguard
|
||
|
|
|
||
|
|
from .dynamic_shapes import _process_constraints, Constraint
|
||
|
|
from .exported_program import (
|
||
|
|
_disable_prexisiting_fake_mode,
|
||
|
|
ExportedProgram,
|
||
|
|
InputKind,
|
||
|
|
ModuleCallEntry,
|
||
|
|
ModuleCallSignature,
|
||
|
|
)
|
||
|
|
from .graph_signature import (
|
||
|
|
_sig_to_specs,
|
||
|
|
ArgumentSpec,
|
||
|
|
ConstantArgument,
|
||
|
|
CustomObjArgument,
|
||
|
|
ExportGraphSignature,
|
||
|
|
SymIntArgument,
|
||
|
|
TensorArgument,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
log = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclasses.dataclass
|
||
|
|
class ExportDynamoConfig:
|
||
|
|
"""
|
||
|
|
Manage Export-specific configurations of Dynamo.
|
||
|
|
"""
|
||
|
|
|
||
|
|
allow_rnn: bool = True
|
||
|
|
reorderable_logging_functions: Set[Callable] = dataclasses.field(
|
||
|
|
default_factory=set
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
DEFAULT_EXPORT_DYNAMO_CONFIG = ExportDynamoConfig()
|
||
|
|
DEFAULT_EXPORT_DYNAMO_CONFIG.reorderable_logging_functions = {
|
||
|
|
logging.critical,
|
||
|
|
logging.debug,
|
||
|
|
logging.error,
|
||
|
|
logging.exception,
|
||
|
|
logging.info,
|
||
|
|
logging.log,
|
||
|
|
logging.warning,
|
||
|
|
print,
|
||
|
|
warnings.warn,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def _ignore_backend_decomps():
|
||
|
|
orig_mkldnn_flag = torch.backends.mkldnn.set_flags(False)
|
||
|
|
orig_nnpack_flag = torch.backends.nnpack.set_flags(False)
|
||
|
|
try:
|
||
|
|
yield
|
||
|
|
finally:
|
||
|
|
torch.backends.mkldnn.set_flags(*orig_mkldnn_flag)
|
||
|
|
torch.backends.nnpack.set_flags(*orig_nnpack_flag)
|
||
|
|
|
||
|
|
|
||
|
|
def _convert_input_to_fake(gm, args, kwargs):
|
||
|
|
params_buffers = _get_params_buffers(gm)
|
||
|
|
fake_inps: List[torch.Tensor] = []
|
||
|
|
for node in gm.graph.nodes:
|
||
|
|
if node.op == "placeholder" and "val" in node.meta:
|
||
|
|
fake_val = node.meta["val"]
|
||
|
|
if fake_val is not None and isinstance(fake_val, torch.Tensor):
|
||
|
|
fake_inps.append(fake_val)
|
||
|
|
|
||
|
|
if detected_fake_mode := detect_fake_mode(fake_inps):
|
||
|
|
fake_mode = detected_fake_mode
|
||
|
|
else:
|
||
|
|
fake_mode = FakeTensorMode(shape_env=ShapeEnv())
|
||
|
|
|
||
|
|
if len(args) == 0 and len(kwargs) == 0:
|
||
|
|
return (), {}, params_buffers, fake_mode
|
||
|
|
|
||
|
|
count = 0
|
||
|
|
|
||
|
|
def convert_to_fake(x):
|
||
|
|
nonlocal count
|
||
|
|
val = fake_inps[count]
|
||
|
|
count += 1
|
||
|
|
return val
|
||
|
|
|
||
|
|
fake_args = pytree.tree_map_only(torch.Tensor, convert_to_fake, args)
|
||
|
|
# TODO properly use the cached fake tensor
|
||
|
|
fake_kwargs = pytree.tree_map_only(torch.Tensor, fake_mode.from_tensor, kwargs)
|
||
|
|
fake_params_buffers = pytree.tree_map_only(
|
||
|
|
torch.Tensor,
|
||
|
|
functools.partial(fake_mode.from_tensor, static_shapes=True),
|
||
|
|
params_buffers,
|
||
|
|
)
|
||
|
|
return fake_args, fake_kwargs, fake_params_buffers, fake_mode
|
||
|
|
|
||
|
|
|
||
|
|
def _replace_param_buffer_names(param_buffer_table, sig):
|
||
|
|
for spec in sig.input_specs:
|
||
|
|
if spec.kind in (
|
||
|
|
InputKind.PARAMETER,
|
||
|
|
InputKind.BUFFER,
|
||
|
|
):
|
||
|
|
spec.target = param_buffer_table[spec.target]
|
||
|
|
for spec in sig.output_specs:
|
||
|
|
if spec.kind in (
|
||
|
|
OutputKind.BUFFER_MUTATION,
|
||
|
|
OutputKind.GRADIENT_TO_PARAMETER,
|
||
|
|
):
|
||
|
|
spec.target = param_buffer_table[spec.target]
|
||
|
|
|
||
|
|
|
||
|
|
def _convert_to_positional_args(orig_arg_names, args, kwargs):
|
||
|
|
assert len(orig_arg_names) == len(args) + len(kwargs), (
|
||
|
|
f"Total number of arg names is expected to be {len(orig_arg_names)} "
|
||
|
|
f"but got {len(args)} positional args, {len(kwargs)} kwargs."
|
||
|
|
)
|
||
|
|
reordered_kwargs = [kwargs[kw_name] for kw_name in orig_arg_names[len(args) :]]
|
||
|
|
return (
|
||
|
|
*args,
|
||
|
|
*reordered_kwargs,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _normalize_nn_module_stack(gm_torch_level, root_cls):
|
||
|
|
# Append a root module to every nn_module_stack.
|
||
|
|
root = "L['self']"
|
||
|
|
root_key = re.sub(r"[^a-zA-Z0-9]", "_", root)
|
||
|
|
for gm in gm_torch_level.modules():
|
||
|
|
if not isinstance(gm, torch.fx.GraphModule):
|
||
|
|
continue
|
||
|
|
for node in gm.graph.nodes:
|
||
|
|
if node.op in ["placeholder", "output"]:
|
||
|
|
continue
|
||
|
|
add_root = True
|
||
|
|
if nn_module_stack := node.meta.get("nn_module_stack", {}):
|
||
|
|
path, ty = next(iter(nn_module_stack.values()))
|
||
|
|
# After deserializing the class `ty` might not exist anymore so
|
||
|
|
# it could be a string
|
||
|
|
if inspect.isclass(ty) and issubclass(ty, torch.nn.Module):
|
||
|
|
# TODO Figure out why sometimes we have root sometimes we don't.
|
||
|
|
if path == root and ty is root_cls:
|
||
|
|
add_root = False
|
||
|
|
else:
|
||
|
|
assert isinstance(ty, str)
|
||
|
|
if add_root:
|
||
|
|
|
||
|
|
def normalize_path(path):
|
||
|
|
try:
|
||
|
|
parts = []
|
||
|
|
|
||
|
|
class Path:
|
||
|
|
def __getattr__(self, name):
|
||
|
|
parts.append(name)
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __getitem__(self, idx):
|
||
|
|
parts.append(str(idx))
|
||
|
|
return self
|
||
|
|
|
||
|
|
eval(path, {"L": {"self": Path()}})
|
||
|
|
return ".".join(parts)
|
||
|
|
except Exception: # TODO(zhxchen17) Remove this.
|
||
|
|
return path
|
||
|
|
|
||
|
|
nn_module_stack = {root_key: (root, root_cls), **nn_module_stack}
|
||
|
|
node.meta["nn_module_stack"] = {
|
||
|
|
key: (normalize_path(path), ty)
|
||
|
|
for key, (path, ty) in nn_module_stack.items()
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _get_param_buffer_mapping(
|
||
|
|
original_module: torch.nn.Module,
|
||
|
|
traced_module: torch.nn.Module,
|
||
|
|
) -> Dict[str, str]:
|
||
|
|
"""
|
||
|
|
Returns a mapping of parameter/buffer names from the new module to the
|
||
|
|
original model. This is to help with restoring the FQN for parameter/buffers
|
||
|
|
of a traced module to what the original module contains.
|
||
|
|
"""
|
||
|
|
|
||
|
|
param_lookup: Dict[int, List[str]] = {}
|
||
|
|
buffer_lookup: Dict[int, List[str]] = {}
|
||
|
|
for name, param in original_module.named_parameters(remove_duplicate=False):
|
||
|
|
param_lookup.setdefault(id(param), []).append(name)
|
||
|
|
for name, buffer in original_module.named_buffers(remove_duplicate=False):
|
||
|
|
buffer_lookup.setdefault(id(buffer), []).append(name)
|
||
|
|
|
||
|
|
param_buffer_table: Dict[str, str] = {}
|
||
|
|
for dynamo_name, dynamo_param in traced_module.named_parameters(
|
||
|
|
remove_duplicate=False
|
||
|
|
):
|
||
|
|
assert dynamo_name not in param_buffer_table
|
||
|
|
if id(dynamo_param) in param_lookup:
|
||
|
|
param_buffer_table[dynamo_name] = param_lookup[id(dynamo_param)].pop()
|
||
|
|
|
||
|
|
for dynamo_name, dynamo_buffer in traced_module.named_buffers(
|
||
|
|
remove_duplicate=False
|
||
|
|
):
|
||
|
|
assert dynamo_name not in param_buffer_table
|
||
|
|
if id(dynamo_buffer) in buffer_lookup:
|
||
|
|
param_buffer_table[dynamo_name] = buffer_lookup[id(dynamo_buffer)].pop()
|
||
|
|
|
||
|
|
return param_buffer_table
|
||
|
|
|
||
|
|
|
||
|
|
def _remap_constants(
|
||
|
|
orig_constant_attrs: ConstantAttrMap,
|
||
|
|
graph_signature: ExportGraphSignature,
|
||
|
|
constants: Dict[str, Union[torch.Tensor, torch.ScriptObject]],
|
||
|
|
) -> None:
|
||
|
|
"""Rewrite the graph signature and constants table to use the FQN from the original module."""
|
||
|
|
remap_table: Dict[str, str] = {}
|
||
|
|
for name, value in constants.items():
|
||
|
|
if value in orig_constant_attrs:
|
||
|
|
remap_table[name] = orig_constant_attrs[value]
|
||
|
|
|
||
|
|
for spec in graph_signature.input_specs:
|
||
|
|
if spec.kind in (
|
||
|
|
InputKind.CONSTANT_TENSOR,
|
||
|
|
InputKind.CUSTOM_OBJ,
|
||
|
|
):
|
||
|
|
orig_target = spec.target
|
||
|
|
assert orig_target is not None
|
||
|
|
spec.target = remap_table.get(orig_target, orig_target)
|
||
|
|
|
||
|
|
constant = constants[orig_target]
|
||
|
|
del constants[orig_target]
|
||
|
|
constants[spec.target] = constant
|
||
|
|
|
||
|
|
|
||
|
|
def _restore_state_dict(
|
||
|
|
original_module: torch.nn.Module, traced_module: torch.fx.GraphModule
|
||
|
|
) -> None:
|
||
|
|
"""
|
||
|
|
Restores the state dict of the traced module to that of the original module.
|
||
|
|
"""
|
||
|
|
param_buffer_table = _get_param_buffer_mapping(original_module, traced_module)
|
||
|
|
# Since the graph module is flattened (no module heirarchy), we
|
||
|
|
# need to noramlize the module by replacing "." with "_". If we
|
||
|
|
# don't, it will try to save the weight to a submodule which no
|
||
|
|
# longer exists.
|
||
|
|
for name, fqn in param_buffer_table.items():
|
||
|
|
param_buffer_table[name] = fqn.replace(".", "_")
|
||
|
|
|
||
|
|
# Replace state dict attr names with the fqn
|
||
|
|
for name, fqn in param_buffer_table.items():
|
||
|
|
if not hasattr(traced_module, name):
|
||
|
|
continue
|
||
|
|
|
||
|
|
attr = getattr(traced_module, name)
|
||
|
|
if isinstance(attr, torch.Tensor) and not isinstance(attr, torch.nn.Parameter):
|
||
|
|
traced_module.register_buffer(fqn, attr)
|
||
|
|
else:
|
||
|
|
setattr(traced_module, fqn, attr)
|
||
|
|
delattr(traced_module, name)
|
||
|
|
|
||
|
|
# Replace graph getattr nodes with the correct name
|
||
|
|
for node in traced_module.graph.nodes:
|
||
|
|
if node.op == "get_attr":
|
||
|
|
attr_name = node.target
|
||
|
|
if attr_name in param_buffer_table:
|
||
|
|
node.target = param_buffer_table[attr_name]
|
||
|
|
|
||
|
|
traced_module.recompile()
|
||
|
|
|
||
|
|
|
||
|
|
def _export_to_torch_ir(
|
||
|
|
f: Callable,
|
||
|
|
args: Tuple[Any, ...],
|
||
|
|
kwargs: Optional[Dict[str, Any]] = None,
|
||
|
|
constraints: Optional[List[Constraint]] = None,
|
||
|
|
*,
|
||
|
|
preserve_module_call_signature: Tuple[str, ...] = (),
|
||
|
|
disable_constraint_solver: bool = False,
|
||
|
|
restore_fqn: bool = True,
|
||
|
|
_log_export_usage: bool = True,
|
||
|
|
) -> torch.fx.GraphModule:
|
||
|
|
"""
|
||
|
|
Traces either an nn.Module's forward function or just a callable with PyTorch
|
||
|
|
operations inside and produce a torch.fx.GraphModule in torch IR.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if _log_export_usage:
|
||
|
|
log_export_usage(event="export.private_api", flags={"_export_to_torch_ir"})
|
||
|
|
|
||
|
|
kwargs = kwargs or {}
|
||
|
|
|
||
|
|
if not isinstance(args, tuple):
|
||
|
|
raise UserError(
|
||
|
|
UserErrorType.INVALID_INPUT,
|
||
|
|
f"Expecting `args` to be a tuple of example positional inputs, got {type(args)}",
|
||
|
|
)
|
||
|
|
|
||
|
|
with torch._dynamo.config.patch(dataclasses.asdict(DEFAULT_EXPORT_DYNAMO_CONFIG)):
|
||
|
|
try:
|
||
|
|
module_call_specs: Dict[str, Dict[str, pytree.TreeSpec]] = {}
|
||
|
|
with _wrap_submodules(
|
||
|
|
f, preserve_module_call_signature, module_call_specs
|
||
|
|
), _ignore_backend_decomps():
|
||
|
|
gm_torch_level, _ = torch._dynamo.export(
|
||
|
|
f,
|
||
|
|
constraints=constraints, # type: ignore[arg-type]
|
||
|
|
assume_static_by_default=True,
|
||
|
|
tracing_mode="symbolic",
|
||
|
|
disable_constraint_solver=disable_constraint_solver,
|
||
|
|
_log_export_usage=_log_export_usage,
|
||
|
|
)(
|
||
|
|
*args,
|
||
|
|
**kwargs,
|
||
|
|
)
|
||
|
|
except (ConstraintViolationError, ValueRangeError) as e:
|
||
|
|
raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e)) # noqa: TRY200
|
||
|
|
except GuardOnDataDependentSymNode as e:
|
||
|
|
raise UserError( # noqa: TRY200
|
||
|
|
UserErrorType.ANTI_PATTERN,
|
||
|
|
f"Consider annotating your code using torch._constrain_as_*(). {str(e)}",
|
||
|
|
case_name="constrain_as_size_example",
|
||
|
|
)
|
||
|
|
|
||
|
|
gm_torch_level.meta["module_call_specs"] = module_call_specs
|
||
|
|
|
||
|
|
if isinstance(f, torch.nn.Module) and restore_fqn:
|
||
|
|
_restore_state_dict(f, gm_torch_level)
|
||
|
|
|
||
|
|
return gm_torch_level
|
||
|
|
|
||
|
|
|
||
|
|
def _gather_constant_attrs(m: torch.nn.Module) -> ConstantAttrMap:
|
||
|
|
"""Search the module hierarchy, gathering up all tensor and ScriptObject constants.
|
||
|
|
|
||
|
|
Returns a dictionary mapping hash(value) to the name of the constant. We
|
||
|
|
have to abuse `hash` here unfortunately, see: [ScriptObject hash].
|
||
|
|
"""
|
||
|
|
constants = ConstantAttrMap()
|
||
|
|
buffers_parameters = set(m.buffers())
|
||
|
|
buffers_parameters.update(m.parameters())
|
||
|
|
|
||
|
|
def inner(m: torch.nn.Module, prefix_atoms: List[str], constants):
|
||
|
|
for k, v in m.__dict__.items():
|
||
|
|
if isinstance(v, (torch.Tensor, torch.ScriptObject)):
|
||
|
|
if v in buffers_parameters:
|
||
|
|
# filter out buffers and parameters, leaving only constants
|
||
|
|
continue
|
||
|
|
|
||
|
|
fqn = ".".join(prefix_atoms + [k])
|
||
|
|
if v in constants:
|
||
|
|
raise ValueError(
|
||
|
|
f"Duplicate reference to constant attribute found: '{constants[v]}' and '{fqn}'."
|
||
|
|
)
|
||
|
|
|
||
|
|
constants[v] = fqn
|
||
|
|
for k, v in m.named_children():
|
||
|
|
inner(v, prefix_atoms + [k], constants)
|
||
|
|
|
||
|
|
inner(m, [], constants)
|
||
|
|
return constants
|
||
|
|
|
||
|
|
|
||
|
|
def _export_non_strict(
|
||
|
|
mod: torch.nn.Module,
|
||
|
|
fake_args,
|
||
|
|
fake_kwargs,
|
||
|
|
fake_params_buffers,
|
||
|
|
constant_attrs: ConstantAttrMap,
|
||
|
|
*,
|
||
|
|
transform=lambda x: x, # TODO(zhxchen17) Revisit if this is needed later.
|
||
|
|
pre_dispatch=False,
|
||
|
|
):
|
||
|
|
# [NOTE] If the user is exporting under training mode, we want to detect if there is any
|
||
|
|
# state change in the autograd global state and error. If the user is exporting under inference
|
||
|
|
# mode, we don't care.
|
||
|
|
is_grad_enabled = torch._C.is_grad_enabled()
|
||
|
|
grad_safe_guard = (
|
||
|
|
AutogradStateOpsFailSafeguard() if is_grad_enabled else nullcontext()
|
||
|
|
)
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def _compiling_state_context():
|
||
|
|
old_value = torch.compiler._is_compiling_flag
|
||
|
|
try:
|
||
|
|
torch.compiler._is_compiling_flag = True
|
||
|
|
yield
|
||
|
|
finally:
|
||
|
|
torch.compiler._is_compiling_flag = old_value
|
||
|
|
|
||
|
|
# This _reparametrize_module makes sure inputs and module.params/buffers have the same fake_mode,
|
||
|
|
# otherwise aot_export_module will error out because it sees a mix of fake_modes.
|
||
|
|
# And we want aot_export_module to use the fake_tensor mode in dynamo to keep the pipeline easy to reason about.
|
||
|
|
with torch.nn.utils.stateless._reparametrize_module(
|
||
|
|
mod, fake_params_buffers
|
||
|
|
), grad_safe_guard, _ignore_backend_decomps(), _compiling_state_context(): # type: ignore[attr-defined]
|
||
|
|
gm, graph_signature = transform(aot_export_module)(
|
||
|
|
mod,
|
||
|
|
fake_args,
|
||
|
|
trace_joint=False,
|
||
|
|
pre_dispatch=pre_dispatch,
|
||
|
|
kwargs=fake_kwargs,
|
||
|
|
)
|
||
|
|
# TODO unfortunately preserving graph-level metadata is not
|
||
|
|
# working well with aot_export. So we manually copy it.
|
||
|
|
# (The node-level meta is addressed above.)
|
||
|
|
if isinstance(mod, torch.fx.GraphModule) and hasattr(mod, "meta"):
|
||
|
|
gm.meta.update(mod.meta)
|
||
|
|
|
||
|
|
if pre_dispatch:
|
||
|
|
from torch._export.passes.replace_set_grad_with_hop_pass import (
|
||
|
|
replace_set_grad_with_hop_pass,
|
||
|
|
)
|
||
|
|
|
||
|
|
gm = replace_set_grad_with_hop_pass(gm)
|
||
|
|
|
||
|
|
# NOTE: aot_export adds symint metadata for placeholders with int values;
|
||
|
|
# since these become specialized, we replace such metadata with the original values
|
||
|
|
flat_args = pytree.tree_leaves((fake_args, fake_kwargs))
|
||
|
|
index = 0
|
||
|
|
total_non_user_inputs = (
|
||
|
|
len(graph_signature.parameters)
|
||
|
|
+ len(graph_signature.buffers)
|
||
|
|
+ len(graph_signature.input_tokens)
|
||
|
|
)
|
||
|
|
for node in gm.graph.nodes:
|
||
|
|
if node.op == "placeholder":
|
||
|
|
if index >= total_non_user_inputs:
|
||
|
|
user_arg = flat_args[index - total_non_user_inputs]
|
||
|
|
if not isinstance(user_arg, torch.Tensor):
|
||
|
|
node.meta["val"] = user_arg
|
||
|
|
index += 1
|
||
|
|
|
||
|
|
is_joint = graph_signature.backward_signature is not None
|
||
|
|
|
||
|
|
def make_argument_spec(node) -> ArgumentSpec:
|
||
|
|
if isinstance(node, (int, bool, float, type(None))):
|
||
|
|
# For const outputs we just directly return this
|
||
|
|
return ConstantArgument(value=node)
|
||
|
|
|
||
|
|
assert (
|
||
|
|
"val" in node.meta
|
||
|
|
), f"{node} is not a constant or a node with a 'val' metadata field"
|
||
|
|
val = node.meta["val"]
|
||
|
|
if isinstance(val, FakeTensor):
|
||
|
|
return TensorArgument(name=node.name)
|
||
|
|
elif isinstance(val, torch.SymInt):
|
||
|
|
return SymIntArgument(name=node.name)
|
||
|
|
elif isinstance(val, torch.ScriptObject):
|
||
|
|
return CustomObjArgument(
|
||
|
|
name=node.name, class_fqn=val._type().qualified_name() # type: ignore[attr-defined]
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
# TODO: this branch is likely wrong, all permissible ConstantArgument type
|
||
|
|
# should have been handled already
|
||
|
|
return ConstantArgument(value=val)
|
||
|
|
|
||
|
|
input_specs, output_specs = _sig_to_specs(
|
||
|
|
user_inputs=set(graph_signature.user_inputs),
|
||
|
|
inputs_to_parameters=graph_signature.inputs_to_parameters, # type: ignore[arg-type]
|
||
|
|
inputs_to_buffers=graph_signature.inputs_to_buffers, # type: ignore[arg-type]
|
||
|
|
user_outputs=set(graph_signature.user_outputs), # type: ignore[arg-type]
|
||
|
|
buffer_mutations=graph_signature.buffers_to_mutate, # type: ignore[arg-type]
|
||
|
|
user_input_mutations=graph_signature.user_inputs_to_mutate, # type: ignore[arg-type]
|
||
|
|
grad_params=graph_signature.backward_signature.gradients_to_parameters if is_joint else {}, # type: ignore[arg-type, union-attr]
|
||
|
|
grad_user_inputs=graph_signature.backward_signature.gradients_to_user_inputs if is_joint else {}, # type: ignore[arg-type, union-attr]
|
||
|
|
loss_output=graph_signature.backward_signature.loss_output if is_joint else None, # type: ignore[arg-type, union-attr]
|
||
|
|
inputs=[
|
||
|
|
make_argument_spec(node)
|
||
|
|
for node in gm.graph.nodes
|
||
|
|
if node.op == "placeholder"
|
||
|
|
],
|
||
|
|
outputs=[
|
||
|
|
make_argument_spec(node)
|
||
|
|
for node in pytree.tree_leaves(next(iter(reversed(gm.graph.nodes))).args)
|
||
|
|
],
|
||
|
|
input_tokens=graph_signature.input_tokens,
|
||
|
|
output_tokens=graph_signature.output_tokens,
|
||
|
|
)
|
||
|
|
export_graph_signature = ExportGraphSignature(
|
||
|
|
input_specs=input_specs, output_specs=output_specs
|
||
|
|
)
|
||
|
|
|
||
|
|
constants = rewrite_script_object_meta(gm)
|
||
|
|
constants.update(lift_constants_pass(gm, export_graph_signature, constant_attrs))
|
||
|
|
|
||
|
|
@dataclasses.dataclass
|
||
|
|
class _ExportedProgramNonStrict:
|
||
|
|
gm: torch.fx.GraphModule
|
||
|
|
sig: ExportGraphSignature
|
||
|
|
constants: Dict[str, Union[torch.Tensor, torch._C.ScriptObject]]
|
||
|
|
|
||
|
|
return _ExportedProgramNonStrict(
|
||
|
|
gm,
|
||
|
|
export_graph_signature,
|
||
|
|
constants,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _get_params_buffers(mod: torch.nn.Module) -> Dict[str, torch.Tensor]:
|
||
|
|
params_buffers: Dict[str, torch.Tensor] = {}
|
||
|
|
for name, param in mod.named_parameters(remove_duplicate=False):
|
||
|
|
params_buffers[name] = param
|
||
|
|
|
||
|
|
for name, buffer in mod.named_buffers(remove_duplicate=False):
|
||
|
|
params_buffers[name] = buffer
|
||
|
|
return params_buffers
|
||
|
|
|
||
|
|
|
||
|
|
def _rewrite_dynamo_tensor_constants(
|
||
|
|
orig_mod_buffers: Set[torch.Tensor],
|
||
|
|
traced_mod_buffers: Dict[str, torch.Tensor],
|
||
|
|
graph_signature: ExportGraphSignature,
|
||
|
|
constants: Dict[str, Union[torch.Tensor, torch.ScriptObject]],
|
||
|
|
):
|
||
|
|
"""Dynamo erroneously marks tensor attributes on modules as a buffers.
|
||
|
|
|
||
|
|
Rewrite them to be tensor constants.
|
||
|
|
"""
|
||
|
|
for spec in graph_signature.input_specs:
|
||
|
|
if spec.kind == InputKind.BUFFER:
|
||
|
|
assert spec.target is not None
|
||
|
|
value = traced_mod_buffers[spec.target]
|
||
|
|
if value not in orig_mod_buffers:
|
||
|
|
# This was a tensor constant erroneously marked as a buffer.
|
||
|
|
# Convert it int oa constant in the graph signature, and add its
|
||
|
|
# value to the constants table.
|
||
|
|
spec.kind = InputKind.CONSTANT_TENSOR
|
||
|
|
constants[spec.target] = value
|
||
|
|
|
||
|
|
|
||
|
|
def _rewrite_non_persistent_buffers(
|
||
|
|
orig_mod: torch.nn.Module,
|
||
|
|
graph_signature: ExportGraphSignature,
|
||
|
|
constants: Dict[str, Union[torch.Tensor, torch.ScriptObject]],
|
||
|
|
):
|
||
|
|
"""Dynamo erroneously drops the persistent flag on buffers.
|
||
|
|
|
||
|
|
Rewrite non-persistent buffers to reflect the original module.
|
||
|
|
"""
|
||
|
|
state_dict = orig_mod.state_dict()
|
||
|
|
for spec in graph_signature.input_specs:
|
||
|
|
if spec.kind == InputKind.BUFFER:
|
||
|
|
assert spec.target is not None
|
||
|
|
if spec.target not in state_dict:
|
||
|
|
assert spec.target not in constants
|
||
|
|
spec.persistent = False
|
||
|
|
constants[spec.target] = orig_mod.get_buffer(spec.target)
|
||
|
|
|
||
|
|
|
||
|
|
def get_ep_stats(ep: ExportedProgram) -> Dict[str, Any]:
|
||
|
|
op_count = 0
|
||
|
|
op_set = set()
|
||
|
|
for m in ep.graph_module.modules():
|
||
|
|
if not isinstance(m, torch.fx.GraphModule):
|
||
|
|
continue
|
||
|
|
for node in m.graph.nodes:
|
||
|
|
if node.op != "call_function":
|
||
|
|
continue
|
||
|
|
op_count += 1
|
||
|
|
assert hasattr(node.target, "__module__")
|
||
|
|
assert hasattr(node.target, "__name__")
|
||
|
|
op_set.add(f"{node.target.__module__}.{node.target.__name__}")
|
||
|
|
return {"op_count": op_count, "op_set": op_set}
|
||
|
|
|
||
|
|
|
||
|
|
_EXPORT_FLAGS: Optional[Set[str]] = None
|
||
|
|
|
||
|
|
|
||
|
|
def _log_export_wrapper(fn):
|
||
|
|
@functools.wraps(fn)
|
||
|
|
def wrapper(*args, **kwargs):
|
||
|
|
global _EXPORT_FLAGS
|
||
|
|
try:
|
||
|
|
start = time.time()
|
||
|
|
ep = fn(*args, **kwargs)
|
||
|
|
end = time.time()
|
||
|
|
log_export_usage(
|
||
|
|
event="export.time",
|
||
|
|
metrics=end - start,
|
||
|
|
flags=_EXPORT_FLAGS,
|
||
|
|
**get_ep_stats(ep),
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
t = type(e)
|
||
|
|
error_type = t.__module__ + "." + t.__qualname__
|
||
|
|
log_export_usage(
|
||
|
|
event="export.error",
|
||
|
|
type=error_type,
|
||
|
|
message=str(e),
|
||
|
|
flags=_EXPORT_FLAGS,
|
||
|
|
)
|
||
|
|
raise e
|
||
|
|
finally:
|
||
|
|
_EXPORT_FLAGS = None
|
||
|
|
|
||
|
|
return ep
|
||
|
|
|
||
|
|
return wrapper
|
||
|
|
|
||
|
|
|
||
|
|
@_log_export_wrapper
|
||
|
|
@_disable_prexisiting_fake_mode
|
||
|
|
def _export(
|
||
|
|
mod: torch.nn.Module,
|
||
|
|
args: Tuple[Any, ...],
|
||
|
|
kwargs: Optional[Dict[str, Any]] = None,
|
||
|
|
dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any], List[Any]]] = None,
|
||
|
|
*,
|
||
|
|
strict: bool = True,
|
||
|
|
preserve_module_call_signature: Tuple[str, ...] = (),
|
||
|
|
pre_dispatch: bool = False,
|
||
|
|
) -> ExportedProgram:
|
||
|
|
"""
|
||
|
|
Traces either an nn.Module's forward function or just a callable with PyTorch
|
||
|
|
operations inside and produce a ExportedProgram.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
f: the `nn.Module` to trace.
|
||
|
|
|
||
|
|
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.
|
||
|
|
|
||
|
|
preserve_module_call_signature: A list of submodule paths for which the original
|
||
|
|
calling conventions are preserved as metadata.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
An ExportedProgram containing the traced method.
|
||
|
|
"""
|
||
|
|
from .dynamic_shapes import _process_dynamic_shapes
|
||
|
|
|
||
|
|
global _EXPORT_FLAGS
|
||
|
|
flags = set()
|
||
|
|
flags.add("strict" if strict else "non_strict")
|
||
|
|
flags.add("pre_dispatch" if pre_dispatch else "aot_dispatch")
|
||
|
|
log_export_usage(event="export.enter", flags=flags)
|
||
|
|
_EXPORT_FLAGS = flags
|
||
|
|
|
||
|
|
constraints = _process_dynamic_shapes(mod, args, kwargs, dynamic_shapes) or []
|
||
|
|
|
||
|
|
kwargs = kwargs or {}
|
||
|
|
|
||
|
|
constant_attrs = _gather_constant_attrs(mod)
|
||
|
|
|
||
|
|
flat_args, orig_in_spec = pytree.tree_flatten((args, kwargs))
|
||
|
|
|
||
|
|
if not strict:
|
||
|
|
out_spec = None
|
||
|
|
|
||
|
|
module_call_specs: Dict[str, Dict[str, pytree.TreeSpec]] = {}
|
||
|
|
|
||
|
|
def strip_root(x):
|
||
|
|
if isinstance(x, str) and x.startswith("_export_root"):
|
||
|
|
stripped = x[len("_export_root") :]
|
||
|
|
return stripped[1:] if stripped.startswith(".") else stripped
|
||
|
|
return x
|
||
|
|
|
||
|
|
def fixup_key(x):
|
||
|
|
return "L__self__" + strip_root(x)
|
||
|
|
|
||
|
|
def _tuplify_outputs(aot_export):
|
||
|
|
def _aot_export_non_strict(mod, args, kwargs=None, **flags):
|
||
|
|
kwargs = kwargs or {}
|
||
|
|
|
||
|
|
class Wrapper(torch.nn.Module):
|
||
|
|
def __init__(self, mod):
|
||
|
|
super().__init__()
|
||
|
|
self._export_root = mod
|
||
|
|
|
||
|
|
def forward(self, *args, **kwargs):
|
||
|
|
nonlocal out_spec
|
||
|
|
if isinstance(self._export_root, torch.fx.GraphModule):
|
||
|
|
with torch.fx.traceback.preserve_node_meta():
|
||
|
|
tree_out = torch.fx.Interpreter(self._export_root).run(
|
||
|
|
*args, **kwargs
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
tree_out = self._export_root(*args, **kwargs)
|
||
|
|
flat_outs, out_spec = pytree.tree_flatten(tree_out)
|
||
|
|
return tuple(flat_outs)
|
||
|
|
|
||
|
|
wrapped_mod = Wrapper(mod)
|
||
|
|
# Patch export_root to the signatures so that wrapper module correctly populates the
|
||
|
|
# in/out spec
|
||
|
|
new_preserved_call_signatures = [
|
||
|
|
"_export_root." + i for i in preserve_module_call_signature
|
||
|
|
]
|
||
|
|
with _wrap_submodules(
|
||
|
|
wrapped_mod, new_preserved_call_signatures, module_call_specs
|
||
|
|
):
|
||
|
|
gm, sig = aot_export(wrapped_mod, args, kwargs=kwargs, **flags)
|
||
|
|
|
||
|
|
sig.parameters = pytree.tree_map(strip_root, sig.parameters)
|
||
|
|
sig.buffers = pytree.tree_map(strip_root, sig.buffers)
|
||
|
|
sig.inputs_to_buffers = pytree.tree_map(
|
||
|
|
strip_root, sig.inputs_to_buffers
|
||
|
|
)
|
||
|
|
sig.inputs_to_parameters = pytree.tree_map(
|
||
|
|
strip_root, sig.inputs_to_parameters
|
||
|
|
)
|
||
|
|
sig.buffers_to_mutate = pytree.tree_map(
|
||
|
|
strip_root, sig.buffers_to_mutate
|
||
|
|
)
|
||
|
|
for node in gm.graph.nodes:
|
||
|
|
if "nn_module_stack" in node.meta:
|
||
|
|
nn_module_stack = node.meta["nn_module_stack"]
|
||
|
|
node.meta["nn_module_stack"] = {
|
||
|
|
fixup_key(key): val
|
||
|
|
for key, val in pytree.tree_map(
|
||
|
|
strip_root, nn_module_stack
|
||
|
|
).items()
|
||
|
|
}
|
||
|
|
|
||
|
|
return gm, sig
|
||
|
|
|
||
|
|
return _aot_export_non_strict
|
||
|
|
|
||
|
|
(
|
||
|
|
fake_mode,
|
||
|
|
fake_args,
|
||
|
|
fake_kwargs,
|
||
|
|
equalities_inputs,
|
||
|
|
original_signature,
|
||
|
|
) = make_fake_inputs(mod, args, kwargs, constraints)
|
||
|
|
|
||
|
|
fake_params_buffers = make_fake_params_buffers(
|
||
|
|
fake_mode, _get_params_buffers(mod)
|
||
|
|
)
|
||
|
|
with fake_mode:
|
||
|
|
ep_non_strict = _export_non_strict(
|
||
|
|
mod,
|
||
|
|
fake_args,
|
||
|
|
fake_kwargs,
|
||
|
|
fake_params_buffers,
|
||
|
|
constant_attrs,
|
||
|
|
pre_dispatch=pre_dispatch,
|
||
|
|
transform=_tuplify_outputs,
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
range_constraints = make_constraints(
|
||
|
|
fake_mode,
|
||
|
|
equalities_inputs,
|
||
|
|
original_signature,
|
||
|
|
ep_non_strict.gm,
|
||
|
|
)
|
||
|
|
except (ConstraintViolationError, ValueRangeError) as e:
|
||
|
|
raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e)) # noqa: TRY200
|
||
|
|
|
||
|
|
assert out_spec is not None
|
||
|
|
|
||
|
|
gm = ep_non_strict.gm
|
||
|
|
|
||
|
|
module_call_signatures = {
|
||
|
|
strip_root(fqn): ModuleCallSignature(inputs=[], outputs=[], **specs)
|
||
|
|
for fqn, specs in module_call_specs.items()
|
||
|
|
}
|
||
|
|
|
||
|
|
if len(preserve_module_call_signature) > 0:
|
||
|
|
for node in gm.graph.nodes:
|
||
|
|
if node.target == torch.ops.higher_order._export_tracepoint:
|
||
|
|
if "path" in node.kwargs:
|
||
|
|
path = strip_root(node.kwargs["path"])
|
||
|
|
with gm.graph.inserting_before(node):
|
||
|
|
new_node = gm.graph.create_node(
|
||
|
|
"call_function",
|
||
|
|
torch.ops.higher_order._export_tracepoint,
|
||
|
|
args=node.args,
|
||
|
|
kwargs={
|
||
|
|
"path": path,
|
||
|
|
"kind": node.kwargs["kind"],
|
||
|
|
},
|
||
|
|
)
|
||
|
|
node.replace_all_uses_with(new_node)
|
||
|
|
gm.graph.erase_node(node)
|
||
|
|
|
||
|
|
res = CollectTracepointsPass(module_call_signatures, ep_non_strict.sig)(gm)
|
||
|
|
assert res is not None
|
||
|
|
gm = res.graph_module
|
||
|
|
|
||
|
|
_rewrite_non_persistent_buffers(mod, ep_non_strict.sig, ep_non_strict.constants)
|
||
|
|
return ExportedProgram(
|
||
|
|
root=gm,
|
||
|
|
graph=gm.graph,
|
||
|
|
graph_signature=ep_non_strict.sig,
|
||
|
|
state_dict=mod.state_dict(keep_vars=True),
|
||
|
|
range_constraints=range_constraints,
|
||
|
|
module_call_graph=[
|
||
|
|
ModuleCallEntry(
|
||
|
|
"",
|
||
|
|
ModuleCallSignature(
|
||
|
|
inputs=[], outputs=[], in_spec=orig_in_spec, out_spec=out_spec
|
||
|
|
),
|
||
|
|
)
|
||
|
|
]
|
||
|
|
+ [
|
||
|
|
ModuleCallEntry(fqn, sig) for fqn, sig in module_call_signatures.items()
|
||
|
|
],
|
||
|
|
example_inputs=(args, kwargs),
|
||
|
|
constants=ep_non_strict.constants,
|
||
|
|
)
|
||
|
|
|
||
|
|
gm_torch_level = _export_to_torch_ir(
|
||
|
|
mod,
|
||
|
|
args,
|
||
|
|
kwargs,
|
||
|
|
constraints,
|
||
|
|
preserve_module_call_signature=preserve_module_call_signature,
|
||
|
|
restore_fqn=False, # don't need to restore because we will do it later
|
||
|
|
_log_export_usage=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
# We detect the fake_mode by looking at gm_torch_level's placeholders, this is the fake_mode created in dynamo.
|
||
|
|
(
|
||
|
|
fake_args,
|
||
|
|
fake_kwargs,
|
||
|
|
fake_params_buffers,
|
||
|
|
dynamo_fake_mode,
|
||
|
|
) = _convert_input_to_fake(gm_torch_level, args, kwargs)
|
||
|
|
|
||
|
|
# First, we want to pass through the graph to try populating
|
||
|
|
# val field for getattr if there is anything missing.
|
||
|
|
# This can happen when quantization adds extra params and forgets
|
||
|
|
# to update "val"
|
||
|
|
for node in gm_torch_level.graph.nodes:
|
||
|
|
if node.op == "get_attr" and "val" not in node.meta:
|
||
|
|
attr = getattr(gm_torch_level, node.target)
|
||
|
|
# Checks if it is not a HigherOrderOp branch or a module
|
||
|
|
if not isinstance(attr, torch.nn.Module):
|
||
|
|
assert (
|
||
|
|
dynamo_fake_mode is not None
|
||
|
|
), "Cannot find dynamo_fake_mode. This could be due to the exported graph module have no placeholders."
|
||
|
|
node.meta["val"] = dynamo_fake_mode.from_tensor(
|
||
|
|
attr, static_shapes=True
|
||
|
|
)
|
||
|
|
|
||
|
|
# When aot_export lifts the params, we lose the nn_module_stack
|
||
|
|
# and source_fn from the param nodes as they are treated as fresh inputs
|
||
|
|
# Therefore, we manually extract them before calling into aot_export
|
||
|
|
params_buffers_to_node_meta = {}
|
||
|
|
for node in gm_torch_level.graph.nodes:
|
||
|
|
target = node.target
|
||
|
|
meta = node.meta
|
||
|
|
if node.op == "call_module":
|
||
|
|
submodule = getattr(gm_torch_level, target)
|
||
|
|
if isinstance(submodule, torch.nn.Module):
|
||
|
|
for name, _ in submodule.named_parameters(
|
||
|
|
recurse=True, remove_duplicate=False
|
||
|
|
):
|
||
|
|
params_buffers_to_node_meta[target + "." + name] = meta
|
||
|
|
|
||
|
|
for name, _ in submodule.named_buffers(
|
||
|
|
recurse=True, remove_duplicate=False
|
||
|
|
):
|
||
|
|
params_buffers_to_node_meta[target + "." + name] = meta
|
||
|
|
|
||
|
|
if node.op == "get_attr":
|
||
|
|
submodule = getattr(gm_torch_level, target)
|
||
|
|
if not isinstance(submodule, torch.fx.GraphModule):
|
||
|
|
params_buffers_to_node_meta[target] = meta
|
||
|
|
|
||
|
|
# If the call_function uses param as input, we also need to update params' meta
|
||
|
|
# with this call_function node's meta.
|
||
|
|
# This is basically the same flow as torch.fx.traceback.preserve_meta()
|
||
|
|
if node.op == "call_function" and not isinstance(
|
||
|
|
node.target, torch._ops.HigherOrderOperator
|
||
|
|
):
|
||
|
|
for arg in node._input_nodes:
|
||
|
|
if arg.op == "get_attr":
|
||
|
|
for entry in torch.fx.proxy._COPY_META_FIELDS:
|
||
|
|
if entry in meta:
|
||
|
|
params_buffers_to_node_meta[arg.target][entry] = meta[entry]
|
||
|
|
|
||
|
|
# Fix the graph output signature to be tuple if scalar
|
||
|
|
out_spec = orig_out_spec = gm_torch_level._out_spec
|
||
|
|
assert out_spec is not None
|
||
|
|
# aot_export expect the return type to always be a tuple.
|
||
|
|
if out_spec.type not in (list, tuple):
|
||
|
|
out_spec = pytree.TreeSpec(tuple, None, [out_spec])
|
||
|
|
|
||
|
|
orig_arg_names = gm_torch_level.graph._codegen.pytree_info.orig_args # type: ignore[attr-defined]
|
||
|
|
|
||
|
|
gm_torch_level.graph._codegen = _PyTreeCodeGen(
|
||
|
|
_PyTreeInfo(
|
||
|
|
orig_arg_names,
|
||
|
|
gm_torch_level._in_spec,
|
||
|
|
out_spec,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
gm_torch_level.recompile()
|
||
|
|
|
||
|
|
_normalize_nn_module_stack(gm_torch_level, type(mod))
|
||
|
|
|
||
|
|
# NOTE: graph module expects only positional args
|
||
|
|
ep_non_strict = _export_non_strict(
|
||
|
|
gm_torch_level,
|
||
|
|
_convert_to_positional_args(orig_arg_names, fake_args, fake_kwargs),
|
||
|
|
{},
|
||
|
|
fake_params_buffers,
|
||
|
|
constant_attrs,
|
||
|
|
pre_dispatch=pre_dispatch,
|
||
|
|
)
|
||
|
|
|
||
|
|
gm = ep_non_strict.gm
|
||
|
|
export_graph_signature = ep_non_strict.sig
|
||
|
|
constants = ep_non_strict.constants
|
||
|
|
|
||
|
|
# After aot_export, set the param/buffer metadata back into placeholders
|
||
|
|
# Technically, users can still construct this data from param names
|
||
|
|
# without relying on this metadata
|
||
|
|
for node in gm.graph.nodes:
|
||
|
|
if node.op == "placeholder":
|
||
|
|
if node.target in export_graph_signature.inputs_to_parameters:
|
||
|
|
param_name = export_graph_signature.inputs_to_parameters[node.target]
|
||
|
|
if param_name in params_buffers_to_node_meta:
|
||
|
|
for k, v in params_buffers_to_node_meta[param_name].items():
|
||
|
|
node.meta[k] = v
|
||
|
|
if node.target in export_graph_signature.inputs_to_buffers:
|
||
|
|
buffer_name = export_graph_signature.inputs_to_buffers[node.target]
|
||
|
|
if buffer_name in params_buffers_to_node_meta:
|
||
|
|
for k, v in params_buffers_to_node_meta[buffer_name].items():
|
||
|
|
node.meta[k] = v
|
||
|
|
|
||
|
|
# The unbacked symint symbols are updated in aot_export
|
||
|
|
# so we serialize them here instead of inside dynamo
|
||
|
|
|
||
|
|
gm.meta["inline_constraints"] = {
|
||
|
|
k: v
|
||
|
|
for k, v in dynamo_fake_mode.shape_env.var_to_range.items()
|
||
|
|
if free_unbacked_symbols(k)
|
||
|
|
}
|
||
|
|
|
||
|
|
num_lifted = next(
|
||
|
|
(
|
||
|
|
i
|
||
|
|
for i, s in enumerate(export_graph_signature.input_specs)
|
||
|
|
if s.kind == InputKind.USER_INPUT
|
||
|
|
),
|
||
|
|
len(export_graph_signature.input_specs),
|
||
|
|
)
|
||
|
|
range_constraints = _process_constraints(
|
||
|
|
dynamo_fake_mode,
|
||
|
|
gm,
|
||
|
|
num_lifted,
|
||
|
|
flat_args,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Do some cleanups on the graph module to restore the state dict to the
|
||
|
|
# expected form. Each of these steps should probably get fixed upstream.
|
||
|
|
# 1. Remove tensor constants that were added as buffers.
|
||
|
|
_rewrite_dynamo_tensor_constants(
|
||
|
|
orig_mod_buffers=set(mod.buffers()),
|
||
|
|
traced_mod_buffers=dict(gm_torch_level.named_buffers()),
|
||
|
|
graph_signature=ep_non_strict.sig,
|
||
|
|
constants=ep_non_strict.constants,
|
||
|
|
)
|
||
|
|
# 2. Restore FQN of param/buffers
|
||
|
|
param_buffer_table: Dict[str, str] = _get_param_buffer_mapping(mod, gm_torch_level)
|
||
|
|
_replace_param_buffer_names(param_buffer_table, export_graph_signature)
|
||
|
|
|
||
|
|
# 3. Remove non-persistent buffers from the graph signature
|
||
|
|
_rewrite_non_persistent_buffers(mod, ep_non_strict.sig, ep_non_strict.constants)
|
||
|
|
|
||
|
|
# 4. Rewrite constants to have the same FQN as the original module.
|
||
|
|
_remap_constants(constant_attrs, export_graph_signature, constants)
|
||
|
|
|
||
|
|
module_call_signatures = {
|
||
|
|
fqn: ModuleCallSignature(inputs=[], outputs=[], **specs)
|
||
|
|
for fqn, specs in gm_torch_level.meta["module_call_specs"].items()
|
||
|
|
}
|
||
|
|
|
||
|
|
if len(preserve_module_call_signature) > 0:
|
||
|
|
res = CollectTracepointsPass(module_call_signatures, export_graph_signature)(gm)
|
||
|
|
assert res is not None
|
||
|
|
gm = res.graph_module
|
||
|
|
|
||
|
|
assert orig_out_spec is not None
|
||
|
|
exported_program = ExportedProgram(
|
||
|
|
root=gm,
|
||
|
|
graph=gm.graph,
|
||
|
|
graph_signature=export_graph_signature,
|
||
|
|
state_dict=mod.state_dict(keep_vars=True),
|
||
|
|
range_constraints=range_constraints,
|
||
|
|
module_call_graph=[
|
||
|
|
ModuleCallEntry(
|
||
|
|
"",
|
||
|
|
ModuleCallSignature(
|
||
|
|
inputs=[], outputs=[], in_spec=orig_in_spec, out_spec=orig_out_spec
|
||
|
|
),
|
||
|
|
)
|
||
|
|
]
|
||
|
|
+ [ModuleCallEntry(fqn, sig) for fqn, sig in module_call_signatures.items()],
|
||
|
|
example_inputs=(args, kwargs),
|
||
|
|
constants=constants,
|
||
|
|
)
|
||
|
|
log.debug("Exported program from AOTAutograd:\n%s", exported_program)
|
||
|
|
|
||
|
|
if len(range_constraints) > 0:
|
||
|
|
exported_program = exported_program._transform_do_not_use(
|
||
|
|
_AddRuntimeAssertionsForInlineConstraintsPass(range_constraints)
|
||
|
|
)
|
||
|
|
|
||
|
|
return exported_program
|