Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,637 @@
|
||||
"""
|
||||
Utilities for reproducing and debugging issues in Dynamo after graph capture.
|
||||
|
||||
This file provides tools and infrastructure for debugging problems that occur
|
||||
after Dynamo has captured the graph but before/during backend compilation.
|
||||
Key components include:
|
||||
|
||||
- Minification tools to reduce large graphs to minimal failing examples
|
||||
- Accuracy testing to validate compiled graph outputs match eager mode
|
||||
- Repro generation to create standalone reproduction scripts
|
||||
- Debug backends for capturing and analyzing failures
|
||||
- Utilities for saving/loading graph states and inputs
|
||||
|
||||
The tools here focus specifically on the post-graph-capture stage, making them
|
||||
useful for debugging backend compilation issues, AOTAutograd problems, and
|
||||
accuracy discrepancies between compiled and eager execution.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import textwrap
|
||||
from collections.abc import Callable, Sequence
|
||||
from importlib import import_module
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.fx as fx
|
||||
from torch._dynamo.debug_utils import (
|
||||
AccuracyError,
|
||||
backend_accuracy_fails,
|
||||
BUCK_CMD_PREFIX,
|
||||
BuckTargetWriter,
|
||||
extra_imports,
|
||||
generate_config_string,
|
||||
generate_env_vars_string,
|
||||
helper_for_dump_minify,
|
||||
InputReader,
|
||||
InputWriter,
|
||||
minifier_dir,
|
||||
NNModuleToString,
|
||||
NopInputReader,
|
||||
run_fwd_maybe_bwd,
|
||||
same_two_models,
|
||||
)
|
||||
from torch.fx.experimental.symbolic_shapes import fx_placeholder_targets
|
||||
from torch.hub import tqdm
|
||||
|
||||
from .. import config
|
||||
from ..backends.registry import CompilerFn, lookup_backend, register_debug_backend
|
||||
from ..debug_utils import clone_inputs_retaining_gradness
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
inductor_config = import_module("torch._inductor.config")
|
||||
use_buck = inductor_config.is_fbcode()
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# MAIN ENTRY POINT
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def _accuracy_fails(
|
||||
gm: torch.fx.GraphModule,
|
||||
example_inputs: Sequence[Any],
|
||||
compiler_fn: Callable[[torch.fx.GraphModule, list[Any]], torch.fx.GraphModule],
|
||||
) -> bool:
|
||||
return backend_accuracy_fails(
|
||||
gm,
|
||||
example_inputs,
|
||||
compiler_fn,
|
||||
only_fwd=config.repro_forward_only,
|
||||
ignore_non_fp=config.repro_ignore_non_fp,
|
||||
)
|
||||
|
||||
|
||||
class WrapBackendDebug:
|
||||
def __init__(
|
||||
self, unconfigured_compiler_fn: CompilerFn, compiler_name: str | None
|
||||
) -> None:
|
||||
functools.wraps(unconfigured_compiler_fn)(self)
|
||||
self._torchdynamo_orig_backend = unconfigured_compiler_fn
|
||||
self._compiler_name = compiler_name
|
||||
if hasattr(unconfigured_compiler_fn, "__name__"):
|
||||
self.__name__ = unconfigured_compiler_fn.__name__
|
||||
if hasattr(unconfigured_compiler_fn, "compiler_name"):
|
||||
self.__name__ = unconfigured_compiler_fn.compiler_name # type: ignore[attr-defined]
|
||||
if hasattr(unconfigured_compiler_fn, "get_compiler_config"):
|
||||
self.get_compiler_config = unconfigured_compiler_fn.get_compiler_config # type: ignore[attr-defined]
|
||||
|
||||
def __call__(
|
||||
self, gm: torch.fx.GraphModule, example_inputs: list[Any], **kwargs: Any
|
||||
) -> torch.fx.GraphModule:
|
||||
compiler_fn = functools.partial(self._torchdynamo_orig_backend, **kwargs)
|
||||
assert config.repro_after in ("dynamo", "aot", None)
|
||||
|
||||
if config.repro_after == "dynamo":
|
||||
|
||||
def add_paths(exc: Exception) -> None:
|
||||
exc.minifier_path = os.path.join(minifier_dir(), "minifier_launcher.py") # type: ignore[attr-defined]
|
||||
if use_buck:
|
||||
exc.buck_command = " ".join( # type: ignore[attr-defined]
|
||||
BUCK_CMD_PREFIX
|
||||
+ [BuckTargetWriter(exc.minifier_path).cmd_line_path] # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
if config.repro_level == 3:
|
||||
dump_to_minify_after_dynamo(gm, example_inputs, self._compiler_name)
|
||||
|
||||
# Check for either accuracy (level 4) or other type of failures.
|
||||
if config.repro_level == 4:
|
||||
# Check Accuracy
|
||||
compiled_gm = compiler_fn(copy.deepcopy(gm), example_inputs)
|
||||
if _accuracy_fails(gm, example_inputs, compiler_fn): # type: ignore[arg-type]
|
||||
log.warning(
|
||||
"Accuracy failed for the TorchDynamo produced graph. Creating script to minify the error."
|
||||
)
|
||||
dump_to_minify_after_dynamo(
|
||||
fx.GraphModule(gm, copy.deepcopy(gm.graph)),
|
||||
example_inputs,
|
||||
self._compiler_name,
|
||||
)
|
||||
exc = AccuracyError("Bad accuracy detected.")
|
||||
add_paths(exc)
|
||||
raise exc
|
||||
else:
|
||||
try:
|
||||
compiled_gm = compiler_fn(copy.deepcopy(gm), example_inputs)
|
||||
run_fwd_maybe_bwd(compiled_gm, example_inputs) # type: ignore[arg-type]
|
||||
except Exception as exc:
|
||||
log.warning(
|
||||
"Compiled Fx GraphModule failed. Creating script to minify the error."
|
||||
)
|
||||
if config.repro_level == 1:
|
||||
dump_state_fn = functools.partial(
|
||||
dump_backend_state, compiler_name=self._compiler_name
|
||||
)
|
||||
dump_state_fn(
|
||||
fx.GraphModule(gm, copy.deepcopy(gm.graph)), example_inputs
|
||||
)
|
||||
elif config.repro_level == 2:
|
||||
dump_to_minify_after_dynamo(
|
||||
fx.GraphModule(gm, copy.deepcopy(gm.graph)),
|
||||
example_inputs,
|
||||
self._compiler_name,
|
||||
)
|
||||
add_paths(exc)
|
||||
raise
|
||||
else:
|
||||
compiled_gm = compiler_fn(gm, example_inputs)
|
||||
|
||||
return compiled_gm # type: ignore[return-value]
|
||||
|
||||
|
||||
def wrap_backend_debug(
|
||||
unconfigured_compiler_fn: CompilerFn, compiler_name: str | None
|
||||
) -> WrapBackendDebug:
|
||||
"""
|
||||
A minifier decorator that wraps the TorchDynamo produced Fx graph modules.
|
||||
As opposed to wrap_compiler_debug, this wrapper intercepts at the
|
||||
TorchDynamo produced Fx Graph Module. This makes it backend-agnostic to some
|
||||
level, e.g., it is useful for minifying issues related to Aot Autograd
|
||||
tracing. If an error is found, we minify and save the minified repro in
|
||||
repro.tar.gz.
|
||||
"""
|
||||
return WrapBackendDebug(unconfigured_compiler_fn, compiler_name)
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# REPRO DUMPERS
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def generate_dynamo_fx_repro_string(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: Sequence[Any],
|
||||
compiler_name: str | None,
|
||||
check_accuracy: bool = False,
|
||||
*,
|
||||
stable_output: bool = False,
|
||||
save_dir: str | None = None,
|
||||
command: str = "run",
|
||||
) -> str:
|
||||
"""
|
||||
Generate a repro string for backend-agnostic minified version.
|
||||
"""
|
||||
|
||||
model_str = NNModuleToString.convert(gm)
|
||||
|
||||
# TODO: Figure out why torch.compile'd hash isn't work on this codepath
|
||||
writer = InputWriter(save_dir, stable_hash=True)
|
||||
for placeholder, arg in zip(fx_placeholder_targets(gm), args):
|
||||
if isinstance(arg, (int, torch.SymInt)):
|
||||
writer.symint(placeholder, arg)
|
||||
elif isinstance(arg, torch.Tensor):
|
||||
# TODO: improve these names with FQN
|
||||
writer.tensor(placeholder, arg)
|
||||
else:
|
||||
raise TypeError(f"arg is neither SymInt/int nor torch.Tensor, {arg}")
|
||||
load_args = "\n".join(writer.lines())
|
||||
|
||||
return textwrap.dedent(
|
||||
f"""
|
||||
{generate_env_vars_string(stable_output=stable_output)}
|
||||
from math import inf
|
||||
import torch
|
||||
from torch import tensor, device
|
||||
import torch.fx as fx
|
||||
import torch._dynamo
|
||||
from torch._dynamo.testing import rand_strided
|
||||
from torch._dynamo.debug_utils import run_fwd_maybe_bwd
|
||||
|
||||
{generate_config_string(stable_output=stable_output)}
|
||||
|
||||
{extra_imports}
|
||||
|
||||
{model_str}
|
||||
mod = Repro()
|
||||
|
||||
{load_args}
|
||||
|
||||
if __name__ == '__main__':
|
||||
from torch._dynamo.repro.after_dynamo import run_repro
|
||||
run_repro(mod, load_args, accuracy={check_accuracy!r}, command={command!r},
|
||||
save_dir={save_dir!r}, autocast={torch.is_autocast_enabled()!r}, backend={compiler_name!r})
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def dump_backend_repro_as_file(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: Sequence[Any],
|
||||
compiler_name: str | None,
|
||||
check_accuracy: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Saves the repro to a repro.py file
|
||||
"""
|
||||
curdir = os.getcwd()
|
||||
subdir = os.path.join(os.getcwd(), "checkpoints")
|
||||
if not os.path.exists(subdir):
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
file_name = os.path.join(subdir, f"minified_{len(gm.graph.nodes)}_nodes.py")
|
||||
log.warning(
|
||||
"Writing checkpoint with %s nodes to %s", len(gm.graph.nodes), file_name
|
||||
)
|
||||
|
||||
with open(file_name, "w") as fd:
|
||||
fd.write(
|
||||
generate_dynamo_fx_repro_string(
|
||||
gm, args, compiler_name, check_accuracy, save_dir=subdir
|
||||
)
|
||||
)
|
||||
latest_repro = os.path.join(curdir, "repro.py")
|
||||
log.warning("Copying %s to %s for convenience", file_name, latest_repro)
|
||||
|
||||
if use_buck:
|
||||
BuckTargetWriter(latest_repro).write()
|
||||
|
||||
shutil.copyfile(file_name, latest_repro)
|
||||
|
||||
|
||||
def dump_backend_state(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: Sequence[Any],
|
||||
compiler_name: str | None,
|
||||
check_accuracy: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Dumps the dynamo graph to repro the issue.
|
||||
1) It tries to convert Fx GraphModule to a string. If we can, it writes to a
|
||||
repro.py file.
|
||||
2) If we can't convert Fx GraphModule to a string, we use to_folder to save
|
||||
the module and save a tar file.
|
||||
"""
|
||||
assert NNModuleToString.can_convert_to_string(gm)
|
||||
return dump_backend_repro_as_file(gm, args, compiler_name, check_accuracy)
|
||||
# return dump_backend_repro_as_tarfile(gm, args, compiler_name)
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# MINIFIER DUMPER
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def dump_to_minify_after_dynamo(
|
||||
gm: torch.fx.GraphModule, args: Sequence[Any], compiler_name: str | None
|
||||
) -> None:
|
||||
# TODO: factor this out
|
||||
subdir = os.path.join(minifier_dir(), "checkpoints")
|
||||
if not os.path.exists(subdir):
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
helper_for_dump_minify(
|
||||
generate_dynamo_fx_repro_string(
|
||||
gm,
|
||||
args,
|
||||
compiler_name,
|
||||
check_accuracy=config.repro_level == 4,
|
||||
save_dir=subdir,
|
||||
command="minify",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# MINIFIER BACKENDS
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
@register_debug_backend # type: ignore[arg-type]
|
||||
def dynamo_minifier_backend(
|
||||
gm: fx.GraphModule, example_inputs: Sequence[Any], compiler_name: str | None
|
||||
) -> fx.GraphModule:
|
||||
from functorch.compile import minifier
|
||||
|
||||
compiler_fn = lookup_backend(compiler_name) # type: ignore[arg-type]
|
||||
|
||||
# TODO: It's inconsistent to pass SymInt inputs but REAL tensors.
|
||||
# We should pass ints and look at the GraphModule placeholders
|
||||
# to resolve them to SymInt (if necessary)
|
||||
example_inputs = [
|
||||
i.node.hint if isinstance(i, torch.SymInt) else i for i in example_inputs
|
||||
]
|
||||
|
||||
try:
|
||||
compiled_gm = compiler_fn(gm, example_inputs)
|
||||
run_fwd_maybe_bwd(compiled_gm, example_inputs) # type: ignore[arg-type]
|
||||
raise ValueError("No issue was detected")
|
||||
except Exception as exc:
|
||||
orig_failure = str(exc)
|
||||
log.warning(
|
||||
"Compiled Fx GraphModule failed. Creating script to minify the error."
|
||||
)
|
||||
dump_state_fn = functools.partial(
|
||||
dump_backend_state, compiler_name=compiler_name
|
||||
)
|
||||
dump_state_fn(fx.GraphModule(gm, copy.deepcopy(gm.graph)), example_inputs)
|
||||
fails_fn = functools.partial(
|
||||
backend_fails,
|
||||
compiler_fn=compiler_fn,
|
||||
orig_failure=orig_failure,
|
||||
)
|
||||
minifier(
|
||||
gm,
|
||||
example_inputs,
|
||||
module_fails=fails_fn,
|
||||
dump_state=dump_state_fn,
|
||||
)
|
||||
return gm
|
||||
|
||||
|
||||
@register_debug_backend # type: ignore[arg-type]
|
||||
def dynamo_accuracy_minifier_backend(
|
||||
gm: fx.GraphModule, example_inputs: Sequence[Any], compiler_name: str | None
|
||||
) -> fx.GraphModule:
|
||||
from functorch.compile import minifier
|
||||
|
||||
compiler_fn = lookup_backend(compiler_name) # type: ignore[arg-type]
|
||||
|
||||
# Set the eval mode to remove randomness.
|
||||
gm.eval()
|
||||
|
||||
# Check Accuracy
|
||||
if _accuracy_fails(gm, example_inputs, compiler_fn): # type: ignore[arg-type]
|
||||
log.warning("Accuracy failed for the TorchDynamo produced graph")
|
||||
dump_state_fn = functools.partial(
|
||||
dump_backend_state, compiler_name=compiler_name, check_accuracy=True
|
||||
)
|
||||
fails_fn = functools.partial(
|
||||
_accuracy_fails,
|
||||
compiler_fn=compiler_fn, # type: ignore[arg-type]
|
||||
)
|
||||
dump_state_fn(fx.GraphModule(gm, copy.deepcopy(gm.graph)), example_inputs)
|
||||
minifier(
|
||||
gm,
|
||||
example_inputs,
|
||||
module_fails=fails_fn,
|
||||
dump_state=dump_state_fn,
|
||||
)
|
||||
else:
|
||||
log.error("Input graph does not fail accuracy testing")
|
||||
return gm
|
||||
|
||||
|
||||
def backend_fails(
|
||||
gm: fx.GraphModule,
|
||||
example_inputs: Sequence[Any],
|
||||
compiler_fn: CompilerFn,
|
||||
orig_failure: Sequence[Any],
|
||||
) -> bool:
|
||||
"""
|
||||
Minifier uses this function to identify if the minified graph module fails
|
||||
with the same error.
|
||||
|
||||
One caveat is that minifier can potentially go into a wrong direction when
|
||||
the resulting graph module fails for a different reason. To avoid this, we
|
||||
save the string for the original exception and check similarity between new
|
||||
and old exception. They can be somewhat different in some cases, when the
|
||||
exception string depends on the failing node information. So, we have a
|
||||
loose similarity metric to guide the minifier path.
|
||||
"""
|
||||
from difflib import SequenceMatcher
|
||||
|
||||
try:
|
||||
# Run the original gm to check eager validity
|
||||
run_fwd_maybe_bwd(gm, clone_inputs_retaining_gradness(example_inputs))
|
||||
compiled_gm = compiler_fn(gm, example_inputs) # type: ignore[arg-type]
|
||||
run_fwd_maybe_bwd(compiled_gm, clone_inputs_retaining_gradness(example_inputs)) # type: ignore[arg-type]
|
||||
except Exception as e:
|
||||
new_failure = str(e)
|
||||
if SequenceMatcher(None, orig_failure, new_failure).ratio() > 0.5:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# REPRO MAIN
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def run_load_args(options: Any, mod: torch.nn.Module, load_args: Any) -> list[Any]:
|
||||
if not hasattr(load_args, "_version"):
|
||||
log.warning(
|
||||
"load_args does not have a _version attribute, please file a bug to PyTorch "
|
||||
"and describe how you generate this repro script"
|
||||
)
|
||||
else:
|
||||
if load_args._version > 0:
|
||||
log.warning(
|
||||
"load_args is version %s, but this version of PyTorch only supports "
|
||||
"version 0. We will try to run it anyway but there may be an incompatibility; "
|
||||
"if so, try upgrading your version of PyTorch.",
|
||||
load_args._version,
|
||||
)
|
||||
|
||||
nop_reader = NopInputReader()
|
||||
load_args(nop_reader)
|
||||
|
||||
with tqdm(desc="Loading inputs", total=nop_reader.total) as pbar:
|
||||
input_reader = InputReader(save_dir=options.save_dir, pbar=pbar)
|
||||
load_args(input_reader)
|
||||
args = input_reader.args
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def repro_minify(options: Any, mod: torch.nn.Module, load_args: Any) -> None:
|
||||
args = run_load_args(options, mod, load_args)
|
||||
|
||||
# Setup debug minifier compiler
|
||||
if not options.accuracy:
|
||||
compiler_fn = lookup_backend("dynamo_minifier_backend")
|
||||
else:
|
||||
compiler_fn = lookup_backend("dynamo_accuracy_minifier_backend")
|
||||
|
||||
if options.backend is None:
|
||||
raise RuntimeError(
|
||||
"Compiler name is None - this likely means that a custom compiler "
|
||||
"was called by torchdynamo. Please remove this error, import your "
|
||||
"custom compiler function, and replace the backend=None "
|
||||
"line in run_repro to backend=<my_imported_custom_function>"
|
||||
)
|
||||
|
||||
dynamo_minifier_backend = functools.partial(
|
||||
compiler_fn,
|
||||
compiler_name=options.backend, # type: ignore[call-arg]
|
||||
)
|
||||
opt_mod = torch._dynamo.optimize(dynamo_minifier_backend)(mod)
|
||||
|
||||
with torch.amp.autocast("cuda", enabled=options.autocast):
|
||||
opt_mod(*args)
|
||||
|
||||
|
||||
def repro_run(options: Any, mod: torch.nn.Module, load_args: Any) -> None:
|
||||
opt_mod = torch._dynamo.optimize(options.backend)(mod)
|
||||
|
||||
if options.accuracy != "":
|
||||
mod.eval()
|
||||
opt_mod.eval() # type: ignore[union-attr]
|
||||
|
||||
with torch.amp.autocast("cuda", enabled=options.autocast):
|
||||
# TODO: disable clone
|
||||
args = run_load_args(options, mod, load_args)
|
||||
assert same_two_models(mod, mod, args), "Eager itself failed" # type: ignore[arg-type]
|
||||
if not same_two_models(
|
||||
mod, # type: ignore[arg-type]
|
||||
opt_mod, # type: ignore[arg-type]
|
||||
args,
|
||||
only_fwd=config.repro_forward_only,
|
||||
ignore_non_fp=config.repro_ignore_non_fp,
|
||||
):
|
||||
raise AccuracyError("Dynamo failed")
|
||||
else:
|
||||
with torch.amp.autocast("cuda", enabled=options.autocast):
|
||||
args = run_load_args(options, mod, load_args)
|
||||
run_fwd_maybe_bwd(mod, args, only_fwd=options.only_fwd, disable_clone=True) # type: ignore[arg-type]
|
||||
del args
|
||||
|
||||
args = run_load_args(options, mod, load_args)
|
||||
run_fwd_maybe_bwd(
|
||||
opt_mod, # type: ignore[arg-type]
|
||||
args,
|
||||
only_fwd=options.only_fwd,
|
||||
disable_clone=True, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
def run_repro(
|
||||
mod: torch.nn.Module,
|
||||
load_args: Any,
|
||||
*,
|
||||
command: str = "run",
|
||||
accuracy: bool | str = "",
|
||||
save_dir: str | None = None,
|
||||
autocast: bool = False,
|
||||
backend: str = "inductor",
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
for k in kwargs:
|
||||
log.warning(
|
||||
"Unrecognized kwarg %s; perhaps this repro was made on a newer version of PyTorch",
|
||||
k,
|
||||
)
|
||||
|
||||
if accuracy is True:
|
||||
accuracy = "accuracy"
|
||||
elif accuracy is False:
|
||||
accuracy = ""
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description=f"""\
|
||||
An after_dynamo repro script, typically triggering a bug in Dynamo or
|
||||
AOTAutograd. When run with no arguments, this script defaults to running
|
||||
'{command}'. Extra flags may be available; to find out more, try '{command}
|
||||
--help'. There are also alternate subcommands available, see below.
|
||||
|
||||
default settings on this script:
|
||||
{accuracy=}
|
||||
{save_dir=}
|
||||
""",
|
||||
formatter_class=argparse.RawTextHelpFormatter,
|
||||
)
|
||||
|
||||
def common_flags(parser: argparse.ArgumentParser) -> None:
|
||||
accuracy_group = parser.add_mutually_exclusive_group()
|
||||
accuracy_group.add_argument(
|
||||
"--no-accuracy",
|
||||
dest="accuracy",
|
||||
action="store_const",
|
||||
const="",
|
||||
default=accuracy,
|
||||
help="do not test accuracy, just run the module and see if it errors",
|
||||
)
|
||||
accuracy_group.add_argument(
|
||||
"--accuracy",
|
||||
action="store_const",
|
||||
const="accuracy",
|
||||
default=accuracy,
|
||||
help="test accuracy",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-dir",
|
||||
type=str,
|
||||
default=save_dir,
|
||||
metavar="DIR",
|
||||
help="directory where saved inputs live",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-save-dir",
|
||||
dest="save_dir",
|
||||
action="store_const",
|
||||
const=None,
|
||||
help="don't use any directory for saved inputs",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-isolate",
|
||||
dest="isolate",
|
||||
action="store_false",
|
||||
default=False,
|
||||
help="no isolate (doesn't do anything for after_dynamo)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--autocast",
|
||||
default=autocast,
|
||||
action="store_true",
|
||||
help="use torch.cuda.amp.autocast",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-autocast",
|
||||
dest="autocast",
|
||||
action="store_false",
|
||||
help="don't use torch.cuda.amp.autocast",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--backend",
|
||||
type=str,
|
||||
default=backend,
|
||||
metavar="BACKEND",
|
||||
help="torch.compile backend to use",
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(
|
||||
dest="command", metavar="{run,minify}", required=True
|
||||
)
|
||||
|
||||
parser_run = subparsers.add_parser(
|
||||
"run",
|
||||
help="just run the repro",
|
||||
)
|
||||
common_flags(parser_run)
|
||||
parser_run.add_argument(
|
||||
"--only-fwd",
|
||||
action="store_true",
|
||||
help="don't run backwards compilation for testing",
|
||||
)
|
||||
|
||||
parser_minify = subparsers.add_parser(
|
||||
"minify", help="run the minifier on the repro"
|
||||
)
|
||||
common_flags(parser_minify)
|
||||
|
||||
args = None
|
||||
if len(sys.argv) <= 1:
|
||||
args = [command, *sys.argv[1:]]
|
||||
|
||||
options = parser.parse_args(args)
|
||||
COMMAND_FNS = {
|
||||
"minify": repro_minify,
|
||||
"run": repro_run,
|
||||
}
|
||||
COMMAND_FNS[options.command](options, mod, load_args)
|
||||
@@ -0,0 +1,662 @@
|
||||
"""
|
||||
Utilities for debugging and reproducing issues in Ahead of Time with Inductor (AOTI) compilation.
|
||||
|
||||
This file provides tools and utilities for:
|
||||
- Generating minimal reproducible test cases (minification)
|
||||
- Handling exported programs and graph modules
|
||||
- Creating debug repros for AOTI compilation issues
|
||||
- Supporting both accuracy testing and error reproduction
|
||||
- Managing configuration and environment for repro cases
|
||||
|
||||
The main components include:
|
||||
- Minification tools to reduce test cases while preserving errors
|
||||
- Repro generation utilities for exported programs
|
||||
- Error handling specific to AOTI compilation
|
||||
- Command-line interface for running and managing repros
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import functools
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
import textwrap
|
||||
from collections.abc import Sequence
|
||||
from importlib import import_module
|
||||
from typing import Any, IO
|
||||
|
||||
import torch
|
||||
from torch._dynamo.debug_utils import (
|
||||
_cuda_system_info_comment,
|
||||
BuckTargetWriter,
|
||||
extra_imports,
|
||||
generate_config_string,
|
||||
generate_env_vars_string,
|
||||
helper_for_dump_minify,
|
||||
InputReader,
|
||||
minifier_dir,
|
||||
NNModuleToString,
|
||||
NopInputReader,
|
||||
)
|
||||
from torch.export import ExportedProgram
|
||||
from torch.hub import tqdm
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
inductor_config = import_module("torch._inductor.config")
|
||||
use_buck = inductor_config.is_fbcode()
|
||||
|
||||
|
||||
class AOTIMinifierError(Exception):
|
||||
def __init__(self, original_exception: str | Exception) -> None:
|
||||
additional_message = "This error is caused by a bug in the AOTI minifier, please report a bug to PyTorch"
|
||||
full_message = f"{additional_message}: {str(original_exception)}"
|
||||
super().__init__(full_message)
|
||||
self.original_exception = original_exception
|
||||
|
||||
|
||||
def dump_to_minify(
|
||||
exported_program: ExportedProgram,
|
||||
compiler_name: str,
|
||||
command: str = "minify",
|
||||
options: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
If command is "minify":
|
||||
Dump exported_program to `debug_dir/minifier/minifier_launcher.py`, with minify command.
|
||||
If command is "run":
|
||||
Dump exported_program to `cwd/repro.py`, with run command.
|
||||
"""
|
||||
assert command in ["minify", "run"]
|
||||
|
||||
subdir = os.path.join(minifier_dir(), "checkpoints")
|
||||
if not os.path.exists(subdir):
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
|
||||
if command == "minify":
|
||||
out = io.StringIO()
|
||||
save_graph_repro_ep(
|
||||
out,
|
||||
compiler_name,
|
||||
exported_program=exported_program,
|
||||
save_dir=subdir,
|
||||
command="minify",
|
||||
config_patches=options,
|
||||
)
|
||||
return helper_for_dump_minify(out.getvalue())
|
||||
else:
|
||||
curdir = os.getcwd()
|
||||
file_name = os.path.join(curdir, "repro.py")
|
||||
try:
|
||||
with open(file_name, "w") as fd:
|
||||
save_graph_repro_ep(
|
||||
fd,
|
||||
compiler_name,
|
||||
exported_program=exported_program,
|
||||
config_patches=options,
|
||||
save_dir=subdir,
|
||||
command="run",
|
||||
module_in_comment=True,
|
||||
)
|
||||
log.warning("Writing repro file to %s", file_name)
|
||||
if use_buck:
|
||||
BuckTargetWriter(file_name).write()
|
||||
except OSError:
|
||||
log.warning("No write permissions for %s", file_name)
|
||||
|
||||
|
||||
def get_module_string(gm: torch.fx.GraphModule) -> str:
|
||||
def _convert_to_comment(s_: str) -> str:
|
||||
s = s_.split("\n")
|
||||
if len(s) == 1:
|
||||
return "# " + s_
|
||||
first = s.pop(0)
|
||||
for i in range(len(s)):
|
||||
line = s[i]
|
||||
if line.strip() != "":
|
||||
s[i] = "# " + line
|
||||
else:
|
||||
s[i] = ""
|
||||
s = "\n".join(s)
|
||||
s = first + "\n" + s
|
||||
return s
|
||||
|
||||
module_string = NNModuleToString.convert(gm)
|
||||
return _convert_to_comment(module_string)
|
||||
|
||||
|
||||
def save_graph_repro_ep(
|
||||
fd: IO[Any],
|
||||
compiler_name: str,
|
||||
*,
|
||||
exported_program: ExportedProgram | None = None,
|
||||
gm: torch.nn.Module | None = None,
|
||||
args: tuple[Any] | None = None,
|
||||
config_patches: dict[str, str] | None = None,
|
||||
stable_output: bool = False,
|
||||
save_dir: str | None = None,
|
||||
command: str = "run",
|
||||
accuracy: str | bool | None = None,
|
||||
check_str: str | None = None,
|
||||
module_in_comment: bool = False,
|
||||
strict: bool = False,
|
||||
) -> None:
|
||||
# Save graph for reproducing the error.
|
||||
# Either exported_program or gm will be saved, depending on which one is defined.
|
||||
# Only one of exported_program and gm should be defined.
|
||||
|
||||
if exported_program is None and gm is None:
|
||||
raise AOTIMinifierError("One of exported_program and gm must be defined")
|
||||
if exported_program is not None and gm is not None:
|
||||
raise AOTIMinifierError("Only one of exported_program and gm can be defined")
|
||||
if gm is not None and args is None:
|
||||
raise AOTIMinifierError("If gm is defined, args should also be defined")
|
||||
|
||||
if exported_program is None:
|
||||
assert gm is not None
|
||||
assert args is not None
|
||||
exported_program = torch.export.export(gm, args, strict=strict)
|
||||
elif gm is None:
|
||||
gm = exported_program.module(check_guards=False)
|
||||
|
||||
# save a graph preview using gm
|
||||
module_string = get_module_string(gm) # type: ignore[arg-type]
|
||||
fd.write(module_string)
|
||||
|
||||
# save a graph repro using exported_program
|
||||
fd.write(
|
||||
generate_compiler_repro_exported_program(
|
||||
exported_program,
|
||||
options=config_patches,
|
||||
stable_output=stable_output,
|
||||
save_dir=save_dir,
|
||||
)
|
||||
)
|
||||
if accuracy is None:
|
||||
accuracy = "_accuracy" in compiler_name
|
||||
fd.write("if __name__ == '__main__':\n")
|
||||
fd.write(" from torch._dynamo.repro.aoti import run_repro\n")
|
||||
fd.write(
|
||||
f" with torch.no_grad():\n"
|
||||
f" run_repro(exported_program, config_patches=config_patches, accuracy={accuracy!r}, command={command!r}, "
|
||||
f"save_dir={save_dir!r}, check_str={check_str!r})\n"
|
||||
)
|
||||
|
||||
|
||||
def dump_compiler_graph_state(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: Sequence[Any],
|
||||
compiler_name: str,
|
||||
*,
|
||||
config_patches: dict[str, str] | None = None,
|
||||
accuracy: str | bool | None = None,
|
||||
strict: bool = False,
|
||||
) -> None:
|
||||
subdir = os.path.join(minifier_dir(), "checkpoints")
|
||||
if not os.path.exists(subdir):
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
file_name = os.path.join(subdir, f"{len(gm.graph.nodes)}.py")
|
||||
log.warning(
|
||||
"Writing checkpoint with %s nodes to %s", len(gm.graph.nodes), file_name
|
||||
)
|
||||
with open(file_name, "w") as fd:
|
||||
save_graph_repro_ep(
|
||||
fd,
|
||||
compiler_name,
|
||||
gm=gm,
|
||||
args=tuple(args),
|
||||
config_patches=config_patches,
|
||||
save_dir=subdir,
|
||||
accuracy=accuracy,
|
||||
module_in_comment=True,
|
||||
strict=strict,
|
||||
)
|
||||
curdir = os.getcwd()
|
||||
repro_path = os.path.join(curdir, "repro.py")
|
||||
try:
|
||||
shutil.copyfile(file_name, repro_path)
|
||||
log.warning("Copying repro file for convenience to %s", repro_path)
|
||||
if use_buck:
|
||||
BuckTargetWriter(file_name).write()
|
||||
except OSError:
|
||||
log.warning("No write permissions for %s", repro_path)
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# DUMP REPROS
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def generate_compiler_repro_exported_program(
|
||||
exported_program: ExportedProgram,
|
||||
*,
|
||||
options: dict[str, str] | None = None,
|
||||
stable_output: bool = False,
|
||||
save_dir: str | None = None,
|
||||
) -> str:
|
||||
model_str = textwrap.dedent(
|
||||
f"""
|
||||
{generate_env_vars_string(stable_output=stable_output)}
|
||||
import torch
|
||||
import torch._inductor.inductor_prims
|
||||
|
||||
{generate_config_string(stable_output=stable_output)}
|
||||
|
||||
isolate_fails_code_str = None
|
||||
|
||||
{extra_imports}
|
||||
|
||||
"""
|
||||
)
|
||||
if not stable_output:
|
||||
model_str += f"# torch version: {torch.version.__version__}\n"
|
||||
if hasattr(torch.version, "cuda"):
|
||||
model_str += f"# torch cuda version: {torch.version.cuda}\n"
|
||||
if hasattr(torch.version, "git_version"):
|
||||
model_str += f"# torch git version: {torch.version.git_version}\n\n\n"
|
||||
model_str += _cuda_system_info_comment()
|
||||
if save_dir:
|
||||
ep_path = os.path.join(save_dir, "exported_program.pt2")
|
||||
else:
|
||||
ep_path = "exported_program.pt2"
|
||||
torch.export.save(exported_program, ep_path)
|
||||
|
||||
model_str += f"exported_program = torch.export.load('{ep_path}')\n"
|
||||
model_str += "# print(exported_program.graph)\n"
|
||||
model_str += f"config_patches={options}\n"
|
||||
return model_str
|
||||
|
||||
|
||||
def repro_load_args(load_args: Any, save_dir: str | None) -> tuple[Any]:
|
||||
if not hasattr(load_args, "_version"):
|
||||
log.warning(
|
||||
"load_args does not have a _version attribute, please file a bug to PyTorch "
|
||||
"and describe how you generate this repro script"
|
||||
)
|
||||
else:
|
||||
if load_args._version > 0:
|
||||
log.warning(
|
||||
"load_args is version %s, but this version of PyTorch only supports "
|
||||
"version 0. We will try to run it anyway but there may be an incompatibility; "
|
||||
"if so, try upgrading your version of PyTorch.",
|
||||
load_args._version,
|
||||
)
|
||||
|
||||
nop_reader = NopInputReader()
|
||||
load_args(nop_reader)
|
||||
|
||||
with tqdm(desc="Loading inputs", total=nop_reader.total) as pbar:
|
||||
input_reader = InputReader(save_dir=save_dir, pbar=pbar)
|
||||
load_args(input_reader)
|
||||
args = input_reader.args
|
||||
|
||||
return tuple(args)
|
||||
|
||||
|
||||
def repro_common(
|
||||
options: Any, exported_program: ExportedProgram
|
||||
) -> tuple[torch.fx.GraphModule, Any, Any]:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
torch._inductor.config.generate_intermediate_hooks = True
|
||||
mod = exported_program.module(check_guards=False)
|
||||
args, kwargs = exported_program.example_inputs
|
||||
return mod, args, kwargs # type: ignore[return-value]
|
||||
|
||||
|
||||
def repro_get_args(
|
||||
options: Any,
|
||||
exported_program: ExportedProgram,
|
||||
config_patches: dict[str, Any] | None,
|
||||
) -> tuple[torch.fx.GraphModule, Any, Any]:
|
||||
mod, args, kwargs = repro_common(options, exported_program)
|
||||
return mod, args, kwargs
|
||||
|
||||
|
||||
def repro_run(
|
||||
options: Any,
|
||||
exported_program: ExportedProgram,
|
||||
config_patches: dict[str, Any] | None,
|
||||
) -> None:
|
||||
from torch._inductor import _aoti_compile_and_package_inner
|
||||
|
||||
gm, args, kwargs = repro_common(options, exported_program)
|
||||
|
||||
from torch.cuda import synchronize
|
||||
|
||||
_aoti_compile_and_package_inner(
|
||||
gm,
|
||||
args,
|
||||
kwargs,
|
||||
load_and_run=True,
|
||||
check_accuracy=options.accuracy,
|
||||
inductor_configs=config_patches,
|
||||
)
|
||||
|
||||
need_sync = False
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor) and arg.is_cuda:
|
||||
need_sync = True
|
||||
break
|
||||
|
||||
if need_sync:
|
||||
synchronize() # ensure segfaults are surfaced
|
||||
|
||||
|
||||
def export_for_aoti_minifier(
|
||||
gm: torch.nn.Module,
|
||||
tuple_inputs: tuple[Any],
|
||||
strict: bool = False,
|
||||
skip_export_error: bool = True,
|
||||
) -> torch.nn.Module | None:
|
||||
# Some graphs cannot be used for AOTI/export (illegal graphs), these should be
|
||||
# considered as graphs that don't fail in the minifier, so the minifier keeps searching.
|
||||
# In these case, we return None. Otherwise, we return the exported graph module.
|
||||
# This won't affect the minifier result because the minifier is only responsible for catching
|
||||
# errors in AOTI, not export.
|
||||
#
|
||||
# Please add to this list of illegal graphs if you change the implementation here.
|
||||
# - graph output is not allowed by export
|
||||
#
|
||||
# If skip_export_error=True, then the errors in export will not be raised, and the minifier
|
||||
# will keep exploring and ignore this graph.
|
||||
from torch._dynamo.exc import UserError, UserErrorType
|
||||
|
||||
try:
|
||||
ep = torch.export.export(gm, tuple_inputs, strict=strict)
|
||||
gm = ep.module(check_guards=False)
|
||||
return gm
|
||||
except Exception as e:
|
||||
if skip_export_error:
|
||||
return None
|
||||
if isinstance(e, UserError) and e.error_type == UserErrorType.INVALID_OUTPUT:
|
||||
# graph output is not allowed by export when strict=True
|
||||
return None
|
||||
if isinstance(e, RuntimeError):
|
||||
# graph output is not allowed by export when strict=False
|
||||
pattern = r"Found .* in output, which is not a known type\."
|
||||
if re.search(pattern, str(e)) is not None:
|
||||
return None
|
||||
raise AOTIMinifierError(e) from e
|
||||
# we should never reach here
|
||||
# pyrefly: ignore [unreachable]
|
||||
return None
|
||||
|
||||
|
||||
def repro_minify(
|
||||
options: Any,
|
||||
exported_program: ExportedProgram,
|
||||
config_patches: dict[str, Any] | None,
|
||||
) -> None:
|
||||
from functorch.compile import minifier
|
||||
from torch._inductor import _aoti_compile_and_package_inner
|
||||
from torch._inductor.compile_fx import _aoti_flatten_inputs
|
||||
|
||||
mod, args, kwargs = repro_common(options, exported_program)
|
||||
|
||||
# update serialized_in_spec and serialized_out_spec
|
||||
flat_example_inputs, inductor_configs = _aoti_flatten_inputs(
|
||||
mod, args, kwargs, options=config_patches
|
||||
)
|
||||
compiler_name = "aot_inductor"
|
||||
assert options.minifier_export_mode in ["dynamo", "python"]
|
||||
strict = options.minifier_export_mode == "dynamo"
|
||||
skip_export_error = options.skip_export_error
|
||||
|
||||
from torch.cuda import synchronize
|
||||
|
||||
need_sync = False
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor) and arg.is_cuda:
|
||||
need_sync = True
|
||||
break
|
||||
|
||||
def module_fails(
|
||||
gm: torch.fx.GraphModule,
|
||||
flat_example_inputs: list[Any],
|
||||
check_str: str | None = None,
|
||||
) -> bool:
|
||||
# Need to export first so the in_spec and out_spec are populated
|
||||
tuple_inputs = tuple(flat_example_inputs)
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
gm = export_for_aoti_minifier(
|
||||
gm, tuple_inputs, strict=strict, skip_export_error=skip_export_error
|
||||
)
|
||||
|
||||
# Some graphs cannot be used for AOTI/export (illegal graphs), these should be
|
||||
# considered as graphs that don't fail in the minifier, so the minifier keeps searching.
|
||||
if gm is None:
|
||||
return False
|
||||
|
||||
assert isinstance(gm, torch.fx.GraphModule)
|
||||
|
||||
try:
|
||||
_aoti_compile_and_package_inner(
|
||||
gm,
|
||||
tuple_inputs,
|
||||
load_and_run=True,
|
||||
check_accuracy=options.accuracy,
|
||||
inductor_configs=inductor_configs,
|
||||
)
|
||||
if need_sync:
|
||||
synchronize() # ensure segfaults are surfaced
|
||||
return False
|
||||
except Exception as e:
|
||||
if check_str is not None and check_str not in repr(e):
|
||||
return False
|
||||
return True
|
||||
|
||||
minifier(
|
||||
mod,
|
||||
flat_example_inputs,
|
||||
module_fails=functools.partial(module_fails, check_str=options.check_str),
|
||||
dump_state=functools.partial(
|
||||
dump_compiler_graph_state,
|
||||
compiler_name=compiler_name,
|
||||
config_patches=config_patches,
|
||||
accuracy=options.accuracy,
|
||||
strict=strict,
|
||||
),
|
||||
save_dir=options.save_dir,
|
||||
offload_to_disk=options.offload_to_disk,
|
||||
skip_offload=options.skip_saving_eager_intermediates,
|
||||
skip_sanity=options.skip_sanity,
|
||||
max_granularity=options.max_granularity,
|
||||
)
|
||||
|
||||
|
||||
def run_repro(
|
||||
exported_program: ExportedProgram,
|
||||
*,
|
||||
config_patches: dict[str, str] | None = None,
|
||||
command: str = "run",
|
||||
accuracy: bool | str = "",
|
||||
save_dir: str | None = None,
|
||||
tracing_mode: str | None = None,
|
||||
check_str: str | None = None,
|
||||
minifier_export_mode: str = "python",
|
||||
skip_export_error: bool = True,
|
||||
**more_kwargs: Any,
|
||||
) -> Any:
|
||||
for k in more_kwargs:
|
||||
log.warning(
|
||||
"Unrecognized kwarg %s; perhaps this repro was made on a newer version of PyTorch",
|
||||
k,
|
||||
)
|
||||
|
||||
if accuracy is True:
|
||||
accuracy = "accuracy"
|
||||
elif accuracy is False:
|
||||
accuracy = ""
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description=f"""\
|
||||
An AOTI repro script, typically triggering a bug in PyTorch AOTInductor.
|
||||
When run with no arguments, this script defaults to running '{command}'.
|
||||
Extra flags may be available; to find out more, try '{command} --help'.
|
||||
There are also alternate subcommands available, see below.
|
||||
|
||||
default settings on this script:
|
||||
{accuracy=}
|
||||
{tracing_mode=}
|
||||
{save_dir=}
|
||||
{check_str=}
|
||||
""",
|
||||
formatter_class=argparse.RawTextHelpFormatter,
|
||||
)
|
||||
|
||||
def common_flags(parser: argparse.ArgumentParser) -> None:
|
||||
accuracy_group = parser.add_mutually_exclusive_group()
|
||||
accuracy_group.add_argument(
|
||||
"--no-accuracy",
|
||||
dest="accuracy",
|
||||
action="store_const",
|
||||
const="",
|
||||
default=accuracy,
|
||||
help="do not test accuracy, just run the module and see if it errors",
|
||||
)
|
||||
accuracy_group.add_argument(
|
||||
"--accuracy",
|
||||
action="store_const",
|
||||
const="accuracy",
|
||||
default=accuracy,
|
||||
help="""\
|
||||
test if the RMSE between the compiled module and the fp64 reference is greater
|
||||
than eager and the fp64 reference. This is usually more reliable than the
|
||||
standard allclose test, as we expect numeric differences from compiling, often
|
||||
improving accuracy over eager. RMSE test allows for compiled module to
|
||||
diverge greatly from eager, as long as this divergence moves it closer to the
|
||||
'true' mathematical value of the network. Caveats: (1) double precision can
|
||||
still suffer from rounding error, so it is not a perfect reference (see for
|
||||
example 'Herbie: Automatically Improving Floating Point Accuracy') for
|
||||
approaches that detect the necessary working precision and compute it in
|
||||
arbitrary precision floating point; unfortunately, this is not practical for
|
||||
tensor computation; (2) if there are not enough samples in the output being
|
||||
compared, we may get unlucky and have an unlucky greater RMSE than eager; this
|
||||
could be overcome by applying a more rigorous statistical test at some
|
||||
p-value, which we leave for future work.
|
||||
""",
|
||||
)
|
||||
accuracy_group.add_argument(
|
||||
"--strict-accuracy",
|
||||
dest="accuracy",
|
||||
action="store_const",
|
||||
const="strict_accuracy",
|
||||
default=accuracy,
|
||||
help="""\
|
||||
by default, when doing accuracy minification we will reject reductions which
|
||||
change the divergence from a floating point divergence to a integral/boolean
|
||||
divergence. This is because some operations like ReLU involve temporarily
|
||||
sharp boundaries that smooth out again afterwards; without requiring
|
||||
divergence on floating point, the minifier will often fixate on divergent
|
||||
boolean tensor even though this is not the true source of the divergence.
|
||||
However, rejecting these reductions makes it more difficult for the minifier
|
||||
to make process. Using this option will let the minifier progress for ALL
|
||||
divergences--you just might not end up with a useful repro in the end.""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--save-dir",
|
||||
type=str,
|
||||
default=save_dir,
|
||||
metavar="DIR",
|
||||
help="directory where saved inputs live",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-save-dir",
|
||||
dest="save_dir",
|
||||
action="store_const",
|
||||
const=None,
|
||||
help="don't use any directory for saved inputs",
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(
|
||||
dest="command", metavar="{run,minify}", required=True
|
||||
)
|
||||
|
||||
parser_run = subparsers.add_parser(
|
||||
"run",
|
||||
help="just run the repro",
|
||||
)
|
||||
common_flags(parser_run)
|
||||
|
||||
parser_minify = subparsers.add_parser(
|
||||
"minify", help="run the minifier on the repro"
|
||||
)
|
||||
common_flags(parser_minify)
|
||||
parser_get_args = subparsers.add_parser("get_args", help="get the args")
|
||||
common_flags(parser_get_args)
|
||||
parser_minify.add_argument(
|
||||
"--skip-saving-eager-intermediates",
|
||||
action="store_true",
|
||||
help="skip saving eager intermediates on --minify",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--offload-to-disk",
|
||||
action="store_true",
|
||||
help="during minification, offload delta debugging intermediates to disk. Use if you're OOMing",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--skip-sanity",
|
||||
action="store_true",
|
||||
help="skip sanity check at beginning of minification on original graph",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--max-granularity",
|
||||
type=int,
|
||||
default=None,
|
||||
help="start at this granularity and work down; must be power of 2",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--check-str",
|
||||
type=str,
|
||||
default=check_str,
|
||||
help="require minified program to fail with error containing this string",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--minifier-export-mode",
|
||||
type=str,
|
||||
default=minifier_export_mode,
|
||||
help=(
|
||||
"The export mode used in minifier, either dynamo or python."
|
||||
"`dynamo` corresponds to strict=True, and `python` corresponds to strict=False."
|
||||
),
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--skip-export-error",
|
||||
type=bool,
|
||||
default=skip_export_error,
|
||||
help="Skip intermediate graphs that cannot be exported.",
|
||||
)
|
||||
|
||||
# Run the repro in the context of minification, inverting exit code meaning
|
||||
parser_minifier_query = subparsers.add_parser(
|
||||
"minifier-query",
|
||||
)
|
||||
common_flags(parser_minifier_query)
|
||||
parser_minifier_query.add_argument(
|
||||
"--check-str",
|
||||
type=str,
|
||||
default=check_str,
|
||||
help="require minified program to fail with error containing this string",
|
||||
)
|
||||
|
||||
args = None
|
||||
if len(sys.argv) <= 1:
|
||||
args = [command, *sys.argv[1:]]
|
||||
|
||||
options = parser.parse_args(args)
|
||||
COMMAND_FNS = {
|
||||
"minify": repro_minify,
|
||||
"run": repro_run,
|
||||
"get_args": repro_get_args,
|
||||
}
|
||||
return COMMAND_FNS[options.command](
|
||||
options, exported_program, config_patches=config_patches
|
||||
)
|
||||
Reference in New Issue
Block a user