Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,3 @@
|
||||
from .checkpoint_activation import checkpoint
|
||||
from .contract import _get_registry, contract
|
||||
from .replicate import replicate
|
||||
+135
@@ -0,0 +1,135 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Generator
|
||||
from contextlib import AbstractContextManager, contextmanager, nullcontext
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.checkpoint import (
|
||||
_checkpoint_without_reentrant_generator,
|
||||
_DEFAULT_DETERMINISM_MODE,
|
||||
)
|
||||
|
||||
from .contract import _State, contract
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _no_hook(module: nn.Module, user_ctx: AbstractContextManager | None = None):
|
||||
r"""
|
||||
Disable hooks installed by checkpoint to avoid unintentional recursion
|
||||
during backward recomputation.
|
||||
"""
|
||||
|
||||
with user_ctx if user_ctx else nullcontext():
|
||||
orig_enable_hook = checkpoint.state(module).enable_hook
|
||||
checkpoint.state(module).enable_hook = False
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
checkpoint.state(module).enable_hook = orig_enable_hook
|
||||
|
||||
|
||||
class _CheckpointState(_State):
|
||||
enable_hook: bool = False
|
||||
_ac_generator: Generator[None, None, None] | None
|
||||
|
||||
|
||||
@contract(_CheckpointState)
|
||||
def checkpoint(module: nn.Module, **kwargs) -> nn.Module:
|
||||
r"""
|
||||
This is a composable activation checkpointing API. Unlike functional
|
||||
activation checkpointing APIs, this one does not require changing model
|
||||
source code. Unlike ``nn.Module`` wrapper activation checkpointing APIs,
|
||||
this one does not modify model structure or fully-qualified names either.
|
||||
Under the hood, it registers activation checkpointing logic as pre- and
|
||||
post-forward hooks. Hence, this API can be easily applied to any model or
|
||||
sub-modules in the model.
|
||||
|
||||
Args:
|
||||
module (nn.Module): the target model or sub-module to apply activation
|
||||
checkpointing.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> import torch.nn as nn
|
||||
>>>
|
||||
>>> class MyModel(nn.Module):
|
||||
>>> def __init__(self) -> None:
|
||||
>>> super().__init__()
|
||||
>>> self.l1 = nn.Linear(10, 10)
|
||||
>>> self.l2 = nn.Linear(10, 10)
|
||||
>>>
|
||||
>>> def forward(self, x):
|
||||
>>> return self.l2(self.l1(x))
|
||||
>>>
|
||||
>>> model = MyModel()
|
||||
>>> checkpoint(model.l1) # apply activation checkpointing only to l1
|
||||
>>> model(torch.zeros(2, 10)).sum().backward()
|
||||
|
||||
"""
|
||||
torch._C._log_api_usage_once("torch.distributed.checkpoint")
|
||||
|
||||
use_reentrant = kwargs.pop("use_reentrant", False)
|
||||
if use_reentrant:
|
||||
raise NotImplementedError(
|
||||
"use_reentrant=True is not supported in composable checkpoint. "
|
||||
"Please use torch.utils.checkpoint.checkpoint instead."
|
||||
)
|
||||
preserve_rng_state = kwargs.pop("preserve_rng_state", True)
|
||||
user_context_fns = kwargs.pop("context_fn", None)
|
||||
determinism_check = kwargs.pop("determinism_check", _DEFAULT_DETERMINISM_MODE)
|
||||
debug = kwargs.pop("debug", False)
|
||||
early_stop = kwargs.pop("early_stop", True)
|
||||
|
||||
if kwargs:
|
||||
raise ValueError(
|
||||
"Unexpected keyword arguments: " + ",".join(arg for arg in kwargs)
|
||||
)
|
||||
|
||||
def forward_pre_hook(
|
||||
module: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]
|
||||
) -> None:
|
||||
if checkpoint.state(module).enable_hook:
|
||||
|
||||
def context_fns():
|
||||
if user_context_fns is not None:
|
||||
ctx1, ctx2 = user_context_fns()
|
||||
return ctx1, _no_hook(module, ctx2)
|
||||
else:
|
||||
return nullcontext(), _no_hook(module)
|
||||
|
||||
gen = _checkpoint_without_reentrant_generator(
|
||||
module,
|
||||
preserve_rng_state,
|
||||
context_fns,
|
||||
determinism_check,
|
||||
debug,
|
||||
early_stop,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
checkpoint.state(module)._ac_generator = gen
|
||||
next(gen)
|
||||
|
||||
def forward_hook(module: nn.Module, inputs: tuple[Any, ...], output: Any) -> Any:
|
||||
if checkpoint.state(module).enable_hook:
|
||||
try:
|
||||
gen = checkpoint.state(module)._ac_generator
|
||||
if gen is None:
|
||||
raise AssertionError
|
||||
next(gen)
|
||||
except StopIteration:
|
||||
pass
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Expected non-reentrant activation checkpoint generator to be exhausted, but it was not!"
|
||||
)
|
||||
|
||||
# Ensure that we no longer hold on to the generator. always_call=True helps ensure we
|
||||
# clear this even in the case of exception in fwd pass.
|
||||
checkpoint.state(module)._ac_generator = None
|
||||
|
||||
checkpoint.state(module).enable_hook = True
|
||||
module.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
|
||||
module.register_forward_hook(forward_hook, prepend=True, always_call=True)
|
||||
return module
|
||||
@@ -0,0 +1,250 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import uuid
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
from typing import Concatenate, Generic, Protocol
|
||||
from typing_extensions import ParamSpec, TypeVar
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributed._composable_state import _State
|
||||
from torch.distributed.utils import _get_root_modules
|
||||
|
||||
|
||||
_T = TypeVar("_T", covariant=True)
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
|
||||
def generate_state_key(string="__composable_api_state_key"):
|
||||
return f"{string}_{str(uuid.uuid4())}"
|
||||
|
||||
|
||||
STATE_KEY = generate_state_key()
|
||||
REGISTRY_KEY = generate_state_key()
|
||||
|
||||
|
||||
# TODO: we can add additional info to RegistryItem to share across APIs. E.g.,
|
||||
# we can add args and kwargs here, and then we can detect whether fully_shard
|
||||
# is combined with reentrant activation checkpointing and error out with a clear
|
||||
# message.
|
||||
class RegistryItem:
|
||||
pass
|
||||
|
||||
|
||||
_TState = TypeVar("_TState", bound="_State", covariant=True)
|
||||
_M = TypeVar("_M", nn.Module, list[nn.Module])
|
||||
|
||||
|
||||
class _ContractFn(Protocol, Generic[_P, _T, _TState]):
|
||||
def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _T: ...
|
||||
|
||||
def state(self, module: nn.Module) -> _TState: ...
|
||||
|
||||
|
||||
def contract(
|
||||
state_cls: type[_TState] = _State, # type: ignore[assignment]
|
||||
) -> Callable[
|
||||
[Callable[Concatenate[_M, _P], _M]],
|
||||
_ContractFn[Concatenate[_M, _P], _M, _TState],
|
||||
]:
|
||||
r"""
|
||||
Decorate a function as a composable distributed API, where the first
|
||||
argument of the function must be an :class:`nn.Module` instance or sequence
|
||||
of :class:`nn.Module` instances.
|
||||
|
||||
The decorator verifies that the decorated function does not modify
|
||||
fully-qualified names (FQNs) for parameters, buffers, or modules. The
|
||||
decorated function can return different module instances than the input
|
||||
modules; the FQN invariant will be enforced following the input order.
|
||||
|
||||
When a function ``func`` is decorated by ``@contract()``, a
|
||||
``.state(module: nn.Module)`` method will be installed to the decorated
|
||||
function. Then you can retrieve and modify the state on a module by calling
|
||||
``func.state(module)``.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> import torch.nn as nn
|
||||
>>>
|
||||
>>> class MyModel(nn.Module):
|
||||
>>> def __init__(self) -> None:
|
||||
>>> super().__init__()
|
||||
>>> self.l1 = nn.Linear(10, 10)
|
||||
>>> self.l2 = nn.Linear(10, 10)
|
||||
>>>
|
||||
>>> def forward(self, x):
|
||||
>>> return self.l2(self.l1(x))
|
||||
>>>
|
||||
>>> @contract()
|
||||
>>> def my_feature(module: nn.Module) -> nn.Module:
|
||||
>>> my_feature.state(module).some_state = "any value"
|
||||
>>> return module
|
||||
>>>
|
||||
>>> model = MyModel()
|
||||
>>> my_feature(model.l1)
|
||||
>>> assert my_feature.state(model.l1).some_state == "any value"
|
||||
>>> my_feature(model.l2)
|
||||
>>> model(torch.randn(2, 10)).sum().backward()
|
||||
"""
|
||||
|
||||
# wraps will make functions decorated with contract() pickleable - needed for integration with torch.package
|
||||
@wraps(state_cls) # type: ignore[arg-type]
|
||||
def inner(
|
||||
func: Callable[Concatenate[_M, _P], _M],
|
||||
) -> _ContractFn[Concatenate[_M, _P], _M, _TState]:
|
||||
@wraps(func)
|
||||
def wrapper(
|
||||
module: _M,
|
||||
*args: _P.args,
|
||||
**kwargs: _P.kwargs,
|
||||
) -> _M:
|
||||
inp_module = module
|
||||
modules: list[nn.Module]
|
||||
if isinstance(module, nn.Module):
|
||||
modules = [module]
|
||||
else:
|
||||
# If the user passes a sequence of modules, then we assume that
|
||||
# we only need to insert the state object on the root modules
|
||||
# (i.e. those without a parent) among the passed-in modules.
|
||||
# pyrefly: ignore [no-matching-overload]
|
||||
modules = _get_root_modules(list(module))
|
||||
state = state_cls() # shared across all modules
|
||||
registry_item = RegistryItem() # shared across all modules
|
||||
|
||||
# `func` is allowed to return different module instances than the
|
||||
# input modules as long as FQNs are preserved following the input
|
||||
# module order
|
||||
all_orig_named_params: list[dict[str, nn.Parameter]] = []
|
||||
all_orig_named_buffers: list[dict[str, torch.Tensor]] = []
|
||||
all_orig_named_modules: list[dict[str, nn.Module]] = []
|
||||
|
||||
for module in modules:
|
||||
default_all_state: dict[Callable, _State] = OrderedDict()
|
||||
default_registry: dict[str, RegistryItem] = OrderedDict()
|
||||
all_state: dict[Callable, _State] = module.__dict__.setdefault( # type: ignore[call-overload]
|
||||
STATE_KEY, default_all_state
|
||||
)
|
||||
if not isinstance(all_state, dict):
|
||||
raise AssertionError(
|
||||
f"Distributed composable API states corrupted: {all_state}"
|
||||
)
|
||||
registry: dict[str, RegistryItem] = module.__dict__.setdefault( # type: ignore[call-overload]
|
||||
REGISTRY_KEY, default_registry
|
||||
)
|
||||
if not isinstance(registry, dict):
|
||||
raise AssertionError(
|
||||
f"Distributed composable API registry corrupted: {registry}"
|
||||
)
|
||||
if func in all_state or func.__name__ in registry:
|
||||
raise AssertionError(
|
||||
"Each distinct composable distributed API can only be applied to a "
|
||||
f"module once. {func.__name__} has already been applied to the "
|
||||
f"following module:\n{module}"
|
||||
)
|
||||
all_state.setdefault(func, state)
|
||||
registry.setdefault(func.__name__, registry_item)
|
||||
|
||||
all_orig_named_params.append(OrderedDict(module.named_parameters()))
|
||||
all_orig_named_buffers.append(OrderedDict(module.named_buffers()))
|
||||
all_orig_named_modules.append(OrderedDict(module.named_modules()))
|
||||
|
||||
updated = func(inp_module, *args, **kwargs)
|
||||
if updated is None:
|
||||
updated = inp_module # type: ignore[assignment]
|
||||
updated_modules: list[nn.Module]
|
||||
if isinstance(updated, nn.Module):
|
||||
updated_modules = [updated]
|
||||
else:
|
||||
updated_modules = _get_root_modules(list(inp_module)) # type: ignore[arg-type, call-overload]
|
||||
|
||||
all_new_named_params: list[dict[str, nn.Parameter]] = []
|
||||
all_new_named_buffers: list[dict[str, torch.Tensor]] = []
|
||||
all_new_named_modules: list[dict[str, nn.Module]] = []
|
||||
for module in updated_modules:
|
||||
all_new_named_params.append(OrderedDict(module.named_parameters()))
|
||||
all_new_named_buffers.append(OrderedDict(module.named_buffers()))
|
||||
all_new_named_modules.append(OrderedDict(module.named_modules()))
|
||||
|
||||
num_orig_modules = len(all_orig_named_modules)
|
||||
num_new_modules = len(all_new_named_modules)
|
||||
if num_orig_modules != num_new_modules:
|
||||
raise AssertionError(
|
||||
f"{func.__name__} should return the same number of modules as input modules"
|
||||
f"Inputs: {num_orig_modules} modules\n"
|
||||
f"Outputs: {num_new_modules} modules"
|
||||
)
|
||||
|
||||
def check_fqn(orig_fqns: list[str], new_fqns: list[str], check_key: str):
|
||||
if orig_fqns == new_fqns:
|
||||
return
|
||||
|
||||
orig_fqn_set, new_fqn_set = set(orig_fqns), set(new_fqns)
|
||||
orig_only = orig_fqn_set - new_fqn_set
|
||||
new_only = new_fqn_set - orig_fqn_set
|
||||
if len(orig_only) or len(new_only):
|
||||
raise RuntimeError(
|
||||
f"{check_key}"
|
||||
"Composable distributed API implementations cannot modify FQNs.\n"
|
||||
f"FQNs only in original: {orig_only}\n"
|
||||
f"FQNs only in new: {new_only}"
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"{check_key}"
|
||||
"Composable distributed API implementations cannot modify "
|
||||
"the order of FQNs.\n"
|
||||
f"Original FQNs: {orig_only}\n"
|
||||
f"New FQNs: {new_only}"
|
||||
)
|
||||
|
||||
for orig_named_params, new_named_params in zip(
|
||||
all_orig_named_params, all_new_named_params
|
||||
):
|
||||
check_fqn(
|
||||
list(orig_named_params.keys()),
|
||||
list(new_named_params.keys()),
|
||||
"Checking parameters: ",
|
||||
)
|
||||
for orig_named_buffers, new_named_buffers in zip(
|
||||
all_orig_named_buffers, all_new_named_buffers
|
||||
):
|
||||
check_fqn(
|
||||
list(orig_named_buffers.keys()),
|
||||
list(new_named_buffers.keys()),
|
||||
"Checking buffers: ",
|
||||
)
|
||||
for orig_named_modules, new_named_modules in zip(
|
||||
all_orig_named_modules, all_new_named_modules
|
||||
):
|
||||
check_fqn(
|
||||
list(orig_named_modules.keys()),
|
||||
list(new_named_modules.keys()),
|
||||
"Checking modules: ",
|
||||
)
|
||||
|
||||
# TODO: verify that installed distributed paradigms are compatible with
|
||||
# each other.
|
||||
|
||||
return updated
|
||||
|
||||
def get_state(module: nn.Module) -> _State:
|
||||
return module.__dict__.setdefault( # type: ignore[call-overload]
|
||||
STATE_KEY,
|
||||
{}, # TODO(@yhcharles): this is a temporary fix, need a better way
|
||||
).get(func) # type: ignore[call-overload]
|
||||
|
||||
wrapper.state = get_state # type: ignore[attr-defined]
|
||||
|
||||
return wrapper # type: ignore[return-value]
|
||||
|
||||
return inner # type: ignore[return-value]
|
||||
|
||||
|
||||
def _get_registry(module: nn.Module) -> dict[str, RegistryItem] | None:
|
||||
r"""
|
||||
Get an ``OrderedDict`` of composable APIs that have been applied to the
|
||||
``module``, indexed by the API name. If no API has been applied, then this
|
||||
returns ``None``.
|
||||
"""
|
||||
return getattr(module, REGISTRY_KEY, None)
|
||||
+3
@@ -0,0 +1,3 @@
|
||||
from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy
|
||||
|
||||
from .fully_shard import FSDPModule, fully_shard, register_fsdp_forward_method
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
# TODO: For backward compatibility, we are importing the public objects
|
||||
# originally from this file.
|
||||
from torch.distributed.fsdp import ( # noqa: F401
|
||||
FSDPModule,
|
||||
fully_shard,
|
||||
register_fsdp_forward_method,
|
||||
UnshardHandle,
|
||||
)
|
||||
@@ -0,0 +1,257 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import weakref
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, NoReturn
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributed._composable_state import _State
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
|
||||
from .contract import _get_registry, contract
|
||||
|
||||
|
||||
_ROOT_MODULE_PREFIX = ""
|
||||
|
||||
|
||||
class _ReplicateState(_State):
|
||||
_ddp_weakref: weakref.ref
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.module: nn.Module = nn.ParameterList()
|
||||
self.has_initialized: bool = False
|
||||
self._param_list: nn.ParameterList = nn.ParameterList()
|
||||
# TODO(@fegin): this variable is originally create for testing, we
|
||||
# should remove this if possible.
|
||||
self._orig_module = self.module
|
||||
self._param_names: list[str] = []
|
||||
self._no_sync: bool = False
|
||||
self._init_args: tuple[Any, ...] | None = None
|
||||
self._init_kwargs: dict[str, Any] = {}
|
||||
self._comm_hook_args: list[Any] = []
|
||||
|
||||
def _collect_params(
|
||||
self,
|
||||
module: nn.Module,
|
||||
ignored_modules: set[nn.Module],
|
||||
ignored_params: set[nn.Parameter],
|
||||
prefix: str = _ROOT_MODULE_PREFIX,
|
||||
) -> None:
|
||||
# skip if managed by fully_sharded API
|
||||
if _is_fully_sharded(module):
|
||||
return
|
||||
|
||||
# if a module is ignored, all descendants of the module are ignored.
|
||||
if module in ignored_modules:
|
||||
return
|
||||
|
||||
recurse_prefix = (
|
||||
f"{prefix}." if prefix != _ROOT_MODULE_PREFIX else _ROOT_MODULE_PREFIX
|
||||
)
|
||||
|
||||
for n, p in module.named_parameters(recurse=False):
|
||||
if p not in ignored_params:
|
||||
self._param_list.append(p)
|
||||
self._param_names.append(f"{recurse_prefix}{n}")
|
||||
|
||||
for name, child_module in module.named_children():
|
||||
self._collect_params(
|
||||
child_module,
|
||||
ignored_modules,
|
||||
ignored_params,
|
||||
prefix=f"{recurse_prefix}{name}",
|
||||
)
|
||||
|
||||
def lazy_init(self) -> None:
|
||||
@torch._disable_dynamo(recursive=True)
|
||||
def _lazy_init():
|
||||
if self._init_args is None:
|
||||
raise AssertionError
|
||||
self.init(*self._init_args, **self._init_kwargs)
|
||||
self.register_comm_hook()
|
||||
self._init_args = ()
|
||||
self._init_kwargs = {}
|
||||
|
||||
_lazy_init()
|
||||
|
||||
def init(
|
||||
self,
|
||||
module: nn.Module,
|
||||
ignored_modules: set[nn.Module],
|
||||
**kwargs,
|
||||
) -> None:
|
||||
if self.has_initialized:
|
||||
return
|
||||
|
||||
self.has_initialized = True
|
||||
self.module = module
|
||||
ignored_params = {p for m in ignored_modules for p in m.parameters()}
|
||||
for submodule in module.modules():
|
||||
if _is_fully_sharded(submodule):
|
||||
ignored_params.update(submodule.parameters())
|
||||
from torch.distributed.tensor.parallel.ddp import _localize_dtensor
|
||||
|
||||
_localize_dtensor(module, ignored_params=ignored_params)
|
||||
self._collect_params(module, ignored_modules, ignored_params)
|
||||
|
||||
if "device_id" in kwargs:
|
||||
# replicate() supports a small usability enhancement where
|
||||
# user can pass in device_id as a Union[int, torch.device] even for
|
||||
# CPU devices so users don't have to change code for CPU/GPU runs.
|
||||
# We derive the right device_ids to feed into DDP to support this.
|
||||
if kwargs["device_id"] is not None:
|
||||
device_id = kwargs["device_id"]
|
||||
# Convert to device_ids that DDP expects.
|
||||
if isinstance(device_id, torch.device) and device_id.type == "cpu":
|
||||
# CPU modules receive device_ids None
|
||||
kwargs["device_ids"] = None
|
||||
else:
|
||||
# GPU modules expect device_ids=[cuda_device]
|
||||
kwargs["device_ids"] = [device_id]
|
||||
else:
|
||||
kwargs["device_ids"] = None
|
||||
kwargs.pop("device_id")
|
||||
|
||||
self._ddp = DistributedDataParallel(self._param_list, **kwargs)
|
||||
# Weakref to the DDP instance is currently only used for testing.
|
||||
replicate.state(self.module)._ddp_weakref = weakref.ref(self._ddp)
|
||||
|
||||
def register_comm_hook(self) -> None:
|
||||
for comm_args, comm_kwargs in self._comm_hook_args:
|
||||
self._ddp.register_comm_hook(*comm_args, **comm_kwargs)
|
||||
self._comm_hook_args.clear()
|
||||
|
||||
def record_init_args(self, *args, **kwargs) -> None:
|
||||
self._init_args = args
|
||||
self._init_kwargs = kwargs
|
||||
|
||||
def forward_pre_hook(
|
||||
self, module: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]
|
||||
) -> Any:
|
||||
if self._init_args or self._init_kwargs:
|
||||
self.lazy_init()
|
||||
self._ddp.require_backward_grad_sync = not self._no_sync
|
||||
DistributedDataParallel._active_ddp_module = self._ddp
|
||||
return self._ddp._pre_forward(*args, **kwargs)
|
||||
|
||||
def forward_post_hook(
|
||||
self,
|
||||
module: nn.Module,
|
||||
input: tuple[torch.Tensor],
|
||||
output: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
DistributedDataParallel._active_ddp_module = None
|
||||
return self._ddp._post_forward(output)
|
||||
|
||||
|
||||
def unimplemented_deepcopy(*args: Any, **kwargs: Any) -> NoReturn:
|
||||
raise AssertionError(
|
||||
"DDP does not support deepcopy. Please use state dict for serialization."
|
||||
)
|
||||
|
||||
|
||||
# Follow the same pattern as FSDP/fully_shard
|
||||
class DDP:
|
||||
def __new__(cls, *args, **kwargs):
|
||||
"""
|
||||
Override ``__new__`` to remove the DDP class and directly construct
|
||||
the original class for cases like indexing into a container module.
|
||||
"""
|
||||
# Use index 2 since 0 is the dynamically constructed `DDP<...>` class
|
||||
# and index 1 is the `DDP` class itself
|
||||
orig_cls = cls.__mro__[2]
|
||||
return orig_cls.__new__(orig_cls, *args, **kwargs)
|
||||
|
||||
def set_requires_gradient_sync(self, requires_gradient_sync: bool) -> None:
|
||||
"""
|
||||
Sets if the module should sync gradients. This can be used to implement
|
||||
gradient accumulation without communication.
|
||||
|
||||
Args:
|
||||
requires_gradient_sync (bool): Whether to reduce gradients for the
|
||||
module's parameters.
|
||||
"""
|
||||
replicate.state(self)._no_sync = not requires_gradient_sync # type: ignore[arg-type]
|
||||
|
||||
def register_comm_hook(self, *args, **kwargs) -> None:
|
||||
replicate.state(self)._comm_hook_args.append((args, kwargs)) # type: ignore[arg-type]
|
||||
|
||||
|
||||
@contract(state_cls=_ReplicateState)
|
||||
def replicate(
|
||||
module: nn.Module,
|
||||
ignored_modules: Iterable[torch.nn.Module] | None = None,
|
||||
**kwargs,
|
||||
) -> nn.Module:
|
||||
r"""Replicates a module
|
||||
|
||||
Args:
|
||||
module (torch.nn.Module): module to replicate
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +REQUIRES(module:torch._C._distributed_c10d)
|
||||
>>> module = nn.Linear(3, 3)
|
||||
>>> replicate(module)
|
||||
"""
|
||||
torch._C._log_api_usage_once("torch.distributed.replicate")
|
||||
|
||||
# TODO(fegin): using kwargs is not a good idea if we would like to make
|
||||
# replicate a formal API to replace DDP.
|
||||
if "device_id" in kwargs:
|
||||
if not isinstance(kwargs["device_id"], (int, torch.device)):
|
||||
raise RuntimeError(
|
||||
"Expected device_id to be int or torch.device, "
|
||||
f"but got {type(kwargs['device_id'])}"
|
||||
)
|
||||
|
||||
if _is_fully_sharded(module):
|
||||
raise RuntimeError(
|
||||
"Cannot apply `replicate()` on a Module already managed by `fully_shard`"
|
||||
)
|
||||
|
||||
if ignored_modules is None:
|
||||
ignored_modules = {}
|
||||
else:
|
||||
ignored_modules = set(ignored_modules)
|
||||
|
||||
state = replicate.state(module)
|
||||
module.register_forward_pre_hook(state.forward_pre_hook, with_kwargs=True)
|
||||
device_mesh = kwargs.get("device_mesh")
|
||||
if device_mesh is not None:
|
||||
root_mesh = device_mesh._get_root_mesh()
|
||||
# if a root mesh is not the same as device_mesh,
|
||||
# meaning the device_mesh is sliced out from the root mesh.
|
||||
if root_mesh != device_mesh:
|
||||
# TODO: This is a temporary work around to enable DDP + TP.
|
||||
# We should do the logic in DDP so that the 2D implementation is
|
||||
# sound and the state_dict works out of the box.
|
||||
#
|
||||
# This won't conflict with what is done in DDP class as the module
|
||||
# replicate is going to pass is NOT the original module.
|
||||
from torch.distributed.tensor.parallel.ddp import (
|
||||
_localize_dtensor,
|
||||
_reconstruct_dtensor,
|
||||
)
|
||||
|
||||
module.register_forward_pre_hook(_reconstruct_dtensor)
|
||||
module.register_forward_hook(_localize_dtensor)
|
||||
|
||||
module.register_forward_hook(state.forward_post_hook) # type: ignore[arg-type]
|
||||
|
||||
state.record_init_args(module, ignored_modules, **kwargs)
|
||||
|
||||
# Place DDP leftmost for highest priority in the method resolution order
|
||||
cls = module.__class__
|
||||
dct = {"__deepcopy__": unimplemented_deepcopy}
|
||||
new_cls = type(f"DDP{cls.__name__}", (DDP, cls), dct)
|
||||
module.__class__ = new_cls
|
||||
return module
|
||||
|
||||
|
||||
def _is_fully_sharded(module: nn.Module) -> bool:
|
||||
r"""Check if module is marked with fully_shard."""
|
||||
registry = _get_registry(module)
|
||||
if registry is None:
|
||||
return False
|
||||
return "fully_shard" in registry
|
||||
+204
@@ -0,0 +1,204 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import overload, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributed._composable_state import _get_module_state
|
||||
from torch.distributed.fsdp._fully_shard._fsdp_api import (
|
||||
MixedPrecisionPolicy,
|
||||
OffloadPolicy,
|
||||
)
|
||||
from torch.distributed.fsdp._fully_shard._fsdp_common import DDPMeshInfo
|
||||
from torch.distributed.fsdp._fully_shard._fsdp_init import (
|
||||
_apply_to_module,
|
||||
_get_device_from_mesh,
|
||||
_get_mesh_info,
|
||||
_get_modules_and_states,
|
||||
_init_default_mesh,
|
||||
_init_param_group,
|
||||
_validate_mesh as _validate_mesh_common,
|
||||
_validate_module as _validate_module_common,
|
||||
)
|
||||
from torch.distributed.fsdp._fully_shard._fsdp_state import FSDPState, FSDPStateContext
|
||||
from torch.distributed.fsdp._fully_shard._fully_shard import (
|
||||
_unimplemented_deepcopy,
|
||||
FSDPModule,
|
||||
)
|
||||
|
||||
from .contract import _get_registry, contract
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.distributed.fsdp._fully_shard._fsdp_api import DataParallelMeshDims
|
||||
from torch.distributed.tensor import DeviceMesh
|
||||
|
||||
|
||||
cls_to_replicate_cls: dict[type, type] = {}
|
||||
|
||||
logger = logging.getLogger("torch.distributed._composable.replicate_with_fsdp")
|
||||
|
||||
|
||||
class _ReplicateStateContext(FSDPStateContext["_ReplicateState"]):
|
||||
"""
|
||||
State shared across Replicate states.
|
||||
|
||||
This is a typed subclass of FSDPStateContext parameterized with _ReplicateState,
|
||||
providing correct type annotations (e.g., all_states: list[_ReplicateState]).
|
||||
It also allows call sites to differentiate between Replicate and FSDP contexts
|
||||
via isinstance checks if needed.
|
||||
"""
|
||||
|
||||
|
||||
def _get_module_replicate_state(module: nn.Module) -> _ReplicateState | None:
|
||||
state = _get_module_state(module)
|
||||
if isinstance(state, _ReplicateState):
|
||||
return state
|
||||
return None
|
||||
|
||||
|
||||
class _ReplicateState(FSDPState):
|
||||
_state_name: str = "Replicate"
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._state_ctx = _ReplicateStateContext()
|
||||
|
||||
def _get_state_for_module(self, module: nn.Module) -> FSDPState | None:
|
||||
return _get_module_replicate_state(module)
|
||||
|
||||
def init(
|
||||
self,
|
||||
modules: tuple[nn.Module, ...],
|
||||
device: torch.device,
|
||||
mp_policy: MixedPrecisionPolicy,
|
||||
auto_reshard_after_forward: bool = False,
|
||||
) -> None:
|
||||
super().init(modules, device, mp_policy, auto_reshard_after_forward)
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def replicate(
|
||||
module: nn.Module,
|
||||
*,
|
||||
mesh: DeviceMesh | None = ...,
|
||||
mp_policy: MixedPrecisionPolicy = ...,
|
||||
offload_policy: OffloadPolicy = ...,
|
||||
ignored_params: set[nn.Parameter] | None = ...,
|
||||
dp_mesh_dims: DataParallelMeshDims | None = ...,
|
||||
) -> ReplicateModule: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def replicate(
|
||||
module: list[nn.Module],
|
||||
*,
|
||||
mesh: DeviceMesh | None = ...,
|
||||
mp_policy: MixedPrecisionPolicy = ...,
|
||||
offload_policy: OffloadPolicy = ...,
|
||||
ignored_params: set[nn.Parameter] | None = ...,
|
||||
dp_mesh_dims: DataParallelMeshDims | None = ...,
|
||||
) -> list[ReplicateModule]: ...
|
||||
|
||||
|
||||
@contract(state_cls=_ReplicateState) # type: ignore[misc]
|
||||
def replicate(
|
||||
module: nn.Module,
|
||||
*,
|
||||
mesh: DeviceMesh | None = None,
|
||||
mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
|
||||
offload_policy: OffloadPolicy = OffloadPolicy(),
|
||||
ignored_params: set[nn.Parameter] | None = None,
|
||||
dp_mesh_dims: DataParallelMeshDims | None = None,
|
||||
):
|
||||
r"""Replicates a module
|
||||
|
||||
Args:
|
||||
module (torch.nn.Module): module to replicate
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +REQUIRES(module:torch._C._distributed_c10d)
|
||||
>>> module = nn.Linear(3, 3)
|
||||
>>> replicate(module)
|
||||
"""
|
||||
torch._C._log_api_usage_once("torch.distributed._composable.replicate_with_fsdp")
|
||||
_validate_module(module)
|
||||
mesh = mesh or _init_default_mesh(mesh_dim_names=("replicate",))
|
||||
if dp_mesh_dims is not None:
|
||||
_validate_mesh_common(mesh, dp_mesh_dims)
|
||||
mesh_info = _get_mesh_info(mesh, dp_mesh_dims)
|
||||
if not isinstance(mesh_info, DDPMeshInfo):
|
||||
raise ValueError(
|
||||
"replicate() with dp_mesh_dims requires replicate-only "
|
||||
"dims (no shard dims). Use fully_shard() for sharding."
|
||||
)
|
||||
else:
|
||||
_validate_mesh(mesh)
|
||||
mesh_info = DDPMeshInfo(mesh, replicate_mesh_dim=0)
|
||||
device = _get_device_from_mesh(mesh)
|
||||
# managed_modules (3rd return) and buffers (5th return) are unused:
|
||||
# - managed_modules: FSDP uses this to set Dynamo-specific attributes
|
||||
# (_is_fsdp_managed_module, _fsdp_use_orig_params), which replicate doesn't need
|
||||
# - buffers: already moved to device by _get_modules_and_states; replicate
|
||||
# doesn't need to track them separately
|
||||
arg_module, modules, _, params, _ = _get_modules_and_states(
|
||||
module,
|
||||
device,
|
||||
ignored_params,
|
||||
is_composable_fn=is_composable_with_replicate,
|
||||
get_state_fn=_get_module_replicate_state,
|
||||
)
|
||||
state = replicate.state(modules[0]) # type: ignore[attr-defined]
|
||||
state.init(modules, device, mp_policy)
|
||||
|
||||
_init_param_group(
|
||||
state,
|
||||
params,
|
||||
modules,
|
||||
mesh_info,
|
||||
None, # post_forward_mesh_info
|
||||
device,
|
||||
None, # shard_placement_fn
|
||||
mp_policy,
|
||||
offload_policy,
|
||||
)
|
||||
|
||||
# Place Replicate leftmost for highest priority in the method resolution order
|
||||
_apply_to_module(
|
||||
modules,
|
||||
cls_to_replicate_cls,
|
||||
ReplicateModule,
|
||||
"Replicate",
|
||||
_unimplemented_deepcopy,
|
||||
)
|
||||
return arg_module
|
||||
|
||||
|
||||
class ReplicateModule(FSDPModule):
|
||||
# Index in MRO where the original class is found.
|
||||
# For Replicate: [Replicate<Orig>, ReplicateModule, FSDPModule, Orig, ...] -> index 3
|
||||
_orig_cls_mro_index: int = 3
|
||||
|
||||
|
||||
def is_composable_with_replicate(module: nn.Module) -> bool:
|
||||
registry = _get_registry(module)
|
||||
if registry is None:
|
||||
return True
|
||||
return "fully_shard" not in registry
|
||||
|
||||
|
||||
def _validate_module(module: nn.Module) -> None:
|
||||
if not is_composable_with_replicate(module):
|
||||
raise RuntimeError(
|
||||
"Cannot apply `replicate()` on a Module already managed by `fully_shard`"
|
||||
)
|
||||
_validate_module_common(module, "replicate")
|
||||
|
||||
|
||||
def _validate_mesh(mesh: DeviceMesh) -> None:
|
||||
if mesh.ndim != 1:
|
||||
raise ValueError(f"replicate expects a 1D DeviceMesh but got {mesh}")
|
||||
Reference in New Issue
Block a user