Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,812 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import contextlib
|
||||
import io
|
||||
from collections.abc import Callable
|
||||
from typing import Any, TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
import torch
|
||||
from torch._higher_order_ops.invoke_subgraph import NestedCompileRegionOptions
|
||||
|
||||
from . import config
|
||||
from ._cache import CacheInfo
|
||||
|
||||
|
||||
__all__ = [
|
||||
"compile",
|
||||
"config",
|
||||
"assume_constant_result",
|
||||
"reset",
|
||||
"allow_in_graph",
|
||||
"substitute_in_graph",
|
||||
"list_backends",
|
||||
"disable",
|
||||
"set_stance",
|
||||
"set_enable_guard_collectives",
|
||||
"cudagraph_mark_step_begin",
|
||||
"load_compiled_function",
|
||||
"wrap_numpy",
|
||||
"is_compiling",
|
||||
"is_dynamo_compiling",
|
||||
"is_exporting",
|
||||
"save_cache_artifacts",
|
||||
"load_cache_artifacts",
|
||||
"keep_portable_guards_unsafe",
|
||||
"skip_guard_on_inbuilt_nn_modules_unsafe",
|
||||
"skip_guard_on_all_nn_modules_unsafe",
|
||||
"keep_tensor_guards_unsafe",
|
||||
"skip_guard_on_globals_unsafe",
|
||||
"skip_all_guards_unsafe",
|
||||
"nested_compile_region",
|
||||
]
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
FuncType = Callable[..., Any]
|
||||
F = TypeVar("F", bound=FuncType)
|
||||
|
||||
|
||||
def compile(*args, **kwargs):
|
||||
"""
|
||||
See :func:`torch.compile` for details on the arguments for this function.
|
||||
"""
|
||||
# pyrefly: ignore [not-iterable]
|
||||
return torch.compile(*args, **kwargs)
|
||||
|
||||
|
||||
def reset() -> None:
|
||||
"""
|
||||
This function clears all compilation caches and restores the system to its initial state.
|
||||
It is recommended to call this function, especially after using operations like `torch.compile(...)`
|
||||
to ensure a clean state before another unrelated compilation
|
||||
"""
|
||||
import torch._dynamo
|
||||
|
||||
torch._dynamo.reset()
|
||||
|
||||
|
||||
def allow_in_graph(fn):
|
||||
"""
|
||||
Tells the compiler frontend (Dynamo) to skip symbolic introspection of the function
|
||||
and instead directly write it to the graph when encountered.
|
||||
|
||||
If you are using :func:`torch.compile` (with backend="inductor" (the default)), or
|
||||
:func:`torch.export.export`, and trying to black-box a Python function throughout
|
||||
all tracing, do not use this API.
|
||||
Instead, please create a custom operator (see `PyTorch Custom Operators Landing Page
|
||||
<https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html>`_)
|
||||
|
||||
.. warning::
|
||||
|
||||
If you're a typical torch.compile user (e.g. you're applying torch.compile to
|
||||
a model to make it run faster), you probably don't want to use this function.
|
||||
:func:`allow_in_graph` is a footgun because it skips the compiler frontend
|
||||
(Dynamo) that is responsible for doing safety checks (graph breaks, handling
|
||||
closures, etc). Incorrect usage will lead to difficult-to-debug silent
|
||||
incorrectness issues.
|
||||
|
||||
Given a Python function with no allow_in_graph decorator, regular execution
|
||||
of torch.compile traces through the function. :func:`allow_in_graph` changes
|
||||
it so that the frontend does not trace inside the function, but the compiler
|
||||
backend still traces through it. Compare this to custom operators, which
|
||||
treats a function as a black box throughout the torch.compile stack. The following
|
||||
table compares these mechanisms.
|
||||
|
||||
+------------------------+-----------------------+--------------------------------+
|
||||
| Mechanism | Frontend (Dynamo) | Backend (AOTAutograd+Inductor) |
|
||||
+========================+=======================+================================+
|
||||
| no decorator | trace inside | trace inside |
|
||||
+------------------------+-----------------------+--------------------------------+
|
||||
| allow_in_graph | opaque callable | trace inside |
|
||||
+------------------------+-----------------------+--------------------------------+
|
||||
| custom op | opaque callable | opaque callable |
|
||||
+------------------------+-----------------------+--------------------------------+
|
||||
|
||||
One common use case for :func:`allow_in_graph()` is as an escape hatch for the compiler
|
||||
frontend: if you know the function works w.r.t. to the downstream components of the
|
||||
compilation stack (AOTAutograd and Inductor) but there is a Dynamo bug that prevents it from
|
||||
symbolically introspecting the function properly (or if your code is in C/C++ and
|
||||
therefore cannot be introspected with Dynamo), then one can decorate said function
|
||||
with :func:`allow_in_graph` to bypass Dynamo.
|
||||
|
||||
We require that ``fn`` adhere to the following restrictions. Failure to adhere
|
||||
results in undefined behavior:
|
||||
|
||||
- The inputs to ``fn`` must be Proxy-able types in the FX graph. Valid types include:
|
||||
Tensor/int/bool/float/None/List[Tensor?]/List[int?]/List[float?]
|
||||
Tuple[Tensor?, ...]/Tuple[int?, ...]/Tuple[float?, ...]/torch.dtype/torch.device
|
||||
- The outputs to ``fn`` must be Proxy-able types in the FX graph (see previous bullet)
|
||||
- all Tensors used inside of ``fn`` must be passed directly as inputs to ``fn``
|
||||
(as opposed to being captured variables).
|
||||
|
||||
Args:
|
||||
fn: A callable representing the function to be included in the graph.
|
||||
If ``fn`` is a list or tuple of callables it recursively applies
|
||||
:func:`allow_in_graph()` to each function and returns a new list or
|
||||
tuple containing the modified functions.
|
||||
|
||||
Example::
|
||||
|
||||
torch.compiler.allow_in_graph(my_custom_function)
|
||||
|
||||
|
||||
@torch.compile(...)
|
||||
def fn(x):
|
||||
x = torch.add(x, 1)
|
||||
x = my_custom_function(x)
|
||||
x = torch.add(x, 1)
|
||||
return x
|
||||
|
||||
|
||||
fn(...)
|
||||
|
||||
Will capture a single graph containing ``my_custom_function()``.
|
||||
|
||||
"""
|
||||
import torch._dynamo
|
||||
|
||||
return torch._dynamo.allow_in_graph(fn)
|
||||
|
||||
|
||||
def substitute_in_graph(
|
||||
original_fn: Callable[_P, _R],
|
||||
*,
|
||||
can_constant_fold_through: bool = False,
|
||||
skip_signature_check: bool = False,
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
|
||||
"""
|
||||
Register a polyfill handler for a function, usually a C function from the C extension, to be
|
||||
used in place of the original function when inlining the original function in the graph.
|
||||
|
||||
.. note::
|
||||
|
||||
The polyfill handler is only used when inlining the original function. It is not used when
|
||||
the original function is called directly. In the eager mode, the decorated function calls
|
||||
the performant C function rather than the polyfill handler.
|
||||
|
||||
The polyfill handler is a function that will be called in place of the original function when
|
||||
inlining the original function. The polyfill handler should have the same signature and the same
|
||||
behavior as the original function.
|
||||
|
||||
Args:
|
||||
original_fn (callable): The original function, usually a C function, to register a polyfill
|
||||
handler for.
|
||||
can_constant_fold_through (bool, optional): Whether the polyfill handler can be constant
|
||||
folded through. That is, if the polyfill handler is a pure function and its arguments
|
||||
are constant, the result of the polyfill handler can be constant folded during the
|
||||
compilation. Defaults to ``False``.
|
||||
skip_signature_check (bool, optional): Whether to skip the signature check between the
|
||||
original function and the polyfill handler. Defaults to ``False``.
|
||||
|
||||
Returns:
|
||||
A decorator that registers the polyfill handler for the original function.
|
||||
|
||||
Example::
|
||||
|
||||
>>> import operator
|
||||
>>> operator.indexOf([1, 2, 3, 4, 5], 3)
|
||||
2
|
||||
>>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3)
|
||||
... # xdoctest: +SKIP("Long tracebacks")
|
||||
Traceback (most recent call last):
|
||||
...
|
||||
torch._dynamo.exc.Unsupported: ...
|
||||
|
||||
>>> @torch.compiler.substitute_in_graph(operator.indexOf)
|
||||
... def indexOf(a, b, /):
|
||||
... for i, item in enumerate(a):
|
||||
... if item is b or item == b:
|
||||
... return i
|
||||
... raise ValueError("sequence.index(x): x not in sequence")
|
||||
>>>
|
||||
>>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3)
|
||||
2
|
||||
"""
|
||||
import torch._dynamo
|
||||
|
||||
return torch._dynamo.substitute_in_graph(
|
||||
original_fn,
|
||||
can_constant_fold_through=can_constant_fold_through,
|
||||
skip_signature_check=skip_signature_check,
|
||||
)
|
||||
|
||||
|
||||
def list_backends(exclude_tags=("debug", "experimental")) -> list[str]:
|
||||
"""
|
||||
Return valid strings that can be passed to `torch.compile(..., backend="name")`.
|
||||
|
||||
Args:
|
||||
exclude_tags(optional): A tuple of strings representing tags to exclude.
|
||||
"""
|
||||
import torch._dynamo
|
||||
|
||||
return torch._dynamo.list_backends(exclude_tags)
|
||||
|
||||
|
||||
def assume_constant_result(fn):
|
||||
"""
|
||||
This function is used to mark a function `fn` as having a constant result.
|
||||
This allows the compiler to optimize away your function.
|
||||
Returns The same function `fn`
|
||||
|
||||
Args:
|
||||
fn: The function to be marked as having a constant result.
|
||||
|
||||
.. warning::
|
||||
`assume_constant_result` can if invalid cause safety and soundness issues, :func:`torch.compile`
|
||||
will not attempt to validate whether the constant assumption is true or not
|
||||
|
||||
"""
|
||||
import torch._dynamo
|
||||
|
||||
return torch._dynamo.assume_constant_result(fn)
|
||||
|
||||
|
||||
def disable(fn=None, recursive=True, *, reason=None):
|
||||
"""
|
||||
This function provides a decorator to disable compilation on a function.
|
||||
It also provides the option of recursively disabling called functions.
|
||||
|
||||
Args:
|
||||
fn (optional): The function to disable
|
||||
recursive (optional): A boolean value indicating whether the disabling should be recursive.
|
||||
reason (optional): A string value indicating the reason for disabling the function.
|
||||
"""
|
||||
import torch._dynamo
|
||||
|
||||
return torch._dynamo.disable(fn, recursive, reason=reason)
|
||||
|
||||
|
||||
def set_stance(
|
||||
stance: str = "default",
|
||||
*,
|
||||
skip_guard_eval_unsafe: bool = False,
|
||||
force_backend: str | Callable[..., Any] | None = None,
|
||||
):
|
||||
"""
|
||||
Set the current stance of the compiler.
|
||||
Can be used as a function, context manager, or decorator.
|
||||
Do not use this function inside a `torch.compile` region - an error will be raised otherwise.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@torch.compile
|
||||
def foo(x): ...
|
||||
|
||||
|
||||
@torch.compiler.set_stance("force_eager")
|
||||
def bar():
|
||||
# will not be compiled
|
||||
foo(...)
|
||||
|
||||
|
||||
bar()
|
||||
|
||||
with torch.compiler.set_stance("force_eager"):
|
||||
# will also not be compiled
|
||||
foo(...)
|
||||
|
||||
torch.compiler.set_stance("force_eager")
|
||||
# will also not be compiled
|
||||
foo(...)
|
||||
torch.compiler.set_stance("default")
|
||||
|
||||
# will be compiled
|
||||
foo(...)
|
||||
|
||||
Args:
|
||||
stance: The stance to set the compiler to. Valid values are:
|
||||
|
||||
- "default": The default stance, used for normal compilation.
|
||||
- "force_eager": Ignore all `torch.compile` directives.
|
||||
- "eager_on_recompile": Run code eagerly when a recompile is necessary.
|
||||
If there is cached compiled code valid for the input, it will still be used.
|
||||
- "fail_on_recompile": Raise an error when recompiling a function.
|
||||
- "eager_then_compile": Run the first invocation in eager mode, then compile on
|
||||
subsequent calls. This is beneficial for dynamic shapes as it allows inferring
|
||||
dynamism from the first two invocations instead of wasting a static compile on
|
||||
the first invocation.
|
||||
- "aot_eager_then_compile": Run the first invocation with AOT eager to get memory
|
||||
benefits from activation checkpointing, then compile on subsequent calls. Like
|
||||
eager_then_compile, this improves handling of dynamic shapes by avoiding an
|
||||
initial static compile.
|
||||
|
||||
|
||||
skip_guard_eval_unsafe: A flag to run only differentiating guards.
|
||||
CAUTION - This flag is unsafe and should only be used if your setup
|
||||
meets the following conditions.
|
||||
|
||||
torch.compile uses a guard system to support recompilations and
|
||||
choose which compiled artifact to run at runtime. These guards,
|
||||
though efficient, add some overhead, which may impact performance in
|
||||
scenarios where you need to optimize for minimal guard processing
|
||||
time. This API enables you to disable guard evaluation, assuming
|
||||
that you have warmed up the compiled model with a sufficient variety
|
||||
of inputs. This assumption means that, after the warmup phase, no
|
||||
further recompilations will be necessary. If this assumption fails,
|
||||
there is a risk of silently producing incorrect results (hence the
|
||||
term "unsafe" in the API name).
|
||||
|
||||
force_backend: If `stance` is "default", this argument can be used to force `torch.compile`
|
||||
to use a specific backend. Otherwise, an error is raised.
|
||||
"""
|
||||
import torch._dynamo
|
||||
|
||||
return torch._dynamo.set_stance(
|
||||
stance,
|
||||
skip_guard_eval_unsafe=skip_guard_eval_unsafe,
|
||||
force_backend=force_backend,
|
||||
)
|
||||
|
||||
|
||||
# forbid in graph
|
||||
set_stance._dynamo_forbidden = True # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def set_enable_guard_collectives(enabled: bool):
|
||||
"""
|
||||
Enables use of collectives *during* guard evaluation to synchronize behavior
|
||||
across ranks. This is expensive: we have to issue a collective every time
|
||||
we enter a compiled code region, even if no rank actually would need to
|
||||
compile. This can help prevent NCCL hangs by ensuring that we never have a
|
||||
situation where one rank starts recompiling while other ranks don't compile;
|
||||
it is especially useful in conjunction with enable_compiler_collectives
|
||||
where such a situation would immediately cause a hang (as it is necessary
|
||||
for all ranks to compile at the same time to run compiler collectives). Like
|
||||
compiler collectives, you can only run this on SPMD programs; you will hang
|
||||
otherwise. Note that a guard collective is only issued if there is any
|
||||
compiled code to guard on; if this the first time we encounter a frame or
|
||||
the frame is skipped, we don't issue collectives.
|
||||
|
||||
Returns the previous setting of enabled.
|
||||
"""
|
||||
from torch._C._dynamo.eval_frame import set_guard_complete_hook # noqa: F401
|
||||
from torch._dynamo.eval_frame import guard_collectives_hook
|
||||
|
||||
if enabled:
|
||||
return set_guard_complete_hook(guard_collectives_hook) is not None # type: ignore[arg-type]
|
||||
else:
|
||||
return set_guard_complete_hook(None) is not None
|
||||
|
||||
|
||||
set_enable_guard_collectives._dynamo_forbidden = True # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def cudagraph_mark_step_begin():
|
||||
"""
|
||||
Indicates that a new iteration of inference or training is about to begin.
|
||||
|
||||
CUDA Graphs will free tensors of a prior iteration. A new iteration is started on each invocation of
|
||||
torch.compile, so long as there is not a pending backward that has not been called.
|
||||
|
||||
If that heuristic is wrong, such as in the following example, manually mark it with this api.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@torch.compile(mode="reduce-overhead")
|
||||
def rand_foo():
|
||||
return torch.rand([4], device="cuda")
|
||||
|
||||
|
||||
for _ in range(5):
|
||||
torch.compiler.cudagraph_mark_step_begin()
|
||||
rand_foo() + rand_foo()
|
||||
|
||||
For more details, see `torch.compiler_cudagraph_trees <https://docs.pytorch.org/docs/main/user_guide/torch_compiler/torch.compiler_cudagraph_trees.html>`__ # noqa: B950
|
||||
"""
|
||||
from torch._inductor import cudagraph_trees
|
||||
|
||||
cudagraph_trees.mark_step_begin()
|
||||
|
||||
|
||||
def wrap_numpy(fn):
|
||||
r"""Decorator that turns a function from ``np.ndarray``s to ``np.ndarray``s into a function
|
||||
from ``torch.Tensor``s to ``torch.Tensor``s.
|
||||
|
||||
It is designed to be used with :func:`torch.compile` with ``fullgraph=True``. It allows to
|
||||
compile a NumPy function as if it were a PyTorch function. This allows you to run NumPy code
|
||||
on CUDA or compute its gradients.
|
||||
|
||||
.. note::
|
||||
|
||||
This decorator does not work without :func:`torch.compile`.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> # Compile a NumPy function as a Tensor -> Tensor function
|
||||
>>> @torch.compile(fullgraph=True)
|
||||
>>> @torch.compiler.wrap_numpy
|
||||
>>> def fn(a: np.ndarray):
|
||||
>>> return np.sum(a * a)
|
||||
>>> # Execute the NumPy function using Tensors on CUDA and compute the gradients
|
||||
>>> x = torch.arange(6, dtype=torch.float32, device="cuda", requires_grad=True)
|
||||
>>> out = fn(x)
|
||||
>>> out.backward()
|
||||
>>> print(x.grad)
|
||||
tensor([ 0., 2., 4., 6., 8., 10.], device='cuda:0')
|
||||
"""
|
||||
from torch._dynamo.external_utils import wrap_numpy as wrap
|
||||
|
||||
return wrap(fn)
|
||||
|
||||
|
||||
_is_compiling_flag: bool = False
|
||||
_is_exporting_flag: bool = False
|
||||
_is_non_strict_tracing_flag: bool = False
|
||||
|
||||
|
||||
def is_compiling() -> bool:
|
||||
"""
|
||||
Indicates whether a graph is executed/traced as part of torch.compile() or torch.export().
|
||||
|
||||
Note that there are 2 other related flags that should deprecated eventually:
|
||||
* torch._dynamo.external_utils.is_compiling()
|
||||
* torch._utils.is_compiling()
|
||||
|
||||
Example::
|
||||
|
||||
>>> def forward(self, x):
|
||||
>>> if not torch.compiler.is_compiling():
|
||||
>>> pass # ...logic that is not needed in a compiled/traced graph...
|
||||
>>>
|
||||
>>> # ...rest of the function...
|
||||
"""
|
||||
if torch.jit.is_scripting():
|
||||
return False
|
||||
else:
|
||||
return _is_compiling_flag
|
||||
|
||||
|
||||
def _is_non_strict_tracing() -> bool:
|
||||
"""
|
||||
Indicates whether we are inside a non-strict make_fx-based tracing session.
|
||||
"""
|
||||
return _is_non_strict_tracing_flag
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _non_strict_tracing_context():
|
||||
"""Context manager that sets the non-strict tracing flag."""
|
||||
global _is_non_strict_tracing_flag
|
||||
old = _is_non_strict_tracing_flag
|
||||
try:
|
||||
_is_non_strict_tracing_flag = True
|
||||
yield
|
||||
finally:
|
||||
_is_non_strict_tracing_flag = old
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _patch_autograd_grad():
|
||||
"""Patch autograd.grad for non-strict make_fx tracing.
|
||||
|
||||
This patch installs autograd hooks so traced backward nodes preserve
|
||||
stack trace, seq_nr, and autograd_backward metadata before delegating to
|
||||
the real torch.autograd.grad.
|
||||
"""
|
||||
import functools
|
||||
|
||||
import torch.autograd
|
||||
from torch._functorch._aot_autograd.logging_utils import (
|
||||
setup_stacktrace_preservation_hooks_from_tensors,
|
||||
)
|
||||
|
||||
_orig_grad = torch.autograd.grad
|
||||
|
||||
@functools.wraps(_orig_grad)
|
||||
def _patched_grad(outputs, inputs, *args, **kwargs):
|
||||
if not _is_non_strict_tracing():
|
||||
raise AssertionError(
|
||||
"_patch_autograd_grad() must be used under "
|
||||
"_non_strict_tracing_context()"
|
||||
)
|
||||
|
||||
setup_stacktrace_preservation_hooks_from_tensors(outputs)
|
||||
return _orig_grad(outputs, inputs, *args, **kwargs)
|
||||
|
||||
torch.autograd.grad = _patched_grad
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.autograd.grad = _orig_grad
|
||||
|
||||
|
||||
def is_dynamo_compiling() -> bool:
|
||||
"""
|
||||
Indicates whether a graph is traced via TorchDynamo.
|
||||
|
||||
It's stricter than is_compiling() flag, as it would only be set to True when
|
||||
TorchDynamo is used.
|
||||
|
||||
Example::
|
||||
|
||||
>>> def forward(self, x):
|
||||
>>> if not torch.compiler.is_dynamo_compiling():
|
||||
>>> pass # ...logic that is not needed in a TorchDynamo-traced graph...
|
||||
>>>
|
||||
>>> # ...rest of the function...
|
||||
"""
|
||||
return False
|
||||
|
||||
|
||||
def is_exporting() -> bool:
|
||||
"""
|
||||
Indicated whether we're under exporting.
|
||||
|
||||
It's stricter than is_compiling() flag, as it would only be set to True when
|
||||
torch.export is used.
|
||||
|
||||
Example::
|
||||
|
||||
>>> def forward(self, x):
|
||||
>>> if not torch.compiler.is_exporting():
|
||||
>>> pass # ...logic that is not needed in export...
|
||||
>>>
|
||||
>>> # ...rest of the function...
|
||||
"""
|
||||
return _is_exporting_flag
|
||||
|
||||
|
||||
def save_cache_artifacts() -> tuple[bytes, CacheInfo] | None:
|
||||
"""
|
||||
Serializes all the cache artifacts that were created during the compilation
|
||||
|
||||
Example:
|
||||
|
||||
- Execute torch.compile
|
||||
- Call torch.compiler.save_cache_artifacts()
|
||||
"""
|
||||
from ._cache import CacheArtifactManager
|
||||
|
||||
if torch._dynamo.config.caching_precompile:
|
||||
from torch._dynamo.precompile_context import PrecompileContext
|
||||
|
||||
PrecompileContext.save_to_dynamo_cache()
|
||||
|
||||
return CacheArtifactManager.serialize()
|
||||
|
||||
|
||||
def load_cache_artifacts(serialized_artifacts: bytes) -> CacheInfo | None:
|
||||
"""
|
||||
Hot loads cache artifacts that were previously serialized via
|
||||
save_cache_artifacts
|
||||
|
||||
Example:
|
||||
|
||||
# From a previous invocation
|
||||
artifacts = torch.compiler.save_cache_artifacts()
|
||||
|
||||
torch.compiler.load_cache_artifacts(artifacts[0])
|
||||
"""
|
||||
from ._cache import CacheArtifactManager, CacheInfo
|
||||
|
||||
artifacts = CacheArtifactManager.deserialize(serialized_artifacts)
|
||||
if artifacts is not None:
|
||||
return CacheArtifactManager.populate_caches(artifacts)
|
||||
return None
|
||||
|
||||
|
||||
def keep_portable_guards_unsafe(guard_entries):
|
||||
"""
|
||||
A common function to only keep guards that can be used in both Python and non-Python environments.
|
||||
This includes:
|
||||
- Tensor metadata and dynamic shape information.
|
||||
- Global contexts state (e.g. autocast, no_grad, etc.)
|
||||
|
||||
This is unsafe to use by default.
|
||||
To use this API, use guard_filter_fn argument while calling torch.compile
|
||||
|
||||
>> opt_mod = torch.compile(
|
||||
>> mod,
|
||||
>> options={"guard_filter_fn": torch.compiler.keep_global_context_and_tensor_guards_unsafe},
|
||||
>> )
|
||||
"""
|
||||
return [
|
||||
(
|
||||
g.guard_type in ("GLOBAL_STATE", "SHAPE_ENV")
|
||||
or (g.guard_type == "TENSOR_MATCH" and not g.is_global)
|
||||
)
|
||||
for g in guard_entries
|
||||
]
|
||||
|
||||
|
||||
def skip_guard_on_inbuilt_nn_modules_unsafe(guard_entries):
|
||||
"""
|
||||
A common function to skip guards on the inbuilt nn modules like
|
||||
torch.nn.Linear. This is unsafe to use by default. But for majority of
|
||||
torch.compile users, the model code does not modify the inbuilt nn module
|
||||
attributes. They can benefit from reduction in guard latency overhead using
|
||||
this API.
|
||||
|
||||
To use this API, use guard_filter_fn argument while calling torch.compile
|
||||
|
||||
>> opt_mod = torch.compile(
|
||||
>> mod,
|
||||
>> options={"guard_filter_fn": torch.compiler.skip_guard_on_all_nn_modules_unsafe},
|
||||
>> )
|
||||
"""
|
||||
return [
|
||||
not entry.orig_guard.source.is_unspecialized_builtin_nn_module()
|
||||
for entry in guard_entries
|
||||
]
|
||||
|
||||
|
||||
def skip_guard_on_all_nn_modules_unsafe(guard_entries):
|
||||
"""
|
||||
A common function to skip guards on all nn modules, both user defined as
|
||||
well inbuilt nn modules (like torch.nn.Linear). This is unsafe to use by
|
||||
default. But for majority of torch.compile users, the model code does not
|
||||
modify the nn module attributes. They can benefit from reduction in guard
|
||||
latency overhead using this API.
|
||||
|
||||
To use this API, use guard_filter_fn argument while calling torch.compile
|
||||
|
||||
>> opt_mod = torch.compile(
|
||||
>> mod,
|
||||
>> options={"guard_filter_fn": torch.compiler.skip_guard_on_all_nn_modules_unsafe},
|
||||
>> )
|
||||
"""
|
||||
|
||||
return [
|
||||
not entry.orig_guard.source.is_unspecialized_nn_module()
|
||||
for entry in guard_entries
|
||||
]
|
||||
|
||||
|
||||
def keep_tensor_guards_unsafe(guard_entries, keep_parameters=False):
|
||||
"""
|
||||
A common function to keep tensor guards on all tensors. This is unsafe to
|
||||
use by default. But if you don't expect any changes in the model code, you
|
||||
can just keep the tensor guards.
|
||||
|
||||
|
||||
>> opt_mod = torch.compile(
|
||||
>> mod,
|
||||
>> options={"guard_filter_fn": torch.compiler.keep_tensor_guards},
|
||||
>> )
|
||||
"""
|
||||
|
||||
keep_flags = []
|
||||
for entry in guard_entries:
|
||||
if entry.guard_type == "TENSOR_MATCH":
|
||||
if not isinstance(entry.value, torch.nn.Parameter):
|
||||
keep_flags.append(True)
|
||||
elif keep_parameters:
|
||||
keep_flags.append(True)
|
||||
else:
|
||||
keep_flags.append(False)
|
||||
else:
|
||||
keep_flags.append(False)
|
||||
return keep_flags
|
||||
|
||||
|
||||
def skip_guard_on_globals_unsafe(guard_entries):
|
||||
"""
|
||||
A common function to skip guards on all globals. This is unsafe to use by
|
||||
default. But if you don't expect any changes in the globals, you can just
|
||||
keep the tensor guards.
|
||||
|
||||
>> opt_mod = torch.compile(
|
||||
>> mod,
|
||||
>> options={"guard_filter_fn": torch.compiler.skip_guard_on_globals},
|
||||
>> )
|
||||
"""
|
||||
|
||||
return [not entry.is_global for entry in guard_entries]
|
||||
|
||||
|
||||
def skip_all_guards_unsafe(guard_entries):
|
||||
"""
|
||||
A function for skipping all guards on a compiled function.
|
||||
|
||||
WARNING: This function will drop all the safety guarantees from Dynamo
|
||||
compiled function. Use this with caution.
|
||||
|
||||
To use this API, use guard_filter_fn argument while calling torch.compile
|
||||
|
||||
>> opt_mod = torch.compile(
|
||||
>> mod,
|
||||
>> options={"guard_filter_fn": torch.compiler.skip_all_guards_unsafe},
|
||||
>> )
|
||||
"""
|
||||
return [False for entry in guard_entries]
|
||||
|
||||
|
||||
def nested_compile_region(
|
||||
fn=None,
|
||||
*,
|
||||
options: NestedCompileRegionOptions | None = None,
|
||||
max_reuse_entries: int = 8,
|
||||
reuse_hash_fn=None,
|
||||
):
|
||||
"""
|
||||
Tells **``torch.compile``** that the marked set of operations forms a nested
|
||||
compile region (which is often repeated in the full model) whose code can be
|
||||
compiled once and safely reused. ``nested_compile_region`` can also be used
|
||||
as a decorator.
|
||||
|
||||
During **``torch.compile``** tracing, the compiler applies *hierarchical
|
||||
compilation* with ``nested_compile_region``: it emits optimized code for the
|
||||
marked region the first time it is encountered and re-emits (or "stamps
|
||||
out") the previously compiled code on every subsequent invocation. This can
|
||||
substantially reduce overall compile time for deeply-stacked,
|
||||
structurally-identical components such as the transformer layers of a
|
||||
large-language-model (LLM).
|
||||
|
||||
Outside a ``torch.compile`` context—i.e., in standard eager execution—the
|
||||
call is a no-op, so existing workflows remain unaffected.
|
||||
|
||||
Note that ``nested_compile_region`` **does not** promise that a region will
|
||||
be compiled exactly once. If the compiler detects that new input conditions
|
||||
(shape, dtype, device, stride, globals etc.) make the cached version invalid
|
||||
to reuse, it will transparently re-compile the region. Using it is
|
||||
therefore *safe*: correctness is always preserved, and you pay the extra
|
||||
compilation cost only when required.
|
||||
|
||||
Args:
|
||||
fn: The function to wrap
|
||||
options: Optional backend to use for compiling the subgraph.
|
||||
Warning: this is an experimental feature under development and
|
||||
not ready for use yet.
|
||||
max_reuse_entries: Maximum number of reuse cache entries per function
|
||||
before raising an error. If this limit is hit, guards keep failing
|
||||
across invocations and hierarchical compilation is not effective.
|
||||
reuse_hash_fn: Optional callable that takes the same ``*args, **kwargs``
|
||||
as the wrapped function and returns an integer hash key. When
|
||||
provided, Dynamo traces this function to obtain a constant integer
|
||||
and uses it as the cache key for subgraph reuse, bypassing the
|
||||
automatic fingerprint/guard machinery. Two calls that produce the
|
||||
same hash key reuse the same cached subgraph. The hash function
|
||||
must be fully traceable (no graph breaks) and must return a
|
||||
constant integer.
|
||||
"""
|
||||
|
||||
if options is not None:
|
||||
from torch._dynamo import config as dynamo_config
|
||||
|
||||
if not dynamo_config.enable_invoke_subgraph_regional_compile:
|
||||
raise RuntimeError(
|
||||
"nested_compile_region config is an experimental feature for testing only."
|
||||
)
|
||||
|
||||
from torch._higher_order_ops.invoke_subgraph import (
|
||||
mark_compile_region as _mark_compile_region,
|
||||
)
|
||||
|
||||
return _mark_compile_region(
|
||||
fn,
|
||||
options=options,
|
||||
max_reuse_entries=max_reuse_entries,
|
||||
reuse_hash_fn=reuse_hash_fn,
|
||||
)
|
||||
|
||||
|
||||
def load_compiled_function(
|
||||
file: io.IOBase,
|
||||
*,
|
||||
f_globals: dict[str, object] | None = None,
|
||||
external_data: dict[str, Any] | None = None,
|
||||
) -> Callable[..., Any]:
|
||||
"""
|
||||
Load an aot-compiled function from a file.
|
||||
|
||||
.. warning::
|
||||
|
||||
This API is currently experimental and subject to change.
|
||||
|
||||
Args:
|
||||
file: A file-like object containing the serialized compiled function.
|
||||
f_globals: Optional global scope enclosing the compiled function.
|
||||
external_data: Optional data to be loaded into the runtime environment
|
||||
of the compiled function. This should contains the same
|
||||
data as AOTCompileResult.external_data returned from save_compiled_function() call.
|
||||
|
||||
Returns:
|
||||
A torch-compiled function with compilation preloaded from disk.
|
||||
"""
|
||||
from torch._dynamo.aot_compile import AOTCompiledFunction
|
||||
|
||||
data = file.read()
|
||||
return AOTCompiledFunction.deserialize(data, f_globals, external_data)
|
||||
@@ -0,0 +1,325 @@
|
||||
import copy
|
||||
import dataclasses
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import defaultdict
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from itertools import chain
|
||||
from typing import Any
|
||||
|
||||
from torch.utils._appending_byte_serializer import (
|
||||
AppendingByteSerializer,
|
||||
BytesReader,
|
||||
BytesWriter,
|
||||
)
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class CacheArtifact(ABC):
|
||||
"""
|
||||
Data for each cache artifact that will be serialized and deserialized
|
||||
"""
|
||||
|
||||
key: str
|
||||
content: bytes = dataclasses.field(repr=False) # Do not display potential binary
|
||||
|
||||
@staticmethod
|
||||
def serialize(writer: BytesWriter, cls: "CacheArtifact") -> None:
|
||||
writer.write_str(cls.key)
|
||||
writer.write_bytes(cls.content)
|
||||
|
||||
@staticmethod
|
||||
def deserialize(artifact_type: str, reader: BytesReader) -> "CacheArtifact":
|
||||
key = reader.read_str()
|
||||
content = reader.read_bytes()
|
||||
return CacheArtifactFactory.create(artifact_type, key, content)
|
||||
|
||||
@staticmethod
|
||||
def encode(content: Any) -> bytes:
|
||||
if not isinstance(content, bytes):
|
||||
raise AssertionError(f"Expected bytes, got {type(content)}")
|
||||
return content
|
||||
|
||||
@abstractmethod
|
||||
def populate_cache(self) -> None:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def type() -> str:
|
||||
"""
|
||||
Returns the type of the artifact. Must be unique across all CacheArtifact classes.
|
||||
|
||||
CacheArtifactFactory.register will add property method to CacheInfo based on this (def {type}_artifacts)
|
||||
that returns all artifacts for specific cache.
|
||||
"""
|
||||
raise RuntimeError("CacheArtifact is an abstract class, please use a subclass")
|
||||
|
||||
|
||||
class CacheArtifactFactory:
|
||||
"""
|
||||
Factory for creating CacheArtifact objects based on their type
|
||||
"""
|
||||
|
||||
_artifact_types: dict[str, type[CacheArtifact]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, artifact_cls: type[CacheArtifact]) -> type[CacheArtifact]:
|
||||
artifact_type_key = artifact_cls.type()
|
||||
if artifact_cls.type() in cls._artifact_types:
|
||||
raise AssertionError(
|
||||
f"Artifact of type={artifact_type_key} already registered in mega-cache artifact factory"
|
||||
)
|
||||
cls._artifact_types[artifact_type_key] = artifact_cls
|
||||
setattr(
|
||||
CacheInfo,
|
||||
f"{artifact_type_key}_artifacts",
|
||||
property(lambda self: self.artifacts[artifact_type_key]),
|
||||
)
|
||||
return artifact_cls
|
||||
|
||||
@classmethod
|
||||
def _get_artifact_type(cls, artifact_type_key: str) -> type[CacheArtifact]:
|
||||
if artifact_type_key not in cls._artifact_types:
|
||||
raise AssertionError(
|
||||
f"Artifact of type={artifact_type_key} not registered in mega-cache artifact factory"
|
||||
)
|
||||
return cls._artifact_types[artifact_type_key]
|
||||
|
||||
@classmethod
|
||||
def create(cls, artifact_type_key: str, key: str, content: bytes) -> CacheArtifact:
|
||||
artifact_cls = cls._get_artifact_type(artifact_type_key)
|
||||
# pyrefly: ignore [bad-instantiation]
|
||||
return artifact_cls(key, content)
|
||||
|
||||
@classmethod
|
||||
def encode_create(
|
||||
cls, artifact_type_key: str, key: str, content: Any
|
||||
) -> CacheArtifact:
|
||||
artifact_cls = cls._get_artifact_type(artifact_type_key)
|
||||
# pyrefly: ignore [bad-instantiation]
|
||||
return artifact_cls(key, artifact_cls.encode(content))
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class CacheInfo:
|
||||
"""
|
||||
Return value of serialization and deserialization for the purpose of
|
||||
instrumentation
|
||||
"""
|
||||
|
||||
artifacts: defaultdict[str, list[str]] = dataclasses.field(
|
||||
default_factory=lambda: defaultdict(list)
|
||||
)
|
||||
|
||||
# Methods set by CacheArtifactFactory.register based on CacheArtifact.type()
|
||||
@property
|
||||
def inductor_artifacts(self) -> list[str]: # type: ignore[empty-body]
|
||||
...
|
||||
|
||||
@property
|
||||
def autotune_artifacts(self) -> list[str]: # type: ignore[empty-body]
|
||||
...
|
||||
|
||||
@property
|
||||
def aot_autograd_artifacts(self) -> list[str]: # type: ignore[empty-body]
|
||||
...
|
||||
|
||||
@property
|
||||
def pgo_artifacts(self) -> list[str]: # type: ignore[empty-body]
|
||||
...
|
||||
|
||||
@property
|
||||
def precompile_artifacts(self) -> list[str]: # type: ignore[empty-body]
|
||||
...
|
||||
|
||||
def add(self, artifact: CacheArtifact) -> None:
|
||||
self.artifacts[artifact.type()].append(artifact.key)
|
||||
|
||||
def clear(self) -> None:
|
||||
self.artifacts.clear()
|
||||
|
||||
def empty(self) -> bool:
|
||||
return not self.artifacts
|
||||
|
||||
|
||||
def _serialize_single_cache(
|
||||
writer: BytesWriter, cls: "tuple[str, list[CacheArtifact]]"
|
||||
) -> None:
|
||||
writer.write_str(cls[0])
|
||||
writer.write_uint64(len(cls[1]))
|
||||
for artifact in cls[1]:
|
||||
CacheArtifact.serialize(writer, artifact)
|
||||
|
||||
|
||||
def _deserialize_single_cache(
|
||||
reader: BytesReader,
|
||||
) -> "tuple[str, list[CacheArtifact]]":
|
||||
artifacts = []
|
||||
artifact_type_key = reader.read_str()
|
||||
num_artifacts = reader.read_uint64()
|
||||
for _ in range(num_artifacts):
|
||||
artifacts.append(CacheArtifact.deserialize(artifact_type_key, reader))
|
||||
|
||||
return artifact_type_key, artifacts
|
||||
|
||||
|
||||
CacheArtifactsResult = dict[str, list[CacheArtifact]]
|
||||
|
||||
|
||||
class CacheArtifactManager:
|
||||
"""
|
||||
Lightweight manager class for collecting and processing cache artifacts for
|
||||
hot loading
|
||||
|
||||
Intended Lifecycle:
|
||||
- Execute code via torch.compile, this will call
|
||||
CacheArtifactManager.record_artifact on each cache artifact
|
||||
- Call CacheArtifactManager.serialize to convert all the cache artifacts
|
||||
to portable format
|
||||
- Call CacheArtifactManager.deserialize to hot load the cache artifacts on
|
||||
a potentially different process
|
||||
|
||||
NOTE: There's no FB/FC guarantees, results of cache artifacts will not be
|
||||
used unless code version matches.
|
||||
"""
|
||||
|
||||
# Protected by the compile_lock
|
||||
_new_cache_artifacts: CacheArtifactsResult = defaultdict(list)
|
||||
# Keep a separate seen artifacts list to make avoid unnecessary duplicates
|
||||
# This list will not be cleared between serialize() calls
|
||||
_seen_artifacts: OrderedSet[CacheArtifact] = OrderedSet()
|
||||
# When serialize() is called, artifacts are transferred from _cache_artifacts to
|
||||
# internal data structure of the _serializer
|
||||
# This allows us to only pay the cost of serialization if serialize() is called
|
||||
_serializer: AppendingByteSerializer[tuple[str, list[CacheArtifact]]] = (
|
||||
AppendingByteSerializer(serialize_fn=_serialize_single_cache)
|
||||
)
|
||||
_cache_info: CacheInfo = CacheInfo()
|
||||
|
||||
@classmethod
|
||||
def clear(cls) -> None:
|
||||
cls._new_cache_artifacts.clear()
|
||||
cls._seen_artifacts.clear()
|
||||
cls._serializer.clear()
|
||||
cls._cache_info.clear()
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def with_fresh_cache(cls) -> Generator[None, None, None]:
|
||||
original_new_cache_artifacts = cls._new_cache_artifacts
|
||||
original_seen_artifacts = cls._seen_artifacts
|
||||
original_serializer = cls._serializer
|
||||
original_cache_info = cls._cache_info
|
||||
|
||||
cls._new_cache_artifacts = defaultdict(list)
|
||||
cls._seen_artifacts = OrderedSet()
|
||||
cls._serializer = AppendingByteSerializer(serialize_fn=_serialize_single_cache)
|
||||
cls._cache_info = cls._cache_info.__class__()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
cls._new_cache_artifacts = original_new_cache_artifacts
|
||||
cls._seen_artifacts = original_seen_artifacts
|
||||
cls._serializer = original_serializer
|
||||
cls._cache_info = original_cache_info
|
||||
|
||||
@classmethod
|
||||
def record_artifact(
|
||||
cls,
|
||||
artifact_type: str,
|
||||
key: str,
|
||||
content: Any,
|
||||
) -> None:
|
||||
"""
|
||||
Called from each caching operation to record the artifact in this
|
||||
"mega" list
|
||||
"""
|
||||
artifact = CacheArtifactFactory.encode_create(artifact_type, key, content)
|
||||
if artifact in cls._seen_artifacts:
|
||||
return
|
||||
log.debug("Recording %s", artifact)
|
||||
cls._new_cache_artifacts[artifact_type].append(artifact)
|
||||
cls._seen_artifacts.add(artifact)
|
||||
|
||||
@classmethod
|
||||
def need_serialize(cls) -> bool:
|
||||
"""
|
||||
Have we seen new artifacts since last serialize call?
|
||||
"""
|
||||
return len(cls._new_cache_artifacts) != 0
|
||||
|
||||
@classmethod
|
||||
def serialize(cls) -> tuple[bytes, CacheInfo] | None:
|
||||
"""
|
||||
Converts the "mega" list into portable format
|
||||
"""
|
||||
for artifact in chain(*cls._new_cache_artifacts.values()):
|
||||
log.debug("saving: %s", artifact)
|
||||
cls._cache_info.add(artifact)
|
||||
|
||||
if cls._cache_info.empty():
|
||||
# If there are not artifacts, dont just return bytes with
|
||||
# version.
|
||||
return None
|
||||
|
||||
try:
|
||||
# We deep copy cls._cache_info since later compilations
|
||||
# can keep adding to cache_info
|
||||
info = copy.deepcopy(cls._cache_info)
|
||||
cls._serializer.extend(cls._new_cache_artifacts.items())
|
||||
artifact_bytes = cls._serializer.to_bytes()
|
||||
cls._new_cache_artifacts.clear()
|
||||
return artifact_bytes, info
|
||||
except Exception:
|
||||
log.warning("Failed to pickle cache artifacts", exc_info=True)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def deserialize(serialized_artifacts: bytes) -> CacheArtifactsResult | None:
|
||||
"""
|
||||
Converts the portable format back into CacheArtifacts
|
||||
"""
|
||||
try:
|
||||
CacheArtifactManager._ensure_cache_artifacts_registered()
|
||||
artifacts = dict(
|
||||
AppendingByteSerializer.to_list(
|
||||
serialized_artifacts,
|
||||
deserialize_fn=_deserialize_single_cache,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
log.warning("Failed to un-pickle cache artifacts", exc_info=True)
|
||||
return None
|
||||
|
||||
return artifacts
|
||||
|
||||
@staticmethod
|
||||
def populate_caches(artifacts: CacheArtifactsResult) -> CacheInfo:
|
||||
info = CacheInfo()
|
||||
for artifact in chain(*artifacts.values()):
|
||||
log.debug("writing: %s", artifact)
|
||||
info.add(artifact)
|
||||
artifact.populate_cache()
|
||||
|
||||
return info
|
||||
|
||||
@classmethod
|
||||
def _ensure_cache_artifacts_registered(cls) -> None:
|
||||
"""When deserializing caches in fresh process, we need to ensure that all
|
||||
cache artifacts are registered in the cache registry. This is done by
|
||||
simply importing all the cache artifacts already wrapped with register call.
|
||||
"""
|
||||
from torch._dynamo.package import PrecompileCacheArtifact # noqa: F401
|
||||
from torch._dynamo.pgo import PGOCacheArtifact # noqa: F401
|
||||
from torch._functorch._aot_autograd.autograd_cache import ( # noqa: F401
|
||||
AOTAutogradCacheArtifact,
|
||||
)
|
||||
from torch._inductor.codecache import InductorCacheArtifact # noqa: F401
|
||||
from torch._inductor.runtime.autotune_cache import ( # noqa: F401
|
||||
AutotuneCacheArtifact,
|
||||
)
|
||||
@@ -0,0 +1,280 @@
|
||||
"""
|
||||
This is the top-level configuration module for the compiler, containing
|
||||
cross-cutting configuration options that affect all parts of the compiler
|
||||
stack.
|
||||
|
||||
You may also be interested in the per-component configuration modules, which
|
||||
contain configuration options that affect only a specific part of the compiler:
|
||||
|
||||
* :mod:`torch._dynamo.config`
|
||||
* :mod:`torch._inductor.config`
|
||||
* :mod:`torch._functorch.config`
|
||||
* :mod:`torch.fx.experimental.config`
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
from torch.utils._config_module import Config, install_config_module
|
||||
|
||||
|
||||
__all__ = [
|
||||
"job_id",
|
||||
"dynamic_shapes",
|
||||
"assume_static_by_default",
|
||||
"automatic_dynamic_shapes",
|
||||
"recompile_limit",
|
||||
"accumulated_recompile_limit",
|
||||
"verbose",
|
||||
"capture_scalar_outputs",
|
||||
"capture_dynamic_output_shape_ops",
|
||||
"log_file_name",
|
||||
"fail_on_recompile_limit_hit",
|
||||
"allow_unspec_int_on_nn_module",
|
||||
"skip_tensor_guards_with_matching_dict_tags",
|
||||
"enable_cpp_symbolic_shape_guards",
|
||||
"wrap_top_frame",
|
||||
"reorderable_logging_functions",
|
||||
"force_disable_caches",
|
||||
]
|
||||
|
||||
|
||||
# NB: Docblocks go UNDER variable definitions! Use spacing to make the
|
||||
# grouping clear.
|
||||
|
||||
# FB-internal note: you do NOT have to specify this explicitly specify this if
|
||||
# you run on MAST, we will automatically default this to
|
||||
# mast:MAST_JOB_NAME:MAST_JOB_VERSION.
|
||||
job_id: str | None = Config(
|
||||
env_name_default=["TORCH_COMPILE_JOB_ID", "TORCH_COMPILE_STICKY_PGO_KEY"],
|
||||
default=None,
|
||||
)
|
||||
"""
|
||||
Semantically, this should be an identifier that uniquely identifies, e.g., a
|
||||
training job. You might have multiple attempts of the same job, e.g., if it was
|
||||
preempted or needed to be restarted, but each attempt should be running
|
||||
substantially the same workload with the same distributed topology. You can
|
||||
set this by environment variable with :envvar:`TORCH_COMPILE_JOB_ID`.
|
||||
|
||||
Operationally, this controls the effect of profile-guided optimization related
|
||||
persistent state. PGO state can affect how we perform compilation across
|
||||
multiple invocations of PyTorch, e.g., the first time you run your program we
|
||||
may compile twice as we discover what inputs are dynamic, and then PGO will
|
||||
save this state so subsequent invocations only need to compile once, because
|
||||
they remember it is dynamic. This profile information, however, is sensitive
|
||||
to what workload you are running, so we require you to tell us that two jobs
|
||||
are *related* (i.e., are the same workload) before we are willing to reuse
|
||||
this information. Notably, PGO does nothing (even if explicitly enabled)
|
||||
unless a valid ``job_id`` is available. In some situations, PyTorch can
|
||||
configured to automatically compute a ``job_id`` based on the environment it
|
||||
is running in.
|
||||
|
||||
Profiles are always collected on a per rank basis, so different ranks may have
|
||||
different profiles. If you know your workload is truly SPMD, you can run with
|
||||
:data:`torch._dynamo.config.enable_compiler_collectives` to ensure nodes get
|
||||
consistent profiles across all ranks.
|
||||
"""
|
||||
|
||||
pgo_extra_read_key: str | None = Config(
|
||||
env_name_default="TORCH_COMPILE_STICKY_PGO_READ", default=None
|
||||
)
|
||||
pgo_extra_write_key: str | None = Config(
|
||||
env_name_default="TORCH_COMPILE_STICKY_PGO_WRITE", default=None
|
||||
)
|
||||
"""
|
||||
Additional read/write keys for PGO.
|
||||
Write key: Besides writing to the default local/remote PGO state, this also writes to the specified key.
|
||||
Read key: Besides reading from the default state, this also reads from the specified key (if written to before)
|
||||
and merges it with the default state.
|
||||
"""
|
||||
|
||||
|
||||
cache_key_tag: str = Config(env_name_default="TORCH_COMPILE_CACHE_KEY_TAG", default="")
|
||||
"""
|
||||
Tag to be included in the cache key generation for all torch compile caching.
|
||||
A common use case for such a tag is to break caches.
|
||||
"""
|
||||
|
||||
force_disable_caches: bool = Config(
|
||||
justknob="pytorch/remote_cache:force_disable_caches",
|
||||
env_name_force=[
|
||||
"TORCHINDUCTOR_FORCE_DISABLE_CACHES",
|
||||
"TORCH_COMPILE_FORCE_DISABLE_CACHES",
|
||||
],
|
||||
default=False,
|
||||
)
|
||||
"""
|
||||
Force disables all caching -- This will take precedence over and override any other caching flag
|
||||
"""
|
||||
|
||||
dynamic_sources: str = Config(
|
||||
env_name_default="TORCH_COMPILE_DYNAMIC_SOURCES", default=""
|
||||
)
|
||||
"""
|
||||
Comma delimited list of sources that should be marked as dynamic. Primarily useful for large
|
||||
models with graph breaks where you need intermediate tensors and ints to be marked dynamic.
|
||||
|
||||
This whitelist is dominant over all other flags dynamic=False, force_nn_module_property_static_shapes
|
||||
and force_parameter_static_shapes.
|
||||
"""
|
||||
|
||||
unbacked_sources: str = Config(
|
||||
env_name_default="TORCH_COMPILE_UNBACKED_SOURCES", default=""
|
||||
)
|
||||
"""
|
||||
Comma delimited list of sources that should be marked as unbacked. Primarily useful for large
|
||||
models with graph breaks where you need intermediate tensors marked unbacked.
|
||||
|
||||
This whitelist is dominant over all other flags dynamic=False, force_nn_module_property_static_shapes
|
||||
and force_parameter_static_shapes.
|
||||
"""
|
||||
|
||||
# force a python GC before recording cudagraphs
|
||||
force_cudagraph_gc: bool = Config(env_name_default="TORCH_CUDAGRAPH_GC", default=False)
|
||||
"""
|
||||
If True (the backward-compatible behavior) then gc.collect() before recording
|
||||
any cudagraph.
|
||||
"""
|
||||
|
||||
|
||||
# Cross-cutting configuration options that affect the entire compilation pipeline
|
||||
|
||||
dynamic_shapes: bool = Config(alias="torch._dynamo.config.dynamic_shapes")
|
||||
"""
|
||||
Controls whether the compilation pipeline supports dynamic tensor shapes.
|
||||
When enabled, the compiler can handle tensors with varying dimensions across
|
||||
different invocations. This is a cross-cutting setting that affects shape
|
||||
inference, guard generation, and code generation across the entire compilation
|
||||
stack.
|
||||
"""
|
||||
|
||||
assume_static_by_default: bool = Config(
|
||||
alias="torch._dynamo.config.assume_static_by_default"
|
||||
)
|
||||
"""
|
||||
When enabled, all tensor dimensions are assumed to be static unless explicitly
|
||||
marked as dynamic or detected as changing. This compilation-wide behavior affects
|
||||
how the entire stack handles shape specialization and can improve performance
|
||||
for static workloads.
|
||||
"""
|
||||
|
||||
automatic_dynamic_shapes: bool = Config(
|
||||
alias="torch._dynamo.config.automatic_dynamic_shapes"
|
||||
)
|
||||
"""
|
||||
Enables automatic detection and handling of dynamic shapes. When a tensor's
|
||||
shape changes between compilations, the system automatically marks those
|
||||
dimensions as dynamic rather than requiring manual specification. This
|
||||
cross-cutting optimization improves the user experience by reducing recompilations.
|
||||
"""
|
||||
|
||||
recompile_limit: int = Config(alias="torch._dynamo.config.recompile_limit")
|
||||
"""
|
||||
Maximum number of recompilations allowed for a single function before falling
|
||||
back to eager execution. This compilation performance control prevents excessive
|
||||
recompilation overhead that can degrade overall performance.
|
||||
"""
|
||||
|
||||
accumulated_recompile_limit: int = Config(
|
||||
alias="torch._dynamo.config.accumulated_recompile_limit"
|
||||
)
|
||||
"""
|
||||
Global limit on total recompilations across all compiled functions to prevent
|
||||
runaway recompilation scenarios. This safeguard protects against compilation
|
||||
performance issues that could affect the entire program.
|
||||
"""
|
||||
|
||||
verbose: bool = Config(alias="torch._dynamo.config.verbose")
|
||||
"""
|
||||
Enables verbose debugging output for Dynamo. When enabled, provides detailed
|
||||
information about Dynamo's compilation decisions, optimizations, and potential
|
||||
issues.
|
||||
"""
|
||||
|
||||
|
||||
# TorchDynamo-specific configuration options
|
||||
|
||||
capture_scalar_outputs: bool = Config(
|
||||
alias="torch._dynamo.config.capture_scalar_outputs"
|
||||
)
|
||||
"""
|
||||
Controls whether TorchDynamo captures operations that return scalar values (like .item())
|
||||
into the FX graph. When disabled, these operations cause graph breaks. This is a
|
||||
TorchDynamo-specific tracing behavior that affects how the tracer handles
|
||||
scalar-returning operations.
|
||||
"""
|
||||
|
||||
capture_dynamic_output_shape_ops: bool = Config(
|
||||
alias="torch._dynamo.config.capture_dynamic_output_shape_ops"
|
||||
)
|
||||
"""
|
||||
Controls whether TorchDynamo captures operations with dynamic output shapes (like
|
||||
nonzero, unique) into the FX graph. When disabled, these operations cause graph breaks.
|
||||
This is a TorchDynamo-specific setting for handling operations with unpredictable
|
||||
output shapes during tracing.
|
||||
"""
|
||||
|
||||
log_file_name: str | None = Config(alias="torch._dynamo.config.log_file_name")
|
||||
"""
|
||||
Specifies a file path for TorchDynamo-specific logging output. When set, internal
|
||||
TorchDynamo debug information is written to this file rather than stdout. This is
|
||||
useful for debugging TorchDynamo's internal tracing behavior.
|
||||
"""
|
||||
|
||||
fail_on_recompile_limit_hit: bool = Config(
|
||||
alias="torch._dynamo.config.fail_on_recompile_limit_hit"
|
||||
)
|
||||
"""
|
||||
Raises a hard error when recompile limits are exceeded instead of falling back
|
||||
to eager execution. This is useful for detecting excessive recompilation in
|
||||
performance-critical deployments where you want to ensure compilation overhead
|
||||
is kept under control.
|
||||
"""
|
||||
|
||||
allow_unspec_int_on_nn_module: bool = Config(
|
||||
alias="torch._dynamo.config.allow_unspec_int_on_nn_module"
|
||||
)
|
||||
"""
|
||||
Allows integer attributes of nn.Module instances to be unspecialized through
|
||||
the dynamic shape mechanism. By default, TorchDynamo specializes on all integer
|
||||
module attributes, but this can cause excessive recompilation when integers
|
||||
like step counters change frequently.
|
||||
"""
|
||||
|
||||
skip_tensor_guards_with_matching_dict_tags: bool = Config(
|
||||
alias="torch._dynamo.config.skip_tensor_guards_with_matching_dict_tags"
|
||||
)
|
||||
"""
|
||||
Optimizes guard generation by treating tensors as immutable when they are
|
||||
dictionary values with consistent dictionary tags across invocations. This
|
||||
reduces guard overhead for tensors stored in persistent data structures.
|
||||
"""
|
||||
|
||||
enable_cpp_symbolic_shape_guards: bool = Config(
|
||||
alias="torch._dynamo.config.enable_cpp_symbolic_shape_guards"
|
||||
)
|
||||
"""
|
||||
Uses C++ implementation for symbolic shape guard evaluation to improve performance.
|
||||
The C++ guard manager can significantly speed up guard checking for symbolic shapes
|
||||
in shape-polymorphic compilations.
|
||||
"""
|
||||
|
||||
wrap_top_frame: bool = Config(alias="torch._dynamo.config.wrap_top_frame")
|
||||
"""
|
||||
Wraps the top-level decorated function/module in a frame wrapper to ensure
|
||||
nn.Module hooks are compiled within the same frame as the main function. This
|
||||
improves compilation coverage for models that rely on hooks.
|
||||
"""
|
||||
|
||||
reorderable_logging_functions: set = Config(
|
||||
alias="torch._dynamo.config.reorderable_logging_functions"
|
||||
)
|
||||
"""
|
||||
A set of logging functions that can be reordered to execute after the compiled
|
||||
portion of the graph, allowing larger graphs to be captured. Functions in this
|
||||
set will have their execution deferred to avoid graph breaks, though this may
|
||||
affect the timing of log output. In particular, mutated values will not be logged
|
||||
at the right time, leading to incorrect logging.
|
||||
"""
|
||||
|
||||
|
||||
install_config_module(sys.modules[__name__])
|
||||
Reference in New Issue
Block a user