Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,888 @@
|
||||
import abc
|
||||
import contextlib
|
||||
import contextvars
|
||||
import functools
|
||||
import logging
|
||||
import threading
|
||||
from collections import defaultdict, deque
|
||||
from collections.abc import (
|
||||
Callable,
|
||||
Generator,
|
||||
Iterable,
|
||||
Iterator,
|
||||
MutableMapping,
|
||||
Sequence,
|
||||
)
|
||||
from typing import (
|
||||
Any,
|
||||
cast,
|
||||
Literal,
|
||||
NamedTuple,
|
||||
Optional,
|
||||
TYPE_CHECKING,
|
||||
TypeAlias,
|
||||
Union,
|
||||
)
|
||||
from weakref import WeakKeyDictionary, WeakValueDictionary
|
||||
|
||||
import torch
|
||||
from torch.autograd.variable import Variable
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
from torch.utils.hooks import RemovableHandle
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._ops import OpOverload
|
||||
|
||||
|
||||
__all__ = [
|
||||
"saved_tensors_hooks",
|
||||
"save_on_cpu",
|
||||
"disable_saved_tensors_hooks",
|
||||
"register_multi_grad_hook",
|
||||
"allow_mutation_on_saved_tensors",
|
||||
"Node",
|
||||
"GradientEdge",
|
||||
"get_gradient_edge",
|
||||
"increment_version",
|
||||
"set_warn_on_accumulate_grad_stream_mismatch",
|
||||
]
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Node(abc.ABC):
|
||||
@abc.abstractmethod
|
||||
def name(self) -> str:
|
||||
r"""Return the name.
|
||||
|
||||
Example::
|
||||
|
||||
>>> import torch
|
||||
>>> a = torch.tensor([0., 0., 0.], requires_grad=True)
|
||||
>>> b = a.clone()
|
||||
>>> assert isinstance(b.grad_fn, torch.autograd.graph.Node)
|
||||
>>> print(b.grad_fn.name())
|
||||
CloneBackward0
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def next_functions(self) -> tuple[tuple[Optional["Node"], int], ...]:
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def metadata(self) -> dict:
|
||||
r"""Return the metadata."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def _sequence_nr(self) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def _input_metadata(self) -> list[Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def _register_hook_dict(self, tensor: torch.Tensor) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def register_hook(self, fn: Callable[..., Any]) -> RemovableHandle:
|
||||
r"""Register a backward hook.
|
||||
|
||||
The hook will be called every time a gradient with respect to the
|
||||
Node is computed. The hook should have the following signature::
|
||||
|
||||
hook(grad_inputs: Tuple[Tensor], grad_outputs: Tuple[Tensor]) -> Tuple[Tensor] or None
|
||||
|
||||
|
||||
The hook should not modify its argument, but it can optionally return
|
||||
a new gradient which will be used in place of :attr:`grad_inputs`.
|
||||
|
||||
This function returns a handle with a method ``handle.remove()``
|
||||
that removes the hook from the module.
|
||||
|
||||
.. note::
|
||||
See :ref:`backward-hooks-execution` for more information on how when this hook
|
||||
is executed, and how its execution is ordered relative to other hooks.
|
||||
|
||||
.. note::
|
||||
In the rare case where the hook is registered while the Node has already
|
||||
begun execution, there is no longer any guarantee on :attr:`grad_outputs`
|
||||
content (it might be as usual or empty depending on other factors). The
|
||||
hook can still optionally return a new gradient to be used in place of
|
||||
:attr:`grad_inputs` independent of :attr:`grad_outputs`.
|
||||
|
||||
Example::
|
||||
|
||||
>>> import torch
|
||||
>>> a = torch.tensor([0., 0., 0.], requires_grad=True)
|
||||
>>> b = a.clone()
|
||||
>>> assert isinstance(b.grad_fn, torch.autograd.graph.Node)
|
||||
>>> handle = b.grad_fn.register_hook(lambda gI, gO: (gO[0] * 2,))
|
||||
>>> b.sum().backward(retain_graph=True)
|
||||
>>> print(a.grad)
|
||||
tensor([2., 2., 2.])
|
||||
>>> handle.remove() # Removes the hook
|
||||
>>> a.grad = None
|
||||
>>> b.sum().backward(retain_graph=True)
|
||||
>>> print(a.grad)
|
||||
tensor([1., 1., 1.])
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def register_prehook(self, fn: Callable[..., Any]) -> RemovableHandle:
|
||||
r"""Register a backward pre-hook.
|
||||
|
||||
The hook will be called every time a gradient with respect to the
|
||||
Node is computed. The hook should have the following signature::
|
||||
|
||||
hook(grad_outputs: Tuple[Tensor]) -> Tuple[Tensor] or None
|
||||
|
||||
The hook should not modify its argument, but it can optionally return
|
||||
a new gradient which will be used in place of :attr:`grad_outputs`.
|
||||
|
||||
This function returns a handle with a method ``handle.remove()``
|
||||
that removes the hook from the module.
|
||||
|
||||
.. note::
|
||||
See :ref:`backward-hooks-execution` for more information on how when this hook
|
||||
is executed, and how its execution is ordered relative to other hooks.
|
||||
|
||||
Example::
|
||||
|
||||
>>> a = torch.tensor([0., 0., 0.], requires_grad=True)
|
||||
>>> b = a.clone()
|
||||
>>> assert isinstance(b.grad_fn, torch.autograd.graph.Node)
|
||||
>>> handle = b.grad_fn.register_prehook(lambda gI: (gI[0] * 2,))
|
||||
>>> b.sum().backward(retain_graph=True)
|
||||
>>> print(a.grad)
|
||||
tensor([2., 2., 2.])
|
||||
>>> handle.remove()
|
||||
>>> a.grad = None
|
||||
>>> b.sum().backward(retain_graph=True)
|
||||
>>> print(a.grad)
|
||||
tensor([1., 1., 1.])
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def __subclasshook__(cls, subclass: type) -> bool:
|
||||
if cls is Node and (
|
||||
(
|
||||
subclass is not None
|
||||
and subclass is getattr(torch._C._functions, subclass.__name__, None)
|
||||
)
|
||||
or issubclass(subclass, torch.autograd.function.BackwardCFunction)
|
||||
):
|
||||
return True
|
||||
return NotImplemented
|
||||
|
||||
|
||||
def _get_grad_fn_or_grad_acc(t: Union[torch.Tensor, "GradientEdge"]) -> Node:
|
||||
if isinstance(t, GradientEdge):
|
||||
return t.node
|
||||
if t.requires_grad and t.grad_fn is None:
|
||||
with torch.enable_grad():
|
||||
node = t.view_as(t).grad_fn.next_functions[0][0] # type: ignore[union-attr]
|
||||
else:
|
||||
node = t.grad_fn
|
||||
if node is None:
|
||||
raise AssertionError("Expected gradient function to be set")
|
||||
return node
|
||||
|
||||
|
||||
class GradientEdge(NamedTuple):
|
||||
"""Object representing a given gradient edge within the autograd graph.
|
||||
|
||||
To get the gradient edge where a given Tensor gradient will be computed,
|
||||
you can do ``edge = autograd.graph.get_gradient_edge(tensor)``.
|
||||
"""
|
||||
|
||||
node: Node
|
||||
output_nr: int
|
||||
# This token can be used to ensure the graph stays alive when it cannot be
|
||||
# done via the node field
|
||||
ownership_token: Node | None = None
|
||||
|
||||
|
||||
def get_gradient_edge(tensor: torch.Tensor) -> GradientEdge:
|
||||
"""Get the gradient edge for computing the gradient of the given Tensor.
|
||||
|
||||
In particular, it is equivalent to call
|
||||
``g = autograd.grad(loss, input)`` and ``g = autograd.grad(loss, get_gradient_edge(input))``.
|
||||
"""
|
||||
if not tensor.requires_grad:
|
||||
raise RuntimeError(
|
||||
"It is not possible to get the gradient edge for a Tensor "
|
||||
"that does not require gradients",
|
||||
)
|
||||
grad_fn = _get_grad_fn_or_grad_acc(tensor)
|
||||
|
||||
# Python-based Node are owned by the C++ side meaning the python grad_fn
|
||||
# object we hold here does NOT keep the C++ graph alive.
|
||||
# Create an ownership token by creating a new C++ node that own the graph
|
||||
# we care about here.
|
||||
token = None
|
||||
if isinstance(grad_fn, torch._C._FunctionBase):
|
||||
with torch.enable_grad():
|
||||
token = tensor.view_as(tensor).grad_fn
|
||||
|
||||
# Note that output_nr default to 0 which is the right value
|
||||
# for the AccumulateGrad node.
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
return GradientEdge(grad_fn, tensor.output_nr, ownership_token=token)
|
||||
|
||||
|
||||
def increment_version(tensor: torch.Tensor | Iterable[torch.Tensor]) -> None:
|
||||
"""Update autograd metadata tracking whether the given Tensor was modified in place.
|
||||
|
||||
This is to enable more accurate error checking within the autograd engine.
|
||||
It is already done automatically by PyTorch functions and within custom Function
|
||||
when mark_dirty() is called appropriately so you only need to call this explicitly
|
||||
if you are doing inplace operation on the Tensor data in a way that Pytorch doesn't
|
||||
know about. For example a custom kernel that reads the Tensor data_ptr and modifies
|
||||
the memory inplace based on this pointer. Can accept either a tensor, or a list of tensors.
|
||||
|
||||
Note that incrementing the version counter multiple times for a single inplace operation
|
||||
is not problematic.
|
||||
|
||||
Note that if you pass in tensor constructed under torch.inference_mode(),
|
||||
we will not bump its version counter (because your tensor does not have one).
|
||||
"""
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
tensor = (tensor,)
|
||||
torch._C._increment_version(tensor)
|
||||
|
||||
|
||||
class saved_tensors_hooks:
|
||||
"""Context-manager that sets a pair of pack / unpack hooks for saved tensors.
|
||||
|
||||
Use this context-manager to define how intermediary results of an operation
|
||||
should be packed before saving, and unpacked on retrieval.
|
||||
|
||||
In that context, the ``pack_hook`` function will be called every time an
|
||||
operation saves a tensor for backward (this includes intermediary results
|
||||
saved using
|
||||
:func:`~torch.autograd.function._ContextMethodMixin.save_for_backward` but
|
||||
also those recorded by a PyTorch-defined operation). The output of
|
||||
``pack_hook`` is then stored in the computation graph instead of the
|
||||
original tensor.
|
||||
|
||||
The ``unpack_hook`` is called when the saved tensor needs to be accessed,
|
||||
namely when executing :func:`torch.Tensor.backward()` or
|
||||
:func:`torch.autograd.grad()`. It takes as argument the *packed* object
|
||||
returned by ``pack_hook`` and should return a tensor which has the same
|
||||
content as the original tensor (passed as input to the corresponding
|
||||
``pack_hook``).
|
||||
|
||||
The hooks should have the following signatures:
|
||||
|
||||
pack_hook(tensor: Tensor) -> Any
|
||||
|
||||
unpack_hook(Any) -> Tensor
|
||||
|
||||
where the return value of ``pack_hook`` is a valid input to ``unpack_hook``.
|
||||
|
||||
In general, you want ``unpack_hook(pack_hook(t))`` to be equal to ``t`` in terms
|
||||
of value, size, dtype and device.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
|
||||
>>> def pack_hook(x):
|
||||
... print("Packing", x)
|
||||
... return x.detach()
|
||||
>>>
|
||||
>>> def unpack_hook(x):
|
||||
... print("Unpacking", x)
|
||||
... return x
|
||||
>>>
|
||||
>>> a = torch.ones(5, requires_grad=True)
|
||||
>>> b = torch.ones(5, requires_grad=True) * 2
|
||||
>>> with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
|
||||
... y = a * b
|
||||
Packing tensor([1., 1., 1., 1., 1.], requires_grad=True)
|
||||
Packing tensor([2., 2., 2., 2., 2.], grad_fn=<MulBackward0>)
|
||||
>>> y.sum().backward()
|
||||
Unpacking tensor([1., 1., 1., 1., 1.], requires_grad=True)
|
||||
Unpacking tensor([2., 2., 2., 2., 2.], grad_fn=<MulBackward0>)
|
||||
|
||||
.. warning ::
|
||||
Performing an inplace operation on the input to either hooks may lead
|
||||
to undefined behavior.
|
||||
|
||||
.. warning ::
|
||||
Only one pair of hooks is allowed at a time. When recursively nesting this
|
||||
context-manager, only the inner-most pair of hooks will be applied.
|
||||
|
||||
.. warning ::
|
||||
To avoid reference cycle, the return value of ``pack_hook`` cannot hold a
|
||||
reference to the input tensor. For example, use `lambda x: x.detach()`
|
||||
instead of `lambda x: x` as the pack hook.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pack_hook: Callable[[torch.Tensor], Any],
|
||||
unpack_hook: Callable[[Any], torch.Tensor],
|
||||
) -> None:
|
||||
self.pack_hook = pack_hook
|
||||
self.unpack_hook = unpack_hook
|
||||
|
||||
def __enter__(self) -> None:
|
||||
torch._C._autograd._push_saved_tensors_default_hooks(
|
||||
self.pack_hook, self.unpack_hook
|
||||
)
|
||||
|
||||
def __exit__(self, *args: object) -> None:
|
||||
torch._C._autograd._pop_saved_tensors_default_hooks()
|
||||
|
||||
|
||||
class save_on_cpu(saved_tensors_hooks):
|
||||
"""Context manager under which tensors saved by the forward pass will be stored on cpu, then retrieved for backward.
|
||||
|
||||
When performing operations within this context manager, intermediary
|
||||
results saved in the graph during the forward pass will be moved to CPU,
|
||||
then copied back to the original device when needed for the backward pass.
|
||||
If the graph was already on CPU, no tensor copy is performed.
|
||||
|
||||
Use this context-manager to trade compute for GPU memory usage (e.g.
|
||||
when your model doesn't fit in GPU memory during training).
|
||||
|
||||
.. warning::
|
||||
|
||||
When ``pin_memory=True``, the GPU to CPU copy during packing is
|
||||
asynchronous. Accessing saved tensors on CPU (e.g. via
|
||||
``grad_fn._saved_self``) before the CUDA stream has finished may
|
||||
yield incorrect data. Call :func:`torch.cuda.synchronize` first
|
||||
if you need to read them.
|
||||
|
||||
Args:
|
||||
pin_memory (bool): If ``True`` tensors will be saved to CPU pinned memory
|
||||
during packing and copied to GPU asynchronously during both
|
||||
packing and unpacking.
|
||||
Defaults to ``False``.
|
||||
Also see :ref:`cuda-memory-pinning`.
|
||||
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
|
||||
>>> a = torch.randn(5, requires_grad=True, device="cuda")
|
||||
>>> b = torch.randn(5, requires_grad=True, device="cuda")
|
||||
>>> c = torch.randn(5, requires_grad=True, device="cuda")
|
||||
>>>
|
||||
>>> def f(a, b, c):
|
||||
... prod_1 = a * b # a and b are saved on GPU
|
||||
... with torch.autograd.graph.save_on_cpu():
|
||||
... prod_2 = prod_1 * c # prod_1 and c are saved on CPU
|
||||
... y = prod_2 * a # prod_2 and a are saved on GPU
|
||||
... return y
|
||||
>>>
|
||||
>>> y = f(a, b, c)
|
||||
>>> del a, b, c # for illustration only
|
||||
>>> # the content of a, b, and prod_2 are still alive on GPU
|
||||
>>> # the content of prod_1 and c only live on CPU
|
||||
>>> y.sum().backward() # all CPU tensors are moved back to GPU, for backward
|
||||
>>> # all intermediary tensors are released (deleted) after the call to backward
|
||||
"""
|
||||
|
||||
def __init__(self, pin_memory: bool = False, device_type: str = "cuda") -> None:
|
||||
device_module = getattr(torch, device_type, torch.cuda)
|
||||
|
||||
def pack_to_cpu(tensor: torch.Tensor) -> tuple[torch.device, torch.Tensor]:
|
||||
if not pin_memory:
|
||||
return (tensor.device, tensor.cpu())
|
||||
is_pinnable = device_module.is_available() and not tensor.is_sparse
|
||||
packed = torch.empty(
|
||||
tensor.size(),
|
||||
dtype=tensor.dtype,
|
||||
layout=tensor.layout,
|
||||
pin_memory=is_pinnable,
|
||||
)
|
||||
packed.copy_(tensor, non_blocking=is_pinnable)
|
||||
return (tensor.device, packed)
|
||||
|
||||
def unpack_from_cpu(packed: tuple[torch.device, torch.Tensor]) -> torch.Tensor:
|
||||
device, tensor = packed
|
||||
return tensor.to(device, non_blocking=pin_memory)
|
||||
|
||||
super().__init__(pack_to_cpu, unpack_from_cpu)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def disable_saved_tensors_hooks(error_message: str) -> Generator[None, None, None]:
|
||||
"""Context-manager that disables the saved tensors default hooks feature.
|
||||
|
||||
Useful for if you are creating a feature that does not work with saved
|
||||
tensors default hooks.
|
||||
|
||||
Args:
|
||||
error_message (str): When saved tensors default hooks are used when they
|
||||
have been are disabled, a RuntimeError with this
|
||||
error message gets raised.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +SKIP(failing)
|
||||
>>> message = "saved tensors default hooks are disabled"
|
||||
>>> with torch.autograd.graph.disable_saved_tensors_hooks(message):
|
||||
... # Raises RuntimeError: saved tensors default hooks are disabled
|
||||
... with torch.autograd.graph.save_on_cpu():
|
||||
... pass
|
||||
"""
|
||||
maybe_prev_message = None
|
||||
try:
|
||||
maybe_prev_message = (
|
||||
torch._C._autograd._saved_tensors_hooks_get_disabled_error_message()
|
||||
)
|
||||
torch._C._autograd._saved_tensors_hooks_disable(error_message)
|
||||
yield
|
||||
finally:
|
||||
# See NOTE: [disabled_error_message invariant]
|
||||
if maybe_prev_message is None:
|
||||
torch._C._autograd._saved_tensors_hooks_enable()
|
||||
else:
|
||||
torch._C._autograd._saved_tensors_hooks_disable(maybe_prev_message)
|
||||
|
||||
|
||||
def set_warn_on_accumulate_grad_stream_mismatch(enabled: bool) -> None:
|
||||
"""Whether to warn when the AccumulateGrad node's stream does not match the stream
|
||||
of the node that produced the incoming gradient.
|
||||
"""
|
||||
return torch._C._set_warn_on_accumulate_grad_stream_mismatch(enabled)
|
||||
|
||||
|
||||
class _MultiHandle(RemovableHandle):
|
||||
handles: tuple[RemovableHandle, ...]
|
||||
|
||||
def __init__(self, handles: tuple[RemovableHandle, ...]) -> None:
|
||||
self.handles = handles
|
||||
|
||||
def remove(self) -> None:
|
||||
for handle in self.handles:
|
||||
handle.remove()
|
||||
|
||||
def __getstate__(self) -> tuple[RemovableHandle, ...]:
|
||||
return self.handles
|
||||
|
||||
def __setstate__(self, state: tuple[RemovableHandle, ...]) -> None:
|
||||
self.handles = state
|
||||
|
||||
|
||||
def register_multi_grad_hook(
|
||||
tensors: Sequence[torch.Tensor],
|
||||
fn: Callable[[Sequence[torch.Tensor | None]], None]
|
||||
| Callable[[torch.Tensor], None],
|
||||
*,
|
||||
mode: Literal["all", "any"] = "all",
|
||||
) -> RemovableHandle:
|
||||
r"""Register a multi-grad backward hook.
|
||||
|
||||
There are two supported modes: ``"all"`` and ``"any"``.
|
||||
|
||||
Under the ``"all"`` mode, the hook will be called after gradients with respect to every tensor in
|
||||
:attr:`tensors` have been computed. If a tensor is in :attr:`tensors` but
|
||||
is not part of the graph, or if a tensor is not needed to compute the gradients
|
||||
for any ``inputs`` specified for the current ``.backward()`` or ``.grad()`` call,
|
||||
this tensor will be ignored and the hook will not wait for its gradient to be
|
||||
computed.
|
||||
|
||||
After every non-ignored tensor's gradient has been computed, :attr:`fn` will be
|
||||
called with those gradients. ``None`` will be passed for tensors that did not
|
||||
have their gradients computed.
|
||||
|
||||
Under the ``"any"`` mode, the hook will be called after the first gradient
|
||||
with respect to a tensor in :attr:`tensors` has been computed. The hook
|
||||
will be called with that gradient as its argument.
|
||||
|
||||
The hook should not modify its arguments.
|
||||
|
||||
This function returns a handle with a method ``handle.remove()`` that removes the hook.
|
||||
|
||||
.. note::
|
||||
See :ref:`backward-hooks-execution` for more information on how when this hook
|
||||
is executed, and how its execution is ordered relative to other hooks.
|
||||
|
||||
Example::
|
||||
|
||||
>>> import torch
|
||||
>>>
|
||||
>>> a = torch.rand(2, 3, requires_grad=True)
|
||||
>>> b = torch.rand(2, 3, requires_grad=True)
|
||||
>>> c = a * b
|
||||
>>> d = a * b
|
||||
>>>
|
||||
>>> def fn(grads):
|
||||
... print([g is not None for g in grads])
|
||||
...
|
||||
>>> torch.autograd.graph.register_multi_grad_hook((a, b, c, d), fn)
|
||||
>>>
|
||||
>>> c.sum().backward(retain_graph=True)
|
||||
[True, True, True, False]
|
||||
>>> c.sum().backward(inputs=(a,), retain_graph=True)
|
||||
[True, False, True, False]
|
||||
>>>
|
||||
"""
|
||||
supported_modes = ("all", "any")
|
||||
lock = threading.Lock()
|
||||
|
||||
if mode not in supported_modes:
|
||||
raise ValueError(f"Expects mode to be one of {supported_modes} but got {mode}")
|
||||
|
||||
if mode == "all":
|
||||
count: dict[int, int] = {}
|
||||
nb_calls = None
|
||||
buffer: dict[int, list[torch.Tensor | None]] = {}
|
||||
|
||||
grad_fns = list(map(_get_grad_fn_or_grad_acc, tensors))
|
||||
len_tensors = len(tensors)
|
||||
|
||||
def get_inner_hook(idx: int) -> Callable[[torch.Tensor], None]:
|
||||
def inner_hook(grad: torch.Tensor) -> None:
|
||||
nonlocal count, nb_calls, buffer, fn
|
||||
id = torch._C._current_graph_task_id()
|
||||
if id == -1:
|
||||
raise AssertionError(
|
||||
"expected this hook to be called inside a backward call"
|
||||
)
|
||||
count[id] = count.get(id, 0)
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
buffer[id] = buffer.get(id, [None] * len_tensors)
|
||||
|
||||
with lock:
|
||||
curr_count, count[id] = count[id], count[id] + 1
|
||||
|
||||
if curr_count == 0:
|
||||
# On the first call, compute the actual nb_calls and buffer
|
||||
nb_calls = sum(
|
||||
map(torch._C._will_engine_execute_node, grad_fns)
|
||||
)
|
||||
|
||||
buffer[id][idx] = grad
|
||||
|
||||
if nb_calls is None:
|
||||
raise AssertionError("Expected nb_calls to be set")
|
||||
if curr_count == nb_calls - 1:
|
||||
fn = cast(Callable[[Sequence[torch.Tensor | None]], None], fn)
|
||||
fn(buffer[id])
|
||||
del count[id]
|
||||
del buffer[id]
|
||||
|
||||
return inner_hook
|
||||
|
||||
handles = tuple(
|
||||
t.register_hook(get_inner_hook(i)) for i, t in enumerate(tensors)
|
||||
)
|
||||
elif mode == "any":
|
||||
fn = cast(Callable[[torch.Tensor], None], fn)
|
||||
ran_hook: dict[int, bool] = defaultdict(bool)
|
||||
|
||||
@functools.wraps(fn)
|
||||
def wrapped_fn(grad: torch.Tensor) -> None:
|
||||
nonlocal ran_hook
|
||||
id = torch._C._current_graph_task_id()
|
||||
if id == -1:
|
||||
raise AssertionError(
|
||||
"expected this hook to be called inside a backward call"
|
||||
)
|
||||
with lock:
|
||||
prev, ran_hook[id] = ran_hook[id], True
|
||||
if prev:
|
||||
return
|
||||
fn(grad)
|
||||
|
||||
handles = tuple(
|
||||
tensor.register_hook(wrapped_fn)
|
||||
for tensor in tensors
|
||||
if tensor.requires_grad
|
||||
)
|
||||
|
||||
return _MultiHandle(handles) # type: ignore[possibly-undefined]
|
||||
|
||||
|
||||
# NOTE [Allow mutation on tensors saved for backward]
|
||||
#
|
||||
# 1. Tensor gets saved for backward
|
||||
# - remember the python object id and the version of the tensor
|
||||
# - remember aliasing information (data_ptr of base + version)
|
||||
# - save the original so we control its lifetime
|
||||
# 2. Any time a tensor gets in-placed
|
||||
# - for each tensor aliased to it:
|
||||
# - check using its object id and version to see if it has been saved
|
||||
# - if it has been saved, clone it
|
||||
# - delete the reference to the original
|
||||
# 3. during backward
|
||||
# - if the clone exists, the tensor must've been modified in-place
|
||||
_allow_mutation_on_saved_tensors_enabled: bool = False
|
||||
|
||||
|
||||
_TID: TypeAlias = tuple[int, int, int]
|
||||
_SID: TypeAlias = tuple[int, int]
|
||||
|
||||
|
||||
def _get_tid(tensor: torch.Tensor) -> _TID:
|
||||
# FIXME: This is almost definitely a bug.
|
||||
if isinstance(
|
||||
tensor,
|
||||
(
|
||||
torch._subclasses.fake_tensor.FakeTensor,
|
||||
torch._subclasses.functional_tensor.FunctionalTensor,
|
||||
),
|
||||
):
|
||||
data_ptr = 0
|
||||
else:
|
||||
data_ptr = tensor.data_ptr()
|
||||
return (id(tensor), data_ptr, tensor._version)
|
||||
|
||||
|
||||
def _get_sid(tensor: torch.Tensor) -> _SID:
|
||||
# FIXME: This is almost definitely a bug.
|
||||
if isinstance(
|
||||
tensor,
|
||||
(
|
||||
torch._subclasses.fake_tensor.FakeTensor,
|
||||
torch._subclasses.functional_tensor.FunctionalTensor,
|
||||
),
|
||||
):
|
||||
data_ptr = 0
|
||||
else:
|
||||
data_ptr = tensor.data_ptr()
|
||||
return (data_ptr, tensor._version)
|
||||
|
||||
|
||||
class _Handle:
|
||||
pass
|
||||
|
||||
|
||||
class _swap_with_cloned(saved_tensors_hooks):
|
||||
def __init__(self, ctx: "_AllowMutationOnSavedContext") -> None:
|
||||
def pack_hook(tensor: torch.Tensor) -> _Handle:
|
||||
tid = _get_tid(tensor)
|
||||
sid = _get_sid(tensor)
|
||||
# Tensors saved for backward have an entry in _tid_to_weakhandle
|
||||
handle: _Handle | None = None
|
||||
|
||||
# Save aliasing information
|
||||
ctx.sid_to_tid[sid].add(tid)
|
||||
|
||||
# NB: The same tensor (of the same version) can be saved multiple times
|
||||
if tid not in ctx.tid_to_weakhandle:
|
||||
handle = _Handle()
|
||||
ctx.tid_to_weakhandle[tid] = handle
|
||||
ctx.original[handle] = tensor
|
||||
else:
|
||||
# Store an additional strong reference to the handle
|
||||
handle = ctx.tid_to_weakhandle[tid]
|
||||
return handle
|
||||
|
||||
def unpack_hook(handle: _Handle) -> torch.Tensor:
|
||||
error_msg = (
|
||||
"Trying to backward outside of the 'allow_mutation_on_saved_tensors' context"
|
||||
"in which the graph was originally recorded."
|
||||
)
|
||||
if not _allow_mutation_on_saved_tensors_enabled:
|
||||
raise AssertionError(error_msg)
|
||||
if handle in ctx.cloned:
|
||||
res = ctx.cloned[handle]
|
||||
else:
|
||||
if handle not in ctx.original:
|
||||
raise AssertionError(error_msg)
|
||||
res = ctx.original[handle]
|
||||
return res
|
||||
|
||||
super().__init__(pack_hook, unpack_hook)
|
||||
|
||||
|
||||
class _CloneArgBeforeMutateMode(TorchDispatchMode):
|
||||
def __init__(self, ctx: "_AllowMutationOnSavedContext") -> None:
|
||||
self.ctx = ctx
|
||||
|
||||
def __torch_dispatch__(
|
||||
self,
|
||||
func: "OpOverload",
|
||||
types: Iterable[type],
|
||||
args: tuple[Any, ...] = (),
|
||||
kwargs: dict[Any, Any] | None = None,
|
||||
) -> Any:
|
||||
kwargs = kwargs or {}
|
||||
|
||||
def maybe_clone(t: torch.Tensor) -> None:
|
||||
tid = _get_tid(t)
|
||||
sid = _get_sid(t)
|
||||
ctx = self.ctx
|
||||
if sid in ctx.sid_to_tid:
|
||||
for tid in ctx.sid_to_tid[sid]:
|
||||
if tid not in ctx.tid_to_weakhandle:
|
||||
# We know that if tid is in sid_to_tid, then it must also be in
|
||||
# tid_to_weakhandle. However, it is possible for the tensor to be
|
||||
# saved at one point, but cleared by backward before it is modified
|
||||
# in-place. Consider the following example:
|
||||
#
|
||||
# >>> a = torch.randn(2, 3, requires_grad=True).clone()
|
||||
# >>> out = (a**2).sum()
|
||||
# >>> out.backward()
|
||||
# >>> a.sin_()
|
||||
continue
|
||||
handle = ctx.tid_to_weakhandle[tid]
|
||||
if handle in ctx.cloned:
|
||||
# The same exact tensor has been cloned already
|
||||
continue
|
||||
ctx.cloned[handle] = ctx.original[handle].clone()
|
||||
del ctx.original[handle]
|
||||
|
||||
for idx, arg in enumerate(func._schema.arguments):
|
||||
if arg.alias_info is not None and arg.alias_info.is_write:
|
||||
if arg.is_out:
|
||||
maybe_clone(kwargs["out"])
|
||||
elif isinstance(args[idx], list):
|
||||
# Foreach case. (Possible optimization: if most of the
|
||||
# tensors need to be cloned, use a for each clone?)
|
||||
for t in args[idx]:
|
||||
maybe_clone(t)
|
||||
else:
|
||||
maybe_clone(args[idx])
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
|
||||
class _AllowMutationOnSavedContext:
|
||||
def __init__(self) -> None:
|
||||
self.cloned: MutableMapping[_Handle, torch.Tensor] = WeakKeyDictionary()
|
||||
self.original: MutableMapping[_Handle, torch.Tensor] = WeakKeyDictionary()
|
||||
self.tid_to_weakhandle: MutableMapping[_TID, _Handle] = WeakValueDictionary()
|
||||
self.sid_to_tid: dict[_SID, set[_TID]] = defaultdict(set)
|
||||
|
||||
def clear(self) -> None:
|
||||
self.cloned.clear()
|
||||
self.original.clear()
|
||||
self.tid_to_weakhandle.clear()
|
||||
self.sid_to_tid.clear()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def allow_mutation_on_saved_tensors() -> Generator[
|
||||
_AllowMutationOnSavedContext, None, None
|
||||
]:
|
||||
"""Context manager under which mutating tensors saved for backward is allowed.
|
||||
|
||||
Under this context manager, tensors saved for backward are cloned on mutation,
|
||||
so the original version can still be used during backward. Normally, mutating a tensor
|
||||
saved for backward will result in an error raised when it's used during backward.
|
||||
|
||||
To ensure the correct behavior, both the forward and backward should be run under
|
||||
the same context manager.
|
||||
|
||||
Returns:
|
||||
An _AllowMutationOnSavedContext object storing the state managed by this
|
||||
context manager. This object can be useful for debugging purposes. The state
|
||||
managed by the context manager is automatically cleared upon exiting.
|
||||
|
||||
Example::
|
||||
|
||||
>>> import torch
|
||||
>>> with torch.autograd.graph.allow_mutation_on_saved_tensors():
|
||||
... # forward
|
||||
... a = torch.ones(2, 3, requires_grad=True)
|
||||
... b = a.clone()
|
||||
... out = (b**2).sum()
|
||||
... b.sin_()
|
||||
... # backward
|
||||
... out.sum().backward()
|
||||
...
|
||||
tensor([[0.8415, 0.8415, 0.8415],
|
||||
[0.8415, 0.8415, 0.8415]], grad_fn=<SinBackward0>)
|
||||
"""
|
||||
global _allow_mutation_on_saved_tensors_enabled
|
||||
|
||||
ctx = _AllowMutationOnSavedContext()
|
||||
|
||||
with _swap_with_cloned(ctx), _CloneArgBeforeMutateMode(ctx):
|
||||
try:
|
||||
if _allow_mutation_on_saved_tensors_enabled:
|
||||
raise RuntimeError(
|
||||
"allow_mutation_on_saved_tensors contexts cannot be nested"
|
||||
)
|
||||
_allow_mutation_on_saved_tensors_enabled = True
|
||||
yield ctx
|
||||
finally:
|
||||
ctx.clear()
|
||||
_allow_mutation_on_saved_tensors_enabled = False
|
||||
|
||||
|
||||
def _register_logging_hooks_on_whole_graph(
|
||||
t_outputs: Sequence[torch.Tensor | GradientEdge],
|
||||
) -> Callable[[], None]:
|
||||
grad_fns = list(map(_get_grad_fn_or_grad_acc, t_outputs))
|
||||
|
||||
def iter_graph(roots: list[Node]) -> Iterator[Node]:
|
||||
if not roots:
|
||||
return
|
||||
seen: set[Node] = set()
|
||||
q: deque[Node] = deque()
|
||||
for node in roots:
|
||||
if node is not None:
|
||||
seen.add(node)
|
||||
q.append(node)
|
||||
|
||||
while q:
|
||||
node = q.popleft()
|
||||
for fn, _ in node.next_functions:
|
||||
if fn in seen or fn is None:
|
||||
continue
|
||||
seen.add(fn)
|
||||
q.append(fn)
|
||||
|
||||
yield node
|
||||
|
||||
def fmt(t: torch.Tensor | None) -> str:
|
||||
# Avoid circular import
|
||||
from torch.utils._dtype_abbrs import dtype_abbrs
|
||||
|
||||
if t is None:
|
||||
return "None"
|
||||
return f"{dtype_abbrs[t.dtype]}[{', '.join(map(str, t.shape))}]"
|
||||
|
||||
def prehook(grad_outputs: Sequence[torch.Tensor | None]) -> None:
|
||||
node = torch._C._current_autograd_node()
|
||||
grad_outputs_str = f"[{','.join(fmt(t) for t in grad_outputs)}]"
|
||||
log_str = f"Executing: {node} with grad_outputs: {grad_outputs_str}"
|
||||
log.debug(log_str)
|
||||
|
||||
handles = [node.register_prehook(prehook) for node in iter_graph(grad_fns)]
|
||||
|
||||
def unregister_hooks() -> None:
|
||||
for handle in handles:
|
||||
handle.remove()
|
||||
|
||||
return unregister_hooks
|
||||
|
||||
|
||||
def _engine_run_backward(
|
||||
t_outputs: Sequence[torch.Tensor | GradientEdge],
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
attach_logging_hooks = log.getEffectiveLevel() <= logging.DEBUG
|
||||
if attach_logging_hooks:
|
||||
unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
|
||||
|
||||
# Need to save the context so compiler config will be visible in device threads
|
||||
torch._C._stash_obj_in_tls("context", contextvars.copy_context())
|
||||
|
||||
try:
|
||||
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
|
||||
t_outputs, *args, **kwargs
|
||||
) # Calls into the C++ engine to run the backward pass
|
||||
finally:
|
||||
if attach_logging_hooks:
|
||||
unregister_hooks() # type: ignore[possibly-undefined]
|
||||
torch._C._stash_obj_in_tls("context", None)
|
||||
Reference in New Issue
Block a user