diff --git a/lib/python3.12/site-packages/torch/_refs/__pycache__/_conversions.cpython-312.pyc b/lib/python3.12/site-packages/torch/_refs/__pycache__/_conversions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c98892ceffc79ea7a6a9994bb17de1ab4283619 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_refs/__pycache__/_conversions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_refs/__pycache__/fft.cpython-312.pyc b/lib/python3.12/site-packages/torch/_refs/__pycache__/fft.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..679bb2d9210064bd0ee7f99d825799a283fef072 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_refs/__pycache__/fft.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_refs/linalg/__init__.py b/lib/python3.12/site-packages/torch/_refs/linalg/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6585f57e3d643bfd58c0b208bf33e3a1ee732d8c --- /dev/null +++ b/lib/python3.12/site-packages/torch/_refs/linalg/__init__.py @@ -0,0 +1,309 @@ +# mypy: allow-untyped-defs +from functools import partial +from typing import Optional, Tuple, Union + +import torch +import torch._prims as prims +import torch._prims_common as utils +import torch._refs as refs +import torch._refs.linalg as linalg +from torch import Tensor +from torch._prims_common import ( + check_fp_or_complex, + check_is_matrix, + Dim, + DimsType, + ELEMENTWISE_TYPE_PROMOTION_KIND, + IntLike, + TensorLikeType, +) +from torch._prims_common.wrappers import ( + _maybe_convert_to_dtype, + elementwise_type_promotion_wrapper, + out_wrapper, +) + + +__all__ = [ + "diagonal", + "matrix_norm", + "norm", + "svd", + "svdvals", + "vector_norm", + "vecdot", + "cross", +] + + +def _check_norm_dtype(dtype: Optional[torch.dtype], x_dtype: torch.dtype, fn_name: str): + """ + Checks related to the dtype kwarg in `linalg.*norm` functions + """ + if dtype is not None: + torch._check( + utils.is_float_dtype(dtype) or utils.is_complex_dtype(dtype), + lambda: f"{fn_name}: dtype should be floating point or complex. Got {dtype}", + ) + torch._check( + utils.is_complex_dtype(dtype) == utils.is_complex_dtype(x_dtype), + lambda: "{fn_name}: dtype should be {d} for {d} inputs. Got {dtype}".format( + fn_name=fn_name, + d="complex" if utils.is_complex_dtype(x_dtype) else "real", + dtype=dtype, + ), + ) + torch._check( + utils.get_higher_dtype(dtype, x_dtype) == dtype, + lambda: f"{fn_name}: the dtype of the input ({x_dtype}) should be convertible " + "without narrowing to the specified dtype ({dtype})", + ) + + +import operator + +# Utilities should come BEFORE this import +from torch._decomp import register_decomposition +from torch._decomp.decompositions import pw_cast_for_opmath + + +@register_decomposition(torch._ops.ops.aten.linalg_cross) +@out_wrapper() +@pw_cast_for_opmath +def cross(a: Tensor, b: Tensor, dim: int = -1): + torch._check( + a.ndim == b.ndim, + lambda: "linalg.cross: inputs must have the same number of dimensions.", + ) + torch._check( + a.size(dim) == 3 and b.size(dim) == 3, + lambda: f"linalg.cross: inputs dim {dim} must have length 3, got {a.size(dim)} and {b.size(dim)}", + ) + a, b = torch.broadcast_tensors(a, b) + dim = utils.canonicalize_dim(a.ndim, dim) + idx = torch.arange(3, device=a.device) + return a.index_select(dim, (idx + 1) % 3) * b.index_select( + dim, (idx + 2) % 3 + ) - a.index_select(dim, (idx + 2) % 3) * b.index_select(dim, (idx + 1) % 3) + + +def diagonal( + input: TensorLikeType, + *, + offset: int = 0, + dim1: int = -2, + dim2: int = -1, +) -> TensorLikeType: + return torch.diagonal(input, offset=offset, dim1=dim1, dim2=dim2) + + +@register_decomposition(torch._ops.ops.aten.linalg_vector_norm) +@out_wrapper(exact_dtype=True) +def vector_norm( + x: TensorLikeType, + ord: Union[float, int] = 2, + dim: Optional[DimsType] = None, + keepdim: bool = False, + *, + dtype: Optional[torch.dtype] = None, +) -> Tensor: + from torch.fx.experimental.symbolic_shapes import guard_size_oblivious + + # Checks + check_fp_or_complex(x.dtype, "linalg.vector_norm") + + if isinstance(dim, Dim): + dim = [dim] # type: ignore[assignment] + + if guard_size_oblivious(x.numel() == 0) and (ord < 0.0 or ord == float("inf")): + torch._check( + dim is not None and len(dim) != 0, + lambda: f"linalg.vector_norm cannot compute the {ord} norm on an empty tensor " + "because the operation does not have an identity", + ) + shape = x.shape + assert dim is not None # mypy does not seem to be able to see through check? + for d in dim: + torch._check( + shape[d] != 0, + lambda: f"linalg.vector_norm cannot compute the {ord} norm on the " + f"dimension {d} because this dimension is empty and the " + "operation does not have an identity", + ) + _check_norm_dtype(dtype, x.dtype, "linalg.vector_norm") + + computation_dtype, result_dtype = utils.reduction_dtypes( + x, utils.REDUCTION_OUTPUT_TYPE_KIND.COMPLEX_TO_FLOAT, dtype + ) + + to_result_dtype = partial(_maybe_convert_to_dtype, dtype=result_dtype) + + # Implementation + if ord == 0.0: + return torch.sum(torch.ne(x, 0.0), dim=dim, keepdim=keepdim, dtype=result_dtype) + elif ord == float("inf"): + return to_result_dtype(torch.amax(torch.abs(x), dim=dim, keepdim=keepdim)) # type: ignore[return-value,arg-type] + elif ord == float("-inf"): + return to_result_dtype(torch.amin(torch.abs(x), dim=dim, keepdim=keepdim)) # type: ignore[return-value,arg-type] + else: + # From here on the computation dtype is important as the reduction is non-trivial + x = _maybe_convert_to_dtype(x, computation_dtype) # type: ignore[assignment] + reduce_sum = partial(torch.sum, dim=dim, keepdim=keepdim) + + is_ord_even = ord % 2 == 0 if isinstance(ord, IntLike) else ord % 2.0 == 0.0 + if not (is_ord_even and utils.is_float_dtype(x.dtype)): + x = torch.abs(x) + return to_result_dtype(torch.pow(reduce_sum(torch.pow(x, ord)), 1.0 / ord)) # type: ignore[return-value] + + +def _backshift_permutation(dim0, dim1, ndim): + # Auxiliary function for matrix_norm + # Computes the permutation that moves the two given dimensions to the back + ret = [i for i in range(ndim) if i != dim0 and i != dim1] + ret.extend((dim0, dim1)) + return ret + + +def _inverse_permutation(perm): + # Given a permutation, returns its inverse. It's equivalent to argsort on an array + return [i for i, j in sorted(enumerate(perm), key=operator.itemgetter(1))] + + +# CompositeImplicitAutograd +@out_wrapper(exact_dtype=True) +def matrix_norm( + A: TensorLikeType, + ord: Union[float, str] = "fro", + dim: DimsType = (-2, -1), + keepdim: bool = False, + *, + dtype: Optional[torch.dtype] = None, +) -> TensorLikeType: + # shape + check_is_matrix(A, "linalg.matrix_norm") + # dim + dim = utils.canonicalize_dims(A.ndim, dim) + if isinstance(dim, Dim): + dim = (dim,) # type: ignore[assignment] + torch._check( + len(dim) == 2, lambda: "linalg.matrix_norm: dim must be a 2-tuple. Got {dim}" + ) + torch._check( + dim[0] != dim[1], + lambda: "linalg.matrix_norm: dims must be different. Got ({dim[0]}, {dim[1]})", + ) + # dtype arg + _check_norm_dtype(dtype, A.dtype, "linalg.matrix_norm") + + if isinstance(ord, str): + # ord + torch._check( + ord in ("fro", "nuc"), + lambda: "linalg.matrix_norm: Order {ord} not supported.", + ) + # dtype + check_fp_or_complex( + A.dtype, "linalg.matrix_norm", allow_low_precision_dtypes=ord != "nuc" + ) + + if ord == "fro": + return vector_norm(A, 2, dim, keepdim, dtype=dtype) + else: # ord == "nuc" + if dtype is not None: + A = _maybe_convert_to_dtype(A, dtype) # type: ignore[assignment] + perm = _backshift_permutation(dim[0], dim[1], A.ndim) + result = torch.sum(svdvals(prims.transpose(A, perm)), -1, keepdim) + if keepdim: + inv_perm = _inverse_permutation(perm) + result = prims.transpose(torch.unsqueeze(result, -1), inv_perm) + return result + else: + # ord + abs_ord = abs(ord) + torch._check( + abs_ord in (2, 1, float("inf")), + lambda: "linalg.matrix_norm: Order {ord} not supported.", + ) + # dtype + check_fp_or_complex( + A.dtype, "linalg.matrix_norm", allow_low_precision_dtypes=ord != 2 + ) + + max_min = partial(torch.amax if ord > 0.0 else torch.amin, keepdim=keepdim) + + if abs_ord == 2.0: + if dtype is not None: + A = _maybe_convert_to_dtype(A, dtype) # type: ignore[assignment] + perm = _backshift_permutation(dim[0], dim[1], A.ndim) + result = max_min(svdvals(prims.transpose(A, perm)), dim=-1) + if keepdim: + inv_perm = _inverse_permutation(perm) + result = prims.transpose(torch.unsqueeze(result, -1), inv_perm) + return result + else: # 1, -1, inf, -inf + dim0, dim1 = dim + if abs_ord == float("inf"): + dim0, dim1 = dim1, dim0 + if not keepdim and (dim0 < dim1): + dim1 -= 1 + return max_min( + vector_norm(A, 1.0, dim=dim0, keepdim=keepdim, dtype=dtype), dim1 + ) + + +# CompositeImplicitAutograd +@out_wrapper(exact_dtype=True) +def norm( + A: TensorLikeType, + ord: Optional[Union[float, str]] = None, + dim: Optional[DimsType] = None, + keepdim: bool = False, + *, + dtype: Optional[torch.dtype] = None, +) -> TensorLikeType: + if dim is not None: + if isinstance(dim, Dim): + dim = (dim,) # type: ignore[assignment] + torch._check( + len(dim) in (1, 2), + lambda: "linalg.norm: If dim is specified, it must be of length 1 or 2. Got {dim}", + ) + elif ord is not None: + torch._check( + A.ndim in (1, 2), + lambda: "linalg.norm: If dim is not specified but ord is, the input must be 1D or 2D. Got {A.ndim}D", + ) + + if ord is not None and ( + (dim is not None and len(dim) == 2) or (dim is None and A.ndim == 2) + ): + if dim is None: + dim = (0, 1) + return matrix_norm(A, ord, dim, keepdim, dtype=dtype) + else: + if ord is None: + ord = 2.0 + return vector_norm(A, ord, dim, keepdim, dtype=dtype) # type: ignore[arg-type] + + +# CompositeImplicitAutograd +@out_wrapper("U", "S", "Vh", exact_dtype=True) +def svd(A: TensorLikeType, full_matrices: bool = True) -> Tuple[Tensor, Tensor, Tensor]: + return prims.svd(A, full_matrices=full_matrices) + + +# CompositeImplicitAutograd +@out_wrapper(exact_dtype=True) +def svdvals(A: TensorLikeType) -> Tensor: + return svd(A, full_matrices=False)[1] + + +# CompositeImplicitAutograd +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("x", "y"), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def vecdot(x: Tensor, y: Tensor, dim: int = -1) -> Tensor: + check_fp_or_complex(x.dtype, "linalg.vecdot") + return (x.conj() * y).sum(dim=dim) diff --git a/lib/python3.12/site-packages/torch/_refs/linalg/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_refs/linalg/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e01158fb973e44b2ca2aa943ffd2438e66c8f4dc Binary files /dev/null and b/lib/python3.12/site-packages/torch/_refs/linalg/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_refs/nn/__init__.py b/lib/python3.12/site-packages/torch/_refs/nn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b7414d43515a6eebe92ab2c204b334cc345fff51 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_refs/nn/__init__.py @@ -0,0 +1,4 @@ +from typing import List + + +__all__: List[str] = [] diff --git a/lib/python3.12/site-packages/torch/_refs/nn/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_refs/nn/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fb62974a00a7034b09e35556f084ea12e47fc48 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_refs/nn/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_refs/nn/functional/__init__.py b/lib/python3.12/site-packages/torch/_refs/nn/functional/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..58aaf505e3842a49d56b75ebf9483fc84780be5a --- /dev/null +++ b/lib/python3.12/site-packages/torch/_refs/nn/functional/__init__.py @@ -0,0 +1,1280 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +import math +from functools import wraps +from typing import Callable, Optional, Union + +import torch +import torch._prims as prims +import torch._prims_common as utils +import torch._refs as refs +from torch._decomp import register_decomposition +from torch._prims_common import ( + ELEMENTWISE_TYPE_PROMOTION_KIND, + NumberType, + ShapeType, + TensorLike, + TensorLikeType, +) +from torch._prims_common.wrappers import ( + elementwise_type_promotion_wrapper, + elementwise_unary_scalar_wrapper, + out_wrapper, +) +from torch._refs import _make_inplace + + +__all__ = [ + "alpha_dropout", + "celu", + "celu_", + "channel_shuffle", + "dropout", + "elu", + "elu_", + "gelu", + "glu", + "group_norm", + "hardshrink", + "hardtanh", + "hinge_embedding_loss", + "huber_loss", + "l1_loss", + "layer_norm", + "leaky_relu", + "log_softmax", + "margin_ranking_loss", + "mish", + "mish_", + "mse_loss", + "nll_loss", + "pairwise_distance", + "pdist", + "poisson_nll_loss", + "prelu", + "relu", + "relu6", + "selu", + "selu_", + "smooth_l1_loss", + "softmax", + "softmin", + "softplus", + "softshrink", + "tanhshrink", + "threshold", + "threshold_", + "triplet_margin_loss", +] + +Tensor = torch.Tensor +aten = torch._ops.ops.aten +DispatchKey = torch._C.DispatchKey # type: ignore[attr-defined] + + +def _dropout_helper( + self: TensorLikeType, + val: float, +) -> TensorLikeType: + """ + Helper function for all dropout-type operators. During training, + some of the elements of the input tensor are randomly masked. + + Returns the masked tensor of the boolean values. + + """ + + return ( + refs._uniform_helper( + self.shape, low=0.0, high=1.0, dtype=torch.float32, device=self.device + ) + < val + ) + + +@register_decomposition(aten.alpha_dropout) +def alpha_dropout( + self: TensorLikeType, p: float = 0.5, training: bool = False, inplace: bool = False +) -> TensorLikeType: + if inplace: + raise NotImplementedError + + if not training: + return self + + torch._check( + p <= 1 and p >= 0, + lambda: f"dropout probability has to be between 0 and 1, but got, {p}", + ) + + if p == 1: + return torch.zeros_like(self) + + if p == 0: + return self + + dropout_mask = _dropout_helper(self, 1 - p) + + # From paper: Self-Normalizing Neural Networks (https://arxiv.org/pdf/1706.02515.pdf) + # alpha = - SELU.alpha * SELU.scale, here + # SELU.alpha = 1.6732632423543772848170429916717 and + # SELU.scale = 1.0507009873554804934193349852946 + alpha = -1.7580993408473766 + + a = 1.0 / math.sqrt((alpha * alpha * p + 1) * (1 - p)) + b = torch.logical_not(dropout_mask) + b = b * (alpha * a) + alpha * a * p + dropout_mask = a * dropout_mask + + return self * dropout_mask + b + + +def _inplace_wrapper(fn): + """ + Given a nn.functional non-linearity, implements its `inplace: bool` argument + """ + + # nb. We use the name of the first argument used in the unary references + @wraps(fn) + def _fn(a, *args, inplace=False, **kwargs): + if inplace: + torch._check( + "out" not in kwargs, + lambda: "Cannot set inplace=True and pass out= at the same time", + ) + return fn(a, *args, inplace=False, out=a, **kwargs) + else: + return fn(a, *args, inplace=False, **kwargs) + + return _fn + + +# celu is implemented specially because it has an alpha argument +# celu is very similar to elu +@register_decomposition(aten.celu) +@_inplace_wrapper +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def celu( + a: TensorLikeType, alpha: Optional[NumberType] = None, inplace: bool = False +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.celu + """ + + if inplace: + raise NotImplementedError + + rhs: TensorLikeType + if alpha is not None: + python_type = utils.dtype_to_type(a.dtype) + if not utils.is_weakly_lesser_type(type(alpha), python_type): + msg = f"alpha argument of type {type(alpha)} cannot be safely cast to type {python_type}!" + raise ValueError(msg) + rhs = alpha * torch.expm1(torch.true_divide(a, alpha)) # type: ignore[arg-type] + else: + rhs = torch.expm1(a) + + return torch.where(a > 0, a, rhs) + + +@_inplace_wrapper +@out_wrapper() +def dropout( + a: TensorLikeType, p: float = 0.5, training: bool = True, inplace: bool = False +) -> TensorLikeType: + if inplace: + raise NotImplementedError + + if not training: + return a + + torch._check( + p <= 1 and p >= 0, + lambda: f"dropout probability has to be between 0 and 1, but got, {p}", + ) + + if p == 1: + return torch.zeros_like(a) + + if p == 0: + return a + + scale = 1 / (1 - p) + dropout_mask = _dropout_helper(a, 1 - p) + + return a * dropout_mask * scale + + +@register_decomposition(aten.elu) +@_inplace_wrapper +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def elu( + a: TensorLikeType, + alpha: NumberType = 1.0, + scale: NumberType = 1.0, + input_scale: NumberType = 1.0, + inplace: bool = False, +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.elu + """ + if inplace: + raise NotImplementedError + + # nb. This should be factored out into a can_cast aux function + python_type = utils.dtype_to_type(a.dtype) + torch._check( + utils.is_weakly_lesser_type(type(input_scale), python_type), + lambda: f"input_scale argument of type {type(input_scale)} cannot be safely cast to type {python_type}!", + ) + torch._check( + utils.is_weakly_lesser_type(type(scale), python_type), + lambda: f"scale argument of type {type(scale)} cannot be safely cast to type {python_type}!", + ) + torch._check( + utils.is_weakly_lesser_type(type(alpha), python_type), + lambda: f"alpha argument of type {type(alpha)} cannot be safely cast to type {python_type}!", + ) + + return torch.where(a > 0, scale * a, (alpha * scale) * torch.expm1(a * input_scale)) + + +@register_decomposition(aten.relu) +@_inplace_wrapper +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def relu(a: TensorLikeType, inplace: bool = False) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.relu + """ + + if inplace: + raise NotImplementedError + + return torch.where(torch.le(a, 0), 0, a) + + +@register_decomposition(aten.channel_shuffle) +@out_wrapper() +def channel_shuffle(input: TensorLikeType, groups: int) -> TensorLikeType: + """ + Reference implementation of :func:`torch.nn.functional.channel_shuffle`. + """ + from torch._meta_registrations import device_hint + + torch._check( + input.dim() > 2, + lambda: f"channel_shuffle expects input with > 2 dims, but got input with sizes {list(input.size())}", + ) + c = input.shape[1] + torch._check( + groups > 0, + lambda: f"Number of groups to divide channels in must be positive. Value of groups:{groups}", + ) + torch._check( + (c % groups) == 0, + lambda: f"Number of channels must be divisible by groups. Got {c} channels and {groups} groups.", + ) + n = input.shape[0] + cg = c // groups + dhw = input.shape[2:] + + if input.numel() == 0 or ( + device_hint(input) == "cuda" and (groups == 1 or groups == c) + ): + return input.view(input.shape) + + return ( + input.reshape(n, groups, cg, *dhw) + .transpose(1, 2) + .reshape(input.shape) + .contiguous() + ) + + +def group_norm( + input: Tensor, + num_groups: int, + weight: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + eps: float = 1e-5, +) -> Tensor: + """ + Reference implementation of :func:`torch.nn.functional.group_norm`. + """ + torch._check( + input.ndim >= 2, + lambda: f"Expected at least 2 dimensions for input tensor but received {input.ndim}", + ) + + batch_size = input.shape[0] + num_channels = input.shape[1] + torch._check( + num_channels % num_groups == 0, + lambda: "Expected number of channels in input to be divisible by num_groups, " + + f"but got input of shape {input.shape} and num_groups = {num_groups}", + ) + + # input shape is (N, C, *), so we flatten all inner dimensions except (N, C) + flattened_inner_size = 1 + for dim_length in input.shape[2:]: + flattened_inner_size *= dim_length + + return torch.native_group_norm( + input, + weight, + bias, + batch_size, + num_channels, + flattened_inner_size, + num_groups, + eps, + )[0] + + +def layer_norm( + input: Tensor, + normalized_shape: ShapeType, + weight: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + eps: float = 1e-5, +) -> Tensor: + """ + Reference implementation of :func:`torch.nn.functional.layer_norm`. + """ + return torch.native_layer_norm(input, normalized_shape, weight, bias, eps)[0] + + +@register_decomposition(aten.leaky_relu) +@_inplace_wrapper +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def leaky_relu( + a: TensorLikeType, negative_slope: float = 0.01, inplace: bool = False +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.leaky_relu + """ + + if inplace: + raise NotImplementedError + + python_type = utils.dtype_to_type(a.dtype) + if not utils.is_weakly_lesser_type(type(negative_slope), python_type): + msg = f"negative_slope argument of type {type(negative_slope)} cannot be safely cast to type {python_type}!" + raise ValueError(msg) + return torch.where(torch.gt(a, 0), a, torch.mul(a, negative_slope)) + + +@register_decomposition(aten.mish) +@_inplace_wrapper +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def mish(a: TensorLikeType, inplace: bool = False) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.mish + """ + + if inplace: + raise NotImplementedError + return a * torch.tanh(torch.nn.functional.softplus(a)) + + +@register_decomposition(aten.selu) +@_inplace_wrapper +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def selu(a: TensorLikeType, inplace: bool = False) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.selu + """ + if inplace: + raise NotImplementedError + + alpha = 1.6732632423543772848170429916717 + scale = 1.0507009873554804934193349852946 + + rhs = alpha * torch.expm1(a) + + return scale * torch.where(a > 0, a, rhs) + + +# Forwarding alias: the functional variant doesn't support the out kwarg +# CompositeImplicitAutograd - don't register decomp +def softmax( + a: TensorLikeType, + dim: Optional[int] = None, + _stacklevel: int = 3, # for compat when using TorchRefsMode(strict=True) + dtype: Optional[torch.dtype] = None, +) -> TensorLikeType: + # The error is for compat with regular PyTorch, which has this behavior + # deprecated. For PrimTorch, it's fine to drop support for deprecated + # behavior because it requires explicit opt in. This error is to inform + # users how to update their calls. + torch._check(dim is not None, lambda: "implicit dim not supported, use dim=X") + return torch.softmax(a=a, dim=dim, dtype=dtype) # type: ignore[call-overload] + + +# CompositeImplicitAutograd - don't register decomp +def softmin( + a: TensorLikeType, + dim: Optional[int] = None, + _stacklevel: int = 3, # for compat when using TorchRefsMode(strict=True) + dtype: Optional[torch.dtype] = None, +) -> TensorLikeType: + # The error is for compat with regular PyTorch, which has this behavior + # deprecated. For PrimTorch, it's fine to drop support for deprecated + # behavior because it requires explicit opt in. This error is to inform + # users how to update their calls. + torch._check(dim is not None, lambda: "implicit dim not supported, use dim=X") + return torch.softmax(a=-a, dim=dim, dtype=dtype) # type: ignore[call-overload] + + +# softplus is implemented specially because it has beta and threshold arguments +@register_decomposition(aten.softplus) +@_inplace_wrapper +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def softplus( + a: TensorLikeType, + beta: Optional[NumberType] = None, + threshold: NumberType = 20, + inplace: bool = False, +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.softplus + """ + + if inplace: + raise NotImplementedError + + rhs: TensorLikeType + if beta is not None: + python_type = utils.dtype_to_type(a.dtype) + if not utils.is_weakly_lesser_type(type(beta), python_type): + msg = f"beta argument of type {type(beta)} cannot be safely cast to type {python_type}!" + raise ValueError(msg) + scaled_input = a * beta + rhs = torch.true_divide(torch.log1p(torch.exp(scaled_input)), beta) # type: ignore[arg-type] + + else: + scaled_input = a + rhs = torch.log1p(torch.exp(scaled_input)) + + return torch.where(scaled_input > threshold, a, rhs) + + +@aten.hardshrink.default.py_impl(DispatchKey.Autograd) +@register_decomposition(aten.hardshrink) +@out_wrapper() +def hardshrink(a: TensorLikeType, lambd: float = 0.5): + # Formula for reference, + # hardshrink(x) = x if x > lambd + # = x if x < -lambd + # = 0 otherwise + return torch.where(torch.abs(a) <= lambd, 0, a) + + +@aten.softshrink.default.py_impl(DispatchKey.Autograd) +@register_decomposition(aten.softshrink) +@out_wrapper() +def softshrink(a: TensorLikeType, lambd: float = 0.5): + # Formula for reference, + # softshrink(x) = x - lambd if x > lambd + # = x + lambd if x < -lambd + # = 0 otherwise + torch._check( + lambd >= 0, + lambda: f"lambda must be greater or equal to 0, but found to be {lambd}", + ) + # We implement this in one torch.where to generate better code in the backward + # see https://github.com/pytorch/pytorch/pull/107052#discussion_r1293748211 + # We multiply by 0 for dealing with nans + return torch.where(torch.abs(a) > lambd, a - torch.sign(a) * lambd, a * 0) + + +# Losses +def _reduction_int_to_str(reduction: int) -> str: + from torch._decomp.decompositions import Reduction + + if reduction == Reduction.NONE.value: + return "none" + elif reduction == Reduction.MEAN.value: + return "mean" + elif reduction == Reduction.SUM.value: + return "sum" + else: + raise ValueError(f"{reduction} is not a valid value for reduction") + + +def _apply_loss_reduction(loss: TensorLikeType, reduction: str) -> TensorLikeType: + if reduction == "sum": + return torch.sum(loss) + elif reduction == "mean": + return torch.mean(loss) + else: # reduction == "none" + return loss + + +def _check_reduction_value(reduction: str): + if reduction not in ("mean", "sum", "none"): + raise ValueError(f"{reduction} is not a valid value for reduction") + + +# This helper function maps depreciated arguments, "size_average" and "reduce" +# to their corresponding "reduction" string argument +def _get_string_reduction_arg( + *, size_average: Optional[bool], reduce: Optional[bool] +) -> str: + if size_average is None: + size_average = True + if reduce is None: + reduce = True + if size_average and reduce: + ret = "mean" + elif reduce: + ret = "sum" + else: + ret = "none" + return ret + + +# CompositeImplicitAutograd - don't register decomp +@elementwise_type_promotion_wrapper( + type_promoting_args=("input", "target"), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT, +) +def l1_loss( + input: TensorLikeType, + target: TensorLikeType, + size_average: Optional[bool] = None, + reduce: Optional[bool] = None, + reduction: str = "mean", +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.l1_loss + """ + if size_average is not None or reduce is not None: + # TODO: Raise exception instead of converting value. This is only for + # primTorch since it can drop support for deprecated arguments. + # msg = "size_average and reduce args are deprecated, please use reduction argument." + reduction = _get_string_reduction_arg(size_average=size_average, reduce=reduce) + _check_reduction_value(reduction) + loss = torch.abs(input - target) + return _apply_loss_reduction(loss, reduction) + + +@elementwise_type_promotion_wrapper( + type_promoting_args=("input", "target"), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT, +) +def smooth_l1_loss( + input: TensorLikeType, + target: TensorLikeType, + size_average: Optional[bool] = None, + reduce: Optional[bool] = None, + reduction: str = "mean", + beta: float = 1.0, +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.smooth_l1_loss + """ + if size_average is not None or reduce is not None: + # TODO: Raise exception instead of converting value. This is only for + # primTorch since it can drop support for deprecated arguments. + # msg = "size_average and reduce args are deprecated, please use reduction argument." + reduction = _get_string_reduction_arg(size_average=size_average, reduce=reduce) + _check_reduction_value(reduction) + + if beta == 0.0: + return torch.nn.functional.l1_loss( + input, target, size_average=size_average, reduce=reduce, reduction=reduction + ) + else: + loss = torch.abs(input - target) + loss = torch.where(loss < beta, 0.5 * loss**2 / beta, loss - 0.5 * beta) + return _apply_loss_reduction(loss, reduction) + + +# Forwarding alias: the functional variant doesn't support the out kwarg +# CompositeImplicitAutograd - don't register decomp +def log_softmax( + a: TensorLikeType, + dim: Optional[int] = None, + _stacklevel: int = 3, # for compat when using TorchRefsMode(strict=True) + dtype: Optional[torch.dtype] = None, +) -> TensorLikeType: + # The error is for compat with regular PyTorch, which has this behavior + # deprecated. For PrimTorch, it's fine to drop support for deprecated + # behavior because it requires explicit opt in. This error is to inform + # users how to update their calls. + torch._check(dim is not None, lambda: "implicit dim not supported, use dim=X") + return torch.log_softmax(a=a, dim=dim, dtype=dtype) # type: ignore[call-overload] + + +@register_decomposition(aten.margin_ranking_loss) +def margin_ranking_loss( + input1: TensorLikeType, + input2: TensorLikeType, + target: TensorLikeType, + margin: float = 0.0, + reduction: str = "mean", +) -> TensorLikeType: + # loss_without_reduction = max(0, -target * (input1 - input2) + margin) + if input1.ndim != input2.ndim or input1.ndim != target.ndim: + raise RuntimeError( + "margin_ranking_loss : All input tensors should have same dimension but got sizes: " + f"input1: {input1.shape}, input2: {input2.shape}, target: {target.shape} " + ) + _check_reduction_value(reduction) + loss = torch.clamp_min(-target * (input1 - input2) + margin, 0) + return _apply_loss_reduction(loss, reduction) + + +@elementwise_type_promotion_wrapper( + type_promoting_args=("input", "target"), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT, +) +def mse_loss( + input: TensorLikeType, + target: TensorLikeType, + size_average: Optional[bool] = None, + reduce: Optional[bool] = None, + reduction: str = "mean", +) -> TensorLikeType: + if size_average is not None or reduce is not None: + # TODO: Raise exception instead of converting value. This is only for + # primTorch since it can drop support for deprecated arguments. + # msg = "size_average and reduce args are deprecated, please use reduction argument." + reduction = _get_string_reduction_arg(size_average=size_average, reduce=reduce) + _check_reduction_value(reduction) + loss = torch.pow(input - target, 2) + return _apply_loss_reduction(loss, reduction) + + +@register_decomposition(aten.hinge_embedding_loss) +def hinge_embedding_loss( + input: TensorLikeType, + target: TensorLikeType, + margin: float = 1.0, + reduction: str = "mean", +) -> TensorLikeType: + # loss_without_reduction = input if y == 1 + # = max(0, margin - input) if y == -1 + _check_reduction_value(reduction) + margin_clamp = torch.clamp_min(margin - input, 0) + output_margin = torch.where(target != 1, margin_clamp, 0) + output_self = torch.where(target != -1, input, 0) + loss = output_margin + output_self + return _apply_loss_reduction(loss, reduction) + + +def _nll_loss_nd( + input: TensorLikeType, + target: TensorLikeType, + weight: Optional[TensorLikeType], + reduction: str, + ignore_index: int, +) -> TensorLikeType: + torch._check( + input.ndim > 0 and input.ndim <= 3, + lambda: f"Expected input dimension to be either [1, 2, 3] but received {input.ndim}.", + ) + + torch._check( + (input.ndim == 1) or (input.shape[0] == target.shape[0]), + lambda: f"Expected input batch size {input.shape[0]} to match target batch size {target.shape[0]}.", + ) + + _check_reduction_value(reduction) + + flat_target = torch.flatten(target) + ignore_classes_mask = torch.eq(flat_target, ignore_index) + + # TODO: Enable data-dependent checks with debug mode + # TODO: This check does not work with FakeTensor inputs; See Issue #85834 + # Explicit cast for class_check to bool; See Issue #78071 + """ + from torch._subclasses.fake_tensor import FakeTensor + num_classes = input.shape[1] if input.ndim > 1 else input.shape[0] + valid_classes_mask = torch.logical_and( + (flat_target >= 0), (flat_target < num_classes) + ) + class_check = torch.all(torch.logical_or(ignore_classes_mask, valid_classes_mask)) + torch._check( + isinstance(target, FakeTensor) or bool(class_check.item()), + lambda: "A target class is out-of-bounds and not the ignore index.", + ) + """ + + ignore_class_weight = torch.scalar_tensor(0, dtype=input.dtype, device=input.device) + class_weight = ( + torch.scalar_tensor(1, dtype=input.dtype, device=input.device) + if weight is None + else weight[flat_target] + ) + current_weight = torch.where( + ignore_classes_mask, + ignore_class_weight, + class_weight, + ) + + if input.ndim == 1: + # implicit batch size = 1 + # input (1 batch size, C classes) + loss = -input[target] * current_weight + elif input.ndim == 2: + # input (N batch size, C classes) + batch_size = input.shape[0] + loss = -input[torch.arange(batch_size), target] * current_weight + else: + # 3D case (N batch size, C classe, K dimensions) + # input (N batch size, C classes, K) + batch_size = input.shape[0] + extent = input.shape[2] + numel = batch_size * extent + indices = torch.arange(numel) + bdx = indices // extent + kdx = indices % extent + loss = -input[bdx, flat_target, kdx] * current_weight + loss = torch.reshape(loss, target.shape) + + if reduction == "none": + return loss + elif reduction == "sum": + return torch.sum(loss) + else: + # calculate weighted mean of the loss function + return torch.sum(loss) / torch.sum(current_weight) + + +@register_decomposition(aten.nll_loss) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("input",), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def nll_loss( + input: TensorLikeType, + target: TensorLikeType, + weight: Optional[TensorLikeType] = None, + size_average: Optional[bool] = None, + ignore_index: int = -100, + reduce: Optional[bool] = None, + reduction: str = "mean", +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.nll_loss + """ + torch._check( + input.ndim > 0, + lambda: f"Expected input tensor to have 1 or more dimensions (got {input.ndim})", + ) + + # TODO: raise exception instead of converting value + # msg = "size_average and reduce args are deprecated, please use reduction argument." + # Convert these options for consistency with the eager mode + if size_average is not None or reduce is not None: + reduction = _get_string_reduction_arg(size_average=size_average, reduce=reduce) + + # The expected behavior when the target and input have zero elements: + # reduction = 'none' --- tensor([]) + # reduction = 'sum' --- tensor(0.) + # reduction = 'mean' --- tensor(nan) + # Mean reduction on empty tensors produces NaN. See the discussion in + # https://github.com/pytorch/pytorch/pull/64572#issuecomment-926504162 + if input.numel() == 0 and target.numel() == 0: + if reduction == "none": + return torch.zeros_like(target) + elif reduction == "sum": + return torch.empty_like(target) + else: + return torch.full_like(target, float("nan")) + + # The _nll_loss_nd helper function handles the most common cases. + # ndim == 1 (Single Example) + # => Batch Size: 1, Input: (C), Target: () + # ndim == 2 (k = 1) + # => Batch Size: N, Input: (N, C), Target: (N) + # ndim == 3 (k > 1) + # => Batch Size: N, Input: (N, C, K), Target: (N, K) + if input.ndim <= 3: + return _nll_loss_nd(input, target, weight, reduction, ignore_index) + + # For ndim > 3, we reshape the input and target to 3-D case. + # Input (N batch-size, C classes, k-dimensions) + # Target (N batch-size, k-dimensions) + torch._check( + input.ndim > 0 and target.ndim > 0 and target.shape[1:] == input.shape[2:], + lambda: ( + "Expected input and target to both have ndim > 0 and " + "target.shape[1:] == input.shape[2:], but got " + f"target.shape {target.shape} and input.shape {input.shape}" + ), + ) + + batch_size = input.shape[0] + num_classes = input.shape[1] + out_size = [batch_size] + list(target.shape[1:]) + + input = torch.reshape(input, [batch_size, num_classes, -1]) + target = torch.reshape(target, [batch_size, -1]) + if reduction != "none": + return _nll_loss_nd(input, target, weight, reduction, ignore_index) + else: + result = _nll_loss_nd(input, target, weight, reduction, ignore_index) + # reshape flattened inner-dim to original k-dimensions + return torch.reshape(result, out_size) + + +# TODO: This ref supports int reduction and out kwarg to be compatible with ATen: +# https://github.com/pytorch/pytorch/issues/83931 +# TODO: Could be rewritten to support complex: +# https://github.com/pytorch/pytorch/pull/85041 +@register_decomposition(aten.huber_loss) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("input", "target"), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def huber_loss( + input: TensorLikeType, + target: TensorLikeType, + reduction: Union[str, int] = "mean", + delta: float = 1.0, +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.huber_loss + """ + if type(reduction) is int: + reduction = _reduction_int_to_str(reduction) + _check_reduction_value(reduction) # type: ignore[arg-type] + torch._check( + delta > 0, + lambda: "huber_loss does not support non-positive values for delta.", + ) + z = (input - target).abs() + loss = torch.where(z < delta, 0.5 * z * z, delta * (z - 0.5 * delta)) + return _apply_loss_reduction(loss, reduction) # type: ignore[arg-type] + + +# tanhshrink does not use _make_elementwise_unary_reference because it does not support out +@elementwise_unary_scalar_wrapper +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def tanhshrink(a: TensorLikeType) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.tanhshrink + """ + if not isinstance(a, TensorLike): + raise RuntimeError( + "Expected a tensor input for an elementwise unary operation!" + ) + return a - torch.tanh(a) + + +@register_decomposition(aten.threshold) +@_inplace_wrapper +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def threshold( + a: TensorLikeType, + threshold: NumberType, + value: Union[bool, int, float], + inplace: bool = False, +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.threshold + """ + + if inplace: + raise NotImplementedError + + return torch.where(a <= threshold, value, a) + + +# CompositeImplicitAutograd - don't register decomp +# No elementwise type promotion - core op doesn't explicitly type promote +def triplet_margin_loss( + anchor: TensorLikeType, + positive: TensorLikeType, + negative: TensorLikeType, + margin: float = 1.0, + p: float = 2, + eps: float = 1e-6, + swap: bool = False, + size_average: Optional[bool] = None, + reduce: Optional[bool] = None, + reduction: str = "mean", +) -> TensorLikeType: + if size_average is not None or reduce is not None: + # TODO: Raise exception instead of converting value. This is only for + # primTorch since it can drop support for deprecated arguments. + # msg = "size_average and reduce args are deprecated, please use reduction argument." + reduction = _get_string_reduction_arg(size_average=size_average, reduce=reduce) + + if margin <= 0: + raise ValueError(f"margin must be greater than 0, got {margin}") + + # torch.nn.functional.triplet_margin_with_distance_loss has no ref defined + # since it's a pure Python implementation. Use this helper instead. + return _triplet_margin_with_distance_loss( + anchor=anchor, + positive=positive, + negative=negative, + distance_function=lambda x, y: torch.pairwise_distance(x, y, p, eps), + margin=margin, + swap=swap, + reduction=reduction, + ) + + +# Pure Python impl - don't register decomp and don't add a ref. Defined as a +# helper here since triplet_margin_loss can be nicely implemented with it. +def _triplet_margin_with_distance_loss( + anchor: TensorLikeType, + positive: TensorLikeType, + negative: TensorLikeType, + *, + distance_function: Optional[ + Callable[[TensorLikeType, TensorLikeType], TensorLikeType] + ] = None, + margin: float = 1.0, + swap: bool = False, + reduction: str = "mean", +) -> TensorLikeType: + _check_reduction_value(reduction) + + a_dim = anchor.ndim + p_dim = positive.ndim + n_dim = negative.ndim + torch._check( + a_dim == p_dim and p_dim == n_dim, + lambda: ( + f"The anchor, positive, and negative tensors are expected to have " + f"the same number of dimensions, but got: anchor {a_dim}D, " + f"positive {p_dim}D, and negative {n_dim}D inputs" + ), + ) + + if distance_function is None: + distance_function = torch.pairwise_distance + + dist_pos = distance_function(anchor, positive) + dist_neg = distance_function(anchor, negative) + # The distance swap is described in the paper "Learning shallow + # convolutional feature descriptors with triplet losses" by V. Balntas, E. + # Riba et al. If True, and if the positive example is closer to the + # negative example than the anchor is, swaps the positive example and the + # anchor in the loss computation. + if swap: + dist_swap = distance_function(positive, negative) + dist_neg = torch.minimum(dist_neg, dist_swap) + loss = torch.clamp_min(margin + dist_pos - dist_neg, 0) + return _apply_loss_reduction(loss, reduction) + + +@register_decomposition(aten.hardtanh) +@_inplace_wrapper +@out_wrapper() +@elementwise_unary_scalar_wrapper +@elementwise_type_promotion_wrapper( + type_promoting_args=("a"), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def hardtanh( + a: TensorLikeType, + min_val: NumberType = -1, + max_val: NumberType = 1, + inplace: bool = False, +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.hardtanh + """ + if inplace: + raise NotImplementedError + if utils.is_boolean_dtype(a.dtype): + raise RuntimeError("Bool inputs not supported for hardtanh") + + # preserve legacy behavior of boundaries not causing type promotion + if utils.is_integer_dtype(a.dtype): + min_val = int(min_val) # type: ignore[arg-type] + max_val = int(max_val) # type: ignore[arg-type] + if not (a.dtype != torch.uint8 or (min_val >= 0 and max_val >= 0)): + raise RuntimeError( + "Cannot do hardtanh on an unsigned type with negative limits" + ) + + if min_val > max_val: # type: ignore[operator] + raise ValueError("min_val cannot be greater than max_val") + + return torch.clamp(a, min_val, max_val) # type: ignore[arg-type] + + +@register_decomposition(aten.gelu) +@out_wrapper() +@elementwise_unary_scalar_wrapper +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def gelu(a: TensorLikeType, approximate: str = "none") -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.gelu + """ + if not isinstance(a, TensorLike): + raise RuntimeError( + "Expected a tensor input for an elementwise unary operation!" + ) + M_SQRT2 = 1.41421356237309504880 + M_SQRT1_2 = 0.70710678118654752440 + M_2_SQRTPI = 1.12837916709551257390 + if approximate == "tanh": + kBeta = M_SQRT2 * M_2_SQRTPI * 0.5 + kKappa = 0.044715 + a_cube = a * a * a + inner = kBeta * (a + kKappa * a_cube) + return 0.5 * a * (1 + torch.tanh(inner)) + elif approximate == "none": + kAlpha = M_SQRT1_2 + return a * 0.5 * (1 + torch.erf(a * kAlpha)) + else: + raise RuntimeError("approximate argument must be either none or tanh.") + + +# CompositeImplicitAutograd - don't register decomp +@elementwise_type_promotion_wrapper( + type_promoting_args=("input", "target"), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def poisson_nll_loss( + input: TensorLikeType, + target: TensorLikeType, + log_input: bool = True, + full: bool = False, + size_average: Optional[bool] = None, + eps: float = 1e-8, + reduce: Optional[bool] = None, + reduction: str = "mean", +) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.poisson_nll_loss + """ + if size_average is not None or reduce is not None: + # TODO: Raise exception instead of converting value. This is only for + # primTorch since it can drop support for deprecated arguments. + # msg = "size_average and reduce args are deprecated, please use reduction argument." + reduction = _get_string_reduction_arg(size_average=size_average, reduce=reduce) + _check_reduction_value(reduction) + if log_input: + loss = torch.exp(input) - target * input + else: + loss = input - target * torch.log(input + eps) + + if full: + stirling_term = ( + target * torch.log(target) - target + 0.5 * torch.log(2 * torch.pi * target) + ) + # avoid inplace add + loss = loss + stirling_term.masked_fill(target <= 1, 0) + return _apply_loss_reduction(loss, reduction) + + +@register_decomposition(aten.prelu) +@elementwise_type_promotion_wrapper( + type_promoting_args=("a", "weight"), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def prelu(a: TensorLikeType, weight: TensorLikeType) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.prelu + """ + torch._check( + isinstance(a, TensorLike), + lambda: f"prelu: Expected `a` to be tensor, but got: {type(a)}", + ) + torch._check( + isinstance(weight, TensorLike), + lambda: f"prelu: Expected `weight` to be tensor, but got: {type(weight)}", + ) + + if weight.numel() != 1: + torch._check(a.ndim > 0, lambda: "Not allow zero-dim input tensor.") + channel_size = a.shape[1] if a.ndim >= 2 else 1 + torch._check( + weight.numel() == channel_size, + lambda: f"Mismatch of parameter numbers and input channel size. Found parameter numbers =" + f" {weight.numel()} and channel size = {channel_size}.", + ) + + torch._check( + weight.ndim == 0 or weight.ndim == 1, + lambda: f"prelu: Expected `weight` to be a scalar or 1D tensor, but got: " + f"ndim = {weight.ndim}", + ) + if a.ndim == 0: + weight = weight[0] if weight.ndim == 1 else weight + else: + weight = prims.broadcast_in_dim( + weight, a.shape, () if weight.ndim == 0 else (0 if a.ndim == 1 else 1,) + ) + + return torch.where(a > 0, a, a * weight) + + +@register_decomposition(aten.relu6) +@_inplace_wrapper +@out_wrapper() +def relu6(a: TensorLikeType, inplace: bool = False) -> TensorLikeType: + """ + Reference implementation of torch.nn.functional.relu6 + """ + if inplace: + raise NotImplementedError + + # See https://github.com/pytorch/pytorch/pull/81142#discussion_r918220126 + # It may be better to use clamp here, but we use hardtanh to replicate + # the behavior of the existing implementation + return torch.nn.functional.hardtanh(a, 0, 6) + + +@register_decomposition(aten.glu) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def glu(a: TensorLikeType, dim: int = -1) -> TensorLikeType: + dim = utils.canonicalize_dims(a.ndim, dim) + torch._check( + a.shape[dim] % 2 == 0, + lambda: f"Halving dimension must be even, but dimension {dim} is size {a.shape[dim]}", + ) + b, c = torch.tensor_split(a, 2, dim) + + return b * torch.sigmoid(c) + + +@register_decomposition(aten.pairwise_distance) +@out_wrapper() +def pairwise_distance( + x1: TensorLikeType, + x2: TensorLikeType, + p: NumberType = 2.0, + eps: NumberType = 1e-6, + keepdim=False, +) -> TensorLikeType: + return torch.linalg.vector_norm(x1 - x2 + eps, ord=p, dim=-1, keepdim=keepdim) + + +@register_decomposition(aten.pdist) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, +) +def pdist(a: TensorLikeType, p: float = 2) -> TensorLikeType: + torch._check(a.ndim == 2, lambda: f"pdist only supports 2D tensors, got: {a.ndim}D") + torch._check(p >= 0, lambda: "pdist only supports non-negative p values") + # For p == 2 we can use an efficient implementation, but other values of p + # require creating a much bigger tensor for an intermediate step + if p == 2: + aTa = torch.mm(a, a.T) + aTa_diag = torch.diag(aTa) + t = torch.sqrt(torch.clamp(aTa_diag + aTa_diag.unsqueeze(-1) - 2 * aTa, min=0)) + else: + t = torch.linalg.vector_norm(a.unsqueeze(1) - a, ord=p, dim=2) + i = torch.triu_indices(t.shape[0], t.shape[1], offset=1, device=a.device) + return t.flatten().index_select(0, i[0] * t.shape[0] + i[1]) + + +@register_decomposition(aten.pixel_shuffle) +@out_wrapper() +def pixel_shuffle(self: Tensor, upscale_factor: int): + torch._check( + self.dim() >= 3, + lambda: f"pixel_shuffle expects input to have at least 3 dimensions, but got input with {self.dim} dimension(s)", + ) + batch = self.shape[:-3] + C_out = self.shape[-3] // upscale_factor**2 + HW_out = (self.shape[-2] * upscale_factor, self.shape[-1] * upscale_factor) + n = len(batch) + B_dims = range(n) + C_dim, r1_dim, r2_dim, H_dim, W_dim = range(n, n + 5) + return ( + self.view( + *batch, + C_out, + upscale_factor, + upscale_factor, + self.shape[-2], + self.shape[-1], + ) + .permute(*B_dims, C_dim, H_dim, r1_dim, W_dim, r2_dim) + .reshape(*batch, C_out, *HW_out) + .clone(memory_format=utils.suggest_memory_format(self)) + ) + + +@register_decomposition(aten.pixel_unshuffle) +@out_wrapper() +def pixel_unshuffle(self: Tensor, downscale_factor: int): + torch._check( + self.dim() >= 3, + lambda: f"pixel_unshuffle expects input to have at least 3 dimensions, but got input with {self.dim} dimension(s)", + ) + batch = self.shape[:-3] + C_out = self.shape[-3] * downscale_factor**2 + HW_out = (self.shape[-2] // downscale_factor, self.shape[-1] // downscale_factor) + n = len(batch) + B_dims = range(n) + C_dim, H_dim, r1_dim, W_dim, r2_dim = range(n, n + 5) + return ( + self.view( + *batch, + self.shape[-3], + HW_out[0], + downscale_factor, + HW_out[1], + downscale_factor, + ) + .permute(*B_dims, C_dim, r1_dim, r2_dim, H_dim, W_dim) + .reshape(*batch, C_out, *HW_out) + .clone(memory_format=utils.suggest_memory_format(self)) + ) + + +# Needed as aten.{celu_,elu_...} exist (even if they don't have the in-place kwarg) +celu_ = _make_inplace(celu) +elu_ = _make_inplace(elu) +mish_ = _make_inplace(mish) +selu_ = _make_inplace(selu) +threshold_ = _make_inplace(threshold) diff --git a/lib/python3.12/site-packages/torch/_refs/nn/functional/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_refs/nn/functional/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7896c5a0d1f43f8628f18caa6e65308868b8cfa2 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_refs/nn/functional/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_refs/special/__init__.py b/lib/python3.12/site-packages/torch/_refs/special/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3599f2fd85c1406c006640ea3e994b9864b605bb --- /dev/null +++ b/lib/python3.12/site-packages/torch/_refs/special/__init__.py @@ -0,0 +1,236 @@ +# mypy: allow-untyped-defs +import math +from typing import Optional, Union + +import torch +import torch._prims as prims +import torch._prims_common as utils +import torch._refs as refs +from torch import Tensor +from torch._decomp import register_decomposition +from torch._prims_common import ( + ELEMENTWISE_TYPE_PROMOTION_KIND, + Number, + NumberType, + TensorLike, + TensorLikeType, +) +from torch._prims_common.wrappers import elementwise_type_promotion_wrapper, out_wrapper +from torch._refs import ( + _make_alias, + _make_elementwise_binary_reference, + _make_elementwise_unary_reference, +) + + +__all__ = [ + "bessel_j0", + "bessel_j1", + "entr", + "erfcx", + "expit", + "i0e", + "i1", + "i1e", + "log_ndtr", + "logit", + "log_softmax", + "multigammaln", + "ndtr", + "ndtri", + "softmax", + "spherical_bessel_j0", + "xlog1py", + "zeta", +] +aten = torch._ops.ops.aten + + +@_make_elementwise_unary_reference( + ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def bessel_j0(a: TensorLikeType) -> TensorLikeType: + return prims.bessel_j0(a) + + +@_make_elementwise_unary_reference( + ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def bessel_j1(a: TensorLikeType) -> TensorLikeType: + return prims.bessel_j1(a) + + +@register_decomposition(aten.special_entr) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def entr(a: TensorLikeType) -> TensorLikeType: + return torch.where( + torch.isnan(a), + a, + torch.where(a > 0, -a * torch.log(a), torch.where(a == 0, 0, -torch.inf)), + ) + + +@register_decomposition(aten.special_erfcx) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def erfcx(a: TensorLikeType) -> TensorLikeType: + return prims.erfcx(a) + + +# alias for sigmoid +expit = _make_alias(torch.sigmoid, "expit") + + +@_make_elementwise_unary_reference( + ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def i0e(a: TensorLikeType) -> TensorLikeType: + return prims.bessel_i0e(a) + + +@_make_elementwise_unary_reference( + ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def i1(a: TensorLikeType) -> TensorLikeType: + return prims.bessel_i1(a) + + +@_make_elementwise_unary_reference( + ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def i1e(a: TensorLikeType) -> TensorLikeType: + return prims.bessel_i1e(a) + + +@register_decomposition(aten.special_log_ndtr) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def log_ndtr(a: TensorLikeType) -> TensorLikeType: + # Note: M_SQRT1_2 is the value of 1 / sqrt(2) + M_SQRT1_2 = 0.707106781186547524400844362104849039 + t = a * M_SQRT1_2 + return torch.where( + a < 1.0, + torch.log(torch.special.erfcx(-t) / 2) - t * t, + torch.log1p(-torch.erfc(t) / 2), + ) + + +@register_decomposition(aten.logit) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("self",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def logit(self: TensorLikeType, eps: Optional[float] = None) -> TensorLikeType: + if eps is None: + eps = -1.0 + lo = eps + hi = 1 - eps + self = torch.clamp(self, lo, hi) + return torch.log(torch.true_divide(self, torch.sub(1, self))) + + +@register_decomposition(aten.special_xlog1py) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a", "b"), + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def xlog1py(a: Union[TensorLikeType, NumberType], b: Union[TensorLikeType, NumberType]): + torch._check( + isinstance(a, TensorLike) or isinstance(b, TensorLike), + lambda: 'Expected either argument a or b to be a Tensor"', + ) + + # Operations like eq and log do not handle scalar values, so we convert them to scalar_tensors. + if isinstance(a, TensorLike) and isinstance(b, Number): + b = refs.scalar_tensor(b, dtype=a.dtype, device=a.device) + elif isinstance(b, TensorLike) and isinstance(a, Number): + a = refs.scalar_tensor(a, dtype=b.dtype, device=b.device) + + # mypy: expected "Tensor" + assert isinstance(a, TensorLike) + assert isinstance(b, TensorLike) + rhs = torch.where(torch.eq(a, 0), 0, torch.mul(a, torch.log1p(b))) + return torch.where(torch.isnan(b), float("nan"), rhs) + + +@register_decomposition(aten.mvlgamma) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def multigammaln(a: TensorLikeType, p: int) -> TensorLikeType: + c = 0.25 * p * (p - 1) * math.log(math.pi) + b = 0.5 * torch.arange(start=(1 - p), end=1, step=1, dtype=a.dtype, device=a.device) + return torch.sum(torch.lgamma(a.unsqueeze(-1) + b), dim=-1) + c + + +@register_decomposition(aten.special_ndtr) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def ndtr(a: TensorLikeType) -> TensorLikeType: + # Note: M_SQRT1_2 is the value of 1 / sqrt(2) + M_SQRT1_2 = 0.707106781186547524400844362104849039 + a_sqrt_2 = a * M_SQRT1_2 + return (1 + torch.erf(a_sqrt_2)) * 0.5 + + +@register_decomposition(aten.special_ndtri) +@out_wrapper() +@elementwise_type_promotion_wrapper( + type_promoting_args=("a",), + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def ndtri(a: TensorLikeType) -> TensorLikeType: + return prims.ndtri(a) + + +# Forwarding alias: the special variant doesn't support the out kwarg +# CompositeImplicitAutograd - don't register decomp +def log_softmax( + a: TensorLikeType, + dim: int, + dtype: Optional[torch.dtype] = None, +) -> TensorLikeType: + return torch.log_softmax(a=a, dim=dim, dtype=dtype) # type: ignore[call-overload] + + +# Forwarding alias: the special variant doesn't support the out kwarg +# CompositeImplicitAutograd - don't register decomp +def softmax( + a: TensorLikeType, + dim: int, + dtype: Optional[torch.dtype] = None, +) -> TensorLikeType: + return torch.softmax(a=a, dim=dim, dtype=dtype) # type: ignore[call-overload] + + +@_make_elementwise_unary_reference( + ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def spherical_bessel_j0(a: TensorLikeType) -> TensorLikeType: + return prims.spherical_bessel_j0(a) + + +# TODO: add docstring +@_make_elementwise_binary_reference( + type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, +) +def zeta(a: TensorLikeType, b: TensorLikeType) -> TensorLikeType: + return prims.zeta(a, b) diff --git a/lib/python3.12/site-packages/torch/_refs/special/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_refs/special/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c1fc051abbc119bca0c4fc029dc6a788cfdb8c33 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_refs/special/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_strobelight/__init__.py b/lib/python3.12/site-packages/torch/_strobelight/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/_strobelight/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_strobelight/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f6a0b9d94160d0730be4993d715b649cb7ffffda Binary files /dev/null and b/lib/python3.12/site-packages/torch/_strobelight/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_strobelight/__pycache__/cli_function_profiler.cpython-312.pyc b/lib/python3.12/site-packages/torch/_strobelight/__pycache__/cli_function_profiler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40071ee1977a17621a60811337b644acc974809c Binary files /dev/null and b/lib/python3.12/site-packages/torch/_strobelight/__pycache__/cli_function_profiler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_strobelight/__pycache__/compile_time_profiler.cpython-312.pyc b/lib/python3.12/site-packages/torch/_strobelight/__pycache__/compile_time_profiler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f406a3979d7ccbec05f0a570dfb219ffe967854 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_strobelight/__pycache__/compile_time_profiler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_strobelight/cli_function_profiler.py b/lib/python3.12/site-packages/torch/_strobelight/cli_function_profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..72e44ec42ac3a47d447ababc8185d35636f8360a --- /dev/null +++ b/lib/python3.12/site-packages/torch/_strobelight/cli_function_profiler.py @@ -0,0 +1,311 @@ +# mypy: disallow-untyped-defs + +import functools +import logging +import os +import re +import subprocess +import time +from threading import Lock +from timeit import default_timer as timer +from typing import Any, List, Optional, Sequence + + +logger = logging.getLogger("strobelight_function_profiler") + +console_handler = logging.StreamHandler() +formatter = logging.Formatter( + "%(name)s, line %(lineno)d, %(asctime)s, %(levelname)s: %(message)s" +) +console_handler.setFormatter(formatter) + +logger.addHandler(console_handler) +logger.setLevel(logging.INFO) +logger.propagate = False + + +class StrobelightCLIProfilerError(Exception): + """ + Raised when an error happens during strobelight profiling + """ + + +def _pid_namespace_link(pid: Optional[int] = None) -> str: + """Returns the link to the process's namespace, example: pid:[4026531836]""" + PID_NAMESPACE_PATH = "/proc/{}/ns/pid" + pid = pid or os.getpid() + return os.readlink(PID_NAMESPACE_PATH.format(pid)) + + +def _pid_namespace(pid: Optional[int] = None) -> int: + """Returns the process's namespace id""" + pid = pid or os.getpid() + link = _pid_namespace_link(pid) + return int(link[link.find("[") + 1 : -1]) + + +def _command_to_string(command: Sequence[str]) -> str: + return " ".join(command) + + +class StrobelightCLIFunctionProfiler: + """ + Note: this is a Meta only tool. + + StrobelightCLIFunctionProfiler can be used to profile a python function and + generate a strobelight link with the results. It works on meta servers but + does not requries an fbcode target. + When stop_at_error is false(default), error during profiling does not prevent + the work function from running. + + Check function_profiler_example.py for an example. + """ + + # This lock is used to make sure only one thread is running the profiler at any point. + _lock = Lock() + + def __init__( + self, + *, + stop_at_error: bool = False, + max_profile_duration_sec: int = 60 * 10, + sample_each: float = 1e7, # sample each sample_each cycles. + run_user_name: str = "pytorch-strobelight-ondemand", + timeout_wait_for_running_sec: int = 60, + timeout_wait_for_finished_sec: int = 60, + recorded_env_variables: Optional[List[str]] = None, + sample_tags: Optional[List[str]] = None, + stack_max_len: int = 127, + async_stack_max_len: int = 127, + ): + self.stop_at_error = stop_at_error + self.max_profile_duration_sec = max_profile_duration_sec + self.sample_each = sample_each + self.run_user_name = run_user_name + self.timeout_wait_for_running_sec = timeout_wait_for_running_sec + self.timeout_wait_for_finished_sec = timeout_wait_for_finished_sec + # Results of the most recent run. + # Tracks the strobelight run id of the most recent run + self.current_run_id: Optional[int] = None + self.profile_result: Optional[List[str]] = None + self.sample_tags = sample_tags + + def _run_async(self) -> None: + processId = os.getpid() + namespace = _pid_namespace(processId) + command = [ + "strobeclient", + "run", + "--profiler", + "pyperf", + "--event", + "cycles", + "--async", + "--sample-interval", + f"{int(self.sample_each)}", + "--duration-ms", + f"{int(self.max_profile_duration_sec * 1000)}", + "--pid", + f"{namespace}:{processId}", + ] + + if self.sample_tags: + command.append("--sample-tags") + command.append(",".join(self.sample_tags)) + + logger.debug("running command: %s", _command_to_string(command)) + result = subprocess.run(command, capture_output=True) + output = result.stderr.decode("utf-8") + logger.debug("output:\n{%s}", output) + + if result.returncode != 0: + raise StrobelightCLIProfilerError( + f"failed to start strobelight profiling, error in run_async:{output}" + ) + + if match := re.search(r"INFO Run Id: (-?\d+)", output): + self.current_run_id = int(match.group(1)) + return + + raise StrobelightCLIProfilerError( + f"failed to start strobelight profiling, unexpected result {output}" + ) + + def _wait_for_running(self, counter: int = 0) -> None: + if counter > 20: + raise StrobelightCLIProfilerError( + "wait_for_running called more than 20 times" + ) + + command = ["strobeclient", "getRunStatus", "--run-id", f"{self.current_run_id}"] + logger.debug("running command: %s", _command_to_string(command)) + result = subprocess.run(command, capture_output=True) + output = result.stderr.decode("utf-8") + logger.debug("output:\n{%s}", output) + + if result.returncode != 0: + raise StrobelightCLIProfilerError( + f"failed to start strobelight profiling, error in wait_for_running:{output}" + ) + + if match := re.search("Profile run status: (.*)", output): + current_status = match.group(1) + if current_status == "RUNNING": + return + elif current_status == "PREPARING": + time.sleep(10) + self._wait_for_running(counter + 1) + return + else: + raise StrobelightCLIProfilerError(f"unexpected {current_status} phase") + + raise StrobelightCLIProfilerError(f"unexpected output\n: {output} ") + + def _stop_run(self) -> None: + command = ["strobeclient", "stopRun", "--run-id", str(self.current_run_id)] + logger.debug("running command: %s", _command_to_string(command)) + result = subprocess.run(command, capture_output=True) + output = result.stderr.decode("utf-8") + logger.debug("output:\n{%s}", output) + + if result.returncode != 0: + raise StrobelightCLIProfilerError( + f"failed to stop strobelight profiling, return code is not 0 :{output}" + ) + + if match := re.search("INFO ::1:(.*)", output): + current_status = match.group(1) + if current_status.__contains__("Success!"): + return + else: + raise StrobelightCLIProfilerError( + f"failed to stop strobelight profiling, got {current_status} result" + ) + + raise StrobelightCLIProfilerError(f"unexpected output\n: {output} ") + + def _get_results(self) -> None: + command = ["strobeclient", "getRunStatus", "--run-id", str(self.current_run_id)] + logger.debug("running command: %s", _command_to_string(command)) + result = subprocess.run(command, capture_output=True) + output = result.stderr.decode("utf-8") + logger.debug("output:\n{%s}", output) + + if result.returncode != 0: + raise StrobelightCLIProfilerError( + f"failed to extract profiling results, return code is not 0 : {output}" + ) + + if match := re.search("INFO ::1:(.*)", output): + current_status = match.group(1) + if current_status.__contains__("Profile run status: PROCESSING"): + time.sleep(10) + self._get_results() + return + elif not current_status.__contains__("Profile run finished with SUCCESS"): + raise StrobelightCLIProfilerError( + f"failed to extract profiling results, unexpected response {output}" + ) + + self.profile_result = [] + for item in re.findall( + r"(Total samples(.*)|GraphProfiler(.*)|Icicle view \(python stack\)(.*))", + output, + ): + self.profile_result += item[0] + logger.info(item[0]) + + def _stop_strobelight_no_throw( + self, + collect_results: bool, + ) -> None: + try: + # call stop run + self._stop_run() + logger.info("strobelight profiling stopped") + + logger.debug("collection stopped") + + if not collect_results: + return + + self._get_results() + except Exception as error: + logger.warning("error during stop_strobelight", exc_info=True) + + # Return true if strobelight started and is running. Never throw. + def _start_strobelight(self) -> bool: + strobelight_started = False + try: + self._run_async() + strobelight_started = True + logger.info("strobelight run id is: %s", self.current_run_id) + self._wait_for_running() + logger.info("strobelight profiling running") + return True + + except Exception as error: + logger.warning("error during start_strobelight:", exc_info=True) + if strobelight_started: + self._stop_strobelight_no_throw(collect_results=False) + return False + + def profile(self, work_function: Any, *args: Any, **kwargs: Any) -> Any: + self.current_run_id = None + self.profile_result = None + + if locked := StrobelightCLIFunctionProfiler._lock.acquire(False): + if not locked: + if self.stop_at_error: + raise StrobelightCLIProfilerError("concurrent runs not supported") + + logger.warning("concurrent runs not supported") + return work_function(*args, **kwargs) + + started = self._start_strobelight() + if not started: + if self.stop_at_error: + StrobelightCLIFunctionProfiler._lock.release() + raise StrobelightCLIProfilerError( + "failed to start strobelight profiling" + ) + result = work_function(*args, **kwargs) + StrobelightCLIFunctionProfiler._lock.release() + return result + + try: + logger.debug("collection started") + start = timer() + result = work_function(*args, **kwargs) + end = timer() + total_time = end - start # Time in seconds, e.g. 5.38091952400282 + logger.info("work function took %s seconds", total_time) + self._stop_strobelight_no_throw(collect_results=True) + StrobelightCLIFunctionProfiler._lock.release() + return result + except Exception as error: + logger.warning("work function throw exception", exc_info=True) + self._stop_strobelight_no_throw(collect_results=False) + StrobelightCLIFunctionProfiler._lock.release() + raise error + + +# A function decorator that wraps profile, if no profiler is provided one with +# default args is created. A function can be annotated as: +# @strobelight() +# @strobelight(profiler = StrobelightFunctionProfiler(stop_at_error=True,..)) +# @strobelight(stop_at_error=True,...) +def strobelight( + profiler: Optional[StrobelightCLIFunctionProfiler] = None, **kwargs: Any +) -> Any: + if not profiler: + profiler = StrobelightCLIFunctionProfiler(**kwargs) + + def strobelight_inner(work_function: Any) -> Any: + @functools.wraps(work_function) + def wrapper_function(*args: Any, **kwargs: Any) -> Any: + return profiler.profile(work_function, *args, **kwargs) + + return wrapper_function + + return strobelight_inner diff --git a/lib/python3.12/site-packages/torch/_strobelight/compile_time_profiler.py b/lib/python3.12/site-packages/torch/_strobelight/compile_time_profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..81ebef2df6b13c3bedb7c9f8b53fd87dce1ca3d4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_strobelight/compile_time_profiler.py @@ -0,0 +1,195 @@ +# mypy: disallow-untyped-defs + +import json +import logging +import os +import subprocess +from datetime import datetime +from socket import gethostname +from typing import Any, Optional + +from torch._strobelight.cli_function_profiler import StrobelightCLIFunctionProfiler + + +logger = logging.getLogger("strobelight_compile_time_profiler") + +console_handler = logging.StreamHandler() +formatter = logging.Formatter( + "%(name)s, line %(lineno)d, %(asctime)s, %(levelname)s: %(message)s" +) +console_handler.setFormatter(formatter) + +logger.addHandler(console_handler) +logger.setLevel(logging.INFO) +logger.propagate = False + + +def get_fburl(url: str) -> str: + short_url = url + # Attempt to shorten the URL + try: + result = subprocess.run( + ["fburl", url], capture_output=True, stdin=subprocess.DEVNULL + ) + if result.returncode == 0: + short_url = result.stdout.decode("utf-8") + except Exception as e: + logger.warning("URL shortening failed: %s, using long URL", repr(e)) + return short_url + + +def get_strobelight_url(identifier: str) -> str: + scuba_json = { + "aggregateList": [], + "aggregation_field": "async_stack_complete", + "b_constraints": [[]], + "c_constraints": [[]], + "cols": ["namespace_id", "namespace_process_id"], + "compare": "none", + "constraints": [ + [{"column": "sample_tags", "op": "all", "value": [f'["{identifier}"]']}] + ], + "derivedCols": [], + "end": "now", + "enumCols": [], + "filterMode": "DEFAULT", + "hideEmptyColumns": "false", + "ignoreGroupByInComparison": "false", + "is_timeseries": "false", + "mappedCols": [], + "metric": "count", + "modifiers": [], + "order": "weight", + "order_desc": "true", + "param_dimensions": [ + {"dim": "py_async_stack", "op": "edge", "param": "0", "anchor": "0"} + ], + "purposes": [], + "return_remainder": "false", + "samplingRatio": "1", + "should_pivot": "false", + "start": "-30 days", + "timezone": "America/Los_Angeles", + "top": 10000, + } + scuba_url_prefix = "https://www.internalfb.com/intern/scuba/query/?dataset=pyperf_experimental/on_demand&drillstate=" + scuba_url_suff = "&view=GraphProfilerView&&normalized=1726332703&pool=uber" + long_url = scuba_url_prefix + json.dumps(scuba_json) + scuba_url_suff + return get_fburl(long_url) + + +class StrobelightCompileTimeProfiler: + success_profile_count: int = 0 + failed_profile_count: int = 0 + ignored_profile_runs: int = 0 + inside_profile_compile_time: bool = False + enabled: bool = False + # A unique identifier that is used as the run_user_name in the strobelight profile to + # associate all compile time profiles together. + identifier: Optional[str] = None + + current_phase: Optional[str] = None + + profiler: Optional[Any] = None + + max_stack_length: int = int( + os.environ.get("COMPILE_STROBELIGHT_MAX_STACK_LENGTH", 500) + ) + max_profile_time: int = int( + os.environ.get("COMPILE_STROBELIGHT_MAX_PROFILE_TIME", 60 * 30) + ) + # Collect sample each x cycles. + sample_each: int = int( + float(os.environ.get("COMPILE_STROBELIGHT_SAMPLE_RATE", 1e7)) + ) + + @classmethod + def enable(cls, profiler_class: Any = StrobelightCLIFunctionProfiler) -> None: + if cls.enabled: + logger.info("compile time strobelight profiling already enabled") + return + + logger.info("compile time strobelight profiling enabled") + + if profiler_class is StrobelightCLIFunctionProfiler: + import shutil + + if not shutil.which("strobeclient"): + logger.info( + "strobeclient not found, cant enable compile time strobelight profiling, seems" + "like you are not on a FB machine." + ) + return + + cls.enabled = True + cls._cls_init() + # profiler_class should have public API similar to that of StrobelightCLIFunctionProfiler. + # we have pass different functionProfilerClass for meta-internal fbcode targets. + # NB: the actual implementation in Meta is at + # fbcode/caffe2/fb/strobelight/function_profiler.py + cls.profiler = profiler_class( + sample_each=cls.sample_each, + max_profile_duration_sec=cls.max_profile_time, + stack_max_len=cls.max_stack_length, + async_stack_max_len=cls.max_stack_length, + run_user_name="pt2-profiler/" + + os.environ.get("USER", os.environ.get("USERNAME", "")), + sample_tags={cls.identifier}, + ) + + @classmethod + def _cls_init(cls) -> None: + cls.identifier = "{date}{pid}{hostname}".format( + date=datetime.now().strftime("%Y-%m-%d-%H:%M:%S"), + pid=os.getpid(), + hostname=gethostname(), + ) + + logger.info("Unique sample tag for this run is: %s", cls.identifier) + logger.info( + "URL to access the strobelight profile at the end of the run: %s", + get_strobelight_url(cls.identifier), + ) + + @classmethod + def _log_stats(cls) -> None: + logger.info( + "%s strobelight success runs out of %s non-recursive compilation events.", + cls.success_profile_count, + cls.success_profile_count + cls.failed_profile_count, + ) + + # TODO use threadlevel meta data to tags to record phases. + @classmethod + def profile_compile_time( + cls, func: Any, phase_name: str, *args: Any, **kwargs: Any + ) -> Any: + if not cls.enabled: + return func(*args, **kwargs) + + if cls.profiler is None: + logger.error("profiler is not set") + return + + if cls.inside_profile_compile_time: + cls.ignored_profile_runs += 1 + logger.info( + "profile_compile_time is requested for phase: %s while already in running phase: %s, recursive call ignored", + phase_name, + cls.current_phase, + ) + return func(*args, **kwargs) + + cls.inside_profile_compile_time = True + cls.current_phase = phase_name + + work_result = cls.profiler.profile(func, *args, **kwargs) + + if cls.profiler.profile_result is not None: + cls.success_profile_count += 1 + else: + cls.failed_profile_count += 1 + + cls._log_stats() + cls.inside_profile_compile_time = False + return work_result diff --git a/lib/python3.12/site-packages/torch/compiler/__init__.py b/lib/python3.12/site-packages/torch/compiler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ca5123314a8b58428b9fa89d108a773025753686 --- /dev/null +++ b/lib/python3.12/site-packages/torch/compiler/__init__.py @@ -0,0 +1,404 @@ +# mypy: allow-untyped-defs +from typing import Any, Callable, List, TypeVar + +import torch + + +__all__ = [ + "compile", + "assume_constant_result", + "reset", + "allow_in_graph", + "substitute_in_graph", + "list_backends", + "disable", + "set_stance", + "cudagraph_mark_step_begin", + "wrap_numpy", + "is_compiling", + "is_dynamo_compiling", +] + + +_F = TypeVar("_F", bound=Callable[..., Any]) + + +def compile(*args, **kwargs): + """ + See :func:`torch.compile` for details on the arguments for this function. + """ + 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 + `_) + + .. 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: _F, + *, + can_constant_fold_through: bool = False, + skip_signature_check: bool = False, +) -> Callable[[_F], _F]: + """ + 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): + """ + 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. + """ + import torch._dynamo + + return torch._dynamo.disable(fn, recursive) + + +def set_stance( + stance: str = "default", *, skip_guard_eval_unsafe=False, force_backend=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. + + 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 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 `__ + """ + 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 + + +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_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 diff --git a/lib/python3.12/site-packages/torch/compiler/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/compiler/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca6546812e7b83215a88b7a6724af6de17471d66 Binary files /dev/null and b/lib/python3.12/site-packages/torch/compiler/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/compiler/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/torch/compiler/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8ffd3c0c63c1f04a7c53478878739a8d6aadf78 Binary files /dev/null and b/lib/python3.12/site-packages/torch/compiler/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/compiler/config.py b/lib/python3.12/site-packages/torch/compiler/config.py new file mode 100644 index 0000000000000000000000000000000000000000..4f69d42e22d6b92b8c58dafff8163168457f1b82 --- /dev/null +++ b/lib/python3.12/site-packages/torch/compiler/config.py @@ -0,0 +1,67 @@ +""" +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 os +import sys +from typing import Optional + + +__all__ = [ + "job_id", +] + + +# 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: Optional[str] = os.environ.get("TORCH_COMPILE_JOB_ID", 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. +""" + +cache_key_tag: str = os.environ.get("TORCH_COMPILE_CACHE_KEY_TAG", "") +""" +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. +""" + +from torch.utils._config_module import install_config_module + + +install_config_module(sys.modules[__name__]) diff --git a/lib/python3.12/site-packages/torch/distributed/__init__.py b/lib/python3.12/site-packages/torch/distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..181fab713febe834bd42bbd0881e61dec15938db --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/__init__.py @@ -0,0 +1,157 @@ +# mypy: allow-untyped-defs +import logging +import pdb +import sys +import traceback +import typing + +import torch + + +log = logging.getLogger(__name__) + + +def is_available() -> bool: + """ + Return ``True`` if the distributed package is available. + + Otherwise, + ``torch.distributed`` does not expose any other APIs. Currently, + ``torch.distributed`` is available on Linux, MacOS and Windows. Set + ``USE_DISTRIBUTED=1`` to enable it when building PyTorch from source. + Currently, the default value is ``USE_DISTRIBUTED=1`` for Linux and Windows, + ``USE_DISTRIBUTED=0`` for MacOS. + """ + return hasattr(torch._C, "_c10d_init") + + +if is_available() and not torch._C._c10d_init(): + raise RuntimeError("Failed to initialize torch.distributed") + +# Custom Runtime Errors thrown from the distributed package +DistError = torch._C._DistError +DistBackendError = torch._C._DistBackendError +DistNetworkError = torch._C._DistNetworkError +DistStoreError = torch._C._DistStoreError + +if is_available(): + from torch._C._distributed_c10d import ( + _broadcast_coalesced, + _compute_bucket_assignment_by_size, + _ControlCollectives, + _DEFAULT_FIRST_BUCKET_BYTES, + _make_nccl_premul_sum, + _register_builtin_comm_hook, + _register_comm_hook, + _StoreCollectives, + _test_python_store, + _verify_params_across_processes, + Backend as _Backend, + BuiltinCommHookType, + DebugLevel, + FileStore, + get_debug_level, + GradBucket, + Logger, + PrefixStore, + ProcessGroup as ProcessGroup, + Reducer, + set_debug_level, + set_debug_level_from_env, + Store, + TCPStore, + Work as _Work, + ) + + class _DistributedPdb(pdb.Pdb): + """ + Supports using PDB from inside a multiprocessing child process. + + Usage: + _DistributedPdb().set_trace() + """ + + def interaction(self, *args, **kwargs): + _stdin = sys.stdin + try: + sys.stdin = open("/dev/stdin") + pdb.Pdb.interaction(self, *args, **kwargs) + finally: + sys.stdin = _stdin + + _breakpoint_cache: typing.Dict[int, typing.Any] = {} + + def breakpoint(rank: int = 0, skip: int = 0): + """ + Set a breakpoint, but only on a single rank. All other ranks will wait for you to be + done with the breakpoint before continuing. + + Args: + rank (int): Which rank to break on. Default: ``0`` + skip (int): Skip the first ``skip`` calls to this breakpoint. Default: ``0``. + """ + if skip > 0: + key = hash(str(traceback.format_exc())) + counter = _breakpoint_cache.get(key, 0) + 1 + _breakpoint_cache[key] = counter + if counter <= skip: + log.warning("Skip the breakpoint, counter=%d", counter) + return + + if get_rank() == rank: + pdb = _DistributedPdb() + pdb.message( + "\n!!! ATTENTION !!!\n\n" + f"Type 'up' to get to the frame that called dist.breakpoint(rank={rank})\n" + ) + pdb.set_trace() + # If Meta/Python keys are in the TLS, we want to make sure that we ignore them + # and hit the (default) CPU/CUDA implementation of barrier. + meta_in_tls = torch._C._meta_in_tls_dispatch_include() + guard = torch._C._DisableTorchDispatch() # type: ignore[attr-defined] + torch._C._set_meta_in_tls_dispatch_include(False) + try: + barrier() + finally: + torch._C._set_meta_in_tls_dispatch_include(meta_in_tls) + del guard + + if sys.platform != "win32": + from torch._C._distributed_c10d import HashStore + + from .device_mesh import DeviceMesh, init_device_mesh + + # Variables prefixed with underscore are not auto imported + # See the comment in `distributed_c10d.py` above `_backend` on why we expose + # this. + from .distributed_c10d import * # noqa: F403 + from .distributed_c10d import ( + _all_gather_base, + _coalescing_manager, + _CoalescingManager, + _create_process_group_wrapper, + _get_process_group_name, + _rank_not_in_group, + _reduce_scatter_base, + get_node_local_rank, + ) + from .remote_device import _remote_device + from .rendezvous import ( + _create_store_from_options, + register_rendezvous_handler, + rendezvous, + ) + + set_debug_level_from_env() + +else: + # This stub is sufficient to get + # python test/test_public_bindings.py -k test_correct_module_names + # working even when USE_DISTRIBUTED=0. Feel free to add more + # stubs as necessary. + # We cannot define stubs directly because they confuse pyre + + class _ProcessGroupStub: + pass + + sys.modules["torch.distributed"].ProcessGroup = _ProcessGroupStub # type: ignore[attr-defined] diff --git a/lib/python3.12/site-packages/torch/distributed/_checkpointable.py b/lib/python3.12/site-packages/torch/distributed/_checkpointable.py new file mode 100644 index 0000000000000000000000000000000000000000..409d1f8a473bb6b55933b5306c6bbfac03eaf19a --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_checkpointable.py @@ -0,0 +1,38 @@ +# mypy: allow-untyped-defs +# Copyright (c) Meta Platforms, Inc. and affiliates +from typing import Any, Protocol, runtime_checkable + +import torch + + +@runtime_checkable +class _Checkpointable(Protocol): # noqa: PYI046 + """ + Interface for checkpointable objects. + Implemented as a protocol, implicit subtyping is supported so subclasses do not need to inherit this explicitly. + This is to allow arbitrary objects/tensor subclasses to hook into DCP seamlessly through implementing the interface. + """ + + def __create_write_items__(self, fqn: str, object: Any): + """ + Return a list of WriteItems based on object's contents. + """ + raise NotImplementedError( + "_Checkpointable._create_write_items is not implemented" + ) + + def __create_chunk_list__(self): + """ + Return a list of `ChunkStorageMetadata` based on object's contents. + """ + raise NotImplementedError( + "_Checkpointable._create_chunk_list is not implemented" + ) + + def __get_tensor_shard__(self, index) -> torch.Tensor: + """ + Return a 'torch.Tensor' shard based on 'MetadataIndex'. + """ + raise NotImplementedError( + "_Checkpointable._get_tensor_shard is not implemented" + ) diff --git a/lib/python3.12/site-packages/torch/distributed/_composable_state.py b/lib/python3.12/site-packages/torch/distributed/_composable_state.py new file mode 100644 index 0000000000000000000000000000000000000000..6d2b8baed766ffaf7d2b1fccaaebe608616f66c1 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_composable_state.py @@ -0,0 +1,44 @@ +import weakref +from typing import cast, Optional + +import torch.nn as nn + + +class _State: + pass + + +_module_state_mapping: weakref.WeakKeyDictionary[ + nn.Module, weakref.ReferenceType[_State] +] = weakref.WeakKeyDictionary() + + +def _insert_module_state(module: nn.Module, state: _State) -> None: + global _module_state_mapping + assert module not in _module_state_mapping, f"Inserting {module} more than once." + _module_state_mapping[module] = weakref.ref(state) + + +def _get_module_state(module: nn.Module) -> Optional[_State]: + """ + Return the ``_State`` in ``model``. + + Given a ``module``, this API finds out if the module is also a ``_State`` + instance or if the module is managed by a composable API. If the module + is also a ``_State``, ``module`` will be casted to ``_State` and returned. + If it is managed by a composable API, the corresponding ``_State`` will + be returned. + """ + global _module_state_mapping + if isinstance(module, _State): + return cast(_State, module) + else: + # https://github.com/pytorch/pytorch/issues/107054 + if module in _module_state_mapping: + state_ref = _module_state_mapping[module] + state = state_ref() + if state is None: + raise AssertionError("State has already been garbage collected") + return state + else: + return None diff --git a/lib/python3.12/site-packages/torch/distributed/_functional_collectives.py b/lib/python3.12/site-packages/torch/distributed/_functional_collectives.py new file mode 100644 index 0000000000000000000000000000000000000000..a413f43956e32be8405a6c44b0fe8a7a0f4988d9 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_functional_collectives.py @@ -0,0 +1,1191 @@ +# mypy: allow-untyped-defs +import contextlib +import sys +import warnings +from typing import Any, cast, List, Optional, Tuple, Type, TYPE_CHECKING, Union + +import torch +import torch.distributed as dist +import torch.distributed.distributed_c10d as c10d +from torch.distributed.device_mesh import DeviceMesh +from torch.fx.experimental.proxy_tensor import get_proxy_mode + +from . import _functional_collectives_impl as fun_col_impl + + +try: + from torch.utils._cxx_pytree import tree_map_only +except ImportError: + from torch.utils._pytree import tree_map_only # type: ignore[no-redef] + + +if torch._running_with_deploy(): + + def is_torchdynamo_compiling(): + """Can't import torchdynamo in torchdeploy builds currently.""" + return False + +else: + try: + from torch.compiler import is_dynamo_compiling as is_torchdynamo_compiling + except Exception: + warnings.warn( + "Unable to import torchdynamo util `is_torchdynamo_compiling`, so won't support torchdynamo correctly" + ) + + def is_torchdynamo_compiling(): + return False + + +""" +New traceable, functional collectives. +RFC: https://github.com/pytorch/pytorch/issues/93173 + + compiler: trace these ops with plain-old-data schemas, then choose how to lower them. + eager: execute these 'functional' ops which in eager return AsyncCollectiveTensor subclasses, + automatically calling .wait() on underlying/hidden async 'work' obj only when fed to + a downstream op. + +Issues: +* Where should these ops live? Couldn't `import torch` if putting these ops in existing torch.distributed files +* Proper support for eager requires inplace ops. We should explore having it as an option for the API. +""" + +""" +Functional collectives are asynchronous only and we perform implicit stream synchronization +on behalf of the user. + +We use AsyncCollectiveTensor to wrap the result tensor of a collective and it lets us witness +first usage of the tensor and insert cross stream sync at the right place. + +The above are the easy bits, the hard one is how we match the Work object returned by +c10d and the tensor AsyncCollectiveTensor wraps. We alloc the tensor inside the collective +op implementation (see ``clone()`` call in ``_all_reduce``) and then it's handled by the +dispatcher which might call other implementations that are allowed to change the returned +tensor - even return a tensor with a different shape (see ``torch.vmap``). + +This means the caller of our ops receives a Tensor that is not guaranteed to be the same +allocated by our implementations and that makes pairing The AsyncTensor to the original +tensor a lot harder. This pairing is needed so we can lookup the Work object to use. + +Originally, we tried WeakKeyDictionary to map from Tensor to Work, but because Tensor's +identity is not stable across dispatch, the op caller would end up with a different Tensor +instance that would not match any in the dictionary. + +With Tensor identity out of the question, we decided use the tensor data pointer, which +should be stable across all the Tensor changes done during dispatch. + +We have a dictionary of tensor::data_ptr -> Work that we insert right after we call into c10d. + +We use this dictionary when AsyncCollectiveTensor is used to invoke Work::wait() + +Finally, we setup a finalizer against the tensor wrapper to observe it getting collected so we +can clean up stale entries in the dictionary. + +To eliminate the possibility of races we have a global version counter that is used by the finalizer. + +As a wise man said once: Don't cross the streams (https://www.youtube.com/watch?v=wyKQe_i9yyo) + +""" + +""" +Functional collectives can accept any of these types to describe the ranks participating in collectives. + +The different types will be desugared to a canonical format +""" +RANK_TYPES = Union[ + List[int], + List[List[int]], + dist.ProcessGroup, + DeviceMesh, + Tuple["dist.tensor.DeviceMesh", int], + str, +] + + +""" +User facing APIs for functional collectives +------------------------------------------- + +These apis are called by user code and expected to work both in eager execution and compilation, +but there are significant differences to how the two modes are implemented underneath. + +Eager execution is 'optimized' using a tensor subclass that schedules the synchronization (via wait_tensor() op) +just before the tensor is first used. Compiled tracing currently relies on the compiler to perform this optimization, +and cannot yet correctly trace the AsyncTensor wrapper class. In the future, these paths may be unified +if sufficient subclass support is added in dynamo. + +Example: all_reduce is an entrypoint API, and other collectives follow a similar pattern. + +Here's how it works under torch.compile/dynamo: +all_reduce(...) + |--> _expand_group(...) - desugars processgroup into canonical/traceable format + |--> c10d_functional.all_reduce(...) - dynamo captures this op call, doesn't trace deeper + |--> _maybe_wrap_tensor(...) - wait_tensor() op is immediately called, no AsyncTensor subclass needed + +And under eager execution: +all_reduce(...) + |--> _expand_group(...) - same as above, but less critical for eager + |--> c10d_functional.all_reduce(...) - dispatches to real kernel OR records op in trace + |--> _maybe_wrap_tensor(...) - AsyncTensor wrapper applied to returned tensor, + which issues wait_tensor() at the time of first use +""" + + +def wait_tensor(tensor): + """ + Wait on a tensor returned by the collectives ops. + + Waiting follows device semantics, which means blocking on CPU and synchronizing streams on CUDA. + """ + return torch.ops._c10d_functional.wait_tensor(tensor) # type: ignore[attr-defined] + + +def broadcast(self: torch.Tensor, src: int, group: RANK_TYPES, tag: str = ""): + """ + Broadcasts the tensor to all processes in the given process group. + + Args: + src (int): Source rank + group (ProcessGroup or List[int]): The process group to work on. + tag (str, optional): A unique identifier for the collective. Default: empty string + """ + group_name = _resolve_group_name(group, tag) + tensor = torch.ops._c10d_functional.broadcast(self, src, group_name) + return _maybe_wrap_tensor(tensor) + + +def all_reduce(self: torch.Tensor, reduceOp: str, group: RANK_TYPES, tag: str = ""): + """ + Reduces the tensor data across all machines in such a way that all get + the final result. + + The input tensor is left unmodified. + + Group can be one of: + List[int]: ranks participating in the collective. + List[List[int]]: 2D mesh of ranks taking part of this collective in MPMD. + ProcessGroup: Will perform a collective using the ranks and tag of the PG. + DeviceMesh: Do a SPMD collective over all ranks of the mesh + (DeviceMesh, int): Do a MPMD collective over one dimension of the DeviceMesh + + :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover + that information and perform collective algebraic optimization. Use other forms of input for that. + """ + group_name = _resolve_group_name(group, tag) + tensor = torch.ops._c10d_functional.all_reduce(self, reduceOp.lower(), group_name) + return _maybe_wrap_tensor(tensor) + + +def all_gather_tensor( + self: torch.Tensor, + gather_dim: int, + group: RANK_TYPES, + tag: str = "", +): + """ + Gather tensor data across from all machines and concatenate over ``gather_dim``. + + Note that it currently only supports gather_dim = 0. + + The input tensor is left unmodified. + Group can be one of: + List[int]: ranks participating in the collective. + List[List[int]]: 2D mesh of ranks taking part of this collective in MPMD. + ProcessGroup: Will perform a collective using the ranks and tag of the PG. + DeviceMesh: Do a SPMD collective over all ranks of the mesh + (DeviceMesh, int): Do a MPMD collective over one dimension of the DeviceMesh + + :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover + that information and perform collective algebraic optimization. Use other forms of input for that. + """ + assert self.is_contiguous() + group_name = _resolve_group_name(group, tag) + group_size = c10d._get_group_size_by_name(group_name) + tensor = torch.ops._c10d_functional.all_gather_into_tensor( + self, group_size, group_name + ) + res = _maybe_wrap_tensor(tensor) + # TODO this should be done inside AsyncCollectiveTensor to delay the wait() call + if gather_dim != 0: + # torch.cat access the data so we already need to wait here, first do wait + # and then chunk + cat avoid us going through ACT dispatching logic again + if isinstance(res, AsyncCollectiveTensor): + res = res.wait() # type: ignore[attr-defined] + res = torch.cat(torch.chunk(res, group_size, dim=0), dim=gather_dim) + return res + + +def all_gather_tensor_autograd( + self: torch.Tensor, + gather_dim: int, + group: RANK_TYPES, + tag: str = "", +): + """ + Gather tensor data across from all machines and concatenate over ``gather_dim``. + + Note that it currently only supports gather_dim = 0. + + This function is the same as all_gather_tensor but will propagate the + backwards gradient across workers. + + See all_gather_tensor for more details on usage. + """ + group_name = _resolve_group_name(group, tag) + group_size = c10d._get_group_size_by_name(group_name) + + tensor = torch.ops._c10d_functional_autograd.all_gather_into_tensor( + self, group_size, group_name + ) + res = _FromTorchTensor.apply(tensor) + # TODO this should be done inside AsyncCollectiveTensor to delay the wait() call + if gather_dim != 0: + # torch.cat access the data so we already need to wait here, first do wait + # and then chunk + cat avoid us going through ACT dispatching logic again + if isinstance(res, AsyncCollectiveTensor): + res = res.wait() # type: ignore[attr-defined] + res = torch.cat(torch.chunk(res, group_size, dim=0), dim=gather_dim) + return res + + +def reduce_scatter_tensor( + self: torch.Tensor, + reduceOp: str, + scatter_dim: int, + group: RANK_TYPES, + tag: str = "", +): + """ + Reduces the tensor data across all machines in such a way that all get + the final result, then scatter the results to corresponding ranks. + + + The input tensor is left unmodified. + Group can be one of: + List[int]: ranks participating in the collective. + List[List[int]]: 2D mesh of ranks taking part of this collective in MPMD. + ProcessGroup: Will perform a collective using the ranks and tag of the PG. + DeviceMesh: Do a SPMD collective over all ranks of the mesh + (DeviceMesh, int): Do a MPMD collective over one dimension of the DeviceMesh + :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover + that information and perform collective algebraic optimization. Use other forms of input for that. + """ + group_name = _resolve_group_name(group, tag) + group_size = c10d._get_group_size_by_name(group_name) + + assert ( + self.size(scatter_dim) % group_size == 0 + ), f"input dimension 0 ({self.size(0)} must be a multiple of group_size {group_size}" + if scatter_dim != 0: + tensor_list = torch.chunk(self, group_size, dim=scatter_dim) + self = torch.cat(tensor_list) + + tensor = torch.ops._c10d_functional.reduce_scatter_tensor( + self, + reduceOp.lower(), + group_size, + group_name, # type: ignore[possibly-undefined] + ) + res = _maybe_wrap_tensor(tensor) + return res + + +def reduce_scatter_tensor_autograd( + self: torch.Tensor, + reduceOp: str, + scatter_dim: int, + group: RANK_TYPES, + tag: str = "", +): + """ + Reduces the tensor data across all machines in such a way that all get + the final result, then scatter the results to corresponding ranks. + + This function is the same as reduce_scatter_tensor but will propagate the + backwards gradient across workers. + + Currently only the "sum" reduceOp is supported. + + See reduce_scatter_tensor for more details on usage. + """ + + group_name = _resolve_group_name(group, tag) + group_size = c10d._get_group_size_by_name(group_name) + + assert ( + self.size(scatter_dim) % group_size == 0 + ), f"input dimension 0 ({self.size(0)} must be a multiple of group_size {group_size}" + if scatter_dim != 0: + tensor_list = torch.chunk(self, group_size, dim=scatter_dim) + self = torch.cat(tensor_list) + + tensor = torch.ops._c10d_functional_autograd.reduce_scatter_tensor( + self, + reduceOp.lower(), + group_size, + group_name, # type: ignore[possibly-undefined] + ) + res = _FromTorchTensor.apply(tensor) + return res + + +def all_reduce_coalesced( + self: List[torch.Tensor], reduceOp: str, group: RANK_TYPES, tag: str = "" +) -> List[torch.Tensor]: + """ + Reduces a list of tensors across all machines in such a way that all get + the final result. + + The all tensors in the input list are left unmodified. + + Group can be one of: + List[int]: ranks participating in the collective. + List[List[int]]: 2D mesh of ranks taking part of this collective in MPMD. + ProcessGroup: Will perform a collective using the ranks and tag of the PG. + DeviceMesh: Do a SPMD collective over all ranks of the mesh + (DeviceMesh, int): Do a MPMD collective over one dimension of the DeviceMesh + + :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover + that information and perform collective algebraic optimization. Use other forms of input for that. + """ + group_name = _resolve_group_name(group, tag) + tensor_list = torch.ops._c10d_functional.all_reduce_coalesced( # type: ignore[attr-defined] + self, + reduceOp.lower(), + group_name, + ) + return list(map(_maybe_wrap_tensor, tensor_list)) + + +def all_gather_into_tensor_coalesced( + self: List[torch.Tensor], group: RANK_TYPES, tag: str = "" +) -> List[torch.Tensor]: + """ + Gather a list of tensors across from all machines. + + Note that it currently only supports gather_dim = 0. + + The input tensor is left unmodified. + Group can be one of: + List[int]: ranks participating in the collective. + List[List[int]]: 2D mesh of ranks taking part of this collective in MPMD. + ProcessGroup: Will perform a collective using the ranks and tag of the PG. + DeviceMesh: Do a SPMD collective over all ranks of the mesh + (DeviceMesh, int): Do a MPMD collective over one dimension of the DeviceMesh + + :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover + that information and perform collective algebraic optimization. Use other forms of input for that. + """ + group_name = _resolve_group_name(group, tag) + group_size = c10d._get_group_size_by_name(group_name) + tensor_list = torch.ops._c10d_functional.all_gather_into_tensor_coalesced( # type: ignore[attr-defined] + self, + group_size, + group_name, + ) + return list(map(_maybe_wrap_tensor, tensor_list)) + + +def reduce_scatter_tensor_coalesced( + inputs: List[torch.Tensor], + reduceOp: str, + scatter_dim: List[int], + group: RANK_TYPES, + tag: str = "", +) -> List[torch.Tensor]: + """ + Reduces a list of tensors across all machines in such a way that all get + the final result, then scatter the results to corresponding ranks. + + The input tensors are left unmodified. + Group can be one of: + List[int]: ranks participating in the collective. + List[List[int]]: 2D mesh of ranks taking part of this collective in MPMD. + ProcessGroup: Will perform a collective using the ranks and tag of the PG. + DeviceMesh: Do a SPMD collective over all ranks of the mesh + (DeviceMesh, int): Do a MPMD collective over one dimension of the DeviceMesh + + :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover + that information and perform collective algebraic optimization. Use other forms of input for that. + """ + group_name = _resolve_group_name(group, tag) + group_size = c10d._get_group_size_by_name(group_name) + + assert len(scatter_dim) == len(inputs) + for idx, (dim, tensor) in enumerate(zip(scatter_dim, inputs)): + assert ( + tensor.size(dim) % group_size == 0 + ), f"input dimension {dim} ({tensor.size(dim)} must be a multiple of group_size {group_size} for tensor at index {idx}" + if dim != 0: + tensor_list = torch.chunk(tensor, group_size, dim=dim) + inputs[idx] = torch.cat(tensor_list) + + tensor_list = torch.ops._c10d_functional.reduce_scatter_tensor_coalesced( # type: ignore[attr-defined] + inputs, + reduceOp.lower(), + group_size, + group_name, # type: ignore[possibly-undefined] + ) + + return list(map(_maybe_wrap_tensor, tensor_list)) + + +# This is a bit unsafe: it checks if the first argument in the schema reports as a non-mutable alias. +# Today, this maps 1:1 with "aten ops that are views". +def _is_view_op(tgt): + assert isinstance(tgt, torch._ops.OpOverload) + schema = tgt._schema + if len(schema.arguments) > 0: + first_arg = schema.arguments[0] + # check if op is a view + return first_arg.alias_info is not None and not first_arg.alias_info.is_write + + +def all_to_all_single( + self: torch.Tensor, + output_split_sizes: Optional[List[int]], + input_split_sizes: Optional[List[int]], + group: RANK_TYPES, + tag: str = "", +) -> torch.Tensor: + """ + Each process splits input tensor and then scatters the split list + to all processes in a group. Then concatenate the received tensors from all + the processes in the group and return single output tensor. + + Group can be one of: + List[int]: ranks participating in the collective. + List[List[int]]: 2D mesh of ranks taking part of this collective in MPMD. + ProcessGroup: Will perform a collective using the ranks and tag of the PG. + DeviceMesh: Do a SPMD collective over all ranks of the mesh + (DeviceMesh, int): Do a MPMD collective over one dimension of the DeviceMesh + + :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover + that information and perform collective algebraic optimization. Use other forms of input for that. + """ + if output_split_sizes is not None: + assert all( + isinstance(size, (int, torch.SymInt)) for size in output_split_sizes + ), output_split_sizes + if input_split_sizes is not None: + assert all( + isinstance(size, (int, torch.SymInt)) for size in input_split_sizes + ), input_split_sizes + group_name = _resolve_group_name(group, tag) + group_size = c10d._get_group_size_by_name(group_name) + if output_split_sizes is None or input_split_sizes is None: + assert output_split_sizes is None and input_split_sizes is None, ( + "output_split_sizes and input_split_sizes must either be " + "specified together or both set to None" + ) + output_split_sizes = [self.shape[0] // group_size] * group_size + input_split_sizes = output_split_sizes + tensor = torch.ops._c10d_functional.all_to_all_single( # type: ignore[attr-defined] + self, + output_split_sizes, + input_split_sizes, + group_name, + ) + return _maybe_wrap_tensor(tensor) + + +def all_to_all_single_autograd( + self: torch.Tensor, + output_split_sizes: Optional[List[int]], + input_split_sizes: Optional[List[int]], + group: RANK_TYPES, + tag: str = "", +) -> torch.Tensor: + """ + Same as all_to_all_single but supports autograd. + """ + if output_split_sizes is not None: + assert all( + isinstance(size, (int, torch.SymInt)) for size in output_split_sizes + ), output_split_sizes + if input_split_sizes is not None: + assert all( + isinstance(size, (int, torch.SymInt)) for size in input_split_sizes + ), input_split_sizes + + group_name = _resolve_group_name(group, tag) + group_size = c10d._get_group_size_by_name(group_name) + if output_split_sizes is None or input_split_sizes is None: + assert output_split_sizes is None and input_split_sizes is None, ( + "output_split_sizes and input_split_sizes must either be " + "specified together or both set to None" + ) + output_split_sizes = [self.shape[0] // group_size] * group_size + input_split_sizes = output_split_sizes + tensor = torch.ops._c10d_functional_autograd.all_to_all_single( # type: ignore[attr-defined] + self, + output_split_sizes, + input_split_sizes, + group_name, + ) + return _FromTorchTensor.apply(tensor) + + +def permute_tensor( + self: torch.Tensor, + src_dst: List[int], + group: RANK_TYPES, + tag: str = "", +) -> torch.Tensor: + """ + Permutes the elements of the tensor according to the given source/destination pairs. `src_dst` should + be defined such that src_dst[m] == n means m sends to n. + + Group can be one of: + List[int]: ranks participating in the collective. + List[List[int]]: 2D mesh of ranks taking part of this collective in MPMD. + ProcessGroup: Will perform a collective using the ranks and tag of the PG. + DeviceMesh: Do a SPMD collective over all ranks of the mesh + (DeviceMesh, int): Do a MPMD collective over one + """ + t, rankset, group_size = _expand_group(group, tag) + local_pg = c10d._find_or_create_pg_by_ranks_and_tag(t, rankset, group_size) + + output_split_sizes = [0] * group_size + input_split_sizes = [0] * group_size + for src, dst in enumerate(src_dst): + if src == dist.get_rank(local_pg): + input_split_sizes[dst] = self.numel() + if dst == dist.get_rank(local_pg): + output_split_sizes[src] = self.numel() + + return all_to_all_single(self, output_split_sizes, input_split_sizes, group, tag) + + +class AsyncCollectiveTensor(torch.Tensor): + r""" + A Tensor wrapper subclass that is used to trigger a call to wait + prior to first use of the underlying tensor. + Use it inside functional collective pytorch wrappers like the following: + def functional_collective(self, group, tag): + tag, rankset, group_size = _expand_group(group, tag) + tensor = torch.ops.c10d_functional.{collective}(self, tag, rankset, group_size) + return _maybe_wrap_tensor(tensor) + """ + elem: torch.Tensor + completed: bool + + __slots__ = ["elem", "completed"] + + @staticmethod + def __new__(cls, elem: torch.Tensor): + r = torch.Tensor._make_wrapper_subclass( # type: ignore[attr-defined] + cls, + elem.size(), + strides=elem.stride(), + storage_offset=elem.storage_offset(), + dtype=elem.dtype, + layout=elem.layout, + device=elem.device, + requires_grad=elem.requires_grad, + ) + r.elem = elem + r.completed = False + return r + + def __tensor_flatten__(self): + return ["elem"], None + + def tolist(self): + return self.trigger_wait().tolist() + + @staticmethod + def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride): + assert meta is None + elem = inner_tensors["elem"] + return AsyncCollectiveTensor(elem) + + def __coerce_same_metadata_as_tangent__( + self, expected_metadata: Any, expected_type: Optional[Type] = None + ): + if expected_type is not torch.Tensor: + return None + + return self.trigger_wait() + + def __repr__(self) -> str: # type: ignore[override] + return f"AsyncCollectiveTensor({self.trigger_wait()})" + + def trigger_wait(self): + if not self.completed: + out = wait_tensor(self.elem) + self.completed = True + return out + else: + return self.elem + + def wait(self) -> torch.Tensor: + return wait_tensor(self.elem) + + def _get_acs_underlying_tensor(self): + """This method enables _functional_collectives_impl to test if a tensor is an ACS""" + return self.elem + + @classmethod + def __torch_dispatch__(cls, func, types, args=(), kwargs=None): + if func == torch.ops.aten.view.default: + # Fast handle aten.view as a lot of view related op goes to aten.view + # eventually, this avoids pytree slowdown + res = func(args[0].elem, args[1]) + wrapper_res = AsyncCollectiveTensor(res) + return wrapper_res + + is_view_op = _is_view_op(func) + + def unwrap(e: AsyncCollectiveTensor): + # wait_tensor is idepotent and will do stream sync only once + if not is_view_op: + return e.trigger_wait() + return e.elem + + def wrap(e: torch.Tensor): + # wait_tensor is idepotent and will do stream sync only once + assert not isinstance(e, AsyncCollectiveTensor) + res = AsyncCollectiveTensor(e) + return res + + unwrapped_args = tree_map_only(AsyncCollectiveTensor, unwrap, args) + unwrapped_kwargs = tree_map_only(AsyncCollectiveTensor, unwrap, kwargs) + + # we don't wrap the result as it doesn't need to be waited on. + out = func(*unwrapped_args, **unwrapped_kwargs) + + # View ops dont require a sync, so we should re-wrap the outputs. + if is_view_op: + out = tree_map_only(torch.Tensor, wrap, out) + + return out + + def numpy(self): # type: ignore[override] + return self.wait().numpy() + + +""" +Utils and infrastructure for tracing support +""" + + +def _expand_group(group: RANK_TYPES, tag: str = "") -> Tuple[str, List[int], int]: + """ + _expand_group desugars the different RANK_TYPES types into a canonical format that is traceable. + + By having this be part of the explicit eager codepath, we avoid having to specialize behavior inside + torchdynamo and can still interoperate with processgroup objects or other untraceable forms. + """ + # had to define this hack _inside_ expand_group to avoid + # graph_break [('torch.* op returned non-Tensor int + # caused by 'cast_*` functions being treated as 'torch.*' ops (iiuc) + if TYPE_CHECKING: + + def cast_listlistint(x): + return cast(List[List[int]], x) + + def cast_listint(x): + return cast(List[int], x) + + else: + # fake cast op for use at runtime since dynamo doesn't support real cast + # also, dynamo didn't like encountering 'typing' objects () + # NotImplementedError: argument of type: + def cast_listlistint(x): + return x + + def cast_listint(x): + return x + + rankset: List[int] + if isinstance(group, list): + if isinstance(group[0], list): + nested_list = cast_listlistint(group) + rankset = [] + group_size = -1 + for rs in nested_list: + rankset.extend(rs) + if group_size != -1 and group_size != len(rs): + raise ValueError( + f"group sizes must be identical found {group_size} and {len(rs)}" + ) + group_size = len(rs) + else: + rankset = cast_listint(group) + group_size = len(rankset) + elif isinstance(group, dist.ProcessGroup): + rankset = dist.get_process_group_ranks(group) + group_size = len(rankset) + tag = tag or c10d._get_group_tag(group) + elif isinstance(group, DeviceMesh): + assert ( + group.ndim == 1 + ), "Only 1D mesh is supported, pass in (DeviceMesh, int) together if mesh > 1D" + # TODO: it should run collective in the whole mesh instead of dim 0 + tag, rankset, _ = group._dim_group_infos[0] + group_size = len(rankset) + elif isinstance(group, tuple): + if ( + len(group) == 2 + and isinstance(group[0], DeviceMesh) + and isinstance(group[1], int) + ): + dmesh = group[0] + dim = group[1] + tag, rankset, _ = dmesh._dim_group_infos[dim] + group_size = len(rankset) + else: + raise ValueError("Invalid tuple for group must be (DeviceMesh, int)") + else: + raise ValueError( + "Invalid type for group, must be one of List, Processgroup, DeviceMesh or (DeviceMesh, int)." + ) + + return (tag, rankset, group_size) + + +def _resolve_group_name(group: RANK_TYPES, tag: str = "") -> str: + """ + Given group in RANK_TYPES, return the group name. + """ + # `tag` will be deprecated. See details in: + # https://github.com/pytorch/pytorch/issues/93173#issuecomment-1907095208 + if isinstance(group, dist.ProcessGroup): + return group.group_name + elif isinstance(group, str): + return group + elif isinstance(group, DeviceMesh): + assert ( + group.ndim == 1 + ), "Only 1D mesh is supported, pass in (DeviceMesh, int) together if mesh > 1D" + return group._dim_group_infos[0][2] + elif isinstance(group, tuple): + if ( + len(group) == 2 + and isinstance(group[0], DeviceMesh) + and isinstance(group[1], int) + ): + dmesh = group[0] + dim = group[1] + return dmesh._dim_group_infos[dim][2] + else: + raise ValueError("Invalid tuple for group must be (DeviceMesh, int)") + elif isinstance(group, list): + if not is_torchdynamo_compiling(): + warnings.warn( + "The combination of ranks + tag as process group " + "identifier has been deprecated. Please switch to " + "using ProcessGroup, DeviceMesh, or group name instead.", + FutureWarning, + stacklevel=3, + ) + return c10d._resolve_group_name_by_ranks_and_tag(cast(List[int], group), tag) + else: + raise ValueError(f"Unsupported group type: {type(group)}, {group}") + + +class _FromTorchTensor(torch.autograd.Function): + """ + _FromTorchTensor allows autograd to propagate from a normal Tensor to an + AsyncCollectiveTensor. + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, # pyre-ignore[2]: Parameter must be annotated. + input: torch.Tensor, + ) -> torch.Tensor: + return _maybe_wrap_tensor(input) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: # type: ignore[override] + return grad_output + + +def _are_we_tracing() -> bool: + if is_torchdynamo_compiling(): + return True + # If fake mode is turned on, we are almost definitely compiling/tracing. + if torch._C._get_dispatch_mode(torch._C._TorchDispatchModeKey.FAKE) is not None: + return True + return get_proxy_mode() is not None + + +def _maybe_wrap_tensor(self) -> torch.Tensor: + if _are_we_tracing(): + return wait_tensor(self) + res = AsyncCollectiveTensor(self) + return cast(torch.Tensor, res) + + +@contextlib.contextmanager +def allow_inflight_collective_as_graph_input_ctx(value: bool = True): + """ + Context manager to temporarily set whether inflight collectives are allowed as torch.compile graph inputs. + Common use case is when the collective is issued in eager (with `async_op=True`) but waited in compiled region: + ``` + def all_reduce_eager(x): + y = x * x + req = dist.all_reduce(y, op=dist.ReduceOp.SUM, async_op=True) + return y + + @torch.compile(fullgraph=True) + def all_reduce_wait_compiled(y): + torch.ops.c10d_functional.wait_tensor(y) + return y * y + + x = torch.ones(1280, 1280, device="cuda") + self.rank + # the context manager ensures that `wait_tensor(y)` will wait on the correct work object + with allow_inflight_collective_as_graph_input_ctx(): + y = all_reduce_eager(x) + z = all_reduce_wait_compiled(y) + ``` + With this context manager, when a collective is called, under the hood the work object of the collective + will be registered in the work registry, and the wait_tensor() in compiled region called on + the output tensor of the collective will wait on the correct work object. + """ + previous = torch._C._distributed_c10d._allow_inflight_collective_as_graph_input() + + try: + torch._C._distributed_c10d._set_allow_inflight_collective_as_graph_input(value) + yield + finally: + torch._C._distributed_c10d._set_allow_inflight_collective_as_graph_input( + previous + ) + + +def _all_gather_into_tensor_coalesced_meta(self, tag, rankset, group_size): + def mk_out_tensor(shard): + out_size = list(shard.size()) + out_size[0] *= group_size + out_tensor = shard.new_empty(out_size) + return out_tensor + + return [mk_out_tensor(t) for t in self] + + +# We now register meta kernels to deal with tracing +def _broadcast_meta(self, *args): + return torch.empty_like(self) + + +def _all_reduce_meta(self, *args): + return torch.empty_like(self) + + +def _wait_tensor_meta(self, *args): + return torch.empty_like(self) + + +def _all_gather_into_tensor_meta(shard, tag, rankset, group_size): + out_size = list(shard.size()) + out_size[0] *= group_size + return shard.new_empty(out_size) + + +def _reduce_scatter_tensor_meta(input, reduce_op, tag, rankset, group_size): + out_size = list(input.size()) + out_size[0] //= group_size + return input.new_empty(out_size) + + +def _all_reduce_coalesced_meta(self, *args): + return [torch.empty_like(t) for t in self] + + +def _all_reduce__meta(inp, *args): + return inp + + +def _broadcast__meta(inp, *args): + return inp + + +def _all_reduce_coalesced__meta(inputs, *args): + return inputs + + +def _reduce_scatter_tensor_coalesced_meta(inputs, reduceOp, tag, rankset, group_size): + def mk_out_tensor(input): + out_size = list(input.size()) + out_size[0] //= group_size + out_tensor = input.new_empty(out_size) + return out_tensor + + return [mk_out_tensor(t) for t in inputs] + + +# NB: We often say all_to_all has dynamic output size, but this is not +# technically true: instead, what typically happens is you manually +# communicate the output_split_sizes ahead of time (which is dynamic), +# but then you pass those sizes explicitly, and the all to all itself +# isn't dynamic, it just follows the specified output splits +def _all_to_all_single_meta( + input, output_split_sizes, input_split_sizes, *args, **kwargs +): + if output_split_sizes is None: + return input.new_empty(input.size()) + else: + for s in output_split_sizes: + torch._check_is_size(s) + out_size = list(input.size()) + out_size[0] = sum(output_split_sizes) + return input.new_empty(out_size) + + +def _all_gather_into_tensor_out_native_meta(input, group_size, group_name, *, out): + shape = list(input.size()) + shape[0] *= group_size + return input.new_empty(shape) + + +def _all_gather_into_tensor_native_meta(input, group_size, group_name): + shape = list(input.size()) + shape[0] *= group_size + return input.new_empty(shape) + + +def _all_gather_into_tensor_coalesced_native_meta(inputs, group_size, group_name): + return [ + _all_gather_into_tensor_native_meta(input, group_size, group_name) + for input in inputs + ] + + +def _reduce_scatter_tensor_native_meta(inp, reduce_op, group_size, group_name): + shape = list(inp.size()) + shape[0] //= group_size + return inp.new_empty(shape) + + +def _reduce_scatter_tensor_coalesced_native_meta( + inputs, reduce_op, group_size, group_name +): + return [ + _reduce_scatter_tensor_native_meta(inp, reduce_op, group_size, group_name) + for inp in inputs + ] + + +if not torch._running_with_deploy(): + # Library MUST be defined at module scope or it doesn't work + # Creating a "DEF" Library always crashes torch::deploy so we create our + # Library instances here guarded against running inside it + lib_impl = torch.library.Library("_c10d_functional", "IMPL") + lib_impl.impl("all_reduce", _all_reduce_meta, "Meta") + lib_impl.impl("all_reduce_", _all_reduce__meta, "Meta") + lib_impl.impl("all_reduce_coalesced", _all_reduce_coalesced_meta, "Meta") + lib_impl.impl("all_reduce_coalesced_", _all_reduce_coalesced__meta, "Meta") + lib_impl.impl("wait_tensor", _wait_tensor_meta, "Meta") + lib_impl.impl( + "all_gather_into_tensor_out", _all_gather_into_tensor_out_native_meta, "Meta" + ) + lib_impl.impl("all_gather_into_tensor", _all_gather_into_tensor_native_meta, "Meta") + lib_impl.impl( + "all_gather_into_tensor_coalesced", + _all_gather_into_tensor_coalesced_native_meta, + "Meta", + ) + lib_impl.impl("reduce_scatter_tensor", _reduce_scatter_tensor_native_meta, "Meta") + lib_impl.impl( + "reduce_scatter_tensor_coalesced", + _reduce_scatter_tensor_coalesced_native_meta, + "Meta", + ) + lib_impl.impl("all_to_all_single", _all_to_all_single_meta, "Meta") + lib_impl.impl("broadcast", _broadcast_meta, "Meta") + lib_impl.impl("broadcast_", _broadcast__meta, "Meta") + + # mark these ops has side effect so that they won't be removed by DCE + torch.fx.node.has_side_effect(torch.ops._c10d_functional.wait_tensor.default) + torch.fx.node.has_side_effect(torch.ops._c10d_functional.wait_tensor) + + # Register legacy ops for backward compatibility + # TODO(yifu): remove these in functional collective beta release + legacy_lib = torch.library.Library("c10d_functional", "DEF") + legacy_lib_impl = torch.library.Library("c10d_functional", "IMPL") + ops_defs = [ + "broadcast(Tensor self, int src, str tag, int[] ranks, int group_size) -> Tensor", + "all_reduce(Tensor self, str reduceOp, str tag, int[] ranks, int group_size) -> Tensor", + "all_reduce_coalesced(Tensor[] self, str reduceOp, str tag, int[] ranks, int group_size) -> Tensor[]", + "wait_tensor(Tensor self) -> Tensor", + "all_gather_into_tensor(Tensor shard, str tag, int[] ranks, int group_size) -> Tensor", + "all_gather_into_tensor_coalesced(Tensor[] input, str tag, int[] ranks, int group_size) -> Tensor[]", + "reduce_scatter_tensor(Tensor input, str reduceOp, str tag, int[] ranks, int group_size) -> Tensor", + "reduce_scatter_tensor_coalesced(Tensor[] inputs, str reduceOp, str tag, int[] ranks, int group_size) -> Tensor[]", + "all_to_all_single(Tensor input, SymInt[]? output_split_sizes, SymInt[]? input_split_sizes, str tag, int[] ranks, int group_size) -> Tensor", # noqa: B950 + ] + + my_module = sys.modules[__name__] + for op_def in ops_defs: + op_name = op_def[0 : op_def.index("(")] + backend_impl = getattr(fun_col_impl, f"_{op_name}") + legacy_lib.define(op_def, tags=torch.Tag.pt2_compliant_tag) + legacy_lib_impl.impl(op_name, backend_impl, "CompositeImplicitAutograd") + +else: + warnings.warn( + "PyTorch Distributed functional collectives do not work with torch::deploy." + ) + + +""" +Dynamo Remappings allow seamless translation from non-functional collectives of supportable form into +functional collective calls followed by inplace copy ops, allowing them to be traced into a functional graph. + +We implement this by writing a decomposition and teaching dynamo how to associate it to a corresponding op via +the mapping dict below. + +These schemas intentionally match torch.distributed.distributed_c10d.* ops that we are trying to remap from +""" + + +def all_gather_tensor_inplace( + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group, # TODO add a type, + async_op: bool = False, + tag: str = "", + gather_dim: int = 0, +): + assert ( + not async_op + ), "Can't remap async version of inplace op to functional collective" + + group = group or dist.group.WORLD + assert group is not None + + return output_tensor.copy_(all_gather_tensor(input_tensor, gather_dim, group, tag)) + + +def reduce_scatter_tensor_inplace( + output: torch.Tensor, + input: torch.Tensor, + op: str = "sum", # TODO type is actually c10d ReduceOp. is this ok? + group=None, # TODO add a type + async_op: bool = False, + scatter_dim: int = 0, + tag: str = "", +): + assert ( + not async_op + ), "Can't remap async version of inplace op to functional collective" + + group = group or dist.group.WORLD + assert group is not None + + return output.copy_(reduce_scatter_tensor(input, op, scatter_dim, group, tag)) + + +REDUCE_OP_TO_STR = { + dist.ReduceOp.SUM: "sum", + dist.ReduceOp.AVG: "avg", + dist.ReduceOp.PRODUCT: "product", + dist.ReduceOp.MIN: "min", + dist.ReduceOp.MAX: "max", + dist.ReduceOp.BAND: "band", + dist.ReduceOp.BOR: "bor", + dist.ReduceOp.BXOR: "bxor", +} + + +def all_reduce_inplace( + tensor: torch.Tensor, + op: str = "sum", + group=None, + async_op: bool = False, + tag: str = "", +): + assert ( + not async_op + ), "Can't remap async version of inplace op to functional collective" + + group = group or dist.group.WORLD + assert group is not None + + return tensor.copy_(all_reduce(tensor, op, group, tag)) + + +def all_to_all_inplace( + output: torch.Tensor, + input: torch.Tensor, + output_split_sizes=None, + input_split_sizes=None, + group=None, + async_op=False, + tag: str = "", +): + assert ( + not async_op + ), "Can't remap async version of inplace op to functional collective" + + group = group or dist.group.WORLD + assert group is not None + + return output.copy_( + all_to_all_single( + input, + output_split_sizes, + input_split_sizes, + group, + tag, + ) + ) + + +def all_gather_inplace( + tensor_list: List[torch.Tensor], + tensor: torch.Tensor, + group=None, + async_op=False, + tag: str = "", +): + assert ( + not async_op + ), "Can't remap async version of inplace op to functional collective" + assert all( + t.size(0) == tensor.size(0) for t in tensor_list + ), "Remapping variable size all_gather is not yet supported" + + group = group or dist.group.WORLD + assert group is not None + + output = all_gather_tensor(tensor, 0, group, tag) + + # Use aten.slice instead of aten.split because the latter causes + # tensor.shape(0) to be unnecessarily baked in when it's a SymInt. + output_splits = [] + offset = 0 + for t in tensor_list: + output_splits.append(output[offset : offset + t.size(0)]) + offset += t.size(0) + for dst, src in zip(tensor_list, output_splits): + dst.copy_(src) + return tensor_list + + +from torch.distributed.distributed_c10d import ( + _all_gather_base as legacy_all_gather_base, + _reduce_scatter_base as legacy_reduce_scatter_base, + all_gather as legacy_all_gather, + all_gather_into_tensor as legacy_allgather, + all_reduce as legacy_allreduce, + all_to_all_single as legacy_all_to_all_single, + reduce_scatter_tensor as legacy_reducescatter, +) + + +# This dict should contain sets of functions that dynamo is allowed to remap. +# Functions in this set should accept the same args/kwargs 1:1 as their mapping. +traceable_collective_remaps = { + legacy_allgather: all_gather_tensor_inplace, + legacy_reducescatter: reduce_scatter_tensor_inplace, + legacy_allreduce: all_reduce_inplace, + legacy_all_to_all_single: all_to_all_inplace, + legacy_all_gather: all_gather_inplace, + legacy_reduce_scatter_base: reduce_scatter_tensor_inplace, + legacy_all_gather_base: all_gather_tensor_inplace, +} diff --git a/lib/python3.12/site-packages/torch/distributed/_functional_collectives_impl.py b/lib/python3.12/site-packages/torch/distributed/_functional_collectives_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..4bd193d662bd6de80e87707455d0d1916410bf84 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_functional_collectives_impl.py @@ -0,0 +1,117 @@ +# mypy: allow-untyped-defs +from typing import List, Optional + +import torch +import torch.distributed.distributed_c10d as c10d + + +""" +This file contains the op impls for the legacy (c10d_functional) functional collectives. +These impls simply call into the native (_c10d_functional) functional collectives. +""" + + +def _broadcast(input, src, tag, ranks, group_size): + group_name = c10d._resolve_group_name_by_ranks_and_tag(ranks, tag) + return torch.ops._c10d_functional.broadcast( + input, + src, + group_name, + ) + + +def _all_reduce(input, reduce_op, tag, ranks, group_size): + group_name = c10d._resolve_group_name_by_ranks_and_tag(ranks, tag) + return torch.ops._c10d_functional.all_reduce( + input, + reduce_op, + group_name, + ) + + +def _all_reduce_coalesced(inputs, reduce_op, tag, ranks, group_size): + group_name = c10d._resolve_group_name_by_ranks_and_tag(ranks, tag) + return torch.ops._c10d_functional.all_reduce_coalesced( + inputs, + reduce_op, + group_name, + ) + + +def _all_gather_into_tensor(input, tag, ranks, group_size): + group_name = c10d._resolve_group_name_by_ranks_and_tag(ranks, tag) + return torch.ops._c10d_functional.all_gather_into_tensor( + input, + group_size, + group_name, + ) + + +def _all_gather_into_tensor_coalesced(input, tag, ranks, group_size): + group_name = c10d._resolve_group_name_by_ranks_and_tag(ranks, tag) + return torch.ops._c10d_functional.all_gather_into_tensor_coalesced( + input, + group_size, + group_name, + ) + + +def _reduce_scatter_tensor( + input: torch.Tensor, + reduce_op: str, + tag: str, + ranks: List[int], + group_size: int, +): + group_name = c10d._resolve_group_name_by_ranks_and_tag(ranks, tag) + return torch.ops._c10d_functional.reduce_scatter_tensor( + input, + reduce_op, + group_size, + group_name, + ) + + +def _reduce_scatter_tensor_coalesced( + inputs: List[torch.Tensor], + reduce_op: str, + tag: str, + ranks: List[int], + group_size: int, +): + group_name = c10d._resolve_group_name_by_ranks_and_tag(ranks, tag) + return torch.ops._c10d_functional.reduce_scatter_tensor_coalesced( + inputs, + reduce_op, + group_size, + group_name, + ) + + +def _all_to_all_single( + input: torch.Tensor, + output_split_sizes: Optional[List[int]], + input_split_sizes: Optional[List[int]], + tag: str, + ranks: List[int], + group_size: int, +): + if output_split_sizes is None or input_split_sizes is None: + assert output_split_sizes is None and input_split_sizes is None, ( + "output_split_sizes and input_split_sizes must either be " + "specified together or both set to None" + ) + output_split_sizes = [input.shape[0] // group_size] * group_size + input_split_sizes = output_split_sizes + + group_name = c10d._resolve_group_name_by_ranks_and_tag(ranks, tag) + return torch.ops._c10d_functional.all_to_all_single( + input, + output_split_sizes, + input_split_sizes, + group_name, + ) + + +def _wait_tensor(tensor: torch.Tensor) -> torch.Tensor: + return torch.ops._c10d_functional.wait_tensor(tensor) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/__init__.py b/lib/python3.12/site-packages/torch/distributed/_shard/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..85a313c779e7aa87726f425146048fcd37efd261 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/__init__.py @@ -0,0 +1 @@ +from .api import _shard_tensor, load_with_process_group, shard_module, shard_parameter diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c89843037cc0786a63a8415d85adf47fc9458681 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8cbf692475d3f590d4aae2b6afb5d0903c963dc Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56575284b0b26fe17c7f1a0c1030e831fc780088 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/common_op_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/common_op_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..034560dac23b00c54bc28e6982db84cc69721807 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/common_op_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/metadata.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/metadata.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..399a2e192e075fe7bebd18dda6281d039963a23b Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/metadata.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/op_registry_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/op_registry_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c0efcf4b0dc7d116bb027ee16504b808459e8a7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/op_registry_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/sharder.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/sharder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23d6458e9bd8d7c1b18dd382d7a3647548f3c405 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/__pycache__/sharder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/_utils.py b/lib/python3.12/site-packages/torch/distributed/_shard/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d06fc4dc961447dc974d2d79d7a4e418a7938541 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/_utils.py @@ -0,0 +1,32 @@ +from typing import Sequence + +import torch +from torch.distributed._shard.metadata import ShardMetadata + + +DEPRECATE_MSG = "Please use DTensor instead and we are deprecating ShardedTensor." + + +def narrow_tensor_by_index( + tensor: torch.Tensor, + offsets: Sequence[int], + sizes: Sequence[int], +) -> torch.Tensor: + """ + Narrow the tensor according to ``offsets`` and ``sizes``. + """ + narrowed_tensor = tensor + for idx, (offset, size) in enumerate(zip(offsets, sizes)): + if size < tensor.size(idx): + # Reshape to get shard for this rank and we don't want autograd + # recording here for the narrow op and 'local_shard' should be a + # leaf variable in the autograd graph. + narrowed_tensor = narrowed_tensor.narrow(idx, offset, size) + return narrowed_tensor + + +def narrow_tensor(tensor: torch.Tensor, metadata: ShardMetadata) -> torch.Tensor: + """ + Narrow the tensor according to the metadata + """ + return narrow_tensor_by_index(tensor, metadata.shard_offsets, metadata.shard_sizes) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/api.py b/lib/python3.12/site-packages/torch/distributed/_shard/api.py new file mode 100644 index 0000000000000000000000000000000000000000..975f499023d132b51a3bb0a31d5d2278f1e8cae8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/api.py @@ -0,0 +1,306 @@ +# mypy: allow-untyped-defs +from contextlib import contextmanager +from typing import Optional + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed import distributed_c10d +from torch.distributed._shard.sharded_tensor import ShardedTensor + +from .sharder import Sharder +from .sharding_plan import ShardingPlan +from .sharding_spec import ChunkShardingSpec, ShardingSpec + + +def _shard_tensor( + tensor: torch.Tensor, sharding_spec: ShardingSpec, src_rank=0, process_group=None +) -> ShardedTensor: + """ + Given a :class:`torch.Tensor`, it shards that tensor according to the provided + ``sharding_spec``. ``src_rank`` denotes the source rank which would be + used as the ground truth of the data which would be scattered as shards + across the rest of the ranks. + + Args: + tensor (:class:`torch.Tensor`): Tensor needs to be sharded. + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The specification + describing how to shard the Tensor. + + Keyword args: + src_rank (int, optional): The source rank which is used as the ground truth of + the data for the parameter that would be sharded and scattered + across the rest of the ranks. + Default: 0. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + + Returns: + A :class:`ShardedTensor` sharded from the given tensor. + + .. warning:: + Only :class:`torch.distributed._shard.sharding_spec.ChunkShardingSpec` is + currently supported as the ``sharding_spec``. + """ + if not tensor.is_contiguous(): + raise ValueError("input tensor is not a contiguous Tensor") + + pg = ( + process_group + if process_group is not None + else distributed_c10d._get_default_group() + ) + world_size = dist.get_world_size(pg) + current_rank = dist.get_rank(pg) + + # Validate src_rank and sharding_spec are same across all ranks. + gathered_list = [None] * world_size + dist.all_gather_object(gathered_list, (src_rank, sharding_spec), group=pg) + + for idx, entry in enumerate(gathered_list): + if src_rank != entry[0]: # type: ignore[index] + raise ValueError( + f"src_rank={src_rank} on rank: {current_rank} does not " # type: ignore[index] + f"match with src_rank={entry[0]} on rank: {idx}" # type: ignore[index] + ) + if sharding_spec != entry[1]: # type: ignore[index] + raise ValueError( + f"sharding_spec={sharding_spec} on rank: {current_rank} does not " # type: ignore[index] + f"match with sharding_spec={entry[1]} on rank: {idx}" # type: ignore[index] + ) + + st = sharding_spec.shard(tensor, src_rank=src_rank, process_group=pg) + + return st + + +def shard_parameter( + module: torch.nn.Module, + param_name: str, + sharding_spec: ShardingSpec, + src_rank=0, + process_group=None, +): + """ + Given a :class:`torch.nn.Module`, a ``param_name`` for a parameter in that + module, it shards that parameter according to the provided + ``sharding_spec``. ``src_rank`` denotes the source rank which would be + used as the ground truth of the data which would be scattered as shards + across the rest of the ranks. + + This method replaces ``module.param_name`` with a + :class:`torch.distributed._sharded_tensor.ShardedTensor` + + Args: + module (:class:`torch.nn.Module`): Module whose parameter needs to be sharded. + param_name (str): Name of the parameter of ``module`` that needs to be sharded. + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The specification + describing how to shard the Tensor. + + Keyword args: + src_rank (int, optional): The source rank which is used as the ground truth of + the data for the parameter that would be sharded and scattered + across the rest of the ranks. + Default: 0. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + + .. warning:: + Only :class:`torch.distributed._shard.sharding_spec.ChunkShardingSpec` is + currently supported as the ``sharding_spec``. + """ + # Perform some validation first. + if not hasattr(module, param_name): + raise AttributeError(f"{module._get_name()} has no attribute `{param_name}`") + + tensor = getattr(module, param_name) + if not isinstance(tensor, torch.Tensor): + raise ValueError( + f"Expected {type(module).__name__}.{param_name} to be a Tensor, but found {type(tensor).__name__}" + ) + + if not tensor.is_contiguous(): + raise ValueError(f"param: {param_name} is not a contiguous Tensor") + + st = _shard_tensor(tensor, sharding_spec, src_rank, process_group) + + # Replace param with ShardedTensor. + module.register_parameter(param_name, nn.Parameter(st)) + + +# Tracks the current process group in the load context manager. +_CURRENT_PROCESS_GROUP: Optional[dist.ProcessGroup] = None + + +@contextmanager +def load_with_process_group(process_group): + """ + Context manager to set the process group with which to load a ShardedTensor. + """ + global _CURRENT_PROCESS_GROUP + if _CURRENT_PROCESS_GROUP is not None: + raise RuntimeError( + 'ProcessGroup already set by previous "load_with_process_group" ' + "context manager" + ) + _CURRENT_PROCESS_GROUP = process_group + try: + yield process_group + finally: + _CURRENT_PROCESS_GROUP = None + + +def _get_current_process_group(): + """ + Retrieves the current process group set by ``load_with_process_group``. + If not set, it just returns the default group. + """ + global _CURRENT_PROCESS_GROUP + if _CURRENT_PROCESS_GROUP is None: + return distributed_c10d._get_default_group() + else: + return _CURRENT_PROCESS_GROUP + + +def _reshard_output( + module: torch.nn.Module, resharding_spec: ShardingSpec +) -> torch.nn.Module: + """ + Hook a module with output resharding in the forward pass according + to the given ``resharding_spec``. + + Args: + module (:class:`torch.nn.Module`): Module whose output needs to be resharded. + resharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): + The specification describing how the output of the module will be resharded. + + Returns: + A :class:`torch.nn.Module` object with reshard API hooked. + """ + + def hook_func(_module, _input, output): + if isinstance(output, ShardedTensor): + return output.reshard(resharding_spec) + return output + + module.register_forward_hook(hook_func) + return module + + +def _collect_local_shard(module: torch.nn.Module) -> torch.nn.Module: + """ + Hook a module with local shards collection in the forward pass. + + This API is typically used to convert a sharded representation back to data parallel + representation. In particular, it returns the local tensor for this Shard. If the + size along the sharding dimension for the local tensor is 1, this dimension is removed + from the final result. For example a [4, 16] ShardedTensor across 4 ranks is typically + a local Tensor of size [16] across each rank and not [1, 16] across each rank. + + Args: + module (:class:`torch.nn.Module`): Module whose output is ShardedTensor and the + local tensor value needs to be returned. + + Returns: + A :class:`torch.nn.Module` object with collection API hooked. + """ + + def hook_func(_module, _input, output): + if isinstance(output, ShardedTensor): + local_tensor = output.local_tensor() + # Squeeze the # of dimensions manually, only applicable to ChunkShardingSpec + sharding_spec = output._sharding_spec + if ( + isinstance(sharding_spec, ChunkShardingSpec) + and local_tensor.size(sharding_spec.dim) == 1 # type: ignore[attr-defined, arg-type] + ): + local_tensor = local_tensor.squeeze( + output._sharding_spec.dim # type: ignore[attr-defined] + ) + return local_tensor + + module.register_forward_hook(hook_func) + return module + + +def shard_module(module: nn.Module, plan: ShardingPlan, src_rank=0, process_group=None): + """ + Shards a given module according to the provided sharding `plan`. This method + first shards all the parameters according to the given sharding `plan`. Then if + `output_plan` and `return_local_tensor` are specified in the sharding `plan`, it + will tag the output of modules according `output_plan`, convert the module's + output back to data parallel according to `return_local_tensor`. + + Needs to be called on all ranks in an SPMD fashion. + + Args: + module (:class:`torch.nn.Module`): The module to apply sharding to + plan (:class:`torch.distributed._shard.sharding_plan.ShardingPlan`): + The ShardingPlan which specified param name to ShardingSpec to apply to + each parameter. + + Keyword args: + src_rank (int, optional): The source rank which is used as the ground truth of + the data for the module that would be sharded and scattered across the rest + of the ranks. + Default: 0. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + """ + # record Sharder paths for sanity check on the plan to ensure items in the plan + # does not conflict with the submodule tree that the Sharder is working with + sharder_paths = [] + for name, spec in plan.plan.items(): + if isinstance(spec, Sharder): + sharder_paths.append(name) + + # shard the parameter according to the ShardingPlan + for name, spec in plan.plan.items(): + if isinstance(spec, ShardingSpec): + # if found a sharding spec, try to shard the parameter + module_path, _, param_name = name.rpartition(".") + + for sharder_path in sharder_paths: + if module_path.startswith(sharder_path): + raise RuntimeError( + f"ShardingPlan is in-valid, trying to shard a parameter: {name}," + f" but there's already a Sharder entry for module {sharder_path}," + f" parameter sharding should not conflict with the submodule tree" + f" that a Sharder is working with!" + ) + + mod = module.get_submodule(module_path) + shard_parameter( + mod, param_name, spec, src_rank=src_rank, process_group=process_group + ) + elif isinstance(spec, Sharder): + parent_mod_path, _, _mod_name = name.rpartition(".") + if name == "": + raise KeyError("Module path must not be empty for custom sharder!") + mod = module.get_submodule(name) + parent_mod = module.get_submodule(parent_mod_path) + sharded_mod = spec.shard(mod) + # swap this submodule with the sharded module + parent_mod.mod_name = sharded_mod + else: + raise TypeError( + f"Only `ShardingSpec` and `Sharder` are supported to shard '{name}'" + ) + + # reshard output if there's an entry in `reshard_output` for this module + if plan.output_plan is not None: + for module_path, output_spec in plan.output_plan.items(): + if isinstance(output_spec, ShardingSpec): + mod = module.get_submodule(module_path) + _reshard_output(mod, output_spec) + else: + raise TypeError( + f"Only `ShardingSpec` is supported as output_plan for '{module_path}'" + ) + # convert the output back to data parallel for the modules appears in + # `return_local_tensor` of the plan, we will call `_collect_local_shard` + # to collect the local tensor for output of modules + if plan.return_local_tensor is not None: + for module_path in plan.return_local_tensor: + mod = module.get_submodule(module_path) + _collect_local_shard(mod) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/checkpoint/__init__.py b/lib/python3.12/site-packages/torch/distributed/_shard/checkpoint/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..85915636a014640d8fff5a29db602c4a114f1b1d --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/checkpoint/__init__.py @@ -0,0 +1,19 @@ +# Keep old package for BC purposes, this file should be removed once +# everything moves to the `torch.distributed.checkpoint` package. +import sys +import warnings + +import torch +from torch.distributed.checkpoint import * # noqa: F403 + + +with warnings.catch_warnings(): + warnings.simplefilter("always") + warnings.warn( + "`torch.distributed._shard.checkpoint` will be deprecated, " + "use `torch.distributed.checkpoint` instead", + DeprecationWarning, + stacklevel=2, + ) + +sys.modules["torch.distributed._shard.checkpoint"] = torch.distributed.checkpoint diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/checkpoint/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/checkpoint/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36342170e8b399b1b2b599dad2ee9343dd504e39 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/checkpoint/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/common_op_utils.py b/lib/python3.12/site-packages/torch/distributed/_shard/common_op_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e2573998712b5f0f73805d89566ef035abac52d7 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/common_op_utils.py @@ -0,0 +1,65 @@ +# mypy: allow-untyped-defs +from typing import Optional + +import torch +from torch.utils import _pytree as pytree + + +def _basic_validation(op, args=(), kwargs=None): + """ + Common validation across all ops go in here. + """ + from torch.distributed._shard.sharded_tensor import ShardedTensor + + if len(args) == 0 and (kwargs is None or len(kwargs) == 0): + raise ValueError(f" No input for '{op.__name__}'!") + + # Validate types + has_distributed_tensor = False + + def is_distributed_tensor(e): + nonlocal has_distributed_tensor + if isinstance(e, ShardedTensor): + has_distributed_tensor = True + + pytree.tree_map_(is_distributed_tensor, args) + pytree.tree_map_(is_distributed_tensor, kwargs) + + if not has_distributed_tensor: + raise TypeError( + f"torch function '{op.__name__}', with args: {args} and " + f"kwargs: {kwargs} are called without any distributed tensor!" + ) + + # Validate all distributed tensors use the same PG. + cur_pg: Optional[torch.distributed.ProcessGroup] = None + + def validate_pg(e): + nonlocal cur_pg + if isinstance(e, ShardedTensor): + if cur_pg is not None and e._process_group is not cur_pg: + raise RuntimeError( + "All distributed tensors should use the " + "same ProcessGroup if used together in an op." + ) + cur_pg = e._process_group + + pytree.tree_map_(validate_pg, args) + pytree.tree_map_(validate_pg, kwargs) + + +def _register_default_op(op, decorator): + @decorator(op) + def tensor_default_op(types, args=(), kwargs=None, pg=None): + """ + Handles ``__torch_function__`` dispatch for the default tensor ops that + behave the same as ``torch.Tensor`` such as ``torch.Tensor.shape`` or + ``torch.Tensor.dtype``. We simply lower to the real op call with + DisableTorchFunctionSubclass context like ``torch.Tensor.__torch_function__`` + to avoid recursions. + """ + if kwargs is None: + kwargs = {} + + with torch._C.DisableTorchFunctionSubclass(): + return op(*args, **kwargs) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/metadata.py b/lib/python3.12/site-packages/torch/distributed/_shard/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..2611d13ef3aafec39973cc619bf76b3fff9f41b0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/metadata.py @@ -0,0 +1,64 @@ +# mypy: allow-untyped-defs +from dataclasses import dataclass +from functools import reduce +from typing import List, Optional, Union + +from torch.distributed.remote_device import _remote_device + + +@dataclass +class ShardMetadata: + """ + Represents a shard of the overall Tensor including its + offsets, lengths and device placement. + + Args: + shard_offsets(List[int]): Offsets in the original tensor indicating + the start offsets for this shard. Should have the same rank as + the original tensor. + shard_sizes(List[int]): Integers indicating the size of each + dimension for this shard. Should have the same rank as the + original tensor. + placement(:class:`torch.distributed._remote_device`): + Specifies the placement of this shard. + """ + + __slots__ = ["shard_offsets", "shard_sizes", "placement"] + + shard_offsets: List[int] + shard_sizes: List[int] + placement: Optional[_remote_device] + + def __init__( + self, + shard_offsets: List[int], + shard_sizes: List[int], + placement: Optional[Union[str, _remote_device]] = None, + ): + self.shard_offsets = shard_offsets + self.shard_sizes = shard_sizes + if isinstance(placement, str): + self.placement = _remote_device(placement) + else: + self.placement = placement + if len(self.shard_offsets) != len(self.shard_sizes): + raise ValueError( + f"shard_offsets and shard_sizes should have " + f"the same number of elements, found {len(self.shard_offsets)} " + f"and {self.shard_sizes} respectively" + ) + + for i in range(len(self.shard_offsets)): + if self.shard_offsets[i] < 0: + raise ValueError("shard_offsets should be >=0") + if self.shard_sizes[i] < 0: + raise ValueError("shard_sizes should be >= 0") + + def __hash__(self): + def _hash_reduce(a, b): + return (a << 8) + hash(b) + + res = reduce(_hash_reduce, self.shard_offsets, 37) + res = reduce(_hash_reduce, self.shard_sizes, res) + res = _hash_reduce(res, self.placement) + return res diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/op_registry_utils.py b/lib/python3.12/site-packages/torch/distributed/_shard/op_registry_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..12e0b1895e2f053e6c4a975cb6d3c0118baf50bb --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/op_registry_utils.py @@ -0,0 +1,41 @@ +# mypy: allow-untyped-defs +import functools +from inspect import signature + +from .common_op_utils import _basic_validation + + +""" +Common utilities to register ops on ShardedTensor +and PartialTensor. +""" + + +def _register_op(op, func, op_table): + """ + Performs basic validation and registers the provided op in the given + op_table. + """ + if len(signature(func).parameters) != 4: + raise TypeError( + f"Custom sharded op function expects signature: " + f"(types, args, kwargs, process_group), but received " + f"signature: {signature(func)}" + ) + + op_table[op] = func + + +def _decorator_func(wrapped_func, op, op_table): + """ + Decorator function to register the given ``op`` in the provided + ``op_table`` + """ + + @functools.wraps(wrapped_func) + def wrapper(types, args, kwargs, process_group): + _basic_validation(op, args, kwargs) + return wrapped_func(types, args, kwargs, process_group) + + _register_op(op, wrapper, op_table) + return wrapper diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/__init__.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d1508208c169070bf12fb866e27e3acaab9a2eec --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/__init__.py @@ -0,0 +1,52 @@ +from typing import Iterator, Tuple, Union + +import torch.nn as nn +from torch.distributed._shard.sharded_tensor import ShardedTensor + +from .api import ShardedOptimizer + + +def named_params_with_sharded_tensor( + module: nn.Module, + prefix: str = "", + recurse: bool = True, +) -> Iterator[Tuple[str, Union[nn.Parameter, ShardedTensor]]]: + r"""Returns an iterator over module parameters (together with the + ShardedTensor parameters), yielding both the name of the parameter + as well as the parameter itself. This is typically passed to a + :class:torch.distributed._shard.sharded_optim.ShardedOptimizer + + Args: + prefix (str): prefix to prepend to all parameter names. + recurse (bool): if True, then yields parameters of this module + and all submodules. Otherwise, yields only parameters that + are direct members of this module. + + Yields: + (str, Union[Tensor, ShardedTensor]): Tuple containing + the name and parameter (or ShardedTensor parameter) + + Example:: + + >>> # xdoctest: +SKIP + >>> model = torch.nn.Linear(*linear_size) + >>> shard_parameter(model, "weight", spec) + >>> for name, param in named_params_with_sharded_tensor(model): + >>> if name in ['weight']: + >>> print(param.size()) + + """ + modules = module.named_modules(prefix=prefix) if recurse else [(prefix, module)] + + memo = set() + for mod_prefix, mod in modules: + # find all sharded tensor params + for name, val in vars(mod).items(): + if isinstance(val, ShardedTensor) and val not in memo: + memo.add(val) + name = mod_prefix + ("." if mod_prefix else "") + name + yield name, val + + # find all nn.Parameters + for name, val in module.named_parameters(): + yield name, val diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c024aa7c970264ade08321e885ba66b9d353ac2 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97ed384f00f95fb1a9b6be559b29f4b133cbdd70 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/api.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/api.py new file mode 100644 index 0000000000000000000000000000000000000000..c0ddbece766b01f2a1b0e17544fe8061dd0cf805 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_optim/api.py @@ -0,0 +1,101 @@ +# mypy: allow-untyped-defs +from typing import Any, Dict, List, Mapping, Union + +import torch.optim as optim +from torch import Tensor +from torch.distributed._shard.sharded_tensor import ShardedTensor + + +class ShardedOptimizer(optim.Optimizer): + def __init__( + self, + named_params: Mapping[str, Union[Tensor, ShardedTensor]], + optimizer_class, + *optimizer_args, + **optimizer_kwargs, + ): + """ + ShardedOptimizer collects all tensors and local shard tensors of + ShardedTensor, then use these tensors as ``params`` for optimizers + + Args: + named_params (Dict[str, Union[Tensor, ShardedTensor]]) : a Dict + of parameters, where key is the parameter key, value is either + Tensor or ShardedTensor parameter. + optimizer_class (torch.optim.Optimizer): the Optimizer to use + locally, i.e. torch.optim.SGD, torch.optim.Adagrad, etc. + *optimizer_args: the arguments to initialize the optimizer. + **optimizer_kwargs: the key-word arguments to initialize the optimizer. + + """ + tensors: List[Tensor] = [] + for value in named_params.values(): + if isinstance(value, ShardedTensor): + tensors.extend( + local_shard.tensor for local_shard in value.local_shards() + ) + else: + tensors.append(value) + + self.named_params = named_params + self._optim = optimizer_class(tensors, *optimizer_args, **optimizer_kwargs) + self.param_groups = self._optim.param_groups + self.state = self._optim.state + + def zero_grad(self, set_to_none: bool = True): # type: ignore[override] + r"""Resets the gradients of all optimized :class:`torch.Tensor` s. + + Args: + set_to_none (bool): instead of setting to zero, set the grads to None. + This will in general have lower memory footprint, and can modestly improve performance. + However, it changes certain behaviors. For example: + 1. When the user tries to access a gradient and perform manual ops on it, + a None attribute or a Tensor full of 0s will behave differently. + 2. If the user requests ``zero_grad(set_to_none=True)`` followed by a backward pass, ``.grad``\ s + are guaranteed to be None for params that did not receive a gradient. + 3. ``torch.optim`` optimizers have a different behavior if the gradient is 0 or None + (in one case it does the step with a gradient of 0 and in the other it skips + the step altogether). + """ + self._optim.zero_grad(set_to_none) + + def step(self, closure=None): + r"""Performs a single optimization step (parameter update). + + Args: + closure (Callable): A closure that reevaluates the model and + returns the loss. Optional for most optimizers. + + .. note:: + Unless otherwise specified, this function should not modify the + ``.grad`` field of the parameters. + """ + self._optim.step(closure) + + def state_dict(self) -> Dict[str, Any]: + """ + Returned state and param_groups will contain parameter keys + instead of parameter indices like torch.optim.Optimizer. + This allows for advanced functionality like optimizer re-sharding to be implemented. + """ + # TODO: implement state_dict + raise NotImplementedError("ShardedOptimizer state_dict not implemented yet!") + + def load_state_dict(self, state_dict: Mapping[str, Any]): + r"""Loads the ShardedOptimizer state. + + Args: + state_dict (dict): ShardedOptimizer state. Should be an object returned + from a call to :meth:`state_dict`. + """ + # TODO: implement load_state_dict + raise NotImplementedError( + "ShardedOptimizer load_state_dict not implemented yet!" + ) + + def add_param_group(self, param_group: Any): + r"""Add a new param group""" + # TODO: implement add_param_group + raise NotImplementedError( + "ShardedOptimizer add_param_group not implemented yet!" + ) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__init__.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db7090820ea0a960446ef80fabd50085013e5dc4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__init__.py @@ -0,0 +1,490 @@ +# mypy: allow-untyped-defs +import functools +from typing import List, TYPE_CHECKING + +import torch +from torch.distributed._shard.op_registry_utils import _decorator_func + +from .api import ( + _CUSTOM_SHARDED_OPS, + _SHARDED_OPS, + Shard, + ShardedTensor, + ShardedTensorBase, + ShardedTensorMetadata, + TensorProperties, +) +from .metadata import ShardMetadata # noqa: F401 + + +if TYPE_CHECKING: + from torch.distributed._shard.sharding_spec import ShardingSpec +else: + ShardingSpec = "ShardingSpec" + + +def empty( + sharding_spec: ShardingSpec, + *size, + dtype=None, + layout=torch.strided, + requires_grad=False, + pin_memory=False, + memory_format=torch.contiguous_format, + process_group=None, + init_rrefs=False, +) -> ShardedTensor: + """ + Returns a :class:`ShardedTensor` filled with uninitialized data. + Needs to be called on all ranks in an SPMD fashion. + + Args: + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The specification + describing how to shard the Tensor. + size (int...): a sequence of integers defining the shape of the output + tensor. Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.contiguous_format``. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + init_rrefs (bool, optional): Whether or not to initialize + :class:`torch.distributed.rpc.RRef`s pointing to remote shards. + Need to initialize the RPC Framework if specified as ``True``. + Default: ``False``. + + Returns: + A :class:`ShardedTensor` object on each rank + """ + return ShardedTensor( + sharding_spec, + *size, + dtype=dtype, + layout=layout, + requires_grad=requires_grad, + pin_memory=pin_memory, + memory_format=memory_format, + process_group=process_group, + init_rrefs=init_rrefs, + ) + + +def ones( + sharding_spec: ShardingSpec, + *size, + dtype=None, + layout=torch.strided, + requires_grad=False, + pin_memory=False, + memory_format=torch.contiguous_format, + process_group=None, + init_rrefs=False, +) -> ShardedTensor: + """ + Returns a :class:`ShardedTensor` with the scalar value 1. + Needs to be called on all ranks in an SPMD fashion. + + Args: + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The specification + describing how to shard the Tensor. + size (int...): a sequence of integers defining the shape of the output + tensor. Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + init_rrefs (bool, optional): Whether or not to initialize + :class:`torch.distributed.rpc.RRef`s pointing to remote shards. + Need to initialize the RPC Framework if specified as ``True``. + Default: ``False``. + + Returns: + A :class:`ShardedTensor` object on each rank + """ + return full( + sharding_spec, + size, + fill_value=1, + dtype=dtype, + layout=layout, + requires_grad=requires_grad, + pin_memory=pin_memory, + memory_format=memory_format, + process_group=process_group, + init_rrefs=init_rrefs, + ) + + +def zeros( + sharding_spec: ShardingSpec, + *size, + dtype=None, + layout=torch.strided, + requires_grad=False, + pin_memory=False, + memory_format=torch.contiguous_format, + process_group=None, + init_rrefs=False, +) -> ShardedTensor: + """ + Returns a :class:`ShardedTensor` filled with the scalar value 0. + Needs to be called on all ranks in an SPMD fashion. + + Args: + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The specification + describing how to shard the Tensor. + size (int...): a sequence of integers defining the shape of the output + tensor. Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + init_rrefs (bool, optional): Whether or not to initialize + :class:`torch.distributed.rpc.RRef`s pointing to remote shards. + Need to initialize the RPC Framework if specified as ``True``. + Default: ``False``. + + Returns: + A :class:`ShardedTensor` object on each rank + """ + return full( + sharding_spec, + size, + fill_value=0, + dtype=dtype, + layout=layout, + requires_grad=requires_grad, + pin_memory=pin_memory, + memory_format=memory_format, + process_group=process_group, + init_rrefs=init_rrefs, + ) + + +def full( + sharding_spec: ShardingSpec, + size, + fill_value, + *, + dtype=None, + layout=torch.strided, + requires_grad=False, + pin_memory=False, + memory_format=torch.contiguous_format, + process_group=None, + init_rrefs=False, +) -> ShardedTensor: + """ + Creates a :class:`ShardedTensor` filled with fill_value. The tensor's dtype + is inferred from fill_value. If dtype is specified, it will override the + inferred type from fill_value. Needs to be called on all ranks in an SPMD fashion. + Args: + sharding_spec (:class:`torch.distributed._sharding_spec.ShardingSpec`): The specification + describing how to shard the Tensor. + size (int...): a list, tuple, or `torch.Size` of integers defining the shape of the + output tensor. + fill_value (Scalar) - the value to fill the output tensor with. + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + init_rrefs (bool, optional): Whether or not to initialize + :class:`torch.distributed.rpc.RRef`s pointing to remote shards. + Need to initialize the RPC Framework if specified as ``True``. + Default: ``False``. + Returns: + A :class:`ShardedTensor` object on each rank + """ + sharded_tensor = ShardedTensor( + sharding_spec, + *size, + dtype=dtype, + layout=layout, + requires_grad=requires_grad, + pin_memory=pin_memory, + memory_format=memory_format, + process_group=process_group, + init_rrefs=init_rrefs, + ) + torch.nn.init.constant_(sharded_tensor, fill_value) # type: ignore[arg-type] + return sharded_tensor + + +def rand( + sharding_spec: ShardingSpec, + *size, + dtype=None, + layout=torch.strided, + requires_grad=False, + pin_memory=False, + memory_format=torch.contiguous_format, + process_group=None, + init_rrefs=False, +) -> ShardedTensor: + """ + Creates a :class:`ShardedTensor` filled with random numbers from a uniform distribution + on the interval :math:`[0, 1)`. The shape of the tensor is defined by the + variable argument `size`. Needs to be called on all ranks in an SPMD fashion. + + Args: + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The specification + describing how to shard the Tensor. + size (int...): a list, tuple, or `torch.Size` of integers defining the shape of the + output tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + init_rrefs (bool, optional): Whether or not to initialize + :class:`torch.distributed.rpc.RRef`s pointing to remote shards. + Need to initialize the RPC Framework if specified as ``True``. + Default: ``False``. + + Returns: + A :class:`ShardedTensor` object on each rank + """ + sharded_tensor = ShardedTensor( + sharding_spec, + *size, + dtype=dtype, + layout=layout, + requires_grad=requires_grad, + pin_memory=pin_memory, + memory_format=memory_format, + process_group=process_group, + init_rrefs=init_rrefs, + ) + torch.nn.init.uniform_(sharded_tensor, 0, 1) # type: ignore[arg-type] + return sharded_tensor + + +def randn( + sharding_spec: ShardingSpec, + *size, + dtype=None, + layout=torch.strided, + requires_grad=False, + pin_memory=False, + memory_format=torch.contiguous_format, + process_group=None, + init_rrefs=False, +) -> ShardedTensor: + """ + Creates a :class:`ShardedTensor` filled with random numbers from a uniform distribution + with mean `0` and variance `1` (also called standard normal distribution). The shape + of the tensor is defined by the variable argument `size`. Needs to be called on all ranks + in an SPMD fashion. + + Args: + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The specification + describing how to shard the Tensor. + size (int...): a list, tuple, or `torch.Size` of integers defining the shape of the + output tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + init_rrefs (bool, optional): Whether or not to initialize + :class:`torch.distributed.rpc.RRef`s pointing to remote shards. + Need to initialize the RPC Framework if specified as ``True``. + Default: ``False``. + + Returns: + A :class:`ShardedTensor` object on each rank + """ + sharded_tensor = ShardedTensor( + sharding_spec, + *size, + dtype=dtype, + layout=layout, + requires_grad=requires_grad, + pin_memory=pin_memory, + memory_format=memory_format, + process_group=process_group, + init_rrefs=init_rrefs, + ) + torch.nn.init.normal_(sharded_tensor, 0, 1) # type: ignore[arg-type] + return sharded_tensor + + +def init_from_local_shards( + local_shards: List[Shard], *global_size, process_group=None, init_rrefs=False +) -> ShardedTensor: + """ + Creates an :class:`ShardedTensor` from local shards and the global metadata. + Needs to be called on all ranks in an SPMD fashion. + + Args: + local_shards (List[:class `torch.distributed._shard.sharded_tensor.Shard`]): A list + of shards that represent the local shards on this rank. + global_size (int...): a list, tuple, or `torch.Size` of integers defining the + shape of the overall sharded tensor. + + Keyword args: + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + init_rrefs (bool, optional): Whether or not to initialize + :class:`torch.distributed.rpc.RRef`s pointing to remote shards. + Need to initialize the RPC Framework if specified as ``True``. + Default: ``False``. + + Returns: + A :class:`ShardedTensor` object handle on this rank + + + Examples: + Suppose we want construct a sharded tensor on two ranks, global size = (10, 5), + each shard have a (5, 5) local tensor, we can do it like below: + + on rank 0: + >>> # xdoctest: +SKIP("not distributed") + >>> local_shard_metadata = ShardMetadata( + >>> shard_offsets=[0, 0], + >>> shard_lengths=[5, 5], + >>> placement="rank:0/cuda:0" + >>> ) + >>> local_shards = [Shard(torch.randn(5, 5), local_shard_metadata)] + >>> sharded_tensor = init_from_local_shards(local_shards, [10, 5]) + + on rank 1: + >>> # xdoctest: +SKIP("not distributed") + >>> local_shard_metadata = ShardMetadata( + >>> shard_offsets=[5, 0], + >>> shard_lengths=[5, 5], + >>> placement="rank:1/cuda:1" + >>> ) + >>> local_shards = [Shard(torch.randn(5, 5), local_shard_metadata)] + >>> sharded_tensor = init_from_local_shards(local_shards, [10, 5]) + """ + return ShardedTensor._init_from_local_shards( + local_shards, *global_size, process_group=process_group, init_rrefs=init_rrefs + ) + + +def state_dict_hook(module, destination, prefix, local_metadata): + """ + Hook to add ShardedTensor to Module's ``state_dict``. Needs to be + registered to the Module using + :meth:`torch.nn.Module._register_state_dict_hook`. + """ + for submodule_name, submodule in module.named_modules(): + for attr_name, attr in submodule.__dict__.items(): + if isinstance(attr, ShardedTensor): + mod_prefix = prefix + submodule_name + key = mod_prefix + ("." if mod_prefix else "") + attr_name + destination[key] = attr + + +def pre_load_state_dict_hook( + module, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, +): + """ + Pre-load state dict hook to add ShardedTensor to the module. + """ + for submodule_name, submodule in module.named_modules(): + for attr_name in submodule.__dict__.keys(): + mod_prefix = prefix + submodule_name + key = mod_prefix + ("." if mod_prefix else "") + attr_name + if key in state_dict: + if isinstance(state_dict[key], ShardedTensor): + setattr(submodule, attr_name, state_dict[key]) + + +def custom_sharded_op_impl(func): + """ + Provides a way for users to write their own custom sharded operator. This + can be used to override existing ShardedTensor operators or write a new + one not supported by ShardedTensor. If the operator in question is covered + by ``__torch_function__`` dispatch and has a ShardedTensor as any of its + parameters, the function provided will be invoked for that operator. + + Example:: + >>> # xdoctest: +SKIP + >>> @custom_sharded_op_impl(torch.nn.functional.linear) + >>> def my_custom_sharded_linear(types, args, kwargs, process_group): + >>> ... + >>> # xdoctest: +SKIP("Undefined variables") + >>> input = torch.rand(10, 32) + >>> weight = sharded_tensor.rand(32, 16) + >>> bias = torch.rand(16) + >>> # This will call 'my_custom_sharded_linear' + >>> torch.nn.functional.linear(input, weight, bias) + + The types, args and kwargs parameters are the same parameters that are + passed to ``__torch_function__`` dispatch API + (https://pytorch.org/docs/stable/notes/extending.html#extending-torch). + There is an additional ``process_group`` parameter which is the + process_group used for the ShardedTensor and can be used by + implementations for communications within a sharded implementation. + + Args: + func(Callable): Torch function for which we want to provide a sharded + implementation (ex: torch.nn.functional.linear) + """ + return functools.partial(_decorator_func, op=func, op_table=_CUSTOM_SHARDED_OPS) + + +def _sharded_op_impl(func): + """ + Decorator to register a default sharded op. + """ + return functools.partial(_decorator_func, op=func, op_table=_SHARDED_OPS) + + +# Import all builtin sharded ops +from ._ops import * # noqa: F403 diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a6df6ac77dc2d9749f2089b18ac3fd00ed1823f Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b40d461c357f765f6643e1917d593d68379b9c09 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/logger.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..61787ad1c4095e7dd51a655dec000a0125c2b390 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/logging_handlers.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/logging_handlers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4ec91cf6e6e97917eca2fb11d27aa4b14e74637 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/logging_handlers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/metadata.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/metadata.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..adb1ff1114d287271fc0ad087960abcb1265ac4f Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/metadata.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/reshard.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/reshard.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30c77625759cf756ec303c01a292ade5c411b15d Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/reshard.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/shard.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/shard.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca2a03c29ddaa10cad05d4df6bff41360ee16eec Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/shard.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..313e4a02a28da7da02eec013f0a7d9a0820ecadd Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__init__.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..be6d01fc8e54ee214fafa847c9261db375d8b87e --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__init__.py @@ -0,0 +1,13 @@ +import torch.distributed._shard.sharded_tensor._ops.misc_ops +import torch.distributed._shard.sharded_tensor._ops.tensor_ops + +# Import all ChunkShardingSpec ops +from torch.distributed._shard.sharding_spec.chunk_sharding_spec_ops.embedding import ( + sharded_embedding, +) +from torch.distributed._shard.sharding_spec.chunk_sharding_spec_ops.embedding_bag import ( + sharded_embedding_bag, +) + +from .binary_cmp import allclose, equal +from .init import constant_, kaiming_uniform_, normal_, uniform_ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d5f80ef2a4b966f00df2defe9a0aeb8474408bf3 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/_common.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7522535f51ed81a905479d01d7b697da67acb505 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/_common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/binary_cmp.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/binary_cmp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd79201c8f00ad2c00f31cd290c66f18ac1d9b7b Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/binary_cmp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/init.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/init.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41b12c74d14f44207dabb24b76836315a2bf5cef Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/init.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/misc_ops.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/misc_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c13c12093ef07e43eb778f134f1c58af3e004c95 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/misc_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/tensor_ops.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/tensor_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a27aec62f6d527bd685a9776d5d9a110adb9f9b6 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/__pycache__/tensor_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/_common.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..502e0ac9a8552dae35c3fa45f52d73d1c2e82067 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/_common.py @@ -0,0 +1,113 @@ +# mypy: allow-untyped-defs +import functools + +from torch.distributed._shard.common_op_utils import _basic_validation +from torch.distributed._shard.sharded_tensor import ( + _sharded_op_impl, + Shard, + ShardedTensor, +) + + +def _sharded_op_common(op, early_stop_func, extra_check): + """ + Inject sharded tensor op registration with common logics executed before + different behaviors are done on either local shards or a local tensor. + + Example:: + >>> # xdoctest: +SKIP("Undefined variables") + >>> op = torch.transpose + >>> @_sharded_op_impl(op) + >>> @_sharded_op_common(op, early_stop_func, extra_check) + >>> def sharded_tensor_op(types, args, kwargs, process_group): + >>> ... + >>> + >>> st = sharded_tensor.rand(32, 16) + >>> st.transpose(1, 2) + >>> # This will call '_sharded_op_common' + + Args: + op: The op to be registered and applied to all shards of the st. + early_stop_func (Callable, optional): the func for early stop. + Default: if ``None``, no early stop. + extra_check (Callable, optional): the func for extra condition check. + Default: if ``None``, no extra check. + + Return: + func (Callable): Torch function for which we want to provide a sharded + implementation (ex: torch.transpose) + """ + + def decorator_sharded_func(wrapped_func): + @functools.wraps(wrapped_func) + def wrapper(types, args=(), kwargs=None, pg=None): + _basic_validation(op, args, kwargs) + + st = args[0] + if kwargs is None: + kwargs = {} + if extra_check: + extra_check(*args, **kwargs) + if early_stop_func: + early_stop = early_stop_func(*args, **kwargs) + if early_stop: + return st + return wrapped_func(types, args, kwargs, pg) + + return wrapper + + return decorator_sharded_func + + +def _register_sharded_op_on_local_shards( + op, early_stop_func=None, extra_check=None, customized_func=None +): + """ + Handles ``__torch_function__`` dispatch for ops which are performed on + each shard of the sharded tensor such as elementwise op like + ``torch.nn.functional.gelu`` or ``torch.nn.functional.relu``. + + For more complicated ops, a customized func can be used to generate + the new shards and sharded tensor size. + + This function expects that the original ShardingSpec for the ShardedTensor + is preserved irrespective of whether or not a customized function is used. + + Args: + op: The op to be registered and applied to all shards of the st. + early_stop_func (Callable, optional): the func for early stop. + Default: if ``None``, no early stop. + extra_check (Callable, optional): the func for extra condition check. + Default: if ``None``, no extra check. + customized_func (Callable, optional): the func for customized logic + to generate new shards and sharded tensor size. + Default: if ``None``, we simply lower to the real op call with + all local shards of the st. + + Return: + func (Callable): registered implementation for sharded op for + ``__torch_function__`` dispatch. + """ + + @_sharded_op_impl(op) + @_sharded_op_common(op, early_stop_func, extra_check) + def sharded_tensor_op_on_local_shards(types, args=(), kwargs=None, pg=None): + st = args[0] + st_metadata = st.metadata() + local_shards = st.local_shards() + local_shards_new = [] + if customized_func: + local_shards_new, st_metadata = customized_func(args, kwargs, pg) + else: + for local_shard in local_shards: + args = (local_shard.tensor, *args[1:]) + local_shards_new.append( + Shard(op(*args, **kwargs), local_shard.metadata) + ) + return ShardedTensor._init_from_local_shards_and_global_metadata( + local_shards_new, + st_metadata, + process_group=pg, + init_rrefs=st._init_rrefs, + sharding_spec=st.sharding_spec(), + ) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/binary_cmp.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/binary_cmp.py new file mode 100644 index 0000000000000000000000000000000000000000..0548b81fb90af087593d05695418664c6d109f2d --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/binary_cmp.py @@ -0,0 +1,78 @@ +# mypy: allow-untyped-defs +import torch +import torch.distributed as dist +import torch.distributed.distributed_c10d as distributed_c10d +from torch.distributed._shard.sharded_tensor import _sharded_op_impl, ShardedTensor + + +def _communicate_result(result, pg): + # Gather results from all ranks. + if result: + result_tensor = torch.ones(1, device=torch.device(torch.cuda.current_device())) + else: + result_tensor = torch.zeros(1, device=torch.device(torch.cuda.current_device())) + + dist.all_reduce(result_tensor, group=pg) + + expected_result = torch.ones( + 1, device=torch.device(torch.cuda.current_device()) + ) * dist.get_world_size(pg) + + return torch.equal(result_tensor, expected_result) + + +def binary_cmp(cmp_fun, types, args, kwargs=None, process_group=None): + if len(args) != 2: + raise ValueError(f"Expected two arguments for torch.{cmp_fun.__name__}") + + st1 = args[0] + st2 = args[1] + if not (isinstance(st1, ShardedTensor) and isinstance(st2, ShardedTensor)): + raise TypeError( + f"Both arguments to torch.{cmp_fun.__name__} need to be of type ShardedTensor" + ) + + # Verify same PG + if st1._process_group != st2._process_group: + return False + + if distributed_c10d._rank_not_in_group( + st1._process_group + ) or distributed_c10d._rank_not_in_group(st2._process_group): + return distributed_c10d._rank_not_in_group( + st1._process_group + ) == distributed_c10d._rank_not_in_group(st2._process_group) + + # Verify metadata + if st1.metadata() != st2.metadata(): + return _communicate_result(False, st1._process_group) + + # Verify number of local shards + st1_local_shards = st1.local_shards() + st2_local_shards = st2.local_shards() + if len(st1_local_shards) != len(st2_local_shards): + return _communicate_result(False, st1._process_group) + + # kwargs must be dict-like + if kwargs is None: + kwargs = {} + # Verify each local shard + for idx in range(len(st1_local_shards)): + if st1_local_shards[idx].metadata != st2_local_shards[idx].metadata: + return _communicate_result(False, st1._process_group) + if not cmp_fun( + st1_local_shards[idx].tensor, st2_local_shards[idx].tensor, **kwargs + ): + return _communicate_result(False, st1._process_group) + + return _communicate_result(True, st1._process_group) + + +@_sharded_op_impl(torch.equal) +def equal(types, args, kwargs, process_group): + return binary_cmp(torch.equal, types, args, kwargs, process_group) + + +@_sharded_op_impl(torch.allclose) +def allclose(types, args, kwargs, process_group): + return binary_cmp(torch.allclose, types, args, kwargs, process_group) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/init.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/init.py new file mode 100644 index 0000000000000000000000000000000000000000..71a9c20b45352cd1526ab00f3b4463bdaefe982a --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/init.py @@ -0,0 +1,151 @@ +# mypy: allow-untyped-defs +import torch +import torch.distributed._shard.sharded_tensor as sharded_tensor +from torch.distributed._shard.sharded_tensor import _sharded_op_impl + + +def validate_param(param, param_name): + if param is None: + raise ValueError(f"param: {param_name} shouldn't be None!") + + +@_sharded_op_impl(torch.nn.init.uniform_) +def uniform_(types, args=(), kwargs=None, pg=None): + r""" + Fills the Tensor in tensor.local_shards with values drawn from the uniform + distribution :math:`\mathcal{U}(a, b)`. + Args: + tensor: tensor sharded across devices + a: the lower bound of the uniform distribution + b: the upper bound of the uniform distribution + """ + validate_param(kwargs, "kwargs") + sharded_tensor = kwargs["tensor"] + validate_param(sharded_tensor, "tensor") + a = kwargs["a"] + validate_param(a, "a") + b = kwargs["b"] + validate_param(b, "b") + + for shard in sharded_tensor.local_shards(): + torch.nn.init.uniform_(shard.tensor, a=a, b=b) + return sharded_tensor + + +@_sharded_op_impl(torch.nn.init.normal_) +def normal_(types, args=(), kwargs=None, pg=None): + r""" + Fills the Tensors in tensor.local_shards with values drawn from the normal + distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`. + Args: + tensor: tensor sharded across devices + mean: the mean of the normal distribution + std: the standard deviation of the normal distribution + """ + validate_param(kwargs, "kwargs") + sharded_tensor = kwargs["tensor"] + validate_param(sharded_tensor, "tensor") + mean = kwargs["mean"] + validate_param(mean, "mean") + std = kwargs["std"] + validate_param(std, "std") + + for shard in sharded_tensor.local_shards(): + torch.nn.init.normal_(shard.tensor, mean=mean, std=std) + return sharded_tensor + + +@_sharded_op_impl(torch.nn.init.kaiming_uniform_) +def kaiming_uniform_(types, args=(), kwargs=None, pg=None): + r""" + Fills the Tensors in tensor.local_shards with values according to the method + described in `Delving deep into rectifiers: Surpassing human-level + performance on ImageNet classification` - He, K. et al. (2015), using a + uniform distribution. The resulting tensor will have values sampled from + :math:`\mathcal{U}(-\text{bound}, \text{bound})` where + .. math:: + \text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}} + Also known as He initialization. + Args: + tensor: tensor sharded across devices + a: the negative slope of the rectifier used after this layer (only + used with ``'leaky_relu'``) + mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'`` + preserves the magnitude of the variance of the weights in the + forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the + backwards pass. + nonlinearity: the non-linear function (`nn.functional` name), + recommended to use only with ``'relu'`` or ``'leaky_relu'`` (default). + """ + validate_param(kwargs, "kwargs") + sharded_tensor = kwargs["tensor"] + validate_param(sharded_tensor, "tensor") + a = kwargs["a"] + validate_param(a, "a") + mode = kwargs["mode"] + validate_param(mode, "mode") + nonlinearity = kwargs["nonlinearity"] + validate_param(nonlinearity, "nonlinearity") + + for shard in sharded_tensor.local_shards(): + torch.nn.init.kaiming_uniform_( + shard.tensor, a=a, mode=mode, nonlinearity=nonlinearity + ) + return sharded_tensor + + +@_sharded_op_impl(torch.nn.init.constant_) +def constant_(types, args=(), kwargs=None, pg=None): + r""" + Fills the input ShardedTensor with the value \text{val}val. + Args: + tensor: tensor sharded across devices + val: the value to fill the tensor with + """ + validate_param(kwargs, "kwargs") + sharded_tensor = kwargs["tensor"] + validate_param(sharded_tensor, "tensor") + val = kwargs["val"] + validate_param(val, "val") + for shard in sharded_tensor.local_shards(): + torch.nn.init.constant_(shard.tensor, val=val) + return sharded_tensor + + +tensor_like_creation_op_map = { + torch.full_like: sharded_tensor.full, + torch.empty_like: sharded_tensor.empty, + torch.zeros_like: sharded_tensor.zeros, + torch.ones_like: sharded_tensor.ones, + torch.rand_like: sharded_tensor.rand, + torch.randn_like: sharded_tensor.randn, +} + + +# tensor ops that behave the same as the default tensor +def register_tensor_creation_op(op): + @_sharded_op_impl(op) + def tensor_creation_op(types, args=(), kwargs=None, pg=None): + """ + Handles ``__torch_function__`` dispatch for tensor creation ops that + takes a ShardedTensor as argument, such as ``torch.zeros_like`` or + ``torch.full_like``. + """ + creation_op = tensor_like_creation_op_map.get(op, None) + if creation_op is None: + raise RuntimeError(f"Tensor creation {op} not supported!") + if kwargs is None: + kwargs = {} + + st = args[0] + + new_st = creation_op(st.sharding_spec(), st.size(), *args[1:], **kwargs) # type: ignore[operator] + return new_st + + +register_tensor_creation_op(torch.full_like) +register_tensor_creation_op(torch.empty_like) +register_tensor_creation_op(torch.zeros_like) +register_tensor_creation_op(torch.ones_like) +register_tensor_creation_op(torch.rand_like) +register_tensor_creation_op(torch.randn_like) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/misc_ops.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/misc_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..8b84c1684c32456989e3998b3d4c30c34cb5dbf4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/misc_ops.py @@ -0,0 +1,12 @@ +# mypy: allow-untyped-defs +import torch +from torch.distributed._shard.sharded_tensor import _sharded_op_impl + + +# This is used by `_apply()` within module.py to set new +# parameters after apply a certain method, we should follow +# the future behavior of overwriting the existing tensor +# instead of doing in-place change using `.data = `. +@_sharded_op_impl(torch._has_compatible_shallow_copy_type) +def tensor_has_compatible_shallow_copy_type(types, args=(), kwargs=None, pg=None): + return False diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..19c475fe817969b7da4811c0b51e86f6cf102896 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py @@ -0,0 +1,219 @@ +# mypy: allow-untyped-defs +import copy + +import torch +from torch.distributed._shard.common_op_utils import _register_default_op +from torch.distributed._shard.sharded_tensor import ( + _sharded_op_impl, + Shard, + ShardedTensor, +) + +from ._common import _register_sharded_op_on_local_shards + + +# Tensor properties access +_register_default_op(torch.Tensor.shape.__get__, _sharded_op_impl) # type: ignore[attr-defined] +_register_default_op(torch.Tensor.dtype.__get__, _sharded_op_impl) # type: ignore[attr-defined] +_register_default_op(torch.Tensor.layout.__get__, _sharded_op_impl) # type: ignore[attr-defined] +_register_default_op(torch.Tensor.size, _sharded_op_impl) +_register_default_op(torch.Tensor.dim, _sharded_op_impl) +_register_default_op(torch.Tensor.ndim.__get__, _sharded_op_impl) # type: ignore[attr-defined] +_register_default_op(torch.Tensor.is_contiguous, _sharded_op_impl) +_register_default_op(torch.Tensor.contiguous, _sharded_op_impl) +_register_default_op(torch.Tensor.is_floating_point, _sharded_op_impl) + +# __reduce_ex__ to dispatch to get_state/set_state +_register_default_op(torch.Tensor.__reduce_ex__, _sharded_op_impl) + +# autograd related properties +_register_default_op(torch.Tensor.requires_grad.__get__, _sharded_op_impl) # type: ignore[attr-defined] +# TODO: set grad with a ShardedTensor that consists of all local grads +_register_default_op(torch.Tensor.grad.__get__, _sharded_op_impl) # type: ignore[union-attr] +_register_default_op(torch.Tensor.grad_fn.__get__, _sharded_op_impl) # type: ignore[union-attr] +_register_default_op(torch.Tensor.is_leaf.__get__, _sharded_op_impl) # type: ignore[attr-defined] + + +# device property is ambiguous as from a global prospective, +# ShardedTensor.device consists of multiple devices (might even across hosts) +# We choose to return the current device of the local tensor to represent +# the device property on each rank +@_sharded_op_impl(torch.Tensor.device.__get__) +def tensor_device(types, args=(), kwargs=None, pg=None): + self_st = args[0] + # Validate types + if not isinstance(self_st, ShardedTensor): + raise TypeError("input needs to be a ShardedTensor") + dev: torch.device + if self_st._local_shards: + dev = self_st._local_shards[0].tensor.device + elif pg and pg._get_backend_name() == "gloo": + dev = torch.device("cpu") + else: + dev = torch.device(torch.cuda.current_device()) + return dev + + +@_sharded_op_impl(torch.Tensor.is_meta.__get__) # type: ignore[attr-defined] +def st_is_meta(types, args=(), kwargs=None, pg=None): + return args[0].local_tensor().is_meta + + +def sharded_type_as_check(*args, **kwargs): + """ + Perform extra checks for the sharded_type_as op such as the input needs to + be either a Tensor or ShardedTensor. + + Args: same as ``torch.Tensor.type_as``. + + Return: None + """ + if len(args) < 2: + raise ValueError("Needs to give a tensor to cast type as!") + if not isinstance(args[1], torch.Tensor) and not isinstance(args[1], ShardedTensor): + raise ValueError("Needs to give a Tensor or ShardedTensor to cast type as!") + + +def same_dtype(*args, **kwargs): + """ + When the dtype is the same, return the original ShardedTensor. + + Args: same as ``torch.Tensor.type_as``. + + Return (bool): Whether to return early or not. + """ + return args[0].dtype == args[1].dtype + + +def sharded_type_as(args, kwargs, pg): + """ + Handles ``__torch_function__`` dispatch for the ``torch.Tensor.type_as`` op. + + Args: same as ``torch.Tensor.type_as``. + + Return: + new_local_shards (List[Shard]): Local shards for the new sharded tensor. + st_meta (ShardedTensorMetadata): Metadata of the new sharded tensor. + """ + st = args[0] + tensor = args[1] + if isinstance(tensor, ShardedTensor): + tensor = tensor.local_tensor() + new_local_shards = [ + Shard(shard.tensor.type_as(tensor), shard.metadata) + for shard in st.local_shards() + ] + st_meta = copy.deepcopy(st._metadata) + st_meta.tensor_properties.dtype = tensor.dtype + return new_local_shards, st_meta + + +_register_sharded_op_on_local_shards( + torch.Tensor.type_as, + early_stop_func=same_dtype, + extra_check=sharded_type_as_check, + customized_func=sharded_type_as, +) + + +def sharded_deepcopy(args, kwargs, pg): + # NOTE: we directly implement deepcopy magic method + # instead of using the default tensor.__deepcopy__ + # and implement clone(). This is because the default + # tensor deepcopy copies every attribute, but the + # process_group in ShardedTensor cannot be deep copied. + self_st = args[0] + new_local_shards = copy.deepcopy(self_st.local_shards()) + new_metadata = copy.deepcopy(self_st.metadata()) + return new_local_shards, new_metadata + + +_register_sharded_op_on_local_shards( + torch.Tensor.__deepcopy__, + customized_func=sharded_deepcopy, +) + + +@_sharded_op_impl(torch.Tensor.copy_) +def sharded_inplace_copy(types, args, kwargs, pg): + # NOTE: inplace op don't need to rewrap + kwargs = {} if kwargs is None else kwargs + self_st = args[0] + new_st = args[1] + nonblocking = kwargs.get("non_blocking", False) + for local_shard, new_shard in zip(self_st.local_shards(), new_st.local_shards()): + if local_shard.metadata != new_shard.metadata: + raise RuntimeError( + "inplace copy can only happen between two ShardedTensor with same metadata!" + ) + for local_shard, new_shard in zip(self_st.local_shards(), new_st.local_shards()): + local_shard.tensor.copy_(new_shard.tensor, nonblocking) + + return self_st + + +def sharded_clone(args, kwargs, pg): + self_st = args[0] + desire_memory_format = kwargs.get("memory_format", None) + if desire_memory_format and desire_memory_format != torch.preserve_format: + raise RuntimeError("Only support torch.preserve_format for ShardedTensor!") + cloned_local_shards = [ + Shard( + local_shard.tensor.clone(memory_format=desire_memory_format), + metadata=copy.deepcopy(local_shard.metadata), + ) + for local_shard in self_st.local_shards() + ] + new_metadata = copy.deepcopy(self_st.metadata()) + return cloned_local_shards, new_metadata + + +_register_sharded_op_on_local_shards( + torch.Tensor.clone, + customized_func=sharded_clone, +) + + +def sharded_detach(args, kwargs, pg): + self_st = args[0] + detached_local_shards = [ + Shard( + local_shard.tensor.detach(), + metadata=copy.deepcopy(local_shard.metadata), + ) + for local_shard in self_st.local_shards() + ] + new_metadata = copy.deepcopy(self_st.metadata()) + new_metadata.tensor_properties.requires_grad = False + return detached_local_shards, new_metadata + + +_register_sharded_op_on_local_shards( + torch.Tensor.detach, + customized_func=sharded_detach, +) + + +@_sharded_op_impl(torch.Tensor.requires_grad_) +def tensor_requires_grad_set(types, args=(), kwargs=None, pg=None): + self_st = args[0] + # Validate types + if not isinstance(self_st, ShardedTensor): + raise TypeError("input needs to be a ShardedTensor") + + if kwargs is None: + kwargs = {} + + requires_grad = args[1] if len(args) > 1 else kwargs.get("requires_grad", True) + if requires_grad == self_st.requires_grad: + return self_st + + for local_shard in self_st.local_shards(): + local_shard.tensor.requires_grad_(requires_grad) + + # update the wrapper class property + with torch._C.DisableTorchFunctionSubclass(): + self_st.requires_grad_(requires_grad) + # update the metadata in the meanwhile + self_st._metadata.tensor_properties.requires_grad = requires_grad + return self_st diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/api.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/api.py new file mode 100644 index 0000000000000000000000000000000000000000..23a0d2d21f95349d3a2381072b7caa73ad3b6a9c --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/api.py @@ -0,0 +1,1304 @@ +# mypy: allow-untyped-defs +from __future__ import annotations # type: ignore[attr-defined] + +import copy +import operator +import threading +import warnings +import weakref +from dataclasses import dataclass +from functools import reduce +from typing import Callable, cast, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING +from typing_extensions import deprecated + +import torch +import torch.distributed as dist +import torch.distributed._shard.sharding_spec as shard_spec +from torch.distributed import distributed_c10d, rpc +from torch.distributed._shard._utils import DEPRECATE_MSG +from torch.distributed._shard.sharding_spec._internals import ( + check_tensor, + validate_non_overlapping_shards_metadata, +) +from torch.distributed._shard.sharding_spec.api import ( + _dispatch_custom_op, + _has_custom_op, +) +from torch.distributed.remote_device import _remote_device +from torch.utils import _pytree as pytree + +from .metadata import ShardedTensorMetadata, TensorProperties +from .reshard import reshard_local_shard, reshuffle_local_shard +from .shard import Shard +from .utils import ( + _flatten_tensor_size, + _parse_and_validate_remote_device, + _validate_output_tensor_for_gather, + build_global_metadata, + build_metadata_from_local_shards, +) + + +if TYPE_CHECKING: + from torch.distributed._shard.metadata import ShardMetadata + + +# Tracking for sharded tensor objects. +_sharded_tensor_lock = threading.Lock() +_sharded_tensor_current_id = 0 +_sharded_tensor_map: Dict[int, weakref.ReferenceType[ShardedTensor]] = {} + +# Default sharded ops +_SHARDED_OPS: Dict[Callable, Callable] = {} + +# Customized user ops +_CUSTOM_SHARDED_OPS: Dict[Callable, Callable] = {} + + +def _register_remote_shards( + sharded_tensor_id: int, rrefs: List[rpc.RRef[Shard]], rpc_rank: int +): + with _sharded_tensor_lock: + if sharded_tensor_id not in _sharded_tensor_map: + raise RuntimeError( + f"Could not find sharded_tensor_id: {sharded_tensor_id} in map: {_sharded_tensor_map.keys()}" + ) + + sharded_tensor = _sharded_tensor_map[sharded_tensor_id]() + if sharded_tensor is None: + raise RuntimeError("ShardedTensor weakref has been deallocated") + else: + sharded_tensor._register_remote_shards(rrefs, rpc_rank) + + +class ShardedTensorBase(torch.Tensor): + _sharding_spec: shard_spec.ShardingSpec + _metadata: ShardedTensorMetadata + _local_shards: List[Shard] + + def __new__(cls, sharding_spec: shard_spec.ShardingSpec, *size, **kwargs): + # Use __new__ to construct a wrapper tensor, for recording tensor + # properties and logging purposes. + torch._C._log_api_usage_once("torch.distributed._shard.sharded_tensor") + + # check sharding spec and build sharded tensor metadata + if not isinstance(sharding_spec, shard_spec.ShardingSpec): + raise ValueError(f"Expecting ShardingSpec but got: {type(sharding_spec)}") + + sizes = _flatten_tensor_size(size) + dtype = kwargs["dtype"] + layout = kwargs["layout"] + pin_memory = kwargs["pin_memory"] + requires_grad = kwargs["requires_grad"] + + if dtype is None: + dtype = torch.get_default_dtype() + + tensor_properties = TensorProperties( + dtype, layout, requires_grad, pin_memory=pin_memory + ) + sharded_tensor_metadata = sharding_spec.build_metadata( + sizes, tensor_properties=tensor_properties + ) + + r = torch.Tensor._make_wrapper_subclass( # type: ignore[attr-defined] + cls, + sizes, + dtype=dtype, + layout=layout, + pin_memory=pin_memory, + requires_grad=requires_grad, + ) + # set sharding spec + r._sharding_spec = sharding_spec + # set metadata + r._metadata = sharded_tensor_metadata + # set local shards + r._local_shards = [] + return r + + def metadata(self) -> ShardedTensorMetadata: + """ + Returns a :class:`ShardedTensorMetadata` object corresponding to the + metadata for the entire tensor. + """ + return self._metadata + + def local_shards(self) -> List[Shard]: + """ + Returns a list of :class:`Shard' corresponding to the + local shards for this rank. Returns an empty list if the current rank + does not host any shards for this Tensor. + """ + return self._local_shards + + @classmethod + def _init_from_local_shards_and_global_metadata( + cls, + local_shards: List[Shard], + sharded_tensor_metadata: ShardedTensorMetadata, + sharding_spec=None, + ) -> ShardedTensorBase: + """ + Initialize a ShardedTensorBase with local shards and a global + ShardedTensorMetadata built on each rank. + Warning: This API is experimental and subject to change. It does + not do cross rank validations, and fully rely on the user + for the correctness of sharded_tensor_metadata on each rank + """ + shards_metadata = sharded_tensor_metadata.shards_metadata + tensor_properties = sharded_tensor_metadata.tensor_properties + + if len(shards_metadata) == 0: + raise ValueError("shards_metadata must not be empty!") + + if tensor_properties.layout != torch.strided: + raise ValueError("Only torch.strided layout is currently supported") + + if sharding_spec is None: + spec = shard_spec._infer_sharding_spec_from_shards_metadata(shards_metadata) + else: + spec = sharding_spec + + sharded_tensor_base = ShardedTensorBase.__new__( + ShardedTensor, + spec, + sharded_tensor_metadata.size, + dtype=tensor_properties.dtype, + layout=tensor_properties.layout, + pin_memory=tensor_properties.pin_memory, + requires_grad=tensor_properties.requires_grad, + ) + + # check if shards_metadata have overlap shards + validate_non_overlapping_shards_metadata(shards_metadata) + + # check if the shards_metadata is compatible with overall size of the sharded tensor. + check_tensor(shards_metadata, list(sharded_tensor_metadata.size)) + + # done validation, add local_shards + sharded_tensor_base._local_shards = local_shards + return sharded_tensor_base + + @classmethod + def __torch_dispatch__(cls, func, types, args=(), kwargs=None): + raise RuntimeError( + f"A {cls.__name__} object is being used from c++ while calling {func.__module__}.{func.__name__} " + "but the there is no custom __torch_dispatch__ implementation for it." + ) + + +class ShardedTensor(ShardedTensorBase): + """ + ShardedTensor is an torch.Tensor subclass to represent Tensors that are sharded + across multiple devices and multiple processes. + + ShardedTensor is initialized in an SPMD like fashion where each rank + initializes the ShardedTensor. The ShardedTensor object on each rank + then only stores the local shard for the Tensor and provides global + metadata for all the shards. + + ShardedTensor doesn't provide any Tensor like operations but is a wrapper + providing the Tensor representing the local shard and the global metadata. + Using these, users can build their custom distributed._sharded computations + on top of this primitive. The local shards are all initialized using the + create_op specified by tensor_init_params.create_op, e.g., torch.ones, or + torch.empty + + Args: + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The specification + describing how to shard the Tensor. + size (int...): a sequence of integers defining the shape of the output + tensor. Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.contiguous_format``. + init_rrefs (bool, optional): Whether or not to initialize + :class:`torch.distributed.rpc.RRef`s pointing to remote shards. + Need to initialize the RPC Framework if specified as ``True``. + Default: ``False``. + + .. note:: ShardedTensor uses collectives to do various operations, i.e. it + uses all_gather to do cross rank validations. For NCCL-based process + groups, internal tensor representations of objects must be moved to the + GPU device before communication takes place. In this case, the device + used is given by ``torch.cuda.current_device()`` and it is the user's + responsibility to ensure that this is set so that each rank has an + individual GPU, via ``torch.cuda.set_device()`` + + """ + + def __new__(cls, sharding_spec: shard_spec.ShardingSpec, *size, **kwargs): + self = super().__new__(cls, sharding_spec, *size, **kwargs) + return self + + def __init__( + self, + sharding_spec: shard_spec.ShardingSpec, + *size, + dtype=None, + layout=torch.strided, + requires_grad=False, + pin_memory=False, + memory_format=torch.contiguous_format, + process_group=None, + init_rrefs=False, + ): + # prepare initialization, initialize fields like + # _process_group, _local_shards, etc. + self._prepare_init(process_group=process_group, init_rrefs=init_rrefs) + + if layout != torch.strided: + raise ValueError("Only torch.strided layout is currently supported") + + if memory_format != torch.contiguous_format: + raise ValueError( + "Only torch.contiguous_format memory_format is currently supported" + ) + + self._metadata.tensor_properties.memory_format = memory_format + + current_rank = dist.get_rank() # global rank + + for shard_metadata in self._metadata.shards_metadata: + rank, device = _parse_and_validate_remote_device( + self._process_group, shard_metadata.placement + ) + if rank == current_rank: + local_tensor = _create_tensor_from_params( + shard_metadata.shard_sizes, + local_device=device, + tensor_properties=self._metadata.tensor_properties, + ) + self._local_shards.append(Shard(local_tensor, shard_metadata)) + + # do post initialization (i.e. register sharded_tensor_id, initialize_rpc) + self._post_init() + + def _prepare_init(self, process_group=None, init_rrefs=False): + self._init_rrefs = init_rrefs + self._sharded_tensor_id = None + + self._process_group = self._normalize_pg(process_group) + self._remote_shards: Dict[int, List[rpc.RRef[Shard]]] = {} + + def _post_init(self): + # Initialize RPC if available. + if self._init_rrefs: + with _sharded_tensor_lock: + global _sharded_tensor_current_id, _sharded_tensor_map + self._sharded_tensor_id = _sharded_tensor_current_id + _sharded_tensor_map[self._sharded_tensor_id] = weakref.ref(self) + _sharded_tensor_current_id += 1 + + if not rpc._is_current_rpc_agent_set(): + raise RuntimeError( + "RPC Framework needs to be initialized using" + " torch.distributed.rpc.init_rpc if init_rrefs is set to True" + ) + self._init_rpc() + + def __del__(self): + # Clean up the global map. + with _sharded_tensor_lock: + global _sharded_tensor_current_id, _sharded_tensor_map + if ( + hasattr(self, "_sharded_tensor_id") + and self._sharded_tensor_id in _sharded_tensor_map + ): + _sharded_tensor_map.pop(self._sharded_tensor_id) # type: ignore[call-overload] + + def _init_rpc(self): + # Validate PG and RPC ranks match. + pg_rank = dist.get_rank() + rpc_rank = rpc.get_worker_info().id + if pg_rank != rpc_rank: + raise ValueError( + f"Default ProcessGroup and RPC ranks must be " + f"the same for ShardedTensor, found process group rank: " + f"{pg_rank} and RPC rank: {rpc_rank}" + ) + + self._remote_shards = {} + + # Gather all the sharded tensor ids. + worker_infos = rpc._get_current_rpc_agent().get_worker_infos() + rank_to_name = {} + name_to_rank = {} + + for worker_info in worker_infos: + rank_to_name[worker_info.id] = worker_info.name + name_to_rank[worker_info.name] = worker_info.id + + all_tensor_ids = rpc.api._all_gather(self._sharded_tensor_id) + + # Share the local shards to the entire world. + futs = [] + rpc_rank = rpc.get_worker_info().id + for rank in range(dist.get_world_size()): + # Skip self. + if rank == dist.get_rank(): + continue + + if len(self.local_shards()) != 0: + rrefs: List[rpc.RRef[Shard]] = [ + rpc.RRef(shard) for shard in self.local_shards() + ] + fut = rpc.rpc_async( + rank, + _register_remote_shards, + args=(all_tensor_ids[rank_to_name[rank]], rrefs, rpc_rank), + ) + futs.append(fut) + + torch.futures.wait_all(futs) + + # Barrier for all RPCs to finish on all ranks. + rpc.api._all_gather(None) + + def _get_preferred_device(self) -> torch.device: + """ + Return the preferred device to be used when creating tensors for collectives. + This method takes into account the associated process group + """ + if dist.get_backend(self._process_group) == dist.Backend.NCCL: + return torch.device(torch.cuda.current_device()) + return torch.device("cpu") + + def gather( # type: ignore[override] + self, + dst: int = 0, + out: Optional[torch.Tensor] = None, + enforce_dtype: bool = False, + dtype: Optional[torch.dtype] = None, + ) -> None: + """ + Creates a full :class:`Tensor` on rank ``dst`` by gathering all shards of the + sharded tensor. + + The API needs to be called on all ranks in SPMD fashion. All ranks should have + the same ``dst``. ``out`` should be a tensor of the same size as the overall + size of the sharded tensor on ``dst`` and ``None`` on all other ranks. + + Args: + dst(int): The rank where full tensor is constructed. + Default: 0 + out (:class `torch.Tensor`, optional): The output full tensor. + Must to be provided ONLY on ``dst`` rank. + Default: ``None`` + enforce_dtype (bool): Deprecated, please use dtype instead. Force the + gathered tensors to be the same type as input and output. + dtype (torch.dtype): Force the gathered tensors to be this dtype. + Default: ``None`` + """ + + def shard_size(shard_md): + return reduce(operator.mul, shard_md.shard_sizes) # type: ignore[attr-defined] + + if enforce_dtype: + warnings.warn( + "`enforce_dtype` is deprecated. Please use `dtype` instead.", + FutureWarning, + stacklevel=2, + ) + + rank = dist.get_rank(self._process_group) + full_size = self.metadata().size + _validate_output_tensor_for_gather(rank, dst, full_size, out) + + local_shards = self.local_shards() + world_size = dist.get_world_size(self._process_group) + rank_sizes = [0 for _ in range(world_size)] + max_rank_size = 0 + shard_placement: Dict[ShardMetadata, Tuple[int, int]] = {} + # collect sizes + for shard_md in self.metadata().shards_metadata: + shard_rank = cast(_remote_device, shard_md.placement).rank() + assert shard_rank is not None + + shard_placement[shard_md] = (shard_rank, rank_sizes[shard_rank]) + rank_sizes[shard_rank] += shard_size(shard_md) + max_rank_size = max(max_rank_size, rank_sizes[shard_rank]) + + gather_list: Optional[List[torch.Tensor]] + if rank == dst: + assert out is not None + if enforce_dtype: + # enforce_dtype is deprecated. Do it for backward compatibility. + dtype = out.dtype + # TODO make it as a view of out tensor + gather_list = [ + torch.empty((max_rank_size,), device=out.device, dtype=dtype) + for _ in range(world_size) + ] + else: + gather_list = None + + with torch.no_grad(): + if enforce_dtype and len(local_shards) > 0: + # enforce_dtype is deprecated. Do it for backward compatibility. + dtype = local_shards[0].tensor.dtype + data = torch.empty( + max_rank_size, device=self._get_preferred_device(), dtype=dtype + ) + + for shard in local_shards: + src = shard.tensor.flatten() + if src.nelement() == 0: + warnings.warn( + "Gathering a tensor with zero elements on rank " + str(rank) + ) + return + shard_offset = shard_placement[shard.metadata][1] + data[shard_offset : shard_offset + src.numel()].copy_(src) + + dist.gather( + tensor=data, + gather_list=gather_list, + dst=dst, + group=self._process_group, + ) + if rank != dst: + return + # In _validate_output_tensor_for_gather, we raise if out == None and rank == dst + out = cast(torch.Tensor, out) + assert gather_list is not None + + full_size = self.metadata().size + dims = len(full_size) + for shard_md in self.metadata().shards_metadata: + rank, rank_offset = shard_placement[shard_md] + tensor = gather_list[rank] + tensor = tensor[rank_offset : rank_offset + shard_size(shard_md)] + tensor = tensor.view(shard_md.shard_sizes) + + out_narrow_view = out + for dim in range(dims): + out_narrow_view = out_narrow_view.narrow( + dim, + shard_md.shard_offsets[dim], + shard_md.shard_sizes[dim], + ) + + out_narrow_view.copy_(tensor) + + def cpu( + self, memory_format=torch.preserve_format, process_group=None + ) -> ShardedTensor: + """ + Returns a copy of this object in CPU memory. + + If this ShardedTensor is already on CPU memory, then no copy is + performed and original object is returned. + + .. note:: When moving a ShardedTensor from GPU to CPU, the ShardedTensor might + need to be managed by a different type of ProcessGroup(i.e. ProcessGroupGloo), + it is the user's responsiblity to explicitly pass in a new process_group that + is compatible with CPU. + """ + # TODO: make this a __torch_function__ op once ShardedTensor becomes a + # torch.Tensor subclass, see https://github.com/pytorch/pytorch/issues/75402 + if ( + memory_format != torch.preserve_format + and memory_format != torch.contiguous_format + ): + raise RuntimeError( + "Only `torch.contiguous_format` or " + "`torch.preserve_format` is supported!" + ) + all_on_cpu = True + for meta in self.metadata().shards_metadata: + all_on_cpu &= meta.placement.device().type == "cpu" # type: ignore[union-attr] + + # if every shard is already on CPU, return the original object + if all_on_cpu: + return self + + # if not, returns a copy of this object on CPU + list_shards: List[Shard] = [] + # move all local shards to cpu, and change metadata + for shard in self._local_shards: + cpu_tensor = shard.tensor.cpu(memory_format=memory_format) # type: ignore[call-arg] + metadata = copy.deepcopy(shard.metadata) + metadata.placement._device = torch.device("cpu") # type: ignore[union-attr] + list_shards.append(Shard(cpu_tensor, metadata)) + + st_meta = copy.deepcopy(self.metadata()) + for meta in st_meta.shards_metadata: + if meta.placement.device().type != "cpu": # type: ignore[union-attr] + meta.placement._device = torch.device("cpu") # type: ignore[union-attr] + + pg = self._process_group if process_group is None else process_group + st_cpu = ShardedTensor._init_from_local_shards_and_global_metadata( + list_shards, + sharded_tensor_metadata=st_meta, + process_group=pg, + init_rrefs=self._init_rrefs, + ) + return st_cpu + + def cuda( + self, + device=None, + non_blocking=False, + memory_format=torch.preserve_format, + process_group=None, + ) -> ShardedTensor: + """ + Returns a copy of this object in CUDA memory, if the original ShardedTensor + is on CPU, we will move the local shard to the current GPU device of each + process in a SPMD fashion. + If this ShardedTensor is already on CUDA memory and local shards on each rank are + already on current device, we still returns a new ShardedTensor object with new + metadata, but no underlying data movements are performed. + .. note:: When moving a ShardedTensor from CPU to GPU, the ShardedTensor might + need to be managed by a different type of ProcessGroup(i.e. ProcessGroupNCCL), + it is the user's responsiblity to explicitly pass in a new process_group that + is compatible with GPU. + """ + if ( + memory_format != torch.preserve_format + and memory_format != torch.contiguous_format + ): + raise RuntimeError( + "Only `torch.contiguous_format` or " + "`torch.preserve_format` is supported!" + ) + + if device is not None: + device = torch.device(device) if isinstance(device, str) else device + assert ( + isinstance(device, torch.device) + and device.index == torch.cuda.current_device() + ), """Only device without device id (e.g. "cpu" or "cuda") is expected for ShardedTensor!""" + + current_device = torch.device(torch.cuda.current_device()) + # returns a copy of ShardedTensor on CUDA current device + list_shards: List[Shard] = [] + # move all local shards to current device, and change metadata + # if local shards already on the current device, there's no + # real data movement, only the metadata are copied. + for shard in self._local_shards: + cuda_tensor = shard.tensor.cuda( + device=current_device, + non_blocking=non_blocking, + memory_format=memory_format, + ) # type: ignore[call-arg] + metadata = copy.deepcopy(shard.metadata) + metadata.placement._device = current_device # type: ignore[union-attr] + + list_shards.append(Shard(cuda_tensor, metadata)) + + st_meta = copy.deepcopy(self.metadata()) + for meta in st_meta.shards_metadata: + if meta.placement.device().type != "cuda": # type: ignore[union-attr] + meta.placement._device = current_device # type: ignore[union-attr] + + pg = self._process_group if process_group is None else process_group + # we need to use `init_from_local_shards` to communicate between ranks + # and update the sharding spec/shards metadata. + st_cuda = ShardedTensor._init_from_local_shards_and_global_metadata( + list_shards, + sharded_tensor_metadata=st_meta, + process_group=pg, + init_rrefs=self._init_rrefs, + ) + return st_cuda + + def to(self, *args, **kwargs) -> ShardedTensor: + current_device: torch.device + if self._local_shards: + current_device = self._local_shards[0].tensor.device + elif self._process_group._get_backend_name() == "gloo": + current_device = torch.device("cpu") + else: + current_device = torch.device(torch.cuda.current_device()) + current_dtype = self.dtype + device_to = current_device + dtype_to = current_dtype + if len(args) == 1: + if isinstance(args[0], torch.dtype): + dtype_to = args[0] + elif isinstance(args[0], torch.device): + device_to = args[0] + elif isinstance(args[0], (str, int)): + device_to = torch.device(args[0]) + elif isinstance(args[0], torch.Tensor): + dtype_to = args[0].dtype + device_to = args[0].device + else: + raise RuntimeError(f"ShardedTensor.to() have wrong arguments: {args}") + elif len(args) == 2: + device_to, dtype_to = args + else: + dtype_to = kwargs.get("dtype", current_dtype) + device_to = kwargs.get("device", current_device) + + device_to = ( + torch.device(device_to) if isinstance(device_to, (str, int)) else device_to + ) + + if device_to.type == "cuda": + # if device_to set to cuda, set to current device even + # if user specify the device index. + current_idx = torch.cuda.current_device() + if device_to.index != current_idx: + warnings.warn( + "ShardedTensor.to only move tensor to its current device" + "If you want to put to different device, use `reshard` instead." + ) + device_to = torch.device(current_idx) + + copy_tensor = kwargs.get("copy", False) + non_blocking = kwargs.get("non_blocking", False) + memory_format = kwargs.get("memory_format", torch.preserve_format) + process_group = kwargs.get("process_group", None) + + if ( + not copy_tensor + and dtype_to == current_dtype + and device_to == current_device + ): + # already have correct dtype and device, return itself + return self + + # returns a copy of ShardedTensor on CUDA current device + list_shards: List[Shard] = [] + + for shard in self._local_shards: + new_tensor = shard.tensor.to( # type: ignore[call-overload] + device=device_to, + dtype=dtype_to, + non_blocking=non_blocking, + copy=copy_tensor, + memory_format=memory_format, + ) + metadata = copy.deepcopy(shard.metadata) + if metadata.placement is not None: + metadata.placement._device = device_to + list_shards.append(Shard(new_tensor, metadata)) + + # update metadata + st_meta = copy.deepcopy(self.metadata()) + st_meta.tensor_properties.dtype = dtype_to + for meta in st_meta.shards_metadata: + meta.placement._device = device_to # type: ignore[union-attr] + + pg = self._process_group if process_group is None else process_group + # we need to use `init_from_local_shards` to communicate between ranks + # and update the sharding spec/shards metadata. + st_to = ShardedTensor._init_from_local_shards_and_global_metadata( + list_shards, + sharded_tensor_metadata=st_meta, + process_group=pg, + init_rrefs=self._init_rrefs, + ) + return st_to + + @classmethod + def _normalize_pg( + cls, process_group: Optional[dist.ProcessGroup] + ) -> dist.ProcessGroup: + if process_group is not None: + return process_group + return distributed_c10d._get_default_group() + + @classmethod + def _init_from_local_shards( + cls, + local_shards: List[Shard], + *global_size, + process_group=None, + init_rrefs=False, + ): + # STEP 1: Validate the Shardmetadatas locally + process_group = cls._normalize_pg(process_group) + current_rank = dist.get_rank() # intentional to get global rank + world_size = dist.get_world_size(process_group) + + local_sharded_tensor_metadata: Optional[ShardedTensorMetadata] = None + global_tensor_size = _flatten_tensor_size(global_size) + + if len(local_shards) > 0: + local_sharded_tensor_metadata = build_metadata_from_local_shards( + local_shards, global_tensor_size, current_rank, process_group + ) + + # STEP 2. Validate metadata across ranks, and build a global sharded tensor + # metadata by gathering local ShardedTensorMetadata + gathered_metadatas: List[Optional[ShardedTensorMetadata]] = [] + if world_size > 1: + gathered_metadatas = [None for _ in range(world_size)] + + dist.all_gather_object( + gathered_metadatas, local_sharded_tensor_metadata, group=process_group + ) + else: + gathered_metadatas = [local_sharded_tensor_metadata] + + global_sharded_tensor_metadata = build_global_metadata(gathered_metadatas) + tensor_properties = global_sharded_tensor_metadata.tensor_properties + + # STEP 3: Validation done, create the actual ShardedTensor and populate fields + # prepare initialization + spec = shard_spec._infer_sharding_spec_from_shards_metadata( + global_sharded_tensor_metadata.shards_metadata + ) + sharded_tensor = cls.__new__( + cls, + spec, + global_sharded_tensor_metadata.size, + dtype=tensor_properties.dtype, + layout=tensor_properties.layout, + pin_memory=tensor_properties.pin_memory, + requires_grad=tensor_properties.requires_grad, + ) + sharded_tensor._prepare_init(process_group=process_group, init_rrefs=init_rrefs) + + # attach local_shards to the ShardedTensor created + sharded_tensor._local_shards = local_shards + + # run post initialization, i.e. map registration, rpc initialization + sharded_tensor._post_init() + return sharded_tensor + + @classmethod + @deprecated(DEPRECATE_MSG, category=FutureWarning) + def _init_from_local_tensor( + cls, + local_tensor: torch.Tensor, + sharding_spec: shard_spec.ShardingSpec, + *global_size: Sequence[int], + process_group: Optional[dist.ProcessGroup] = None, + init_rrefs=False, + ) -> ShardedTensor: + """ + Initialize a ShardedTensor given only one local tensor, global sharded tensor + size and sharding spec on each rank. + + Args: + local_tensor (Tensor): Single tensor of local shard stored in each rank. + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): + The specification describing how to shard the Tensor. + global_size (Sequence[int]): Size of the sharded tensor. + process_group (ProcessGroup, optional): The process group to aggregate on. + Default: None + init_rrefs (bool, optional): Whether or not to initialize + :class:`torch.distributed.rpc.RRef`s pointing to remote shards. + Need to initialize the RPC Framework if specified as ``True``. + Default: ``False``. + + Returns: + A :class:`ShardedTensor` sharded based on the given sharding_spec with local + tensor stored in the current rank. + + Examples: + >>> # xdoctest: +SKIP + >>> # All tensors below are of torch.int64 type. + >>> # We have 2 process groups, 2 ranks. + >>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank + >>> local_tensor = torch.unsqueeze(torch.cat([tensor, tensor + 2])) + >>> local_tensor + tensor([[1, 2, 3, 4]]) # Rank 0 + tensor([[3, 4, 5, 6]]) # Rank 1 + >>> sharding_dim = 0 + >>> sharding_spec = ChunkShardingSpec( + dim=sharding_dim, + placements=[ + "rank:0/cuda:0", + "rank:1/cuda:1", + ], + ) + >>> st = ShardedTensor._init_from_local_tensor(local_tensor, sharding_spec, [2, 4]) + >>> st + ShardedTensor( + ShardedTensorMetadata( + shards_metadata=[ + ShardMetadata(shard_offsets=[0, 0], shard_sizes=[1, 4], placement=rank:0/cuda:0), + ShardMetadata(shard_offsets=[1, 0], shard_sizes=[1, 4], placement=rank:1/cuda:1), + ], + size=torch.Size([2, 4]) + ) + >>> st.local_tensor() + tensor([1, 2, 3, 4]) # Rank 0 + tensor([3, 4, 5, 6]) # Rank 1 + + Warning: This API is experimental and subject to change. It lacks of a fully across + rank validations, and we only validate the local shard on the current rank. + We fully rely on the user to ensure local tensor is sharded based on the + sharding spec. + """ + if not local_tensor.is_contiguous(): + raise ValueError("local_tensor is not a contiguous Tensor.") + + global_tensor_size = _flatten_tensor_size(global_size) + tensor_properties = TensorProperties( + dtype=local_tensor.dtype, + layout=local_tensor.layout, + requires_grad=local_tensor.requires_grad, + memory_format=torch.contiguous_format, + pin_memory=local_tensor.is_pinned(), + ) + sharded_tensor_metadata = sharding_spec.build_metadata( + global_tensor_size, tensor_properties + ) + + process_group = cls._normalize_pg(process_group) + current_rank = dist.get_rank() # intentional to get global rank + + local_shards: List[Shard] = [] + for shard_metadata in sharded_tensor_metadata.shards_metadata: + rank, _device = _parse_and_validate_remote_device( + process_group, shard_metadata.placement + ) + if rank == current_rank: + local_shards.append(Shard(local_tensor, shard_metadata)) + + # TODO: figure out what the API should behave when some rank have no shard + # see https://github.com/pytorch/pytorch/issues/7313 + return ShardedTensor._init_from_local_shards_and_global_metadata( + local_shards, + sharded_tensor_metadata, + process_group=process_group, + init_rrefs=init_rrefs, + sharding_spec=sharding_spec, + ) + + @classmethod + def _init_from_local_shards_and_global_metadata( # type: ignore[override] + cls, + local_shards: List[Shard], + sharded_tensor_metadata: ShardedTensorMetadata, + process_group=None, + init_rrefs=False, + sharding_spec=None, + ) -> ShardedTensor: + """ + Initialize a ShardedTensor with local shards and a global + ShardedTensorMetadata built on each rank. + + Warning: This API is experimental and subject to change. It does + not do cross rank validations, and fully rely on the user + for the correctness of sharded_tensor_metadata on each rank + """ + process_group = cls._normalize_pg(process_group) + current_rank = dist.get_rank() # intentional to get global rank + + shards_metadata = sharded_tensor_metadata.shards_metadata + + local_shard_metadatas = [] + + # collect local shard metadatas from the global sharded_tensor_metadata + for shard_metadata in shards_metadata: # type: ignore[attr-defined] + rank, local_device = _parse_and_validate_remote_device( + process_group, shard_metadata.placement + ) + + if current_rank == rank: + local_shard_metadatas.append(shard_metadata) + + if len(local_shards) != len(local_shard_metadatas): + raise RuntimeError( + f"Number of local shards ({len(local_shards)}) does not match number of local " + f"shards metadata in sharded_tensor_metadata ({len(local_shard_metadatas)}) " + f"on rank ({current_rank}) " + ) + + shards_metadata = sharded_tensor_metadata.shards_metadata + tensor_properties = sharded_tensor_metadata.tensor_properties + + if len(shards_metadata) == 0: + raise ValueError("shards_metadata must not be empty!") + + if tensor_properties.layout != torch.strided: + raise ValueError("Only torch.strided layout is currently supported") + + if sharding_spec is None: + spec = shard_spec._infer_sharding_spec_from_shards_metadata(shards_metadata) + else: + spec = sharding_spec + + sharded_tensor = ShardedTensor.__new__( + ShardedTensor, + spec, + sharded_tensor_metadata.size, + dtype=tensor_properties.dtype, + layout=tensor_properties.layout, + pin_memory=tensor_properties.pin_memory, + requires_grad=tensor_properties.requires_grad, + ) + + def _raise_if_mismatch(expected, actual, prop_name, rank, is_property=False): + tensor_property_or_metadata = ( + "tensor property" if is_property else "local ShardMetadata" + ) + if expected != actual: + raise ValueError( + f"Local shards' tensor {prop_name} property is incompatible with " + f"{tensor_property_or_metadata} on rank {rank}: " + f"{tensor_property_or_metadata} {prop_name}={expected}, " + f"local shard tensor {prop_name}={actual}." + ) + + for shard in local_shards: + shard_meta = shard.metadata + local_shard_tensor = shard.tensor + placement = shard_meta.placement + assert placement is not None, "Must specify placement for `Shard`!" + rank = placement.rank() + local_device = placement.device() + + _raise_if_mismatch( + tensor_properties.layout, + local_shard_tensor.layout, + "layout", + rank, + True, + ) + if not local_shard_tensor.is_contiguous(): + raise ValueError( + "Only torch.contiguous_format memory_format is currently supported" + ) + + _raise_if_mismatch( + shard_meta.shard_sizes, + list(local_shard_tensor.size()), + "size", + rank, + ) + _raise_if_mismatch( + tensor_properties.pin_memory, + local_shard_tensor.is_pinned(), + "pin_memory", + rank, + True, + ) + _raise_if_mismatch(local_device, local_shard_tensor.device, "device", rank) + _raise_if_mismatch( + tensor_properties.dtype, + local_shard_tensor.dtype, + "dtype", + rank, + True, + ) + _raise_if_mismatch( + tensor_properties.requires_grad, + local_shard_tensor.requires_grad, + "requires_grad", + rank, + True, + ) + + # check if shards_metadata have overlap shards + validate_non_overlapping_shards_metadata(shards_metadata) + + # check if the shards_metadata is compatible with overall size of the sharded tensor. + check_tensor(shards_metadata, list(sharded_tensor_metadata.size)) + + # done validation, add local_shards + sharded_tensor._local_shards = local_shards + sharded_tensor._prepare_init(process_group=process_group, init_rrefs=init_rrefs) + + # run post initialization, i.e. map registration, rpc initialization + sharded_tensor._post_init() + return sharded_tensor + + def sharding_spec(self) -> shard_spec.ShardingSpec: + """ + Returns the ShardingSpec for the tensor. + """ + return self._sharding_spec + + @deprecated(DEPRECATE_MSG, category=FutureWarning) + def reshard(self, resharding_spec: shard_spec.ShardingSpec) -> ShardedTensor: + """ + Reshard a sharded tensor given the ``resharding_spec``. For now, we only support + single local shard. + + If ``resharding_spec`` is same as the original one, this becomes a no-op. + If only ``resharding_spec`` shares the same sharding dim with the original one, + we swap local shards directly. + For more generic cases, we merge different shards across different ranks and split + the local shards based on the ``resharding_spec`` via `all_to_all` collective API. + + Args: + resharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The + specification describing how the tensor is sharded. + + Returns: + A :class:`ShardedTensor` object whose local shards are resharded. + + Examples: + >>> # xdoctest: +SKIP + >>> # We have 2 process groups, 2 ranks. + >>> tensor = torch.arange(4, dtype=torch.int64) + 1 + 2 * rank + >>> tensor = torch.stack([tensor, tensor]) + >>> tensor + tensor([[1, 2, 3, 4], [1, 2, 3, 4]]) # Rank 0 + tensor([[3, 4, 5, 6], [3, 4, 5, 6]]) # Rank 1 + tensor([[5, 6, 7, 8], [5, 6, 7, 8]]) # Rank 2 + tensor([[7, 8, 9, 10], [7, 8, 9, 10]]) # Rank 3 + >>> sharding_dim = 0 + >>> spec = ChunkShardingSpec( + dim=sharding_dim, + placements=[ + "rank:0/cuda:0", + "rank:1/cuda:1", + "rank:2/cuda:2", + "rank:3/cuda:3", + ], + ) + >>> current_offsets = [0] * 2 + >>> current_offsets[0] = rank * 2 + >>> shard_metadata = ShardMetadata( + shard_offsets=copy.deepcopy(current_offsets), + shard_sizes=tensor.size(), + placement=spec.placements[rank], + ) + >>> local_shards = [ + Shard( + tensor=tensor, + metadata=shard_metadata, + ) + ] + >>> st = ShardedTensor._init_from_local_shards(local_shards, tensor.size()) + >>> sharding_dim = 1 + >>> resharding_spec = ChunkShardingSpec( + dim=sharding_dim, + placements=[ + "rank:0/cuda:0", + "rank:1/cuda:1", + "rank:2/cuda:2", + "rank:3/cuda:3", + ], + ) + >>> st.reshard(resharding_spec) + >>> tensor = st.local_shards()[0].tensor + >>> tensor + tensor([[1], [1], [3], [3], [5], [5], [7], [7]]) # Rank 0 + tensor([[2], [2], [4], [4], [6], [6], [8], [8]]) # Rank 1 + tensor([[3], [3], [5], [5], [7], [7], [9], [9]]) # Rank 2 + tensor([[4], [4], [6], [6], [8], [8], [10], [10]]) # Rank 3 + """ + if not isinstance( + resharding_spec, shard_spec.ChunkShardingSpec + ) or not isinstance(self._sharding_spec, shard_spec.ChunkShardingSpec): + raise NotImplementedError("Only ChunkShardingSpec supported for reshard.") + if len(self.local_shards()) != 1: + raise NotImplementedError("Only single local shard supported for reshard.") + + if self._sharding_spec.dim == resharding_spec.dim: # type: ignore[attr-defined] + if self._sharding_spec.placements == resharding_spec.placements: # type: ignore[attr-defined] + return self + else: + local_shards, shards_metadata = reshuffle_local_shard( + self.local_tensor(), + self.size(), # type: ignore[arg-type] + self._sharding_spec, + resharding_spec, + self._process_group, + ) + else: + local_shards, shards_metadata = reshard_local_shard( + self.local_tensor(), + self.size(), # type: ignore[arg-type] + self._sharding_spec, + resharding_spec, + self._process_group, + ) + self._local_shards = local_shards + self._metadata.shards_metadata = shards_metadata + self._sharding_spec = resharding_spec + return self + + def local_tensor(self) -> torch.Tensor: + """ + Return local tensor for a sharded_tensor. For now we only support single local shard. + + Returns: + A :class:`torch.Tensor` of the local shard. + """ + if len(self.local_shards()) != 1: + raise NotImplementedError("Only single local shard is supported.") + return self.local_shards()[0].tensor + + @classmethod + @deprecated(DEPRECATE_MSG, category=FutureWarning) + def __torch_function__(cls, func, types, args=(), kwargs=None): + def dispatch(st: ShardedTensor, func: Callable): + # Dispatch to custom user provided op first if it exists. + if func in _CUSTOM_SHARDED_OPS: + return _CUSTOM_SHARDED_OPS[func](types, args, kwargs, st._process_group) + + # Dispatch to custom sharding spec op if it has one. + if _has_custom_op(st._sharding_spec, func): + return _dispatch_custom_op( + st._sharding_spec, func, types, args, kwargs, st._process_group + ) + + if func in _SHARDED_OPS: + return _SHARDED_OPS[func](types, args, kwargs, st._process_group) + + raise RuntimeError( + f"torch function '{func.__name__}', with args: {args} and " + f"kwargs: {kwargs} not supported for ShardedTensor!" + ) + + # Find ShardedTensor instance to get process_group and sharding_spec. + st_instance = None + + def find_sharded_tensor(e): + nonlocal st_instance + if st_instance is None and isinstance(e, ShardedTensor): + st_instance = e + + pytree.tree_map_(find_sharded_tensor, args) + pytree.tree_map_(find_sharded_tensor, kwargs) + + if st_instance is not None: + return dispatch(st_instance, func) + + raise RuntimeError( + f"torch function '{func.__name__}', with args: {args} and " + f"kwargs: {kwargs} not supported for ShardedTensor!" + ) + + def is_pinned(self) -> bool: # type: ignore[override] + """ + Returns True if the sharded tensor (each local shard) resides in pinned memory. + """ + return self._metadata.tensor_properties.pin_memory + + def _register_remote_shards( + self, remote_shards: List[rpc.RRef[Shard]], rpc_rank: int + ): + self._remote_shards[rpc_rank] = remote_shards + + def remote_shards(self) -> Dict[int, List[rpc.RRef[Shard]]]: + """ + Returns a Dict[int, RRef] with keys being the RPC rank and values + being RRefs to shards on that rank. Need to initialize the + RPC framework for this functionality. + + Raises an exception if ShardedTensor was created with ``init_rrefs=False`` + """ + if not self._init_rrefs: + raise RuntimeError( + "ShardedTensor created with init_rrefs=False, no RRefs to remote shards available" + ) + return self._remote_shards + + def __hash__(self): + return id(self) + + def __repr__(self) -> str: # type: ignore[override] + return f"ShardedTensor({self._metadata})" + + @dataclass + class ProcessGroupState: + """ + State for ser-de of process group + """ + + local_rank: int + global_rank: int + local_world_size: int + global_world_size: int + + def __getstate__(self): + pg_state = ShardedTensor.ProcessGroupState( + distributed_c10d.get_rank(self._process_group), + distributed_c10d.get_rank(), + distributed_c10d.get_world_size(self._process_group), + distributed_c10d.get_world_size(), + ) + + return ( + self._local_shards, + self._metadata, + pg_state, + self._sharding_spec, + self._init_rrefs, + ) + + def __setstate__(self, state): + self._sharded_tensor_id = None + if not distributed_c10d.is_initialized(): + raise RuntimeError( + "Need to initialize default process group using " + '"init_process_group" before loading ShardedTensor' + ) + + ( + self._local_shards, + self._metadata, + pg_state, + self._sharding_spec, + self._init_rrefs, + ) = state + + # Setup process group + from torch.distributed._shard.api import _get_current_process_group + + self._process_group = _get_current_process_group() + + # Validate process group. + local_rank = distributed_c10d.get_rank(self._process_group) + if pg_state.local_rank != local_rank: + raise RuntimeError( + f"Local rank at save time was {pg_state.local_rank}, but at " + f"load time was {local_rank}" + ) + + global_rank = distributed_c10d.get_rank() + if pg_state.global_rank != global_rank: + raise RuntimeError( + f"Global rank at save time was {pg_state.global_rank}, but at " + f"load time was {global_rank}" + ) + + local_world_size = distributed_c10d.get_world_size(self._process_group) + if pg_state.local_world_size != local_world_size: + raise RuntimeError( + f"Local world size at save time was {pg_state.local_world_size}, " + f"but at load time was {local_world_size}" + ) + + global_world_size = distributed_c10d.get_world_size() + if pg_state.global_world_size != global_world_size: + raise RuntimeError( + f"Global world size at save time was {pg_state.global_world_size}, " + f"but at load time was {global_world_size}" + ) + + self._post_init() + + +def _create_tensor_from_params( + *size, local_device, tensor_properties: TensorProperties +): + """Helper to construct tensor from size, device and common params.""" + dtype = tensor_properties.dtype + layout = tensor_properties.layout + requires_grad = tensor_properties.requires_grad + memory_format = tensor_properties.memory_format + pin_memory = tensor_properties.pin_memory + + return torch.empty( + *size, + dtype=dtype, + layout=layout, + device=local_device, + requires_grad=requires_grad, + memory_format=memory_format, + pin_memory=pin_memory, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/logger.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..ebb749dc7d5c7315028d6edd4ee6552360b48594 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/logger.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 + +# Copyright (c) Facebook, Inc. and its affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from typing import List, Tuple + +from torch.distributed._shard.sharded_tensor.logging_handlers import _log_handlers + + +__all__: List[str] = [] + + +def _get_or_create_logger() -> logging.Logger: + logging_handler, log_handler_name = _get_logging_handler() + logger = logging.getLogger(f"sharding-spec-{log_handler_name}") + logger.setLevel(logging.DEBUG) + formatter = logging.Formatter( + "%(asctime)s %(filename)s:%(lineno)s %(levelname)s p:%(processName)s t:%(threadName)s: %(message)s" + ) + logging_handler.setFormatter(formatter) + logger.propagate = False + logger.addHandler(logging_handler) + return logger + + +def _get_logging_handler( + destination: str = "default", +) -> Tuple[logging.Handler, str]: + log_handler = _log_handlers[destination] + log_handler_name = type(log_handler).__name__ + return (log_handler, log_handler_name) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/logging_handlers.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/logging_handlers.py new file mode 100644 index 0000000000000000000000000000000000000000..021ad100f06a89fa944d7e6dba18f1a9558d88f5 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/logging_handlers.py @@ -0,0 +1,17 @@ +#!/usr/bin/env python3 + +# Copyright (c) Facebook, Inc. and its affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from typing import Dict, List + + +__all__: List[str] = [] + +_log_handlers: Dict[str, logging.Handler] = { + "default": logging.NullHandler(), +} diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/metadata.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..e53ac25fa55d9c2a73dce027cf8e0785b8cba930 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/metadata.py @@ -0,0 +1,95 @@ +# mypy: allow-untyped-defs +from dataclasses import dataclass, field +from enum import Enum +from typing import List + +import torch +from torch.distributed._shard.metadata import ShardMetadata + + +class MEM_FORMAT_ENCODING(Enum): + TORCH_CONTIGUOUS_FORMAT = 0 + TORCH_CHANNELS_LAST = 1 + TORCH_PRESERVE_FORMAT = 2 + + +@dataclass +class TensorProperties: + """Properties used to create :class:`Tensor`""" + + # Regular tensor fields + dtype: torch.dtype = field(default=torch.get_default_dtype()) + layout: torch.layout = field(default=torch.strided) + requires_grad: bool = False + memory_format: torch.memory_format = field(default=torch.contiguous_format) + pin_memory: bool = False + + def __getstate__(self): + # Since torch.memory_format cannot be pickled! + memory_format = self.memory_format + if memory_format == torch.contiguous_format: + mem_format_encoding = MEM_FORMAT_ENCODING.TORCH_CONTIGUOUS_FORMAT + elif memory_format == torch.channels_last: + mem_format_encoding = MEM_FORMAT_ENCODING.TORCH_CHANNELS_LAST + elif memory_format == torch.preserve_format: + mem_format_encoding = MEM_FORMAT_ENCODING.TORCH_PRESERVE_FORMAT + else: + raise RuntimeError(f"Invalid torch.memory_format: {memory_format}") + + return ( + self.dtype, + self.layout, + self.requires_grad, + mem_format_encoding, + self.pin_memory, + ) + + def __setstate__( + self, + state, + ): + ( + self.dtype, + self.layout, + self.requires_grad, + mem_format_encoding, + self.pin_memory, + ) = state + + if mem_format_encoding == MEM_FORMAT_ENCODING.TORCH_CONTIGUOUS_FORMAT: + memory_format = torch.contiguous_format + elif mem_format_encoding == MEM_FORMAT_ENCODING.TORCH_CHANNELS_LAST: + memory_format = torch.channels_last + elif mem_format_encoding == MEM_FORMAT_ENCODING.TORCH_PRESERVE_FORMAT: + memory_format = torch.preserve_format + else: + raise RuntimeError( + f"Invalid torch.memory_format encoding: {mem_format_encoding}" + ) + + self.memory_format = memory_format + + @staticmethod + def create_from_tensor(tensor: torch.Tensor) -> "TensorProperties": + return TensorProperties( + dtype=tensor.dtype, + layout=tensor.layout, + requires_grad=tensor.requires_grad, + memory_format=torch.contiguous_format, + pin_memory=tensor.is_pinned(), + ) + + +@dataclass +class ShardedTensorMetadata: + """ + Represents metadata for :class:`ShardedTensor` + """ + + # Metadata about each shard of the Tensor + shards_metadata: List[ShardMetadata] = field(default_factory=list) + + # Size of each dim of the overall Tensor. + size: torch.Size = field(default=torch.Size([])) + + tensor_properties: TensorProperties = field(default_factory=TensorProperties) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/reshard.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/reshard.py new file mode 100644 index 0000000000000000000000000000000000000000..07c09029498b52888f08a1379ef07dd4ba06c903 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/reshard.py @@ -0,0 +1,246 @@ +# mypy: allow-untyped-defs +import copy +from typing import List, Tuple + +import torch +import torch.distributed as dist +import torch.distributed._shard.sharding_spec as shard_spec +from torch._C._distributed_c10d import ProcessGroup +from torch.distributed._shard.metadata import ShardMetadata +from torch.distributed._shard.sharding_spec._internals import ( + get_chunked_dim_size, + get_split_size, +) +from torch.distributed.nn.functional import all_to_all, all_to_all_single + +from .shard import Shard + + +def get_idx_from_placements(placements, current_rank) -> int: + """ + Return the position of the current rank in the given placements. + + Args: + placements(List[Union[_remote_device, str]]): + Specifies the placement of each shard of the Tensor. The size of + the list represents the number of shards to be created. This could + be a list of + :class:`torch.distributed._remote_device`'s. This list + could also contain a string which represents remote + device as accepted by + :class:`torch.distributed._remote_device` + current_rank (int): number of current device. + + Returns: + A int which contains the position of current device in the placement list. + """ + for idx, placement in enumerate(placements): # type: ignore[attr-defined] + if current_rank == placement.rank(): # type: ignore[union-attr] + return idx + raise RuntimeError("current_rank not in the placement.") + + +def build_reshard_metadata( + st_size: torch.Size, + sharding_spec: shard_spec.ShardingSpec, + world_size: int, +) -> Tuple[List[ShardMetadata], List[int]]: + """ + Based the given sharding spec, we calculate the offset and local shard size. + We then build a ShardMetadata on top of the calculation result. + + Args: + st_size (torch.Size): The size of the sharded tensor. + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The + specification describing how the tensor is sharded. + world_size (int): number of ranks. + + Returns: + A Tuple of the followings: + A List[`ShardMetadata`] which contains the metadata for the shard, including + offsets, lengths and device placement. + A List[int] which contains the ranks in the order of placement. + """ + shard_dim = int(sharding_spec.dim) # type: ignore[attr-defined] + shards_metadata = [None] * world_size + ranks = [] + offsets = [0] * len(st_size) + split_size = get_split_size(st_size[shard_dim], world_size) + for idx, placement in enumerate(sharding_spec.placements): # type: ignore[attr-defined] + ranks.append(placement.rank()) + sharded_dim_size = get_chunked_dim_size(st_size[shard_dim], split_size, idx) + local_tensor_size = list(st_size) + local_tensor_size[shard_dim] = sharded_dim_size + shards_metadata[placement.rank()] = ShardMetadata( # type: ignore[call-overload] + shard_offsets=copy.deepcopy(offsets), + shard_sizes=local_tensor_size, + placement=placement, + ) + offsets[shard_dim] += sharded_dim_size + return shards_metadata, ranks # type: ignore[return-value] + + +def reshuffle_local_shard( + local_shard: torch.Tensor, + st_size: torch.Size, + sharding_spec: shard_spec.ShardingSpec, + resharding_spec: shard_spec.ShardingSpec, + pg: ProcessGroup, +) -> Tuple[List[Shard], List[ShardMetadata]]: + """ + Reshuffle the local shard directly when the reshard dim is same as the original + sharding dim. Logically we do this in two step: + 1. To collect all shards based on original sharding spec. + 2. Reshard the tensor based on the given resharding spec. + + In reality, we consolidate the two steps into one by sending the local tensor to + the new shard directly based on the resharding spec. + + Args: + local_shard (Tensor): Local tensor stored in the current rank. + st_size (torch.Size): The size of the sharded tensor. + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The + specification describing how the tensor is sharded originally. + resharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The + specification describing how the tensor will be resharded. + pg (ProcessGroup): The process group to aggregate on. + + Returns: + A Tuple of the followings: + A List[`Shard`] which contains the local tensor and its metadata. + A List[`ShardMetadata`] which contains the metadata for the shard, including + offsets, lengths and device placement. + """ + current_rank = dist.get_rank(pg) + world_size = dist.get_world_size(pg) + # Build shards_metadata first. + shards_metadata, ranks = build_reshard_metadata( + st_size, resharding_spec, world_size + ) + # Get input split size for all2all. + reshard_dim = int(resharding_spec.dim) # type: ignore[attr-defined] + split_size = get_split_size(st_size[reshard_dim], world_size) + input_split_sizes = [0] * world_size + idx = get_idx_from_placements(sharding_spec.placements, current_rank) # type: ignore[attr-defined] + new_rank = resharding_spec.placements[idx].rank() # type: ignore[union-attr, attr-defined] + input_split_sizes[new_rank] = local_shard.size(reshard_dim) + # Get output split size for all2all. + output_split_sizes = [0] * world_size + new_idx = ranks.index(current_rank) + sharded_dim_size = get_chunked_dim_size(st_size[reshard_dim], split_size, new_idx) + output_split_sizes[new_rank] = sharded_dim_size + # Get gathered_input for all2all. + local_shard = local_shard.transpose(0, reshard_dim).contiguous() + gathered_input_size = list(local_shard.size()) + gathered_input_size[0] = sharded_dim_size + gathered_input = torch.empty( + gathered_input_size, device=local_shard.device, dtype=local_shard.dtype + ) + # all2all. + local_shard = all_to_all_single( + gathered_input, + local_shard, + input_split_sizes=input_split_sizes, + output_split_sizes=output_split_sizes, + group=pg, + ) + local_tensor = local_shard.transpose(0, reshard_dim).contiguous() + local_shards = [Shard(local_tensor, shards_metadata[current_rank])] + return local_shards, shards_metadata + + +def reshard_local_shard( + local_tensor: torch.Tensor, + st_size: torch.Size, + sharding_spec: shard_spec.ShardingSpec, + resharding_spec: shard_spec.ShardingSpec, + pg: ProcessGroup, +) -> Tuple[List[Shard], List[ShardMetadata]]: + """ + Reshard a sharded tensor given the ``resharding_spec``. When the reshard dim is + different from the original sharding dim, we need to do two steps logically: + 1. To collect all shards based on original sharding spec. + 2. Reshard the tensor based on the given resharding spec. + + In reality, we consolidate the two steps into one by sending each rank the new + shard based on the resharding spec. + + Args: + local_tensor (Tensor): Local tensor stored in the current rank. + st_size (torch.Size): The size of the sharded tensor. + sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The + specification describing how the tensor is sharded originally. + resharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`): The + specification describing how the tensor will be resharded. + pg (ProcessGroup): The process group to aggregate on. + + Returns: + A Tuple of the followings: + A List[`Shard`] which contains the local tensor and its metadata. + A List[`ShardMetadata`] which contains the metadata for the shard, including + offsets, lengths and device placement. + """ + current_rank = dist.get_rank(pg) + world_size = dist.get_world_size(pg) + current_sharding_dim = int(sharding_spec.dim) # type: ignore[attr-defined] + reshard_dim = int(resharding_spec.dim) # type: ignore[attr-defined] + + # Build shards_metadata first. + shards_metadata, ranks = build_reshard_metadata( + st_size, resharding_spec, world_size + ) + + # Compute expected size + input_split_sizes = [ + metadata.shard_sizes[reshard_dim] for metadata in shards_metadata + ] + rearrange_input = any(ranks[i] > ranks[i + 1] for i in range(len(ranks) - 1)) + + if rearrange_input: + # Need to re-arrange reshard_dim of local_tensor before all2all. + indices: List[int] = [] + for metadata in shards_metadata: + offset_start_idx = metadata.shard_offsets[reshard_dim] + split_size = metadata.shard_sizes[reshard_dim] + indices += range(offset_start_idx, offset_start_idx + split_size) + local_tensor = local_tensor.index_select( + reshard_dim, torch.tensor(indices, device=local_tensor.device) + ) + + # Because reshard_dim != original shard_dim. We need to compute the + # size of tensor from each rank. + output_tensor_list = [torch.tensor(1)] * world_size + split_size = get_split_size(st_size[current_sharding_dim], world_size) + rearrange_output_list = False + indices = [] + for idx, placement in enumerate(sharding_spec.placements): # type: ignore[attr-defined] + sharded_dim_size = get_chunked_dim_size( + st_size[current_sharding_dim], split_size, idx + ) + output_tensor_size = list(st_size) + output_tensor_size[current_sharding_dim] = sharded_dim_size + output_tensor_size[reshard_dim] = input_split_sizes[current_rank] + output_tensor_list[ + placement.rank() + ] = torch.empty( # type: ignore[union-attr, index] + output_tensor_size, device=local_tensor.device, dtype=local_tensor.dtype + ) + indices.append(placement.rank()) # type: ignore[union-attr, index, arg-type] + if idx != placement.rank(): # type: ignore[union-attr] + rearrange_output_list = True + + # Perform autograd enabled all2all. + input_tensor_tuple = torch.split(local_tensor, input_split_sizes, dim=reshard_dim) + input_tensor_list = [tensor.contiguous() for tensor in input_tensor_tuple] + output_tensor_list = all_to_all( + output_tensor_list, + input_tensor_list, + group=pg, + ) + + if rearrange_output_list: + # Need to re-arrange original shard_dim of output_tensor_list. + output_tensor_list = [output_tensor_list[idx] for idx in indices] # type: ignore[call-overload] + local_tensor = torch.cat(output_tensor_list, dim=current_sharding_dim) + local_shards = [Shard(local_tensor, shards_metadata[current_rank])] + return local_shards, shards_metadata diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/shard.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/shard.py new file mode 100644 index 0000000000000000000000000000000000000000..dcb6b3b5d626748911c8aa8ce00605ce0a6d319c --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/shard.py @@ -0,0 +1,63 @@ +# mypy: allow-untyped-defs +from dataclasses import dataclass +from typing import List + +import torch +from torch.distributed._shard.metadata import ShardMetadata +from torch.distributed.remote_device import _remote_device + + +@dataclass +class Shard: + """ + Container which holds the data for a shard as a Tensor and also + the associated metadata for that shard. + + Args: + tensor(torch.Tensor): Local tensor for the shard. + metadata(:class `torch.distributed._shard.sharded_tensor.ShardMetadata`): + The metadata for the shard, including offsets, lengths and device placement. + """ + + __slots__ = ["tensor", "metadata"] + tensor: torch.Tensor + metadata: ShardMetadata + + def __post_init__(self): + # verification between local tensor and metadata + if list(self.tensor.size()) != self.metadata.shard_sizes: + raise ValueError( + "Shard tensor size does not match with metadata.shard_lengths! " + f"Found shard tensor size: {list(self.tensor.size())}, " + f"metadata.shard_lengths: {self.metadata.shard_sizes}, " + ) + placement_device = self.metadata.placement + if ( + placement_device is not None + and placement_device.device() != self.tensor.device + ): + raise ValueError( + f"Local shard tensor device does not match with local Shard's placement! " + f"Found local shard tensor device: {self.tensor.device}, " + f"local shard metadata placement device: {placement_device.device()}" + ) + + @classmethod + def from_tensor_and_offsets( + cls, tensor: torch.Tensor, shard_offsets: List[int], rank: int + ): + """ + Creates a Shard of a ShardedTensor from a local torch.Tensor, shard_offsets and rank. + + Args: + tensor(torch.Tensor): Local tensor for the shard. + shard_offsets(List[int]): List of integers specify the offset + of the shard on each dimension. + rank(int): Specify the rank for the shard. + """ + shard_sizes = list(tensor.size()) + placement = _remote_device(f"rank:{rank}/{str(tensor.device)}") + shard_meta = ShardMetadata( + shard_offsets=shard_offsets, shard_sizes=shard_sizes, placement=placement + ) + return Shard(tensor, shard_meta) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/utils.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a6954813f82b27a8850dde4dc5d64ef257f9a5bc --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharded_tensor/utils.py @@ -0,0 +1,267 @@ +# mypy: allow-untyped-defs +import collections.abc +import copy +from typing import List, Optional, Sequence, TYPE_CHECKING + +import torch +from torch.distributed import distributed_c10d as c10d, rpc +from torch.distributed._shard.sharding_spec._internals import ( + check_tensor, + validate_non_overlapping_shards_metadata, +) + +from .metadata import ShardedTensorMetadata, TensorProperties +from .shard import Shard + + +if TYPE_CHECKING: + from torch.distributed._shard.metadata import ShardMetadata + + +def _parse_and_validate_remote_device(pg, remote_device): + if remote_device is None: + raise ValueError("remote device is None") + + worker_name = remote_device.worker_name() + rank = remote_device.rank() + device = remote_device.device() + + # Validate rank, skip validation if rank is not part of process group. + if rank is not None and not c10d._rank_not_in_group(pg): + pg_global_ranks = c10d.get_process_group_ranks(pg) + if rank not in pg_global_ranks: + raise ValueError( + f"Global rank {rank} does not exist in input process group: {pg_global_ranks}" + ) + + if worker_name is not None: + if not rpc._is_current_rpc_agent_set(): + raise RuntimeError( + f"RPC framework needs to be initialized for using worker names: {worker_name}" + ) + + workers = rpc._get_current_rpc_agent().get_worker_infos() + for worker in workers: + if worker.name == worker_name: + return worker.id, device + + raise ValueError(f"Invalid worker name: {worker_name}") + + return rank, device + + +def _validate_output_tensor_for_gather( + my_rank: int, + dst_rank: int, + size: torch.Size, + dst_tensor: Optional[torch.Tensor], +) -> None: + if dst_rank == my_rank: + if dst_tensor is None: + raise ValueError( + f"Argument ``dst_tensor`` must be specified on destination rank {dst_rank}" + ) + if tuple(size) != (dst_tensor.size()): + raise ValueError( + f"Argument ``dst_tensor`` have size {tuple(dst_tensor.size())}," + f"but should be {tuple(size)}" + ) + elif dst_tensor: + raise ValueError( + "Argument ``dst_tensor`` must NOT be specified " "on non-destination ranks." + ) + + +def _flatten_tensor_size(size) -> torch.Size: + """ + Checks if tensor size is valid, then flatten/return a torch.Size object. + """ + if len(size) == 1 and isinstance(size[0], collections.abc.Sequence): + dims = list(*size) + else: + dims = list(size) + + for dim in dims: + if not isinstance(dim, int): + raise TypeError(f"size has to be a sequence of ints, found: {dims}") + + return torch.Size(dims) + + +def _raise_if_mismatch(expected, actual, prop_name, ranks, is_local=True): + if is_local: + assert isinstance(ranks, int) + if expected != actual: + raise ValueError( + f"Local shards' tensor {prop_name} property need to be the same on rank:{ranks}! " + f"Found one local shard tensor {prop_name}={expected}, " + f"the other local shard tensor {prop_name}={actual}." + ) + else: + # compare failure check across ranks, ranks list should have two rank + assert len(ranks) == 2 + if expected != actual: + raise ValueError( + f"ShardedTensor {prop_name} property does not match from different ranks! " + f"Found {prop_name}={expected} on rank:{ranks[0]}, " + f"and {prop_name}={actual} on rank:{ranks[1]}." + ) + + +def build_metadata_from_local_shards( + local_shards: List[Shard], + global_size: torch.Size, + current_rank: int, + pg: c10d.ProcessGroup, +) -> ShardedTensorMetadata: + assert len(local_shards) > 0, "must have local shards!" + local_shard_metadatas: List[ShardMetadata] = [] + + first_shard_dtype = local_shards[0].tensor.dtype + first_shard_layout = local_shards[0].tensor.layout + first_shard_requires_grad = local_shards[0].tensor.requires_grad + first_shard_is_pinned = local_shards[0].tensor.is_pinned() + + # 1). Validate local tensors and associated metadatas + for local_shard in local_shards: + local_shard_tensor = local_shard.tensor + local_shard_meta = local_shard.metadata + local_shard_metadatas.append(local_shard_meta) + rank, local_device = _parse_and_validate_remote_device( + pg, local_shard_meta.placement + ) + + if ( + local_shard_tensor.layout != torch.strided + or local_shard_tensor.layout != first_shard_layout + ): + raise ValueError( + f"Only torch.strided layout is currently supported, but found " + f"{local_shard_tensor.layout} on rank:{current_rank}!" + ) + + if not local_shard_tensor.is_contiguous(): + raise ValueError( + "Only torch.contiguous_format memory_format is currently supported!" + ) + + if rank != current_rank: + raise ValueError( + f"Local shard metadata's rank does not match with the rank in its process group! " + f"Found current rank in the process group: {current_rank}, " + f"local ShardMetadata placement's rank: {rank}" + ) + if local_shard_tensor.device != local_device: + raise ValueError( + f"Local shard tensor device does not match with local Shard's placement! " + f"Found local shard tensor device: {local_shard_tensor.device}, " + f"local shard metadata placement device: {local_device}" + ) + + _raise_if_mismatch( + local_shard_meta.shard_sizes, + list(local_shard_tensor.size()), + "size", + current_rank, + ) + _raise_if_mismatch( + local_shard_tensor.is_pinned(), + first_shard_is_pinned, + "pin_memory", + current_rank, + ) + _raise_if_mismatch( + local_shard_tensor.dtype, first_shard_dtype, "dtype", current_rank + ) + _raise_if_mismatch( + local_shard_tensor.requires_grad, + first_shard_requires_grad, + "requires_grad", + current_rank, + ) + + # 2). Build a "local" ShardedTensorMetadata with all local shards on this rank, then + # do all_gather to collect local_sharded_tensor_metadata from all ranks + local_tensor_properties = TensorProperties( + dtype=first_shard_dtype, + layout=first_shard_layout, + requires_grad=first_shard_requires_grad, + memory_format=torch.contiguous_format, + pin_memory=first_shard_is_pinned, + ) + + local_sharded_tensor_metadata = ShardedTensorMetadata( + shards_metadata=local_shard_metadatas, + size=global_size, + tensor_properties=local_tensor_properties, + ) + + return local_sharded_tensor_metadata + + +def build_global_metadata( + gathered_metadatas: Sequence[Optional[ShardedTensorMetadata]], +): + global_sharded_tensor_metadata = None + global_metadata_rank = 0 + + for rank, rank_metadata in enumerate(gathered_metadatas): + if rank_metadata is None: + continue + + if global_sharded_tensor_metadata is None: + global_sharded_tensor_metadata = copy.deepcopy(rank_metadata) + global_metadata_rank = rank + else: + _raise_if_mismatch( + global_sharded_tensor_metadata.size, + rank_metadata.size, + "global_size", + [global_metadata_rank, rank], + is_local=False, + ) + + # don't need to check layout and memory format as we already checked in local shards validation stage + _raise_if_mismatch( + global_sharded_tensor_metadata.tensor_properties.dtype, + rank_metadata.tensor_properties.dtype, + "dtype", + [global_metadata_rank, rank], + is_local=False, + ) + + _raise_if_mismatch( + global_sharded_tensor_metadata.tensor_properties.requires_grad, + rank_metadata.tensor_properties.requires_grad, + "requires_grad", + [global_metadata_rank, rank], + is_local=False, + ) + + _raise_if_mismatch( + global_sharded_tensor_metadata.tensor_properties.pin_memory, + rank_metadata.tensor_properties.pin_memory, + "pin_memory", + [global_metadata_rank, rank], + is_local=False, + ) + # pass all validations, extend shards metadata + global_sharded_tensor_metadata.shards_metadata.extend( + rank_metadata.shards_metadata + ) + + if global_sharded_tensor_metadata is not None: + # check if shards_metadata have overlap shards + validate_non_overlapping_shards_metadata( + global_sharded_tensor_metadata.shards_metadata + ) + + # check if the shards_metadata is compatible with global size of the sharded tensor. + check_tensor( + global_sharded_tensor_metadata.shards_metadata, + global_sharded_tensor_metadata.size, + ) + else: + raise ValueError("ShardedTensor have no local shards on all ranks!") + + return global_sharded_tensor_metadata diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharder.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharder.py new file mode 100644 index 0000000000000000000000000000000000000000..5d91ec15775bea870b81c4b10fb1443a3fba0977 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharder.py @@ -0,0 +1,29 @@ +import abc + +import torch.nn as nn + + +class Sharder(abc.ABC): + """ + This is an interface which allows user to create more advanced + sharding strategies that are not easily be composed by the + `ShardingSpec`. + + :class:`torch.distributed._shard.sharding_plan.ShardingPlan` could + take an object of the `Sharder` and call `shard` to shard the module, + then replace the original module with sharded module returned. + """ + + @abc.abstractmethod + def shard(self, module: nn.Module) -> nn.Module: + """ + Shard a module base on the implementation of this method, and + return the sharded version of the module. + + Args: + module (:class:`torch.nn.Module`): + The module to apply sharding to. + Returns: + A :class:`torch.nn.Module` object that represents a module + that's already been sharded. + """ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/__init__.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..325f7d7eb47b96a79fdc10cc2d1f072cdec9b4ce --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/__init__.py @@ -0,0 +1 @@ +from .api import ShardingPlan, ShardingPlanner diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc57037146c89f878baff438e2002970e8d92c88 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..756c332f2505e3e5e89f92da65cf0a5dd849565c Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/api.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/api.py new file mode 100644 index 0000000000000000000000000000000000000000..217ef8dab1e90fa7683ad2f5267cba6592985afe --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_plan/api.py @@ -0,0 +1,87 @@ +import abc +from dataclasses import dataclass +from typing import Dict, List, Optional, Union + +import torch.nn as nn +from torch.distributed._shard.sharder import Sharder +from torch.distributed._shard.sharding_spec import ShardingSpec + + +@dataclass +class ShardingPlan: + """ + Representation of a sharding plan, describes how to shard a module + across hosts. `plan` is used to shard module parameters according to the spec provided, + `output_plan` and `return_local_tensor` are optional, they are used to specify the output + layout of a module with a spec, and when to convert back to data parallel fashion. + + Args: + plan (Dict[str, Union[:class:`torch.distributed._shard.sharding_spec.ShardingSpec`, + :class:`torch.distributed._shard.sharder.Sharder`]): + a dict describes how to shard a module, there're currently two ways to shard a module: + 1. directly shard a module parameter by a `ShardingSpec`, keyed by the name of + a parameter to a `ShardingSpec`. + 2. shard a submodule by applying a `Sharder` on it, keyed by the name of a module + to a `Sharder` object. + output_plan (Dict[str, :class:`torch.distributed._shard.sharding_spec.ShardingSpec`), optional): + a dict specifies the layout of a module's output which produces a ShardedTensor, + keyed by the name of module to ShardingSpec("" in key means the root module). + Default: `None` + return_local_tensor (List[str], optional): a list of string, each element enables + a module's sharded output to be returned as a Tensor from its local shards to + ensure further processing in a data parallel fashion. ("" in list means the + root module). + Default: None + Example: + Suppose we want to shard a module with two linear layers and then run it with DDP, we also + want to convert the output of the second linear layer back to DDP, we can do it as follows: + + >>> # xdoctest: +REQUIRES(module:torch._C._distributed_c10d) + >>> class MyModule(nn.Module): + >>> def __init__(self) -> None: + >>> super().__init__() + >>> self.fc1 = nn.Linear() + >>> self.gelu = nn.GELU() + >>> self.fc2 = nn.Linear() + >>> self.relu = nn.Linear() + >>> + >>> def forward(self, input): + >>> return self.relu(self.fc2(self.gelu(self.fc1(input)))) + + + >>> # xdoctest: +SKIP("Undefined spec1, spec2) + >>> sharding_plan = ShardingPlan( + >>> plan={ + >>> "fc1.weight": spec1, + >>> "fc2.weight": spec2 + >>> }, + >>> output_plan={ + >>> "fc2": output_spec + >>> }, + >>> return_local_tensor=["fc2"] + >>> ) + """ + + plan: Dict[str, Union[ShardingSpec, Sharder]] + output_plan: Optional[Dict[str, ShardingSpec]] = None + return_local_tensor: Optional[List[str]] = None + + +class ShardingPlanner(abc.ABC): + """ + Default ShardingPlanner interface, can be extended and + implement advanced sharding strategies. + """ + + @abc.abstractmethod + def build_plan(self, module: nn.Module) -> ShardingPlan: + """ + Given a nn.Module, define how to shard the module across + ranks, return a ShardingPlan + Args: + module (:class:`torch.nn.Module`): + The module to apply sharding to. + Returns: + A :class:`torch.distributed._shard.sharding_plan.ShardingPlan` object that + represents how to shard the module. + """ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__init__.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bfd3f0a7581e8c4352eba843af6d3751bee7f387 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__init__.py @@ -0,0 +1,10 @@ +from torch.distributed._shard.metadata import ShardMetadata + +from .api import ( + _infer_sharding_spec_from_shards_metadata, + DevicePlacementSpec, + EnumerableShardingSpec, + PlacementSpec, + ShardingSpec, +) +from .chunk_sharding_spec import ChunkShardingSpec as ChunkShardingSpec diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b3a5ae67ca4dfdea06e378b055476df848a94c7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/_internals.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/_internals.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86913bf88254063262e30b3406e562cc9f62aef2 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/_internals.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c5d8726ba0df51d726736d202613a970640fcb6 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/chunk_sharding_spec.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/chunk_sharding_spec.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e42e8c773c3a171a04d510e9ffb878199a98f2f0 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/__pycache__/chunk_sharding_spec.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/_internals.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/_internals.py new file mode 100644 index 0000000000000000000000000000000000000000..8a439c447eff077dbb5bf0c0fff254062d346879 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/_internals.py @@ -0,0 +1,217 @@ +# mypy: allow-untyped-defs +from typing import List, Optional, Tuple + +from torch.distributed._shard.metadata import ShardMetadata + + +def _check_shard_metadata_pair_overlap(shard1: ShardMetadata, shard2: ShardMetadata): + """ + Checks if two shards overlap. + """ + + # For each dim of each shard, check if one shard resides on the other + # end of second shard with respect to that dim. As an example for a 2D + # shard, we would check if one shard is above or on the left of the + # other shard. + ndims = len(shard1.shard_offsets) + for i in range(ndims): + if shard1.shard_offsets[i] >= shard2.shard_offsets[i] + shard2.shard_sizes[i]: + return False + if shard2.shard_offsets[i] >= shard1.shard_offsets[i] + shard1.shard_sizes[i]: + return False + + return True + + +def _find_nd_overlapping_shards( + shards: List[ShardMetadata], sharded_dims: List[int] +) -> Optional[Tuple[int, int]]: + # Each rank has len(sharded_dims) tuples. Each tuple represent the + # [begin, end] (inclusive) pair of that dimension. + shard_intervals = [ + [ + (s.shard_offsets[dim], s.shard_offsets[dim] + s.shard_sizes[dim] - 1) + for dim in sharded_dims + ] + for s in shards + ] + + for i in range(len(shards)): + shard_i = shard_intervals[i] + for j in range(i + 1, len(shards)): + shard_j = shard_intervals[j] + # For each dim of each shard, check if one shard resides on the other + # end of second shard with respect to that dim. As an example for a 2D + # shard, we would check if one shard is above or on the left of the + # other shard. + overlap = True + for interval_i, interval_j in zip(shard_i, shard_j): + if interval_i[0] > interval_j[1] or interval_j[0] > interval_i[1]: + overlap = False + break + if overlap: + return (i, j) + return None + + +def _find_1d_overlapping_shards( + shards: List[ShardMetadata], dim: int +) -> Optional[Tuple[int, int]]: + # (begin, end, index_in_shards). Begin and end are inclusive. + intervals = [ + (s.shard_offsets[dim], s.shard_offsets[dim] + s.shard_sizes[dim] - 1, i) + for i, s in enumerate(shards) + ] + intervals.sort() + for i in range(len(shards) - 1): + if intervals[i][1] >= intervals[i + 1][0]: + return (intervals[i][2], intervals[i + 1][2]) + return None + + +def validate_non_overlapping_shards_metadata(shards: List[ShardMetadata]): + """ + Ensures none of the shards overlap with each other. + + Args: + shards(List[ShardMetadata]): List of :class:`ShardMetadata` objects representing + each shard. + Raises: + ``ValueError`` if there's overlap in any two shards. + """ + if not shards or len(shards) == 1: + return + + sharded_dims: List[int] = [] + for dim in range(len(shards[0].shard_offsets)): + for i in range(1, len(shards)): + if ( + shards[i].shard_offsets[dim] != shards[0].shard_offsets[dim] + or shards[i].shard_sizes[dim] != shards[0].shard_sizes[dim] + ): + sharded_dims.append(dim) + break + + pair: Optional[Tuple[int, int]] = None + if len(sharded_dims) == 0: + # All shards are the same, all dims are not partitioned. Choose any 2. + pair = (0, 1) + elif len(sharded_dims) == 1: + # Shards are partitioned over only one dimension. Overlap can be found + # using a O(nlogn) overlapping interval algorithm. + pair = _find_1d_overlapping_shards(shards, sharded_dims[0]) + else: + # Shards are partitioned over more than one dimension. Fall back to + # pair-wise check. Even though O(nlogn) algorithms (line sweep) exist + # for 2D overlap, the implementation is not trivial and may not justify + # the time saving in most cases. + pair = _find_nd_overlapping_shards(shards, sharded_dims) + + if pair: + raise ValueError(f"Shards {shards[pair[0]]} and {shards[pair[1]]} overlap") + + +def check_tensor(shards_metadata, tensor_dims) -> None: + """ + Checks if the shards_metadata is compatible with the provided tensor dims. + + Args: + shards_metadata(List[ShardMetadata]): List of :class:`ShardMetadata` + objects representing each shard of the tensor. + tensor_dims(Sequence of int): Dimensions of tensor to verify + Raises: + ``ValueError`` if not compatible. + """ + + # If the tensor's volume matches the total volume of all shards and + # all shard boundaries are within tensor dims, we have a compatible + # sharding spec for this tensor. Note that we have already verified + # we don't have overlapping shards. + tensor_rank = len(tensor_dims) + shards_rank = len(shards_metadata[0].shard_offsets) + if tensor_rank != shards_rank: + raise ValueError( + f"Rank of tensor is {tensor_rank}, but shards rank is {shards_rank}" + ) + + total_shard_volume = 0 + for shard in shards_metadata: + shard_volume = 1 + for i, shard_length in enumerate(shard.shard_sizes): + shard_volume *= shard_length + if shard.shard_offsets[i] + shard.shard_sizes[i] > tensor_dims[i]: + raise ValueError( + f"Shard offset {shard.shard_offsets[i]} and length " + f"{shard.shard_sizes[i]} exceeds tensor dim: {tensor_dims[i]} for shard {shard}" + ) + total_shard_volume += shard_volume + + tensor_volume = 1 + for size in tensor_dims: + tensor_volume *= size + + if total_shard_volume != tensor_volume: + # TODO: Can we improve this error message to point out the gaps? + raise ValueError( + f"Total volume of shards: {total_shard_volume} " + f"does not match tensor volume: {tensor_volume}, in other words " + f"all the individual shards do not cover the entire tensor" + ) + + +def get_split_size(dim_size, chunks): + """ + Computes the split size inline with ``torch.chunk`` + + Args: + dim_size(int): Size of the dimension being chunked. + chunks(int): Number of chunks to create for ``dim_size``. + + Returns: + An int indicating the split size to use. + """ + return (dim_size + chunks - 1) // chunks + + +def get_chunked_dim_size(dim_size, split_size, idx): + """ + Computes the dim size of the chunk for provided ``idx`` given ``dim_size`` + and ``split_size``. + + Args: + dim_size(int): Size of the dimension being chunked. + split_size(int): The chunk size for each chunk of ``dim_size``. + idx(int): The index of chunk whose dim size is being requested. + + Returns: + An int indicating the dim size of the chunk. + """ + return max(min(dim_size, split_size * (idx + 1)) - split_size * idx, 0) + + +def get_chunk_sharding_params(sharding_dim_size, world_size, spec, rank): + """ + Generate the start pos and offset length for the current rank for + chunk sharding. + + Args: + sharding_dim_size(int): The dimension length which we shard on. + world_size(int): number of ranks. + spec (:class:`torch.distributed._shard.sharding_spec.ChunkShardingSpec`): + sharding spec. + rank(int): # of cuda process. + + Returns: + start_pos(int): start position of sharded tensor on the given rank. + chunk_size(int): chunk size of sharded tensor on the given rank. + """ + split_size = get_split_size(sharding_dim_size, world_size) + current_offsets = 0 + start_pos = current_offsets + for idx, placement in enumerate(spec.placements): + chunk_size = get_chunked_dim_size(sharding_dim_size, split_size, idx) + if rank == placement.rank(): + start_pos = current_offsets + break + current_offsets += chunk_size + return start_pos, chunk_size # type: ignore[possibly-undefined] diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/api.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/api.py new file mode 100644 index 0000000000000000000000000000000000000000..91355e356c33eee8ad841408d6cc695a9e791d18 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/api.py @@ -0,0 +1,263 @@ +# mypy: allow-untyped-defs +import functools +import operator +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Callable, Dict, List, TYPE_CHECKING + +import torch +import torch.distributed._shard.sharded_tensor.metadata as sharded_tensor_meta +from torch.distributed._shard.metadata import ShardMetadata +from torch.distributed._shard.op_registry_utils import _decorator_func + +from ._internals import ( + check_tensor, + get_chunked_dim_size, + get_split_size, + validate_non_overlapping_shards_metadata, +) + + +if TYPE_CHECKING: + # Only include ShardedTensor when do type checking, exclude it + # from run-time to resolve circular dependency. + from torch.distributed._shard.sharded_tensor import ShardedTensor + + +class PlacementSpec(ABC): # noqa: B024 + """ + Base class representing the placement of an entity. Subclasses of this + class can be used to specify customized placements which might not be + covered by existing APIs. + """ + + +@dataclass +class DevicePlacementSpec(PlacementSpec): + """ + Associates placement of an entity with a single device. + + Args: + device(:class:`torch.distributed._remote_device`): The device to place the entity on. + """ + + device: torch.distributed._remote_device + + def __post_init__(self): + if not isinstance(self.device, torch.distributed._remote_device): + self.device = torch.distributed._remote_device(self.device) + + +class ShardingSpec(ABC): + """ + Base class representing sharding specifications. + """ + + @abstractmethod + def build_metadata( + self, + tensor_sizes: torch.Size, + tensor_properties: sharded_tensor_meta.TensorProperties, + ) -> sharded_tensor_meta.ShardedTensorMetadata: + """ + Given a global tensor size, define how to shard a tensor like this shape + across ranks, return ShardedTensorMetadata + Args: + tensor_sizes (:class:`torch.Size`): + The tensor shape to shard on, a `torch.Size` object that represents the + tensor shape to be sharded according to the ShardingSpec. + tensor_properties(:class:`torch.distributed._shard.sharded_tensor.TensorProperties): + Tensor properties used to create a ShardedTensor. + Returns: + A :class:`ShardedTensorMetadata` object that encodes the information about + the layout of the ShardedTensor and its properties. + """ + + @abstractmethod + def shard( + self, tensor: torch.Tensor, src_rank: int = 0, process_group=None + ) -> "ShardedTensor": + """ + Given a global tensor on src_rank, shard this tensor + across ranks within the process group, return a ShardedTensor. + Args: + tensor (:class:`torch.Tensor`): Tensor needs to be sharded. + Keyword args: + src_rank (int, optional): The source rank which is used as the ground truth of + the data for the parameter that would be sharded and scattered + across the rest of the ranks. + Default: 0. + process_group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + Returns: + A :class:`ShardedTensor` sharded from the given tensor. + """ + + +# Ops customized for a particular ShardingSpec. +_CUSTOM_SHARDING_SPEC_OPS: Dict[str, Dict[Callable, Callable]] = {} + + +def _has_custom_op(sharding_spec, op): + """ + Returns whether or not the ShardingSpec has a custom op implementation. + """ + class_name = type(sharding_spec).__qualname__ + return ( + class_name in _CUSTOM_SHARDING_SPEC_OPS + and op in _CUSTOM_SHARDING_SPEC_OPS[class_name] + ) + + +def _dispatch_custom_op( + sharding_spec, op: Callable, types, args, kwargs, process_group +): + """ + Calls the custom op for this ShardingSpec if it exists. + """ + class_name = type(sharding_spec).__qualname__ + if not _has_custom_op(sharding_spec, op): + raise RuntimeError(f"Custom op: {op} not registered for {class_name}") + func = _CUSTOM_SHARDING_SPEC_OPS[class_name][op] + return func(types, args, kwargs, process_group) + + +def custom_sharding_spec_op(sharding_spec_class, func): + """ + Decorator to allow custom registration of ops. + Args: + sharding_spec_class(type): The ShardingSpec for which we need to add this custom op. + func(Callable): The op to override (ex: torch.bmm) + """ + class_name = sharding_spec_class.__qualname__ + if class_name not in _CUSTOM_SHARDING_SPEC_OPS: + _CUSTOM_SHARDING_SPEC_OPS[class_name] = {} + return functools.partial( + _decorator_func, op=func, op_table=_CUSTOM_SHARDING_SPEC_OPS[class_name] + ) + + +@dataclass +class EnumerableShardingSpec(ShardingSpec): + """ + This is a type of PlacementSpec that allows users to specify a generic + sharding scheme by enumerating exactly how each shard is laid out. + + Args: + shards(List[ShardMetadata]): List of :class:`ShardMetadata` objects representing + each shard. Note that none of the shards should overlap. + """ + + shards: List[ShardMetadata] + + def __post_init__(self): + if len(self.shards) == 0: + raise ValueError(f"Empty shard list provided: {self.shards}") + + # Validate each shard has same rank. + rank = -1 + for shard in self.shards: + if rank != -1 and rank != len(shard.shard_offsets): + raise ValueError( + f"Found inconsistent ranks for shards: {rank} and {len(shard.shard_offsets)}" + ) + rank = len(shard.shard_offsets) + + validate_non_overlapping_shards_metadata(self.shards) + + def build_metadata( + self, + tensor_sizes: torch.Size, + tensor_properties: sharded_tensor_meta.TensorProperties, + ) -> sharded_tensor_meta.ShardedTensorMetadata: + # check if shards form a valid tensor + check_tensor(self.shards, tensor_sizes) + return sharded_tensor_meta.ShardedTensorMetadata( + self.shards, tensor_sizes, tensor_properties + ) + + def shard( + self, tensor: torch.Tensor, src_rank: int = 0, process_group=None + ) -> "ShardedTensor": + # TODO: figure out a generic and efficient way to scatter the shards for EnumerableShardingSpec + raise NotImplementedError("EnumerableShardingSpec.shard not implemented yet!") + + +def _infer_sharding_spec_from_shards_metadata(shards_metadata): + """ + Infer the sharding spec from the metadata of each shard of a ShardedTensor. + If the tensor is sharded only on one dimension, we can then verify whether it's + a ChunkShardingSpec or not. The way to verify it is to first get the total length + and perform a chunk sharding with the given placements to see if we can have the + same chunk size as the given shards_metadata. If not, we assume it's enum sharded. + + Args: + shards_metadata (List[ShardMetadata]): List of Metadata of local shards. + + Returns: + A :class:`torch.distributed._shard.sharding_spec.ShardingSpec` object of sharding + spec for one sharded tensor. + """ + placements = [] + chunk_sharding_dim = None + chunk_offset_list = [] + shard_size_list = [] + shard_offset_list = [] + # collect local shard metadatas from the global sharded_tensor_metadata + for shard_metadata in shards_metadata: # type: ignore[attr-defined] + placements.append(shard_metadata.placement) + local_offsets = shard_metadata.shard_offsets + chunk_offset_list.append(sum(local_offsets)) + shard_size_list.append(shard_metadata.shard_sizes) + shard_offset_list.append(shard_metadata.shard_offsets) + shard_dims = [idx for idx, e in enumerate(local_offsets) if e != 0] + # If the offset is [0, 0, ..., 0] (all zeros), + # we cannot decide whether how the tensor is sharded. + if len(shard_dims) == 0: + continue + # If the offset is [0, N, .,0, M, 0, .., 0], + # we are sure it's sharded by more than one dimension. + if len(shard_dims) != 1: + chunk_sharding_dim = None + break + # If the offset is [0, 0, .,0, M, 0, .., 0], aka, it's sharded by just + # one dimension, we need to make sure all ranks share the same dimension. + if not chunk_sharding_dim: + chunk_sharding_dim = shard_dims[0] + elif chunk_sharding_dim != shard_dims[0]: + chunk_sharding_dim = None + break + + if chunk_sharding_dim is not None: + # Ensure we infer the correct placement order from offsets + placements = [ + x + for _, x in sorted( + zip(chunk_offset_list, placements), key=operator.itemgetter(0) + ) + ] + + from .chunk_sharding_spec import ChunkShardingSpec + + chunk_spec = ChunkShardingSpec( + dim=chunk_sharding_dim, + placements=placements, + ) + + shard_sizes = sorted([x[chunk_sharding_dim] for x in shard_size_list]) + shard_total_length = sum(shard_sizes) + shard_offsets = sorted([x[chunk_sharding_dim] for x in shard_offset_list]) + + chunks = len(placements) + split_size = get_split_size(shard_total_length, chunks) + chunk_shard_sizes = sorted( + [ + get_chunked_dim_size(shard_total_length, split_size, idx) + for idx in range(chunks) + ] + ) + # Should match ChunkShardingSpec offsets calculation + chunk_shard_offsets = [split_size * idx for idx in range(chunks)] + if shard_sizes == chunk_shard_sizes and shard_offsets == chunk_shard_offsets: + return chunk_spec + return EnumerableShardingSpec(shards_metadata) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py new file mode 100644 index 0000000000000000000000000000000000000000..5101ed3b8da591750dd6070f2eed0d3058ae3f51 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py @@ -0,0 +1,226 @@ +# mypy: allow-untyped-defs +from dataclasses import dataclass +from typing import cast, List, Optional, TYPE_CHECKING, Union + +import torch +import torch.distributed as dist +import torch.distributed._shard.sharded_tensor.metadata as sharded_tensor_meta +import torch.distributed.distributed_c10d as distributed_c10d +from torch.distributed._shard._utils import narrow_tensor +from torch.distributed._shard.metadata import ShardMetadata +from torch.distributed._shard.sharded_tensor.shard import Shard +from torch.distributed._shard.sharded_tensor.utils import ( + _parse_and_validate_remote_device, +) + +from ._internals import get_chunked_dim_size, get_split_size +from .api import ShardingSpec + + +if TYPE_CHECKING: + # Only include ShardedTensor when do type checking, exclude it + # from run-time to resolve circular dependency. + from torch.distributed._shard.sharded_tensor import ShardedTensor + + +@dataclass +class ChunkShardingSpec(ShardingSpec): + """ + This is a type of PlacementSpec that defines the placement as being sharded + across multiple devices. In particular, it represents sharding a Tensor + along a single dimension into equal chunks (similar to :meth:`torch.chunk`). + + The semantics of how a tensor is partitioned is inline with + :meth:`torch.chunk`, where ``dim`` in torch.chunk corresponds to the + specified ``dim`` and ``chunks`` in torch.chunk is the number of elements + in the placement specified. + + Args: + dim (int or str): + The dimension to shard on, could be an integer representing the + dimension or a string in case of named tensors where dimensions are + named. Note that named tensor support is not added yet. + placement(List[Union[_remote_device, str]]): + Specifies the placement of each shard of the Tensor. The size of + the list represents the number of shards to be created. This could + be a list of + :class:`torch.distributed._remote_device`'s. This list + could also contain a string which represents remote + device as accepted by + :class:`torch.distributed._remote_device` + """ + + ShardingDim = Union[int, str] + + dim: ShardingDim + placements: List[Union[torch.distributed._remote_device, str]] + + def __post_init__(self): + self._verify_dim(self.dim) + for i, remote_device in enumerate(self.placements): + if not isinstance(remote_device, torch.distributed._remote_device): + self.placements[i] = torch.distributed._remote_device(remote_device) + + @staticmethod + def _verify_dim(dim): + # Validate the sharding spec. + # TODO: support named dimension + if isinstance(dim, str): + raise NotImplementedError( + "ChunkShardingSpec does not support named dimension yet!" + ) + + if not isinstance(dim, int): + raise ValueError(f"Sharding dim needs to be an integer, found: {dim}") + + def build_metadata( + self, + tensor_sizes: torch.Size, + tensor_properties: sharded_tensor_meta.TensorProperties, + ) -> sharded_tensor_meta.ShardedTensorMetadata: + tensor_num_dim = len(tensor_sizes) + + self._verify_dim(self.dim) + if self.dim >= tensor_num_dim or self.dim < -tensor_num_dim: # type: ignore[operator] + raise ValueError(f"Invalid sharding dim: {self.dim}") + + shards_metadata = [] + sharding_dim_size = tensor_sizes[self.dim] # type: ignore[index] + chunks = len(self.placements) + split_size = get_split_size(sharding_dim_size, chunks) + for idx, placement in enumerate(self.placements): + # generate ShardMetadata for each placement device + chunked_dim_size = get_chunked_dim_size(sharding_dim_size, split_size, idx) + shard_size = list(tensor_sizes) + current_offsets = [0] * tensor_num_dim + current_offsets[self.dim] = split_size * idx # type: ignore[index] + shard_size[self.dim] = chunked_dim_size # type: ignore[index] + + shard_metadata = ShardMetadata( + shard_offsets=current_offsets, + shard_sizes=shard_size, + placement=placement, + ) + shards_metadata.append(shard_metadata) + + return sharded_tensor_meta.ShardedTensorMetadata( + shards_metadata, tensor_sizes, tensor_properties + ) + + def shard( + self, tensor: torch.Tensor, src_rank: int = 0, process_group=None + ) -> "ShardedTensor": + """ + Args: + src_rank: group rank relative to ``process_group`` + + N.B. If ``process_group`` is None, ``src_rank`` is a global rank. + """ + # relative imports to avoid circular dependency + from torch.distributed._shard.sharded_tensor import ShardedTensor + + tensor_properties = sharded_tensor_meta.TensorProperties( + dtype=tensor.dtype, + layout=tensor.layout, + requires_grad=tensor.requires_grad, + memory_format=torch.contiguous_format, + pin_memory=tensor.is_pinned(), + ) + current_rank = dist.get_rank(process_group) + current_global_rank = dist.get_rank() + tensor_meta = self.build_metadata(tensor.size(), tensor_properties) + local_shards = [] + local_tensor = None + local_metadata = None + + tensors_to_scatter = cast( + List[Optional[torch.Tensor]], + [None] * dist.get_world_size(process_group), + ) + + sharding_dim_size = tensor.size()[self.dim] # type: ignore[index] + chunks = len(self.placements) + split_size = get_split_size(sharding_dim_size, chunks) + scatter_shape = list(tensor.size()) + scatter_shape[self.dim] = split_size # type: ignore[index] + + for shard_meta in tensor_meta.shards_metadata: + remote_global_rank, device = _parse_and_validate_remote_device( + process_group, shard_meta.placement + ) + if current_rank == src_rank: + # Reshape to get shard for this rank and we don't want autograd + # recording here for the narrow op and 'local_shard' should be a + # leaf variable in the autograd graph. + narrowed_tensor = narrow_tensor(tensor, shard_meta) + if shard_meta.shard_sizes[self.dim] < split_size: # type: ignore[index] + # for the last shard that might be smaller to other shards + # resize the narrowed tensor to the same size and use it for + # the scatter collective as dist.scatter requires same size + # inputs on every rank + tensor_to_scatter = ( + narrowed_tensor.detach().clone().resize_(scatter_shape) + ) + else: + tensor_to_scatter = narrowed_tensor.detach().clone().contiguous() + + tensors_to_scatter[ + dist.get_group_rank(process_group, remote_global_rank) + ] = tensor_to_scatter + + if current_global_rank == remote_global_rank: + local_tensor = torch.empty( + scatter_shape, + dtype=tensor.dtype, + layout=tensor.layout, + device=device, + ) + local_metadata = shard_meta + + # each rank should have local_tensor and local_metadata initialized if we build + # the metadata list in a correct way. + assert local_tensor is not None + assert local_metadata is not None + + # Scatter the shards to all ranks in the pg + # scatter takes the global rank as ``src`` + src_for_scatter = src_rank + if ( + process_group is not None + and process_group is not distributed_c10d._get_default_group() + ): + src_for_scatter = distributed_c10d.get_global_rank( + process_group, src_for_scatter + ) + + tensors_to_scatter_: Optional[List[torch.Tensor]] = None + if current_rank == src_rank: + tensors_to_scatter_ = [] + for t in tensors_to_scatter: + assert isinstance(t, torch.Tensor) + tensors_to_scatter_.append(t) + + dist.scatter( + local_tensor, + scatter_list=tensors_to_scatter_, + src=src_for_scatter, + group=process_group, + ) + + if list(local_tensor.size()) != local_metadata.shard_sizes: + # detach again after receiving to ensure local shards remain a leaf node + local_tensor = local_tensor.resize_(local_metadata.shard_sizes).detach() + + # Sync requires_grad to local_shard. + local_tensor.requires_grad = tensor.requires_grad + + local_shards.append(Shard(tensor=local_tensor, metadata=local_metadata)) + + st = ShardedTensor._init_from_local_shards_and_global_metadata( + local_shards, tensor_meta, process_group=process_group + ) + + # Manually set sharding_spec + st._sharding_spec = self + + return st diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__init__.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d3e6eb04197445545a1c374c7d0e519c6174836d Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/_common.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e3789a055665081eb08958c2b567b0abbcd0ce7f Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/_common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/embedding.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/embedding.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce10637f44210d7fa057002f396c1ef459d2267f Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/embedding.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/embedding_bag.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/embedding_bag.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d34da9a2068a5d86045bffcd48b40d3555a6da01 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/embedding_bag.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/_common.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..83d3371c7f904d3a6f0e72e0cf5c10ce22b6a8b4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/_common.py @@ -0,0 +1,348 @@ +# mypy: allow-untyped-defs + +import torch +import torch.distributed as dist +from torch.distributed._shard.sharded_tensor import ShardedTensor +from torch.distributed._shard.sharded_tensor._ops._common import _sharded_op_common +from torch.distributed._shard.sharding_spec import ChunkShardingSpec +from torch.distributed._shard.sharding_spec._internals import ( + get_chunk_sharding_params, + get_chunked_dim_size, + get_split_size, +) +from torch.distributed._shard.sharding_spec.api import custom_sharding_spec_op +from torch.distributed.nn.functional import ( + _all_gather_base, + all_reduce, + all_to_all_single, +) + + +def _chunk_sharding_spec_check(spec, op): + """ + For the given op implementation check if the sharding spec is ChunkShardingSpec. + """ + if not isinstance(spec, ChunkShardingSpec): + raise NotImplementedError( + f"Only ChunkShardingSpec supported for '{op.__name__}'." + ) + + +def _register_sharded_op_on_local_tensor( + op, early_stop_func=None, extra_check=None, customized_func=None +): + """ + Handles ``__torch_function__`` dispatch for ops which are performed on + the single local tensor of the sharded tensor such as op like + ``torch.nn.functional.softmax`` or ``torch.Tensor.view``. + + For more complicated ops, a customized func can be used to generate + the new local tensor, sharding spec and sharded tensor size. + + Args: + op: The op to be registered and applied to all shards of the st. + early_stop_func (Callable, optional): the func for early stop. + Default: if ``None``, no early stop. + extra_check (Callable, optional): the func for extra condition check. + Default: if ``None``, no extra check. + customized_func (Callable, optional): the func for customized logic + to generate the new local tensor, sharding spec and sharded tensor size. + Default: if ``None``, we simply lower to the real op call with + the single local tensor of the st. + + Return: + func (Callable): registered implementation for sharded op for + ``__torch_function__`` dispatch. + """ + + @custom_sharding_spec_op(ChunkShardingSpec, op) + @_sharded_op_common(op, early_stop_func, extra_check) + def sharded_tensor_op_on_local_tensor(types, args=(), kwargs=None, pg=None): + st = args[0] + sharding_spec = st.sharding_spec() + if len(st.local_shards()) != 1: + raise TypeError( + f"torch function '{op.__name__}', with args: {args} and " + f"kwargs: {kwargs} only supported for single local tensor!" + ) + st_size = st.size() + if customized_func: + local_tensor, sharding_spec, st_size = customized_func(args, kwargs, pg) + else: + args = (st.local_tensor(), *args[1:]) + local_tensor = op(*args, **kwargs) + return ShardedTensor._init_from_local_tensor( + local_tensor.contiguous(), + sharding_spec, + st_size, # type: ignore[arg-type] + process_group=pg, + init_rrefs=st._init_rrefs, + ) + + +def _handle_col_wise_sharding_base( + op_func, + col_dim, + input, + world_size, + weight, + local_shard, + pg, + gathered_inputs, + mode=None, + gathered_per_sample_weights=None, + gathered_offsets=None, + padding_idx=None, +): + """ + For col-wise sharding of weight, lots of logic are common. + So we extract the common logic and put in this function: + Step 1. To get input from each rank and + Step 2. To perform the op on the concatenated tensor. + Step 3. To distribute results to each rank with col rearrangement. + Step 4. To concatenate all results from all ranks. + + Args: + op_func: operator which is applied to the input tensor. + col_dim: dim of result tensor after the operation. + input: tensor to be applied op on. + world_size: number of ranks. + weight: sharded weight tensor. + local_shard: col-wise sharded weight tensor. + pg: process group. + gathered_inputs: list of inputs from all ranks. If specified, we + don't need to communicate with each rank any more. + mode: aggregation mode of EmbeddingBag. + gathered_per_sample_weights: per_sample_weights across all ranks. + gathered_offsets: offsets across all ranks. + padding_idx: If specified, the entries at padding_idx do + not contribute to the gradient; therefore, the embedding + vector at padding_idx is not updated during training, + i.e. it remains as a fixed "pad". + Note that the embedding vector at padding_idx is + excluded from the reduction. + + Return: final result of input being applied with the op. + """ + # run the operator's function for all the inputs. + results = [] + for i, inp in enumerate(gathered_inputs): + if op_func == torch.nn.functional.embedding_bag: + result = op_func( + inp, + local_shard, + offsets=gathered_offsets[i] if gathered_offsets is not None else None, + mode=mode, + per_sample_weights=gathered_per_sample_weights[i] + if gathered_per_sample_weights is not None + else None, + padding_idx=padding_idx, + ) + elif op_func == torch.nn.functional.embedding: + result = op_func( + inp, + local_shard, + padding_idx=padding_idx, + ) + else: + result = op_func(inp, local_shard) + results.append(torch.transpose(result, 0, col_dim)) + + # Distribute results to each rank with col rearrangement. + output = _result_distribute_with_col_rearrange( + results, input, world_size, weight, pg + ) + + # transpose the output and return result. + return torch.transpose(output, 0, col_dim) + + +def _result_distribute_with_col_rearrange(results, input, world_size, weight, pg): + """ + For col-wise sharding of weight, we need to distribute + results to each rank. We do them in this function. + Note that, if the index in the Sharding Spec is not equal to + the rank number, we need to do the rearrangement based on the + order given by the Sharding Spec (placement). + + Args: + results: results from ops applied to inputs from all ranks. + We need to distribute them back to their original ranks. + input: tensor to be applied op to. + world_size: number of ranks. + weight: sharded weight tensor. + pg: process group. + + Return: column rearranged result. + """ + # Process results and outputs for all2all. + sharding_dim = weight._sharding_spec.dim + sharding_dim_size = weight.size(sharding_dim) + dims = list(results[0].size()) + dims[0] = sharding_dim_size + combined_results = torch.cat(results) + output = torch.empty( + *dims, device=combined_results.device, dtype=combined_results.dtype + ) + + # Compute output splits + split_size = get_split_size(sharding_dim_size, world_size) + output_split_sizes = [0] * world_size + for idx, placement in enumerate(weight._sharding_spec.placements): + output_split_sizes[placement.rank()] = get_chunked_dim_size( + sharding_dim_size, split_size, idx + ) + + # distribute the outputs using all2all. + output = all_to_all_single( + output, combined_results, output_split_sizes=output_split_sizes, group=pg + ) + + # Check if we need to rearrange columns appropriately for output. + rearrange_columns = any( + idx != placement.rank() + for idx, placement in enumerate(weight._sharding_spec.placements) + ) + if not rearrange_columns: + return output + + indices = [] + for placement in weight._sharding_spec.placements: + dim_size = output_split_sizes[placement.rank()] + start = sum( + split_size if i < placement.rank() else 0 + for i, split_size in enumerate(output_split_sizes) + ) + indices += list(range(start, start + dim_size)) + + return output.index_select(0, torch.tensor(indices, device=output.device)) + + +def _handle_max_norm_col_wise( + max_norm, + norm_type, + local_shard, + input, + world_size, + gathered_inputs, + pg, +): + """ + For col-wise sharding of weight, we need to aggregate the + norm across all ranks before we can perform the proper re-norm. + Note that, the max_norm logic is only applied to the embedding + indices that are looked up and not the whole shard. + + Args: + max_norm: If given, each embedding vector with norm larger + than max_norm is renormalized to have norm max_norm. + Note: this will modify weight in-place. + norm_type: The p in the p-norm to compute for the max_norm option. + local_shard: col-wise shared local weight used for lookup. + input: tensor to be applied op to. + world_size: number of ranks. + gathered_inputs: list of inputs from all ranks. + pg: process group. + + Return: + local_shard_norm_renormed: local_shard re-normed to max_norm if the norm is larger + than it. + + """ + norm_type = norm_type if norm_type is not None else 2.0 + unique_inp = torch.unique(torch.cat(gathered_inputs)) + local_shard_sum = torch.sum( + torch.pow(torch.abs(local_shard), norm_type), dim=1, dtype=local_shard.dtype + ) + # For col-wise sharding, we need to first aggregate the powered sum + # from each rank first and then calculate the norm. + local_shard_sum = all_reduce(local_shard_sum, group=pg) + local_shard_norm = torch.pow(local_shard_sum, 1.0 / norm_type) + max_norm_tensor = torch.full( + (local_shard.size(0),), + float("inf"), + dtype=local_shard.dtype, + device=input.device, + ) + max_norm_tensor[unique_inp] = max_norm + local_shard_t = local_shard.t().contiguous() + normalized_tensor = torch.where( + local_shard_norm > max_norm_tensor, max_norm_tensor, local_shard_norm + ) + # Make sure divisor is not zero. + local_shard_norm[local_shard_norm == 0.0] = 1.0 + local_shard_norm_renormed = ( + torch.div(torch.mul(local_shard_t, normalized_tensor), local_shard_norm) + .t() + .contiguous() + ) + return local_shard_norm_renormed + + +def _all_gather_base_input(input, pg): + """ + Use _all_gather_base to get a concatenated input from each rank. + + Args: + input: tensor to be applied op on. + pg: process group. + + Returns: + gathered_inputs: input gathered from each rank and concat by dim 0. + """ + # allgather the inputs first. + gather_inp_size = list(input.size()) + gather_inp_size[0] = input.size(0) * dist.get_world_size(pg) + gather_inp = torch.empty(gather_inp_size, device=input.device, dtype=input.dtype) + return _all_gather_base(gather_inp, input, group=pg) + + +def _handle_row_wise_mask(gather_inp, padding_idx, weight, world_size, rank): + """ + Mask the input for embedding look-up for IDs which are not stored + on the current rank. This function also adjust the ``padding_idx`` + so that it is only used on the rank where the corresponding row is + stored. + + Note that, with ``max_norm`` flag on, only weights of rows being + looked up will be re-normed. So we need an extra row for masked ID + so that it does not affect the final result and ``max_norm``. + + Args: + gather_inp: tensor to be applied op on gathered from all ranks. + padding_idx: If specified, the entries at padding_idx do + not contribute to the gradient; therefore, the embedding + vector at padding_idx is not updated during training, + i.e. it remains as a fixed "pad". + Note that the embedding vector at padding_idx is + excluded from the reduction. + weight: weight tensor of Embedding look-up table. + world_size: number of ranks. + rank: # of cuda process. + + Returns: + lookup_input: Tensor of masked input. + padding_idx: adjusted padding_idx. + padding_row: The extra row we used during lookup so that + looking up does not affect ``max_norm``. + """ + (start_pos, chunk_size) = get_chunk_sharding_params( + weight.size(0), world_size, weight._sharding_spec, rank + ) + mask = (gather_inp < start_pos) | (gather_inp >= start_pos + chunk_size) + lookup_input = gather_inp.clone() - start_pos + lookup_input[mask] = chunk_size + if ( + padding_idx is not None + and padding_idx >= start_pos + and padding_idx < (start_pos + chunk_size) + ): + padding_idx = padding_idx - start_pos + else: + padding_idx = None + + # When max_norm is set, it will only re-norm the row being looked up. + padding_row = torch.zeros( + 1, weight.size(1), device=gather_inp.device, dtype=weight.dtype + ) + return lookup_input, padding_idx, padding_row diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..117aed79520d9ad78c10bdd2310fb6b032c2a024 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding.py @@ -0,0 +1,294 @@ +# mypy: allow-untyped-defs + +import torch +import torch.distributed as dist +from torch.distributed._shard.sharded_tensor import ShardedTensor +from torch.distributed._shard.sharding_spec import ChunkShardingSpec +from torch.distributed._shard.sharding_spec.api import custom_sharding_spec_op +from torch.distributed.nn.functional import all_gather, reduce_scatter + +from ._common import ( + _all_gather_base_input, + _handle_col_wise_sharding_base, + _handle_max_norm_col_wise, + _handle_row_wise_mask, +) + + +@custom_sharding_spec_op(ChunkShardingSpec, torch.nn.functional.embedding) +def sharded_embedding(types, args, kwargs, pg): + """ + Handles ``__torch_function__`` dispatch for ``torch.nn.functional.embedding``. + This method computes a sharded embedding lookup and has the following limitations: + + 1. Supports only sharding of ``weight``. + 2. Supports only ``ChunkShardingSpec``. + 3. Supports only a single local shard per rank. + 4. Supports all specs except for scale_grad_by_freq, sparse, etc. + + Based on the dimension that the weight is sharded on, there are two + algorithms: + + ROWWISE SHARDING + ================ + For row-wise sharding the weight is sharded on dimension 0. + + The overall algorithm can be best explained with an example. Let's assume + the dims for input are (4 x 6) and W are (10 x 17) and W is sharded across + 4 GPUs creating 3 shard of (3 x 17) and 1 shard of (1 x 17). + The algorithm is as follows: + + 1. First the input is all gathered to all ranks, since this is SPMD and + input is actually sharded across all ranks. The inputs then become a + 4 (4 x 6) tensor on each rank. For example if the given input is + tensor([[6, 5, 2, 9, 6, 3], + [3, 1, 2, 4, 7, 6], + [4, 0, 4, 9, 8, 9], + [8, 6, 6, 4, 6, 1]]) + on rank 0. + Then on every rank, we will have this tensor. + If input itself is already replicated, no all-gather will be done. + 2. Next, we mask the ID which are not stored on that rank. + For example on rank 0, we store ID [0, 1, 2]. We only keep the ID + inside the set of numbers. The rest of them will be masked to an extra row. + The masked matrix will be used for embedding look up and is like: + tensor([[4, 4, 2, 4, 4, 4], + [4, 1, 2, 4, 4, 4], + [4, 0, 4, 4, 4, 4], + [4, 4, 4, 4, 4, 1]]) + The reason of having an extra row (aka, number 4 in the example) is + because when max_norm is specified only weight which has looked will + be re-normed so mask IDs whose embeddings are not stored in current + rank will to an extra row will ensure max_norm still works as expected. + 3. If max_norm is specified, the extra row guarantees that the mask ID will + not affect the behavior of weigh re-norm. + + COLWISE SHARDING + ================ + For col-wise sharding the weight is sharded on dimension 1. + + The overall algorithm can be best explained with an example. Let's assume + the dims for input are (4 x 6) and W are (16 x 17) and W is sharded across + 4 GPUs creating 3 shards of (16 x 5) and 1 shard of (16 x 2). + The algorithm is as follows: + + 1. First the input is broadcasted to all ranks, since this is SPMD we + actually do an all_gather for all the inputs resulting in 4 (4 x 6) + inputs on each rank. + 2. Next we perform local embedding lookup operation by apply each + input (4 x 6) with the local shard (16 x 5) ((16 x 2) for the last). + This results in 4 (5 x 6 x 4) ((2 x 6 x 4) for the last) matrices + on each rank. We transpose dim 0 and dim 2. + 3. Next, we concat these 4 matrices and perform an all2all to share the + appropriate (5 x 6 x 4) or (2 x 6 x 4) matrices to each rank. + 4. Now, each rank receives a (17 x 6 x 4) matrix which is basically the + size of the result we need. + 5. If placements are not in order any appropriate rearrangement of columns + are done for the (17 x 6 x 4) matrix and finally we transpose the + dim 0 and dim 2 again. + 6. If max_norm is specified, we manually sum up the norm and renorm. Because + the renorm must be in place, we need to override the local_shard to mimic + this behavior. + """ + # Validate input params + _validate_embedding_param(args, kwargs) + + input = args[0] + weight = args[1] + max_norm = kwargs.get("max_norm") + norm_type = kwargs.get("norm_type") + padding_idx = kwargs.get("padding_idx") + + local_shard = weight.local_tensor().contiguous() + sharding_dim = weight._sharding_spec.dim + world_size = dist.get_world_size(pg) + rank = dist.get_rank(pg) + + if sharding_dim == 1: + output, local_shard = _handle_col_wise_sharding( + input, world_size, weight, local_shard, max_norm, norm_type, padding_idx, pg + ) + weight.local_shards()[0].tensor = local_shard + return output + elif sharding_dim == 0: + return _handle_row_wise_sharding( + input, + world_size, + weight, + local_shard, + max_norm, + norm_type, + padding_idx, + rank, + pg, + ) + else: + raise RuntimeError( + f"nn.Embedding weight sharded on dim {sharding_dim} not supported!" + ) + + +def _validate_embedding_param(args, kwargs): + """ + Validate input params of sharded embedding op. + + Args: + input: list of ID used for lookup. + weight: sharded weight tensor. + kwargs: same as normal Embedding. + + Return: None. + """ + + input = args[0] + weight = args[1] + max_norm = kwargs.get("max_norm") + scale_grad_by_freq = kwargs.get("scale_grad_by_freq") + sparse = kwargs.get("sparse") + + # Validate types + if not isinstance(input, torch.Tensor): + raise TypeError("input need to be torch.Tensor") + if not isinstance(weight, ShardedTensor): + raise TypeError("weight needs to be ShardedTensor") + weight_size = weight.size() + if len(weight_size) != 2: + raise ValueError("Weight needs to have exactly 2 dims") + if int(torch.min(input).item()) < 0: + raise ValueError( + "Index out of range in Input %d %d", + int(torch.min(input).item()), + weight_size[1], + ) + if int(torch.max(input).item()) >= weight_size[0]: + raise ValueError( + "Index out of range in Input %d %d", + int(torch.max(input).item()), + weight_size[1], + ) + if scale_grad_by_freq: + raise RuntimeError( + 'nn.Embedding weight sharded with flag on "scale_grad_by_freq" not supported!' + ) + if sparse: + raise RuntimeError( + 'nn.Embedding weight sharded with flag on "sparse" not supported!' + ) + if max_norm and max_norm <= 0.0: + raise ValueError('"max_norm" must be larger than zero!') + + if not isinstance(weight._sharding_spec, ChunkShardingSpec): + raise ValueError("Only ChunkShardingSpec supported for ShardedTensor ops!") + if len(weight.local_shards()) != 1: + raise ValueError("Only one local shard supported!") + + +def _handle_col_wise_sharding( + input, world_size, weight, local_shard, max_norm, norm_type, padding_idx, pg +): + """ + Entry-point function to handle the logic of col-wise sharding of weight + for embedding. (Detailed explanations of the logic can be found in + the comment for sharded_embedding.) + + Args: + input: list of ID used for lookup and aggregation. + world_size: number of ranks. + weight: sharded weight tensor. + local_shard: col-wise shared local weight used for lookup. + max_norm: If given, each embedding vector with norm larger + than max_norm is renormalized to have norm max_norm. + Note: this will modify weight in-place. + norm_type: The p in the p-norm to compute for the max_norm option. + padding_idx: If specified, the entries at padding_idx do + not contribute to the gradient; therefore, the embedding + vector at padding_idx is not updated during training, + i.e. it remains as a fixed "pad". + pg: process group. + + Returns: final result of lookup. + """ + # allgather the inputs first for non Replicated Tensor. + gathered_inputs = all_gather(input, group=pg) + + if max_norm is not None: + # max_norm changes the weight in-place + local_shard = _handle_max_norm_col_wise( + max_norm, norm_type, local_shard, input, world_size, gathered_inputs, pg + ) + + output = _handle_col_wise_sharding_base( + torch.nn.functional.embedding, + len(input.size()), + input, + world_size, + weight, + local_shard, + pg, + gathered_inputs, + padding_idx=padding_idx, + ) + return (output, local_shard) + + +def _handle_row_wise_sharding( + input, world_size, weight, local_shard, max_norm, norm_type, padding_idx, rank, pg +): + """ + Entry-point function to handle the logic of row-wise sharding of weight + for embedding. (Detailed explanations of the logic can be found in + the comment for sharded_embedding.) + + Args: + input: list of ID used for lookup and aggregation. + world_size: number of ranks. + weight: sharded weight tensor. + local_shard: row-wise shared local weight used for lookup. + max_norm: If given, each embedding vector with norm larger + than max_norm is renormalized to have norm max_norm. + Note: this will modify weight in-place. + norm_type: The p in the p-norm to compute for the max_norm option. + padding_idx: If specified, the entries at padding_idx do + not contribute to the gradient; therefore, the embedding + vector at padding_idx is not updated during training, + i.e. it remains as a fixed "pad". + rank: # of cuda process. + pg: process group. + + Returns: final result of lookup. + """ + # allgather the inputs first for non Replicated Tensor. + gather_inp = _all_gather_base_input(input, pg) + + # Mask the input according to sharding spec. + lookup_input, padding_idx, padding_row = _handle_row_wise_mask( + gather_inp, padding_idx, weight, world_size, rank + ) + + # When input is a large tensor, the value of weight is changed. + # This is a walk-around for now. GH issue: #81717 + if max_norm is not None: + torch.nn.functional.embedding( + torch.unique(lookup_input)[:-1], + local_shard, + padding_idx=padding_idx, + max_norm=max_norm, + norm_type=norm_type, + ) + max_norm = None + + local_input_embeddings = torch.nn.functional.embedding( + lookup_input, + torch.cat([local_shard, padding_row]), + padding_idx=padding_idx, + max_norm=max_norm, + norm_type=norm_type, + ) + + # TODO: Make the result a PartialTensor. + local_shards = local_input_embeddings.chunk(pg.size()) + return reduce_scatter( + torch.empty_like(local_shards[0]), + list(local_shards), + group=pg, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding_bag.py b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding_bag.py new file mode 100644 index 0000000000000000000000000000000000000000..01a148b5a9a93c184b7d7dda7b80f9eb49f32e7d --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding_bag.py @@ -0,0 +1,477 @@ +# mypy: allow-untyped-defs + +from typing import cast, List + +import torch +import torch.distributed as dist +from torch._C._distributed_c10d import ReduceOp +from torch.distributed._shard.sharded_tensor import ShardedTensor +from torch.distributed._shard.sharding_spec import ChunkShardingSpec +from torch.distributed._shard.sharding_spec.api import custom_sharding_spec_op +from torch.distributed.nn.functional import all_gather, reduce_scatter + +from ._common import ( + _all_gather_base_input, + _handle_col_wise_sharding_base, + _handle_max_norm_col_wise, + _handle_row_wise_mask, +) + + +@custom_sharding_spec_op(ChunkShardingSpec, torch.nn.functional.embedding_bag) +def sharded_embedding_bag(types, args, kwargs, pg): + """ + Handles ``__torch_function__`` dispatch for ``torch.nn.functional.embedding_bag``. + This method computes a sharded embedding bag aggregation and has the following limitations: + + 1. Supports only sharding of ``weight``. + 2. Supports only ``ChunkShardingSpec``. + 3. Supports only a single local shard per rank. + 4. Supports all specs except for scale_grad_by_freq, sparse, etc. + + Based on the dimension that the weight is sharded on, there are two + algorithms: + + ROWWISE SHARDING + ================ + For row-wise sharding the weight is sharded on dimension 0. + + The overall algorithm can be best explained with an example. Let's assume + the dims for input are (4 x 6) and W are (16 x 17) and W is sharded across + 4 GPUs creating 4 shard of (4 x 17). + The algorithm is as follows: + + 1. First the input is all gathered to all ranks, since this is SPMD and + input is actually sharded across all ranks. The inputs then become a + 4 (4 x 6) tensor on each rank. For example if the given input is + tensor([[6, 5, 2, 9, 6, 3], + [3, 1, 2, 4, 7, 6], + [4, 0, 4, 9, 8, 9], + [8, 6, 6, 4, 6, 1]]) + on rank 0. + Then on every rank, we will have this tensor. + If input itself is already replicated, no all-gather will be done. + 2. Next, we mask the ID which are not stored on that rank. + For example on rank 0, we store ID [0, 1, 2]. We only keep the ID + inside the set of numbers. The rest of them will be masked to an extra row. + The masked matrix will be used for embedding look up and is like: + tensor([[4, 4, 2, 4, 4, 4], + [4, 1, 2, 4, 4, 4], + [4, 0, 4, 4, 4, 4], + [4, 4, 4, 4, 4, 1]]) + 3. If ``max_norm`` is specified, the extra row guarantees that the mask ID will + not affect the behavior of weigh re-norm. + 4. The example above only happens in one rank and each rank does a very similar thing. + For "Mean" mode we need to divide by either column size (2D) or the interval length + defined by the offset (excluding the row specified in ``padding_idx``). + We also need to mask the unexisting row to neg Inf so that negative value does not + gets wiped out in the "Max" mode. + + COLWISE SHARDING + ================ + For col-wise sharding the weight is sharded on dimension 1. + + The overall algorithm can be best explained with an example. Let's assume + the dims for input are (4 x 6) and W are (16 x 17) and W is sharded across + 4 GPUs creating 3 shards of (16 x 5) and 1 shard of (16 x 2). + The algorithm is as follows: + + 1. First the input is broadcasted to all ranks, since this is SPMD we + actually do an all_gather for all the inputs resulting in 4 (4 x 6) + inputs on each rank. + 2. Next we perform local embedding bag operation under the given mode by + apply each input (4 x 6) with the local shard (16 x 5) ((16 x 2) for the last). + This results in 4 (5 x 4) ((2 x 4) for the last) matrices on each rank. + We transpose the aggregation result. + 3. Next, we concatenate these 4 matrices and perform an all2all to share the + appropriate (5 x 4) or (2 x 4) matrices to each rank. + 4. Now, each rank receives a (17 x 4) matrix which is basically the + size of the result we need. + 5. If placements are not in order any appropriate rearrangement of columns + are done for the (17 x 4) matrix and finally we transpose the output again. + 6. If max_norm is specified, we manually sum up the norm and renorm. Because + the renorm must be in place, we need to override the local_shard to mimic + this behavior. + """ + # Validate input params + _validate_embedding_bag_param(args, kwargs) + + input = args[0] + weight = args[1] + offsets = kwargs.get("offsets") + per_sample_weights = kwargs.get("per_sample_weights") + mode = kwargs.get("mode") + max_norm = kwargs.get("max_norm") + norm_type = kwargs.get("norm_type") + include_last_offset = kwargs.get("include_last_offset") + padding_idx = kwargs.get("padding_idx") + + local_shard = weight.local_tensor().contiguous() + sharding_dim = weight._sharding_spec.dim + world_size = dist.get_world_size(pg) + rank = dist.get_rank(pg) + if include_last_offset: + offsets = offsets[:-1] + + if sharding_dim == 1: + output, local_shard = _handle_col_wise_sharding( + input, + world_size, + weight, + local_shard, + offsets, + per_sample_weights, + mode, + max_norm, + norm_type, + padding_idx, + pg, + ) + weight.local_shards()[0].tensor = local_shard + return output + elif sharding_dim == 0: + return _handle_row_wise_sharding( + input, + world_size, + weight, + local_shard, + offsets, + per_sample_weights, + mode, + max_norm, + norm_type, + padding_idx, + rank, + pg, + ) + else: + raise RuntimeError( + f"nn.EmbeddingBag weight sharded on dim {sharding_dim} not supported!" + ) + + +def _validate_embedding_bag_param(args, kwargs): + """ + Validate input params of sharded embeddingBag op. + + Args: + input: list of ID used for lookup and aggregation. + weight: sharded weight tensor. + kwargs: same as normal EmbeddingBag. + + Return: None. + """ + + input = args[0] + weight = args[1] + offsets = kwargs.get("offsets") + per_sample_weights = kwargs.get("per_sample_weights") + mode = kwargs.get("mode") + max_norm = kwargs.get("max_norm") + scale_grad_by_freq = kwargs.get("scale_grad_by_freq") + sparse = kwargs.get("sparse") + include_last_offset = kwargs.get("include_last_offset") + + # Validate types + if not isinstance(input, torch.Tensor): + raise TypeError("input need to be torch.Tensor") + if offsets is not None and not isinstance(offsets, torch.Tensor): + raise TypeError("offsets need to be torch.Tensor") + if per_sample_weights is not None and not isinstance( + per_sample_weights, torch.Tensor + ): + raise TypeError("per_sample_weights need to be torch.Tensor") + if not isinstance(weight, ShardedTensor): + raise TypeError("weight needs to be ShardedTensor") + if len(input.size()) > 2: + raise ValueError("Input more than 2 dims not supported") + weight_size = weight.size() + if len(weight_size) != 2: + raise ValueError("Weight needs to have exactly 2 dims") + if int(torch.min(input).item()) < 0: + raise ValueError( + "Index out of range in Input %d %d", + int(torch.min(input).item()), + weight_size[1], + ) + if int(torch.max(input).item()) >= weight_size[0]: + raise ValueError( + "Index out of range in Input %d %d", + int(torch.max(input).item()), + weight_size[1], + ) + if offsets is not None and len(input.size()) != 1: + raise ValueError("Input dimension needs to be exactly 1 dim") + if len(input.size()) == 1 and offsets is None: + raise ValueError("offsets is required for 1D input") + if per_sample_weights is not None and per_sample_weights.size() != input.size(): + raise ValueError( + f"per_sample_weights size {per_sample_weights.size()} not equal to input size {input.size()}" + ) + if mode is None: + mode = "mean" + if mode not in ["sum", "mean", "max"]: + raise ValueError(f"mode '{mode}' is not supported") + if scale_grad_by_freq: + raise RuntimeError( + 'nn.Embedding weight sharded with flag on "scale_grad_by_freq" not supported!' + ) + if sparse: + raise RuntimeError( + 'nn.Embedding weight sharded with flag on "sparse" not supported!' + ) + if include_last_offset and offsets is None: + raise ValueError('offsets is required for flag "include_last_offset"!') + if include_last_offset and cast(List[int], offsets)[-1] != input.size(0): + raise ValueError( + 'offsets need to have the input size in the end when the flag "include_last_offset" is on!' + ) + + if max_norm and max_norm <= 0.0: + raise ValueError('"max_norm" must be larger than zero!') + + if not isinstance(weight._sharding_spec, ChunkShardingSpec): + raise ValueError("Only ChunkShardingSpec supported for ShardedTensor ops!") + if len(weight.local_shards()) != 1: + raise ValueError("Only one local shard supported!") + + +def _handle_col_wise_sharding( + input, + world_size, + weight, + local_shard, + offsets, + per_sample_weights, + mode, + max_norm, + norm_type, + padding_idx, + pg, +): + """ + Entry-point function to handle the logic of col-wise sharding of weight + for embeddingBag. (Detailed explanations of the logic can be found in + the comment for sharded_embedding_bag.) + + Args: + input: list of ID used for lookup and aggregation. + world_size: number of ranks. + weight: sharded weight tensor. + local_shard: col-wise shared local weight used for lookup. + offsets: list of start positions of each bag for 1D input. + per_sample_weights: weights for weighted sum mode. + mode: aggregation method of each bag. + max_norm: If given, each embedding vector with norm larger + than max_norm is renormalized to have norm max_norm. + Note: this will modify weight in-place. + norm_type: The p in the p-norm to compute for the max_norm option. + padding_idx: If specified, the entries at padding_idx do + not contribute to the gradient; therefore, the embedding + vector at padding_idx is not updated during training, + i.e. it remains as a fixed "pad". + Note that the embedding vector at padding_idx is + excluded from the reduction. + pg: process group. + + Return: + output: final result of lookup and aggregation. + local_shard: col-wise shared local weight used for lookup. + If max_norm, this will be the renormed weight. + """ + # allgather the special input of embedding bag first. + ( + gathered_inputs, + gathered_per_sample_weights, + gathered_offsets, + ) = _all_gather_embedding_bag_input(input, per_sample_weights, offsets, pg) + + if max_norm is not None: + # max_norm changes the weight in-place + local_shard = _handle_max_norm_col_wise( + max_norm, norm_type, local_shard, input, world_size, gathered_inputs, pg + ) + + output = _handle_col_wise_sharding_base( + torch.nn.functional.embedding_bag, + 1, + input, + world_size, + weight, + local_shard, + pg, + gathered_inputs, + mode=mode, + gathered_per_sample_weights=gathered_per_sample_weights, + gathered_offsets=gathered_offsets, + padding_idx=padding_idx, + ) + return (output, local_shard) + + +def _handle_row_wise_sharding( + input, + world_size, + weight, + local_shard, + offsets, + per_sample_weights, + mode, + max_norm, + norm_type, + padding_idx, + rank, + pg, +): + """ + Entry-point function to handle the logic of row-wise sharding of weight + for embeddingBag. (Detailed explanations of the logic can be found in + the comment for sharded_embedding_bag.) + + Args: + input: list of ID used for lookup and aggregation. + world_size: number of ranks. + weight: sharded weight tensor. + local_shard: row-wise shared local weight used for lookup. + offsets: list of start positions of each bag for 1D input. + per_sample_weights: weights for weighted sum mode. + mode: aggregation method of each bag. + max_norm: If given, each embedding vector with norm larger + than max_norm is renormalized to have norm max_norm. + Note: this will modify weight in-place. + norm_type: The p in the p-norm to compute for the max_norm option. + padding_idx: If specified, the entries at padding_idx do + not contribute to the gradient; therefore, the embedding + vector at padding_idx is not updated during training, + i.e. it remains as a fixed "pad". + Note that the embedding vector at padding_idx is + excluded from the reduction. + rank: # of cuda process. + pg: process group. + + Returns: + gathered_output: final result of lookup and aggregation. + """ + if input.dim() > 1 and per_sample_weights is None: + # allgather the inputs first for non Replicated Tensor. + gather_inp = _all_gather_base_input(input, pg) + else: + ( + gathered_inputs, + gathered_per_sample_weights, + gathered_offsets, + ) = _all_gather_embedding_bag_input(input, per_sample_weights, offsets, pg) + cat_dim = 0 if input.dim() != 1 else -1 + gather_inp = torch.cat(gathered_inputs, dim=cat_dim) + if per_sample_weights is not None: + per_sample_weights = torch.cat(gathered_per_sample_weights, dim=cat_dim) + offset_add = 0 if input.dim() > 1 else input.size(0) + if offsets is not None: + offsets_list = torch.cat( + [gathered_offsets[i] + (offset_add * i) for i in range(pg.size())], + dim=cat_dim, + ) + + # Mask the input according to sharding spec. + lookup_input, padding_local, padding_row = _handle_row_wise_mask( + gather_inp, padding_idx, weight, world_size, rank + ) + if mode == "max": + padding_row[:] = -float("Inf") + + # When input is a large tensor, the value of weight is changed. + # This is a walk-around for now. GH issue: #81717. + if max_norm is not None: + torch.nn.functional.embedding_bag( + torch.unique(lookup_input)[:-1], + local_shard, + offsets=torch.tensor([0], device=local_shard.device, dtype=torch.long), + mode=mode, + per_sample_weights=None, + max_norm=max_norm, + norm_type=norm_type, + padding_idx=padding_local, + ) + max_norm = None + result = torch.nn.functional.embedding_bag( + lookup_input, + torch.cat([local_shard, padding_row]), + offsets=offsets_list if offsets is not None else offsets, # type: ignore[possibly-undefined] + mode=mode if mode != "mean" else "sum", + per_sample_weights=per_sample_weights, + max_norm=max_norm, + norm_type=norm_type, + padding_idx=padding_local, + ) + + op = ReduceOp.SUM if mode != "max" else ReduceOp.MAX + # TODO: Make the result a PartialTensor and move the logic below there. + local_shards = result.chunk(pg.size()) + result = reduce_scatter( + torch.empty_like(local_shards[0]), + list(local_shards), + op=op, + group=pg, + ) + + # For Mean, we cannot do the division until very end because the sum of means + # not equal to the mean of sum. (Divisor is different) + if mode == "mean": + if input.dim() > 1: + padding_idx = padding_idx if padding_idx is not None else -1 + split_sizes = torch.sum( + torch.ne(input, padding_idx), dim=-1, dtype=local_shard.dtype + ) + else: + split_sizes = torch.cat( + ( + offsets[1 : offsets.size(0)] - offsets[0:-1], + (input.size(0) - offsets[-1]).unsqueeze(0), + ), + dim=-1, + ) + return torch.div(result, split_sizes.unsqueeze(1)) + + # Return the appropriate local result. + return result + + +def _all_gather_embedding_bag_input(input, per_sample_weights, offsets, pg): + """ + In case we need to gather input and all other parameters of embeddingBag + ops, we need to stack all input together to perform ``all_gather`` + collective communication just once. + + Note that since offsets does not share the same size as input and + is always smaller than input, we resize it during the communication. + + Args: + input: tensor to be applied op on. + per_sample_weights: weights for weighted sum mode. + offsets: when input is 1D. offsets determines the starting + index position of each bag (sequence) in input. + pg: process group. + + Returns: + gathered_inputs: list of input tensor gathered from each rank. + gathered_per_sample_weights: list of per_sample_weights from each rank. + gathered_offsets: list of offsets from each rank. + """ + input_to_gather = [input] + if per_sample_weights is not None: + input_to_gather.append(per_sample_weights) + if offsets is not None: + input_to_gather.append(offsets.clone().resize_(input.size())) + gathered_inputs = all_gather(torch.stack(input_to_gather), group=pg) + + gathered_per_sample_weights = None + if per_sample_weights is not None: + gathered_per_sample_weights = [t[1] for t in gathered_inputs] + gathered_offsets = None + if offsets is not None: + idx = 2 if per_sample_weights is not None else 1 + gathered_offsets = [ + t[idx].resize_(offsets.size()).to(offsets.dtype) for t in gathered_inputs + ] + gathered_inputs = [t[0].to(input.dtype) for t in gathered_inputs] + return gathered_inputs, gathered_per_sample_weights, gathered_offsets diff --git a/lib/python3.12/site-packages/torch/distributed/_state_dict_utils.py b/lib/python3.12/site-packages/torch/distributed/_state_dict_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..76286df1ae14b81ffff371c76b2eef2a5138e2f2 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_state_dict_utils.py @@ -0,0 +1,756 @@ +# mypy: allow-untyped-defs +import copy +import io +import math +import weakref +from typing import ( + Any, + Callable, + cast, + Dict, + List, + Mapping, + MutableMapping, + NamedTuple, + Optional, + Tuple, + TYPE_CHECKING, + Union, +) + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed._functional_collectives import AsyncCollectiveTensor + + +if dist.is_available() or TYPE_CHECKING: + from torch.distributed import distributed_c10d + from torch.distributed._shard.sharded_tensor import ShardedTensor + from torch.distributed.tensor import distribute_tensor, DTensor, Replicate + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + + +def _identity_func( + obj: torch.Tensor, + pg: Optional[dist.ProcessGroup], + device: Optional[torch.device], + companion_obj: Any, +) -> torch.Tensor: + return obj + + +def _all_gather_sharded_tensor( + sharded_tensor: "ShardedTensor", + pg: Optional[dist.ProcessGroup] = None, + device: Optional[torch.device] = None, +) -> torch.Tensor: + if pg is None: + pg = distributed_c10d._get_default_group() + world_size = dist.get_world_size(pg) + shards = sharded_tensor.local_shards() + dim_0_size = sharded_tensor.size()[0] # type: ignore[index] + tensor_numel = sharded_tensor.size().numel() # type: ignore[union-attr] + chunk_size = math.ceil(dim_0_size / world_size) * tensor_numel // dim_0_size + pg_device = ( + distributed_c10d._get_pg_default_device(pg) if device is None else device + ) + if shards: + local_tensor = shards[0].tensor.flatten() + if local_tensor.device.type != pg_device.type: + local_tensor = local_tensor.to(pg_device) + num_padding = chunk_size - local_tensor.numel() + if num_padding > 0: + local_tensor = F.pad(local_tensor, [0, num_padding]) + else: + local_tensor = torch.zeros( + chunk_size, dtype=sharded_tensor.dtype, device=pg_device + ) + + tensor = torch.empty( + chunk_size * world_size, + dtype=local_tensor.dtype, + device=pg_device, + ) + dist.all_gather_into_tensor(tensor, local_tensor, group=pg) + + tensor = tensor.narrow(0, 0, tensor_numel).reshape(sharded_tensor.size()) + return tensor + + +class CompanionMismatch(Exception): + ... + + +def _iterate_state_dict( + iter_object: Any, + sharded_tensor_func: Callable, + dtensor_func: Callable, + tensor_func: Callable, + *, + pg: Optional[dist.ProcessGroup] = None, + device: Optional[torch.device] = None, + cpu_offload: bool = False, + companion_obj: Any = None, + ranks_only: Tuple[int, ...] = (), + type_check: bool = True, + non_blocking: bool = True, +) -> Dict[str, Any]: + """Iterate through the state dict, applying the given functions to each tensor type. + + Args: + iter_object (Any): the target state_dict. + sharded_tensor_func (Callable): the function to apply to ShardedTensor + dtensor_func (Callable): the function to apply to DTensor + tensor_func (Callable): the function to apply to Tensor + pg (Optional[dist.ProcessGroup]): process group passed to tensor functions + device (Optional[torch.device]): device passed to tensor functions + cpu_offload (bool): whether to offload the tensors to CPU memory. This option is ignored + if a companion_obj is supplied. + companion_obj (Any): A companion object to the state dict. If this object + is supplied, we attempt to copy the tensor to the companion object. + ranks_only (Tuple[int, ...]): if this tuple is empty, all ranks will + have the same state_dicts. Otherwise only ranks that in ``ranks_only`` + have the same state_dicts. Other ranks will get empty state_dicts. + type_check (bool): check if the instance data type is a supported type + that can be saved by DCP. The current supported data types are + torch.Tensor, DTensor, int, float, str, list, dict, None. + non_blocking (bool): whether to use non-blocking copy when copying to the companion object. + """ + # TODO: should we use pytree? + cpu_device = torch.device("cpu") + if isinstance(iter_object, ShardedTensor): + ret = sharded_tensor_func(iter_object, pg, device, companion_obj) + elif isinstance(iter_object, DTensor): + ret = dtensor_func(iter_object, pg, device, companion_obj) + elif isinstance(iter_object, torch.Tensor): + ret = tensor_func(iter_object, pg, device, companion_obj) + elif ( + isinstance(iter_object, (int, float, str, bytes, io.BytesIO)) + or iter_object is None + ): + ret = iter_object + elif isinstance(iter_object, dict): + if companion_obj is not None and ( + not isinstance(companion_obj, dict) + or set(companion_obj.keys()) != set(iter_object.keys()) + ): + msg = ( + "" + if isinstance(companion_obj, dict) + else f"{set(companion_obj.keys())=} {set(iter_object.keys())=}" + ) + raise CompanionMismatch(msg) + + ret = { + key: _iterate_state_dict( + value, + sharded_tensor_func, + dtensor_func, + tensor_func, + pg=pg, + device=device, + cpu_offload=cpu_offload, + companion_obj=companion_obj[key] if companion_obj is not None else None, + ranks_only=ranks_only, + type_check=type_check, + non_blocking=non_blocking, + ) + for key, value in iter_object.items() + } + elif isinstance(iter_object, (list, tuple)): + if companion_obj is not None and ( + not isinstance(companion_obj, (list, tuple)) + or len(companion_obj) != len(iter_object) + ): + raise CompanionMismatch + + ret = [ + _iterate_state_dict( + v, + sharded_tensor_func, + dtensor_func, + tensor_func, + pg=pg, + device=device, + cpu_offload=cpu_offload, + companion_obj=companion_obj[idx] if companion_obj is not None else None, + ranks_only=ranks_only, + type_check=type_check, + non_blocking=non_blocking, + ) + for idx, v in enumerate(iter_object) + ] + if isinstance(iter_object, tuple): + ret = tuple(ret) + elif not type_check: + ret = copy.deepcopy(iter_object) + else: + raise ValueError(f"Unexpected value type {type(iter_object)}") + + if not ranks_only or dist.get_rank(pg) in ranks_only: + if isinstance(ret, torch.Tensor): + if cpu_offload and companion_obj is None: + ret = ret.to(cpu_device) + + if companion_obj is not None: + # TODO: support DTensor + companion_obj.copy_(ret, non_blocking=non_blocking) + ret = companion_obj + else: + ret = {} if isinstance(ret, dict) else None + + return ret + + +def _gather_state_dict( + state_dict: Dict[str, Any], + *, + pg: Optional[dist.ProcessGroup] = None, + device: Optional[torch.device] = None, + cpu_offload: bool = False, + ranks_only: Tuple[int, ...] = (), + type_check: bool = True, +) -> Dict[str, Any]: + """ + Given a state_dict, this API gathers all the ShardedTensors or DTensors in + the state_dict. + + + Args: + state_dict (Dict[str, Any]): the target sharded state_dict. + pg (Optional[dist.ProcessGroup]): the process group that is used to + gather ShardedTensor. Note that gathering a DTensor will use + the DeviceMesh. So this argument will be ignored when gathering a + DTensor. + device: (Optional[torch.device]): the device that is used to + perform allgather for ShardedTensor. Note that gathering a DTensor + will use the DeviceMesh. So this argument will be ignored when + gathering a DTensor. + cpu_offload (bool): whether to offload the tensors to CPU memory. The + default value is False. + ranks_only: (Tuple[int, ...]): if this tuple is empty, all ranks will + have the same state_dicts. Otherwise only ranks that in ``ranks_only`` + have the same state_dicts. Other ranks will get empty state_dicts. + type_check: (bool): check if the instance data type is a supported type + that can be saved by DCP. The current supported data types are + torch.Tensor, DTensor, int, float, str, list, dict, None. + + Returns: + The gathered state dictionary. + """ + + def sharded_tensor_func(value, pg, device, companion_obj): + # ShardedTensor does not seem to record the original device type. + # So if the tensor is moved to CPU, we won't know the original type. + # As a result, we have to rely on the user to tell us the correct one. + cpu_device = torch.device("cpu") + output_tensor = _all_gather_sharded_tensor(value, pg, device) + local_shard_device = ( + value.local_shards()[0].tensor.device + if value.local_shards() + else cpu_device + ) + if output_tensor.device != local_shard_device: + value = output_tensor.to(local_shard_device) + else: + value = output_tensor + return value + + def dtensor_func(value, pg, device, companion_obj): + if value.device != value.device_mesh.device_type: + value = value.to(value.device_mesh.device_type) + # FSDP all_gather: [Shard(0)] -> [Replicate()] + # HSDP all_gather: [Replicate(), Shard(0)] -> [Replicate(), Replicate()] + # 2D FSDP + TP all_gather: + # - [Shard(0), Shard(n)] -> [Replicate(), Replicate()] + # - [Shard(0), Replicate()] -> [Replicate(), Replicate()] + placements = [Replicate() for _ in value.placements] + value = value.redistribute( + device_mesh=value.device_mesh, + placements=placements, + ) + # Call `wait()` to force the tensor to be synchronous with respect + # to the main stream. + # See the discussion in https://github.com/pytorch/pytorch/pull/117799. + value = value.to_local() + if isinstance(value, AsyncCollectiveTensor): + value = value.wait() + return value + + return _iterate_state_dict( + state_dict, + sharded_tensor_func, + dtensor_func, + _identity_func, + pg=pg, + device=device, + cpu_offload=cpu_offload, + ranks_only=ranks_only, + type_check=type_check, + ) + + +def _offload_state_dict_to_cpu( + state_dict: Dict[str, Any], + *, + ranks_only: Tuple[int, ...] = (), + type_check: bool = True, +) -> Dict[str, Any]: + """ + Given a state_dict, this API offload all the tensors to CPU memory. + + Args: + state_dict (Dict[str, Any]): the target state_dict. + pg (Optional[dist.ProcessGroup]): the process group that is used to + gather ShardedTensor. Note that gathering a DTensor will use + the DeviceMesh. So this argument will be ignored when gathering a + DTensor. + ranks_only: (Tuple[int, ...]): if this tuple is empty, all ranks will + have the same state_dicts. Otherwise only ranks that in ``ranks_only`` + have the same state_dicts. Other ranks will get empty state_dicts. + type_check: (bool): check if the instance data type is a supported type + that can be saved by DCP. The current supported data types are + torch.Tensor, DTensor, int, float, str, list, dict, None. + + Returns: + The gathered state dictionary. + """ + + ret = _iterate_state_dict( + state_dict, + _identity_func, + _identity_func, + _identity_func, + pg=None, + device=None, + cpu_offload=True, + ranks_only=ranks_only, + type_check=type_check, + ) + return ret + + +def _copy_state_dict( + state_dict: Dict[str, Any], + copy_state_dict: Dict[str, Any], + non_blocking: bool = False, + type_check: bool = True, +) -> Dict[str, Any]: + """ + Copies all tensors in a given state dict into a different state_dict with the + same structure. Additionally, a copied state dict with the same value references + is returned. Editing the keys on this state dict will not affect the + passed in copy_state_dict (but the value references are the same). + + .. warning:: + It is expected by this function that state_dict and copy_state_dict share + the same structure and data types. + + .. warning:: + The current supported data types are + torch.Tensor, DTensor, int, float, str, list, dict, None. + + Args: + state_dict (Dict[str, Any]): the target state_dict. + copy_state_dict (Dict[str, Any]): + The state dict we are copying into. This state_dict must have exactly + the same structure as the source `state_dict`. + non_blocking: (bool): Whether copy ops should be performed asynchronously + type_check (bool): check if the instance data type is a supported type + that can be saved by DCP. The current supported data types are + torch.Tensor, DTensor, int, float, str, list, dict, None. + + Returns: + State Dict copy + """ + + return _iterate_state_dict( + state_dict, + _identity_func, + _identity_func, + _identity_func, + pg=None, + device=None, + cpu_offload=False, + ranks_only=(), + companion_obj=copy_state_dict, + type_check=type_check, + non_blocking=non_blocking, + ) + + +def _create_cpu_state_dict( + state_dict: Dict[str, Any], pin_memory: bool = False, share_memory: bool = False +) -> Dict[str, Any]: + """ + Given a state_dict, create another state_dict with the same structure and elements. + However, all tensors in the returned state_dict are new tensors on CPU. These + tensors can be placed on pin_memory or share_memory based on the provided arguments. + + .. warning:: + Setting both `pin_memory` and `share_memory` to True significantly increases the + latency of this method because of the nuances which require us to register memory + as pinned directly as opposed to relying on the pin_memory cache allocator. This + option should only be used for long lived tensors which are required to be shared. + This is not the case as long as at least one of `pin_memory` or `share_memory` is + set to False. + + """ + + def tensor_func( + obj: torch.Tensor, + pg: Optional[dist.ProcessGroup], + device: Optional[torch.device], + _: Any, + ) -> torch.Tensor: + if len(obj.size()) == 0: + return torch.tensor(0, dtype=obj.dtype) + + if share_memory: + t = torch.empty(*tuple(obj.size()), dtype=obj.dtype) + t = t.share_memory_() + if pin_memory: + + def unpin_memory(t): + succ = int(torch.cuda.cudart().cudaHostUnregister(t.data_ptr())) + assert ( + succ == 0 + ), f"Unpinning shared memory failed with error-code: {succ}" + + weakref.finalize(t, unpin_memory, t) + succ = int( + torch.cuda.cudart().cudaHostRegister( + t.data_ptr(), + t.numel() * t.element_size(), + 1, # lines up with 'cudaHostRegisterPortable' + ) + ) + assert ( + succ == 0 + ), f"Pinning shared memory failed with error-code: {succ}" + return t + elif pin_memory: + return torch.empty(*tuple(obj.size()), dtype=obj.dtype).pin_memory() + else: + return torch.empty(*tuple(obj.size()), dtype=obj.dtype) + + ret = _iterate_state_dict( + state_dict, + _identity_func, + _identity_func, + tensor_func, + pg=None, + device=None, + cpu_offload=False, + ranks_only=(), + type_check=False, + ) + return ret + + +def _check_state_dict_similarity( + state_dict: Dict[str, Any], + compared_state_dict: Dict[str, Any], +) -> bool: + """ + Given two state_dicts, check if the structures are the same. And + if a [key, tensor] pair exist in one state_dict there must be + the a corresponding pait, [key, other_tensor], in the other state_dict, + where tensor and other_tensor have the same size and dtype. + + Return the check result. + """ + + def tensor_func( + obj: torch.Tensor, + pg: Optional[dist.ProcessGroup], + device: Optional[torch.device], + companion_obj: Any, + ) -> torch.Tensor: + if companion_obj.dtype != obj.dtype or companion_obj.size() != obj.size(): + raise CompanionMismatch + return obj + + try: + _iterate_state_dict( + state_dict, + _identity_func, + _identity_func, + tensor_func, + pg=None, + device=None, + cpu_offload=False, + ranks_only=(), + companion_obj=compared_state_dict, + type_check=False, + ) + except CompanionMismatch: + return False + + return True + + +class _TensorInfo(NamedTuple): + size: torch.Size + dtype: torch.dtype + + +def _broadcast_tensors( + full_state_dict: Dict[str, Any], + local_state_dict: Dict[str, Any], + keys: List[str], + device: torch.device, + pg: Optional[dist.ProcessGroup] = None, +) -> None: + tensors = [] + for key in keys: + if dist.get_rank() == 0: + full_state = full_state_dict[key] + assert isinstance(full_state, torch.Tensor) + full_tensor = full_state.detach().to(device) + else: + tensor_info = full_state_dict[key] + full_tensor = torch.empty( + size=tensor_info.size, + device=device, + dtype=tensor_info.dtype, + ) + + tensors.append(full_tensor) + local_state = local_state_dict.get(key, None) + if local_state is None: + continue + elif isinstance(local_state, DTensor): + local_state_dict[key] = (local_state, full_tensor) + else: + local_state_dict[key] = full_tensor + + if pg is None: + pg = dist.distributed_c10d._get_default_group() + + if len(tensors) > 1: + dist._broadcast_coalesced(pg, tensors, 500, 0) + else: + dist.broadcast(tensors[0], src=0, group=pg) + + _distribute_tensors(local_state_dict, keys, device, pg) + + +def _distribute_tensors( + local_state_dict: Dict[str, Any], + keys: List[str], + device: torch.device, + pg: Optional[dist.ProcessGroup] = None, +) -> None: + if pg is None: + pg = dist.distributed_c10d._get_default_group() + for key in keys: + _local_state = local_state_dict.get(key, None) + if _local_state is None or torch.is_tensor(_local_state): + continue + + local_state = _local_state[0] + full_tensor = _local_state[1] + + shape, offset = compute_local_shape_and_global_offset( + full_tensor.shape, local_state.device_mesh, local_state.placements + ) + slices = [ + slice(cur_offset, cur_offset + cur_shape) + for cur_shape, cur_offset in zip(shape, offset) + ] + local_tensor = full_tensor[slices] + # TODO: currently, we cannot handle strided sharding if the dp dimension is not even. For example, + # one of the case that is not yet supported is when placements = (Shard(0), _StridedShard(0, sf=2)). + local_state_dict[key] = DTensor.from_local( + local_tensor, + local_state.device_mesh, + local_state.placements, + shape=local_state.shape, + stride=local_state.stride(), + ) + + +def _broadcast_state_dict( + full_state_dict: Dict[str, Any], + local_state_dict: Dict[str, Any], + device: torch.device, + pg: Optional[dist.ProcessGroup] = None, + strict: bool = False, +) -> None: + # Broadcast from rank0's `full_state_dict` to all ranks' `local_state_dict`. + # If strict is True, any keys in `local_state_dict` but not in `full_state_dict` + # will be removed from `local_state_dict`. + ret = {} + if dist.get_rank() == 0: + for key, value in full_state_dict.items(): + if not torch.is_tensor(value): + ret[key] = value + elif value.dim() == 0: + ret[key] = value.cpu() + else: + ret[key] = _TensorInfo(value.size(), value.dtype) + + broadcast_list = [ret] + dist.broadcast_object_list(broadcast_list, src=0, group=pg) + ret = broadcast_list[0] + + # Gather values + keys = [] + local_state_dict_keys = set(local_state_dict.keys()) + global_keys = set() + for key, value in ret.items(): + global_keys.add(key) + if not isinstance(value, _TensorInfo): + if key in local_state_dict: + local_state_dict[key] = value + continue + + if dist.get_rank() == 0: + ret[key] = full_state_dict[key] + + keys.append(key) + # Broadcast every tensor to avoid OOM for now. + if len(keys) >= 1: + _broadcast_tensors(ret, local_state_dict, keys, device, pg) + keys.clear() + + if strict: + if missing_keys := (local_state_dict_keys - global_keys): + for key in missing_keys: + local_state_dict.pop(key) + + if keys: + _broadcast_tensors(ret, local_state_dict, keys, device, pg) + + +def _distribute_state_dict( + full_state_dict: Dict[str, Any], + local_state_dict: Dict[str, Any], + device: torch.device, + pg: Optional[dist.ProcessGroup] = None, +) -> None: + # Full_state_dict = True, broadcast_from_rank0 = False here. Each rank has + # full_state_dict. Skip the broadcast in ``_broadcast_state_dict`` and + # distribute tensors in each rank + for key, value in full_state_dict.items(): + if key not in full_state_dict: + continue + if not torch.is_tensor(value): + local_state_dict[key] = value + elif value.dim() == 0: + local_state_dict[key] = value.cpu() + else: + assert isinstance(value, torch.Tensor) + local_state = local_state_dict.get(key, None) + if local_state is None: + continue + elif isinstance(local_state, DTensor): + local_state_dict[key] = distribute_tensor( + value.detach().to(device), + local_state.device_mesh, + local_state.placements, + ) + else: + local_state_dict[key] = value.detach().to(device) + + +# These APIs are from torch.distributed.checkpoint. +# TODO: We should consolidate the code here as some not all modules can depend on +# DCP. +PATH_ITEM = Union[str, int] +OBJ_PATH = Tuple[PATH_ITEM, ...] +FLATTEN_MAPPING = Dict[str, OBJ_PATH] +STATE_DICT_TYPE = Dict[str, Any] +CONTAINER_TYPE = MutableMapping[PATH_ITEM, Any] + + +def _traverse_state_dict( + state_dict: STATE_DICT_TYPE, + visitor: Callable[[OBJ_PATH, Any], None], +) -> None: + """ + Invoke ``visitor`` for each value recursively in ``state_dict``. + Mapping, list, and tuple will be flattened and other value types are treated + as the terminal values and will invoke ``visitor``. + """ + + def _traverse_obj(path: OBJ_PATH, value: Any) -> None: + if isinstance(value, Mapping): + for k, v in value.items(): + _traverse_obj(path + (str(k),), v) + elif isinstance(value, (list, tuple)): + for i, v in enumerate(value): + _traverse_obj(path + (i,), v) + else: + visitor(path, value) + + for key, value in state_dict.items(): + _traverse_obj((str(key),), value) + + +def _flatten_state_dict( + state_dict: STATE_DICT_TYPE, +) -> Tuple[STATE_DICT_TYPE, FLATTEN_MAPPING]: + """ + Flatten ``state_dict`` made of nested dicts and lists into a top level dictionary. + + Use ``unflatten_state_dict`` to revert this process. + Returns: + A tuple with the flatten state_dict and a mapping from original to new state_dict. + N.B. The new keys are derived from the object paths, joined by dot. + For example: ``{ 'a': {'b':...}}`` results in the key `a.b`. + """ + flattened: STATE_DICT_TYPE = {} + mappings: FLATTEN_MAPPING = {} + + def flat_copy(path: OBJ_PATH, value: Any) -> None: + new_fqn = ".".join(map(str, path)) + if new_fqn in flattened: + raise ValueError(f"duplicated flatten key {new_fqn}") + flattened[new_fqn] = value + mappings[new_fqn] = path + + _traverse_state_dict(state_dict, flat_copy) + return flattened, mappings + + +def _set_element(root_dict: STATE_DICT_TYPE, path: OBJ_PATH, value: Any) -> None: + """Set ``value`` in ``root_dict`` along the ``path`` object path.""" + cur_container = cast(CONTAINER_TYPE, root_dict) + + def extend_list(lst: List[Any], idx: int) -> None: + while len(lst) <= idx: + lst.append(None) + + for i in range(1, len(path)): + prev_key = path[i - 1] + key = path[i] + def_val: Union[CONTAINER_TYPE, List[Any]] = {} if type(key) == str else [] + + if isinstance(cur_container, Mapping): + cur_container = cast( + CONTAINER_TYPE, cur_container.setdefault(prev_key, def_val) + ) + else: + extend_list(cur_container, prev_key) + if cur_container[prev_key] is None: + cur_container[prev_key] = def_val + cur_container = cur_container[prev_key] + + key = path[-1] + if type(key) == int: + extend_list(cast(List[Any], cur_container), key) + + cur_container[key] = value + + +def _unflatten_state_dict( + state_dict: STATE_DICT_TYPE, mapping: FLATTEN_MAPPING +) -> STATE_DICT_TYPE: + """Restore the original nested state_dict according to ``mapping`` and the flattened ``state_dict``.""" + nested: STATE_DICT_TYPE = {} + for key, value in state_dict.items(): + _set_element(nested, mapping[key], value) + return nested diff --git a/lib/python3.12/site-packages/torch/distributed/_tensor/__init__.py b/lib/python3.12/site-packages/torch/distributed/_tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..40f9727015d7644916167e8fdaee10ac37593f34 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_tensor/__init__.py @@ -0,0 +1,44 @@ +""" +NOTICE: DTensor has moved to torch.distributed.tensor + +This file is a shim to redirect to the new location, and +we keep the old import path starts with `_tensor` for +backward compatibility. We will remove this folder once +we resolve all the BC issues. +""" +import sys +from importlib import import_module + + +submodules = [ + # TODO: _shards_wrapper/_utils here mainly for checkpoint BC, remove them + "_shards_wrapper", + "_utils", + "experimental", + "device_mesh", +] + +# Redirect imports +for submodule in submodules: + full_module_name = f"torch.distributed.tensor.{submodule}" + sys.modules[f"torch.distributed._tensor.{submodule}"] = import_module( + full_module_name + ) + +from torch.distributed.tensor import ( # noqa: F401 + DeviceMesh, + distribute_module, + distribute_tensor, + DTensor, + empty, + full, + init_device_mesh, + ones, + Partial, + Placement, + rand, + randn, + Replicate, + Shard, + zeros, +) diff --git a/lib/python3.12/site-packages/torch/distributed/_tensor/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_tensor/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..74d7d2da2fcfccc6ec1679c6b18e5d6e311307c2 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_tensor/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_tensor/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_tensor/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84aca4465f9f0303f709183b5147909455e43126 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_tensor/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_tensor/__pycache__/placement_types.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/_tensor/__pycache__/placement_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..43aa8d65f9be157972e587a3d4570328bc04eeca Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/_tensor/__pycache__/placement_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/_tensor/api.py b/lib/python3.12/site-packages/torch/distributed/_tensor/api.py new file mode 100644 index 0000000000000000000000000000000000000000..9e5742156a86ca511619360038a9028b0efeeaef --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_tensor/api.py @@ -0,0 +1,9 @@ +""" +NOTE: torch.distributed._tensor has been moved to torch.distributed.tensor. +The imports here are purely for backward compatibility. We will remove these +imports in a few releases + +TODO: throw warnings when this module imported +""" + +from torch.distributed.tensor._api import * # noqa: F401, F403 diff --git a/lib/python3.12/site-packages/torch/distributed/_tensor/placement_types.py b/lib/python3.12/site-packages/torch/distributed/_tensor/placement_types.py new file mode 100644 index 0000000000000000000000000000000000000000..6a4e70dbba455471feef2326cae8ba28b32d0304 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/_tensor/placement_types.py @@ -0,0 +1,10 @@ +""" +NOTE: torch.distributed._tensor has been moved to torch.distributed.tensor. +The imports here are purely for backward compatibility. We will remove these +imports in a few releases + +TODO: throw warnings when this module imported +""" + +from torch.distributed.tensor._dtensor_spec import * # noqa: F401, F403 +from torch.distributed.tensor.placement_types import * # noqa: F401, F403 diff --git a/lib/python3.12/site-packages/torch/distributed/argparse_util.py b/lib/python3.12/site-packages/torch/distributed/argparse_util.py new file mode 100644 index 0000000000000000000000000000000000000000..c475eebf21273abb53ab99e3edcbdef18e9f0c8f --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/argparse_util.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +# mypy: allow-untyped-defs + +# Copyright (c) Facebook, Inc. and its affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +import os +from argparse import Action + + +class env(Action): + """ + Get argument values from ``PET_{dest}`` before defaulting to the given ``default`` value. + + For flags (e.g. ``--standalone``) + use ``check_env`` instead. + + .. note:: when multiple option strings are specified, ``dest`` is + the longest option string (e.g. for ``"-f", "--foo"`` + the env var to set is ``PET_FOO`` not ``PET_F``) + + Example: + :: + + parser.add_argument("-f", "--foo", action=env, default="bar") + + ./program -> args.foo="bar" + ./program -f baz -> args.foo="baz" + ./program --foo baz -> args.foo="baz" + PET_FOO="env_bar" ./program -f baz -> args.foo="baz" + PET_FOO="env_bar" ./program --foo baz -> args.foo="baz" + PET_FOO="env_bar" ./program -> args.foo="env_bar" + + parser.add_argument("-f", "--foo", action=env, required=True) + + ./program -> fails + ./program -f baz -> args.foo="baz" + PET_FOO="env_bar" ./program -> args.foo="env_bar" + PET_FOO="env_bar" ./program -f baz -> args.foo="baz" + """ + + def __init__(self, dest, default=None, required=False, **kwargs) -> None: + env_name = f"PET_{dest.upper()}" + default = os.environ.get(env_name, default) + + # ``required`` means that it NEEDS to be present in the command-line args + # rather than "this option requires a value (either set explicitly or default" + # so if we found default then we don't "require" it to be in the command-line + # so set it to False + if default: + required = False + + super().__init__(dest=dest, default=default, required=required, **kwargs) + + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, values) + + +class check_env(Action): + """ + Check whether the env var ``PET_{dest}`` exists before defaulting to the given ``default`` value. + + Equivalent to + ``store_true`` argparse built-in action except that the argument can + be omitted from the commandline if the env var is present and has a + non-zero value. + + .. note:: it is redundant to pass ``default=True`` for arguments + that use this action because a flag should be ``True`` + when present and ``False`` otherwise. + + Example: + :: + + parser.add_argument("--verbose", action=check_env) + + ./program -> args.verbose=False + ./program --verbose -> args.verbose=True + PET_VERBOSE=1 ./program -> args.verbose=True + PET_VERBOSE=0 ./program -> args.verbose=False + PET_VERBOSE=0 ./program --verbose -> args.verbose=True + + Anti-pattern (don't do this): + + :: + + parser.add_argument("--verbose", action=check_env, default=True) + + ./program -> args.verbose=True + ./program --verbose -> args.verbose=True + PET_VERBOSE=1 ./program -> args.verbose=True + PET_VERBOSE=0 ./program -> args.verbose=False + + """ + + def __init__(self, dest, default=False, **kwargs) -> None: + env_name = f"PET_{dest.upper()}" + default = bool(int(os.environ.get(env_name, "1" if default else "0"))) + super().__init__(dest=dest, const=True, default=default, nargs=0, **kwargs) + + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, self.const) diff --git a/lib/python3.12/site-packages/torch/distributed/autograd/__init__.py b/lib/python3.12/site-packages/torch/distributed/autograd/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b1cf0aec6140f8c61e36ca90de5165399447b290 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/autograd/__init__.py @@ -0,0 +1,53 @@ +# mypy: allow-untyped-defs + +import torch + + +def is_available(): + return hasattr(torch._C, "_dist_autograd_init") + + +if is_available() and not torch._C._dist_autograd_init(): + raise RuntimeError("Failed to initialize torch.distributed.autograd") + +if is_available(): + from torch._C._distributed_autograd import ( + _current_context, + _get_debug_info, + _get_max_id, + _init, + _is_valid_context, + _new_context, + _release_context, + _retrieve_context, + backward, + DistAutogradContext, + get_gradients, + ) + + +class context: + """ + Context object to wrap forward and backward passes when using + distributed autograd. The ``context_id`` generated in the ``with`` + statement is required to uniquely identify a distributed backward pass + on all workers. Each worker stores metadata associated with this + ``context_id``, which is required to correctly execute a distributed + autograd pass. + + Example:: + >>> # xdoctest: +SKIP + >>> import torch.distributed.autograd as dist_autograd + >>> with dist_autograd.context() as context_id: + >>> t1 = torch.rand((3, 3), requires_grad=True) + >>> t2 = torch.rand((3, 3), requires_grad=True) + >>> loss = rpc.rpc_sync("worker1", torch.add, args=(t1, t2)).sum() + >>> dist_autograd.backward(context_id, [loss]) + """ + + def __enter__(self): + self.autograd_context = _new_context() + return self.autograd_context._context_id() + + def __exit__(self, type, value, traceback): + _release_context(self.autograd_context._context_id()) diff --git a/lib/python3.12/site-packages/torch/distributed/autograd/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/autograd/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9bd45b8bfcb56a46cb15f70c928b76e364888dec Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/autograd/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/c10d_logger.py b/lib/python3.12/site-packages/torch/distributed/c10d_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..17f822945852a899233e4d391dc8e1d488df2c96 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/c10d_logger.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +# mypy: allow-untyped-defs + +# Copyright (c) Facebook, Inc. and its affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import functools +import logging +from typing import Any, Callable, Dict, List, Tuple, TypeVar +from typing_extensions import ParamSpec + +import torch +import torch.distributed as dist +from torch.distributed.logging_handlers import _log_handlers +from torch.monitor import _WaitCounter + + +__all__: List[str] = [] + +_DEFAULT_DESTINATION = "default" + + +def _get_or_create_logger(destination: str = _DEFAULT_DESTINATION) -> logging.Logger: + logging_handler, log_handler_name = _get_logging_handler(destination) + logger = logging.getLogger(f"c10d-{log_handler_name}") + logger.setLevel(logging.DEBUG) + formatter = logging.Formatter( + "%(asctime)s %(filename)s:%(lineno)s %(levelname)s p:%(processName)s t:%(threadName)s: %(message)s" + ) + logging_handler.setFormatter(formatter) + logger.propagate = False + logger.addHandler(logging_handler) + return logger + + +def _get_logging_handler( + destination: str = _DEFAULT_DESTINATION, +) -> Tuple[logging.Handler, str]: + log_handler = _log_handlers[destination] + log_handler_name = f"{type(log_handler).__name__}-{destination}" + return (log_handler, log_handler_name) + + +global _c10d_logger +_c10d_logger = _get_or_create_logger() + + +def _get_msg_dict(func_name, *args, **kwargs) -> Dict[str, Any]: + if dist.is_initialized(): + group = kwargs.get("group") or kwargs.get("process_group") + msg_dict = { + "func_name": f"{func_name}", + "pg_name": f"{dist._get_process_group_name(kwargs.get('pg'))}", # type: ignore[arg-type] + "backend": f"{dist.get_backend(group)}", + "world_size": f"{dist.get_world_size()}", + "group_size": f"{dist.get_world_size(group)}", + "global_rank": f"{dist.get_rank()}", + "local_rank": f"{dist.get_rank(group)}", + } + if msg_dict["backend"] == "nccl": + nccl_version = torch.cuda.nccl.version() + msg_dict["nccl_version"] = ".".join(str(v) for v in nccl_version) + else: + msg_dict = { + "func_name": f"{func_name}", + } + return msg_dict + + +_T = TypeVar("_T") +_P = ParamSpec("_P") + + +def _exception_logger(func: Callable[_P, _T]) -> Callable[_P, _T]: + @functools.wraps(func) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T: + try: + return func(*args, **kwargs) + except Exception as error: + msg_dict = _get_msg_dict(func.__name__, *args, **kwargs) + msg_dict["error"] = f"{error}" + _c10d_logger.debug(msg_dict) + raise + + return wrapper + + +def _time_logger(func: Callable[_P, _T]) -> Callable[_P, _T]: + @functools.wraps(func) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T: + with _WaitCounter(f"pytorch.wait_counter.c10d.{func.__name__}").guard(): + func_return = func(*args, **kwargs) + return func_return + + return wrapper diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/__init__.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3262acccac1a6d4c1c157233ef6ac519b3629e93 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/__init__.py @@ -0,0 +1,14 @@ +from .api import CheckpointException +from .default_planner import DefaultLoadPlanner, DefaultSavePlanner +from .filesystem import FileSystemReader, FileSystemWriter +from .metadata import ( + BytesStorageMetadata, + ChunkStorageMetadata, + Metadata, + TensorStorageMetadata, +) +from .optimizer import load_sharded_optimizer_state_dict +from .planner import LoadPlan, LoadPlanner, ReadItem, SavePlan, SavePlanner, WriteItem +from .state_dict_loader import load, load_state_dict +from .state_dict_saver import async_save, save, save_state_dict +from .storage import StorageReader, StorageWriter diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/_checkpointer.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/_checkpointer.py new file mode 100644 index 0000000000000000000000000000000000000000..a93fe8197dea66e5e331133f94c0067ff8ecb90f --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/_checkpointer.py @@ -0,0 +1,100 @@ +from concurrent.futures import Future +from typing import Any, Dict, List, Optional + +import torch.distributed as dist +import torch.distributed.checkpoint.state_dict_loader as loader +import torch.distributed.checkpoint.state_dict_saver as saver +from torch.distributed.checkpoint.metadata import Metadata, STATE_DICT_TYPE +from torch.distributed.checkpoint.storage import ( + LoadPlanner, + SavePlanner, + StorageReader, + StorageWriter, +) + + +__all__: List[str] = [] + + +class _Checkpointer: + """This base class specefies a high level API for saving and loading + distributed `state_dict` 's. It provides an abstraction over the low-level APIs + provided by :py:mod:`torch.distributed.checkpoint.storage`, essentially calling + :py:meth: `torch.distributed.state_dict_saver.save` and + :py:meth: `torch.distributed.state_dict_loader.load` with the provided storage + readers and writers. + + .. warning:: + This feature is experimental and subject to removal/change. + + """ + + def __init__( + self, + storage_writer: StorageWriter, + storage_reader: StorageReader, + *, + process_group: Optional[dist.ProcessGroup] = None, + coordinator_rank: int = 0, + no_dist: bool = False, + load_planner: Optional[LoadPlanner] = None, + save_planner: Optional[SavePlanner] = None, + ): + """Initializes the Checkpointer instance. + + Args: + storage_writer: Instance of StorageWrite use to perform writes. + storage_reader: StorageReader used to load data from. + process_group: ProcessGroup to be used for cross-rank synchronization. + coordinator_rank: Rank to use to coordinate the checkpoint. rank0 is used by default. + no_dist: If ``True``, distributed checkpoint will not load in SPMD style. (Default: ``False``) + loader_planner: Instance of LoadPlanner to use when loading. + save_planner: Instance of SavePlanner to use when saving. + """ + self.storage_writer = storage_writer + self.storage_reader = storage_reader + self.process_group = process_group + self.coordinator_rank = coordinator_rank + self.no_dist = no_dist + self.load_planner = load_planner + self.save_planner = save_planner + + def save( + self, + state_dict: STATE_DICT_TYPE, + ) -> Metadata: + """Calls :py:meth: `torch.distributed.state_dict_saver.save`. Utilizing values passed during initialization.""" + return saver.save( + state_dict, + self.storage_writer, + process_group=self.process_group, + coordinator_rank=self.coordinator_rank, + no_dist=self.no_dist, + planner=self.save_planner, + ) + + def async_save( + self, + state_dict: STATE_DICT_TYPE, + ) -> Future: + """ + Calls :py:meth: `torch.distributed.state_dict_saver._async_save`. Utilizing values passed during initialization. + + Returns: + Future: A future holding the resultant Metadata object from `save`. + """ + return saver.async_save( + state_dict, + storage_writer=self.storage_writer, + process_group=self.process_group, + planner=self.save_planner, + ) + + def load(self, state_dict: Dict[str, Any]) -> None: + """Calls :py:meth: `torch.distributed.state_dict_loader.load`. Utilizing values passed during initialization.""" + loader.load( + state_dict, + storage_reader=self.storage_reader, + process_group=self.process_group, + planner=self.load_planner, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/_dedup_save_plans.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/_dedup_save_plans.py new file mode 100644 index 0000000000000000000000000000000000000000..dd37634a0aa64bb5e849375a4fcedb39617b998e --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/_dedup_save_plans.py @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates +import dataclasses +from collections import defaultdict +from typing import Dict, List, Set, TYPE_CHECKING + +from torch.distributed.checkpoint.planner import SavePlan, WriteItem + + +if TYPE_CHECKING: + from torch.distributed.checkpoint.metadata import MetadataIndex + +__all__ = ["dedup_save_plans"] + + +def dedup_save_plans( + all_plans: List[SavePlan], + save_to_lowest_rank: bool = False, +) -> List[SavePlan]: + """ + Removes duplicate entries from appearing on multiple SavePlans. For each duplicate across + a set of SavePlans, only the smallest SavePlan in terms of planned storage keeps the entry. + """ + + write_item_to_plan_indices: Dict[MetadataIndex, Set[int]] = defaultdict(set) + write_item_idx_to_write_item: Dict[MetadataIndex, WriteItem] = {} + for plan_idx, plan in enumerate(all_plans): + for write_item in plan.items: + # map each write item to its plan + write_item_to_plan_indices[write_item.index].add(plan_idx) + write_item_idx_to_write_item[write_item.index] = write_item + + # put item in the plan with the smallest size and remove it from the other plan_indices + to_remove: List[Set] = [set() for _ in range(len(all_plans))] + plan_to_size = [0] * len(all_plans) + for write_item_idx, plan_indices in write_item_to_plan_indices.items(): + if save_to_lowest_rank: + select_plan_idx = min(plan_indices) + else: + select_plan_idx = min( + plan_indices, key=lambda plan_idx: plan_to_size[plan_idx] + ) + + write_item = write_item_idx_to_write_item[write_item_idx] + # essentially ignores the storage size of anything that is not a tensor, since + # we don't know how much storage they represent + plan_to_size[select_plan_idx] += write_item.tensor_storage_size() or 1 + + plan_indices.remove(select_plan_idx) + for plan_idx in plan_indices: + to_remove[plan_idx].add(write_item_idx) + + for plan_idx, remove_set in enumerate(to_remove): + new_items = [ + write_item + for write_item in all_plans[plan_idx].items + if write_item.index not in remove_set + ] + all_plans[plan_idx] = dataclasses.replace(all_plans[plan_idx], items=new_items) + + return all_plans diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/_dedup_tensors.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/_dedup_tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..687afb287b3c74a4454085561528c31df3a5f702 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/_dedup_tensors.py @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates +import dataclasses +import logging +from typing import Dict, List, TYPE_CHECKING + +from torch.distributed.checkpoint.planner import SavePlan + + +if TYPE_CHECKING: + from torch.distributed.checkpoint.metadata import MetadataIndex + +__all__ = ["dedup_tensors"] + + +def init_logger() -> logging.Logger: + logger = logging.getLogger(__name__) + level = logging.INFO + logger.setLevel(level) + console = logging.StreamHandler() + formatter = logging.Formatter( + "%(asctime)s %(filename)s:%(lineno)s %(levelname)s p:%(processName)s t:%(threadName)s: %(message)s" + ) + console.setFormatter(formatter) + console.setLevel(level) + logger.addHandler(console) + logger.propagate = False + return logger + + +logger = init_logger() + + +# TODO add docstring for dedup_tensors +def dedup_tensors(all_plans: List[SavePlan]) -> List[SavePlan]: + all_plans = list(all_plans) + key_to_plan: Dict[MetadataIndex, List[int]] = {} + for plan_idx, plan in enumerate(all_plans): + for write_item in plan.items: + key_to_plan.setdefault(write_item.index, []).append(plan_idx) + + replicated_items = {k: v for k, v in key_to_plan.items() if len(v) > 1} + + # Remove duplicates by always keeping the first entry. + # Compute the per-rank remove set. + plan_to_keys: Dict[int, List[MetadataIndex]] = {} + for key, plans in replicated_items.items(): + for plan_idx in plans[1:]: + plan_to_keys.setdefault(plan_idx, []).append(key) + if len(plan_to_keys) > 0: + logger.info("Duplicate keys to remove: %s", plan_to_keys) + + for plan_idx, keys in plan_to_keys.items(): + key_set = set(keys) + # rewrite items and remove elements + new_items = [ + write_item + for write_item in all_plans[plan_idx].items + if write_item.index not in key_set + ] + all_plans[plan_idx] = dataclasses.replace(all_plans[plan_idx], items=new_items) + + return all_plans diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/_fsspec_filesystem.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/_fsspec_filesystem.py new file mode 100644 index 0000000000000000000000000000000000000000..8da02763066214eac8787c7bd9cc1ee6edfaf36a --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/_fsspec_filesystem.py @@ -0,0 +1,151 @@ +# Mypy will not try inferring the types of any 3rd party libraries installed. +# mypy: ignore-errors + +import io +import os +from contextlib import contextmanager +from pathlib import Path +from typing import Generator, Optional, TYPE_CHECKING, Union + +from fsspec.core import url_to_fs + +from torch.distributed.checkpoint.filesystem import ( + FileSystemBase, + FileSystemReader, + FileSystemWriter, +) + + +if TYPE_CHECKING: + from fsspec import AbstractFileSystem + + +__all__ = [ + "FsspecWriter", + "FsspecReader", +] + + +class FileSystem(FileSystemBase): + def __init__(self) -> None: + self.fs: Optional[AbstractFileSystem] = None + + @contextmanager + def create_stream( + self, path: Union[str, os.PathLike], mode: str + ) -> Generator[io.IOBase, None, None]: + assert self.fs is not None + path = os.fspath(path) + + # fsspec does not support concurrent transactions, and not all + # AbstractFileSystem have working rollback implementations, so + # just manually delete the file if necessary on errors. + with self.fs.open(path, mode) as stream: + try: + yield stream + except: # noqa: B001,E722 + if "w" or "+" or "a" in mode: # cleanup file if not read-only + try: + self.rm_file(path) + except: # noqa: B001,E722 + pass + raise + + def concat_path( + self, path: Union[str, os.PathLike], suffix: str + ) -> Union[str, os.PathLike]: + return os.path.join(path, suffix) + + def init_path(self, path: Union[str, os.PathLike]) -> Union[str, os.PathLike]: + self.fs, _ = url_to_fs(path) + return path + + def rename( + self, path: Union[str, os.PathLike], new_path: Union[str, os.PathLike] + ) -> None: + self.fs.rename(path, new_path) + + def mkdir(self, path: Union[str, os.PathLike]) -> None: + self.fs.makedirs(path, exist_ok=True) + + @classmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + if isinstance(checkpoint_id, Path): + return False + + try: + url_to_fs(checkpoint_id) + except ValueError: + return False + + return True + + def exists(self, path: Union[str, os.PathLike]) -> bool: + return self.fs.exists(path) + + def rm_file(self, path: Union[str, os.PathLike]) -> None: + self.fs.rm(path) + + +# TODO: add the dcp.async_save mixin +class FsspecWriter(FileSystemWriter): + """ + Basic implementation of StorageWriter using FFspec. + + This implementation makes the following assumptions and simplifications: + + * The checkpoint path is an empty or non-existing directory. + * File creation is atomic + + The checkpoint consist of one file per write request plus + a `.metadata` file with the serialized metadata. + + """ + + def __init__( + self, + path: Union[str, os.PathLike], + single_file_per_rank: bool = True, + sync_files: bool = True, + thread_count: int = 1, + per_thread_copy_ahead: int = 10_000_000, + overwrite: bool = True, + ) -> None: + """ + Initialize the writer pointing to `path`. + + Args: + path: directory where the checkpoint will be written to. + single_file_per_rank: Produce one file per rank instead of one file per tensor/blob. Default to True. + sync_files : force files to be synced to permanent storage. Default to True. + thread_count: Number of IO threads to use to write. Default to 1. + per_thread_copy_ahead: How many bytes to copy from the GPU ahead of saving then. Default 10Mb. + overwrite: Whether to allow overwriting existing checkpoints. Defaults to True. + + N. B. If sync_files is disabled, there's no guarantee that the checkpoint will be consistent in the case of a failure. + """ + super().__init__( + path, + single_file_per_rank, + sync_files, + thread_count, + per_thread_copy_ahead, + overwrite=overwrite, + ) + self.fs = FileSystem() + self.path = self.fs.init_path(path) + + @classmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + return FileSystem.validate_checkpoint_id(checkpoint_id) + + +class FsspecReader(FileSystemReader): + def __init__(self, path: Union[str, os.PathLike]) -> None: + super().__init__(path) + self.fs = FileSystem() + self.path = self.fs.init_path(path) + + @classmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + return FileSystem.validate_checkpoint_id(checkpoint_id) diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/_nested_dict.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/_nested_dict.py new file mode 100644 index 0000000000000000000000000000000000000000..b846a1d47d843433e1c3de295c86d318b7157a1e --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/_nested_dict.py @@ -0,0 +1,70 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates +from typing import Dict, Tuple + +from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE + +from . import _version +from ._traverse import ( + OBJ_PATH, + set_element, + STATE_DICT_ITEM, + traverse_state_dict, + traverse_state_dict_v_2_3, +) + + +""" +TODO: +Need to add ability to handle tuple, OrderedDict, NamedTuple. +Update mappings from dict to a class. +Change set_element to recreate the right type for tuple, OrderedDict, and NamedTuple. +""" + + +FLATTEN_MAPPING = Dict[str, OBJ_PATH] + + +# TODO: Update Docstring for nested_dict.py +def flatten_state_dict( + state_dict: STATE_DICT_TYPE, +) -> Tuple[STATE_DICT_TYPE, FLATTEN_MAPPING]: + """ + Flatten ``state_dict`` made of nested dicts and lists into a top level dictionary. + + Use ``unflatten_state_dict`` to revert this process. + Returns: + A tuple with the flatten state_dict and a mapping from original to new state_dict. + N.B. The new keys are derived from the object paths, joined by dot. + For example: ``{ 'a': {'b':...}}`` results in the key `a.b`. + """ + flattened: STATE_DICT_TYPE = {} + mappings: FLATTEN_MAPPING = {} + + def flat_copy(path: OBJ_PATH, value: STATE_DICT_ITEM) -> None: + new_fqn = ".".join(map(str, path)) + if new_fqn in flattened: + raise ValueError(f"duplicated flatten key {new_fqn}") + flattened[new_fqn] = value + mappings[new_fqn] = path + + # We started to flatten dictionary since v2.4. But in order to not break + # the checkpoints that were saved before v2.4, we need to keep the old + # traversal so that we can reconstruct those checkpoints. + use_v_2_3 = ( + _version._derived_version is not None and _version._derived_version == "2_3" + ) + if use_v_2_3: + traverse_state_dict_v_2_3(state_dict, flat_copy) + else: + traverse_state_dict(state_dict, flat_copy) + return flattened, mappings + + +def unflatten_state_dict( + state_dict: STATE_DICT_TYPE, mapping: FLATTEN_MAPPING +) -> STATE_DICT_TYPE: + """Restore the original nested state_dict according to ``mapping`` and the flattened ``state_dict``.""" + nested: STATE_DICT_TYPE = {} + for key, value in state_dict.items(): + set_element(nested, mapping[key], value) + return nested diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/_sharded_tensor_utils.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/_sharded_tensor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a68bcddeb7f9d9ffe6f89056dfe1ccc30cc12eb5 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/_sharded_tensor_utils.py @@ -0,0 +1,107 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates + +import copy +from typing import TYPE_CHECKING + +import torch.distributed as dist +from torch.distributed._shard.sharded_tensor import Shard, ShardedTensor, ShardMetadata +from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE +from torch.distributed.remote_device import _remote_device + +from ._traverse import OBJ_PATH, set_element, STATE_DICT_ITEM, traverse_state_dict +from .utils import _element_wise_add, _normalize_device_info + + +if TYPE_CHECKING: + from torch.distributed._shard.sharded_tensor.metadata import ShardedTensorMetadata + + +# TODO: We need to refactor this code. +def _flatten_sharded_tensors(state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE: + r""" + Transform ``state_dict`` by flattening all nested ShardedTensor instances found. + + The resulting ShardedTensor instances are only correct regarding the local shard and + MUST not be used for any other purpose but checkpointing, as no operator will work with them. + + This function should be used in conjunction with a state_dict produced by FSDP's + StateDictType.SHARDED_STATE_DICT methods. + """ + new_state_dict: STATE_DICT_TYPE = {} + + def rewrite_dict(path: OBJ_PATH, value: STATE_DICT_ITEM) -> None: + if not isinstance(value, ShardedTensor): + set_element(new_state_dict, path, value) + return + shards = value.local_shards() + + if len(shards) == 0: + return + if len(shards) != 1: + set_element(new_state_dict, path, value) + return + + outer_shard = shards[0] + + inner_st = outer_shard.tensor + if not isinstance(inner_st, ShardedTensor): + set_element(new_state_dict, path, value) + return + + if len(inner_st.local_shards()) != 1: + raise ValueError("Cannot handle inner tensor with more than 1 shard") + inner_shard = inner_st.local_shards()[0] + + local_shards = [ + Shard( + tensor=inner_shard.tensor, + metadata=ShardMetadata( + shard_offsets=_element_wise_add( + outer_shard.metadata.shard_offsets, + inner_shard.metadata.shard_offsets, + ), + shard_sizes=inner_shard.metadata.shard_sizes, + placement=f"rank:{dist.get_rank()}/{inner_shard.tensor.device}", + ), + ) + ] + + st_meta: ShardedTensorMetadata = copy.deepcopy(value.metadata()) + other_rank = 0 if dist.get_rank() > 0 else 1 + device_info = _normalize_device_info(inner_shard.tensor.device.type, 0) + + # Remove the outer ST shard the inner ST covers + for i, shard_md in enumerate(st_meta.shards_metadata): + if shard_md.shard_offsets == outer_shard.metadata.shard_offsets: + st_meta.shards_metadata.pop(i) + break + + # Attribute other rank for the other shards + for shard_md in st_meta.shards_metadata: + shard_md.placement = _remote_device(f"rank:{other_rank}/{device_info}") + + # Add other inner shards from the inner tensor + for inner_md in inner_st.metadata().shards_metadata: + if inner_md.shard_offsets != inner_shard.metadata.shard_offsets: + st_meta.shards_metadata.append( + ShardMetadata( + shard_offsets=_element_wise_add( + outer_shard.metadata.shard_offsets, + inner_md.shard_offsets, + ), + shard_sizes=inner_md.shard_sizes, + placement=f"rank:{other_rank}/{device_info}", + ) + ) + + # Finally add this shard + st_meta.shards_metadata.append(local_shards[0].metadata) + + st = ShardedTensor._init_from_local_shards_and_global_metadata( + local_shards=local_shards, + sharded_tensor_metadata=st_meta, + ) + set_element(new_state_dict, path, st) + + traverse_state_dict(state_dict, rewrite_dict) + return new_state_dict diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/_storage_utils.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/_storage_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..194c9c8c4b9b15eb1b2909cfd8a9e8e01e46f54d --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/_storage_utils.py @@ -0,0 +1,49 @@ +import os +from typing import List, Type, Union + +from .filesystem import FileSystemReader, FileSystemWriter +from .storage import StorageReader, StorageWriter + + +def _storage_setup( + storage: Union[StorageReader, StorageWriter, None], + checkpoint_id: Union[str, os.PathLike, None], + reader: bool = False, +) -> Union[None, StorageReader, StorageWriter]: + if storage: + if checkpoint_id is not None: + storage.reset(checkpoint_id) + return storage + + if not checkpoint_id: + raise RuntimeError( + "`checkpoint_id` must be specificed if " + "storage_reader/storage_writer is None." + ) + + targets: List[Type[Union[StorageReader, StorageWriter]]] = [] + if reader: + targets = [ + FileSystemReader, + ] + else: + targets = [ + FileSystemWriter, + ] + try: + from ._fsspec_filesystem import FsspecReader, FsspecWriter + + targets.append(FsspecReader if reader else FsspecWriter) + except Exception: + pass + + for target in targets: + if target.validate_checkpoint_id(checkpoint_id): + storage = target(checkpoint_id) # type: ignore[call-arg] + storage.reset(checkpoint_id) + return storage + + raise RuntimeError( + "Cannot detect which StorageReader or StorageWriter to use. " + "Please specify the storage_reader/storage_writer." + ) diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/_traverse.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/_traverse.py new file mode 100644 index 0000000000000000000000000000000000000000..2c79013abeb2692a7359e77df0ede0237125bdb4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/_traverse.py @@ -0,0 +1,208 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates +from typing import ( + Callable, + cast, + Collection, + List, + Mapping, + MutableMapping, + Optional, + Tuple, + TypeVar, + Union, +) + +import torch +from torch.distributed._shard.sharded_tensor.api import ShardedTensor +from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE +from torch.distributed.tensor import DTensor + + +PATH_ITEM = Union[str, int] +OBJ_PATH = Tuple[PATH_ITEM, ...] +T = TypeVar("T") + +STATE_DICT_ITEM = object +CONTAINER_TYPE = MutableMapping[PATH_ITEM, STATE_DICT_ITEM] + +__all__ = ["traverse_state_dict", "set_element", "get_element", "print_tensor"] + + +def _keep_visiting_tensors(value: STATE_DICT_ITEM) -> bool: + return isinstance(value, torch.Tensor) + + +# TODO: update docstring for traverse.py +def traverse_state_dict( + state_dict: STATE_DICT_TYPE, + visitor: Callable[[OBJ_PATH, STATE_DICT_ITEM], None], + keep_traversing: Callable[[STATE_DICT_ITEM], bool] = _keep_visiting_tensors, +) -> None: + """ + Invoke ``visitor`` for each value recursively in ``state_dict``. + Mapping will be traversed and ``visitor`` will be applied to the leaf elements. + ``visitor`` will only be applied to elements in a list or a tuple, if the + container contains tensors or mappings. + """ + + def _is_terminal(value: STATE_DICT_ITEM) -> bool: + values: Collection[STATE_DICT_ITEM] + if isinstance(value, Mapping): + return False + elif isinstance(value, list): + values = value + else: + return True + + for entry in values: + if isinstance(entry, (Mapping, list)) and not _is_terminal(entry): + return False + if keep_traversing is not None and keep_traversing(entry): + return False + return True + + def _traverse_obj(path: OBJ_PATH, value: STATE_DICT_ITEM) -> None: + if isinstance(value, Mapping): + for k, v in value.items(): + _traverse_obj(path + (str(k),), v) + elif _is_terminal(value): + visitor(path, value) + elif isinstance(value, (list, tuple)): + for i, v in enumerate(value): + _traverse_obj(path + (i,), v) + + for key, value in state_dict.items(): + _traverse_obj((str(key),), value) + + +def traverse_state_dict_v_2_3( + state_dict: STATE_DICT_TYPE, + visitor: Callable[[OBJ_PATH, STATE_DICT_ITEM], None], + keep_traversing: Callable[[STATE_DICT_ITEM], bool] = _keep_visiting_tensors, +) -> None: + """ + Traversal is short-circuited when if finds a collection for which ``keep_visiting_tensors`` evaluates + to false for all elements. + By default, all collections with at least one ``torch.Tensor`` element are traversed. + Visitor takes a path argument that is a tuple of the keys used to reach it. + """ + + # a value is terminal if it has no other containers values inside it + def _is_terminal(value: STATE_DICT_ITEM) -> bool: + values: Collection[STATE_DICT_ITEM] + if isinstance(value, Mapping): + values = value.values() + elif isinstance(value, list): + values = value + else: + return True + + for entry in values: + if isinstance(entry, (Mapping, list)) and not _is_terminal(entry): + return False + if keep_traversing is not None and keep_traversing(entry): + return False + return True + + def _traverse_obj(path: OBJ_PATH, value: STATE_DICT_ITEM) -> None: + if _is_terminal(value): + visitor(path, value) + elif isinstance(value, Mapping): + for k, v in value.items(): + _traverse_obj(path + (str(k),), v) + elif isinstance(value, list): + for i, v in enumerate(value): + _traverse_obj(path + (i,), v) + + for key, value in state_dict.items(): + _traverse_obj((str(key),), value) + + +def set_element( + root_dict: STATE_DICT_TYPE, path: OBJ_PATH, value: STATE_DICT_ITEM +) -> None: + """Set ``value`` in ``root_dict`` along the ``path`` object path.""" + cur_container = cast(CONTAINER_TYPE, root_dict) + + def extend_list(lst: List[STATE_DICT_ITEM], idx: int) -> None: + while len(lst) <= idx: + lst.append(None) + + for i in range(1, len(path)): + prev_key = path[i - 1] + key = path[i] + def_val = cast(STATE_DICT_ITEM, {} if type(key) == str else []) + + if isinstance(cur_container, Mapping): + cur_container = cast( + CONTAINER_TYPE, cur_container.setdefault(prev_key, def_val) + ) + else: + extend_list(cur_container, prev_key) + if cur_container[prev_key] is None: + cur_container[prev_key] = def_val + cur_container = cur_container[prev_key] + + key = path[-1] + if type(key) == int: + extend_list(cast(List[STATE_DICT_ITEM], cur_container), key) + + cur_container[key] = value + + +def get_element( + root_dict: STATE_DICT_TYPE, + path: OBJ_PATH, + default_value: Optional[T] = None, +) -> Optional[T]: + """Retrieve the value at ``path``from ``root_dict``, returning ``default_value`` if not found.""" + cur_value = cast(CONTAINER_TYPE, root_dict) + for part in path: + if type(part) is int: + if not isinstance(cur_value, list) or len(cur_value) < part: + return default_value + elif not isinstance(cur_value, Mapping) or part not in cur_value: + return default_value + + cur_value = cast(CONTAINER_TYPE, cur_value[part]) + return cast(Optional[T], cur_value) + + +def _print_nested( + value: STATE_DICT_ITEM, + prefix: str = "", + print_fun: Callable[[str], None] = print, +) -> None: + if type(value) is ShardedTensor: + print_fun(f"{prefix} ShardedTensor size: {value.size()}") + for shard in value.local_shards(): + _print_nested( + shard.tensor, + f"{shard.metadata.shard_offsets} ", + print_fun=print_fun, + ) + elif type(value) is (DTensor): + print_fun(f"{prefix} DistributedTensor size: {value.size()}") + # TODO: add local offset for _local_tensor in print_nested. + _print_nested( + value._local_tensor, + print_fun=print_fun, + ) + elif isinstance(value, torch.Tensor): + print_fun(f"{prefix} Tensor size: {value.size()}") + else: + print_fun(f"{prefix} Type: {type(value)}") + + +def print_tensor( + path: OBJ_PATH, + value: STATE_DICT_ITEM, + print_fun: Callable[[str], None] = print, +) -> None: + """ + Use this callback with traverse_state_dict to print its content. + + By default the content is printed using the builtin ``print`` but this can + be change by passing a different ``print_fun` callable. + """ + _print_nested(value, prefix=str(path), print_fun=print_fun) diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/_version.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/_version.py new file mode 100644 index 0000000000000000000000000000000000000000..b3065bdfd6a2c141a959ef0ffe30aeafdc2dc54f --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/_version.py @@ -0,0 +1,6 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates + +from typing import Optional + + +_derived_version: Optional[str] = None diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/api.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/api.py new file mode 100644 index 0000000000000000000000000000000000000000..e587580617a1bff43f4d50861dd554595b371c3e --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/api.py @@ -0,0 +1,43 @@ +# mypy: allow-untyped-defs +import traceback as tb +from typing import Any, Dict, Tuple + + +WRAPPED_EXCEPTION = Tuple[BaseException, tb.StackSummary] + +__all__ = ["CheckpointException"] + + +def _wrap_exception(exc: BaseException) -> WRAPPED_EXCEPTION: + return (exc, tb.extract_tb(exc.__traceback__)) + + +def _is_wrapped_exception(obj: Any) -> bool: + if not isinstance(obj, tuple): + return False + if len(obj) != 2: + return False + return isinstance(obj[0], BaseException) and isinstance(obj[1], tb.StackSummary) + + +class CheckpointException(BaseException): + """Exception raised if failure was detected as part of a checkpoint load or save.""" + + def __init__(self, msg: str, failures: Dict[int, WRAPPED_EXCEPTION]): + super().__init__(msg, failures) + self._failures = failures + + @property + def failures(self) -> Dict[int, WRAPPED_EXCEPTION]: + """Return a dictionary mapping node ranks to their associated exceptions in case of failure.""" + return self._failures + + def __str__(self): + str = f"CheckpointException ranks:{self._failures.keys()}\n" + for rank, exc_pair in self._failures.items(): + exc, trace = exc_pair + str += f"Traceback (most recent call last): (RANK {rank})\n" + if trace is not None: + str += "".join(tb.format_list(trace)) + str += "".join(tb.format_exception_only(type(exc), value=exc)) + return str diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/filesystem.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/filesystem.py new file mode 100644 index 0000000000000000000000000000000000000000..079608c7230629e33346c0ed10ff80a915bee016 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/filesystem.py @@ -0,0 +1,769 @@ +# mypy: allow-untyped-defs +import collections +import dataclasses +import io +import operator +import os +import pickle +import queue +import threading +import uuid +import warnings +from abc import ABC, abstractmethod +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import ( + Any, + Callable, + cast, + Dict, + Generator, + IO, + Iterable, + Iterator, + List, + Optional, + Tuple, + Union, +) + +import torch +from torch import Tensor +from torch._utils import _get_available_device_type, _get_device_module +from torch.distributed._shard._utils import narrow_tensor_by_index +from torch.distributed.checkpoint.metadata import ( + Metadata, + MetadataIndex, + STATE_DICT_TYPE, + StorageMeta, +) +from torch.distributed.checkpoint.planner import ( + LoadItemType, + LoadPlan, + LoadPlanner, + ReadItem, + SavePlan, + SavePlanner, + WriteItem, + WriteItemType, +) +from torch.distributed.checkpoint.staging import BlockingAsyncStager +from torch.distributed.checkpoint.storage import ( + StorageReader, + StorageWriter, + WriteResult, +) +from torch.distributed.checkpoint.utils import _create_file_view +from torch.futures import Future + + +__all__ = ["FileSystemWriter", "FileSystemReader", "FileSystem", "FileSystemBase"] + +_metadata_fn: str = ".metadata" + + +@dataclass +class _StorageInfo: + """This is the per entry storage info.""" + + relative_path: str + offset: int + length: int + + +@dataclass +class _StoragePrefix: + prefix: str + + +DEFAULT_SUFFIX = ".distcp" + + +def _generate_uuid() -> str: + return str(uuid.uuid4()) + + +class _TensorLoader(ABC): + @abstractmethod + def add(self, size: int, obj: object) -> None: + pass + + @abstractmethod + def start_loading(self) -> None: + pass + + @abstractmethod + def values(self) -> Iterator[Tuple[torch.Tensor, object]]: + pass + + +class _SerialCpuLoader(_TensorLoader): + def __init__(self, resolve_fun: Callable) -> None: + self.resolve_fun = resolve_fun + self.items: List[Tuple[int, object]] = [] + + def add(self, size: int, obj: object) -> None: + self.items.append((size, obj)) + + def start_loading(self) -> None: + pass + + def values(self) -> Iterator[Tuple[torch.Tensor, object]]: + for _, obj in self.items: + tensor = self.resolve_fun(obj).detach() + tensor = tensor.cpu() + if tensor.storage().size() != tensor.numel(): + tensor = tensor.clone() + yield ( + tensor, + obj, + ) + + +class _OverlappingCpuLoader(_TensorLoader): + def __init__( + self, + resolve_fun: Callable, + stream: Optional[torch.Stream] = None, + inflight_threshhold: int = 1_000_000, + ) -> None: + self.resolve_fun = resolve_fun + self.items: List[Tuple[int, object]] = [] + self.inflight_threshhold = inflight_threshhold + self.in_flight_data = 0 + self.current_items: collections.deque = collections.deque() + self.idx = 0 + self.started = False + self.device_type = ( + stream.device_type if stream else _get_available_device_type() + ) + self.device_module = _get_device_module(self.device_type) + self.stream = cast( + torch.cuda.Stream, stream or self.device_module.current_stream() + ) + if self.stream != self.device_module.current_stream(): + self.stream.wait_stream(self.device_module.current_stream()) + + @property + def _done(self) -> bool: + return self.idx >= len(self.items) + + def _drain(self) -> List[Tuple[torch.Tensor, object]]: + drained = [] + if self.in_flight_data >= self.inflight_threshhold: + self.stream.synchronize() + while self.in_flight_data >= self.inflight_threshhold: + val = self.current_items.popleft() + self.in_flight_data -= val[0].numel() * val[0].element_size() + drained.append(val) + return drained + + def _refill(self) -> None: + with self.device_module.stream(self.stream): + while not self._done and self.in_flight_data < self.inflight_threshhold: + _, obj = self.items[self.idx] + self.idx += 1 + tensor = self.resolve_fun(obj).detach() + if tensor.device.type == self.device_type: + tensor = tensor.to(device="cpu", non_blocking=True) + elif tensor.device == torch.device("cpu"): + if ( + tensor.untyped_storage().size() + != tensor.numel() * tensor.itemsize + ): + # this forces the tensor to be both contiguous and with minimal storage + tensor = tensor.clone() + + self.current_items.append( + ( + tensor, + obj, + ) + ) + self.in_flight_data += tensor.numel() * tensor.element_size() + + def _finish(self) -> Iterable[Tuple[torch.Tensor, object]]: + assert self._done + if len(self.current_items) > 0: + self.stream.synchronize() + return self.current_items + + def add(self, size: int, obj: object) -> None: + if self.started: + raise RuntimeError("cannot add items after loading started") + self.items.append((size, obj)) + + def start_loading(self) -> None: + if self.started: + return + self.started = True + self.items.sort(key=operator.itemgetter(0)) + self._refill() + + def values(self) -> Iterator[Tuple[torch.Tensor, object]]: + self.start_loading() + while not self._done: + drained = self._drain() + self._refill() + yield from drained + + yield from self._finish() + + +def _item_size(item: WriteItem) -> int: + size = 1 + assert item.tensor_data is not None + # can't use math.prod as PT needs to support older python + for s in item.tensor_data.size: + size *= s + + dtype = item.tensor_data.properties.dtype + return size * torch._utils._element_size(dtype) + + +def _split_by_size_and_type(bins: int, items: List[WriteItem]) -> List[List[WriteItem]]: + if bins == 1: + return [items] + + bytes_w = [wi for wi in items if wi.type == WriteItemType.BYTE_IO] + tensor_w = [wi for wi in items if wi.type != WriteItemType.BYTE_IO] + + buckets: List[List[WriteItem]] = [[] for _ in range(bins)] + bucket_sizes = [0 for _ in range(bins)] + + tensor_w.sort(key=_item_size, reverse=True) + + for i, wi in enumerate(bytes_w): + buckets[i % bins].append(wi) + + for wi in tensor_w: + # TODO replace with headq + idx = min(enumerate(bucket_sizes), key=operator.itemgetter(1))[0] + buckets[idx].append(wi) + bucket_sizes[idx] += _item_size(wi) + + return buckets + + +def _write_item( + stream: io.IOBase, + data: Union[io.BytesIO, torch.Tensor], + write_item: WriteItem, + storage_key: str, +) -> WriteResult: + offset = stream.tell() + + if write_item.type == WriteItemType.BYTE_IO: + assert isinstance(data, io.BytesIO) + stream.write(data.getbuffer()) + else: + assert isinstance(data, torch.Tensor) + assert data.device == torch.device("cpu") + torch.save(data, cast(IO[bytes], stream)) + length = stream.tell() - offset + + return WriteResult( + index=write_item.index, + size_in_bytes=length, + storage_data=_StorageInfo(storage_key, offset, length), + ) + + +def _write_files_from_queue( + create_stream: Callable, + file_queue: queue.Queue, + result_queue: queue.Queue, + planner: SavePlanner, + inflight_threshhold: int, + use_fsync: bool, + thread_count: int, +) -> None: + try: + while True: + file_name, storage_key, write_items = file_queue.get_nowait() + loader: _TensorLoader + + custom_backend_name = torch._C._get_privateuse1_backend_name() + custom_device_mod = getattr(torch, custom_backend_name, None) + + # TODO: Using the OverlappingCpuLoader with multiple threads creates significant + # performance degredation, observed as being related to cuda stream syncs. We + # should try to fix this and use _OverlappingCpuLoader for all threaded cases + if ( + thread_count == 1 + and ( + torch.cuda.is_available() + or (custom_device_mod and custom_device_mod.is_available()) + ) + and inflight_threshhold > 0 + ): + loader = _OverlappingCpuLoader( + planner.resolve_data, + inflight_threshhold=inflight_threshhold, + ) + else: + loader = _SerialCpuLoader( + planner.resolve_data, + ) + + tensor_w = [wi for wi in write_items if wi.type != WriteItemType.BYTE_IO] + for write_item in tensor_w: + loader.add(_item_size(write_item), write_item) + loader.start_loading() + + bytes_w = [wi for wi in write_items if wi.type == WriteItemType.BYTE_IO] + write_results = [] + + with create_stream(file_name, "wb") as stream: + for write_item in bytes_w: + data = planner.resolve_data(write_item) + write_results.append( + _write_item(stream, data, write_item, storage_key) + ) + + for tensor, write_item in loader.values(): + assert tensor.is_cpu + write_results.append( + _write_item(stream, tensor, write_item, storage_key) + ) + + if use_fsync: + try: + os.fsync(stream.fileno()) + except AttributeError: + os.sync() + result_queue.put(write_results) + except queue.Empty: + pass + + +class FileSystemBase(ABC): + @contextmanager + @abstractmethod + def create_stream( + self, path: Union[str, os.PathLike], mode: str + ) -> Generator[io.IOBase, None, None]: + ... + + @abstractmethod + def concat_path( + self, path: Union[str, os.PathLike], suffix: str + ) -> Union[str, os.PathLike]: + ... + + @abstractmethod + def rename( + self, path: Union[str, os.PathLike], new_path: Union[str, os.PathLike] + ) -> None: + ... + + @abstractmethod + def init_path(self, path: Union[str, os.PathLike]) -> Union[str, os.PathLike]: + ... + + @abstractmethod + def mkdir(self, path: Union[str, os.PathLike]) -> None: + ... + + @classmethod + @abstractmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + ... + + @abstractmethod + def exists(self, path: Union[str, os.PathLike]) -> bool: + ... + + @abstractmethod + def rm_file(self, path: Union[str, os.PathLike]) -> None: + ... + + +class FileSystem(FileSystemBase): + @contextmanager + def create_stream( + self, path: Union[str, os.PathLike], mode: str + ) -> Generator[io.IOBase, None, None]: + with cast(Path, path).open(mode) as stream: + yield cast(io.IOBase, stream) + + def concat_path( + self, path: Union[str, os.PathLike], suffix: str + ) -> Union[str, os.PathLike]: + return cast(Path, path) / suffix + + def init_path(self, path: Union[str, os.PathLike]) -> Union[str, os.PathLike]: + if not isinstance(path, Path): + path = Path(path) + return path + + def rename( + self, path: Union[str, os.PathLike], new_path: Union[str, os.PathLike] + ) -> None: + cast(Path, path).rename(cast(Path, new_path)) + + def mkdir(self, path: Union[str, os.PathLike]) -> None: + cast(Path, path).mkdir(parents=True, exist_ok=True) + + @classmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + if isinstance(checkpoint_id, Path): + return True + + if "://" in str(checkpoint_id): + return False + + for p in Path(checkpoint_id).parents: + if p.exists() and os.access(str(p), os.W_OK): + return True + + return False + + def exists(self, path: Union[str, os.PathLike]) -> bool: + return cast(Path, path).exists() + + def rm_file(self, path: Union[str, os.PathLike]) -> None: + cast(Path, path).unlink() + + +class _FileSystemWriter(StorageWriter): + """ + Basic implementation of StorageWriter using file IO. + + This implementation makes the following assumptions and simplifications: + + * The checkpoint path is an empty or non-existing directory. + * File creation is atomic + + The checkpoint consist of one file per write request plus + a `.metadata` file with the serialized metadata. + + """ + + def __init__( + self, + path: Union[str, os.PathLike], + single_file_per_rank: bool = True, + sync_files: bool = True, + thread_count: int = 1, + per_thread_copy_ahead: int = 10_000_000, + overwrite: bool = True, + *args: Any, + **kwargs: Any, + ) -> None: + """ + Initialize the writer pointing to `path`. + + Args: + path: directory where the checkpoint will be written to. + single_file_per_rank: Produce one file per rank instead of one file per tensor/blob. Default to True. + sync_files : force files to be synced to permanent storage. Default to True. + thread_count: Number of IO threads to use to write. Default to 1. + per_thread_copy_ahead: How many bytes to copy from the GPU ahead of saving then. Default 10Mb. + overwrite: Whether to allow overwriting existing checkpoints. Defaults to True. + + N. B. If sync_files is disabled, there's no guarantee that the checkpoint will be consistent in the case of a failure. + """ + super().__init__() + self.fs = FileSystem() + self.path = self.fs.init_path(path) + self.single_file_per_rank = single_file_per_rank + self.sync_files = sync_files + self.thread_count = thread_count + self.per_thread_copy_ahead = per_thread_copy_ahead + self.save_id = _generate_uuid() + self.overwrite = overwrite + + def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None: + if checkpoint_id: + self.path = self.fs.init_path(checkpoint_id) + self.save_id = _generate_uuid() + + def set_up_storage_writer(self, is_coordinator: bool) -> None: + pass + + def prepare_local_plan(self, plan: SavePlan) -> SavePlan: + self.fs.mkdir(self.path) + if self.fs.exists(self.metadata_path): + if self.overwrite: + warnings.warn( + f"Detected an existing checkpoint in {self.metadata_path}, overwriting since {self.overwrite=}." + " Past version 2.5 of PyTorch, `overwrite` will default to False. Set this variable to True to" + " maintain this functionality or False to raise when an existing checkpoint is found." + ) + else: + raise RuntimeError(f"Checkpoint already exists and {self.overwrite=}.") + + return plan + + def prepare_global_plan(self, plans: List[SavePlan]) -> List[SavePlan]: + new_plans = [ + dataclasses.replace(plan, storage_data=_StoragePrefix(f"__{i}_")) + for i, plan in enumerate(plans) + ] + return new_plans + + def write_data( + self, + plan: SavePlan, + planner: SavePlanner, + ) -> Future[List[WriteResult]]: + storage_plan: _StoragePrefix = plan.storage_data + file_count = 0 + + def gen_file(): + nonlocal file_count + file_name = f"{storage_plan.prefix}{file_count}{DEFAULT_SUFFIX}" + file_count += 1 + return file_name + + file_queue: queue.Queue = queue.Queue() + if self.single_file_per_rank: + for bucket in _split_by_size_and_type(self.thread_count, plan.items): + file_name = gen_file() + path = self.fs.concat_path(self.path, file_name) + file_queue.put((path, file_name, bucket)) + else: + for item in plan.items: + file_name = gen_file() + path = self.fs.concat_path(self.path, file_name) + file_queue.put((path, file_name, [item])) + + result_queue: queue.Queue = queue.Queue() + + threads = [] + for _ in range(1, self.thread_count): + t = threading.Thread( + target=_write_files_from_queue, + args=( + self.fs.create_stream, + file_queue, + result_queue, + planner, + self.per_thread_copy_ahead, + self.sync_files, + self.thread_count, + ), + ) + t.start() + threads.append(t) + + _write_files_from_queue( + create_stream=self.fs.create_stream, + file_queue=file_queue, + result_queue=result_queue, + planner=planner, + inflight_threshhold=self.per_thread_copy_ahead, + use_fsync=self.sync_files, + thread_count=self.thread_count, + ) + + for t in threads: + t.join() + + res = [] + try: + while True: + res += result_queue.get_nowait() + except queue.Empty: + fut: Future[List[WriteResult]] = Future() + fut.set_result(res) + return fut + + def finish(self, metadata: Metadata, results: List[List[WriteResult]]) -> None: + storage_md = {} + for wr_list in results: + storage_md.update({wr.index: wr.storage_data for wr in wr_list}) + metadata.storage_data = storage_md + + metadata.storage_meta = self.storage_meta() + + tmp_path = cast(Path, self.fs.concat_path(self.path, f"{_metadata_fn}.tmp")) + with self.fs.create_stream(tmp_path, "wb") as metadata_file: + pickle.dump(metadata, metadata_file) + if self.sync_files: + try: + os.fsync(metadata_file.fileno()) + except AttributeError: + os.sync() + + # delete in-case other checkpoints were present. + if self.fs.exists(self.metadata_path): + self.fs.rm_file(self.metadata_path) + + self.fs.rename(tmp_path, self.metadata_path) + + def storage_meta(self) -> Optional[StorageMeta]: + return StorageMeta(checkpoint_id=self.checkpoint_id, save_id=self.save_id) + + @property + def metadata_path(self) -> Union[str, os.PathLike]: + return cast(Path, self.fs.concat_path(self.path, _metadata_fn)) + + @property + def checkpoint_id(self) -> Union[str, os.PathLike]: + """ + return the checkpoint_id that will be used to save the checkpoint. + """ + return self.path + + @classmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + return FileSystem.validate_checkpoint_id(checkpoint_id) + + +class FileSystemReader(StorageReader): + def __init__(self, path: Union[str, os.PathLike]) -> None: + super().__init__() + self.fs = FileSystem() + self.path = self.fs.init_path(path) + self.storage_data: Dict[MetadataIndex, _StorageInfo] = {} + self.load_id = _generate_uuid() + + def _slice_file(self, file, sinfo: _StorageInfo) -> io.IOBase: + return _create_file_view(file, sinfo.offset, sinfo.length) + + def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None: + self.storage_data = {} + if checkpoint_id: + self.path = self.fs.init_path(checkpoint_id) + self.load_id = _generate_uuid() + + def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]: + # group requests by file + per_file: Dict[str, List[ReadItem]] = {} + for read_item in plan.items: + item_md = self.storage_data[read_item.storage_index] + path = item_md.relative_path + per_file.setdefault(path, []).append(read_item) + + for relative_path, reqs in per_file.items(): + new_path = self.fs.concat_path(self.path, relative_path) + with self.fs.create_stream(new_path, "rb") as stream: + # TODO sort by offset and cache the reading + for req in reqs: + item_md = self.storage_data[req.storage_index] + file_slice = self._slice_file(stream, item_md) + if req.type == LoadItemType.BYTE_IO: + read_bytes = io.BytesIO(file_slice.read(item_md.length)) + read_bytes.seek(0) + planner.load_bytes(req, read_bytes) + else: + tensor = cast( + Tensor, + torch.load( + cast(IO[bytes], file_slice), + map_location="cpu", + weights_only=True, + ), + ) + tensor = narrow_tensor_by_index( + tensor, req.storage_offsets, req.lengths + ) + target_tensor = planner.resolve_tensor(req).detach() + + assert ( + target_tensor.size() == tensor.size() + ), f"req {req.storage_index} mismatch sizes {target_tensor.size()} vs {tensor.size()}" + target_tensor.copy_(tensor) + planner.commit_tensor(req, target_tensor) + + fut: Future = Future() + fut.set_result(None) + return fut + + # Implementing the abstract function in StorageReader + def read_metadata(self) -> Metadata: + path = self.fs.concat_path(self.path, ".metadata") + with self.fs.create_stream(path, "rb") as metadata_file: + metadata = pickle.load(metadata_file) + + if getattr(metadata, "storage_meta", None) is None: + metadata.storage_meta = StorageMeta() + metadata.storage_meta.load_id = self.load_id + + return metadata + + def set_up_storage_reader(self, metadata: Metadata, is_coordinator: bool) -> None: + self.storage_data = metadata.storage_data + assert self.storage_data is not None + + def prepare_local_plan(self, plan: LoadPlan) -> LoadPlan: + return plan + + def prepare_global_plan(self, plans: List[LoadPlan]) -> List[LoadPlan]: + return plans + + @property + def checkpoint_id(self) -> Union[str, os.PathLike]: + """ + return the checkpoint_id that will be used to load the checkpoint. + """ + return self.path + + @classmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + return FileSystem.validate_checkpoint_id(checkpoint_id) + + +class FileSystemWriter(_FileSystemWriter, BlockingAsyncStager): + """ + Basic implementation of StorageWriter using file IO. + + This implementation makes the following assumptions and simplifications: + + * The checkpoint path is an empty or non-existing directory. + * File creation is atomic + + The checkpoint consist of one file per write request plus + a `.metadata` file with the serialized metadata. + + """ + + def __init__( + self, + path: Union[str, os.PathLike], + single_file_per_rank: bool = True, + sync_files: bool = True, + thread_count: int = 1, + per_thread_copy_ahead: int = 10_000_000, + cache_staged_state_dict: bool = False, + overwrite: bool = True, + ) -> None: + """ + Initialize the writer pointing to `path`. + + Args: + path: directory where the checkpoint will be written to. + single_file_per_rank: Produce one file per rank instead of one file per tensor/blob. Default to True. + sync_files : force files to be synced to permanent storage. Default to True. + thread_count: Number of IO threads to use to write. Default to 1. + per_thread_copy_ahead: How many bytes to copy from the GPU ahead of saving then. Default 10Mb. + cache_staged_state_dict: Whether to cache the staged state_dict. This option decreases staging latency + at the cost of increases memory usage. Additionally, if this parameter is set to True, it's the expectation + that the stager is maintained and re-used for multiple dcp.async_save calls. Default to False. + overwrite: Whether to allow overwriting existing checkpoints. Defaults to True. + + N. B. If sync_files is disabled, there's no guarantee that the checkpoint will be consistent in the case of a failure. + """ + _FileSystemWriter.__init__( + self, + path=path, + single_file_per_rank=single_file_per_rank, + sync_files=sync_files, + thread_count=thread_count, + per_thread_copy_ahead=per_thread_copy_ahead, + overwrite=overwrite, + ) + BlockingAsyncStager.__init__( + self, + cache_staged_state_dict=cache_staged_state_dict, + ) + + def stage(self, state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE: + """Override of AsyncStager.stage""" + # in the async case, the state dict is already on CPU, so maintaining this + # buffer makes no sense + self.per_thread_copy_ahead = 0 + return super().stage(state_dict) diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/format_utils.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/format_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ec97504e79b58f10771654eff0c04d7f9690cd46 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/format_utils.py @@ -0,0 +1,280 @@ +# mypy: allow-untyped-defs +import argparse +import os +from enum import Enum +from typing import cast, Dict, List, Optional, Union + +import torch +import torch.distributed as dist +from torch.distributed._shard._utils import narrow_tensor_by_index +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter +from torch.distributed.checkpoint._nested_dict import flatten_state_dict +from torch.distributed.checkpoint.default_planner import ( + _EmptyStateDictLoadPlanner, + DefaultLoadPlanner, +) +from torch.distributed.checkpoint.metadata import ( + Metadata, + STATE_DICT_TYPE, + STORAGE_TYPES, + TensorProperties, + TensorStorageMetadata, +) +from torch.distributed.checkpoint.planner import LoadItemType, LoadPlan, LoadPlanner +from torch.distributed.checkpoint.planner_helpers import _create_chunk_list +from torch.distributed.checkpoint.state_dict_loader import _load_state_dict +from torch.distributed.checkpoint.state_dict_saver import _save_state_dict +from torch.distributed.checkpoint.storage import StorageReader +from torch.futures import Future + + +__all__ = [ + "dcp_to_torch_save", + "torch_save_to_dcp", + "BroadcastingTorchSaveReader", + "DynamicMetaLoadPlanner", +] + + +class BroadcastingTorchSaveReader(StorageReader): + """ + StorageReader for reading a Torch Save file. This reader will read the entire checkpoint + on the coordinator rank, and then broadcast and shard each tensor to all ranks. + + . N.B. Intended to be used with DynamicMetaLoadPlanner + + .. warning:: + Current implementation only supports loading Tensors. + + >>> # xdoctest: +SKIP("undefined vars") + >>> sd = {"mode": model} + >>> dcp.load( + >>> sd, + >>> storage_reader=BroadcastingTorchSaveReader(), + >>> planner=DynamicMetaLoadPlanner(), + >>> checkpoint_id="path_to_model.pt" + >>> ) + """ + + def __init__( + self, + checkpoint_id: Optional[Union[str, os.PathLike]] = None, + coordinator_rank: int = 0, + ) -> None: + self.checkpoint_id = checkpoint_id + self.coordinator_rank = coordinator_rank + + def read_metadata(self) -> Metadata: + """Extends the default StorageReader to support building the metadata file""" + # Metadata is built in planner.set_up_planner, since we are not actually reading metadata from + # the disk + return Metadata(state_dict_metadata={}) + + def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]: + """ + Reads torch save data on the coordinator rank, and broadcast afterwards + this incurrs a communication cost, but avoids having to load + the entire checkpoint on each rank, hopefully preventing OOM issues + """ + planner = cast(DefaultLoadPlanner, planner) + + # data is read in on the coordinator rank, and broadcast afterwards + # this incurrs a communication cost, but it avoids having to load + # the entire checkpoint on each rank, hopefully preventing OOM issues + # TODO: read on each host, instead of only the coordinator + if self.is_coordinator: + assert self.checkpoint_id is not None + torch_state_dict = torch.load( + self.checkpoint_id, map_location="cpu", weights_only=False + ) + if planner.flatten_state_dict: + torch_state_dict, _ = flatten_state_dict(torch_state_dict) + else: + torch_state_dict = None + + for req in plan.items: + if req.type == LoadItemType.BYTE_IO: + raise RuntimeError( + f"Non-tensor value identified at {req.storage_index.fqn}. " + f"At this time {type(self).__name__} only supports loading Tensors." + ) + + # Broadcast the tensor from the coordinator rank + if self.is_coordinator: + pg_device = dist.distributed_c10d._get_pg_default_device() + tensor = torch_state_dict[req.storage_index.fqn].to(pg_device) + else: + tensor = torch.empty_like(planner.state_dict[req.storage_index.fqn]) + + dist.broadcast(tensor, src=self.coordinator_rank, async_op=False) + + tensor = narrow_tensor_by_index(tensor, req.storage_offsets, req.lengths) + target_tensor = planner.resolve_tensor(req).detach() + assert target_tensor.size() == tensor.size(), ( + f"req {req.storage_index} mismatch sizes, " + f"{target_tensor.size()} vs {tensor.size()}" + ) + target_tensor.copy_(tensor) + planner.commit_tensor(req, target_tensor) + + fut: Future = Future() + fut.set_result(None) + return fut + + def set_up_storage_reader(self, metadata: Metadata, is_coordinator: bool) -> None: + """Implementation of the StorageReader method""" + self.is_coordinator = is_coordinator + if self.is_coordinator: + assert dist.get_rank() == self.coordinator_rank + + assert self.checkpoint_id is not None + + def prepare_local_plan(self, plan: LoadPlan) -> LoadPlan: + """Implementation of the StorageReader method""" + return plan + + def prepare_global_plan(self, global_plan: List[LoadPlan]) -> List[LoadPlan]: + """Implementation of the StorageReader method""" + return global_plan + + def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None: + """Implementation of the StorageReader method""" + self.checkpoint_id = checkpoint_id + + @classmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + """Implementation of the StorageReader method""" + return os.path.isfile(checkpoint_id) + + +class DynamicMetaLoadPlanner(DefaultLoadPlanner): + """ + Extension of DefaultLoadPlanner, which creates a new Metadata object based on the passed in state dict, + avoiding the need to read metadata from disk. This is useful when reading formats which don't have a + metadata file, like Torch Save files. + + . N.B. Intended to be used with BroadcastingTorchSaveReader + + .. warning:: + Current implementation only supports loading Tensors. + + >>> # xdoctest: +SKIP("undefined vars") + >>> sd = {"mode": model} + >>> dcp.load( + >>> sd, + >>> storage_reader=BroadcastingTorchSaveReader(), + >>> planner=DynamicMetaLoadPlanner(), + >>> checkpoint_id="path_to_model.pt" + >>> ) + """ + + def set_up_planner( + self, + state_dict: STATE_DICT_TYPE, + metadata: Optional[Metadata] = None, + is_coordinator: bool = False, + ) -> None: + """Setups of the planner, extnding default behavior by creating the Metadata object from the state dict""" + super().set_up_planner(state_dict, metadata, is_coordinator) + + state_dict_metadata: Dict[str, STORAGE_TYPES] = {} + for key, tensor in self.state_dict.items(): + if not torch.is_tensor(tensor): + raise RuntimeError( + f"Non-tensor value identified at {key}. " + f"At this time {type(self).__name__} only supports loading Tensors." + ) + + state_dict_metadata[key] = TensorStorageMetadata( + TensorProperties(dtype=tensor.dtype), + tensor.size(), + _create_chunk_list(tensor), + ) + self.metadata = Metadata(state_dict_metadata=state_dict_metadata) + + +def dcp_to_torch_save( + dcp_checkpoint_dir: Union[str, os.PathLike], + torch_save_path: Union[str, os.PathLike], +): + """ + Given a directory containing a DCP checkpoint, this function will convert it into a + Torch save file. + + Args: + dcp_checkpoint_dir: Directory containing the DCP checkpoint. + torch_save_path: Filename to store the converted Torch save file. + + .. warning:: + To avoid OOM, it's recommended to only run this function on a single rank. + """ + sd: STATE_DICT_TYPE = {} + _load_state_dict( + sd, + storage_reader=FileSystemReader(dcp_checkpoint_dir), + planner=_EmptyStateDictLoadPlanner(), + no_dist=True, + ) + torch.save(sd, torch_save_path) + + +def torch_save_to_dcp( + torch_save_path: Union[str, os.PathLike], + dcp_checkpoint_dir: Union[str, os.PathLike], +): + """ + Given the location of a torch save file, converts it into a DCP checkpoint. + + Args: + torch_save_path: Filename of the Torch save file. + dcp_checkpoint_dir: Directory to store the DCP checkpoint. + + .. warning:: + To avoid OOM, it's recommended to only run this function on a single rank. + """ + + state_dict = torch.load(torch_save_path, weights_only=False) + # we don't need stateful behavior here because the expectation is anything loaded by + # torch.load would not contain stateful objects. + _save_state_dict( + state_dict, storage_writer=FileSystemWriter(dcp_checkpoint_dir), no_dist=True + ) + + +if __name__ == "__main__": + + class FormatMode(Enum): + TORCH_TO_DCP = "torch_to_dcp" + DCP_TO_TORCH = "dcp_to_torch" + + # Parse command-line arguments + parser = argparse.ArgumentParser() + parser.add_argument( + "mode", + type=str, + help="Conversion mode", + choices=[m.value for m in FormatMode], + default=FormatMode.TORCH_TO_DCP, + ) + parser.add_argument("src", type=str, help="Path to the source model") + parser.add_argument("dst", type=str, help="Path to the destination model") + args = parser.parse_args() + + print( + f"Converting checkpoint from {args.src} to {args.dst} using method: '{args.mode}'" + ) + checkpoint_missing_warning = ( + f"No checkpoint found at {args.src}. Skipping conversion." + ) + if args.mode == FormatMode.TORCH_TO_DCP.value: + if os.path.isfile(args.src): + torch_save_to_dcp(args.src, args.dst) + else: + print(checkpoint_missing_warning) + elif args.mode == FormatMode.DCP_TO_TORCH.value: + if os.path.isdir(args.src): + dcp_to_torch_save(args.src, args.dst) + else: + print(checkpoint_missing_warning) + else: + raise ValueError(f"Unknown conversion mode: {args.mode}") diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/logging_handlers.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/logging_handlers.py new file mode 100644 index 0000000000000000000000000000000000000000..119dd52b7679750180c6ce1a1b0182762bc3f9e3 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/logging_handlers.py @@ -0,0 +1,15 @@ +import logging +from typing import List + +from torch.distributed.logging_handlers import _log_handlers + + +__all__: List[str] = [] + +DCP_LOGGER_NAME = "dcp_logger" + +_log_handlers.update( + { + DCP_LOGGER_NAME: logging.NullHandler(), + } +) diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/metadata.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..d1f87e2d9cba8a1b5e8c78f795eb7a6a2c10d2c6 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/metadata.py @@ -0,0 +1,182 @@ +# mypy: allow-untyped-defs +import os +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, List, Optional, Sequence, Union + +import torch +from torch.distributed.checkpoint.stateful import StatefulT + + +__all__ = [ + "ChunkStorageMetadata", + "TensorStorageMetadata", + "BytesStorageMetadata", + "Metadata", + "MetadataIndex", + "TensorProperties", + "StorageMeta", +] + + +@dataclass +class ChunkStorageMetadata: + """ + Each chunk is expected to have the same properties of the TensorStorageMetadata + that includes it. + """ + + offsets: torch.Size + sizes: torch.Size + + +class _MEM_FORMAT_ENCODING(Enum): + """Describe the memory format of a tensor.""" + + TORCH_CONTIGUOUS_FORMAT = 0 + TORCH_CHANNELS_LAST = 1 + TORCH_PRESERVE_FORMAT = 2 + + +@dataclass +class TensorProperties: + """Properties used to create :class:`Tensor`""" + + # Regular tensor fields + dtype: torch.dtype = field(default_factory=torch.get_default_dtype) + # This field is deprecated. + layout: torch.layout = field(default=torch.strided) + # This field is deprecated. + requires_grad: bool = False + # This field is deprecated. + memory_format: torch.memory_format = field(default=torch.contiguous_format) + # This field is deprecated. + pin_memory: bool = False + + def __getstate__(self): + # Since torch.memory_format cannot be pickled! + memory_format = self.memory_format + if memory_format == torch.contiguous_format: + mem_format_encoding = _MEM_FORMAT_ENCODING.TORCH_CONTIGUOUS_FORMAT + elif memory_format == torch.channels_last: + mem_format_encoding = _MEM_FORMAT_ENCODING.TORCH_CHANNELS_LAST + elif memory_format == torch.preserve_format: + mem_format_encoding = _MEM_FORMAT_ENCODING.TORCH_PRESERVE_FORMAT + else: + raise RuntimeError(f"Invalid torch.memory_format: {memory_format}") + + return ( + self.dtype, + self.layout, + self.requires_grad, + mem_format_encoding, + self.pin_memory, + ) + + def __setstate__( + self, + state, + ): + ( + self.dtype, + self.layout, + self.requires_grad, + mem_format_encoding, + self.pin_memory, + ) = state + + if mem_format_encoding == _MEM_FORMAT_ENCODING.TORCH_CONTIGUOUS_FORMAT: + memory_format = torch.contiguous_format + elif mem_format_encoding == _MEM_FORMAT_ENCODING.TORCH_CHANNELS_LAST: + memory_format = torch.channels_last + elif mem_format_encoding == _MEM_FORMAT_ENCODING.TORCH_PRESERVE_FORMAT: + memory_format = torch.preserve_format + else: + raise RuntimeError( + f"Invalid torch.memory_format encoding: {mem_format_encoding}" + ) + + self.memory_format = memory_format + + @staticmethod + def create_from_tensor(tensor: torch.Tensor) -> "TensorProperties": + return TensorProperties( + dtype=tensor.dtype, + layout=tensor.layout, + requires_grad=tensor.requires_grad, + memory_format=torch.contiguous_format, + pin_memory=tensor.is_pinned(), + ) + + +@dataclass +class TensorStorageMetadata: + properties: TensorProperties + size: torch.Size + chunks: List[ChunkStorageMetadata] + + +@dataclass +class BytesStorageMetadata: + pass + + +STORAGE_TYPES = Union[TensorStorageMetadata, BytesStorageMetadata] +STATE_DICT_TYPE = Dict[str, Union[StatefulT, Any]] + + +@dataclass +class StorageMeta: + checkpoint_id: Union[str, os.PathLike, None] = None + save_id: Optional[str] = None + load_id: Optional[str] = None + + +@dataclass +class Metadata: + """This class represents the metadata of the checkpoint.""" + + # Keys are the same from the `state_dict` used. + state_dict_metadata: Dict[str, STORAGE_TYPES] + # It is the responsibility of the planner and storage plugins to ensure + # backward compatibility of the planner_data and storage_data. DCP will + # also ensure the backward compatibility of the metadata in this file and + # the metadata of the built-in planner and storage plugins. + planner_data: Any = None + storage_data: Any = None + storage_meta: Optional[StorageMeta] = None + + +@dataclass(frozen=True) +class MetadataIndex: + """This class represents a lookup key for items in a state dict or Metadata.""" + + fqn: str + """Fully Qualified Name of the object""" + + offset: Optional[torch.Size] = None + """If the object is a tensor, offset into the tensor we're looking for""" + + index: Optional[int] = field(hash=False, compare=False, default=None) + """ + Index hint when searching for tensor chunk to speedup lookups (optional) + + A common representation of a sharded tensor is as a list of chunks so to + find the index in such a list you need to linear search it. + + When constructing an instance of MetadataIndex that points to that list, + one can provide the index as a hint and it will be probed first before + the linear search and thus making it significantly faster. + """ + + def __init__( + self, + fqn: str, + offset: Optional[Sequence[int]] = None, + index: Optional[int] = None, + ): + # We must use object.__setattr__ due to frozen=True + object.__setattr__(self, "fqn", fqn) + object.__setattr__(self, "index", index) + if offset is not None: + object.__setattr__(self, "offset", torch.Size(offset)) diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/optimizer.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..5ad170768db9294e63e09b2c99da8f0d6b93241d --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/optimizer.py @@ -0,0 +1,356 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates + +import dataclasses +from typing import cast, Dict, List, Optional, Sequence, Tuple, Union + +import torch +import torch.distributed as dist +from torch._utils import _get_device_module +from torch.distributed._shard.sharded_tensor.api import ShardedTensor +from torch.distributed._shard.sharded_tensor.metadata import ( + TensorProperties as ShardTensorProperties, +) +from torch.distributed._shard.sharded_tensor.shard import Shard +from torch.distributed._shard.sharding_spec.chunk_sharding_spec import ChunkShardingSpec +from torch.distributed.checkpoint._nested_dict import unflatten_state_dict +from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner +from torch.distributed.checkpoint.metadata import ( + BytesStorageMetadata, + ChunkStorageMetadata, + Metadata, + MetadataIndex, + STATE_DICT_TYPE, + TensorProperties, + TensorStorageMetadata, +) +from torch.distributed.checkpoint.planner import LoadPlan, LoadPlanner +from torch.distributed.checkpoint.planner_helpers import ( + _create_read_items, + create_read_items_for_chunk_list, +) +from torch.distributed.checkpoint.state_dict_loader import load_state_dict +from torch.distributed.checkpoint.storage import StorageReader +from torch.distributed.checkpoint.utils import ( + _element_wise_add, + _element_wise_sub, + _normalize_device_info, +) +from torch.distributed.distributed_c10d import _get_default_group +from torch.distributed.fsdp._shard_utils import _create_chunk_sharded_tensor +from torch.distributed.remote_device import _remote_device +from torch.distributed.tensor import DTensor + + +STATE_DICT_2D_LAYOUT = Dict[str, Tuple[Optional[Sequence[int]], Sequence[int]]] + + +# TODO: Update docstrings for optimizer.py +__all__ = [ + "load_sharded_optimizer_state_dict", +] + + +def _gen_rank_device(global_rank: int, device_type: str = "cuda") -> str: + if device_type == "cpu": + return "cpu" + device_module = _get_device_module(device_type) + if device_module.is_available(): + return _normalize_device_info( + device_type, global_rank % device_module.device_count() + ) + return "cpu" + + +def _create_colwise_spec( + pg: Optional[dist.ProcessGroup] = None, +) -> ChunkShardingSpec: + pg_device_type = dist.distributed_c10d._get_pg_default_device(pg).type + if pg is None: + placements = [ + f"rank:{idx}/{_gen_rank_device(idx, pg_device_type)}" + for idx in range(dist.get_world_size()) + ] + else: + placements = [ + f"rank:{idx}/{_gen_rank_device(dist.get_global_rank(pg, idx), pg_device_type)}" + for idx in range(pg.size()) + ] + return ChunkShardingSpec( + dim=0, + placements=cast(List[Union[_remote_device, str]], placements), + ) + + +def _is_nested_tensor(val: torch.Tensor) -> bool: + if type(val) is ShardedTensor: + if len(val.local_shards()) == 0: + return False + if type(val.local_shards()[0].tensor) is ShardedTensor: + return True + if type(val.local_shards()[0].tensor) is DTensor: + raise ValueError("Cannot handle DTensor nested insided ShardedTensor") + elif type(val) is DTensor and ( + type(val._local_tensor) is DTensor or type(val._local_tensor) is ShardedTensor + ): + raise ValueError("Cannot handle nested DTensor") + return False + + +def _alloc_tensor( + props: TensorProperties, size: Sequence[int], device_type: str = "cuda" +) -> torch.Tensor: + if device_type == "cpu": + device = cast(torch.device, _get_device_module(device_type).current_device()) + else: + device = torch.device( + device_type, _get_device_module(device_type).current_device() + ) + + return torch.empty( + size=size, + dtype=props.dtype, + layout=props.layout, + requires_grad=props.requires_grad, + pin_memory=props.pin_memory, + device=device, + ) + + +def _get_state_dict_2d_layout( + state_dict: STATE_DICT_TYPE, +) -> Tuple[STATE_DICT_2D_LAYOUT, Optional[dist.ProcessGroup]]: + """ + Load the right TP slice of the optimizer state. + + This is not easy since the per-tensor slicing can't be inferred from checkpoint metadata. + We take advantage of the model state_dict producing a sliced ST to figure out what we need to load. + This is pretty fragile and it might be easier for FSDP to compute this info for us. + Returns a dictionary where keys are the same of the state_dict and the value is a tuple of + (offset, size) for the current rank TP slice. + N.B. The state_dict *MUST* come from FSDP.sharded_state_dict. + """ + specs: STATE_DICT_2D_LAYOUT = {} + dp_pg: Optional[dist.ProcessGroup] = None + for key, value in state_dict.items(): + specs[key] = (None, value.size()) + if _is_nested_tensor(value): + assert ( + len(value.local_shards()) == 1 + ), "Cannot handle ST with multiple shards" + assert isinstance( + value, ShardedTensor + ), "Can only handle nested ShardedTensor" + shard = value.local_shards()[0] + specs[key] = ( + shard.metadata.shard_offsets, + shard.metadata.shard_sizes, + ) + dp_pg = shard.tensor._process_group # type: ignore[attr-defined] + + return ( + specs, + dp_pg, + ) + + +class _ReaderWithOffset(DefaultLoadPlanner): + translation: Dict[MetadataIndex, MetadataIndex] + state_dict: STATE_DICT_TYPE + metadata: Metadata + + def __init__(self, fqn_to_offset: Dict[str, Sequence[int]]) -> None: + super().__init__() + self.fqn_to_offset = fqn_to_offset + self.metadata = Metadata({}) + self.state_dict = {} + self.translation = {} + + def create_local_plan(self) -> LoadPlan: + requests = [] + self.translation = {} + for fqn, obj in self.state_dict.items(): + md = self.metadata.state_dict_metadata[fqn] + if not isinstance(obj, ShardedTensor): + requests += _create_read_items(fqn, md, obj) + continue + + if fqn not in self.fqn_to_offset: + requests += _create_read_items(fqn, md, obj) + continue + + offset = self.fqn_to_offset[fqn] + + assert len(obj.local_shards()) == 1 + original_shard = obj.local_shards()[0] + local_chunks = [ + ChunkStorageMetadata( + offsets=torch.Size( + _element_wise_add(original_shard.metadata.shard_offsets, offset) + ), + sizes=torch.Size(original_shard.metadata.shard_sizes), + ) + ] + + reqs = create_read_items_for_chunk_list( + fqn, cast(TensorStorageMetadata, md), local_chunks + ) + # TODO: The ReadItems will have a displaced MetadataIndex, fix it. + # TODO: we should change _create_sharded_read_items to have more ergonomic API + for ri in reqs: + assert ri.dest_index.offset is not None + original_offset = _element_wise_sub(ri.dest_index.offset, offset) + original_index = dataclasses.replace( + ri.dest_index, offset=torch.Size(original_offset) + ) + self.translation[ri.dest_index] = original_index + + requests += reqs + return LoadPlan(requests) + + def lookup_tensor(self, index: MetadataIndex) -> torch.Tensor: + return super().lookup_tensor(self.translation.get(index, index)) + + +def load_sharded_optimizer_state_dict( + model_state_dict: STATE_DICT_TYPE, + optimizer_key: str, + storage_reader: StorageReader, + planner: Optional[LoadPlanner] = None, +) -> STATE_DICT_TYPE: + """ + Load a state_dict in conjunction with FSDP sharded optimizer state. + + This is the current recommended way to checkpoint FSDP. + >>> # xdoctest: +SKIP + >>> import torch.distributed.checkpoint as dist_cp + >>> # Save + >>> model: torch.nn.Model + >>> optim_params = model.parameters() + >>> optim = torch.optim.SGD(optim_params, lr=0.01) + >>> # Save + >>> with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT): + >>> state_dict = { + >>> "optimizer": FSDP.optim_state_dict(model, optim), + >>> "model": model.state_dict() + >>> } + >>> dist_cp.save_state_dict( + >>> state_dict=optim_state, + >>> storage_writer=dist_cp.FileSystemWriter("checkpoint"), + >>> planner=dist_cp.DefaultSavePlanner(), + >>> ) + >>> + >>> # Load + >>> with FSDP.state_dict_type(model_tp, StateDictType.SHARDED_STATE_DICT): + >>> model_state_dict = model_tp.state_dict() + >>> checkpoint = { + >>> "model": model_state_dict + >>> } + >>> dist_cp.load_state_dict( + >>> state_dict=checkpoint, + >>> storage_reader=dist_cp.FileSystemReader(checkpoint_file), + >>> planner=dist_cp.DefaultLoadPlanner(), + >>> ) + >>> model.load_state_dict(checkpoint["model_state"]) + >>> + >>> optim_state = dist_cp.load_sharded_optimizer_state_dict( + >>> model_state_dict, + >>> optimizer_key="optimizer", + >>> storage_reader=dist_cp.FileSystemReader("checkpoint"), + >>> ) + >>> + >>> flattened_osd = FSDP.optim_state_dict_to_load( + >>> model, optim, optim_state["optimizer"] + >>> ) + >>> + >>> optim.load_state_dict(flattened_osd) + """ + metadata = storage_reader.read_metadata() + + layout_specs, dp_pg = _get_state_dict_2d_layout(model_state_dict) + dp_pg_device_type = dist.distributed_c10d._get_pg_default_device(dp_pg).type + device_module = _get_device_module(dp_pg_device_type) + + if dp_pg is None: + placements = [] + for i in range(dist.get_world_size()): + device_info = _normalize_device_info( + dp_pg_device_type, i % device_module.device_count() + ) + placements.append(f"rank:{i}/{device_info}") + sharding_spec = ChunkShardingSpec(dim=0, placements=placements) # type: ignore[arg-type] + else: + sharding_spec = _create_colwise_spec(dp_pg) + + # Create a state_dict for optimizer state + state_dict: STATE_DICT_TYPE = {} + + fqn_to_offset: Dict[str, Sequence[int]] = {} + for key, value in metadata.state_dict_metadata.items(): + key_path = metadata.planner_data[key] + if key_path[0] != optimizer_key: + continue + + if isinstance(value, BytesStorageMetadata): + state_dict[key] = "" + continue + + # value: TensorStorageMetadata + if value.size.numel() == 1: + state_dict[key] = _alloc_tensor( + value.properties, value.size, dp_pg_device_type + ) + elif dp_pg is None: + state_dict[key] = _create_chunk_sharded_tensor( + _alloc_tensor(value.properties, value.size, dp_pg_device_type), + rank=dist.get_rank(), + world_size=dist.get_world_size(), + num_devices_per_node=device_module.device_count(), + pg=_get_default_group(), + ) + else: + spec_key = key_path[2] + alloc_size = layout_specs.get(spec_key, (None, value.size))[1] + + properties = ShardTensorProperties( + dtype=value.properties.dtype, + layout=value.properties.layout, + requires_grad=value.properties.requires_grad, + memory_format=value.properties.memory_format, + pin_memory=value.properties.pin_memory, + ) + + st_md = sharding_spec.build_metadata(torch.Size(alloc_size), properties) + local_shards = [] + current_rank = dist.get_rank(dp_pg) + for shard_md in st_md.shards_metadata: + if cast(_remote_device, shard_md.placement).rank() != current_rank: + continue + local_shards.append( + Shard( + tensor=_alloc_tensor( + value.properties, shard_md.shard_sizes, dp_pg_device_type + ), + metadata=shard_md, + ) + ) + + st = ShardedTensor._init_from_local_shards_and_global_metadata( + local_shards, st_md, process_group=dp_pg + ) + + if spec_key in layout_specs and layout_specs[spec_key][0] is not None: + fqn_to_offset[key] = cast(Sequence[int], layout_specs[spec_key][0]) + + state_dict[key] = st + + # Whether we unflatten before or after doesn't matter + load_state_dict( + state_dict=state_dict, + storage_reader=storage_reader, + # FIXME the type of planner is wrong in load_state_dict + planner=_ReaderWithOffset(fqn_to_offset) if dp_pg is not None else planner, + ) + + state_dict = unflatten_state_dict(state_dict, metadata.planner_data) + + return state_dict diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/planner.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/planner.py new file mode 100644 index 0000000000000000000000000000000000000000..75619b627356ad3c0c63fa7cedae4733d9771491 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/planner.py @@ -0,0 +1,417 @@ +import abc +import io +import operator +from dataclasses import dataclass +from enum import auto, Enum +from functools import reduce +from typing import Any, List, Optional, Tuple, Union + +import torch +from torch.distributed.checkpoint.metadata import ( + ChunkStorageMetadata, + Metadata, + MetadataIndex, + STATE_DICT_TYPE, + StorageMeta, + TensorProperties, +) + + +__all__ = [ + "WriteItemType", + "LoadItemType", + "TensorWriteData", + "WriteItem", + "ReadItem", + "SavePlan", + "LoadPlan", + "SavePlanner", + "LoadPlanner", +] + + +class WriteItemType(Enum): + TENSOR = auto() + SHARD = auto() + BYTE_IO = auto() + + +class LoadItemType(Enum): + TENSOR = auto() + BYTE_IO = auto() + + +@dataclass(frozen=True) +class TensorWriteData: + chunk: ChunkStorageMetadata + properties: TensorProperties + size: torch.Size + + +@dataclass(frozen=True) +class WriteItem: + """Dataclass which holds information about what needs to be written to storage.""" + + index: MetadataIndex + type: WriteItemType + + # Value present if it's a tensor write + tensor_data: Optional[TensorWriteData] = None + + def tensor_storage_size(self) -> Optional[int]: + """ + Calculates the storage size of the underlying tensor, or None if this is not a tensor write. + + Returns: + Optional[int] storage size, in bytes of underlying tensor if any. + """ + if self.tensor_data is None: + return None + + numels = reduce(operator.mul, self.tensor_data.size, 1) + dtype_size = torch._utils._element_size(self.tensor_data.properties.dtype) + return numels * dtype_size + + +@dataclass(frozen=True) +class ReadItem: + # Read Item + type: LoadItemType + + # Index into the state_dict + dest_index: MetadataIndex + # Offsets into destination tensor + dest_offsets: torch.Size + + # Index into the checkpoint + storage_index: MetadataIndex + # Offset into the checkpoint data + storage_offsets: torch.Size + + # Size of the hypercube to copy + lengths: torch.Size + + +@dataclass(frozen=True) +class SavePlan: + items: List[WriteItem] + storage_data: Any = None + planner_data: Any = None + + +@dataclass +class LoadPlan: + items: List[ReadItem] + storage_data: Any = None + planner_data: Any = None + + +class SavePlanner(abc.ABC): + """ + Abstract class defining the protocol used by save_state_dict to plan the save process. + + SavePlanners are stateful objects that can be used to customize the whole save process. + + SavePlanner acts as an access proxy to the state_dict, so any transformation done to it + will be visible to the whole process. + + A planner subclass can expect the following sequence of calls during save_state_dict: + + 1) set_up_planner - called on all ranks. + Signals the start of a checkpoint save. + + 2) create_local_plan - called on all ranks. + Process the state_dict and produces a `SavePlan` that will be sent for global planning. + + 3) create_global_plan - called on the coordinator rank only. + Takes the SavePlan from all ranks and make any global decision. + + 4) finish_plan - called on all ranks. + This gives each rank a chance to adjust to global planning decisions. + + 5) resolve_data - called multiple times on each rank + Lookups a value on the `state_dict` for the storage layer to write. + + Users are recommended to extend DefaultSavePlanner instead of this interface directly as + most changes can be expressed by changes in a single method. + + There are 3 usual patterns of extension: + + Rewriting state_dict. This is the simplest way to extend the save process as it + doesn't requite understanding the intrincacies of how SavePlan works: + + >>> # xdoctest: +SKIP("undefined vars") + >>> class RenamePlanner(DefaultSavePlanner): + >>> def set_up_planner( + >>> self, + >>> state_dict: STATE_DICT_TYPE, + >>> storage_meta: Optional[StorageMeta], + >>> is_coordinator: bool, + >>> ) -> None: + >>> # prefix all keys with `foo_`` + >>> super().set_up_planner({"foo_" + k: v for k, v in state_dict.items()}, storage_meta, is_coordinator) + + Modifying local plan and lookup in tandem. This is useful when fine control of how data is persisted + + >>> # xdoctest: +SKIP("undefined vars") + >>> class FP16Planner(DefaultSavePlanner): + >>> def create_local_plan(self): + >>> plan = super().create_local_plan() + >>> for p in plan: + >>> if p.tensor_data is not None: + >>> p.tensor_data.properties.dtype = torch.float16 + >>> return plan + >>> + >>> def resolve_data(self, write_item): + >>> item = super().resolve_data(write_item) + >>> return item if write_item.type == WriteItemType.BYTE_IO else item.to(torch.float16) + + Using the global planning step to make central decisions that can't be made individually by each rank + + >>> # xdoctest: +SKIP("undefined vars") + >>> from itertools import zip_longest + >>> from dataclasses import replace + >>> class DDPLoadBalancingPlanner(DefaultSavePlanner): + >>> # This uses the default local plan behavior of having all non-sharded writes in rank 0 + >>> # This sample doesn't handle ShardedTensors + >>> def create_global_plan(self, all_plans): + >>> iters = [iter(all_plans[0].items)] * len(all_plans) + >>> items_per_rank = [ + >>> [item for item in items if item is not None] + >>> for items in zip(*zip_longest(*iters), strict=True) + >>> ] + >>> all_plans = [ + >>> replace(plan, items=items) + >>> for plan, items in zip(all_plans, items_per_rank, strict=True) + >>> ] + >>> return super().create_global_plan(all_plans) + + Finally, some planners need to save additional metadata in the checkpoint, this is + accomplished by having each rank contribute their data items in the local plan and + the global planner aggregate them: + + >>> # xdoctest: +SKIP("undefined vars") + >>> class SaveExtraDataPlanner(DefaultSavePlanner): + >>> def create_local_plan(self) -> SavePlan: + >>> plan = super().create_local_plan() + >>> return replace(plan, planner_data="per-rank-data") + >>> + >>> def create_global_plan(self, all_plans: List[SavePlan]) -> Tuple[List[SavePlan], Metadata]: + >>> global_plan, metadata = super().create_global_plan(all_plans) + >>> merged_data = [p.planner_data for p in global_plan] + >>> metadata = replace(metadata, planner_data=merged_data) + >>> return global_plan, metadata + """ + + @abc.abstractmethod + def set_up_planner( + self, + state_dict: STATE_DICT_TYPE, + storage_meta: Optional[StorageMeta] = None, + is_coordinator: bool = False, + ) -> None: + """ + Initialize this planner to save ``state_dict``. + + Implementations should save those values as they won't be provided lated in the save process. + + This is called on all ranks. + """ + + @abc.abstractmethod + def create_local_plan(self) -> SavePlan: + """ + Compute the save plan for the current rank. + + This will be aggregated and passed to create_global_plan. + Planner specific data can be passed through SavePlan::planner_data. + + This is called on all ranks. + """ + + @abc.abstractmethod + def create_global_plan( + self, all_plans: List[SavePlan] + ) -> Tuple[List[SavePlan], Metadata]: + """ + Compute the global checkpoint plan and return the local plan of each rank. + + This is called on the coordinator rank only. + """ + + @abc.abstractmethod + def finish_plan(self, new_plan: SavePlan) -> SavePlan: + """ + Merge the plan created by `create_local_plan` and the result of `create_global_plan`. + + This is called on all ranks. + """ + + @abc.abstractmethod + def resolve_data(self, write_item: WriteItem) -> Union[torch.Tensor, io.BytesIO]: + """ + Transform and prepare ``write_item`` from ``state_dict`` for storage, ensuring idempotency and thread-safety. + + Lookup the object associated with ``write_item`` in ``state_dict`` and apply any + transformation (such as serialization) prior to the storage layer consuming it. + + Called on each rank multiple times, at least once per WriteItem in the final SavePlan. + + This method should be idempotent and thread-save. StorageWriter implementations + are free to call it as frequently as they need. + + Any transformation that allocates memory should be lazily done when his method + is called in order to reduce peak memory required by checkpointing. + + When returning tensors, they can be on any device or format, they can be views too. + It's the storage layer responsibility to figure out how to save them. + """ + + +class LoadPlanner: + """ + Abstract class defining the protocol used by load_state_dict to plan the load process. + + LoadPlanner are stateful objects that can be used to customize the whole load process. + + LoadPlanner acts as an access proxy to the state_dict, so any transformation done to it + will be visible to the whole process. + + A planner subclass can expect the following sequence of calls during load_state_dict: + + 1) set_up_planner - called on all ranks. + Signals the start of loading a checkpoint. + + 2) create_local_plan - called on all ranks. + Process the state_dict and produces a `LoadPlan` that will be sent for global planning. + + 3) create_global_plan - called on the coordinator rank only. + Takes the LoadPlan from all ranks and make any global decision. + + 4) load_bytes - called multiple times on each rank + This is called once per non-tensor value in state_dict. + + 5) resolve_tensor and commit_tensor - called multiple times on each rank + They are called in pair for each Tensor value in state_dict. + + Users are recommended to extend DefaultLoadPlanner instead of this interface directly as + most changes can be expressed by changes in a single method. + + There are two usual patterns of extension: + + Rewriting state_dict. This is the simplest way to extend the load process as it + doesn't requite understanding the intrincacies of how LoadPlan works. We need + to keep a reference to the original state_dict as load happens in place so + we need to be able to perform it in place + + >>> # xdoctest: +SKIP("undefined vars") + >>> class RenamePlanner(DefaultLoadPlanner): + >>> def set_up_planner( + >>> self, + >>> state_dict: STATE_DICT_TYPE, + >>> metadata: Metadata, + >>> is_coordinator: bool, + >>> ) -> None: + >>> self.original_state_dict = state_dict + >>> state_dict = {"foo_" + k: v for k, v in state_dict.items()} + >>> + >>> if self.flatten_sharded_tensors: + >>> state_dict = _flatten_sharded_tensors(state_dict) + >>> + >>> if self.flatten_state_dict: + >>> state_dict, self.mappings = flatten_state_dict(state_dict) + >>> + >>> self.state_dict = state_dict + >>> self.metadata = metadata + >>> self.is_coordinator = is_coordinator + >>> + >>> def load_bytes(self, read_item, value): + >>> # Remove the "foo_" prefix + >>> self.original_state_dict[read_item.dest_index.fqn[4:]] = torch.load(value, weights_only=False) + + + Modifying resolve_tensor and commit_tensor to handle load time transformation. + + >>> # xdoctest: +SKIP("undefined vars") + >>> class MetaModelMaterialize(DefaultSavePlanner): + >>> def resolve_tensor(self, read_item): + >>> tensor = super().resolve_tensor(read_item) + >>> return torch.empty_like(tensor, device="cpu") + >>> + >>> def commit_tensor(self, read_item, tensor): + >>> self.state_dict[read_item.dest_index.fqn] = tensor + """ + + @abc.abstractmethod + def set_up_planner( + self, + state_dict: STATE_DICT_TYPE, + metadata: Optional[Metadata] = None, + is_coordinator: bool = False, + ) -> None: + """ + Initialize this instance to load data into ``state_dict``. + + . N.B. This is called on every rank. + """ + + @abc.abstractmethod + def create_local_plan(self) -> LoadPlan: + """ + Create a LoadPlan based on state_dict and metadata provided by set_up_planner. + + . N.B. This is called on every rank. + """ + + @abc.abstractmethod + def create_global_plan(self, global_plan: List[LoadPlan]) -> List[LoadPlan]: + """ + Compute the global load plan and return plans for each rank. + + . N.B. This is called on the coordinator rank only + """ + + @abc.abstractmethod + def finish_plan(self, central_plan: LoadPlan) -> LoadPlan: + """Accept the plan from coordinator and return final LoadPlan.""" + + @abc.abstractmethod + def load_bytes(self, read_item: ReadItem, value: io.BytesIO) -> None: + """ + Load the item described by ``read_item``and ``value``. + + This method is expected to modify in-place the underlying state_dict. + + The contents of ``value`` are defined by the SavePlanner used to produce + the checkpoint being loaded. + """ + + def resolve_bytes(self, read_item: ReadItem) -> io.BytesIO: + """ + Return the BytesIO to be used by the StorageReader to load `read_item`. + + The BytesIO should alias with one on the underlying state_dict as StorageReader will replace its contents. + """ + raise NotImplementedError("LoadPlanner.resolve_bytes is not implemented") + + @abc.abstractmethod + def resolve_tensor(self, read_item: ReadItem) -> torch.Tensor: + """ + Return the tensor described by ``read_item`` to be used by the StorageReader to load `read_item`. + + The tensor should alias with one on the underlying state_dict as StorageReader will replace its contents. + If, for any reason, that's not possible, the planner can use the ``commit_tensor`` method to copy the data + back to the one in state_dict. + """ + + @abc.abstractmethod + def commit_tensor(self, read_item: ReadItem, tensor: torch.Tensor) -> None: + """ + Call once the StorageReader finished loading data into ``tensor``. + + The provided tensor is the same one returned by the call to ``resolve_tensor``. + This method is only needed if this LoadPlanner needs to post process ``tensor`` prior to + copying it back to the one in the state_dict. + + The contents of tensor will follow its device synchronization model. + """ diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/planner_helpers.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/planner_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..06b2482d5d514373cdbc13139d8107a47bdff8b0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/planner_helpers.py @@ -0,0 +1,388 @@ +# mypy: allow-untyped-defs +import io +from typing import Any, Callable, cast, Dict, List + +import torch +import torch.distributed as dist +from torch._utils import _get_device_module +from torch.distributed._shard.metadata import ShardMetadata +from torch.distributed._shard.sharded_tensor import ShardedTensor +from torch.distributed.tensor import DTensor +from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + +from .metadata import ( + BytesStorageMetadata, + ChunkStorageMetadata, + MetadataIndex, + STATE_DICT_TYPE, + STORAGE_TYPES, + TensorProperties, + TensorStorageMetadata, +) +from .planner import ( + LoadItemType, + ReadItem, + SavePlan, + TensorWriteData, + WriteItem, + WriteItemType, +) +from .resharding import ( + _check_shard_metadata_pair_overlap, + _shards_get_overlap_region_wrt_saved_tensor, +) + + +__all__: List[str] = ["create_read_items_for_chunk_list"] + + +def _create_chunk_from_tensor(tensor: torch.Tensor) -> ChunkStorageMetadata: + return ChunkStorageMetadata( + offsets=torch.Size([0] * len(tensor.size())), sizes=tensor.size() + ) + + +def _chunk_for_shard(shard_md: ShardMetadata) -> ChunkStorageMetadata: + return ChunkStorageMetadata( + offsets=torch.Size(shard_md.shard_offsets), + sizes=torch.Size(shard_md.shard_sizes), + ) + + +def _sharded_tensor_metadata( + sharded_tensor: ShardedTensor, shard_md: ShardMetadata +) -> TensorWriteData: + shard_properties = sharded_tensor.metadata().tensor_properties + + properties = TensorProperties( + dtype=shard_properties.dtype, + layout=shard_properties.layout, + requires_grad=shard_properties.requires_grad, + memory_format=shard_properties.memory_format, + pin_memory=shard_properties.pin_memory, + ) + + return TensorWriteData( + chunk=_chunk_for_shard(shard_md), + properties=properties, + size=sharded_tensor.metadata().size, + ) + + +def _create_write_items_for_dtensor(fqn: str, tensor: DTensor) -> WriteItem: + sizes, offsets = compute_local_shape_and_global_offset( + tensor.shape, tensor.device_mesh, tensor.placements + ) + sizes, offsets = torch.Size(sizes), torch.Size(offsets) + + return WriteItem( + index=MetadataIndex(fqn, offsets), + type=WriteItemType.SHARD, + tensor_data=TensorWriteData( + chunk=ChunkStorageMetadata( + offsets=offsets, + sizes=sizes, + ), + properties=TensorProperties.create_from_tensor(tensor.to_local()), + size=tensor.size(), + ), + ) + + +def _create_write_item_for_shard( + fqn: str, sharded_tensor: ShardedTensor, shard_md: ShardMetadata +) -> WriteItem: + offsets = torch.Size(shard_md.shard_offsets) + return WriteItem( + index=MetadataIndex(fqn, offsets), + type=WriteItemType.SHARD, + tensor_data=_sharded_tensor_metadata(sharded_tensor, shard_md), + ) + + +def _create_write_item_for_tensor(fqn: str, tensor: torch.Tensor) -> WriteItem: + offsets = torch.Size([0] * len(tensor.size())) + return WriteItem( + index=MetadataIndex(fqn, offsets), + type=WriteItemType.TENSOR, + tensor_data=TensorWriteData( + chunk=ChunkStorageMetadata(offsets=offsets, sizes=tensor.size()), + properties=TensorProperties.create_from_tensor(tensor), + size=tensor.size(), + ), + ) + + +def _create_write_item_for_bytesio(fqn: str, bytes: Any): + return WriteItem( + index=MetadataIndex(fqn), + type=WriteItemType.BYTE_IO, + ) + + +def _create_read_item_for_byteio( + dest_index, dest_offset, storage_index, storage_offset, length +): + return ReadItem( + type=LoadItemType.BYTE_IO, + dest_index=dest_index, + dest_offsets=torch.Size((dest_offset,)), + storage_index=storage_index, + storage_offsets=torch.Size((storage_offset,)), + lengths=torch.Size((length,)), + ) + + +def _create_read_item_for_tensor( + dest_index, dest_offsets, storage_index, storage_offsets, lengths +): + return ReadItem( + type=LoadItemType.TENSOR, + dest_index=dest_index, + dest_offsets=torch.Size(dest_offsets), + storage_index=storage_index, + storage_offsets=torch.Size(storage_offsets), + lengths=torch.Size(lengths), + ) + + +def create_read_items_for_chunk_list( + fqn: str, + checkpoint_md: TensorStorageMetadata, + local_chunks: List[ChunkStorageMetadata], +) -> List[ReadItem]: + """ + Create a list of ``ReadItem`` based on the checkpoint and local chunks. + + This applies the resharding algorithm and computes the reads needed + to satisfy ``local_chunks`` with a checkpoint described by ``checkpoint_md``. + + Args: + fqn (str) : The state_dict FQN to pass to ``ReadItem``. + checkpoint_md (TensorStorageMetadata): metadata for a given tensor + from a checkpoint. + local_chunks (List[ChunkStorageMetadata]): Local chunks that needs to be + loaded. + + Returns: + A list of ``ReadItem`` that will satisfy all input chunks. + """ + read_items = [] + # this is a naive quadratic algo that can be optimized later + for idx, shard in enumerate(local_chunks): + for storage_idx, storage_md in enumerate(checkpoint_md.chunks): + if not _check_shard_metadata_pair_overlap(shard, storage_md): + continue + + storage_offsets = [] + dest_offsets = [] + lengths = [] + for ( + _dim, + offset_for_saved_tensor, + offset_for_current_tensor, + length, + ) in _shards_get_overlap_region_wrt_saved_tensor( + saved_shard=storage_md, current_shard=shard + ): + storage_offsets.append(offset_for_saved_tensor) + dest_offsets.append(offset_for_current_tensor) + lengths.append(length) + + read_items.append( + _create_read_item_for_tensor( + dest_index=MetadataIndex(fqn, shard.offsets, idx), + dest_offsets=dest_offsets, + storage_index=MetadataIndex(fqn, storage_md.offsets, storage_idx), + storage_offsets=storage_offsets, + lengths=lengths, + ) + ) + return read_items + + +def _create_default_metadata_only_plan(state_dict: STATE_DICT_TYPE) -> SavePlan: + requests = [] + for fqn, obj in state_dict.items(): + if isinstance(obj, DTensor): + requests.append(_create_write_items_for_dtensor(fqn, obj)) + elif isinstance(obj, ShardedTensor): + requests.extend( + _create_write_item_for_shard(fqn, obj, shard_md) + for shard_md in obj.metadata().shards_metadata + ) + elif isinstance(obj, torch.Tensor): + requests.append(_create_write_item_for_tensor(fqn, obj)) + else: + requests.append(_create_write_item_for_bytesio(fqn, obj)) + return SavePlan(requests) + + +def _create_write_items(fqn: str, object: Any) -> List[WriteItem]: + if hasattr(object, "__create_write_items__"): + # DTensor implements _Checkpointable + return object.__create_write_items__(fqn, object) + elif isinstance(object, ShardedTensor): + return [ + _create_write_item_for_shard(fqn, object, shard.metadata) + for shard in object.local_shards() + ] + elif isinstance(object, torch.Tensor): + return [_create_write_item_for_tensor(fqn, object)] + else: + return [_create_write_item_for_bytesio(fqn, object)] + + +def _create_chunk_from_dtensor(tensor: DTensor) -> ChunkStorageMetadata: + sizes, offsets = compute_local_shape_and_global_offset( + tensor.shape, tensor.device_mesh, tensor.placements + ) + sizes, offsets = torch.Size(sizes), torch.Size(offsets) + return ChunkStorageMetadata( + offsets=offsets, + sizes=sizes, + ) + + +def _create_chunk_list(tensor: torch.Tensor) -> List[ChunkStorageMetadata]: + if hasattr(tensor, "__create_chunk_list__"): + # DTensor implements _Checkpointable + local_chunks = tensor.__create_chunk_list__() # type: ignore[attr-defined] + elif isinstance(tensor, ShardedTensor): + local_chunks = [ + _chunk_for_shard(shard.metadata) for shard in tensor.local_shards() + ] + elif isinstance(tensor, torch.Tensor): + local_chunks = [_create_chunk_from_tensor(tensor)] + else: + raise ValueError( + "Unsupported Type, expecting one of [Tensor, DTensor, ShardedTensor] " + f",but got {type(tensor)}" + ) + + return local_chunks + + +def _create_read_items(fqn: str, md: STORAGE_TYPES, obj: Any) -> List[ReadItem]: + if not isinstance(md, BytesStorageMetadata): + try: + local_chunks = _create_chunk_list(obj) + except ValueError as ex: + raise ValueError( + f"Invalid checkpoint metadata for {fqn}, " + + f"expected BytesStorageMetadata but found {type(md)}", + ) from ex + + return create_read_items_for_chunk_list(fqn, md, local_chunks) + else: + return [ + _create_read_item_for_byteio( + dest_index=MetadataIndex(fqn), + dest_offset=0, + storage_index=MetadataIndex(fqn), + storage_offset=0, + length=0, + ) + ] + + +def _init_state_dict(state_dict: Dict[str, Any]) -> Any: + """ + Initializes meta tensor if the meta tensor is DTensor or torch.Tensor. + """ + + def dtensor_func(value: DTensor): + device = getattr(value, "device", None) + if device == torch.device("meta"): + device_type = dist.distributed_c10d._get_pg_default_device().type + device = cast( + torch.device, _get_device_module(device_type).current_device() + ) + new_local_tensor = torch.empty_like(value.to_local(), device=device) + # We need to pass shape and stride explicitly, since DTensor might be + # sharded unevenly. + dtensor = DTensor.from_local( + new_local_tensor, + device_mesh=value.device_mesh, + placements=value.placements, + shape=value.size(), + stride=value.stride(), + ) + return dtensor + else: + return value + + def sharded_tensor_func(value: Any): + device = getattr(value, "device", None) + if device == torch.device("meta"): + raise RuntimeError( + f"Found unsupported type {type(value)} for meta device loading." + ) + else: + return value + + def tensor_func(value: torch.Tensor): + device = getattr(value, "device", None) + if device == torch.device("meta"): + device_type = dist.distributed_c10d._get_pg_default_device().type + device = cast( + torch.device, _get_device_module(device_type).current_device() + ) + tensor = torch.empty_like(value, device=device) + return tensor + else: + return value + + _iterate_state_dict( + state_dict, + dtensor_func, + sharded_tensor_func, + tensor_func, + ) + + +def _iterate_state_dict( + iter_object: Any, + dtensor_func: Callable, + sharded_tensor_func: Callable, + tensor_func: Callable, +): + """ + Iterate through the state dict, applying the given functions to each tensor type + and update the state dict in place. + + Args: + iter_object (Any): the target state_dict. + sharded_tensor_func (Callable): the function to apply to ShardedTensor + dtensor_func (Callable): the function to apply to DTensor + tensor_func (Callable): the function to apply to Tensor + + # TODO: let state_dict_util._iterate_state_dict() to support in place option + so we don't need to have two versions of _iterate_state_dict. + """ + + if isinstance(iter_object, DTensor): + return dtensor_func(iter_object) + elif isinstance(iter_object, ShardedTensor): + return sharded_tensor_func(iter_object) + elif isinstance(iter_object, torch.Tensor): + return tensor_func(iter_object) + elif ( + isinstance(iter_object, (int, float, str, bytes, io.BytesIO)) + or iter_object is None + ): + return iter_object + elif isinstance(iter_object, dict): + for key, value in iter_object.items(): + iter_object[key] = _iterate_state_dict( + value, dtensor_func, sharded_tensor_func, tensor_func + ) + return iter_object + elif isinstance(iter_object, (list, tuple)): + ret = [ + _iterate_state_dict(v, dtensor_func, sharded_tensor_func, tensor_func) + for v in iter_object + ] + if isinstance(iter_object, tuple): + ret = tuple(ret) # type: ignore[assignment] + return ret diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/resharding.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/resharding.py new file mode 100644 index 0000000000000000000000000000000000000000..0e5153df8da0adbe58b2c525c233ce805cf8b1c1 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/resharding.py @@ -0,0 +1,72 @@ +# mypy: allow-untyped-defs +from typing import List, Tuple + +from torch.distributed.checkpoint.metadata import ChunkStorageMetadata + + +__all__: List[str] = [] + + +def _check_shard_metadata_pair_overlap( + shard1: ChunkStorageMetadata, shard2: ChunkStorageMetadata +): + """Check if two shards overlap.""" + # For each dim of each shard, check if one shard resides on the other + # end of second shard with respect to that dim. As an example for a 2D + # shard, we would check if one shard is above or on the left of the + # other shard. + ndims = len(shard1.offsets) + for i in range(ndims): + if shard1.offsets[i] >= shard2.offsets[i] + shard2.sizes[i]: + return False + if shard2.offsets[i] >= shard1.offsets[i] + shard1.sizes[i]: + return False + + return True + + +def _shards_get_overlap_region_wrt_saved_tensor( + saved_shard: ChunkStorageMetadata, current_shard: ChunkStorageMetadata +) -> List[Tuple[int, int, int, int]]: + """ + Return the overlapping region between saved_shard and current_shard. + + There returned list has the same number of elements as the tensor's dimension. + For each element, we produce a tuple with the following contents: + (dimension, `saved_shard` offset, `current_shard` offset, length) + + Offsets are relative to each shard. + """ + narrows = [] + for dim, ( + saved_shard_offset, + current_shard_offset, + saved_shard_size, + current_shard_size, + ) in enumerate( + zip( + saved_shard.offsets, + current_shard.offsets, + saved_shard.sizes, + current_shard.sizes, + ) + ): + min_range_end = min( + saved_shard_offset + saved_shard_size, + current_shard_offset + current_shard_size, + ) + + length = min_range_end - max(current_shard_offset, saved_shard_offset) + + if saved_shard_offset > current_shard_offset: + offset_for_saved_tensor = 0 + offset_for_current_tensor = saved_shard_offset - current_shard_offset + else: + offset_for_saved_tensor = current_shard_offset - saved_shard_offset + offset_for_current_tensor = 0 + + narrows.append( + (dim, offset_for_saved_tensor, offset_for_current_tensor, length) + ) + + return narrows diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/staging.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/staging.py new file mode 100644 index 0000000000000000000000000000000000000000..22eed63c1cc517bc18f777d8943045115b440f06 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/staging.py @@ -0,0 +1,117 @@ +from typing import Optional, runtime_checkable +from typing_extensions import Protocol + +from torch.distributed._state_dict_utils import ( + _copy_state_dict, + _create_cpu_state_dict, + _offload_state_dict_to_cpu, +) +from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE + + +__all__ = ["AsyncStager", "BlockingAsyncStager"] + + +@runtime_checkable +class AsyncStager(Protocol): + """ + This protocol is meant to provide customization and extensibility for dcp.async_save, allowing users + to customize how data is staged previous to executing the usual dcp.save path in parallel. + The expected order of operations (concretely defined in `torch.distributed.state_dict_saver.async_save`) + is the following: + + 1. AsyncStager.stage_data(state_dict): + This call gives the AsyncStager the opportunity to 'stage' + the state_dict. The expectation and purpose of staging in this context is to create a "training-safe" + representation of the state dict, meaning that any updates to module data after staging is complete + should not be reflected in the state dict returned from this method. For example, in the default + case a copy of the entire state dict is created on CPU RAM and returned here, allowing users + to continue training without risking changes to data which is being serialized. + + 2. dcp.save is called on the state_dict returned from stage in parallel. This call is responsible + for serializing the state_dict and writing it to storage. + + 3. If AsyncStager.should_synchronize_after_execute is True, this method will be called immediately after + the serialization thread starts and before returning from dcp.async_save. If this is set to False, + the assumption is the user has defined a custom synchronization point for the the purpose of further + optimizing save latency in the training loop (for example, by overlapping staging with the + forward/backward pass), and it is the respondsibility of the user to call `AsyncStager.synchronize_staging` + at the appropriate time. + + """ + + # default to True since the common case is to stage synchronously + _synchronize_after_execute: bool = True + + @property + def should_synchronize_after_execute(self) -> bool: + """ + Whether to synchronize after executing the stage. + """ + + return self._synchronize_after_execute + + def stage(self, state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE: + """ + Returns a "staged" copy of `state_dict`. The expectation of the staged copy is that it is + innoculated from any updates incurred after the stage call is complete. + """ + raise NotImplementedError( + f"{self.__class__.__name__} must implement stage method" + ) + + def synchronize_staging(self) -> None: + """ + In the case `stage` is async in some way, this method should be called to ensure staging + is complete and it is safe to begin modifying the original `state_dict` + """ + + +class BlockingAsyncStager(AsyncStager): + """ + An implementation of AsyncStager which stages the state_dict on CPU RAM and blocks until the copy is complete. + This implementation also provides an option to optimize stage latency using pinned memory. + + N.B. synchronize_staging is a no-op in this case. + + + """ + + # default to True since the common case is to stage synchronously + _synchronize_after_execute: bool = False + + def __init__( + self, + cache_staged_state_dict: bool = False, + type_check: bool = False, + ): + """ + Initializes the BlockingAsyncStager. + + Args: + cache_staged_state_dict: Whether to cache the staged state_dict. This option decreases staging latency + at the cost of increases memory usage. Additionally, if this parameter is set to True, it's the expectation + that the stager is maintained and re-used for multiple dcp.async_save calls. Default to False. + type_check: Whether to perform a type check during cpu_offload. Defaults to False. + + """ + self.cache_staged_state_dict = cache_staged_state_dict + self.type_check = type_check + self.state_dict_cache: Optional[STATE_DICT_TYPE] = None + + def stage(self, state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE: + """ + Returns a copy of `state_dict` on the CPU. + """ + + if not self.cache_staged_state_dict: + return _offload_state_dict_to_cpu(state_dict, type_check=self.type_check) + + if self.state_dict_cache is None: + self.state_dict_cache = _create_cpu_state_dict(state_dict, pin_memory=True) + return _copy_state_dict(state_dict, self.state_dict_cache) + + def synchronize_staging(self) -> None: + """ + No-op function, since staging is blocking. + """ diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict.py new file mode 100644 index 0000000000000000000000000000000000000000..aae3f48165f3e7d12bcb4b9199f3f2d993f8184a --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict.py @@ -0,0 +1,1426 @@ +# mypy: allow-untyped-defs +import contextlib +import functools +import gc +import warnings +from dataclasses import asdict, dataclass, field +from itertools import chain +from typing import ( + Any, + Callable, + cast, + Dict, + Generator, + Iterable, + List, + no_type_check, + Optional, + Set, + Tuple, + Union, +) + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed._shard.sharded_tensor import ShardedTensor +from torch.distributed._state_dict_utils import ( + _broadcast_state_dict, + _distribute_state_dict, + _flatten_state_dict, + _gather_state_dict, + _offload_state_dict_to_cpu, + _unflatten_state_dict, +) +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + _CHECKPOINT_PREFIX, +) +from torch.distributed.fsdp import ( + FullOptimStateDictConfig, + FullStateDictConfig, + FullyShardedDataParallel as FSDP, + OptimStateDictConfig, + ShardedOptimStateDictConfig, + ShardedStateDictConfig, + StateDictConfig, + StateDictType, +) +from torch.distributed.fsdp._common_utils import ( + _get_module_fsdp_state_if_fully_sharded_module, + FSDP_WRAPPED_MODULE, +) +from torch.distributed.tensor import DTensor +from torch.nn.modules.module import _IncompatibleKeys +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils._pytree import tree_map_only + + +__all__ = [ + "FQNS_T", + "PrimitiveType", + "ValueType", + "DictValueType", + "ListDictValueType", + "OptimizerStateType", + "StateDictOptions", + "get_model_state_dict", + "get_optimizer_state_dict", + "get_state_dict", + "set_model_state_dict", + "set_optimizer_state_dict", + "set_state_dict", +] + + +_FLAT_PARAM = "_flat_param" +_PG = "param_groups" +_PARAMS = "params" +_STATE = "state" + +FQNS_T = Set[str] +PrimitiveType = Union[DTensor, ShardedTensor, torch.Tensor, int, float, str] +ValueType = Union[ + PrimitiveType, List[PrimitiveType], Tuple[PrimitiveType], Dict[str, "ValueType"] +] +DictValueType = Dict[str, ValueType] +ListDictValueType = List[DictValueType] +OptimizerStateType = Dict[str, Union[DictValueType, ListDictValueType]] + + +_patched_state_dict: Set[Callable] = set() + + +@contextlib.contextmanager +def _gc_context(): + is_enabled = gc.isenabled() + gc.disable() + try: + yield + finally: + if is_enabled: + gc.enable() + + +@dataclass +class StateDictOptions: + """ + This dataclass specifies how get_state_dict/set_state_dict will work. + + - ``full_state_dict``: if this is set to True, all the tensors in the + returned state_dict will be gathered. No ShardedTensor and DTensor + will be in the returned state_dict. + + - ``cpu_offload``: offload all the tensors to cpu. To prevent CPU OOM, if + ``full_state_dict`` is also true, then only the rank0 will get the + state_dict and all other ranks will get empty state_dict. + + - ``ignore_frozen_params``: if the value is True, the returned state_dict + won't contain any frozen parameters -- the ``requires_grad`` is False. + The default value is False. + + - ``keep_submodule_prefixes`` (deprecated): when ``submodules`` is not None, this option + indicates whether to keep the submodule prefixes from the state_dict keys. + or example, if the submodule is ``module.pretrain`` and the full FQN of + the parameter is ``pretrain.layer1.weight`` of the param. When this option + is True, the parameter's key in the returned state_dict will be + ``pretrain.layer1.weight``. If the options is False, the key will be + ``layer1.weight``. + Note that if ``keep_submodule_prefixes`` is False, there may be conflicted + FQNs, hence there should be only one submodule in ``submodules``. + + - ``strict``: the ``strict`` option when ``set_state_dict`` calls + model.load_state_dict(). + + - ``broadcast_from_rank0``: when the option is True, rank0 should receive a + full state_dict and will broadcast the tensors in the state_dict/ + optim_state_dict one by one to other ranks. Other ranks will receive + the tensors and shard according to the local shards in the model and + optimizer. ``full_state_dict`` must be set to True when using this option. + This option currently only supports DTensor, not the legacy ShardedTensor. + """ + + full_state_dict: bool = False + cpu_offload: bool = False + ignore_frozen_params: bool = False + keep_submodule_prefixes: bool = True + strict: bool = True + broadcast_from_rank0: bool = False + flatten_optimizer_state_dict: bool = False + + +@dataclass +class _StateDictInfo(StateDictOptions): + fqn_param_mapping: Dict[ + Union[str, torch.Tensor], Union[FQNS_T, torch.Tensor] + ] = field(default_factory=dict) + shared_params_mapping: Dict[ + Union[str, torch.Tensor], Union[FQNS_T, torch.Tensor] + ] = field(default_factory=dict) + submodule_prefixes: Set[str] = field(default_factory=set) + handle_model: bool = True + handle_optim: bool = True + fsdp_context: Callable = contextlib.nullcontext + fsdp_modules: List[nn.Module] = field(default_factory=list) + + +@functools.lru_cache(maxsize=None) +def _get_fqns( + model: nn.Module, + name: str, + skip_ddp_prefix: bool = True, + skip_compiler_prefix: bool = True, +) -> FQNS_T: + """ + This API is used to convert the name of a parameter to the FQNs. For FSDP + without `use_orig_params`, the name of FlatParameter can be mapped to + multiple original parameters. As a result, the return type of this function + is `Set[str]`. + + Args: + module (nn.Module): the root model. + name (str): the name + skip_ddp_prefix (bool): whether to skip DDP's `module` prefix + + Returns: + The canonical FQNs based on the model traversal. + """ + + # Remove the checkpoint prefix, if it exists. + name = name.replace(_CHECKPOINT_PREFIX, "") + if "." not in name: + return {name} + + obj_names = name.split(".") + fqn_obj_names = [] + curr_obj = model + for i, curr_obj_name in enumerate(obj_names): + if isinstance(curr_obj, DDP): + assert curr_obj_name == "module" + curr_obj = curr_obj.module + if not skip_ddp_prefix: + fqn_obj_names.append(curr_obj_name) + elif isinstance(curr_obj, FSDP): + if i < len(obj_names) - 1 and obj_names[i + 1] == _FLAT_PARAM: + prefix = ".".join(fqn_obj_names) + flat_param = getattr(curr_obj, _FLAT_PARAM) + if prefix: + prefix = f"{prefix}." + return {f"{prefix}{fqn}" for fqn in flat_param._fqns} + curr_obj = getattr(curr_obj, FSDP_WRAPPED_MODULE) + if curr_obj_name != FSDP_WRAPPED_MODULE: + fqn_obj_names.append(curr_obj_name) + curr_obj = getattr(curr_obj, curr_obj_name) + elif isinstance(curr_obj, torch._dynamo.eval_frame.OptimizedModule): + assert curr_obj_name == "_orig_mod" + curr_obj = curr_obj._orig_mod + if not skip_compiler_prefix: + fqn_obj_names.append(curr_obj_name) + else: + fqn_obj_names.append(curr_obj_name) + if curr_obj_name == nn.modules.module._EXTRA_STATE_KEY_SUFFIX: + if i != len(obj_names) - 1: + raise RuntimeError("Expect `_extra_state` to be the last obj name") + else: + curr_obj = getattr(curr_obj, curr_obj_name) + + return {".".join(fqn_obj_names).replace(_CHECKPOINT_PREFIX, "")} + + +class _EXTRA_STATE: + pass + + +def _iterate_valid_model_state(model): + visited_modules: Set[nn.Module] = set() + + def recurse(module: nn.Module, curr_fqn: str) -> Generator: + visited_modules.add(module) + + curr_fqn = f"{curr_fqn}." if curr_fqn else "" + for name, submodule in module.named_children(): + if submodule in visited_modules: + continue + new_fqn = f"{curr_fqn}{name}" + yield from recurse(submodule, new_fqn) + + for name, obj in chain( + module.named_buffers(recurse=False), module.named_parameters(recurse=False) + ): + if name in module._non_persistent_buffers_set: + continue + new_fqn = f"{curr_fqn}{name}" + yield new_fqn, obj + + if ( + getattr(module.__class__, "get_extra_state", nn.Module.get_extra_state) + != nn.Module.get_extra_state + ): + new_fqn = f"{curr_fqn}{nn.modules.module._EXTRA_STATE_KEY_SUFFIX}" + yield new_fqn, _EXTRA_STATE() + + yield from recurse(model, "") + + +def _verify_options( + model: nn.Module, + optims: Tuple[torch.optim.Optimizer, ...], + optim_only: bool, + *, + submodules: Optional[Set[nn.Module]] = None, + options: Optional[StateDictOptions] = None, +) -> _StateDictInfo: + """ + Verify the model and options passed by the user and generates _StateDictInfo. + """ + if submodules: + warnings.warn( + "Getting submodules only model/optim state_dict is deprecated and " + "will be removed in 2.5. This feature can be achieved by manually " + "filtering out the state_dict returned from get_state_dict.", + FutureWarning, + ) + if optim_only and not optims: + raise RuntimeError( + "Optimizers are not passed in but optim_only is set to True." + ) + + options = options or StateDictOptions() + + fqn_param_mapping: Dict[ + Union[str, torch.Tensor], Union[Set[str], torch.Tensor] + ] = {} + shared_params_mapping: Dict[ + Union[str, torch.Tensor], Union[Set[str], torch.Tensor] + ] = {} + for name, param in _iterate_valid_model_state(model): + if isinstance(param, _EXTRA_STATE): + continue + + fqns = _get_fqns(model, name) + fqn = fqn_param_mapping.get(param, None) + if fqn is not None: + cast(Set[str], fqn_param_mapping[param]).update(fqns) + shared_params_mapping[param] = fqn_param_mapping[param] + else: + # We need to do copy as _get_fqns is lru_cached + fqn_param_mapping[param] = fqns.copy() + for fqn in fqns: + if not isinstance(param, _EXTRA_STATE): + fqn_param_mapping[fqn] = param + + for param_, fqns_ in list(shared_params_mapping.items()): + for fqn in fqns_: + shared_params_mapping[fqn] = cast(torch.Tensor, param_) + + submodule_prefixes: Set[str] = set() + if submodules: + submodules = set(submodules) + for name, module in model.named_modules(): + if module not in submodules: + continue + fqns = _get_fqns(model, name) + assert len(fqns) == 1, "Submodule FQN should only have 1 instance" + submodule_prefixes.update(f"{fqn}." for fqn in fqns) + + if options.broadcast_from_rank0 and not options.full_state_dict: + raise ValueError( + "full_state_dict must be True when broadcast_from_rank0 is True." + ) + fsdp_modules = FSDP.fsdp_modules(model) + state_dict_config: StateDictConfig + optim_state_dict_config: OptimStateDictConfig + fsdp_context: Callable + if fsdp_modules: + # FSDP API only work if at least one FSDP instance exists. + if options.full_state_dict: + state_dict_config = FullStateDictConfig( + offload_to_cpu=options.cpu_offload, rank0_only=options.cpu_offload + ) + optim_state_dict_config = FullOptimStateDictConfig( + offload_to_cpu=options.cpu_offload, + rank0_only=(options.cpu_offload or options.broadcast_from_rank0), + ) + state_dict_type = StateDictType.FULL_STATE_DICT + else: + state_dict_config = ShardedStateDictConfig( + offload_to_cpu=options.cpu_offload, + ) + optim_state_dict_config = ShardedOptimStateDictConfig( + offload_to_cpu=options.cpu_offload, + ) + state_dict_type = StateDictType.SHARDED_STATE_DICT + + @contextlib.contextmanager + def fsdp_state_dict_type_without_warning( + module, + state_dict_type, + state_dict_config, + optim_state_dict_config, + ): + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", message="FSDP.state_dict_type", category=FutureWarning + ) + with FSDP.state_dict_type( + module=module, + state_dict_type=state_dict_type, + state_dict_config=state_dict_config, + optim_state_dict_config=optim_state_dict_config, + ): + yield + + fsdp_context = functools.partial( + fsdp_state_dict_type_without_warning, + module=model, + state_dict_type=state_dict_type, + state_dict_config=state_dict_config, + optim_state_dict_config=optim_state_dict_config, + ) + else: + fsdp_context = contextlib.nullcontext + + return _StateDictInfo( + **asdict(options), + fqn_param_mapping=fqn_param_mapping, + shared_params_mapping=shared_params_mapping, + submodule_prefixes=submodule_prefixes, + fsdp_context=fsdp_context, + fsdp_modules=cast(List[nn.Module], fsdp_modules), + handle_model=not optim_only, + handle_optim=(len(optims) > 0), + ) + + +def _verify_state_dict( + model_state_dict: Dict[str, ValueType], + optim_state_dict: OptimizerStateType, + info: _StateDictInfo, +) -> None: + for module in info.fsdp_modules: + fsdp_state = _get_module_fsdp_state_if_fully_sharded_module(module) + assert fsdp_state is not None, "Expected a fsdp_state with a fsdp module." + + # Verify if the model_state_dict and optim_state_dict are valid. This API + # should give the users an explicit error message to debug or report. + if ( + info.handle_model + and not model_state_dict + and not info.submodule_prefixes + and not info.ignore_frozen_params + and not (info.cpu_offload and info.full_state_dict) + and info.strict + and not info.broadcast_from_rank0 + ): + raise RuntimeError( + "The option indicates that model state_dict is required to save " + "or load, but model state_dict is empty." + f"rank = {dist.get_rank()=}." + ) + + if info.handle_optim: + if ( + not optim_state_dict + and not (info.cpu_offload and info.full_state_dict) + and (not info.broadcast_from_rank0) + ): + raise RuntimeError( + "The option indicates that model state_dict is required to save, " + f"or load but optim state_dict is empty. {optim_state_dict}" + ) + + for key in model_state_dict.keys(): + if _FLAT_PARAM in key: + raise RuntimeError( + f"{key} contains {_FLAT_PARAM}. This can happen if the model " + "is not the root module." + ) + + +def _state_dict_fn(obj: Union[nn.Module, torch.optim.Optimizer], api: str) -> Callable: + call = getattr(obj, api) + if call in _patched_state_dict: + call = functools.partial(getattr(obj.__class__, api), self=obj) + return call + + +def _maybe_full_or_cpu_state_dict( + state_dict: Dict[str, Any], info: _StateDictInfo +) -> Dict[str, Any]: + if info.full_state_dict: + ranks_only = ( + () + if (not info.cpu_offload or not torch.distributed.is_initialized()) + else (0,) + ) + return _gather_state_dict( + state_dict, cpu_offload=info.cpu_offload, ranks_only=ranks_only + ) + elif info.cpu_offload: + return _offload_state_dict_to_cpu(state_dict) + else: + return state_dict + + +@torch.no_grad() +def _get_model_state_dict( + model: nn.Module, info: _StateDictInfo +) -> Dict[str, ValueType]: + if not info.handle_model: + return {} + + with info.fsdp_context(): + state_dict = _state_dict_fn(model, "state_dict")() + + for key in list(state_dict.keys()): + fqns = _get_fqns(model, key) + assert len(fqns) == 1, (key, fqns) + fqn = next(iter(fqns)) + if fqn != key: + # As we only support FSDP, DDP, and TP, the only cases are + # wrapper-based DDP and compiler. Verify if the assumption + # is correct. + def verify(key, fqn) -> bool: + if len(fqn) >= len(key): + return False + fqn_split = fqn.split(".") + key_split = key.split(".") + fqn_idx = 0 + for key_idx, key_name in enumerate(key_split): + if key_name == fqn_split[fqn_idx]: + fqn_idx += 1 + if fqn_idx == len(fqn_split): + return key_idx == len(key_split) - 1 + elif key_name in ("module", "_orig_mod"): + continue + else: + return False + return True + + if not verify(key, fqn): + raise RuntimeError(f"An unexpected key, {key}, exists. FQN is {fqn}") + state_dict[fqn] = state_dict.pop(key) + + if info.submodule_prefixes: + new_state_dict: Dict[str, ValueType] = {} + # TODO: make this faster. + for fqn in state_dict.keys(): + for prefix in info.submodule_prefixes: + if not fqn.startswith(prefix): + continue + if info.keep_submodule_prefixes: + new_state_dict[fqn] = state_dict[fqn] + else: + new_fqn = fqn[len(prefix) :] + new_state_dict[new_fqn] = state_dict[fqn] + state_dict = new_state_dict + + if info.ignore_frozen_params: + for key, param in model.named_parameters(): + if param.requires_grad: + continue + fqns = _get_fqns(model, key) + for fqn in fqns: + state_dict.pop(fqn) + + for key, p in list(state_dict.items()): + if torch.is_tensor(p) and p.is_meta: + state_dict.pop(key) + + return _maybe_full_or_cpu_state_dict(state_dict, info) + + +@torch.no_grad() +def _load_model_state_dict( + model: nn.Module, + state_dict: Dict[str, ValueType], + info: _StateDictInfo, +) -> _IncompatibleKeys: + if not info.handle_model or (not state_dict and not info.broadcast_from_rank0): + return _IncompatibleKeys({}, {}) + + local_state_dict = {} + for key, value in _iterate_valid_model_state(model): + fqns = _get_fqns(model, key) + fqns_with_prefix = _get_fqns( + model, key, skip_ddp_prefix=False, skip_compiler_prefix=False + ) + + for fqn, fqn_with_prefix in zip(fqns, fqns_with_prefix): + if ( + not info.broadcast_from_rank0 or dist.get_rank() == 0 + ) and fqn != fqn_with_prefix: + state_dict[fqn_with_prefix] = state_dict.pop(fqn) + local_state_dict[fqn_with_prefix] = value + + assign = False + if info.broadcast_from_rank0 or info.full_state_dict: + device = None + for key, value in local_state_dict.items(): + if torch.is_tensor(value) and value.dim() > 0: + if device is None: + device = value.device + else: + assert device == value.device + assert device is not None + if device == torch.device("meta"): + device = dist.distributed_c10d._get_pg_default_device() + assign = True + if info.broadcast_from_rank0: + _broadcast_state_dict( + state_dict, local_state_dict, device=device, strict=info.strict + ) + elif info.full_state_dict: + _distribute_state_dict(state_dict, local_state_dict, device=device) + for fqn, local_state in local_state_dict.items(): + state_dict[fqn] = local_state + + with info.fsdp_context(): + return cast( + _IncompatibleKeys, + _state_dict_fn(model, "load_state_dict")( + state_dict=state_dict, strict=info.strict, assign=assign + ), + ) + + +def _init_optim_state(optim: torch.optim.Optimizer) -> None: + """ + Initialize optim states by calling the step() with zero grads. + """ + if optim.state: + # The optimizer state is initialized. + return + + # There are some stateless optimizers like SGD. These optimizer will + # not return in the above condition. So if gradients exist, we should also + # return. If gradients do not exist, the following initialization should + # not disturb SGD because the gradients and lr are both zero. + for param_group in optim.param_groups: + for param in param_group[_PARAMS]: + if param.grad is not None: + return + + for param_group in optim.param_groups: + for param in param_group[_PARAMS]: + if param.requires_grad: + param.grad = torch.zeros_like(param) + + # Some optimizers will update parameters regardless of grads due to lr, so + # make lr to zero when calling `step()`. + lrs = [] + for param_group in optim.param_groups: + if "lr" in param_group: + lrs.append(param_group["lr"]) + param_group["lr"] = ( + torch.tensor(0.0) + if isinstance(param_group["lr"], torch.Tensor) + else 0.0 + ) + optim.step(closure=None) + # Whether to recover the "lr" should not matter too much as we will + # restore checkpointing later. + for param_group in optim.param_groups: + if "lr" in param_group: + param_group["lr"] = lrs.pop(0) + optim.zero_grad(set_to_none=True) + + +def _flatten_optim_state_dict(state_dict: OptimizerStateType) -> Dict[str, ValueType]: + """ + This API flattens the optimizer state_dict to support optimizer resharding for + MPMD, e.g., pipeline parallelism. + + Without the API, the original optimizer state_dict looks like: + { + "state": { + "layer1.weight": { + "step": 10, "exp_avg": SomeTensor, "exp_avg_sq": SomeTensor + }, + "layer2.weight": { + "step": 10, "exp_avg": SomeTensor, "exp_avg_sq": SomeTensor + }, + }, + "param_group": [ + { + "lr": 0.0, + "betas": (0.9, 0.95), ..., + "params": ["layer1.weight", "layer2.weight"] + } + ] + } + + With this API, the optimizer state_dict looks like: + { + "state.layer1.weight.step": 10, + "state.layer2.weight.step": 10, + "state.layer1.weight.exp_avg": SomeTensor, + "state.layer2.weight.exp_avg": SomeTensor, + "state.layer1.weight.exp_avg_sq": SomeTensor, + "state.layer2.weight.exp_avg_sq": SomeTensor, + "param_group.layer1.weight.lr" : 0.1, + "param_group.layer2.weight.lr" : 0.1, + "param_group.layer1.weight.betas" : (0.9, 0.95), + "param_group.layer2.weight.betas" : (0.9, 0.95), + } + + Note that if any of the value is a container, like the betas in the example, + this API won't flattent it. + """ + + def _raise_if_type_not_supported(v): + if not isinstance(v, (torch.Tensor, int, float)): + raise NotImplementedError( + "Flattening optimizer state_dict only supports " + "tensor, int, float states now. " + f"Type is {type(v)}." + ) + + ret: Dict[str, ValueType] = {} + for fqn, state in cast(DictValueType, state_dict[_STATE]).items(): + for k, v in cast(DictValueType, state).items(): + _raise_if_type_not_supported(v) + ret[f"{_STATE}.{fqn}.{k}"] = v + + for param_group in cast(ListDictValueType, state_dict[_PG]): + fqns = param_group.pop(_PARAMS) + for fqn in cast(List[str], fqns): + for k, v in param_group.items(): + ret[f"{_PG}.{fqn}.{k}"] = v + return ret + + +def _unflatten_optim_state_dict( + optim: torch.optim.Optimizer, + state_dict: Dict[str, ValueType], + info: _StateDictInfo, +) -> OptimizerStateType: + """ + This API unflattens the state_dict generated by _flatten_optim_state_dict(). + See the docstring of _flatten_optim_state_dict() for more detail. + """ + state: DictValueType = {} + pg_state: ListDictValueType = [] + return_osd: OptimizerStateType = {_STATE: state, _PG: pg_state} + + for param_group in optim.param_groups: + pg_state.append({_PARAMS: []}) + for param in param_group[_PARAMS]: + for fqn in info.fqn_param_mapping[param]: + params = pg_state[-1][_PARAMS] + assert isinstance(params, list) # typing + params.append(fqn) + if not param.requires_grad: + continue + state[fqn] = {} + for state_name in optim.state[param].keys(): + cast(DictValueType, state[fqn])[state_name] = state_dict[ + f"{_STATE}.{fqn}.{state_name}" + ] + + first_param_fqn = cast(List[str], pg_state[-1][_PARAMS])[0] + for k in param_group.keys(): + if k == _PARAMS: + continue + value = state_dict[f"{_PG}.{first_param_fqn}.{k}"] + if k not in pg_state[-1]: + pg_state[-1][k] = value + elif pg_state[-1][k] != value: + raise RuntimeError( + "All the parameters in the same parameter group should have " + f"the same saved param_group value. But {first_param_fqn}.{k} " + f"is {value} while other(s) is {pg_state[-1][k]}." + ) + + return return_osd + + +@torch.no_grad() +def _get_optim_state_dict( + model: nn.Module, + optimizers: Tuple[torch.optim.Optimizer, ...], + info: _StateDictInfo, +) -> OptimizerStateType: + if not info.handle_optim: + return {} + + optim_state_dict: OptimizerStateType = {_STATE: {}, _PG: []} + for optim in optimizers: + _init_optim_state(optim) + osd = _state_dict_fn(optim, "state_dict")() + if info.fsdp_modules: + with info.fsdp_context(): + osd = FSDP.optim_state_dict(model, optim, osd) + + # We need to specially handle FlatParameter FSDP as + # FlatParameter FSDP converts the FQNs. + # There are no easy ways to do this conversion systematically. + # We can only use a string replacment without correctness check. + if not osd: + continue + for k in list(osd[_STATE].keys()): + if "_orig_mod" in k: + osd[_STATE][k.replace("_orig_mod.", "")] = osd[_STATE].pop(k) + for g in osd[_PG]: + params = [k.replace("_orig_mod.", "") for k in g[_PARAMS]] + g[_PARAMS] = params + else: + params = list(chain.from_iterable(g[_PARAMS] for g in optim.param_groups)) + param_pid_mapping = dict(zip(params, range(len(params)))) + fqn_pid_mapping = {} + for key, param in model.named_parameters(): + fqns = _get_fqns(model, key) + assert len(fqns) == 1 + fqn = next(iter(fqns)) + if param not in param_pid_mapping: + continue + pid = param_pid_mapping[param] + fqn_pid_mapping[fqn] = pid + fqn_pid_mapping[pid] = fqn + + for key in list(osd[_STATE].keys()): + fqn = fqn_pid_mapping[key] + osd[_STATE][fqn] = osd[_STATE].pop(key) + + for group in osd[_PG]: + group[_PARAMS] = [fqn_pid_mapping[pid] for pid in group[_PARAMS]] + + if not osd: + continue + + cast(DictValueType, optim_state_dict[_STATE]).update(osd[_STATE]) + cast(ListDictValueType, optim_state_dict[_PG]).extend(osd[_PG]) + + if info.flatten_optimizer_state_dict: + optim_state_dict = cast( + OptimizerStateType, _flatten_optim_state_dict(optim_state_dict) + ) + + return _maybe_full_or_cpu_state_dict(optim_state_dict, info) + + +def _split_optim_state_dict( + model: nn.Module, + optim: torch.optim.Optimizer, + optim_state_dict: OptimizerStateType, + info: _StateDictInfo, +) -> OptimizerStateType: + """ + Extract the corresponding optim state_dict from ``optim_state_dict`` for + ``optim`` and return the result optim state_dict. + + Args: + model (nn.Module): the root model. + optim (torch.optim.Optimizer): the optimizer. + optim_state_dict (Dict[str, ValueType]): the superset optim state_dict that + contains the optim state_dict of ``optim``. + info (_StateDictInfo): state dict information. + + Returns: + The optim state_dict of ``optim``. + """ + + state: DictValueType = {} + pg_state: ListDictValueType = [] + return_osd: OptimizerStateType = {_STATE: state, _PG: pg_state} + pg_mapping: Dict[int, int] = {} + + if all( + isinstance(k, int) for k in cast(DictValueType, optim_state_dict[_STATE]).keys() + ): + return optim_state_dict + + for param_group in optim.param_groups: + pg_state.append({_PARAMS: []}) + for param in param_group[_PARAMS]: + for fqn in info.fqn_param_mapping[param]: + if fqn in info.shared_params_mapping: + in_params = False + for loaded_param_group in cast( + ListDictValueType, optim_state_dict[_PG] + ): + if fqn in cast(List[str], loaded_param_group[_PARAMS]): + in_params = True + break + else: + in_params = True + if not in_params: + continue + + params = pg_state[-1][_PARAMS] + assert isinstance(params, list) + params.append(fqn) + if param.requires_grad: + state[fqn] = cast(DictValueType, optim_state_dict[_STATE])[fqn] + for loaded_param_group in cast( + ListDictValueType, optim_state_dict[_PG] + ): + if fqn in cast(List[str], loaded_param_group[_PARAMS]): + pg_mapping[id(loaded_param_group)] = len(return_osd[_PG]) - 1 + + for param_group in cast(ListDictValueType, optim_state_dict[_PG]): + idx = pg_mapping.get(id(param_group), -1) + if idx == -1: + continue + for key, value in param_group.items(): + if key == _PARAMS: + continue + # TODO: check if value is the same if exists. + pg_state[idx][key] = value + + return return_osd + + +@torch.no_grad() +def _load_optim_state_dict( + model: nn.Module, + optimizers: Tuple[torch.optim.Optimizer, ...], + state_dict: OptimizerStateType, + info: _StateDictInfo, +) -> None: + if not info.handle_optim: + return + + for optim in optimizers: + _init_optim_state(optim) + if state_dict: + if _STATE in state_dict: + optim_state_dict = _split_optim_state_dict( + model, optim, state_dict, info + ) + else: + optim_state_dict = _unflatten_optim_state_dict( + optim, cast(Dict[str, ValueType], state_dict), info + ) + else: + optim_state_dict = {} + if info.fsdp_modules: + # We need to specially handle FlatParameter FSDP as + # FlatParameter FSDP converts the FQNs. + for original_fqn, _ in model.named_parameters(): + fqns = _get_fqns(model, original_fqn) + fqns_with_compiler = _get_fqns( + model, original_fqn, skip_compiler_prefix=False + ) + if fqns == fqns_with_compiler: + continue + + assert len(fqns) == 1 + fqn = fqns.pop() + fqn_with_compiler = fqns_with_compiler.pop() + for g in optim_state_dict[_PG]: + val = cast(Dict[str, Any], g) + params = [ + key.replace(fqn, fqn_with_compiler) for key in val[_PARAMS] + ] + val[_PARAMS] = params + osd_state = cast(DictValueType, optim_state_dict[_STATE]) + for k in list(osd_state.keys()): + if fqn in k: + osd_state[k.replace(fqn, fqn_with_compiler)] = osd_state.pop(k) + + with info.fsdp_context(): + optim_state_dict = FSDP.optim_state_dict_to_load( + model, optim, optim_state_dict + ) + elif info.full_state_dict: + info.full_state_dict = False + local_state_dict = _get_optim_state_dict(model, (optim,), info) + info.full_state_dict = True + device = None + + def _device(t): + if t.dim() > 0: + nonlocal device + if device is None: + device = t.device + elif device != t.device: + raise ValueError("Device mismatch") + return t + + _ = tree_map_only(torch.Tensor, _device, local_state_dict) + assert device is not None + flatten_osd, osd_mapping = _flatten_state_dict(optim_state_dict) + flatten_local_osd, local_osd_mapping = _flatten_state_dict(local_state_dict) + if info.broadcast_from_rank0: + _broadcast_state_dict(flatten_osd, flatten_local_osd, device=device) + else: + _distribute_state_dict(flatten_osd, flatten_local_osd, device=device) + # The modifications listed seek to address the problem where optim might possess + # dissimilar parameters in comparison to optim_state_dict. This is achieved by + # incorporating differential parameters within local, which may result in optim + # having additional parameters ultimately. + for optim_key in flatten_osd.keys(): + if optim_key not in flatten_local_osd: + assert optim_key in osd_mapping + flatten_local_osd[optim_key] = flatten_osd[optim_key] + local_osd_mapping[optim_key] = osd_mapping[optim_key] + optim_state_dict = _unflatten_state_dict( + flatten_local_osd, local_osd_mapping + ) + + # Note that we do not have to convert the FQN back to param id here if + # order in optim.param_groups[idx][_PARAMS] is the same as the one in + # optim_state_dict[_PG][idx][_PARAMS]. + _state_dict_fn(optim, "load_state_dict")(state_dict=optim_state_dict) + + +def get_model_state_dict( + model: nn.Module, + *, + submodules: Optional[Set[nn.Module]] = None, + options: Optional[StateDictOptions] = None, +) -> Dict[str, ValueType]: + """ + Return the model state_dict of ``model``. + + See ``get_state_dict`` for the detail usage. + + Args: + model (nn.Module): the nn.Module to the model. + submodules (deprecated): Optional[Set[nn.Module]]: only return the model parameters + that belong to the submodules. + options (StateDictOptions): the options to control how + model state_dict and optimizer state_dict should be returned. See + `StateDictOptions` for the details. + + Returns: + The state_dict for ``model``. + + :rtype: typing.Dict[str, ValueType] + """ + with _gc_context(): + info = _verify_options( + model, + (), + optim_only=False, + submodules=submodules, + options=options, + ) + model_state_dict = _get_model_state_dict(model, info) + _verify_state_dict(model_state_dict, {}, info) + return model_state_dict + + +def get_optimizer_state_dict( + model: nn.Module, + optimizers: Union[torch.optim.Optimizer, Iterable[torch.optim.Optimizer]], + *, + submodules: Optional[Set[nn.Module]] = None, + options: Optional[StateDictOptions] = None, +) -> OptimizerStateType: + """ + Return the combined state_dict for optimizers. + + See ``get_state_dict`` for the detail usage. + + Args: + model (nn.Module): the nn.Module to the model. + optimizers (Union[None, Optimizer, Iterable[Optimizer]]): + The optimizers that are used to optimize ``model``. + submodules (deprecated): Optional[Set[nn.Module]]: only return the model parameters + that belong to the submodules. + options (StateDictOptions): the options to control how + model state_dict and optimizer state_dict should be returned. See + `StateDictOptions` for the details. + + Returns: + The state_dict for ``optimizers``. + + :rtype: OptimizerStateType + """ + with _gc_context(): + optimizers = ( + (optimizers,) + if isinstance(optimizers, torch.optim.Optimizer) + else tuple(optimizers) + ) + info = _verify_options( + model, + optimizers, + optim_only=True, + submodules=submodules, + options=options, + ) + optim_state_dict = _get_optim_state_dict(model, optimizers, info) + _verify_state_dict({}, optim_state_dict, info) + return optim_state_dict + + +def get_state_dict( + model: nn.Module, + optimizers: Union[torch.optim.Optimizer, Iterable[torch.optim.Optimizer]], + *, + submodules: Optional[Set[nn.Module]] = None, + options: Optional[StateDictOptions] = None, +) -> Tuple[Dict[str, ValueType], OptimizerStateType]: + """ + Return the model state_dict and optimizers state_dict. + + ``get_state_dict`` can process any module that is parallelized by PyTorch + FSDP/fully_shard, DDP/replicate, tensor_parallel/parallelize_module, and any + combination of these parallelisms. The main functions of ``get_state_dict`` + are: 1.) returning a model and optimizer state_dict that can be resharded + with a different number of trainers and/or different parallelisms. + 2.) hiding the parallelism-specific state_dict APIs. Users don't have to call + these APIs. + 3.) sanity checking the result state_dict. + + The keys of the result state dictionary are the canonical FQNs (Fully + Qualified Names). A canonical FQN refers to the FQN based on a parameter's + position in an nn.Module hierarchy. More specifically, a canonical FQN to a + parameter is the FQN returned by ``module.named_parameters()`` or + ``module.named_buffers()`` when the module is not distributed by any + parallelisms. Since the optimizer internally uses parameter IDs to represent + a parameter, there will be a conversion from the parameter IDs to the + canonical FQNs when calling this API. + + ``get_state_dict`` can also process a module that is not parallelized. In + such a case, ``get_state_dict`` only performs one function -- converting the + optimizer parameter IDs to the canonical FQNs. + + Example: + >>> # xdoctest: +SKIP + >>> import torch + >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + >>> from torch.nn.parallel import DistributedDataParallel as DDP + >>> from torch.distributed.checkpoint.state_dict import get_state_dict + + >>> fsdp_model = FSDP(copy.deepcopy(model)) + >>> fsdp_optim = torch.optim.Adam(model.parameters(), lr=1e-3) + >>> ddp_model = DDP(copy.deepcopy(model)) + >>> ddp_optim = torch.optim.Adam(model.parameters(), lr=1e-3) + + + >>> ddp_state_dict, ddp_optim_state_dict = get_state_dict(ddp_model, ddp_optim) + >>> fsdp_state_dict, fsdp_optim_state_dict = get_state_dict(fsdp_model, fsdp_optim) + + >>> # if we simply call ddp_model.state_dict() and fsdp_model.state_dict(), + >>> # the asserts will fail. + >>> assert ddp_state_dict == fsdp_state_dict + >>> assert ddp_optim_state == fsdp_optim_state_dict + + + Args: + model (nn.Module): the nn.Module to the model. + optimizers (Union[None, Optimizer, Iterable[Optimizer]]): + The optimizers that are used to optimize ``model``. + submodules (deprecated): Optional[Set[nn.Module]]: only return the model parameters + that belong to the submodules. + options (StateDictOptions): the options to control how + model state_dict and optimizer state_dict should be returned. See + `StateDictOptions` for the details. + + Returns: + ``Tuple`` that contain model state_dict and optimizer state_dict. + + :rtype: typing.Tuple[typing.Dict[str, ValueType], OptimizerStateType] + """ + + with _gc_context(): + optimizers = ( + (optimizers,) + if isinstance(optimizers, torch.optim.Optimizer) + else tuple(optimizers) + ) + info = _verify_options( + model, + optimizers, + optim_only=False, + submodules=submodules, + options=options, + ) + model_state_dict = _get_model_state_dict(model, info) + optim_state_dict = _get_optim_state_dict(model, optimizers, info) + _verify_state_dict(model_state_dict, optim_state_dict, info) + return model_state_dict, optim_state_dict + + +def _unflatten_model_state_dict( + model: nn.Module, + state_dict: Union[Dict[nn.Module, Dict[str, ValueType]], Dict[str, ValueType]], +) -> Dict[str, ValueType]: + if not state_dict: + return {} + + if isinstance(next(iter(state_dict.keys())), nn.Module): + warnings.warn( + "Passing model_state_dict as a ``Dict[nn.Module, Dict[str, Any]]``" + "is deprecated and will be removed in 2.5. If you need this " + "feature, please preprocessing the model_state_dict to achieve the " + "same functionality.", + FutureWarning, + ) + cast_state_dict = cast(Dict[nn.Module, Dict[str, ValueType]], state_dict) + new_state_dict: Dict[str, ValueType] = {} + for submodule, sub_state_dict in cast_state_dict.items(): + for name, m in model.named_modules(): + if m != submodule: + continue + + fqns = _get_fqns(model, name) + assert len(fqns) == 1, "FQNs for a submodule should only have 1 element" + prefix = f"{next(iter(fqns))}." + new_state_dict.update( + {prefix + subfqn: value for subfqn, value in sub_state_dict.items()} + ) + return new_state_dict + else: + return cast(Dict[str, ValueType], state_dict) + + +def set_model_state_dict( + model: nn.Module, + model_state_dict: Dict[str, ValueType], + *, + options: Optional[StateDictOptions] = None, +) -> _IncompatibleKeys: + """Load the model state_dict. + + The counterpart of ``get_model_state_dict`` to set the state_dict to the + model. See ``set_state_dict`` for the detail usage. + + Args: + model (nn.Module): the nn.Module to the model. + model_state_dict: (Dict[str, ValueType]): + the model state_dict to load. If the key of the ``model_state_dict`` + is nn.Module, the key is a submodule of ``model`` and the value should + be the state_dict of the submodule. When loading the state_dict, + the prefix of the submodule will be append to the state_dict. + options (StateDictOptions): the options to control how + model state_dict and optimizer state_dict should be loaded. See + `StateDictOptions` for the details. + + Returns: + ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: + * **missing_keys** is a list of str containing the missing keys + * **unexpected_keys** is a list of str containing the unexpected keys + + :type model_state_dict: typing.Dict[str, ValueType] + """ + model_state_dict: Dict[str, ValueType] = _unflatten_model_state_dict( + model, model_state_dict + ) + with _gc_context(): + info = _verify_options(model, (), optim_only=False, options=options) + + _verify_state_dict(model_state_dict, {}, info) + return _load_model_state_dict(model, model_state_dict, info) + + +def set_optimizer_state_dict( + model: nn.Module, + optimizers: Union[torch.optim.Optimizer, Iterable[torch.optim.Optimizer]], + optim_state_dict: OptimizerStateType, + *, + options: Optional[StateDictOptions] = None, +) -> None: + """Load the optimizers state_dict. + + The counterpart of ``get_optimizer_state_dict`` to set the state_dict to the + optimizers. See ``set_state_dict`` for the detail usage. + + Args: + model (nn.Module): the nn.Module to the model. + optimizers (Union[Optimizer, Iterable[Optimizer]]): + The optimizers that are used to optimize ``model``. + optim_state_dict: OptimizerStateType: + the optimizer state_dict to load. + options (StateDictOptions): the options to control how + model state_dict and optimizer state_dict should be loaded. See + `StateDictOptions` for the details. + + Returns: + None + + :type optim_state_dict: typing.OptimizerStateType + """ + with _gc_context(): + optimizers = ( + (optimizers,) + if isinstance(optimizers, torch.optim.Optimizer) + else tuple(optimizers) + ) + info = _verify_options(model, optimizers, optim_only=True, options=options) + + _verify_state_dict({}, optim_state_dict, info) + _load_optim_state_dict(model, optimizers, optim_state_dict, info) + + +def set_state_dict( + model: nn.Module, + optimizers: Union[torch.optim.Optimizer, Iterable[torch.optim.Optimizer]], + *, + model_state_dict: Dict[str, ValueType], + optim_state_dict: OptimizerStateType, + options: Optional[StateDictOptions] = None, +) -> _IncompatibleKeys: + """Load the model state_dict and optimizers state_dict. + + The counterpart of ``get_state_dict`` to set the state_dict to the model and + optimizers. The given ``model_state_dict`` and ``optim_state_dict`` do not + have to be returned by ``get_state_dict`` but must meet the following + requirements: 1) all FQNs are canonical FQNs as defined in ``get_state_dict``, + 2) if a tensor is sharded, it must be either a ShardedTensor or DTensor, + 3) optimizer state_dict cannot contain the parameter IDs; the keys should be + the canonical FQNs. + + Args: + model (nn.Module): the nn.Module to the model. + optimizers (Union[Optimizer, Iterable[Optimizer]]): + The optimizers that are used to optimize ``model``. + model_state_dict: (Union[Dict[nn.Module, Dict[str, ValueType]], Dict[str, ValueType]]): + the model state_dict to load. If the key of the ``model_state_dict`` + is nn.Module, the key is a submodule of ``model`` and the value should + be the state_dict of the submodule. When loading the state_dict, + the prefix of the submodule will be append to the state_dict. + optim_state_dict: OptimizerStateType: + the optimizer state_dict to load. + options (StateDictOptions): the options to control how + model state_dict and optimizer state_dict should be loaded. See + `StateDictOptions` for the details. + + Returns: + ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: + * **missing_keys** is a list of str containing the missing keys of the model state_dict. + * **unexpected_keys** is a list of str containing the unexpected keys of the model state_dict. + + :type model_state_dict: typing.Dict[str, ValueType] + :type optim_state_dict: typing.OptimizerStateType + """ + + model_state_dict: Dict[str, ValueType] = _unflatten_model_state_dict( + model, model_state_dict + ) + with _gc_context(): + optimizers = ( + (optimizers,) + if isinstance(optimizers, torch.optim.Optimizer) + else tuple(optimizers) + ) + info = _verify_options( + model, optimizers, optim_only=not model_state_dict, options=options + ) + + _verify_state_dict(model_state_dict, optim_state_dict, info) + _load_optim_state_dict(model, optimizers, optim_state_dict, info) + return _load_model_state_dict(model, model_state_dict, info) + + +# TODO: correct the state_dict function signature. +# TODO: this API is not yet fully tested. Make it private +@no_type_check +def _patch_model_state_dict( + model: nn.Module, + *, + options: Optional[StateDictOptions] = None, +) -> None: + """Patch the ``state_dict`` and ``load_state_dict`` attributes of ``model``. + + Patch the ``state_dict`` and ``load_state_dict`` attributes of ``model`` to + be a partial function to call ``get_state_dict`` and ``set_state_dict``. + + Example: + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from torch.distributed.checkpoint.state_dict import patch_model_state_dict + + model = fsdp(model) + patch_model_state_dict(model) + + Args: + model (nn.Module): the nn.Module to the model. + options (StateDictOptions): the options to control how + model state_dict and optimizer state_dict should be loaded. See + `StateDictOptions` for the details. + Returns: + None + """ + + _state_dict_call = functools.partial( + get_model_state_dict, + model=model, + options=options, + ) + + def state_dict_call(): + return _state_dict_call() + + model.state_dict = state_dict_call + + _load_state_dict_call = functools.partial( + set_model_state_dict, + model=model, + options=options, + ) + + def load_state_dict_call(state_dict: Dict[str, Any]): + _load_state_dict_call(model_state_dict=state_dict) + + model.load_state_dict = load_state_dict_call + + _patched_state_dict.add(state_dict_call) + _patched_state_dict.add(load_state_dict_call) + + +# TODO: correct the load_state_dict function signature. +# TODO: this API is not yet fully tested. Make it private +@no_type_check +def _patch_optimizer_state_dict( + model: nn.Module, + *, + optimizers: Tuple[torch.optim.Optimizer, ...], + options: Optional[StateDictOptions] = None, +) -> None: + """Patch the ``state_dict`` and ``load_state_dict`` attributes of ``optimizers``. + + Patch the ``state_dict`` and ``load_state_dict`` attributes of ``optimizers`` to + be a partial function to call ``get_state_dict`` and ``set_state_dict``. + + Note that if there are multiple optimizers, all of the optimizers will be patched. + So users only need to call one of the state_dict() to get the full result. + + Example: + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from torch.distributed.checkpoint.state_dict import patch_model_state_dict + + model = fsdp(model) + patch_model_state_dict(model) + + Args: + model (nn.Module): the nn.Module to the model. + options (StateDictOptions): the options to control how + model state_dict and optimizer state_dict should be loaded. See + `StateDictOptions` for the details. + Returns: + None + """ + + _state_dict_call = functools.partial( + get_optimizer_state_dict, + model=model, + optimizers=optimizers, + options=options, + ) + + def state_dict_call(): + return _state_dict_call() + + _load_state_dict_call = functools.partial( + set_optimizer_state_dict, + model=model, + optimizers=optimizers, + options=options, + ) + + def load_state_dict_call(state_dict: Dict[str, Any]): + _load_state_dict_call(optim_state_dict=state_dict) + + _patched_state_dict.add(state_dict_call) + _patched_state_dict.add(load_state_dict_call) + optimizers = ( + (optimizers,) + if isinstance(optimizers, torch.optim.Optimizer) + else tuple(optimizers) + ) + for optim in optimizers: + optim.state_dict = state_dict_call + optim.load_state_dict = load_state_dict_call diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict_loader.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..6a915b784878018c33519187b387da87a5af1625 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict_loader.py @@ -0,0 +1,323 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +import os +import warnings +from typing import Any, cast, Dict, Optional, Set, Union +from typing_extensions import deprecated + +import torch +import torch.distributed as dist +from torch.distributed.checkpoint.default_planner import _EmptyStateDictLoadPlanner +from torch.distributed.checkpoint.logger import _dcp_method_logger +from torch.distributed.checkpoint.stateful import Stateful + +from ._storage_utils import _storage_setup +from .default_planner import DefaultLoadPlanner +from .planner import LoadPlan, LoadPlanner +from .storage import StorageReader +from .utils import _all_gather_keys, _api_bc_check, _DistWrapper, _profile + + +__all__ = ["load_state_dict", "load"] + + +@deprecated( + "`load_state_dict` is deprecated and will be removed in future versions. " + "Please use `load` instead.", + category=FutureWarning, +) +def load_state_dict( + state_dict: Dict[str, Any], + storage_reader: StorageReader, + process_group: Optional[dist.ProcessGroup] = None, + coordinator_rank: int = 0, + no_dist: bool = False, + planner: Optional[LoadPlanner] = None, +) -> None: + """This method is deprecated. Please switch to 'load'.""" + storage_reader.reset() + with _profile(): + # TODO: test returning `load` here instead. + return _load_state_dict( + state_dict, + storage_reader, + process_group, + coordinator_rank, + no_dist, + planner, + ) + + +@_dcp_method_logger(log_exceptions=True) +@_api_bc_check +def load( + state_dict: Dict[str, Any], + *, + checkpoint_id: Union[str, os.PathLike, None] = None, + storage_reader: Optional[StorageReader] = None, + planner: Optional[LoadPlanner] = None, + process_group: Optional[dist.ProcessGroup] = None, +) -> None: + """ + Load a distributed ``state_dict`` in SPMD style. + + Each rank will try to read the least amount of data necessary + to fullfill the requested `state_dict`. When loading :class:`ShardedTensor` + or :class:`DTensor` instances, each rank only reads data for their local shards. + + For each ``Stateful`` object (having both a ``state_dict`` and a ``load_state_dict``), + load will first call ``state_dict`` before attempting deserialization, followed by + ``load_state_dict`` once the deserialization is complete. + For each non-``Stateful`` object, load will deserailize the object, and then replace + it in the ``state_dict`` with the deserialized object. + + .. warning:: + All tensors in ``state_dict`` must be allocated on their + destination device *prior to* calling this function. + + All non-tensor data is loaded using `torch.load()` and modified in place + on state_dict. + + .. warning:: + Users must call `load_state_dict` on the root module to ensure load + pos-processing and non-tensor data properly propagates. + + .. note: + If no process group is initialized, this function will assume the intent + is to load a checkpoint into the local process. This can be useful in the + case of local inference, and when using regular Tensors (as opposed to DTensor + or ShardedTensor) + + .. note: + Rank 0 is assumed to be the coordinator rank. + + Args: + state_dict (Dict[str, Any]): The state_dict to save. + checkpoint_id (Union[str, os.PathLike, None]): + The ID of this checkpoint instance. The meaning of the checkpoint_id + depends on the storage. It can be a path to a folder or to a file. + It can also be a key if the storage is a key-value store. + (Default: ``None``) + storage_reader (Optional[StorageReader]): + Instance of StorageWriter used to perform reads. If this is not + specified, DCP will automatically infer the reader based on the + checkpoint_id. If checkpoint_id is also None, an exception will + be raised. (Default: ``None``) + planner (Optional[LoadPlanner]): + Instance of LoadPlanner. If this is not specificed, the default + planner will be used. (Default: ``None``) + process_group (Optional[ProcessGroup]): + ProcessGroup to be used for cross-rank synchronization. + (Default: ``None``) + + Returns: + None. + + Examples + >>> # xdoctest: +SKIP + >>> my_model = MyModule() + >>> optimizer = Adagrad(my_model.parameters()) + >>> model_state_dict = my_model.state_dict() + >>> fs_storage_reader = torch.distributed.checkpoint.FileSystemReader("/checkpoint/1") + + >>> torch.distributed.checkpoint.load_state_dict( + >>> state_dict=model_state_dict, + >>> storage_reader=fs_storage_reader, + >>> ) + + >>> # module.load_state_dict() function might have customized steps + >>> # to flush the state_dict, must call it to + >>> # ensure correct behavior. + >>> my_model.load_state_dict(model_state_dict) + + .. note:: + load_state_dict uses collectives to coordinate reads across ranks. + For NCCL-based process groups, internal tensor representations of + objects must be moved to the GPU device before communication takes place. + In this case, the device used is given by ``torch.cuda.current_device()`` + and it is the user's responsibility to ensure that this is set so that each + rank has an individual GPU, via ``torch.cuda.set_device()``. + """ + + no_dist = not (dist.is_available() and dist.is_initialized()) + if no_dist: + warnings.warn( + "torch.distributed is unavailable or uninitialized, assuming the intent is to load in a single process." + ) + + with _profile(): + storage_reader = cast( + StorageReader, _storage_setup(storage_reader, checkpoint_id, reader=True) + ) + + if no_dist: + keys = list(state_dict.keys()) + else: + keys = _all_gather_keys(state_dict, process_group) + if keys != sorted(state_dict.keys()): + warnings.warn( + "Detected mismatched keys in state dict after all gather!" + " This behavior is unsupported and may cause errors may cause errors." + ) + + statetful_sd = {} + for key in keys: + if key not in state_dict: + continue + elem = state_dict[key] + statetful_sd[key] = ( + elem.state_dict() if isinstance(elem, Stateful) else elem + ) + + _load_state_dict( + state_dict=statetful_sd, + storage_reader=storage_reader, + process_group=process_group, + no_dist=no_dist, + planner=planner, + ) + for key in keys: + if key not in state_dict: + continue + elem = state_dict[key] + if isinstance(elem, Stateful): + # If the state_dict is a Stateful object, + # DCP does an in-place load in the original state dict. + elem.load_state_dict(statetful_sd[key]) + else: + # Otherwise, replace the state_dict with the loaded state_dict. + state_dict[key] = statetful_sd[key] + + +def _load_state_dict( + state_dict: Dict[str, Any], + storage_reader: StorageReader, + process_group: Optional[dist.ProcessGroup] = None, + coordinator_rank: int = 0, + no_dist: bool = False, + planner: Optional[LoadPlanner] = None, +) -> None: + torch._C._log_api_usage_once("torch.distributed.checkpoint.load_state_dict") + + distW = _DistWrapper(process_group, not no_dist, coordinator_rank) + if planner is None: + planner = DefaultLoadPlanner() + + ckpt_kwargs = {} + if (ckpt_id := getattr(storage_reader, "checkpoint_id", None)) is not None: + ckpt_kwargs["checkpoint_id"] = ckpt_id + ckpt_kwargs["process_group"] = distW.group + + @_dcp_method_logger(**ckpt_kwargs) + def local_step(): + assert planner is not None + metadata = storage_reader.read_metadata() + planner.set_up_planner(state_dict, metadata, distW.is_coordinator) + storage_reader.set_up_storage_reader(metadata, distW.is_coordinator) + + local_plan = planner.create_local_plan() + local_plan = storage_reader.prepare_local_plan(local_plan) + return local_plan + + @_dcp_method_logger(**ckpt_kwargs) + def global_step(all_local_plans): + assert planner is not None + all_local_plans = planner.create_global_plan(all_local_plans) + all_local_plans = storage_reader.prepare_global_plan(all_local_plans) + return all_local_plans + + central_plan: LoadPlan = distW.reduce_scatter("plan", local_step, global_step) + + @_dcp_method_logger(**ckpt_kwargs) + def read_data(): + assert planner is not None + final_local_plan = planner.finish_plan(central_plan) + all_reads = storage_reader.read_data(final_local_plan, planner) + + all_reads.wait() + return None + + _ = distW.all_gather("read", read_data) + + +def _load_state_dict_from_keys( + keys: Optional[Union[Set[str], str]] = None, + *, + checkpoint_id: Union[str, os.PathLike, None] = None, + storage_reader: Optional[StorageReader] = None, + process_group: Optional[dist.ProcessGroup] = None, +) -> Dict[str, Any]: + """ + Load only the specified keys from the checkpoint, if no keys are specified, the entire + checkpoint will be loaded. Note, this method completely loads the checkpoint into the + current process and is not distributed. + + .. warning:: + + + .. warning:: + + All non-tensor data is loaded using `torch.load()` + + .. note: + As opposed to the usual pattern, this function does not take a state dict as input + and does not load inplace. Instead, a new state dict is directly initialized and read + from file. + + .. note: + If no process group is initialized, this function will assume the intent + is to load a checkpoint into the local process. This can be useful in the + case of local inference, and when using regular Tensors (as opposed to DTensor + or ShardedTensor) + + .. note: + Rank 0 is assumed to be the coordinator rank. + + Args: + keys (Optional[Union[Set[str], str]]): + Loads any key specified in this set. If no keys are specified, the entire checkpoint + is loaded. + checkpoint_id (Union[str, os.PathLike, None]): + The ID of this checkpoint instance. The meaning of the checkpoint_id + depends on the storage. It can be a path to a folder or to a file. + It can also be a key if the storage is a key-value store. + (Default: ``None``) + storage_reader (Optional[StorageReader]): + Instance of StorageWriter used to perform reads. If this is not + specified, DCP will automatically infer the reader based on the + checkpoint_id. If checkpoint_id is also None, an exception will + be raised. (Default: ``None``) + process_group (Optional[ProcessGroup]): + ProcessGroup to be used for cross-rank synchronization. + (Default: ``None``) + + Returns: + State dict from specified keys + """ + torch._C._log_api_usage_once( + "torch.distributed.checkpoint._load_state_dict_from_keys" + ) + + no_dist = not (dist.is_available() and dist.is_initialized()) + if no_dist: + warnings.warn( + "torch.distributed is unavailable or uninitialized, assuming the intent is to load in a single process." + ) + + storage_reader = cast( + StorageReader, _storage_setup(storage_reader, checkpoint_id, reader=True) + ) + + if isinstance(keys, str): + keys = {keys} + + sd: Dict[str, Any] = {} + _load_state_dict( + state_dict=sd, + storage_reader=storage_reader, + process_group=process_group, + no_dist=no_dist, + planner=_EmptyStateDictLoadPlanner(keys=keys or set()), + ) + + return sd diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict_saver.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict_saver.py new file mode 100644 index 0000000000000000000000000000000000000000..8df80c21b42d00eca9b306a0962f474780c99045 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict_saver.py @@ -0,0 +1,334 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +import inspect +import os +import warnings +from concurrent.futures import Future, ThreadPoolExecutor +from typing import cast, Optional, Union +from typing_extensions import deprecated + +import torch +import torch.distributed as dist +from torch.distributed._state_dict_utils import _offload_state_dict_to_cpu +from torch.distributed.checkpoint._storage_utils import _storage_setup +from torch.distributed.checkpoint.default_planner import DefaultSavePlanner +from torch.distributed.checkpoint.logger import _dcp_method_logger +from torch.distributed.checkpoint.metadata import Metadata, STATE_DICT_TYPE +from torch.distributed.checkpoint.planner import SavePlan, SavePlanner +from torch.distributed.checkpoint.staging import AsyncStager +from torch.distributed.checkpoint.stateful import Stateful +from torch.distributed.checkpoint.storage import StorageWriter +from torch.distributed.distributed_c10d import _get_default_group + +from .utils import _api_bc_check, _DistWrapper, _profile + + +__all__ = ["save_state_dict", "save", "async_save"] + + +@deprecated( + "`save_state_dict` is deprecated and will be removed in future versions." + "Please use `save` instead.", + category=FutureWarning, +) +def save_state_dict( + state_dict: STATE_DICT_TYPE, + storage_writer: StorageWriter, + process_group: Optional[dist.ProcessGroup] = None, + coordinator_rank: int = 0, + no_dist: bool = False, + planner: Optional[SavePlanner] = None, +) -> Metadata: + """This method is deprecated. Please switch to 'save'.""" + storage_writer.reset() + + # TODO: test returning `save` here instead. + with _profile(): + return _save_state_dict( + state_dict, + storage_writer, + process_group, + coordinator_rank, + no_dist, + planner, + ) + + +@_dcp_method_logger(log_exceptions=True) # type: ignore[arg-type] +@_api_bc_check +def save( + state_dict: STATE_DICT_TYPE, + *, + checkpoint_id: Union[str, os.PathLike, None] = None, + storage_writer: Optional[StorageWriter] = None, + planner: Optional[SavePlanner] = None, + process_group: Optional[dist.ProcessGroup] = None, +) -> Metadata: + """ + Save a distributed model in SPMD style. + + This function is different from ``torch.save()`` as it handles + ``ShardedTensor`` , and ``DTensor`` by having each rank only save their local shards. + + For each ``Stateful`` object (having both a ``state_dict`` and a ``load_state_dict``), + save will call ``state_dict`` before serialization. + + .. warning:: + There is no guarantees of Backwards Compatibility across PyTorch versions + for saved state_dicts. + + .. warning:: + If using the `process_group` argument, make sure that only its ranks + call `save_state_dict` and that all data in state_dict belong to it. + + .. note:: + When saving checkpoint for FSDP's `ShardingStrategy.HYBRID_SHARD`, only one of + the shard_group should be calling `save_state_dict` and the corresponding process + group needs to be passed in. + + .. note:: + If no process group is available, this function assumes the intention is to save the + state_dict in the local process. + + .. note: + Rank 0 is assumed to be the coordinator rank. + + + Args: + state_dict (Dict[str, Any]): The state_dict to save. + checkpoint_id (Union[str, os.PathLike, None]): + The ID of this checkpoint instance. The meaning of the checkpoint_id + depends on the storage. It can be a path to a folder or to a file. + It can also be a key if the storage is a key-value store. + (Default: ``None``) + storage_writer (Optional[StorageWriter]): + Instance of StorageWriter used to perform writes. If this is not + specified, DCP will automatically infer the writer based on the + checkpoint_id. If checkpoint_id is also None, an exception will + be raised. (Default: ``None``) + planner (Optional[SavePlanner]): + Instance of SavePlanner. If this is not specificed, the default + planner will be used. (Default: ``None``) + process_group (Optional[ProcessGroup]): + ProcessGroup to be used for cross-rank synchronization. + (Default: ``None``) + + Returns: + Metadata: Metadata object for the saved checkpoint. + + Example: + >>> # xdoctest: +SKIP + >>> my_model = MyModule() + + >>> state_dict = {"model": my_model} + + >>> fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter("/checkpoint/1") + >>> torch.distributed.checkpoint.save( + >>> state_dict=state_dict, + >>> storage_writer=fs_storage_writer, + >>> ) + + .. note:: + save_state_dict uses collectives to coordinate writes across ranks. + For NCCL-based process groups, internal tensor representations of + objects must be moved to the GPU device before communication takes place. + In this case, the device used is given by ``torch.cuda.current_device()`` + and it is the user's responsibility to ensure that this is set so that + each rank has an individual GPU, via ``torch.cuda.set_device()``. + """ + torch._C._log_api_usage_once("torch.distributed.checkpoint.save") + + no_dist = not (dist.is_available() and dist.is_initialized()) + if no_dist: + warnings.warn( + "torch.distributed is unavailable or uninitialized, assuming the intent is to save in a single process." + ) + + with _profile(): + storage_writer = cast( + StorageWriter, _storage_setup(storage_writer, checkpoint_id, reader=False) + ) + + return _save_state_dict( + state_dict=_stateful_to_state_dict(state_dict), + storage_writer=storage_writer, + process_group=process_group, + no_dist=no_dist, + planner=planner, + ) + + +@_dcp_method_logger(log_exceptions=True) +def async_save( + state_dict: STATE_DICT_TYPE, + *, + checkpoint_id: Union[str, os.PathLike, None] = None, + storage_writer: Optional[StorageWriter] = None, + planner: Optional[SavePlanner] = None, + process_group: Optional[dist.ProcessGroup] = None, +) -> Future: + """Asynchronous version of ``save``. This code first de-stages the state_dict on to the + staging storage (defaults to CPU memory), and then calls the `save` in a separate thread. + + .. warning:: + This feature is experimental and subject to change. + + Args: + state_dict (Dict[str, Any]): The state_dict to save. + checkpoint_id (Union[str, os.PathLike, None]): + The ID of this checkpoint instance. The meaning of the checkpoint_id + depends on the storage. It can be a path to a folder or to a file. + It can also be a key if the storage is a key-value store. + (Default: ``None``) + storage_writer (Optional[StorageWriter]): + Instance of StorageWriter used to perform 'stage' and 'save'. If + this is not specified, DCP will automatically infer the writer based on the + checkpoint_id. If checkpoint_id is also None, an exception will + be raised. (Default: ``None``) + planner (Optional[SavePlanner]): + Instance of SavePlanner. If this is not specificed, the default + planner will be used. (Default: ``None``) + process_group (Optional[ProcessGroup]): + ProcessGroup to be used for cross-rank synchronization. + (Default: ``None``) + + Returns: + Future: A future holding the resultant Metadata object from `save`. + + Example: + >>> # xdoctest: +SKIP + >>> my_model = MyModule() + + >>> state_dict = {"model": my_model} + + >>> fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter("/checkpoint/1") + >>> checkpoint_future = torch.distributed.checkpoint.async_save( + >>> state_dict=state_dict, + >>> storage_writer=fs_storage_writer, + >>> ) + >>> + >>> # ... do some work ... + >>> + >>> checkpoint_future.result() + + """ + torch._C._log_api_usage_once("torch.distributed.checkpoint.async_save") + + if dist.is_available() and dist.is_initialized(): + pg = process_group or _get_default_group() + assert ( + torch.device("cpu") in pg._device_types # type: ignore[attr-defined] + ), "A CPU backend must be enabled for async save; try initializing process group with 'cpu:gloo,cuda:nccl'" + + storage_writer = cast( + StorageWriter, _storage_setup(storage_writer, checkpoint_id, reader=False) + ) + + state_dict = _stateful_to_state_dict(state_dict) + if isinstance(storage_writer, AsyncStager): + staged_state_dict = storage_writer.stage(state_dict) + else: # provides bwc for storage_writers not implementing AsyncStager + staged_state_dict = _offload_state_dict_to_cpu(state_dict, type_check=False) + + executor = ThreadPoolExecutor(max_workers=1) + f: Future = executor.submit( + save, + staged_state_dict, + checkpoint_id=checkpoint_id, + storage_writer=storage_writer, + planner=planner, + process_group=process_group, + ) + f.add_done_callback(lambda f: executor.shutdown(wait=False)) + + if ( + isinstance(storage_writer, AsyncStager) + and storage_writer.should_synchronize_after_execute + ): + storage_writer.synchronize_staging() + + return f + + +def _stateful_to_state_dict(state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE: + """Creates a shallow copy of `state_dict` where `state_dict` is called for each Stateful object.""" + stateful_state_dict = {} + for key, elem in state_dict.items(): + stateful_state_dict[key] = ( + elem.state_dict() if isinstance(elem, Stateful) else elem + ) + return stateful_state_dict + + +def _save_state_dict( + state_dict: STATE_DICT_TYPE, + storage_writer: StorageWriter, + process_group: Optional[dist.ProcessGroup] = None, + coordinator_rank: int = 0, + no_dist: bool = False, + planner: Optional[SavePlanner] = None, +) -> Metadata: + torch._C._log_api_usage_once("torch.distributed.checkpoint.save_state_dict") + + distW = _DistWrapper(process_group, not no_dist, coordinator_rank) + if planner is None: + planner = DefaultSavePlanner() + assert planner is not None + + global_metadata = None + + ckpt_kwargs = {} + if (ckpt_id := getattr(storage_writer, "checkpoint_id", None)) is not None: + ckpt_kwargs["checkpoint_id"] = ckpt_id + ckpt_kwargs["process_group"] = distW.group + + @_dcp_method_logger(**ckpt_kwargs) + def local_step(): + assert planner is not None + storage_meta = storage_writer.storage_meta() + if "storage_meta" not in inspect.signature(planner.set_up_planner).parameters: + warnings.warn( + "The function definition for SavePlanner.set_up_planner has been updated" + " to include the storage_meta argument. Please update your implementation" + " to include this parameter." + ) + planner.set_up_planner(state_dict, distW.is_coordinator) # type: ignore[call-arg, arg-type] + else: + planner.set_up_planner( + state_dict=state_dict, + storage_meta=storage_meta, + is_coordinator=distW.is_coordinator, + ) + storage_writer.set_up_storage_writer(distW.is_coordinator) + + local_plan = planner.create_local_plan() + local_plan = storage_writer.prepare_local_plan(local_plan) + return local_plan + + @_dcp_method_logger(**ckpt_kwargs) + def global_step(all_local_plans): + nonlocal global_metadata + + assert planner is not None + all_local_plans, global_metadata = planner.create_global_plan(all_local_plans) + all_local_plans = storage_writer.prepare_global_plan(all_local_plans) + return all_local_plans + + central_plan: SavePlan = distW.reduce_scatter("plan", local_step, global_step) + + @_dcp_method_logger(**ckpt_kwargs) + def write_data(): + assert planner is not None + final_local_plan = planner.finish_plan(central_plan) + all_writes = storage_writer.write_data(final_local_plan, planner) + + all_writes.wait() + return all_writes.value() + + @_dcp_method_logger(**ckpt_kwargs) + def finish_checkpoint(all_results): + assert global_metadata is not None + storage_writer.finish(metadata=global_metadata, results=all_results) + return global_metadata + + return distW.all_reduce("write", write_data, finish_checkpoint) diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/storage.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/storage.py new file mode 100644 index 0000000000000000000000000000000000000000..985e91df9e739d1033725ce6173e35ea25296907 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/storage.py @@ -0,0 +1,284 @@ +import abc +import os +from dataclasses import dataclass +from typing import Any, List, Optional, Union + +from torch.distributed.checkpoint.metadata import Metadata, MetadataIndex, StorageMeta +from torch.distributed.checkpoint.planner import ( + LoadPlan, + LoadPlanner, + SavePlan, + SavePlanner, +) +from torch.futures import Future + + +__all__ = ["WriteResult", "StorageWriter", "StorageReader"] + + +@dataclass(frozen=True) +class WriteResult: + index: MetadataIndex + + size_in_bytes: int + storage_data: Any + + +class StorageWriter(abc.ABC): + """ + Interface used by ``save_state_dict`` to write to storage. + + One StorageWriter instance acts as both the coordinator and the follower + in a distributed checkpoint. As part of initialization, each instance + is told its role. + + A subclass should expect the following sequence of calls. + + 0) (all ranks) set checkpoint_id if users pass a valid checkpoint_id. + 1) (all ranks) set_up_storage_writer() + 2) (all ranks) prepare_local_plan() + 3) (coordinator) prepare_global_plan() + 4) (all ranks) write_data() + 5) (coordinator) finish() + """ + + @abc.abstractmethod + def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None: + """ + Calls to indicates a brand new checkpoint write is going to happen. + A checkpoint_id may be present if users set the checkpoint_id for + this checkpoint write. The meaning of the checkpiont_id is + storage-dependent. It can be a path to a folder/file or a key for + a key-value storage. + + Args: + checkpoint_id (Union[str, os.PathLike, None]): + The ID of this checkpoint instance. The meaning of the checkpoint_id + depends on the storage. It can be a path to a folder or to a file. + It can also be a key if the storage is a key-value store. + (Default: ``None``) + """ + ... + + @abc.abstractmethod + def set_up_storage_writer(self, is_coordinator: bool) -> None: + """ + Initialize this instance. + + Args: + is_coordinator (bool): Whether this instance is responsible for coordinating + the checkpoint. + """ + + @abc.abstractmethod + def prepare_local_plan(self, plan: SavePlan) -> SavePlan: + """ + Perform storage-specific local planning. + + While this method can produce a completely different plan, the recommended + way is to store storage specific data in SavePlan::storage_data. + + Args: + plan (SavePlan): The local plan from the ``SavePlanner`` in use. + + Returns: + A transformed ``SavePlan`` after storage local planning + """ + + @abc.abstractmethod + def prepare_global_plan(self, plans: List[SavePlan]) -> List[SavePlan]: + """ + Perform centralized planning of storage. + + This method is only called on the coordinator instance. + + While this method can produce a completely different plan, the preferred + way is to store storage specific data in SavePlan::storage_data. + + Args: + plans: A list of ``SavePlan`` instances, one for each rank. + + Returns: + A list of transformed ``SavePlan`` after storage global planning + """ + + @abc.abstractmethod + def write_data( + self, plan: SavePlan, planner: SavePlanner + ) -> Future[List[WriteResult]]: + """ + Write all items from ``plan`` using ``planner`` to resolve the data. + + A subclass should call ``SavePlanner::resolve_data`` on each item + from the plan to get access to the underlying object to write. + + Subclasses should lazily call `resolve_data` as it can allocate memory. + In case of tensors, make following assumptions: + + - They might be on any device, including not matching the one on ``WriteItem::tensor_data`` + - They might be views or not contiguous. Only the projection needs to be saved. + + Args: + plan (SavePlan): The save plan to execute. + planner (SavePlanner): Planner object to be used to resolve items to data. + + Returns: + A future that completes to a list of WriteResult + """ + + @abc.abstractmethod + def finish(self, metadata: Metadata, results: List[List[WriteResult]]) -> None: + """ + Write the metadata and marks the current checkpoint as successful. + + The actual format/schema used for serializing `metadata` is an + implementation detail. The only requirement is that it's recoverable + in to the same object graph. + + Args: + metadata (Metadata): metadata for the new checkpoint + results: A list of WriteResults from all ranks. + + Returns: + None + """ + + @classmethod + @abc.abstractmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + """ + Check if the given checkpoint_id is supported by the stroage. This allow + us to enable automatic storage selection. + """ + ... + + def storage_meta(self) -> Optional[StorageMeta]: + """ + Return the storage-specific metadata. This is used to store additional information + in a checkpoint that can be useful for providing request-level observability. StorageMeta + is passed to the ``SavePlanner`` during save calls. Returns None by default. + + TODO: provide an example + """ + return None + + +class StorageReader(abc.ABC): + """ + Interface used by ``load_state_dict`` to read from storage. + + One StorageReader instance acts as both the coordinator and the follower + in a distributed checkpoint. As part of initialization, each instance + is told its role. + + A subclass should expected the following sequence of calls by ``load_state_dict``: + + 0) (all ranks) set checkpoint_id if users pass a valid checkpoint_id. + 1) (all ranks) read_metadata() + 2) (all ranks) set_up_storage_reader() + 3) (all ranks) prepare_local_plan() + 4) (coordinator) prepare_global_plan() + 5) (all ranks) read_data() + """ + + @abc.abstractmethod + def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None: + """ + Calls to indicates a brand new checkpoint read is going to happen. + A checkpoint_id may be present if users set the checkpoint_id for + this checkpoint read. The meaning of the checkpiont_id is + storage-dependent. It can be a path to a folder/file or a key for + a key-value storage. + + Args: + checkpoint_id (Union[str, os.PathLike, None]): + The ID of this checkpoint instance. The meaning of the checkpoint_id + depends on the storage. It can be a path to a folder or to a file. + It can also be a key if the storage is more like a key-value store. + (Default: ``None``) + """ + ... + + @abc.abstractmethod + def read_metadata(self) -> Metadata: + """ + Read the checkpoint metadata. + + Returns: + The metadata object associated with the checkpoint being loaded. + + """ + + @abc.abstractmethod + def set_up_storage_reader(self, metadata: Metadata, is_coordinator: bool) -> None: + """ + Initialize this instance. + + Args: + metadata (Metadata): The metadata schema to use. + is_coordinator (bool): Whether this instance is responsible for coordinating + the checkpoint. + """ + + @abc.abstractmethod + def prepare_local_plan(self, plan: LoadPlan) -> LoadPlan: + """ + Perform storage-specific local planning. + + While this method can produce a completely different plan, the recommended + way is to store storage specific data in LoadPlan::storage_data. + + Args: + plan (LoadPlan): The local plan from the ``LoadPlan`` in use. + + Returns: + A transformed ``LoadPlan`` after storage local planning + """ + + @abc.abstractmethod + def prepare_global_plan(self, plans: List[LoadPlan]) -> List[LoadPlan]: + """ + Perform centralized planning of storage loading. + + This method is only called on the coordinator instance. + + While this method can produce a completely different plan, the preferred + way is to store storage specific data in LoadPlan::storage_data. + + Args: + plans: A list of ``LoadPlan`` instances, one for each rank. + + Returns: + A list of transformed ``LoadPlan`` after storage global planning + """ + + @abc.abstractmethod + def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]: + """ + Read all items from ``plan`` using ``planner`` to resolve the data. + + A subclass should call ``LoadPlanner::load_bytes`` to deserialize a BytesIO + object into the right place. + + A subclass should call ``LoadPlanner::resolve_tensor`` to get access to the + tensors that in should load data into. + + It's the StorageLayer responsibility to properly schedule any cross device copies + required. + + Args: + plan (LoadPlan): The local plan to execute on + planner (LoadPlanner): The planner object to use to resolve items. + + Returns: + A future that completes once all reads are finished. + """ + + @classmethod + @abc.abstractmethod + def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool: + """ + Check if the given checkpoint_id is supported by the stroage. This allow + us to enable automatic storage selection. + """ + ... diff --git a/lib/python3.12/site-packages/torch/distributed/checkpoint/utils.py b/lib/python3.12/site-packages/torch/distributed/checkpoint/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cb90dd119121593f5df5bf355f30df98361f5764 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/checkpoint/utils.py @@ -0,0 +1,431 @@ +# mypy: allow-untyped-defs +import cProfile +import inspect +import io +import itertools +import os +import warnings +from contextlib import contextmanager +from functools import wraps +from pstats import Stats +from typing import Any, Callable, cast, Dict, List, Optional, Sequence, TypeVar, Union + +import torch +import torch.distributed as dist +from torch.distributed._shard.sharded_tensor import ShardedTensor +from torch.distributed._shard.sharded_tensor.shard import Shard + +from .api import ( + _is_wrapped_exception, + _wrap_exception, + CheckpointException, + WRAPPED_EXCEPTION, +) +from .metadata import MetadataIndex, STATE_DICT_TYPE + + +__all__ = ["find_tensor_shard", "find_state_dict_object"] + +T = TypeVar("T") +R = TypeVar("R") + + +def _get_failure_dict( + results: List[Union[T, WRAPPED_EXCEPTION]] +) -> Dict[int, WRAPPED_EXCEPTION]: + return cast( + Dict[int, WRAPPED_EXCEPTION], + {i: err for i, err in enumerate(results) if _is_wrapped_exception(err)}, + ) + + +def _all_gather_keys( + local_dict: Dict[Any, Any], group: Optional[dist.ProcessGroup] = None +) -> List[Any]: + """Gathers all keys, and returns them sorted.""" + keys = list(local_dict.keys()) + gathered_keys: List[List[Any]] = [None] * dist.get_world_size(group) # type: ignore[list-item] + + dist.all_gather_object(gathered_keys, keys, group=group) + return sorted(set(itertools.chain.from_iterable(gathered_keys))) + + +class _DistWrapper: + """ + This is a wrapper around PG that provides a series of features around object collectives. + + It works without distributed initialized, where most collectives turns into nops. + + All variants that take functions are exception robust, meaning that if one or more + ranks raise errors, all ranks will observe those. + """ + + def __init__( + self, + group: Optional[dist.ProcessGroup], + use_dist: bool, + coordinator_rank: int, + ): + self.group = group + self.use_dist = use_dist + self.coordinator_rank = coordinator_rank + if self.use_dist: + self.rank = dist.get_rank(group) + self.is_coordinator = self.rank == coordinator_rank + else: + self.rank = 0 + self.is_coordinator = True + + def get_rank(self) -> int: + return self.rank + + def get_world_size(self) -> int: + if self.use_dist: + return dist.get_world_size(self.group) + return 1 + + def broadcast_object(self, object: Optional[T]) -> T: + """Implement functionality similar to c10d::broadcast_object_list but without distributed enabled.""" + object_list = [object] + if self.use_dist: + dist.broadcast_object_list( + object_list=object_list, + group=self.group, + src=self.coordinator_rank, + ) + return cast(T, object_list[0]) + + def gather_object(self, object: T) -> Optional[List[T]]: + """Implement functionality similar to c10d::gather_object but without distributed enabled.""" + if self.use_dist: + gather_objs = ( + cast(List[T], [None] * dist.get_world_size(self.group)) + if self.is_coordinator + else None + ) + + dist.gather_object( + obj=object, + object_gather_list=gather_objs if self.is_coordinator else None, + dst=self.coordinator_rank, + group=self.group, + ) + result = gather_objs + else: + result = [object] + return result + + def all_gather_object(self, object: T) -> List[T]: + """Implement functionality similar to c10d::all_gather_object but without distributed enabled.""" + if self.use_dist: + gather_objs = cast(List[T], [None] * dist.get_world_size(self.group)) + + dist.all_gather_object( + object_list=gather_objs, obj=object, group=self.group + ) + else: + gather_objs = [object] + return gather_objs + + def scatter_object(self, object_list: Optional[List[T]]) -> T: + """Implement functionality similar to c10d::scatter_object but without distributed enabled.""" + if self.use_dist: + gather_result = cast(List[T], [None]) + dist.scatter_object_list( + scatter_object_output_list=gather_result, + scatter_object_input_list=object_list if self.is_coordinator else None, + src=self.coordinator_rank, + group=self.group, + ) + + local_reply = gather_result[0] + else: + assert object_list is not None + local_reply = object_list[0] + return local_reply + + def reduce_scatter( + self, + step: str, + map_fun: Callable[[], T], + reduce_fun: Callable[[List[T]], List[R]], + ) -> R: + """ + Compute a value on each rank, then do centralized reduce on a single rank, followed by a scatter. + + This method operates in the following way: + Run ``map_fun`` on all ranks + Gather results on rank 0 + Call ``reduce_fun`` on all those values + Scatter to each rank part of the result. + """ + local_data: Union[WRAPPED_EXCEPTION, T] + try: + local_data = map_fun() + except BaseException as e: + local_data = _wrap_exception(e) + + all_data = self.gather_object(local_data) + all_results: Optional[List[Union[R, CheckpointException]]] = None + if self.is_coordinator: + assert all_data is not None + node_failures = _get_failure_dict(all_data) + + if len(node_failures) == 0: + try: + # N.B. why can't mypy cast List[R] to List[Union[R, WRAPPED_EXCEPTION]]? + all_results = cast( + List[Union[R, CheckpointException]], + reduce_fun(cast(List[T], all_data)), + ) + except BaseException as e: + node_failures[self.rank] = _wrap_exception(e) + + if len(node_failures) > 0: + all_results = [ + CheckpointException(step, node_failures) + ] * self.get_world_size() + + result = self.scatter_object(all_results) + if isinstance(result, CheckpointException): + raise result + return result + + def all_reduce( + self, + step: str, + map_fun: Callable[[], T], + reduce_fun: Callable[[List[T]], R], + ) -> R: + """ + Compute a value on each rank, then do centralized reduce on a single rank, followed by a broadcast. + + This method operates in the following way: + Run ``map_fun`` on all ranks + Gather results on rank 0 + Call ``reduce_fun`` on all those values + Broadcast the reduced value to all ranks. + """ + local_data: Union[T, WRAPPED_EXCEPTION] + try: + local_data = map_fun() + except BaseException as e: + local_data = _wrap_exception(e) + + all_data = self.gather_object(local_data) + result: Optional[Union[R, CheckpointException]] = None + if self.is_coordinator: + assert all_data is not None + node_failures = _get_failure_dict(all_data) + if len(node_failures) == 0: + try: + result = reduce_fun(cast(List[T], all_data)) + except BaseException as e: + node_failures[self.rank] = _wrap_exception(e) + + if len(node_failures) > 0: + result = CheckpointException(step, node_failures) + + final_result = self.broadcast_object(result) + if isinstance(final_result, CheckpointException): + raise final_result + return cast(R, final_result) + + def all_gather( + self, + step: str, + map_fun: Callable[[], T], + ) -> List[T]: + """ + Compute a value on each rank, then all_gather them. + + This method operates in the following way: + Run ``map_cp`` on all ranks + all_gather the values to all ranks + """ + result: Union[T, WRAPPED_EXCEPTION] + try: + result = map_fun() + except BaseException as e: + result = _wrap_exception(e) + + all_results = self.all_gather_object(result) + + node_failures = _get_failure_dict(all_results) + if len(node_failures) > 0: + raise CheckpointException(step, node_failures) + return cast(List[T], all_results) + + def broadcast( + self, + step: str, + map_fun: Callable[[], T], + ) -> T: + """ + Compute a value on rank 0 and broadcast it. + + This method operates in the following way: + Run ``map_cp`` on rank 0 + broadcast the value + """ + result: Optional[Union[T, CheckpointException]] = None + if self.is_coordinator: + try: + result = map_fun() + except BaseException as e: + result = CheckpointException(step, {self.rank: _wrap_exception(e)}) + final_result = self.broadcast_object(result) + if isinstance(final_result, CheckpointException): + raise final_result + return cast(T, final_result) + + +def _find_shard(tensor: ShardedTensor, index: MetadataIndex) -> Shard: + if index.offset is None: + raise ValueError( + f"Cannot lookup {index.fqn} since its a ShardedTensor and no offset was provided" + ) + + shards = tensor.local_shards() + # index fast path + if index.index is not None: + if ( + len(shards) > index.index + and torch.Size(shards[index.index].metadata.shard_offsets) == index.offset + ): + return shards[index.index] + + for shard in shards: + if torch.Size(shard.metadata.shard_offsets) == index.offset: + return shard + raise ValueError(f"Could not find shard at '{index.offset}' for FQN: '{index.fqn}'") + + +def find_tensor_shard(tensor: torch.Tensor, index: MetadataIndex) -> torch.Tensor: + if hasattr(tensor, "__get_tensor_shard__"): + # DTensor implements _Checkpointable + return tensor.__get_tensor_shard__(index) # type: ignore[attr-defined] + if isinstance(tensor, ShardedTensor): + return _find_shard(tensor, index).tensor + if index.offset is not None: + # special case looking up a tensor by origin + if index.offset == torch.Size([0] * len(tensor.size())): + return tensor + raise ValueError( + f"FQN: '{index.fqn}' is not a ShardedTensor, can't find by offset: '{index.offset}'" + ) + return tensor + + +def find_state_dict_object(state_dict: STATE_DICT_TYPE, index: MetadataIndex) -> Any: + if index.fqn not in state_dict: + raise ValueError(f"Could not find FQN: '{index.fqn}'") + obj = state_dict[index.fqn] + + if isinstance(obj, torch.Tensor): + return find_tensor_shard(obj, index) + elif index.offset is not None: + raise ValueError( + f"FQN: '{index.fqn}' is not a ShardedTensor, can't find by offset: '{index.offset}'" + ) + return obj + + +def _element_wise_add(a: Sequence[int], b: Sequence[int]) -> List[int]: + return [i_a + i_b for i_a, i_b in zip(a, b)] + + +def _element_wise_sub(a: Sequence[int], b: Sequence[int]) -> List[int]: + return [i_a - i_b for i_a, i_b in zip(a, b)] + + +class _ReaderView(io.IOBase): + def __init__(self, base_stream: io.IOBase, offset: int, len: int): + super().__init__() + self.offset = offset + self.len = len + self.base_stream = base_stream + self.seek(0) + + def seek(self, __offset: int, __whence: int = os.SEEK_SET) -> int: + if __whence == os.SEEK_SET: + __offset = self.offset + __offset + elif __whence == os.SEEK_END: + __whence = os.SEEK_SET + __offset = (self.offset + self.len) - __offset + return self.base_stream.seek(__offset, __whence) + + def tell(self) -> int: + return self.base_stream.tell() - self.offset + + def readable(self) -> bool: + return self.base_stream.readable() + + def seekable(self) -> bool: + return self.base_stream.seekable() + + def readinto(self, b): + return self.base_stream.readinto(b) # type: ignore[attr-defined] + + def read(self, size=-1): + return self.base_stream.read(size) + + +def _create_file_view(file: io.IOBase, offset: int, length: int) -> io.IOBase: + # FIXME (kumpera) torch.load fails if we wrap with io.BufferedReader + return _ReaderView(file, offset, length) + + +def _normalize_device_info(device_type: str, device_id: int) -> str: + """Device info normalization.""" + if device_type == "cpu": + return "cpu" + return f"{device_type}:{device_id}" + + +# TODO: integrate with distributed logging flag +ENABLE_PROFILE = False + + +@contextmanager +def _profile(): + # Only log the profiling when it is enable and is on rank0 or dist is not + # avaiable. + if ENABLE_PROFILE and (not dist.is_available() or dist.get_rank() == 0): + profiler = cProfile.Profile() + profiler.enable() + try: + yield + finally: + profiler.disable() + stats = Stats(profiler) + stats.sort_stats("time").print_stats(10) + else: + yield + + +def _api_bc_check(func): + @wraps(func) + def inner_func(*args, **kwargs) -> Any: + if len(args) == 2: + warnings.warn( + f"The argument order of {func.__name__} has been changed. " + "Please check the document to avoid future breakages." + ) + sig = inspect.signature(func) + kwonlyargs = [ + p.name for p in sig.parameters.values() if p.kind == p.KEYWORD_ONLY + ] + if "storage_writer" in kwonlyargs: + assert "storage_writer" not in kwargs, (args, kwargs) + kwargs["storage_writer"] = args[1] + elif "storage_reader" in kwonlyargs: + assert "storage_reader" not in kwargs, (args, kwargs) + kwargs["storage_reader"] = args[1] + else: + raise RuntimeError(f"Unexpected kwonlyargs = {kwonlyargs}") + return func(args[0], **kwargs) + else: + return func(*args, **kwargs) + + return inner_func diff --git a/lib/python3.12/site-packages/torch/distributed/collective_utils.py b/lib/python3.12/site-packages/torch/distributed/collective_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..78199e7a26f22cee28564fb66fd8bbdbc8a1defa --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/collective_utils.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python3 + + +""" +A set of primitive functions for performing collective ops. + +Each should also handle single rank scenario. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, cast, Generic, List, Optional, Tuple, TypeVar, Union + +import torch.distributed as dist + + +T = TypeVar("T") + + +@dataclass +class SyncPayload(Generic[T]): + stage_name: Optional[str] + success: bool + payload: T + exception: Optional[Exception] = None + + +def broadcast( + data_or_fn: Union[T, Callable[[], T]], + *, + success: bool = True, + stage_name: Optional[str] = None, + rank: int = 0, + pg: Optional[dist.ProcessGroup] = None, +) -> T: + """ + Broadcasts the data payload from rank 0 to all other ranks. + Or if a function is passed, execute it in rank 0 and broadcast result to all other ranks. + + Can be used to broadcast a failure signal to stop all ranks. + + If the function raises an exception, all ranks will raise. + + Args: + data_or_fn: the data to broadcast or function to execute and broadcast result. + success: False to stop all ranks. + stage_name: the name of the logical stage for synchronization and debugging + rank: rank to broadcast data or execute function and broadcast resutls. + pg: the process group for sync + Throws: + RuntimeError from original exception trace + Returns: + the value after synchronization + + Example usage: + >> id = broadcast(data_or_fn=allocate_id, rank=0, pg=ext_pg.my_pg) + """ + + if not success and data_or_fn is not None: + raise AssertionError( + "Data or Function is expected to be None if not successful" + ) + + payload: Optional[T] = None + exception: Optional[Exception] = None + # if no pg is passed then execute if rank is 0 + if (pg is None and rank == 0) or (pg is not None and pg.rank() == rank): + # determine if it is an executable function or data payload only + if callable(data_or_fn): + try: + payload = data_or_fn() + except Exception as e: + success = False + exception = e + else: + payload = data_or_fn + + # broadcast the exception type if any to all ranks for failure categorization + sync_obj = SyncPayload( + stage_name=stage_name, + success=success, + payload=payload, + exception=exception, + ) + + if pg is not None: + broadcast_list = [sync_obj] + dist.broadcast_object_list(broadcast_list, src=rank, group=pg) + assert len(broadcast_list) == 1 + sync_obj = broadcast_list[0] + + # failure in any rank will trigger a throw in every rank. + if not sync_obj.success: + error_msg = f"Rank {rank} failed" + if stage_name is not None: + error_msg += f": stage {sync_obj.stage_name}" + if sync_obj.exception is not None: + error_msg += f": exception {sync_obj.exception}" + raise RuntimeError(error_msg) from sync_obj.exception + + return cast(T, sync_obj.payload) + + +def all_gather( + data_or_fn: Union[T, Callable[[], T]], + stage_name: Optional[str] = None, + pg: Optional[dist.ProcessGroup] = None, +) -> List[T]: + """ + A simple all_gather primitive with basic synchronization guard logic, + by checking payload from all ranks has the same stage name. + + Args: + data_or_fn: the data to be all gathered across ranks or function to be executed + stage_name: the sync stage name for out-of-sync protection + pg: the process group for sync + Throws: + RuntimeError from original exception trace + Returns: + a list of synced data from all ranks + + Example usage: + >> all_ids = all_gather(data_or_fn=allocate_id, pg=ext_pg.my_pg) + """ + payload: Optional[T] = None + exception: Optional[Exception] = None + success = True + # determine if it is an executable function or data payload only + if callable(data_or_fn): + try: + payload = data_or_fn() + except Exception as e: + success = False + exception = e + else: + payload = data_or_fn + + sync_obj = SyncPayload( + stage_name=stage_name, + success=success, + payload=payload, + exception=exception, + ) + + if pg is not None: + # List of success/failure across all ranks. + total_list = [None] * dist.get_world_size(pg) + all_gather_object_enforce_type(pg, total_list, sync_obj) + # Each rank will throw RuntimeError in case of failure on any rank. + stage_name = cast(SyncPayload[T], total_list[0]).stage_name + exception_list: List[Tuple[int, Exception]] = [] + ret_list: List[T] = [] + error_msg: str = "" + + for i, sp in enumerate(cast(List[SyncPayload[T]], total_list)): + if sp.stage_name != stage_name: + error_msg += ( + f"Unexpected stage name received from rank {i}: {sp.stage_name} " + ) + continue + if not sp.success and sp.exception is not None: + exception_list.append((i, sp.exception)) + continue + ret_list.append(sp.payload) + + if len(exception_list) > 0: + raise RuntimeError( # type: ignore[misc] + error_msg, exception_list + ) from exception_list[0] + return ret_list + else: + if not sync_obj.success: + raise RuntimeError( + f"all_gather failed with exception {sync_obj.exception}", + ) from sync_obj.exception + return [sync_obj.payload] # type: ignore[list-item] + + +# Note: use Any for typing for now so users can pass in +# either a list of None or target type placeholders +# otherwise pyre would complain +def all_gather_object_enforce_type( + pg: dist.ProcessGroup, + # pyre-fixme[2]: Parameter must have a type that does not contain `Any` + object_list: List[Any], + # pyre-fixme[2]: Parameter must have a type other than `Any` + obj: Any, + # pyre-fixme[2]: Parameter must have a type that does not contain `Any` + type_checker: Callable[[Any, Any], bool] = lambda x, y: type(x) == type(y), +) -> None: + """ + Similar to plain all_gather_object but with additional type checking + AFTER gather is done to ensure basic consistency. + If check does not pass, all ranks will fail with exception. + + This is generally to prevent conditional logic leading to + unexpected messages being received. This is considered fatal code error, + but due to logic stacks this might happen implicitly in practice. + + The default check does not check sub type (considered different) + or covariance (considered same) but users can pass in custom checker + if more complicated check is needed. + """ + dist.all_gather_object(object_list, obj, group=pg) + + # conservative check + list_len = len(object_list) + if list_len == 0: + return + first_obj = object_list[0] + for i in range(1, list_len): + if not type_checker(first_obj, object_list[i]): + raise TypeError( + f"Object type at index {i} is {type(object_list[i])}, " + f"while first object type is {type(first_obj)}" + ) diff --git a/lib/python3.12/site-packages/torch/distributed/constants.py b/lib/python3.12/site-packages/torch/distributed/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..b3754043644b8cc96feee06a4fc1cc9fa824b648 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/constants.py @@ -0,0 +1,26 @@ +from datetime import timedelta +from typing import Optional + +from torch._C._distributed_c10d import _DEFAULT_PG_TIMEOUT + + +__all__ = ["default_pg_timeout", "default_pg_nccl_timeout"] + +# Default process group wide timeout, if applicable. +# This only applies to the non-nccl backends +# To make an attempt at backwards compatibility with THD, we use an +# extraordinarily high default timeout, given that THD did not have timeouts. +default_pg_timeout: timedelta = _DEFAULT_PG_TIMEOUT +# Separate timeout for PGNCCL mainly becuase it's always been that way in the C++ layer, but until recently +# there was one default that applied across all backends in the python layer. +# Later, we could consider merging them back together at the c++ layer if we can align on a same value. +# (only if TORCH_NCCL_BLOCKING_WAIT or TORCH_NCCL_ASYNC_ERROR_HANDLING is set to 1). + +try: + from torch._C._distributed_c10d import _DEFAULT_PG_NCCL_TIMEOUT + + default_pg_nccl_timeout: Optional[timedelta] = _DEFAULT_PG_NCCL_TIMEOUT +except ImportError: + # if C++ NCCL support is not compiled, we don't have access to the default nccl value. + # if anyone is actually trying to use nccl in this state, it should error. + default_pg_nccl_timeout = None diff --git a/lib/python3.12/site-packages/torch/distributed/device_mesh.py b/lib/python3.12/site-packages/torch/distributed/device_mesh.py new file mode 100644 index 0000000000000000000000000000000000000000..c8a1e257873607004c216b585121cb997d95b143 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/device_mesh.py @@ -0,0 +1,1009 @@ +# mypy: allow-untyped-defs +# Copyright (c) Meta Platforms, Inc. and affiliates +import logging +import math +import threading +from functools import reduce +from itertools import chain +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union + +import torch +from torch.distributed import is_available +from torch.utils._typing_utils import not_none + + +__all__ = ["init_device_mesh", "DeviceMesh"] + + +if not is_available(): + import sys + + # We need to create the stubs when distributed is not available. + # Otherwise, we would fail the doc tests (```./.ci/pytorch/docs-test.sh```), + # since it would try to import ``torch.distributed.device_mesh`` or + # ``torch.distributed.init_device_mesh`` but cannot find them. + + class _DeviceMeshStub: + pass + + def _init_device_mesh_stub(): + pass + + sys.modules["torch.distributed.device_mesh"].DeviceMesh = _DeviceMeshStub # type: ignore[attr-defined] + sys.modules[ + "torch.distributed.device_mesh" + ].init_device_mesh = _init_device_mesh_stub # type: ignore[attr-defined] + + +else: + from torch._C._distributed_c10d import Backend as C10dBackend + from torch.distributed.distributed_c10d import ( + _find_pg_by_ranks_and_tag, + _get_default_group, + _get_group_tag, + get_backend, + get_process_group_ranks, + get_rank, + get_world_size, + init_process_group, + is_initialized, + new_group, + ProcessGroup, + split_group, + ) + + logger = logging.getLogger(__name__) + + # only import numpy typing when type checking + if TYPE_CHECKING: + try: + from numpy.typing import ArrayLike + except ImportError: + logger.warning( + "DeviceMesh requires numpy >= 1.21 to be installed for type checking" + ) + + class _MeshEnv(threading.local): + def __init__(self) -> None: + self.mesh_stack: List[DeviceMesh] = [] + self.child_to_root_mapping: Dict[DeviceMesh, DeviceMesh] = {} + self.mesh_dim_group_options: Dict[ + int, Tuple[str, Optional[C10dBackend.Options]] + ] = {} + self.root_to_flatten_mapping: Dict[DeviceMesh, Dict[str, DeviceMesh]] = {} + # Record flatten mesh name to its mesh dim index in root mesh. + self.flatten_name_to_root_dims: Dict[ + DeviceMesh, Dict[str, Tuple[int, ...]] + ] = {} + + def get_current_mesh(self) -> "DeviceMesh": + if len(self.mesh_stack) == 0: + raise RuntimeError("No device mesh is currently active!") + return self.mesh_stack[-1] + + def create_sub_mesh( + self, + device_mesh: "DeviceMesh", + submesh_dim_names: Tuple[str, ...], + submesh_dims: List[Tuple[int, ...]], + ) -> "DeviceMesh": + # Get the submesh dim size from the submesh_dims. + # For example, if we have a 3D mesh with mesh_shape (2, 2, 2) mesh_dim_names ("dp", "cp", "tp") and we want + # to slice out mesh["dp_cp"], then submesh_dims = [(0, 1), (2,)] and submesh_dim_size = [2 * 2, 2] = [4, 2]. + # If we want to slice out mesh["dp", "cp"], then submesh_dims = [(0,), (1,)] and submesh_dim_size = [2, 2]. + slice_dim_size = [ + reduce( + lambda x, y: x * device_mesh.mesh.size(y), + mesh_dim, + 1, + ) + for mesh_dim in submesh_dims + ] + + mesh_tensor = device_mesh.mesh + # slice_dim_idx could be differnt from submesh_dims, as we may need to flatten out some dims. + slice_dim_idx = [] + slice_dim_group_info = [] + # keep track of the number of dims that have been flattened so we can get the correct slice_dim_idx in the + # flattened mesh tensor. + num_dims_flatten = 0 + for mesh_dim_indices, mesh_dim_name in zip(submesh_dims, submesh_dim_names): + # Currently, this only allows slicing out a contiguous flattened dim. + # TODO: we need to handle reconstructing a non-contiguous flattened dim. + if len(mesh_dim_indices) > 1: + # We need to move the start_dim and end_dim to the left if some dims are already flattened. + mesh_tensor = mesh_tensor.flatten( + start_dim=mesh_dim_indices[0] - num_dims_flatten, + end_dim=mesh_dim_indices[-1] - num_dims_flatten, + ) + # If some dims are already flattened, we need to adjust the slice_dim_idx accordingly. + # For example, if the submesh_dims = [(0, 1), (2,), (3, 4)] with 0-1 flattened and 3-4 flattened, + # then the final slice_dim_idx should be [0, 1, 2]. + slice_dim_idx.append(mesh_dim_indices[0] - num_dims_flatten) + num_dims_flatten += len(mesh_dim_indices) - 1 + slice_dim_group_info.append( + self.root_to_flatten_mapping[device_mesh][ + mesh_dim_name + ]._dim_group_infos[0] + ) + else: + slice_dim_idx.append(mesh_dim_indices[0] - num_dims_flatten) + slice_dim_group_info.append( + device_mesh._dim_group_infos[mesh_dim_indices[0]] + ) + + # mesh_tensor has already been flattened if needed. So mesh_tensor.ndim <= device_mesh.mesh.ndim now. + mesh_dims_remained_idx = list(range(mesh_tensor.ndim)) + for idx in slice_dim_idx: + mesh_dims_remained_idx.remove(idx) + + # pg_ranks_by_dim is the size of [number of local ranks of the outermost submesh dimension, *slice_dim_idx] + # This means on each local rank of the outermost slice mesh dim, we have a tensor of submesh size with + # the pg ranks of the submesh. From this, we can extract the submesh mesh tensor contains the current rank. + pg_ranks_by_dim = mesh_tensor.permute( + *mesh_dims_remained_idx, *slice_dim_idx + ).reshape(-1, *slice_dim_size) + + cur_rank = device_mesh.get_rank() + for mesh_nd in pg_ranks_by_dim: + submesh = DeviceMesh( + device_mesh.device_type, + mesh_nd, + mesh_dim_names=submesh_dim_names, + _init_backend=False, + ) + if cur_rank in mesh_nd: + res_submesh = submesh + + res_submesh._dim_group_infos = slice_dim_group_info # type: ignore[possibly-undefined] + self.child_to_root_mapping[res_submesh] = device_mesh + + return res_submesh + + def create_flatten_mesh( + self, device_mesh: "DeviceMesh", mesh_dim_name: Optional[str] = None + ) -> "DeviceMesh": + root_mesh = _mesh_resources.get_root_mesh(device_mesh) + + flatten_dims_in_root = [ + not_none(root_mesh.mesh_dim_names).index(flattened_mesh_dim_name) + for flattened_mesh_dim_name in not_none(device_mesh.mesh_dim_names) + ] + + if not mesh_dim_name: + mesh_dim_name = "_".join( + [ + not_none(root_mesh.mesh_dim_names)[dim] + for dim in flatten_dims_in_root + ] + ) + + # Check whether the mesh_dim_name for flattened mesh is valid. + self.flatten_name_to_root_dims.setdefault(root_mesh, {}) + invalid_dim_names = chain( + *list(not_none(root_mesh.mesh_dim_names)), + *self.flatten_name_to_root_dims[root_mesh].keys(), + ) + if mesh_dim_name in invalid_dim_names: + raise RuntimeError( + f"{mesh_dim_name} already exists for submesh of the {root_mesh}. ", + f"The mesh_dim_names of submesh and flattened mesh are {invalid_dim_names}. " + f"Please specify another valid mesh_dim_name.", + ) + + # Quick return if the flatten mesh has been created before. + # TODO: If we decide to restrict flatten initialization once, we should remove + # this check and throw an error if the flatten mesh is already created before. + if ( + root_mesh in self.root_to_flatten_mapping + and mesh_dim_name in self.root_to_flatten_mapping[root_mesh] + ): + return self.root_to_flatten_mapping[root_mesh][mesh_dim_name] + + flattened_mesh_dim_size = math.prod(device_mesh.mesh.size()) + + remained_dims_in_root = list(range(root_mesh.mesh.ndim)) + for flatten_dim_in_root in flatten_dims_in_root: + remained_dims_in_root.remove(flatten_dim_in_root) + + pg_ranks_by_dim = root_mesh.mesh.permute( + *remained_dims_in_root, *flatten_dims_in_root + ).reshape(-1, flattened_mesh_dim_size) + + cur_rank = root_mesh.get_rank() + for mesh_nd in pg_ranks_by_dim: + # need to init backend here since the flattened pg doesn't exist in root mesh. + flattened_mesh = DeviceMesh( + root_mesh.device_type, + mesh_nd, + mesh_dim_names=(mesh_dim_name,), + ) + if cur_rank in mesh_nd: + res_flattened_mesh = flattened_mesh + self.child_to_root_mapping[res_flattened_mesh] = root_mesh # type: ignore[possibly-undefined] + self.root_to_flatten_mapping.setdefault(root_mesh, {})[mesh_dim_name] = res_flattened_mesh # type: ignore[possibly-undefined] + self.flatten_name_to_root_dims[root_mesh][mesh_dim_name] = tuple(flatten_dims_in_root) # type: ignore[possibly-undefined] + + return res_flattened_mesh + + def get_root_mesh(self, device_mesh: "DeviceMesh") -> "DeviceMesh": + # If a mesh could not be found in the child_to_root_mapping, it is a root mesh itself. + # A root mesh is not created through slicing. + # We considers the root mesh of a root mesh is itself. + root_mesh = self.child_to_root_mapping.get(device_mesh, None) + return device_mesh if not root_mesh else root_mesh + + def get_root_mesh_dim(self, device_mesh: "DeviceMesh") -> Optional[int]: + """ + Returns the index of the mesh dim in the root mesh. + The device_mesh passed in needs to be sliced out from the root mesh + or submesh of the root mesh. + """ + root_mesh = self.get_root_mesh(device_mesh) + child_mesh_dim_names = device_mesh.mesh_dim_names + if root_mesh and child_mesh_dim_names: + assert ( + len(child_mesh_dim_names) == 1 + ), "The submesh can only be a 1D mesh." + child_mesh_dim_name = child_mesh_dim_names[0] + return self.get_mesh_dim_by_name(root_mesh, child_mesh_dim_name) + return None + + @staticmethod + def num_devices_per_host(device_type: str) -> int: + return _get_device_handle(device_type).device_count() + + @staticmethod + def num_hosts(device_type: str) -> int: + # ProcessGroup can't tell us this info so we have to infer it, assume + # homogeneous hardware for now + return get_world_size() // _MeshEnv.num_devices_per_host(device_type) + + def get_mesh_dim_by_name( + self, device_mesh: "DeviceMesh", mesh_dim_name: str + ) -> int: + if ( + device_mesh.mesh_dim_names is None + or len(device_mesh.mesh_dim_names) == 0 + ): + raise KeyError( + "No `mesh_dim_names` found.", + ) + if mesh_dim_name not in device_mesh.mesh_dim_names: + raise KeyError( + f"Mesh dimension '{mesh_dim_name}' does not exist.", + f"Available mesh dimensions are: mesh_dim_names={device_mesh.mesh_dim_names}", + ) + return not_none(device_mesh.mesh_dim_names.index(mesh_dim_name)) + + def _set_mesh_dim_group_options( + self, + dim: int, + backend: str, + pg_options: Optional[C10dBackend.Options] = None, + ) -> None: + self.mesh_dim_group_options[dim] = (backend, pg_options) + + def _get_slice_mesh_dims( + self, device_mesh, mesh_dim_names + ) -> List[Tuple[int, ...]]: + """ + Validate whether the mesh_dim_names is valid for slicing the given device_mesh. + If valid, return dim indexes of the slice mesh in the device mesh. + """ + if device_mesh != self.get_root_mesh(device_mesh): + raise RuntimeError("Cannot create a submesh from a submesh.") + + # The slice mesh_dim_names should consist either the device_mesh's mesh_dim_names + # or its flattened mesh's mesh_dim_names. + self.flatten_name_to_root_dims.setdefault(device_mesh, {}) + flatten_name_to_root_dims = self.flatten_name_to_root_dims[device_mesh] + valid_mesh_dim_names = [ + *device_mesh.mesh_dim_names, + *flatten_name_to_root_dims, + ] + + if not all( + mesh_dim_name in valid_mesh_dim_names + for mesh_dim_name in mesh_dim_names + ): + raise KeyError( + f"Invalid mesh_dim_names {mesh_dim_names} specified. " + f"Valid mesh_dim_names are {valid_mesh_dim_names}." + ) + + # Validate the order of the slice mesh dim indices. + # This needs to be in ascending order. + curr_idx = -1 + slice_mesh_dims = [] + for mesh_dim_name in mesh_dim_names: + if mesh_dim_name in flatten_name_to_root_dims: + mesh_indices = flatten_name_to_root_dims[mesh_dim_name] + # TODO: this doesn't allow non-contiguous slicing with flatten dim yet. next_idx + # should be mesh_indices[0] once we support non-contiguous slicing with flatten dim. + next_idx = mesh_indices[-1] + slice_mesh_dims.append(mesh_indices) + else: + next_idx = device_mesh.mesh_dim_names.index(mesh_dim_name) + slice_mesh_dims.append((next_idx,)) + if next_idx <= curr_idx: + raise KeyError( + f"Invalid mesh_dim_names {mesh_dim_names} specified. ", + f"Found mesh dim indices to slice: {slice_mesh_dims}. ", + "Mesh dim indices should be in ascending order.", + ) + curr_idx = next_idx + + return slice_mesh_dims + + def _get_all_submeshes( + self, device_mesh: "DeviceMesh", mesh_dim_name: str + ) -> List["DeviceMesh"]: + """ + Return all the submeshes of a given mesh dimension of the device mesh. + """ + mesh_dim = self.get_mesh_dim_by_name(device_mesh, mesh_dim_name) + pg_ranks_by_dim = device_mesh.mesh.swapdims(-1, mesh_dim).reshape( + -1, device_mesh.mesh.size(mesh_dim) + ) + + cur_rank = device_mesh.get_rank() + res_submeshes = [] + for mesh_1d in pg_ranks_by_dim: + submesh = DeviceMesh( + device_mesh.device_type, + mesh_1d, + mesh_dim_names=(mesh_dim_name,), + _init_backend=False, + ) + submesh._dim_group_infos = ( + [device_mesh._dim_group_infos[mesh_dim]] + if cur_rank in mesh_1d + else [] + ) + res_submeshes.append(submesh) + + return res_submeshes + + _mesh_resources: _MeshEnv = _MeshEnv() + + def _get_device_handle(device_type: str = "cuda"): + """ + Get the module corresponding to the device_type which is cuda or cuda-like device. + For example, when the device_type is cuda, the module `torch.cuda` is returned. + Return None when there is no corresponding module for device_type, otherwise + return the corresponding module. + """ + return getattr(torch, device_type, None) + + class DeviceMesh: + """ + DeviceMesh represents a mesh of devices, where layout of devices could be + represented as a n-d dimension array, and each value of the n-d dimensional + array is the global id of the default process group ranks. + + DeviceMesh could be used to describe the layout of devices across the cluster, + and serves as a proxy for communication among the device lists within the cluster. + + DeviceMesh can be used as a context manager. + + .. note:: + DeviceMesh follows SPMD programming model, which means the same PyTorch Python program + is running on all processes/ranks in the cluster. Therefore, users need to make sure the + `mesh` array (which describes the layout of devices) should be identical across all ranks. + Inconsistent `mesh` will lead to silent hang. + + Args: + device_type (str): The device type of the mesh. Currently supports: "cpu", "cuda/cuda-like". + mesh (ndarray): A multi-dimensional array or an integer tensor describing the layout + of devices, where the IDs are global IDs of the default process group. + + Returns: + DeviceMesh: A :class:`DeviceMesh` object representing the device layout. + + The following program runs on each process/rank in an SPMD manner. In this example, we have 2 + hosts with 4 GPUs each. + A reduction over the first dimension of mesh will reduce across + columns (0, 4), .. and (3, 7), a reduction over the second dimension + of mesh reduces across rows (0, 1, 2, 3) and (4, 5, 6, 7). + + Example:: + >>> # xdoctest: +SKIP("no rank") + >>> from torch.distributed.device_mesh import DeviceMesh + >>> + >>> # Initialize device mesh as (2, 4) to represent the topology + >>> # of cross-host(dim 0), and within-host (dim 1). + >>> mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]]) + """ + + device_type: str + mesh: torch.Tensor + mesh_dim_names: Optional[Tuple[str, ...]] + + def __init__( + self, + device_type: str, + mesh: Union[torch.Tensor, "ArrayLike"], + *, + mesh_dim_names: Optional[Tuple[str, ...]] = None, + _init_backend: bool = True, + ) -> None: + self.device_type = device_type + if isinstance(mesh, torch.Tensor) and mesh.device.type != "cpu": + raise ValueError(f"`mesh` must be a CPU tensor, got {mesh}") + self.mesh = ( + mesh.detach().to(dtype=torch.int) + if isinstance(mesh, torch.Tensor) + else torch.tensor(mesh, device="cpu", dtype=torch.int) + ) + self.mesh_dim_names = tuple(mesh_dim_names) if mesh_dim_names else None + + # private field to pre-generate DeviceMesh's hash + self._flatten_mesh_list = tuple(self.mesh.flatten().tolist()) + self._thread_id = None + + # Skip process group initialization if xla device or init backend is False + # TODO(yeounoh) implement DeviceMesh backend and register XLA backend. + if device_type != "xla": + # always try to create default (world) pg, even if it is not initialized + # already. The world pg is used for device mesh identity (rank) on each + # process (we need to know if the current global rank is in the mesh or not). + if _init_backend: + self._get_or_create_default_group() + self._init_process_groups() + + if is_initialized() and get_backend() == "threaded": + self._thread_id = threading.get_ident() + + # calculate the coordinates of the current global rank on the mesh + rank_coords = (self.mesh == get_rank()).nonzero() + assert rank_coords.size(0) in (0, 1) + self._coordinate_on_dim: Optional[List[int]] = ( + rank_coords[0].tolist() if rank_coords.size(0) > 0 else None + ) + + def _get_or_create_default_group(self): + default_initialized = is_initialized() + if not default_initialized: + init_process_group() + + world_size = get_world_size() + if self.mesh.numel() > world_size: + raise RuntimeError( + f"Mesh should not be bigger than default world size {world_size}, but found {self.mesh.numel()} ranks!" + ) + + device_handle = _get_device_handle(self.device_type) + # TODO: if user want to pass pg_options, offer a way to do it + if not default_initialized and device_handle: + # automatically set the current cuda/cuda-like device base on num of gpu devices available in each host + # NOTE: This device selection would only work for homogeneous hardware. + num_devices_per_host = device_handle.device_count() + if ( + world_size > num_devices_per_host + and world_size % num_devices_per_host != 0 + ): + raise RuntimeError( + f"DeviceMesh only support homogeneous hardware, but found " + f"{world_size} ranks and {num_devices_per_host} {self.device_type} devices!" + ) + device_handle.set_device(get_rank() % num_devices_per_host) + + return _get_default_group() + + def _init_process_groups(self): + # tag/ranks/group_name associated with each mesh dimension, each + # mesh dimension should have one sub-group per rank + # + # TODO(yifu): remove tag and ranks once we fully migrate to native + # functional collectives. See details in: + # https://github.com/pytorch/pytorch/issues/93173#issuecomment-1907095208 + dim_group_infos: List[Tuple[str, List[int], str]] = [] + default_group = _get_default_group() + + if self.mesh.ndim == 1 and self.mesh.numel() == get_world_size(): + # Append the default pg to the first dim groups only if the default pg is compatible with `self.device_type`. + # Otherwise, create new pg. + ranks = list(range(get_world_size())) + dim_group = ( + new_group( + backend="cpu:gloo,cuda:nccl", + ranks=ranks, + group_desc="mesh_default", + ) + if torch.cuda.is_available() + and get_backend(default_group) == "gloo" + else default_group + ) + dim_group_infos.append( + ( + _get_group_tag(dim_group), + ranks, + dim_group.group_name, + ) + ) + else: + # create sub pgs base on the mesh argument specified + for dim in range(self.mesh.ndim): + # swap the current dim to the last dim + # then reshape to flatten out other dims + pg_ranks_by_dim = self.mesh.swapdims(-1, dim).reshape( + -1, self.mesh.size(dim) + ) + + # Respect dim group options specified via _MeshEnv.set_dim_group_options(). + # Inherit from the parent group if no options are specified for the group. + if dim in _mesh_resources.mesh_dim_group_options: + ( + backend, + pg_options, + ) = _mesh_resources.mesh_dim_group_options[dim] + else: + backend, pg_options = None, None + + # If we have a 2D mesh with mesh_dim_names ("dp", "tp"), the group description + # of the subgroups would be `mesh_dim_dp` and `mesh_name_tp`. + # If the mesh doesn't not have a mesh_dim_names, then the group description of the + # subgroup would be `mesh_dim_0` and `mesh_dim_1`. + group_desc = ( + f"mesh_{self.mesh_dim_names[dim]}" + if self.mesh_dim_names + else f"mesh_dim_{dim}" + ) + + # If bound_device_id exists, it means the nccl communicator has been eagerly initialized + # so that we can use `split_group` to create subgroups through `ncclCommSplit`. + # In this case, we only need to make one API call (`split_group``) for the subgroup creation + # for each mesh dimension. In a 2 * 4 mesh, we only need to make 2 API calls per ranks to create + # all the subgroups. + # Otherwise, we need to make more than one API call (`new_group`) for subgroup creations. The + # numbers of API calls are equal to the number of subgroups for each mesh dimension. In a 2 * 4 + # mesh, we need to make 2 + 4 = 6 API calls per ranks to create all the subgroups. + dim_group = None + if ( + bound_device_id := getattr( + default_group, "bound_device_id", None + ) + ) is not None: + dim_group = split_group( + parent_pg=default_group, + pg_options=pg_options, + split_ranks=pg_ranks_by_dim.tolist(), + group_desc=group_desc, + ) + + # If the subgroup has been already created through `split_group`, we simply loop over `pg_ranks_by_dim` + # and append the `(group_tag, subgroup_ranks, and group_name)` tuple to the `dim_group_infos` list when + # the current rank is in the subgroup. + # Otherwise, we use `new_group` instead of `split_group` to create subgroups by looping over `pg_ranks_by_dim` + # along with appending information to the `dim_group_infos` list whenever necessary. + for dim_mesh in pg_ranks_by_dim: + subgroup_ranks = dim_mesh.tolist() + + # We temporarily revert the re-use subgroup, since it breaks two internal tests. + # Temporarily reverting to resolve test timeout while root-causing. + # TODO: Add two tests to cover internal tests scenarios and re-enable reuse subgroup if exists. + if bound_device_id is None: + dim_group = new_group( + ranks=subgroup_ranks, + backend=backend, + pg_options=pg_options, + group_desc=group_desc, + ) + + # only add to dim_groups if the current rank in the subgroup + if self.get_rank() in subgroup_ranks: + if len(dim_group_infos) > dim: + raise RuntimeError( + f"Each device mesh dimension should get only one process group, but got {self.get_rank()} " + f"in {subgroup_ranks}!" + ) + dim_group_infos.append( + ( + _get_group_tag(not_none(dim_group)), + subgroup_ranks, + dim_group.group_name, + ) + ) + self._dim_group_infos = dim_group_infos + + def __enter__(self) -> "DeviceMesh": + # set this mesh as the current mesh in mesh env + _mesh_resources.mesh_stack.append(self) + return self + + # pyre-fixme[2]: Parameter must be annotated. + def __exit__(self, exc_type, exc_value, exc_traceback) -> None: + # pop this mesh from mesh env + _mesh_resources.mesh_stack.pop() + + def __repr__(self) -> str: + device_mesh_repr = ( + f"DeviceMesh('{self.device_type}', {self.mesh.tolist()})" + if not self.mesh_dim_names + else f"DeviceMesh('{self.device_type}', {self.mesh.tolist()}, mesh_dim_names={self.mesh_dim_names})" + ) + return device_mesh_repr + + def __hash__(self): + # lazily compute hash + self._hash = getattr(self, "_hash", None) + if not self._hash: + self._hash = hash( + ( + self._flatten_mesh_list, + self.mesh.shape, + self.device_type, + self.mesh_dim_names, + self._thread_id, + ) + ) + return self._hash + + def __eq__(self, other: object) -> bool: + if not isinstance(other, DeviceMesh): + return False + if id(self) == id(other): + return True + else: + return ( + self._flatten_mesh_list == other._flatten_mesh_list + and self.mesh.shape == other.mesh.shape + and self.device_type == other.device_type + and self.mesh_dim_names == other.mesh_dim_names + and self._thread_id == other._thread_id + ) + + def __getitem__( + self, mesh_dim_names: Union[str, Tuple[str, ...]] + ) -> "DeviceMesh": + """ + Slice the current DeviceMesh based on the mesh_dim_names given to create a submesh. + The submesh created consists of the dimensions and the communicators indicated by + ``mesh_dim_names`` + + Args: + mesh_dim_names (Union[str, Tuple[str]]): the name or the tuple of names of the + mesh dimension of the DeviceMesh to create the submesh for. + Returns: + A :class:`DeviceMesh` object + + The following program runs on each process/rank in an SPMD manner in a world size of 8. + In the first example: + Calling mesh_2d["tp"] on rank 0, 1, 2, 3 returns a 1D submesh of DeviceMesh:([0, 1, 2, 3]). + Calling mesh_2d["tp"] on rank 4, 5, 6, 7 returns a 1D submesh of DeviceMesh:([4, 5, 6, 7]). + Calling mesh_2d["dp"] on rank 0, 4 returns a 1D submesh of DeviceMesh:([0, 4]). + Calling mesh_2d["dp"] on rank 1, 5 returns a 1D submesh of DeviceMesh:([1, 5]). + Calling mesh_2d["dp"] on rank 2, 6 returns a 1D submesh of DeviceMesh:([2, 6]). + Calling mesh_2d["dp"] on rank 3, 7 returns a 1D submesh of DeviceMesh:([3, 7]). + + In the second example: + Calling mesh_3d["dp", "cp"] on rank 0, 1, 4, 5 returns a 2D submesh of DeviceMesh:([[0, 1], [4, 5]]). + Calling mesh_3d["dp", "cp"] on rank 2, 3, 6, 7 returns a 2D submesh of DeviceMesh:([[2, 3], [6, 7]]). + Calling mesh_3d["cp", "dp"] on rank 0, 1, 4, 5 returns a 2D submesh of DeviceMesh:([[0, 4], [1, 5]]). + Calling mesh_3d["cp", "dp"] on rank 2, 3, 6, 7 returns a 2D submesh of DeviceMesh:([[2, 6], [3, 7]]). + + Example:: + >>> # xdoctest: +SKIP("no rank") + >>> from torch.distributed.device_mesh import DeviceMesh + >>> + >>> # Initialize a 2D device mesh as (2, 4) to represent the topology + >>> # of cross-host(dim 0), and within-host (dim 1). + >>> mesh_2d = init_device_mesh(device_type="cuda", (2,4), mesh_dim_names=("dp", "tp")) + >>> tp_mesh = mesh_2d["tp"] + >>> dp_mesh = mesh_2d["dp"] + >>> + >>> # Initialize a 3D mesh. + >>> mesh_3d = init_device_mesh(device_type="cuda", (2,2,2), mesh_dim_names=("dp", "pp", "cp")) + >>> # The order of the mesh_dim_names provided deteremines the order of dimensions in the submesh. + >>> dp_cp_mesh = mesh_3d["dp", "cp"] + >>> cp_dp_mesh = mesh_3d["cp", "dp"] + """ + if not self.mesh_dim_names: + raise RuntimeError("Cannot slice a DeviceMesh without mesh_dim_names!") + + mesh_dim_names = ( + (mesh_dim_names,) if isinstance(mesh_dim_names, str) else mesh_dim_names + ) + + if mesh_dim_names == self.mesh_dim_names: + return self + else: + slice_mesh_dims = _mesh_resources._get_slice_mesh_dims( + self, mesh_dim_names + ) + # When using FakeTensorMode to trace the model, `create_sub_mesh()` will + # fail as it will require a real tensor to manipulate. + # `unset_fake_temporarily()` will allow us to materialize the tensors + # within `_mesh_resources`, which should not affect modling. + # + # Note that this should be orthogonal to torch.compile(). But whether + # we can compile device_mesh `slicing` (no graph break) is not verified + # yet and need a follow-up, + # TODO: compiler + device_mesh slicing. + with torch._subclasses.fake_tensor.unset_fake_temporarily(): + submesh = _mesh_resources.create_sub_mesh( + self, mesh_dim_names, slice_mesh_dims + ) + return submesh + + def get_group(self, mesh_dim: Optional[Union[int, str]] = None) -> ProcessGroup: + """ + Returns the single ProcessGroup specified by mesh_dim, or, if mesh_dim is not specified and the + DeviceMesh is 1-dimensional, returns the only ProcessGroup in the mesh. + + Args: + mesh_dim (str/int, optional): it can be the name of the mesh dimension or the index + of the mesh dimension. Default is None. + + Returns: + A :class:`ProcessGroup` object. + """ + if not hasattr(self, "_dim_group_infos"): + raise RuntimeError("DeviceMesh process groups not initialized!") + + if self.mesh.ndim > 1 and mesh_dim is None: + raise RuntimeError( + f"Found the DeviceMesh have {self.mesh.ndim} dimensions", + "Optional kwarg `mesh_dim` needs to be specified when device_mesh.ndim > 1.", + "If you want to get the list of all the ProcessGroups in the DeviceMesh," + "please use `get_all_groups()` instead.", + ) + + # Quick return if the current device_mesh is a 1D mesh. + if self.mesh.ndim == 1 and mesh_dim is None: + return not_none( + _find_pg_by_ranks_and_tag(*self._dim_group_infos[0][:2]) # type: ignore[index] + ) + + root_mesh = _mesh_resources.get_root_mesh(self) + root_to_flatten_mapping = _mesh_resources.root_to_flatten_mapping.get( + root_mesh, None + ) + if root_to_flatten_mapping and mesh_dim in root_to_flatten_mapping.keys(): + dim_group_infos = root_to_flatten_mapping[mesh_dim]._dim_group_infos[0][:2] # type: ignore[index] + return not_none(_find_pg_by_ranks_and_tag(*dim_group_infos)) + else: + mesh_dim = ( + _mesh_resources.get_mesh_dim_by_name(self, mesh_dim) + if isinstance(mesh_dim, str) + else mesh_dim + ) + return not_none( + _find_pg_by_ranks_and_tag(*self._dim_group_infos[mesh_dim][:2]) # type: ignore[index] + ) + + def get_all_groups(self) -> List[ProcessGroup]: + """ + Returns a list of ProcessGroups for all mesh dimensions. + + Returns: + A list of :class:`ProcessGroup` object. + """ + return [self.get_group(i) for i in range(self.mesh.ndim)] + + @staticmethod + def from_group( + group: Union[ProcessGroup, List[ProcessGroup]], + device_type: str, + mesh: Optional[Union[torch.Tensor, "ArrayLike"]] = None, + *, + mesh_dim_names: Optional[Tuple[str, ...]] = None, + ) -> "DeviceMesh": + """ + Constructs a :class:`DeviceMesh` with ``device_type`` from an + existing :class:`ProcessGroup`. + + The constructed device mesh has number of dimensions equal to the + number of groups passed. If more than one group is passed, then the + ``mesh`` argument is required. + """ + if isinstance(group, ProcessGroup): + group_ranks = get_process_group_ranks(group) + if ( + isinstance(mesh, torch.Tensor) and mesh.tolist() != group_ranks + ) or ( + mesh is not None + and not isinstance(mesh, torch.Tensor) + and mesh != group_ranks + ): + raise ValueError( + f"Invalid mesh {str(mesh)} for ProcessGroup with ranks {group_ranks}" + ) + mesh = torch.tensor(group_ranks, device="cpu", dtype=torch.int) + device_mesh = DeviceMesh( + device_type, + mesh, + mesh_dim_names=mesh_dim_names, + _init_backend=False, + ) + device_mesh._dim_group_infos = [ + (_get_group_tag(group), group_ranks, group.group_name) + ] + return device_mesh + groups = list(group) + if len(groups) == 0: + raise ValueError("Expects at least one ProcessGroup to be passed") + if mesh is None: + raise ValueError("Must pass mesh if passing multiple ProcessGroups") + mesh = ( + mesh.detach().to(dtype=torch.int, device="cpu") + if isinstance(mesh, torch.Tensor) + else torch.tensor(mesh, device="cpu", dtype=torch.int) + ) + if mesh.ndim != len(groups): + raise ValueError( + "Expects mesh with ndim equal to number of ProcessGroups but got " + f"mesh {mesh.tolist()} and {len(groups)} ProcessGroups" + ) + device_mesh = DeviceMesh( + device_type, mesh, mesh_dim_names=mesh_dim_names, _init_backend=False + ) + device_mesh._dim_group_infos = [ + ( + _get_group_tag(group), + get_process_group_ranks(group), + group.group_name, + ) + for group in groups + ] + return device_mesh + + def size(self, mesh_dim: Optional[int] = None) -> int: + return self.mesh.numel() if mesh_dim is None else self.mesh.size(mesh_dim) + + @property + def ndim(self) -> int: + return self.mesh.ndim + + @property + def shape(self) -> Tuple[int, ...]: + return tuple(self.mesh.shape) + + def get_rank(self) -> int: + """ + Returns the current global rank. + """ + return get_rank() + + def get_local_rank(self, mesh_dim: Optional[Union[int, str]] = None) -> int: + """ + Returns the local rank of the given mesh_dim of the DeviceMesh. + + Args: + mesh_dim (str/int, optional): it can be the name of the mesh dimension or the index + of the mesh dimension. Default is None. + + Returns: + An integer denotes the local rank. + + The following program runs on each process/rank in an SPMD manner. In this example, we have 2 + hosts with 4 GPUs each. + Calling mesh_2d.get_local_rank(mesh_dim=0) on rank 0, 1, 2, 3 would return 0. + Calling mesh_2d.get_local_rank(mesh_dim=0) on rank 4, 5, 6, 7 would return 1. + Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 0, 4 would return 0. + Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 1, 5 would return 1. + Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 2, 6 would return 2. + Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 3, 7 would return 3. + + Example:: + >>> # xdoctest: +SKIP("no rank") + >>> from torch.distributed.device_mesh import DeviceMesh + >>> + >>> # Initialize device mesh as (2, 4) to represent the topology + >>> # of cross-host(dim 0), and within-host (dim 1). + >>> mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]]) + """ + if self.ndim > 1 and mesh_dim is None: + raise RuntimeError( + f"Found the DeviceMesh have {self.mesh.ndim} dimensions", + "Optional kwarg `mesh_dim` needs to be specified when device_mesh.ndim > 1.", + ) + elif mesh_dim is None: + mesh_dim = 0 + + mesh_dim_group = not_none(self.get_group(mesh_dim)) + assert isinstance( + mesh_dim_group, ProcessGroup + ), "We expect ProcessGroup before calling `get_rank`!" + return not_none(get_rank(mesh_dim_group)) + + def get_coordinate(self) -> Optional[List[int]]: + """ + Return the relative indices of this rank relative to all + dimensions of the mesh. If this rank is not part of the mesh, return None. + """ + return self._coordinate_on_dim if self._coordinate_on_dim else None + + def _flatten(self, mesh_dim_name: Optional[str] = None) -> "DeviceMesh": + """ + Returns a 1D DeviceMesh by flattening the current DeviceMesh. + + If no mesh_dim_name is provided, the default is a string concatentaing the mesh_dim_names of the + given submesh with each mesh_dim_name separated by "_". For example, if we have a 3D mesh + DeviceMesh([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], mesh_dim_names=("dp", "cp", "tp")), calling + mesh_3d["dp", "cp"]._flatten() will create a 1D submesh DeviceMesh([0, 1, 2, 3], mesh_dim_names=("dp_cp",)) + on rank 0, 1, 2, 3 and a 1D submesh DeviceMesh([4, 5, 6, 7], mesh_dim_names=("dp_cp",)) on rank 4, 5, 6, 7. + + After the flattened dimension is created, to access the flattened dimesnion in mesh_3d, one can use the + existing slicing method to obtain the flattened mesh through calling mesh_3d["dp_cp"]. + """ + if not self.mesh_dim_names: + raise RuntimeError( + "Cannot flatten a DeviceMesh without mesh_dim_names!" + ) + + return _mesh_resources.create_flatten_mesh(self, mesh_dim_name) + + def init_device_mesh( + device_type: str, + mesh_shape: Tuple[int, ...], + *, + mesh_dim_names: Optional[Tuple[str, ...]] = None, + ) -> DeviceMesh: + """ + Initializes a `DeviceMesh` based on `device_type`, `mesh_shape`, and `mesh_dim_names` parameters. + + This creates a DeviceMesh with an n-dimensional array layout, where `n` is the length of `mesh_shape`. + If `mesh_dim_names` is provided, each dimension is labeled as `mesh_dim_names[i]`. + + .. note:: + `init_device_mesh` follows SPMD programming model, meaning the same PyTorch Python program + runs on all processes/ranks in the cluster. Ensure `mesh_shape` (the dimensions of the nD array + describing device layout) is identical across all ranks. Inconsistent `mesh_shape` may lead to hanging. + + .. note:: + If no process group is found, init_device_mesh will initialize distributed process group/groups + required for distributed communications behind the scene. + + Args: + device_type (str): The device type of the mesh. Currently supports: "cpu", "cuda/cuda-like". + Passing in a device type with a GPU index, such as "cuda:0", is not allowed. + mesh_shape (Tuple[int]): A tuple defining the dimensions of the multi-dimensional array + describing the layout of devices. + mesh_dim_names (Tuple[str], optional): A tuple of mesh dimension names to assign to each dimension + of the multi-dimensional array describing the layout of devices. Its length must match the length + of `mesh_shape`. Each string in `mesh_dim_names` must be unique. + + Returns: + DeviceMesh: A :class:`DeviceMesh` object representing the device layout. + + Example:: + >>> # xdoctest: +SKIP("no rank") + >>> from torch.distributed.device_mesh import init_device_mesh + >>> + >>> mesh_1d = init_device_mesh("cuda", mesh_shape=(8,)) + >>> mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8), mesh_dim_names=("dp", "tp")) + + """ + if mesh_dim_names is not None: + if len(set(mesh_dim_names)) != len(mesh_dim_names): + raise RuntimeError( + "Each mesh_dim_name must be unique.", + f"Found repeated mesh_dim_name in mesh_dim_names {mesh_dim_names}", + ) + + if len(mesh_shape) != len(mesh_dim_names): + raise RuntimeError( + "mesh_shape and mesh_dim_names should have same length!", + f"Found len(mesh_dim_names): {len(mesh_dim_names)} and len(mesh_shape):{len(mesh_shape)}.", + ) + + # assume valid device types are all letters + if device_type and not device_type.isalpha(): + raise RuntimeError( + f"Device type with index is not supported but got {device_type}. ", + "If you maintained a 'torch.device' object, it's recommended to pass in 'device.type'.", + ) + + # Always initialize the mesh's tensor on CPU, regardless of what the + # external device type has been set to be (e.g. meta) + with torch.device("cpu"): + mesh = torch.arange(math.prod(mesh_shape), dtype=torch.int).view(mesh_shape) + device_mesh = DeviceMesh( + device_type=device_type, + mesh=mesh, + mesh_dim_names=mesh_dim_names, + ) + + return device_mesh diff --git a/lib/python3.12/site-packages/torch/distributed/distributed_c10d.py b/lib/python3.12/site-packages/torch/distributed/distributed_c10d.py new file mode 100644 index 0000000000000000000000000000000000000000..a60b12ad5df40c9dbdb89fbd3e983ba81855c513 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/distributed_c10d.py @@ -0,0 +1,5378 @@ +# mypy: allow-untyped-defs +"""Distributed Collective Communication (c10d).""" + +import collections.abc +import contextlib +import ctypes +import hashlib +import io +import itertools +import logging +import os +import pickle +import sys +import time +import warnings +from collections import namedtuple +from datetime import timedelta +from typing import Any, Callable, Dict, List, Optional, Tuple, TYPE_CHECKING, Union +from typing_extensions import deprecated + +import torch +from torch._C import _DistStoreError as DistStoreError +from torch._C._distributed_c10d import ( + _DistributedBackendOptions, + _register_process_group, + _resolve_process_group, + _unregister_all_process_groups, + _unregister_process_group, + AllgatherOptions, + AllreduceCoalescedOptions, + AllreduceOptions, + AllToAllOptions, + BarrierOptions, + BroadcastOptions, + DebugLevel, + GatherOptions, + get_debug_level, + PrefixStore, + ProcessGroup, + ReduceOp, + ReduceOptions, + ReduceScatterOptions, + ScatterOptions, + Store, + Work, +) +from torch._utils_internal import set_pytorch_distributed_envs_from_justknobs +from torch.monitor import _WaitCounter +from torch.utils._typing_utils import not_none + +from .c10d_logger import _exception_logger, _time_logger +from .constants import default_pg_nccl_timeout, default_pg_timeout +from .rendezvous import register_rendezvous_handler, rendezvous # noqa: F401 + + +__all__ = [ + "Backend", + "BackendConfig", + "GroupMember", + "P2POp", + "all_gather", + "all_gather_coalesced", + "all_gather_object", + "all_reduce", + "all_reduce_coalesced", + "all_to_all", + "all_to_all_single", + "barrier", + "batch_isend_irecv", + "broadcast", + "send_object_list", + "recv_object_list", + "broadcast_object_list", + "destroy_process_group", + "gather", + "gather_object", + "get_backend_config", + "get_backend", + "get_default_backend_for_device", + "get_rank", + "get_world_size", + "get_pg_count", + "group", + "init_process_group", + "irecv", + "is_gloo_available", + "is_initialized", + "is_mpi_available", + "is_backend_available", + "is_nccl_available", + "is_torchelastic_launched", + "is_ucc_available", + "is_xccl_available", + "isend", + "monitored_barrier", + "new_group", + "new_subgroups", + "new_subgroups_by_enumeration", + "recv", + "reduce", + "reduce_scatter", + "scatter", + "scatter_object_list", + "send", + "supports_complex", + "AllreduceCoalescedOptions", + "AllreduceOptions", + "AllToAllOptions", + "BarrierOptions", + "BroadcastOptions", + "GatherOptions", + "PrefixStore", + "ProcessGroup", + "ReduceOp", + "ReduceOptions", + "ReduceScatterOptions", + "ScatterOptions", + "Store", + "DebugLevel", + "get_debug_level", + "Work", + "default_pg_timeout", + "get_group_rank", + "get_global_rank", + "get_process_group_ranks", + "reduce_op", + "all_gather_into_tensor", + "reduce_scatter_tensor", + "get_node_local_rank", + "split_group", +] + +_MPI_AVAILABLE = True +_NCCL_AVAILABLE = True +_GLOO_AVAILABLE = True +_UCC_AVAILABLE = True +_XCCL_AVAILABLE = True + +_pickler = pickle.Pickler +_unpickler = pickle.Unpickler + + +# Change __module__ of all imported types from torch._C._distributed_c10d that are public +def _export_c_types() -> None: + _public_types_to_change_module = [ + AllreduceCoalescedOptions, + AllreduceOptions, + AllToAllOptions, + BarrierOptions, + BroadcastOptions, + GatherOptions, + PrefixStore, + ProcessGroup, + ReduceOp, + ReduceOptions, + ReduceScatterOptions, + ScatterOptions, + Store, + DebugLevel, + get_debug_level, + Work, + ] + for type in _public_types_to_change_module: + type.__module__ = "torch.distributed.distributed_c10d" + + +_export_c_types() + +try: + from torch._C._distributed_c10d import ProcessGroupMPI + + ProcessGroupMPI.__module__ = "torch.distributed.distributed_c10d" + __all__ += ["ProcessGroupMPI"] +except ImportError: + _MPI_AVAILABLE = False + +try: + from torch._C._distributed_c10d import ProcessGroupNCCL + + ProcessGroupNCCL.__module__ = "torch.distributed.distributed_c10d" + __all__ += ["ProcessGroupNCCL"] +except ImportError: + _NCCL_AVAILABLE = False + +try: + from torch._C._distributed_c10d import _ProcessGroupWrapper, ProcessGroupGloo + + ProcessGroupGloo.__module__ = "torch.distributed.distributed_c10d" + __all__ += ["ProcessGroupGloo"] +except ImportError: + _GLOO_AVAILABLE = False + +try: + from torch._C._distributed_c10d import ProcessGroupUCC + + ProcessGroupUCC.__module__ = "torch.distributed.distributed_c10d" + __all__ += ["ProcessGroupUCC"] +except ImportError: + _UCC_AVAILABLE = False + +try: + from torch._C._distributed_c10d import ProcessGroupXCCL + + ProcessGroupXCCL.__module__ = "torch.distributed.distributed_c10d" + __all__ += ["ProcessGroupXCCL"] +except ImportError: + _XCCL_AVAILABLE = False + +logger = logging.getLogger(__name__) + +PG_WRAPPER_STORE_PREFIX = "pg_wrapper" + + +# Some reduce ops are not supported by complex numbers and will result in an error. +# We currently provide complex support to the distributed API by viewing +# complex tensors as real (torch.view_as_real), meaning that calling +# these unsupported ops will return garbage values rather than error out. +# (e.g. max(2+3i, 3+2i) = 3+3i) +# We'd like calls to unsupported ops to error out accordingly, +# rather than returning garbage values. +def supports_complex(reduceOp: ReduceOp) -> bool: + """Return true if reduce ops is supported. False otherwise.""" + denyList = [ + ReduceOp.MAX, + ReduceOp.MIN, + ReduceOp.PRODUCT, + ReduceOp.BAND, + ReduceOp.BOR, + ReduceOp.BXOR, + ] + return reduceOp not in denyList + + +class Backend(str): + """ + An enum-like class for backends. + + Available backends: GLOO, NCCL, UCC, MPI, XCCL, and other registered backends. + + The values of this class are lowercase strings, e.g., ``"gloo"``. They can + be accessed as attributes, e.g., ``Backend.NCCL``. + + This class can be directly called to parse the string, e.g., + ``Backend(backend_str)`` will check if ``backend_str`` is valid, and + return the parsed lowercase string if so. It also accepts uppercase strings, + e.g., ``Backend("GLOO")`` returns ``"gloo"``. + + .. note:: The entry ``Backend.UNDEFINED`` is present but only used as + initial value of some fields. Users should neither use it directly + nor assume its existence. + """ + + UNDEFINED = "undefined" + GLOO = "gloo" + NCCL = "nccl" + UCC = "ucc" + MPI = "mpi" + XCCL = "xccl" + + _BackendPlugin = namedtuple("_BackendPlugin", ["creator_fn", "extended_api"]) + + _plugins: Dict[str, _BackendPlugin] = {} + + backend_list = [UNDEFINED, GLOO, NCCL, XCCL, UCC, MPI] + + # 3rd-party devices can register the default backend support here + default_device_backend_map: Dict[str, str] = { + "cpu": GLOO, + "cuda": NCCL, + "xpu": XCCL, + } + + backend_capability: Dict[str, List[str]] = { + GLOO: ["cpu", "cuda"], + NCCL: ["cuda"], + XCCL: ["xpu"], + UCC: ["cpu", "cuda"], + MPI: ["cpu", "cuda"], + } + + backend_type_map: Dict[str, ProcessGroup.BackendType] = { + UNDEFINED: ProcessGroup.BackendType.UNDEFINED, + GLOO: ProcessGroup.BackendType.GLOO, + NCCL: ProcessGroup.BackendType.NCCL, + XCCL: ProcessGroup.BackendType.XCCL, + UCC: ProcessGroup.BackendType.UCC, + MPI: ProcessGroup.BackendType.MPI, + } + + def __new__(cls, name: str): + """Create and return a new instance of the class.""" + if not isinstance(name, str): + raise ValueError("Backend constructor parameter must be string-ish") + value = getattr(Backend, name.upper(), Backend.UNDEFINED) + + if value == Backend.UNDEFINED: + value = name.lower() + return value + + @classmethod + def register_backend( + cls, + name, + func, + extended_api=False, + devices: Optional[Union[str, List[str]]] = None, + ) -> None: + """ + Register a new backend with the given name and instantiating function. + + This class method is used by 3rd party ``ProcessGroup`` extension to + register new backends. + + Args: + name (str): Backend name of the ``ProcessGroup`` extension. It + should match the one in ``init_process_group()``. + func (function): Function handler that instantiates the backend. + The function should be implemented in the backend + extension and takes four arguments, including + ``store``, ``rank``, ``world_size``, and ``timeout``. + extended_api (bool, optional): Whether the backend supports extended argument structure. + Default: ``False``. If set to ``True``, the backend + will get an instance of ``c10d::DistributedBackendOptions``, and + a process group options object as defined by the backend implementation. + device (str or list of str, optional): device type this backend + supports, e.g. "cpu", "cuda", etc. If `None`, + assuming both "cpu" and "cuda" + + .. note:: This support of 3rd party backend is experimental and subject to change. + + """ + # Allow UCC plugin if Pytorch is not built with native support. + # TODO: remove this exception once UCC plugin is fully deprecated. + if name != Backend.UCC or (name == Backend.UCC and is_ucc_available()): + assert not hasattr( + Backend, name.upper() + ), f"{name.upper()} c10d backend already exist" + assert ( + name.upper() not in Backend._plugins + ), f"{name.upper()} c10d backend creator function already exist" + + setattr(Backend, name.upper(), name.lower()) + Backend.backend_list.append(name.lower()) + if devices is not None: + for device in devices: + if device != "cpu" and device != "cuda": + Backend.default_device_backend_map[device] = name.lower() + Backend.backend_type_map[name.lower()] = ProcessGroup.BackendType.CUSTOM + + # Update device capability matrix in Backend class + if devices is None: + # This is more of a backward support for groups like `threaded`: + # assume default devices "cpu" and "cuda", but warn + warnings.warn( + f"Device capability of {name} unspecified, assuming `cpu` and " + "`cuda`. Please specify it via the `devices` argument of " + "`register_backend`." + ) + Backend.backend_capability[name.lower()] = ["cpu", "cuda"] + elif isinstance(devices, str): + # Single device string specified. Simply convert to list. + Backend.backend_capability[name.lower()] = [devices] + else: + Backend.backend_capability[name.lower()] = devices + + Backend._plugins[name.upper()] = Backend._BackendPlugin(func, extended_api) + + +class BackendConfig: + """Backend configuration class.""" + + def __init__(self, backend: Backend): + """Init.""" + self.device_backend_map: Dict[str, Backend] = {} + backend = str(backend) + + if backend == Backend.UNDEFINED: + # Detect the accelerator on the machine. If no accelerator is + # available, it returns CPU. + device_type = torch._C._get_accelerator().type + try: + backend_str = Backend.default_device_backend_map[device_type] + self.device_backend_map[device_type] = Backend(backend_str) + except KeyError: + raise ValueError( + f"We detected accelerator {device_type} on your machine. " + f"But we don't know which communication backend to use for this accelerator. " + f"Please specify the `backend` argument in the `init_process_group` call." + ) from None + elif backend.lower() in Backend.backend_list: + # Cases for when backend is a single string (without device types) + # e.g. "nccl", "gloo", "ucc", "mpi" + supported_devices = Backend.backend_capability[backend.lower()] + backend_val = Backend(backend) + self.device_backend_map = dict.fromkeys(supported_devices, backend_val) + elif ":" in backend.lower(): + # Backend specified in "device:backend" format + # make sure the backend string is in the correct format + # "{device_type1}:{backend1},{device_type2}:{backend2}" + # e.g. "cpu:gloo,cuda:nccl" + backend_str_error_message = f"""The custom backend string argument is invalid: {backend}. + Custom backend string is an experimental feature where the backend string must be in the format: + ":,:...". e.g. 'cpu:gloo,cuda:nccl'""" + + # parse the backend string and populate the device_backend_map + for device_backend_pair_str in backend.lower().split(","): + device_backend_pair = device_backend_pair_str.split(":") + if len(device_backend_pair) != 2: + raise ValueError( + f"Invalid device:backend pairing: \ + {device_backend_pair_str}. {backend_str_error_message}" + ) + device, backend = device_backend_pair + if device in self.device_backend_map: + raise ValueError( + f"Duplicate device type {device} \ + in backend string: {backend}. {backend_str_error_message}" + ) + self.device_backend_map[device] = Backend(backend) + else: + # User specified a single backend name whose device capability is + # unknown, assuming it can support the default devices of PyTorch + # (cpu and cuda) + warnings.warn( + f"Device capability of {backend} unknown, assuming `cpu` and " + "`cuda`. You can specify it in `device:backend` format in " + "`init_process_group` call." + ) + backend_val = Backend(backend) + self.device_backend_map = { + "cpu": backend_val, + "cuda": backend_val, + "xpu": backend_val, + } + + logger.info("Using backend config: %s", self.device_backend_map) + + def __repr__(self): + """Return all the device:backend pairs separated by commas.""" + return ",".join( + f"{device}:{backend}" for device, backend in self.device_backend_map.items() + ) + + def get_device_backend_map(self) -> Dict[str, Backend]: + """Return backend map of the device.""" + return self.device_backend_map + + +class _reduce_op: + r""" + Deprecated enum-like class. + + For reduction operations: ``SUM``, ``PRODUCT``, ``MIN``, and ``MAX``. + + :class:`~torch.distributed.ReduceOp` is recommended to use instead. + """ + + def __init__(self) -> None: + # __members__ is a dict storing key-value pairs for enum classes + for k, v in ReduceOp.RedOpType.__members__.items(): + setattr(self, k, v) + self.__members__ = ReduceOp.RedOpType.__members__ + + @deprecated( + "`torch.distributed.reduce_op` is deprecated, " + "please use `torch.distributed.ReduceOp` instead", + category=FutureWarning, + ) + def __getattribute__(self, key): + return object.__getattribute__(self, key) + + +reduce_op = _reduce_op() + + +class P2POp: + """ + A class to build point-to-point operations for ``batch_isend_irecv``. + + This class builds the type of P2P operation, communication buffer, peer rank, + Process Group, and tag. Instances of this class will be passed to + ``batch_isend_irecv`` for point-to-point communications. + + Args: + op (Callable): A function to send data to or receive data from a peer process. + The type of ``op`` is either ``torch.distributed.isend`` or + ``torch.distributed.irecv``. + tensor (Tensor): Tensor to send or receive. + peer (int, optional): Destination or source rank. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + tag (int, optional): Tag to match send with recv. + group_peer (int, optional): Destination or source rank. + """ + + def __init__( + self, + op: Callable, + tensor: torch.Tensor, + peer: Optional[int] = None, + group: Optional[ProcessGroup] = None, + tag: int = 0, + group_peer: Optional[int] = None, + ): + """Init.""" + self.op = op + self.tensor = tensor + self.group = _group_or_default_group(group) + self.peer = _canonicalize_group_rank( + self.group, peer, group_peer, return_global=True + ) + self.tag = tag + self.group_peer = _canonicalize_group_rank(self.group, peer, group_peer) + + def __new__( + cls, + op: Callable, + tensor: torch.Tensor, + peer: Optional[int] = None, + group: Optional[ProcessGroup] = None, + tag: int = 0, + group_peer: Optional[int] = None, + ): + """Create and return a new instance of the class.""" + _check_op(op) + _check_single_tensor(tensor, "tensor") + + return object.__new__(cls) + + def __repr__(self): + my_group_rank = get_rank(self.group) + op_name = self.op.__name__ + group_name = self.group.group_name if self.group else "default_pg" + if "send" in op_name: + s = my_group_rank + d = self.group_peer + elif "recv" in op_name: + s = self.group_peer + d = my_group_rank + else: + return super().__repr__() + + return f"P2POp({op_name} pg={group_name}, group_src={s}, group_dst={d}, {self.tensor.shape}, {self.tensor.dtype})" + + +class _CollOp: + """ + A class to capture collective operations. + + Args: + op (Callable): A collective function, e.g. ``torch.distributed.all_reduce``. + tensor (Tensor): Tensor to operate on. + dst_tensor (Tensor, optional): Provided when source and destinaton tensors are not the same. + redop (ReduceOp, optional): reduce operation. + root (int, optional): root of broadcast or reduce. + """ + + def __init__( + self, + op: Callable, + tensor: torch.Tensor, + dst_tensor: Optional[torch.Tensor] = None, + redop: Optional[ReduceOp] = None, + root: Optional[int] = None, + ): + self.op = op + self.tensor = tensor + self.dst_tensor = dst_tensor + self.redop = redop + self.root = root + + +# DO NOT USE THESE FIELDS DIRECTLY. +# Use them through the _world object to make sure the _world override mechanism +_pg_map: Dict[ProcessGroup, Tuple[str, Store]] = {} +_pg_names: Dict[ProcessGroup, str] = {} +_pg_group_ranks: Dict[ProcessGroup, Dict[int, int]] = {} +# For a pg, it is a map from ProcessGroup to BackendConfig +_pg_backend_config: Dict[ProcessGroup, str] = {} +_group_count = 0 +_tags_to_pg: Dict[str, List[ProcessGroup]] = {} +_pg_to_tag: Dict[ProcessGroup, str] = {} +_backend: Optional[str] = None + + +class _World: + """ + Container class for c10d process group state. + + This is used during registration and lookup of PG state. + + .. warning:: This is an experimental API intended to expose the inner workings + of c10d and is subject to change.. + """ + + def __init__(self) -> None: + self._default_pg = None + self._pg_coalesce_state: Dict[ProcessGroup, List[_CollOp]] = {} + + @property + def default_pg(self) -> Optional[ProcessGroup]: + """ + Process group that includes all ranks of the cluster. + + This default ProcessGroup is used by c10d APIs when a ProcessGroup is needed + but None is provided. + """ + return self._default_pg + + @default_pg.setter + def default_pg(self, value) -> None: + self._default_pg = value + + @property + def pg_map(self) -> Dict[ProcessGroup, Tuple[str, Store]]: + """ + Provide Mapping from ProcessGroup to backend name and store. + + For NCCL and GLOO pg, it is a map from ProcessGroup to (Backend, Store) + For MPI pg, it is a map from ProcessGroup to (Backend, None) + + TODO don't expose the map, expose fine grained ops + """ + global _pg_map + return _pg_map + + @property + def pg_names(self) -> Dict[ProcessGroup, str]: + """ + Process group's names, map from ProcessGroup to str. + + TODO don't expose the map, expose fine grained ops + """ + global _pg_names + return _pg_names + + @property + def pg_group_ranks(self) -> Dict[ProcessGroup, Dict[int, int]]: + """ + Process group's global rank to local rank mapping. + + TODO don't expose the map, expose fine grained ops + """ + global _pg_group_ranks + return _pg_group_ranks + + @property + def pg_backend_config(self) -> Dict[ProcessGroup, str]: + """ + Process group's backend config. + + TODO don't expose the map, expose fine grained ops + """ + global _pg_backend_config + return _pg_backend_config + + @property + def group_count(self) -> int: + """ + Process group count for default naming. + + TODO don't expose group_count, use something else instead + """ + global _group_count + return _group_count + + @group_count.setter + def group_count(self, value: int) -> None: + """Use to compute the name of ProcessGroups when using global synchronization.""" + global _group_count + _group_count = value + + @property + def tags_to_pg(self) -> Dict[str, List[ProcessGroup]]: + global _tags_to_pg + return _tags_to_pg + + @property + def pg_to_tag(self) -> Dict[ProcessGroup, str]: + global _pg_to_tag + return _pg_to_tag + + @property + def pg_coalesce_state(self) -> Dict[ProcessGroup, List[_CollOp]]: + return self._pg_coalesce_state + + @property + def pg_config_info(self) -> List[Dict[str, Any]]: + """ + Return a list of dict with process groups and backends. + + Along with their unique IDs and configurations (types and ranks). + """ + config_info: List[Dict[str, Any]] = [] + default_pg_size = _get_group_size(None) + for pg in self.pg_map.keys(): + ranks = self.pg_group_ranks[pg] + config_info.append( + { + "pg_name": self.pg_names[pg], + "pg_desc": pg.group_desc, + "backend_config": self.pg_backend_config[pg], + "ranks": ( + list(ranks.keys()) if len(ranks) != default_pg_size else [] + ), # 'ranks' is an empty list when all ranks are involved in a pg + "group_size": len(ranks), + "group_count": self.group_count, + } + ) + return config_info + + +_world = _World() +"""Holds the singleton instance of ``_World`` used by c10. Experimental extension point to override it""" + + +class _WorldMeta(type): + """ + Meta class of ``group`` and ``GroupMember``. + + Allows them to have the class property ``WORLD``. + """ + + # Points to the default PG once initialized. + @property + def WORLD(cls) -> Optional[ProcessGroup]: + return _world.default_pg + + @WORLD.setter + def WORLD(cls, pg: Optional[ProcessGroup]): + _world.default_pg = pg + + +class group(metaclass=_WorldMeta): + """Group class. Placeholder.""" + + +class GroupMember(metaclass=_WorldMeta): + """Group member class.""" + + NON_GROUP_MEMBER = -100 + + +def _get_default_timeout(backend: Backend) -> timedelta: + # see note on nccl vs other backend timeout (constants.py) + if backend == Backend.NCCL: + if not isinstance(default_pg_nccl_timeout, timedelta): + # TODO moco benchmark on CPU initializes pgnccl backend today, triggered this assert in CI before it was + # changed to be a warning. We should fix the moco model. + warnings.warn( + "Attempted to get default timeout for nccl backend, but NCCL support is not compiled" + ) + return default_pg_timeout + return default_pg_nccl_timeout + else: + return default_pg_timeout + + +def _check_valid_timeout(timeout: Any) -> None: + if not isinstance(timeout, timedelta): + raise TypeError( + f"Expected timeout argument to be of type datetime.timedelta, got {timeout}" + ) + + +# Default process group state +_default_pg_init_method: Optional[str] = None + +STORE_BASED_BARRIER_PREFIX = "store_based_barrier_key" + + +def _get_object_coll_device(group: Optional[ProcessGroup] = None) -> str: + """ + .. note:: This is an internal helper and does not have backward + compatibility, please use with caution. + + Return the device type to use with ``group`` for object collectives or + barrier. + + There are selection rules: + 1. If user specifies exactly one backend in ``init_process_group`` call: + use that backend + 2. Else if user specifies multiple "device:backend" pairs in init_process_group: + If "cpu" is among those pairs, use "cpu" (because the object is in cpu memory); + Otherwise, use the first backend (sort of a random pick). + + Args: + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + + Returns: + str: The device type to use for object collective with ``group``. + + """ + group = group or _get_default_group() + + if not isinstance(group, ProcessGroup): + warnings.warn( + f"You are using a Backend {type(group)} as a ProcessGroup. " + "This usage is deprecated since PyTorch 2.0. Please use a public API " + "of PyTorch Distributed instead.", + ) + # Provide backward compatibility to cases where `group` passed in is + # actually a Backend (like `ProcessGroupGloo`) rather than a + # `ProcessGroup` in PT 2.0 sense + if isinstance(group, ProcessGroupGloo): + # RPC uses Gloo for object collectives + return "cpu" + else: + raise ValueError(f"Expecting a ProcessGroup, but got a {type(group)}.") + + """ + ``group._device_types`` is a property pybind that returns the devices + ("cpu", "cuda", etc) supported by ``group``. Can be multiple if the + ``group`` supports multiple devices. + """ + devices = group._device_types + + if len(devices) == 1: + # User fixed exactly one backend in `init_process_group` + return devices[0].type + elif len(devices) == 0: + # No backend has been registered with this PG (maybe because no + # collective has been run?) We pick cpu as the default and hopefully + # this would lazily init Gloo or other available cpu backend. + return "cpu" + elif torch.device("cpu") in devices: + # There are multiple backends in this PG and cpu is among them. + # cpu is preferred as the object is in cpu memory. No need for device + # copy. + return "cpu" + else: + # No cpu in the backend list. Randomly pick the first backend + return devices[0].type + + +def _get_pg_default_device(group: Optional[ProcessGroup] = None) -> torch.device: + """ + .. note:: This method will be deprecated, it only stays for + backward-compatiblity reason. Alternatives: + + - If you need to find a device for object collectives, please use + `_get_object_coll_device(group)`. + + - If you need to query the device types supported by group, please use + `_device_capability(group)`. + + Return the device type registered with ``group``. + + For example, if `init_process_group("nccl", ...)` was called, the returned + value would be `torch.device("cuda")`. + + Errors out if no device has been registered. + + Args: + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + + Returns: + torch.device: The device type registered with ``group``. + """ + + warnings.warn( + "`_get_pg_default_device` will be deprecated, it only stays for " + "backward-compatiblity reason. If you need to find a device for object " + "collectives, please use `_get_object_coll_device`. If you need to query " + "the device types supported by group, please use " + "`_device_capability(group)`. " + ) + group = group or _get_default_group() + + if not isinstance(group, ProcessGroup): + # Provide backward compatibility to cases where `group` passed in is + # actually a Backend (like `ProcessGroupGloo`) rather than a + # `ProcessGroup` in PT 2.0 sense + warnings.warn( + f"You are using a Backend {type(group)} as a ProcessGroup. " + "This usage is deprecated since PyTorch 2.0. Please use a public API " + "of PyTorch Distributed instead.", + FutureWarning, + stacklevel=3, + ) + # Most users create Gloo with private API for object collectives + return torch.device("cpu") + + """ + ``group._device_types`` is a property pybind that returns the devices + ("cpu", "cuda", etc) supported by ``group``. Can be multiple if the + ``group`` supports multiple devices. + """ + devices = group._device_types + + if len(devices) == 1: + # User fixed exactly one backend in `init_process_group` + return devices[0] + elif len(devices) == 0: + raise RuntimeError( + "Default device not found, because no backend has been registered " + "with this ProcessGroup." + ) + else: + # There are multiple backends in this PG. + if torch.device("cpu") in devices: + rv = torch.device("cpu") + else: + rv = devices[0] + warnings.warn( + "Multiple backends are registered with this ProcessGroup. We cannot " + f"determine which one is the default. Returning {rv}. " + "Please consider using other APIs." + ) + return rv + + +def _device_capability(group: Optional[ProcessGroup] = None) -> List[str]: + """ + Return the device type(s) supported by ``group``. + + Args: + group (ProcessGroup, optional): The process group to query. If None, + the default process group will be used. + + Returns: + List[str]: A list of device types supported by ``group``. + """ + group = group or _get_default_group() + return [device.type for device in group._device_types] + + +@_time_logger +def _store_based_barrier( + rank, + store, + group_name, + rendezvous_count, + timeout, + logging_interval=timedelta(seconds=10), +) -> None: + """ + Store based barrier for synchronizing processes. + + Barrier based on store which is used for synchronizing processes after + ``init_process_group`` or ``new_group``. Intended to be used only with + those two methods and is not a generic alternative to ``barrier()``. + """ + store_key = f"{STORE_BASED_BARRIER_PREFIX}:{group_name}" + store.add(store_key, 1) + logger.debug("Added key: %s to store for rank: %s", store_key, rank) + + # Now wait for all workers to check in with the store. + world_size = rendezvous_count + worker_count = store.add(store_key, 0) + + last_worker_key = f"{store_key}:last_worker" + if worker_count == world_size: + store.set(last_worker_key, "1") + + # adjust the timeout to be at least 10secs + 1sec per thousand ranks to reduce the odds of timeout + # this value was empirically found while scale testing. + logging_interval = max(logging_interval, timedelta(seconds=10 + world_size / 1000)) + + start = time.time() + while True: + try: + # This will throw an exception after the logging_interval in which we print out + # the status of the group or time out officially, throwing runtime error + store.wait([last_worker_key], logging_interval) + break + except RuntimeError as e: + worker_count = store.add(store_key, 0) + # Print status periodically to keep track. + logger.debug( + "Waiting in store based barrier to initialize process group for %s seconds" + "rank: %s, key: %s (world_size=%s, num_workers_joined=%s, timeout=%s error=%s)", + time.time() - start, + rank, + store_key, + world_size, + worker_count, + timeout, + e, + ) + + if timedelta(seconds=(time.time() - start)) > timeout: + raise DistStoreError( # noqa: B904 + "Timed out initializing process group in store based barrier on " + f"rank {rank}, for key: {store_key} (world_size={world_size}, " + f"num_workers_joined={worker_count}, timeout={timeout} error={e})" + ) + + logger.info( + "Rank %s: Completed store-based barrier for key:%s with %s nodes.", + rank, + store_key, + world_size, + ) + + +def _rank_not_in_group(group: Optional[ProcessGroup]) -> bool: + """Check if the current process's rank is not in a given group.""" + if group is None: + return False + return group == GroupMember.NON_GROUP_MEMBER + + +def _warn_not_in_group(op_name) -> None: + global_rank = -1 if GroupMember.WORLD is None else GroupMember.WORLD.rank() + warnings.warn( + f"Running {op_name} on global rank {global_rank} which does not " + "belong to the given group." + ) + + +def get_group_rank(group: ProcessGroup, global_rank: int) -> int: + """ + Translate a global rank into a group rank. + + ``global_rank`` must be part of ``group`` otherwise this raises RuntimeError. + + Args: + group (ProcessGroup): ProcessGroup to find the relative rank. + global_rank (int): Global rank to query. + + Returns: + Group rank of ``global_rank`` relative to ``group`` + + N.B. calling this function on the default process group returns identity + """ + if group is GroupMember.WORLD: + return global_rank + if group not in _world.pg_group_ranks: + raise ValueError( + f"Group {group} is not registered, please create group with torch.distributed.new_group API" + ) + group_ranks = _world.pg_group_ranks[group] + if global_rank not in group_ranks: + raise ValueError(f"Global rank {global_rank} is not part of group {group}") + + return group_ranks[global_rank] + + +def get_global_rank(group: ProcessGroup, group_rank: int) -> int: + """ + Translate a group rank into a global rank. + + ``group_rank`` must be part of `group` otherwise this raises RuntimeError. + + Args: + group (ProcessGroup): ProcessGroup to find the global rank from. + group_rank (int): Group rank to query. + + Returns: + Global rank of ``group_rank`` relative to ``group`` + + N.B. calling this function on the default process group returns identity + """ + if group is GroupMember.WORLD: + return group_rank + if group not in _world.pg_group_ranks: + raise ValueError( + f"Group {group} is not registered, please create group with torch.distributed.new_group API" + ) + for rank, grp_rank in _world.pg_group_ranks[group].items(): + if grp_rank == group_rank: + return rank + raise ValueError(f"Group rank {group_rank} is not part of group {group}") + + +# TODO: remove this once the ecosystem moves away from it. +@deprecated( + "`torch.distributed.distributed_c10d._get_global_rank` is deprecated, " + "please use `torch.distributed.distributed_c10d.get_global_rank` instead", + category=FutureWarning, +) +def _get_global_rank(group, rank) -> int: + """Use get_global_rank as this method is deprecated.""" + return get_global_rank(group, rank) + + +def get_process_group_ranks(group: ProcessGroup) -> List[int]: + """ + Get all ranks associated with ``group``. + + Args: + group (ProcessGroup): ProcessGroup to get all ranks from. + + Returns: + List of global ranks ordered by group rank. + """ + return list(_world.pg_group_ranks[group].keys()) + + +def _get_group_size(group) -> int: + """Get a given group's world size.""" + if group is GroupMember.WORLD or group is None: + default_pg = _get_default_group() + return default_pg.size() + return group.size() + + +def _get_group_size_by_name(group_name: str) -> int: + group = _resolve_process_group(group_name) + return group.size() + + +def _resolve_group_name_by_ranks_and_tag(ranks: List[int], tag: str) -> str: + # TODO(yifu): remove this function once ranks + tag is not a supported + # identifier for process group for functional collectives. + group = _find_pg_by_ranks_and_tag(tag, ranks) + if group is None: + raise ValueError("") + return group.group_name + + +def _check_single_tensor(param, param_name) -> None: + """Check that the parameter ``param_name`` is a single tensor.""" + if not isinstance(param, torch.Tensor): + raise TypeError( + f"""Invalid function argument. Expected parameter `{param_name}` of type torch.Tensor + but got {type(param)} instead.""" + ) + + +def _check_tensor_list(param, param_name) -> None: + """Check that the parameter ``param_name`` is a list of tensors.""" + if not isinstance(param, list): + raise TypeError( + f"""Invalid function argument. Expected parameter `{param_name}` of type List[torch.Tensor] + but got {type(param)} instead.""" + ) + elif not all(isinstance(p, torch.Tensor) for p in param): + raise TypeError( + f"""Invalid function argument. Expected parameter `{param_name}` of type List[torch.Tensor] + but got {type(param)} with elements of type {[type(p) for p in param]}.""" + ) + + +def _group_or_default_group(group: Optional[ProcessGroup] = None) -> ProcessGroup: + if group is None or group is GroupMember.WORLD: + group = _get_default_group() + return group + + +def _canonicalize_group_rank( + group: ProcessGroup, + global_rank: Optional[int] = None, + group_rank: Optional[int] = None, + return_global: bool = False, +) -> int: + """ + Helper method to take _either_ a global rank or a group rank and produce a group rank. + + If 'return_global' is true, produce a global rank instead of a group rank. + """ + + if group_rank is not None: + if global_rank is not None: + raise ValueError("Can't specify both group_rank and global_rank") + global_rank = get_global_rank(group, group_rank) + else: + if global_rank is None: + raise ValueError("Must specify global_rank or group_rank") + group_rank = get_group_rank(group, global_rank) + return global_rank if return_global else group_rank + + +def _check_not_self_rank(group: ProcessGroup, rank: int, rank_type: str): + if group.rank() == rank: + raise ValueError( + f"Invalid {rank_type} rank: {rank_type} rank should not be the same as " + "the rank of the current process." + ) + + +def _as_iterable(obj) -> collections.abc.Iterable: + return obj if isinstance(obj, list) else (obj,) + + +def _ensure_all_tensors_same_dtype(*tensors) -> None: + last_dtype = None + for tensor in itertools.chain.from_iterable(map(_as_iterable, tensors)): + tensor_dtype = tensor.dtype + # Mixing complex and its element type is allowed + if tensor_dtype.is_complex: + tensor_dtype = ( + torch.float32 if tensor_dtype == torch.complex64 else torch.complex128 + ) + + if last_dtype is None: + last_dtype = tensor_dtype + else: + if last_dtype != tensor_dtype: + raise ValueError( + "Invalid usage of tensors with different dtypes" + f"Found {last_dtype} and {tensor.dtype}" + ) + + +def _check_op(op) -> None: + """Check that the ``op`` is either isend or irecv.""" + if op not in [isend, irecv]: + raise ValueError( + "Invalid ``op``. Expected ``op`` " + "to be of type ``torch.distributed.isend`` or " + "``torch.distributed.irecv``." + ) + + +def _check_p2p_op_list(p2p_op_list) -> None: + """ + Check that the ``p2p_op_list`` is a list of P2POp instances. + + Also, check that all ops use the same group. + """ + if not isinstance(p2p_op_list, list) or not all( + isinstance(p2p_op, P2POp) for p2p_op in p2p_op_list + ): + raise ValueError( + "Invalid ``p2p_op_list``. Each op is expected to " + "to be of type ``torch.distributed.P2POp``." + ) + + group = p2p_op_list[0].group + if not all(group == p2p_op.group for p2p_op in p2p_op_list): + raise ValueError("All ops need to use the same group.") + + +def is_mpi_available() -> bool: + """Check if the MPI backend is available.""" + return _MPI_AVAILABLE + + +def is_nccl_available() -> bool: + """Check if the NCCL backend is available.""" + return _NCCL_AVAILABLE + + +def is_gloo_available() -> bool: + """Check if the Gloo backend is available.""" + return _GLOO_AVAILABLE + + +def is_ucc_available() -> bool: + """Check if the UCC backend is available.""" + return _UCC_AVAILABLE + + +def is_xccl_available() -> bool: + """Check if the XCCL backend is available.""" + return _XCCL_AVAILABLE + + +def is_backend_available(backend: str) -> bool: + """ + Check backend availability. + + Checks if the given backend is available and supports the built-in backends or + third-party backends through function ``Backend.register_backend``. + + Args: + backend (str): Backend name. + Returns: + bool: Returns true if the backend is available otherwise false. + """ + # If the backend has an ``is_backend_available`` function, return the result of that function directly + available_func = getattr(torch.distributed, f"is_{backend.lower()}_available", None) + if available_func: + return available_func() + + return backend.lower() in Backend.backend_list + + +def is_initialized() -> bool: + """Check if the default process group has been initialized.""" + return GroupMember.WORLD is not None + + +def is_torchelastic_launched() -> bool: + """ + Check whether this process was launched with ``torch.distributed.elastic`` (aka torchelastic). + + The existence of ``TORCHELASTIC_RUN_ID`` environment + variable is used as a proxy to determine whether the current process + was launched with torchelastic. This is a reasonable proxy since + ``TORCHELASTIC_RUN_ID`` maps to the rendezvous id which is always a + non-null value indicating the job id for peer discovery purposes.. + """ + return os.getenv("TORCHELASTIC_RUN_ID") is not None + + +def _is_barrier_after_init() -> int: + # Environment variable to control whether process group should perform a + # barrier after its init. Default value is 0, i.e. no barrier. If you + # experience issue with this setting, you may set + # `TORCH_DIST_INIT_BARRIER=1` to add the barrier. + return int(os.getenv("TORCH_DIST_INIT_BARRIER", "0")) + + +def _get_default_group() -> ProcessGroup: + """Get the default process group created by init_process_group.""" + if not is_initialized(): + raise ValueError( + "Default process group has not been initialized, " + "please make sure to call init_process_group." + ) + if TYPE_CHECKING: + return not_none(GroupMember.WORLD) + else: + return GroupMember.WORLD + + +def _get_default_store() -> Store: + """Get the default store created by init_process_group.""" + if not is_initialized(): + raise ValueError( + "Default process group has not been initialized, " + "please make sure to call init_process_group." + ) + default_pg = _get_default_group() + _, default_store = _world.pg_map[default_pg] + return default_store + + +def _update_default_pg(pg) -> None: + _world.default_pg = pg + rank = pg.rank() if pg is not None and pg != GroupMember.NON_GROUP_MEMBER else -1 + torch._C._distributed_c10d._set_global_rank(rank) + + +def get_backend_config(group: Optional[ProcessGroup] = None) -> str: + """ + Return the backend configuration of the given process group. + + Args: + group (ProcessGroup, optional): The process group to work on. The + default is the general main process group. If another specific group + is specified, the calling process must be part of :attr:`group`. + + Returns: + The backend configuration of the given process group as a lower case string. + + """ + pg = group or _get_default_group() + if _rank_not_in_group(pg): + raise ValueError("Invalid process group specified") + backend_config = _world.pg_backend_config.get(pg) + return str(not_none(backend_config)) + + +def get_backend(group: Optional[ProcessGroup] = None) -> Backend: + """ + Return the backend of the given process group. + + Args: + group (ProcessGroup, optional): The process group to work on. The + default is the general main process group. If another specific group + is specified, the calling process must be part of :attr:`group`. + + Returns: + The backend of the given process group as a lower case string. + + """ + pg = group or _get_default_group() + if _rank_not_in_group(pg): + raise ValueError("Invalid process group specified") + pg_store = _world.pg_map[pg] if pg in _world.pg_map else None + return Backend(not_none(pg_store)[0]) + + +def get_default_backend_for_device(device: Union[str, torch.device]) -> str: + """ + Return the default backend for the given device. + + Args: + Union[str, torch.device]: The device to get the default backend for. + + Returns: + The default backend for the given device as a lower case string. + + """ + if isinstance(device, torch.device): + device_str = device.type + else: + device_str = device.split(":")[0] + + backend = Backend.default_device_backend_map.get(device_str) + if backend is None: + raise ValueError(f"Default backend not registered for device : {device}") + + return backend + + +def _get_process_group_uid(pg: ProcessGroup) -> int: + backend = None + try: + backend = pg._get_backend(torch.device("cuda")) + except RuntimeError: + pass + if is_nccl_available() and isinstance(backend, ProcessGroupNCCL): + return backend.uid + return -1 + + +def _get_pg_config(group: Optional[ProcessGroup] = None) -> Dict[str, Any]: + """ + Return the pg configuration of the given process group. + + """ + pg = group or _get_default_group() + return { + "pg_name": _get_process_group_name(pg), + "pg_desc": pg.group_desc, + "backend_config": get_backend_config(pg), + "pg_size": _get_group_size(pg), + "ranks": get_process_group_ranks(pg), + } + + +def _get_all_pg_configs() -> List[Dict[str, Any]]: + """ + Return the pg configuration of all the process groups. + + """ + config_info: List[Dict[str, Any]] = [ + _get_pg_config(pg) for pg in _world.pg_map.keys() + ] + return config_info + + +def get_pg_count() -> int: + """ + Return the number of process groups. + + """ + return _world.group_count + + +def get_node_local_rank(fallback_rank: Optional[int] = None) -> int: + """ + Return the local rank of the current process relative to the node. + + Semantically, this is a useful concept for mapping processes to devices. + For example, on a node with 8 accelerator you could use the node local rank to decide + which accelerator device to bind the process to. + + In practice, the actual assignment of node local ranks is handled by the process launcher outside of pytorch, + and communicated via the `LOCAL_RANK` environment variable. + + Torchrun will automatically populate `LOCAL_RANK`, but other launchers may not. If `LOCAL_RANK` is unspecified, + this API will fall back to the provided kwarg 'fallback_rank' if specified, otherwise it will raise an error. The + intent is to allow writing an application that runs either in single or multi device contexts without error. + + """ + if "LOCAL_RANK" in os.environ: + return int(os.environ["LOCAL_RANK"]) + elif fallback_rank is not None: + return int(fallback_rank) + raise RuntimeError( + "LOCAL_RANK is not in the environment. Consider passing fallback_rank to allow `get_node_local_rank` to work, " + "assuming you are not running in a multi-device context and want the code to run locally instead." + ) + + +def _add_ephemeral_timeout_for_all_pgs(timeout: timedelta) -> None: + """ + This API adds an ephemeral timeout extension for all PGs locally + on one rank. The timeout gets reset when the first collective issued + after API called finished. + NOTE: We only support to set timeout for cuda backends for now. + NOTE: While this feature + provides flexibility in specific scenarios, it introduces statefulness + to timeout setting. Therefore, it is advisable to use this API sparingly + and consider alternative approaches, such as directly setting the timeout + or utilizing a barrier collective (one can set any timeout to the barrier), + whenever feasible. + + Args: + timeout (timedelta): The delta of timeout to extend. + + Returns: + None. + """ + for pg in _world.pg_map.keys(): + devices = pg._device_types + if torch.device("cuda") in devices: + backend = pg._get_backend(torch.device("cuda")) + if is_nccl_available() and isinstance(backend, ProcessGroupNCCL): + backend._add_ephemeral_timeout(timeout) + + +def _set_pg_timeout(timeout: timedelta, group: Optional[ProcessGroup] = None) -> None: + """ + Set the timeout for the given process group when users want to use a different timeout instead of + default values. + + Args: + timeout (timedelta): Timeout for operations executed against the process group which + users want to set. Default value is 10 minutes for NCCL and 30 minutes for other backends. + This is the duration after which collectives will be aborted asynchronously and the process will crash. + This is done since CUDA execution is async and it is no longer safe to continue executing user code since + failed async NCCL operations might result in subsequent CUDA operations running on corrupted data. + When TORCH_NCCL_BLOCKING_WAIT is set, the process will block and wait for this timeout. + + group (ProcessGroup, optional): The process group to work on. The + default is the general main process group. If another specific group + is specified, the calling process must be part of :attr:`group`. + + Returns: + None + """ + if group is None: + group = _get_default_group() + if _rank_not_in_group(group): + raise ValueError("Invalid process group specified") + assert isinstance(group, ProcessGroup) + devices = group._device_types + backends = set() + if torch.device("cpu") in devices and is_gloo_available(): + backend = group._get_backend(torch.device("cpu")) + if isinstance(backend, ProcessGroupGloo): + backends.add(backend) + if torch.device("cuda") in devices: + backend = group._get_backend(torch.device("cuda")) + if is_nccl_available() and isinstance(backend, ProcessGroupNCCL): + backends.add(backend) # type: ignore[arg-type] + elif is_gloo_available() and isinstance(backend, ProcessGroupGloo): + backends.add(backend) # type: ignore[arg-type] + if len(backends) == 0: + warnings.warn("Set timeout is now only supported for either nccl or gloo.") + for backend in backends: + backend._set_default_timeout(timeout) + + +@_exception_logger +@_time_logger +def init_process_group( + backend: Optional[str] = None, + init_method: Optional[str] = None, + timeout: Optional[timedelta] = None, + world_size: int = -1, + rank: int = -1, + store: Optional[Store] = None, + group_name: str = "", + pg_options: Optional[Any] = None, + device_id: Optional[torch.device] = None, +) -> None: + """ + Initialize the default distributed process group. + + This will also initialize the distributed package. + + There are 2 main ways to initialize a process group: + 1. Specify ``store``, ``rank``, and ``world_size`` explicitly. + 2. Specify ``init_method`` (a URL string) which indicates where/how + to discover peers. Optionally specify ``rank`` and ``world_size``, + or encode all required parameters in the URL and omit them. + + If neither is specified, ``init_method`` is assumed to be "env://". + + + Args: + backend (str or Backend, optional): The backend to use. Depending on + build-time configurations, valid values include ``mpi``, ``gloo``, + ``nccl``, ``ucc``, or one that is registered by a third-party + plugin. + Since 2.6, if ``backend`` is not provided, c10d will use a backend + registered for the device type indicated by the `device_id` kwarg + (if provided). The known default registrations today are: ``nccl`` + for ``cuda``, ``gloo`` for ``cpu``. + If neither ``backend`` nor ``device_id`` is provided, c10d will + detect the accelerator on the run-time machine and use a backend + registered for that detected accelerator (or ``cpu``). + This field can be given as a lowercase string (e.g., ``"gloo"``), + which can also be accessed via :class:`Backend` attributes (e.g., + ``Backend.GLOO``). + If using multiple processes per machine with ``nccl`` backend, each + process must have exclusive access to every GPU it uses, as sharing + GPUs between processes can result in deadlock or NCCL invalid usage. + ``ucc`` backend is experimental. + init_method (str, optional): URL specifying how to initialize the + process group. Default is "env://" if no + ``init_method`` or ``store`` is specified. + Mutually exclusive with ``store``. + world_size (int, optional): Number of processes participating in + the job. Required if ``store`` is specified. + rank (int, optional): Rank of the current process (it should be a + number between 0 and ``world_size``-1). + Required if ``store`` is specified. + store(Store, optional): Key/value store accessible to all workers, used + to exchange connection/address information. + Mutually exclusive with ``init_method``. + timeout (timedelta, optional): Timeout for operations executed against + the process group. Default value is 10 minutes for NCCL and 30 minutes for other backends. + This is the duration after which collectives will be aborted asynchronously and the process will crash. + This is done since CUDA execution is async and it is no longer safe to continue executing user code since + failed async NCCL operations might result in subsequent CUDA operations running on corrupted data. + When TORCH_NCCL_BLOCKING_WAIT is set, the process will block and wait for this timeout. + + group_name (str, optional, deprecated): Group name. This argument is ignored + pg_options (ProcessGroupOptions, optional): process group options + specifying what additional options need to be passed in during + the construction of specific process groups. As of now, the only + options we support is ``ProcessGroupNCCL.Options`` for the ``nccl`` + backend, ``is_high_priority_stream`` can be specified so that + the nccl backend can pick up high priority cuda streams when + there're compute kernels waiting. For other availble options to config nccl, + See https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t + device_id (torch.device, optional): a single, specific device + to "bind" this process to, allowing for backend-specific + optimizations. Currently this has two effects, only under + NCCL: the communicator is immediately formed (calling + ``ncclCommInit*`` immediately rather than the normal lazy + call) and sub-groups will use ``ncclCommSplit`` when + possible to avoid unnecessary overhead of group creation. If you + want to know NCCL initialization error early, you can also use this + field. + + .. note:: To enable ``backend == Backend.MPI``, PyTorch needs to be built from source + on a system that supports MPI. + + .. note:: Support for multiple backends is experimental. Currently when no backend is + specified, both ``gloo`` and ``nccl`` backends will be created. The ``gloo`` backend + will be used for collectives with CPU tensors and the ``nccl`` backend will be used + for collectives with CUDA tensors. A custom backend can be specified by passing in + a string with format ":,:", e.g. + "cpu:gloo,cuda:custom_backend". + + """ + + global _world + + global _backend + global _default_pg_init_method + + if GroupMember.WORLD is not None: + raise ValueError("trying to initialize the default process group twice!") + + set_pytorch_distributed_envs_from_justknobs() + + # Depending on the import order, some trace_rules functions may be evaluated + # during the import phase. In such a case, these functions may not correctly + # add the distributed related rules due to import circular dependency. + # We need to clear the lru_cache during the runtime to ensure the correctness + # of these trace_rules. + # + # Since this API must be called before all distributed code being compiled, + # clearing the cache here should be safe. + if "torch._dynamo" in sys.modules: + torch._dynamo.trace_rules.clear_lru_cache() + + assert (store is None) or ( + init_method is None + ), "Cannot specify both init_method and store." + + if store is not None: + assert world_size > 0, "world_size must be positive if using store" + assert rank >= 0, "rank must be non-negative if using store" + elif init_method is None: + init_method = "env://" + + # If user did not provide a backend string but provided a device id, e.g. + # >>> init_process_group(device_id=device) + # we try to figure out the backend name based on the device type. + if backend is None and device_id is not None: + # Note: 3rd-party devices can register default backend through the + # default map below. + backend = Backend.default_device_backend_map.get(device_id.type) + + # If we still cannot figure it out, e.g. + # >>> init_process_group() + # we set it to `undefined` and rely on lazy init. + if backend is None: + backend = "undefined" + + # Convert string into `Backend` type + backend = Backend(backend) + + if timeout is None: + timeout = _get_default_timeout(backend) + + _check_valid_timeout(timeout) + + """ + Group name is not visible to users unless they access + internals of c10d. This means we can ignore the value + they provide as it not exposed in a public way. + """ + group_name = _process_group_name([], use_hashed_name=False) + if backend == Backend.MPI: + if world_size != -1 or rank != -1: + warnings.warn( + f"For MPI backend, world_size ({world_size}) and rank ({rank}) " + "are ignored since they are assigned by the " + "MPI runtime." + ) + + default_pg, _ = _new_process_group_helper( + -1, + -1, + [], + backend, + Store(), # Placeholder value since store cannot be None + group_name, + timeout=timeout, + group_desc="default_pg", + ) + _update_default_pg(default_pg) + else: + # backward compatible API + if store is None: + rendezvous_iterator = rendezvous( + not_none(init_method), rank, world_size, timeout=timeout + ) + store, rank, world_size = next(rendezvous_iterator) + store.set_timeout(timeout) + + # Use a PrefixStore to avoid accidental overrides of keys used by + # different systems (e.g. RPC) in case the store is multi-tenant. + store = PrefixStore("default_pg", store) + + default_pg, _ = _new_process_group_helper( + world_size, + rank, + [], + backend, + store, + group_name, + backend_options=pg_options, + timeout=timeout, + device_id=device_id, + group_desc="default_pg", + ) + _update_default_pg(default_pg) + + _world.pg_group_ranks[GroupMember.WORLD] = {i: i for i in range(GroupMember.WORLD.size())} # type: ignore[attr-defined, index] + _backend = _world.pg_map[not_none(GroupMember.WORLD)][0] + _default_pg_init_method = init_method + + old_hook = sys.excepthook + excepthook_prefix = f"[rank{get_rank()}]" + + def _distributed_excepthook(*args): + old_stderr = sys.stderr + sys.stderr = buf = io.StringIO() + try: + old_hook(*args) + finally: + sys.stderr = old_stderr + msg = buf.getvalue() + msg = "\n".join( + f"{excepthook_prefix}: {s}" if s != "" else "" for s in msg.split("\n") + ) + sys.stderr.write(msg) + sys.stderr.flush() + + sys.excepthook = _distributed_excepthook + + if _is_barrier_after_init() == 1: + # barrier at the end to ensure that once we return from this method, all + # process groups including global variables (if any) are updated + # correctly on all ranks. + # Update 04/2023: for large-scale runs, this barrier (esp. store-based + # barrier) may be costly and/or unscalable. Also, in a lot of cases, + # these barriers may be unnecessary, as proven by a green CI after + # removal. An environment variable `TORCH_DIST_INIT_BARRIER` has been + # added which enables this barrier only when set to 1. + logger.debug( + "Performing barrier after ProcessGroup initialization since " + "TORCH_DIST_INIT_BARRIER = 1" + ) + if backend == Backend.MPI: + # MPI backend doesn't use store. + barrier() + else: + # Use store based barrier here since barrier() used a bunch of + # default devices and messes up NCCL internal state. + _store_based_barrier(rank, store, group_name, world_size, timeout) + + +def _get_split_source(pg): + split_from = None + if pg.bound_device_id: + split_from = pg._get_backend(pg.bound_device_id) + elif pg is _world.default_pg: + try: + split_from = pg._get_backend(torch.device("cuda")) + except RuntimeError: + # no cuda device associated with this backend + pass + + if not split_from or not split_from.supports_splitting: + return None + + # If necessary, find a backend to split from by peeling process + # group wrappers from our potentially wrapped process group. + while _GLOO_AVAILABLE and isinstance(split_from, _ProcessGroupWrapper): + split_from = split_from.wrapped_pg + + return split_from + + +def _shutdown_backend(pg): + """ + Try to shut down the backend of a process group. + Currently, only ProcessGroupNCCL backend is supported. + No op for other backends. + """ + backend = None + try: + backend = pg._get_backend(torch.device("cuda")) + except RuntimeError: + pass + if is_nccl_available() and isinstance(backend, ProcessGroupNCCL): + # explictly call shutdown to ensure that NCCL resources are released + backend._shutdown() + + +def _abort_backend(pg: ProcessGroup): + """ + Abort the backend of a process group. + Currently, only ProcessGroupNCCL backend is supported. + No op for other backends. + """ + try: + backend = pg._get_backend(torch.device("cuda")) + except RuntimeError: + backend = None + if isinstance(backend, ProcessGroupNCCL): + backend.abort() + + +def _new_process_group_helper( + group_size, + group_rank, + global_ranks_in_group, + backend, + store, + group_name, + backend_options=None, + timeout=None, + pg_tag=None, + device_id=None, + group_desc=None, +): + """ + Create a new distributed process group. + + This function must be called by ALL processes in the global group, even if + the calling process is not part of the newly created group. In that case, + this function returns GroupMember.NON_GROUP_MEMBER. + + This function is called with ``global_ranks_in_group == []`` for the default group. + """ + global _world + + if group_name in _world.pg_names.values(): + raise ValueError( + "The specified group name has already been " + "created, please use a different group name" + ) + + if device_id is not None and (device_id.index is None or device_id.type != "cuda"): + raise ValueError( + "init_process_group device_id parameter must be a cuda device with an " + "id, e.g. cuda:0, not just cuda or cpu" + ) + + # Note: _new_process_group_helper is only called from init_process_group, which always provides a timeout value + _check_valid_timeout(timeout) + + if pg_tag not in [None, ""]: + # creating with the same tag and rank set results in the same underlying PG + existing_group = _find_pg_by_ranks_and_tag(pg_tag, global_ranks_in_group) + if existing_group: + _, prefix_store = _world.pg_map[existing_group] + return existing_group, prefix_store + + group_desc = "undefined" if group_desc is None else group_desc + + # The list of group ranks is empty if we're creating the default group. + is_default_group = len(global_ranks_in_group) == 0 + + # nccl and potentially other backends allow creation of + # communicators based on pre-existing ones, which can save + # initialization time. Due to lazy initialization of + # communicators in some backends, we have to be careful and only + # split when we *know* the default PG has already started communicator initialization. + # We know this if we have bound a device id to the default pg (eager initialized). + if is_initialized() and _get_default_group().bound_device_id: + split_from = _get_split_source(_get_default_group()) + else: + split_from = None + + # If this is a subgroup (which means group_ranks is specified), + # we check if the current process is a member of the new group. + if not is_default_group: + global_rank = _get_default_group().rank() + if global_rank not in global_ranks_in_group: + # If we are using `ncclCommSplit` (or similar split from + # other APIs) to create the communicator, we will need to + # call `ncclCommSplit` on *all* ranks in this new group's + # parent group, even those not in the new group. This is + # a requirement of the NCCL API as otherwise we would get + # out of sync. + if split_from: + split_from.perform_nocolor_split(_get_default_group().bound_device_id) + return GroupMember.NON_GROUP_MEMBER, None + + prefix_store = PrefixStore(f"{group_name}/", store) + # The backend for PG will be set later based on what's inside BackendConfig + # and timeout are set in each backend's option. + pg: ProcessGroup = ProcessGroup( + prefix_store, + group_rank, + group_size, + ) + # Set the default backend when only single backend is passed in. + if "," not in str(backend) and ":" not in str(backend): + assert backend in Backend.backend_type_map, f"Unknown backend type {backend}" + pg._set_default_backend(Backend.backend_type_map[backend]) + if device_id: + pg.bound_device_id = device_id + backend_config = BackendConfig(backend) + backend_class: torch._C._distributed_c10d.Backend + for device, backend_str in backend_config.get_device_backend_map().items(): + # Use the group name as prefix in the default store, such that + # a single store can be reused by multiple groups. + backend_prefix_store = PrefixStore(f"{device}/", prefix_store) + + if backend_str == Backend.MPI: + if not is_mpi_available(): + raise RuntimeError( + "Distributed package doesn't have MPI built in." + " MPI is only included if you build PyTorch from" + " source on a host that has MPI installed." + ) + backend_class = ProcessGroupMPI.create(global_ranks_in_group) + backend_type = ProcessGroup.BackendType.MPI + if not backend_class: + return GroupMember.NON_GROUP_MEMBER, None + # create new process group with accurate rank and size + if pg.rank() == -1 and pg.size() == -1: + pg = ProcessGroup( + backend_prefix_store, + backend_class.rank(), + backend_class.size(), + ) + pg._set_default_backend(backend_type) + elif backend_str == Backend.GLOO: + # TODO: remove this check after lazy initialization is supported + # if pg_options is not None: + # raise RuntimeError("GLOO options not supported") + backend_class = ProcessGroupGloo( + backend_prefix_store, group_rank, group_size, timeout=timeout + ) + backend_type = ProcessGroup.BackendType.GLOO + elif backend_str == Backend.NCCL: + if not is_nccl_available(): + raise RuntimeError("Distributed package doesn't have NCCL built in") + if backend_options is not None: + assert isinstance( + backend_options, ProcessGroupNCCL.Options + ), "Expected backend_options argument to be of type ProcessGroupNCCL.Options" + if backend_options._timeout != timeout: + warnings.warn( + "backend_options._timeout was specified, " + "but timeout kwarg has a default value that will always override it. " + ) + else: + # default backend_options for NCCL + backend_options = ProcessGroupNCCL.Options() + backend_options.is_high_priority_stream = False + backend_options._timeout = timeout + + if split_from: + backend_options.split_from = split_from + backend_options.split_color = _process_group_color( + global_ranks_in_group + ) + backend_options.global_ranks_in_group = global_ranks_in_group + backend_options.group_name = group_name + backend_class = ProcessGroupNCCL( + backend_prefix_store, group_rank, group_size, backend_options + ) + backend_type = ProcessGroup.BackendType.NCCL + elif backend_str == Backend.UCC and is_ucc_available(): + # TODO: once UCC plugin is fully deprecated, remove + # is_ucc_available() from above elif-condition and raise + # RuntimeError if is_ucc_available() returns false. + + backend_class = ProcessGroupUCC( + backend_prefix_store, group_rank, group_size, timeout=timeout + ) + backend_type = ProcessGroup.BackendType.UCC + elif backend_str == Backend.XCCL: + if not is_xccl_available(): + raise RuntimeError("Distributed package doesn't have XCCL built in") + backend_class = ProcessGroupXCCL( + backend_prefix_store, group_rank, group_size + ) + backend_type = ProcessGroup.BackendType.XCCL + else: + assert ( + backend_str.upper() in Backend._plugins + ), f"Unknown c10d backend type {backend_str.upper()}" + + backend_plugin = Backend._plugins[backend_str.upper()] + creator_fn = backend_plugin.creator_fn + extended_api = backend_plugin.extended_api + backend_type = ProcessGroup.BackendType.CUSTOM + + if not extended_api: + backend_class = creator_fn( + backend_prefix_store, group_rank, group_size, timeout + ) + else: + dist_backend_opts = _DistributedBackendOptions() + dist_backend_opts.store = backend_prefix_store + dist_backend_opts.group_rank = group_rank + dist_backend_opts.group_size = group_size + dist_backend_opts.timeout = timeout + dist_backend_opts.group_id = group_name + dist_backend_opts.global_ranks_in_group = global_ranks_in_group + + backend_class = creator_fn(dist_backend_opts, backend_options) + + # Set sequence numbers for gloo and nccl backends. + if backend_str == Backend.GLOO: + assert isinstance(backend_class, ProcessGroupGloo) + backend_class._set_sequence_number_for_group() + elif backend_str == Backend.NCCL: + assert isinstance(backend_class, ProcessGroupNCCL) + backend_class._set_sequence_number_for_group() + + # If the type is a subclass of ProcessGroup then return this process group immediately + # TODO: This defaults to the old behavior for PythonProcessGroups which overwrites the + # ProcessGroup instance + if issubclass(type(backend_class), ProcessGroup): + pg = backend_class # type: ignore[assignment] + break + + # Process group wrapper initialization for supported PGs when TORCH_DISTRIBUTED_DEBUG is set + if ( + backend_str in [Backend.GLOO, Backend.NCCL, Backend.UCC] + or backend_str.upper() in Backend._plugins + ): + # In debug mode and if GLOO is available, wrap in a wrapper PG that + # enables enhanced collective checking for debuggability. + if get_debug_level() == DebugLevel.DETAIL: + if not _GLOO_AVAILABLE: + logger.info( + """TORCH_DISTRIBUTED_DEBUG was set to DETAIL, but + GLOO is not available. Build with Gloo to + create a wrapper process group in debug mode + to aid collective desynchronization debugging.""" + ) + else: + backend_class = _create_process_group_wrapper( + wrapped_pg=backend_class, + store_prefix=group_name, + store=backend_prefix_store, + rank=group_rank, + world_size=group_size, + timeout=timeout, + ) + + # register only a single backend when all get_device_backend_map values are the same + if len(set(backend_config.get_device_backend_map().values())) == 1: + for device in backend_config.get_device_backend_map().keys(): + pg._register_backend(torch.device(device), backend_type, backend_class) + + # break out of outer loop to not create any more backends + break + + pg._register_backend(torch.device(device), backend_type, backend_class) + + # set group_name and group_dsec to backend + assert group_name is not None + assert group_desc is not None + pg._set_group_name(group_name) + pg._set_group_desc(group_desc) + + if device_id and pg._get_backend(device_id).supports_splitting: + eager_backend = pg._get_backend(device_id) + eager_backend.eager_connect_single_device(device_id) + + # update global state + _world.pg_map[pg] = (backend, prefix_store) + _world.pg_names[pg] = group_name + _register_process_group(group_name, pg) + + _world.pg_backend_config[pg] = str(backend_config) + # "" is the default tag for user PGs + if pg_tag in [None, ""]: + pg_tag = f"ptd:{group_name}" + _world.tags_to_pg.setdefault("", []).append(pg) + else: + pg_tag = f"user:{pg_tag}" + + _world.tags_to_pg.setdefault(pg_tag, []).append(pg) + _world.pg_to_tag[pg] = pg_tag + return pg, prefix_store + + +def destroy_process_group(group: Optional[ProcessGroup] = None): + """ + Destroy a given process group, and deinitialize the distributed package. + + Args: + group (ProcessGroup, optional): The process group to be destroyed, if + group.WORLD is given, all process + groups including the default one will + be destroyed. + """ + global _world + + if group == GroupMember.NON_GROUP_MEMBER: + return + + if group is None: + pg = GroupMember.WORLD + else: + pg = group + + assert pg is not None + if _world.pg_map.get(pg, None) is None: + raise ValueError("Invalid process group specified") + + # When users register Python onCompletion hooks, those hooks will run on a + # different thread than the main thread. Today, the ProcessGroup dtor does + # wait for that thread. However, the dtor might finish after the Python + # Interpreter exits. After that grabbing the GIL for the Python hook will crash. + # We can either revive the interpreter when running hooks or keep the main one + # alive until all works and hooks are done. The current implementation does the + # latter. Therefore, we explicitly call _wait_for_pending_works() here to wait + # for the pending hooks to finish. + if pg.name().lower() == "nccl" and pg._has_hooks(): + pg._wait_for_pending_works() + + if group is None or group == GroupMember.WORLD: + # shutdown all backends in the order of pg names. shutting down in order because + # ncclCommAbort() was a 'collective' call in some versions of NCCL. + for pg_to_shutdown in sorted( + _world.pg_names, key=lambda x: _world.pg_names[x], reverse=True + ): + _shutdown_backend(pg_to_shutdown) + + _update_default_pg(None) + _world.pg_map.clear() + _world.pg_names.clear() + _world.pg_group_ranks.clear() + _world.pg_backend_config.clear() + _world.pg_to_tag.clear() + _world.tags_to_pg.clear() + _world.pg_coalesce_state.clear() + _unregister_all_process_groups() + + # when process group doesn't have an explicit name (only WORLD (default) + # process group can have an explicit name), we use global _world.group_count + # to generate the name. We need to reset the counter on destruction to + # allow consistent value to be generated when we re-create process + # groups after some trainers recover from failure + # + # We only reset this when WORLD is being destroyed because if this + # process group is in good state, we aren't dealing with failures. + _world.group_count = 0 + else: + _shutdown_backend(pg) + del _world.pg_map[pg] + del _world.pg_names[pg] + del _world.pg_group_ranks[pg] + del _world.pg_backend_config[pg] + if pg in _world.pg_coalesce_state.keys(): + warnings.warn( + "Some coalesced collectives haven't been launched when " + "ProcessGroup is destroyed. They will be cleaned." + ) + del _world.pg_coalesce_state[pg] + + tag = _world.pg_to_tag.get(pg) + del _world.pg_to_tag[pg] + if tag is not None: + try: + _world.tags_to_pg[tag].remove(pg) + if tag.startswith("ptd:"): + _world.tags_to_pg[""].remove(pg) + except Exception: + pass + _unregister_process_group(pg.group_name) + + +def _abort_process_group(group: Optional[ProcessGroup] = None): + """ + Abort a given process group. If group.WORLD (i.e. `None`) is given, all + process groups including the default one will be aborted. + + Args: + group (ProcessGroup, optional): The process group to be aborted. + + .. note:: this API is experimental and currently only works with the NCCL + backend. + + .. note:: this API should be used with `TORCH_NCCL_ASYNC_ERROR_HANDLING` + turned off (i.e. set to 0). Otherwise, ProcessGroupNCCL's watchdog may + automatically handle errors or timeouts for you including aborting the + ProcessGroup. + """ + global _world + + if group == GroupMember.NON_GROUP_MEMBER: + return + + pg = group or GroupMember.WORLD + + assert pg is not None + if _world.pg_map.get(pg, None) is None: + raise ValueError("Invalid process group specified or has been destroyed.") + + try: + backend = pg._get_backend(torch.device("cuda")) + except RuntimeError: + backend = None + + if not isinstance(backend, ProcessGroupNCCL): + logger.warning( + "`abort_process_group` currently only has implementation for ProcessGroupNCCL; " + "however, no NCCL backend is found. This call will be a no-op." + ) + return + + if group == GroupMember.WORLD: + # Abort all backends within a ncclGroupStart|End semantic. + # This ensures that different NCCL communicators' abort calls won't + # deadlock each other. + # For details, please see: https://github.com/pytorch/pytorch/issues/119797 + backend._group_start() + for pg_to_abort in sorted( + _world.pg_names, key=lambda x: _world.pg_names[x], reverse=True + ): + _abort_backend(pg_to_abort) + backend._group_end() + + _update_default_pg(None) + _world.pg_map.clear() + _world.pg_names.clear() + _world.pg_group_ranks.clear() + _world.pg_backend_config.clear() + _world.pg_to_tag.clear() + _world.tags_to_pg.clear() + _world.pg_coalesce_state.clear() + _unregister_all_process_groups() + + # when process group doesn't have an explicit name (only WORLD (default) + # process group can have an explicit name), we use global _world.group_count + # to generate the name. We need to reset the counter on destruction to + # allow consistent value to be generated when we re-create process + # groups after some trainers recover from failure + # + # We only reset this when WORLD is being destroyed because if this + # process group is in good state, we aren't dealing with failures. + _world.group_count = 0 + else: + _abort_backend(pg) + del _world.pg_map[pg] + del _world.pg_names[pg] + del _world.pg_group_ranks[pg] + del _world.pg_backend_config[pg] + if pg in _world.pg_coalesce_state.keys(): + warnings.warn( + "Some coalesced collectives haven't been launched when " + "ProcessGroup is aborted. They will be cleaned." + ) + del _world.pg_coalesce_state[pg] + + tag = _world.pg_to_tag.get(pg) + del _world.pg_to_tag[pg] + if tag is not None: + try: + _world.tags_to_pg[tag].remove(pg) + if tag.startswith("ptd:"): + _world.tags_to_pg[""].remove(pg) + except Exception: + pass + _unregister_process_group(pg.group_name) + + +def get_rank(group: Optional[ProcessGroup] = None) -> int: + """ + Return the rank of the current process in the provided ``group``, default otherwise. + + Rank is a unique identifier assigned to each process within a distributed + process group. They are always consecutive integers ranging from 0 to + ``world_size``. + + Args: + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + + Returns: + The rank of the process group + -1, if not part of the group + + """ + if _rank_not_in_group(group): + return -1 + + default_pg = _get_default_group() + if group is None or group is GroupMember.WORLD: + return default_pg.rank() + + return get_group_rank(group, default_pg.rank()) + + +def get_world_size(group: Optional[ProcessGroup] = None) -> int: + """ + Return the number of processes in the current process group. + + Args: + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + + Returns: + The world size of the process group + -1, if not part of the group + + """ + if _rank_not_in_group(group): + return -1 + + return _get_group_size(group) + + +def isend( + tensor: torch.Tensor, + dst: Optional[int] = None, + group: Optional[ProcessGroup] = None, + tag: int = 0, + group_dst: Optional[int] = None, +) -> Optional[Work]: + """ + Send a tensor asynchronously. + + .. warning:: + Modifying ``tensor`` before the request completes causes undefined + behavior. + + .. warning:: + ``tag`` is not supported with the NCCL backend. + + Unlike send, which is blocking, isend allows src == dst rank, i.e. send to self. + + Args: + tensor (Tensor): Tensor to send. + dst (int): Destination rank on global process group (regardless of ``group`` argument) + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + tag (int, optional): Tag to match send with remote recv + group_dst (int, optional): Destination rank on ``group``. Invalid to specify both ``dst`` and ``group_dst`` + + Returns: + A distributed request object. + None, if not part of the group + + """ + group = _group_or_default_group(group) + group_dst = _canonicalize_group_rank(group, dst, group_dst) + _check_single_tensor(tensor, "tensor") + if _rank_not_in_group(group): + _warn_not_in_group("isend") + return None + + if tensor.is_complex(): + tensor = torch.view_as_real(tensor) + + return group.send([tensor], group_dst, tag) + + +def irecv( + tensor: torch.Tensor, + src: Optional[int] = None, + group: Optional[ProcessGroup] = None, + tag: int = 0, + group_src: Optional[int] = None, +) -> Optional[Work]: + """ + Receives a tensor asynchronously. + + .. warning:: + ``tag`` is not supported with the NCCL backend. + + Unlike recv, which is blocking, irecv allows src == dst rank, i.e. recv from self. + + Args: + tensor (Tensor): Tensor to fill with received data. + src (int, optional): Source rank on global process group (regardless of ``group`` argument). + Will receive from any process if unspecified. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + tag (int, optional): Tag to match recv with remote send + group_src (int, optional): Destination rank on ``group``. Invalid to specify both ``src`` and ``group_src``. + + Returns: + A distributed request object. + None, if not part of the group + + """ + _check_single_tensor(tensor, "tensor") + if _rank_not_in_group(group): + _warn_not_in_group("irecv") + return None + + if tensor.is_complex(): + tensor = torch.view_as_real(tensor) + + group = _group_or_default_group(group) + if src is None and group_src is None: + return group.recv_anysource([tensor], tag) + else: + group_src = _canonicalize_group_rank(group, src, group_src) + return group.recv([tensor], group_src, tag) + + +@_exception_logger +def send( + tensor: torch.Tensor, + dst: Optional[int] = None, + group: Optional[ProcessGroup] = None, + tag: int = 0, + group_dst: Optional[int] = None, +) -> None: + """ + Send a tensor synchronously. + + .. warning:: + ``tag`` is not supported with the NCCL backend. + + Args: + tensor (Tensor): Tensor to send. + dst (int): Destination rank on global process group (regardless of ``group`` argument). + Destination rank should not be the same as the rank of the current process. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + tag (int, optional): Tag to match send with remote recv + group_dst (int, optional): Destination rank on ``group``. Invalid to specify both ``dst`` and ``group_dst``. + + """ + group = _group_or_default_group(group) + group_dst = _canonicalize_group_rank(group, dst, group_dst) + _check_not_self_rank(group, group_dst, "destination") + work = isend(tensor, group=group, tag=tag, group_dst=group_dst) + if work is not None: + work.wait() + + +@_exception_logger +def recv( + tensor: torch.Tensor, + src: Optional[int] = None, + group: Optional[ProcessGroup] = None, + tag: int = 0, + group_src: Optional[int] = None, +) -> int: + """ + Receives a tensor synchronously. + + .. warning:: + ``tag`` is not supported with the NCCL backend. + + Args: + tensor (Tensor): Tensor to fill with received data. + src (int, optional): Source rank on global process group (regardless of ``group`` argument). + Will receive from any process if unspecified. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + tag (int, optional): Tag to match recv with remote send + group_src (int, optional): Destination rank on ``group``. Invalid to specify both ``src`` and ``group_src``. + Returns: + Sender rank + -1, if not part of the group + + """ + work = irecv(tensor, src=src, group=group, tag=tag, group_src=group_src) + if work is None: + return -1 + work.wait() + if src is None: + if group_src is None: + group_src = work._source_rank() + group = _group_or_default_group(group) + _check_not_self_rank(group, group_src, "source") + src = get_global_rank(group, group_src) + return src + + +class _IllegalWork(Work): + def __getattribute__(self, name): + if name in [ + "is_success", + "exception", + "wait", + "source_rank", + "_source_rank", + "result", + "synchronize", + ]: + raise ValueError(f"Illegal to call {name} on IllegalWork object") + + +class _CoalescingManager: + def __init__(self) -> None: + self.works: List[Work] = [] + + def append(self, work: Work): + if work: + self.works.append(work) + + def wait(self): + for work in self.works: + work.wait() + + +@contextlib.contextmanager +def _coalescing_manager( + group: Optional[ProcessGroup] = None, + device: Optional[torch.device] = None, + async_ops: Optional[bool] = False, +): + """ + Context manager used to coalesce collectives or P2P operations when possible. + + Args: + group (`ProcessGroup`, optional): The process group to work on. If None, + the default process group will be used. + device (`torch.device`, optional): Default is None, set to a device if + there isn't a `**_coalesced` implementation by the backend. + async_ops (`bool`, optional): whether the coalesced ops are async ops. + + Examples: + >>> # xdoctest: +SKIP("no rank") + >>> # Synchronous ops + >>> with _coalescing_manager(): + >>> for i in range(num_colls): + >>> dist.all_reduce(tensors[i]) + >>> # Asynchronous ops + >>> with _coalescing_manager(async_ops=True) as cm: + >>> for i in range(num_colls): + >>> dist.all_reduce(tensors[i]) + >>> cm.wait() + + .. warning:: + :func:`_coalescing_manager` currently do not support coalescing + all-reduces with different reduce operators, e.g. `ReduceOp.SUM` mixed + with `ReduceOp.PRODUCT`. + """ + group = group or _get_default_group() + op_list = _world.pg_coalesce_state.setdefault(group, []) + if op_list: + raise ValueError( + "ProcessGroup has non-empty op list at the start of coalescing" + ) + if device: + group._start_coalescing(device) + cm = _CoalescingManager() + yield cm + op_list = _world.pg_coalesce_state.pop(group) + if op_list: + # Collectives supporting "Fast Path" coalescing are captured. + # See implementation in corresponding collective APIs. + # Currently supported: + # - coalesced `all_reduce` + # - coalesced `all_gather_into_tensor` + # - coalesced `reduce_scatter_tensor` + op0 = op_list[0].op + if op0 == all_reduce: + tensors = [op.tensor for op in op_list] + all_reduce_opts = AllreduceCoalescedOptions() + all_reduce_opts.reduceOp = not_none(op_list[0].redop) + work = group.allreduce_coalesced(tensors, all_reduce_opts) + elif op0 == all_gather_into_tensor: + inputs = [] + outputs = [] + for op in op_list: + inputs.append(op.tensor) + outputs.append(not_none(op.dst_tensor)) + work = group.allgather_into_tensor_coalesced(outputs, inputs) + elif op0 == reduce_scatter_tensor: + inputs = [] + outputs = [] + for op in op_list: + inputs.append(op.tensor) + outputs.append(not_none(op.dst_tensor)) + reduce_opts = ReduceScatterOptions() + reduce_opts.reduceOp = not_none(op_list[0].redop) + work = group.reduce_scatter_tensor_coalesced(outputs, inputs, reduce_opts) + else: + raise AssertionError( + f"Coalescing manager does not support fast-path coalescing of {op0}, " + f"yet {op0} is still recorded in op list. This is an internal error of c10d." + ) + + if device: + # Old style of letting each coll inside the context manager to call into C++ counterpart via python binding + work = group._end_coalescing(device) + + if async_ops: + cm.append(work) # type: ignore[possibly-undefined] + else: + work.wait() # type: ignore[possibly-undefined] + + +def batch_isend_irecv(p2p_op_list: List[P2POp]) -> List[Work]: + """ + Send or Receive a batch of tensors asynchronously and return a list of requests. + + Process each of the operations in ``p2p_op_list`` and return the corresponding + requests. NCCL, Gloo, and UCC backend are currently supported. + + Args: + p2p_op_list: A list of point-to-point operations(type of each operator is + ``torch.distributed.P2POp``). The order of the isend/irecv in the list + matters and it needs to match with corresponding isend/irecv on the + remote end. + + Returns: + A list of distributed request objects returned by calling the corresponding + op in the op_list. + + Examples: + >>> # xdoctest: +SKIP("no rank") + >>> send_tensor = torch.arange(2, dtype=torch.float32) + 2 * rank + >>> recv_tensor = torch.randn(2, dtype=torch.float32) + >>> send_op = dist.P2POp(dist.isend, send_tensor, (rank + 1)%world_size) + >>> recv_op = dist.P2POp(dist.irecv, recv_tensor, (rank - 1 + world_size)%world_size) + >>> reqs = batch_isend_irecv([send_op, recv_op]) + >>> for req in reqs: + >>> req.wait() + >>> recv_tensor + tensor([2, 3]) # Rank 0 + tensor([0, 1]) # Rank 1 + + .. note:: Note that when this API is used with the NCCL PG backend, users must set + the current GPU device with `torch.cuda.set_device`, otherwise it will + lead to unexpected hang issues. + + In addition, if this API is the first collective call in the ``group`` + passed to ``dist.P2POp``, all ranks of the ``group`` must participate in + this API call; otherwise, the behavior is undefined. If this API call is + not the first collective call in the ``group``, batched P2P operations + involving only a subset of ranks of the ``group`` are allowed. + """ + _check_p2p_op_list(p2p_op_list) + group = p2p_op_list[0].group + device = p2p_op_list[0].tensor.device + + def peer_kwarg(op: P2POp) -> Dict[str, int]: + key = "group_dst" if op.op == isend else "group_src" + return {key: op.group_peer} + + if device.type == "cuda": + # NCCL style coalescing + with _coalescing_manager(group, device, async_ops=True) as cm: + for p2p_op in p2p_op_list: + p2p_op.op( + p2p_op.tensor, + group=p2p_op.group, + tag=p2p_op.tag, + **peer_kwarg(p2p_op), + ) + + return cm.works + else: + # Backward support for Gloo + reqs = [] + for p2p_op in p2p_op_list: + work = p2p_op.op( + p2p_op.tensor, + group=p2p_op.group, + tag=p2p_op.tag, + **peer_kwarg(p2p_op), + ) + if work: + reqs.append(work) + return reqs + + +@_exception_logger +def broadcast( + tensor: torch.Tensor, + src: Optional[int] = None, + group: Optional[ProcessGroup] = None, + async_op: bool = False, + group_src: Optional[int] = None, +): + """ + Broadcasts the tensor to the whole group. + + ``tensor`` must have the same number of elements in all processes + participating in the collective. + + Args: + tensor (Tensor): Data to be sent if ``src`` is the rank of current + process, and tensor to be used to save received data otherwise. + src (int): Source rank on global process group (regardless of ``group`` argument). + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op + group_src (int): Source rank on ``group``. Must specify one of ``group_src`` + and ``src`` but not both. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + """ + group = _group_or_default_group(group) + group_src = _canonicalize_group_rank(group, src, group_src, return_global=False) + _check_single_tensor(tensor, "tensor") + if _rank_not_in_group(group): + _warn_not_in_group("broadcast") + return + + opts = BroadcastOptions() + opts.rootRank = group_src + opts.rootTensor = 0 + opts.asyncOp = async_op + work = group.broadcast([tensor], opts) + if async_op: + return work + else: + work.wait() + + +@_exception_logger +def all_reduce(tensor, op=ReduceOp.SUM, group=None, async_op=False): + """ + Reduces the tensor data across all machines in a way that all get the final result. + + After the call ``tensor`` is going to be bitwise identical in all processes. + + Complex tensors are supported. + + Args: + tensor (Tensor): Input and output of the collective. The function + operates in-place. + op (optional): One of the values from + ``torch.distributed.ReduceOp`` + enum. Specifies an operation used for element-wise reductions. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + Examples: + >>> # xdoctest: +SKIP("no rank") + >>> # All tensors below are of torch.int64 type. + >>> # We have 2 process groups, 2 ranks. + >>> device = torch.device(f'cuda:{rank}') + >>> tensor = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank + >>> tensor + tensor([1, 2], device='cuda:0') # Rank 0 + tensor([3, 4], device='cuda:1') # Rank 1 + >>> dist.all_reduce(tensor, op=ReduceOp.SUM) + >>> tensor + tensor([4, 6], device='cuda:0') # Rank 0 + tensor([4, 6], device='cuda:1') # Rank 1 + + >>> # All tensors below are of torch.cfloat type. + >>> # We have 2 process groups, 2 ranks. + >>> tensor = torch.tensor([1+1j, 2+2j], dtype=torch.cfloat, device=device) + 2 * rank * (1+1j) + >>> tensor + tensor([1.+1.j, 2.+2.j], device='cuda:0') # Rank 0 + tensor([3.+3.j, 4.+4.j], device='cuda:1') # Rank 1 + >>> dist.all_reduce(tensor, op=ReduceOp.SUM) + >>> tensor + tensor([4.+4.j, 6.+6.j], device='cuda:0') # Rank 0 + tensor([4.+4.j, 6.+6.j], device='cuda:1') # Rank 1 + + """ + _check_single_tensor(tensor, "tensor") + if _rank_not_in_group(group): + _warn_not_in_group("all_reduce") + return + + if tensor.is_complex(): + if not supports_complex(op): + raise ValueError(f"all_reduce does not support {op} on complex tensors") + tensor = torch.view_as_real(tensor) + + opts = AllreduceOptions() + opts.reduceOp = op + if group is None: + group = _get_default_group() + + if group in _world.pg_coalesce_state.keys(): + # We are in coalescing context, do not issue single operation, just append a collective representation + coll = _CollOp(all_reduce, tensor, None, op, None) + _world.pg_coalesce_state[group].append(coll) + if async_op: + return _IllegalWork() + else: + return None + + work = group.allreduce([tensor], opts) + + if async_op: + return work + else: + work.wait() + + +@_exception_logger +@deprecated( + "`torch.distributed.all_reduce_coalesced` will be deprecated. If you must " + "use it, please revisit our documentation later at " + "https://pytorch.org/docs/main/distributed.html#collective-functions", + category=FutureWarning, +) +def all_reduce_coalesced(tensors, op=ReduceOp.SUM, group=None, async_op=False): + """ + WARNING: at this time individual shape checking is not implemented across nodes. + + For example, if the rank 0 node passes [torch.rand(4), torch.rand(2)] and the + rank 1 node passes [torch.rand(2), torch.rand(2), torch.rand(2)], the allreduce + operation will proceed without complaint and return erroneous outputs. This lack + of shape checking results in significant performance improvements but users of this + function should take extra care to ensure that each node passes in tensors whose + shapes match across nodes. + + Reduces each tensor in tensors (residing on the same device) across all machines + in such a way that all get the final result. + + After the call each tensor in tensors is going to bitwise identical + in all processes. + + Complex tensors are supported. + + Args: + tensors (Union[List[Tensor], Tensor]): Input and output of the collective. + The function operates in-place. + op (Optional[ReduceOp]): One of the values from + ``torch.distributed.ReduceOp`` enum. Specifies an operation used for + element-wise reductions. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (Optional[bool]): Whether this op should be an async op. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group. + + """ + if isinstance(tensors, torch.Tensor): + tensors = [tensors] + _check_tensor_list(tensors, "tensor") + _ensure_all_tensors_same_dtype(tensors) + if _rank_not_in_group(group): + _warn_not_in_group("all_reduce_coalesced") + return + + if any(t.is_complex() for t in tensors) and not supports_complex(op): + raise ValueError(f"all_reduce does not support {op} on complex tensors") + + tensors = [t if not t.is_complex() else torch.view_as_real(t) for t in tensors] + + opts = AllreduceCoalescedOptions() + opts.reduceOp = op + group = group or _get_default_group() + work = group.allreduce_coalesced(tensors, opts) + + if async_op: + return work.get_future() + else: + work.wait() + + +@_exception_logger +def reduce( + tensor: torch.Tensor, + dst: Optional[int] = None, + op=ReduceOp.SUM, + group: Optional[ProcessGroup] = None, + async_op: bool = False, + group_dst: Optional[int] = None, +): + """ + Reduces the tensor data across all machines. + + Only the process with rank ``dst`` is going to receive the final result. + + Args: + tensor (Tensor): Input and output of the collective. The function + operates in-place. + dst (int): Destination rank on global process group (regardless of ``group`` argument) + op (optional): One of the values from + ``torch.distributed.ReduceOp`` + enum. Specifies an operation used for element-wise reductions. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op + group_dst (int): Destination rank on ``group``. Must specify one of ``group_dst`` + and ``dst`` but not both. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + """ + group = _group_or_default_group(group) + group_dst = _canonicalize_group_rank(group, dst, group_dst, return_global=False) + _check_single_tensor(tensor, "tensor") + if _rank_not_in_group(group): + _warn_not_in_group("reduce") + return + + opts = ReduceOptions() + opts.reduceOp = op + opts.rootRank = group_dst + work = group.reduce([tensor], opts) + if async_op: + return work + else: + work.wait() + + +def _object_to_tensor(obj, device, group): + with _WaitCounter("pytorch.wait_counter.c10d._object_to_tensor").guard(): + f = io.BytesIO() + _pickler(f).dump(obj) + byte_storage = torch.ByteStorage._from_buffer(f.getvalue()) # type: ignore[attr-defined] + # Do not replace `torch.ByteTensor` or `torch.LongTensor` with torch.tensor and specifying dtype. + # Otherwise, it will casue 100X slowdown. + # See: https://github.com/pytorch/pytorch/issues/65696 + byte_tensor = torch.ByteTensor(byte_storage).to(device) + if get_debug_level() == DebugLevel.DETAIL and is_nccl_available(): + backend = get_backend(group) + if backend == Backend.NCCL: + hash = torch._C._distributed_c10d._hash_tensors([byte_tensor]) + logger.warning( + "_object_to_tensor size: %s hash value: %s", + byte_tensor.numel(), + hash, + ) + local_size = torch.LongTensor([byte_tensor.numel()]).to(device) + return byte_tensor, local_size + + +def _tensor_to_object(tensor, tensor_size, group): + with _WaitCounter("pytorch.wait_counter.c10d._tensor_to_object").guard(): + if get_debug_level() == DebugLevel.DETAIL and is_nccl_available(): + backend = get_backend(group) + if backend == Backend.NCCL: + hash = torch._C._distributed_c10d._hash_tensors([tensor]) + logger.warning( + "_tensor_to_object size: %s hash value: %s", tensor.numel(), hash + ) + tensor = tensor.cpu() + buf = tensor.numpy().tobytes()[:tensor_size] + return _unpickler(io.BytesIO(buf)).load() + + +@_exception_logger +def all_gather_object(object_list, obj, group=None): + """ + Gathers picklable objects from the whole group into a list. + + Similar to :func:`all_gather`, but Python objects can be passed in. + Note that the object must be picklable in order to be gathered. + + Args: + object_list (list[Any]): Output list. It should be correctly sized as the + size of the group for this collective and will contain the output. + obj (Any): Pickable Python object to be broadcast from current process. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. Default is ``None``. + + Returns: + None. If the calling rank is part of this group, the output of the + collective will be populated into the input ``object_list``. If the + calling rank is not part of the group, the passed in ``object_list`` will + be unmodified. + + .. note:: Note that this API differs slightly from the :func:`all_gather` + collective since it does not provide an ``async_op`` handle and thus + will be a blocking call. + + .. note:: For NCCL-based processed groups, internal tensor representations + of objects must be moved to the GPU device before communication takes + place. In this case, the device used is given by + ``torch.cuda.current_device()`` and it is the user's responsiblity to + ensure that this is set so that each rank has an individual GPU, via + ``torch.cuda.set_device()``. + + .. warning:: + :func:`all_gather_object` uses ``pickle`` module implicitly, which is + known to be insecure. It is possible to construct malicious pickle data + which will execute arbitrary code during unpickling. Only call this + function with data you trust. + + .. warning:: + Calling :func:`all_gather_object` with GPU tensors is not well supported + and inefficient as it incurs GPU -> CPU transfer since tensors would be + pickled. Please consider using :func:`all_gather` instead. + + Example:: + >>> # xdoctest: +SKIP("need process group init") + >>> # Note: Process group initialization omitted on each rank. + >>> import torch.distributed as dist + >>> # Assumes world_size of 3. + >>> gather_objects = ["foo", 12, {1: 2}] # any picklable object + >>> output = [None for _ in gather_objects] + >>> dist.all_gather_object(output, gather_objects[dist.get_rank()]) + >>> output + ['foo', 12, {1: 2}] + """ + if _rank_not_in_group(group): + _warn_not_in_group("all_gather_object") + return + + current_device = _get_object_coll_device(group) + input_tensor, local_size = _object_to_tensor(obj, current_device, group) + + # Gather all local sizes. This is so that we can find the max size, and index + # until the correct size when deserializing the tensors. + group_size = get_world_size(group=group) + object_sizes_tensor = torch.zeros( + group_size, dtype=torch.long, device=current_device + ) + object_size_list = [ + object_sizes_tensor[i].unsqueeze(dim=0) for i in range(group_size) + ] + # Allgather tensor sizes + all_gather(object_size_list, local_size, group=group) + max_object_size = int(max(object_size_list).item()) # type: ignore[type-var] + # Resize tensor to max size across all ranks. + input_tensor.resize_(max_object_size) + coalesced_output_tensor = torch.empty( + max_object_size * group_size, dtype=torch.uint8, device=current_device + ) + # Output tensors are nonoverlapping views of coalesced_output_tensor + output_tensors = [ + coalesced_output_tensor[max_object_size * i : max_object_size * (i + 1)] + for i in range(group_size) + ] + all_gather(output_tensors, input_tensor, group=group) + # Deserialize outputs back to object. + for i, tensor in enumerate(output_tensors): + tensor = tensor.type(torch.uint8) + tensor_size = object_size_list[i] + object_list[i] = _tensor_to_object(tensor, tensor_size, group) + + +@_exception_logger +def gather_object( + obj: Any, + object_gather_list: Optional[List[Any]] = None, + dst: Optional[int] = None, + group: Optional[ProcessGroup] = None, + group_dst: Optional[int] = None, +): + """ + Gathers picklable objects from the whole group in a single process. + + Similar to :func:`gather`, but Python objects can be passed in. Note that the + object must be picklable in order to be gathered. + + Args: + obj (Any): Input object. Must be picklable. + object_gather_list (list[Any]): Output list. On the ``dst`` rank, it + should be correctly sized as the size of the group for this + collective and will contain the output. Must be ``None`` on non-dst + ranks. (default is ``None``) + dst (int, optional): Destination rank on global process group (regardless of ``group`` argument). + (If both ``dst`` and ``group_dst`` are None, default is global rank 0) + group: (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. Default is ``None``. + group_dst (int, optional): Destination rank on ``group``. Invalid to specify both ``dst`` and ``group_dst`` + + Returns: + None. On the ``dst`` rank, ``object_gather_list`` will contain the + output of the collective. + + .. note:: Note that this API differs slightly from the gather collective + since it does not provide an async_op handle and thus will be a blocking + call. + + .. note:: For NCCL-based processed groups, internal tensor representations + of objects must be moved to the GPU device before communication takes + place. In this case, the device used is given by + ``torch.cuda.current_device()`` and it is the user's responsiblity to + ensure that this is set so that each rank has an individual GPU, via + ``torch.cuda.set_device()``. + + .. warning:: + :func:`gather_object` uses ``pickle`` module implicitly, which is + known to be insecure. It is possible to construct malicious pickle data + which will execute arbitrary code during unpickling. Only call this + function with data you trust. + + .. warning:: + Calling :func:`gather_object` with GPU tensors is not well supported + and inefficient as it incurs GPU -> CPU transfer since tensors would be + pickled. Please consider using :func:`gather` instead. + + Example:: + >>> # xdoctest: +SKIP("need process group init") + >>> # Note: Process group initialization omitted on each rank. + >>> import torch.distributed as dist + >>> # Assumes world_size of 3. + >>> gather_objects = ["foo", 12, {1: 2}] # any picklable object + >>> output = [None for _ in gather_objects] + >>> dist.gather_object( + ... gather_objects[dist.get_rank()], + ... output if dist.get_rank() == 0 else None, + ... dst=0 + ... ) + >>> # On rank 0 + >>> output + ['foo', 12, {1: 2}] + """ + group = _group_or_default_group(group) + if dst is None and group_dst is None: + dst = 0 + global_dst = _canonicalize_group_rank(group, dst, group_dst, return_global=True) + if _rank_not_in_group(group): + _warn_not_in_group("gather_object") + return + + # Ensure object_gather_list is specified appropriately. + my_global_rank = get_rank() + _validate_output_list_for_rank(my_global_rank, global_dst, object_gather_list) + current_device = _get_object_coll_device(group) + input_tensor, local_size = _object_to_tensor(obj, current_device, group) + + # Gather all local sizes. This is so that we can find the max size, and index + # until the correct size when deserializing the tensors. + group_size = get_world_size(group=group) + object_sizes_tensor = torch.zeros( + group_size, dtype=torch.long, device=current_device + ) + object_size_list = [ + object_sizes_tensor[i].unsqueeze(dim=0) for i in range(group_size) + ] + # Allgather tensor sizes. An all-gather is needed here despite this being a + # gather, since each rank needs to broadcast a tensor of the same (maximal) + # size. + all_gather(object_size_list, local_size, group=group) + max_object_size = int(max(object_size_list).item()) # type: ignore[type-var] + # Resize tensor to max size across all ranks. + input_tensor.resize_(max_object_size) + # Avoid populating output tensors if the result won't be gathered on this rank. + if my_global_rank == global_dst: + coalesced_output_tensor = torch.empty( + max_object_size * group_size, dtype=torch.uint8, device=current_device + ) + # Output tensors are nonoverlapping views of coalesced_output_tensor + output_tensors = [ + coalesced_output_tensor[max_object_size * i : max_object_size * (i + 1)] + for i in range(group_size) + ] + # All ranks call gather with equal-sized tensors. + gather( + input_tensor, + gather_list=output_tensors if my_global_rank == global_dst else None, # type: ignore[possibly-undefined] + dst=global_dst, + group=group, + ) + if my_global_rank != global_dst: + return + + assert object_gather_list is not None, "Must provide object_gather_list on dst rank" + for i, tensor in enumerate(output_tensors): + tensor = tensor.type(torch.uint8) + tensor_size = object_size_list[i] + object_gather_list[i] = _tensor_to_object(tensor, tensor_size, group) + + +@_exception_logger +def send_object_list( + object_list: List[Any], + dst: Optional[int] = None, + group: Optional[ProcessGroup] = None, + device: Optional[torch.device] = None, + group_dst: Optional[int] = None, +): + """ + Sends picklable objects in ``object_list`` synchronously. + + Similar to :func:`send`, but Python objects can be passed in. + Note that all objects in ``object_list`` must be picklable in order to be + sent. + + Args: + object_list (List[Any]): List of input objects to sent. + Each object must be picklable. Receiver must provide lists of equal sizes. + dst (int): Destination rank to send ``object_list`` to. + Destination rank is based on global process group (regardless of ``group`` argument) + group: (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. Default is ``None``. + device (``torch.device``, optional): If not None, the objects are + serialized and converted to tensors which are moved to the + ``device`` before sending. Default is ``None``. + group_dst (int, optional): Destination rank on ``group``. + Must specify one of ``dst`` and ``group_dst`` but not both + Returns: + ``None``. + + .. note:: For NCCL-based process groups, internal tensor representations + of objects must be moved to the GPU device before communication takes + place. In this case, the device used is given by + ``torch.cuda.current_device()`` and it is the user's responsibility to + ensure that this is set so that each rank has an individual GPU, via + ``torch.cuda.set_device()``. + + .. warning:: + :func:`send_object_list` uses ``pickle`` module implicitly, which + is known to be insecure. It is possible to construct malicious pickle + data which will execute arbitrary code during unpickling. Only call this + function with data you trust. + + .. warning:: + Calling :func:`send_object_list` with GPU tensors is not well supported + and inefficient as it incurs GPU -> CPU transfer since tensors would be + pickled. Please consider using :func:`send` instead. + + Example:: + >>> # xdoctest: +SKIP("need process group init") + >>> # Note: Process group initialization omitted on each rank. + >>> import torch.distributed as dist + >>> # Assumes backend is not NCCL + >>> device = torch.device("cpu") + >>> if dist.get_rank() == 0: + >>> # Assumes world_size of 2. + >>> objects = ["foo", 12, {1: 2}] # any picklable object + >>> dist.send_object_list(objects, dst=1, device=device) + >>> else: + >>> objects = [None, None, None] + >>> dist.recv_object_list(objects, src=0, device=device) + >>> objects + ['foo', 12, {1: 2}] + """ + group = _group_or_default_group(group) + group_dst = _canonicalize_group_rank(group, dst, group_dst) + _check_not_self_rank(group, group_dst, "destination") + + if _rank_not_in_group(group): + _warn_not_in_group("send_object_list") + return + + # Current device selection. + # To preserve backwards compatibility, ``device`` is default to ``None`` + # in which case we run current logic of device selection, i.e. + # ``current_device`` is CUDA if backend is NCCL otherwise CPU device. In the + # case it is not ``None`` we move the size and object tensors to be + # sent to this device. + current_device = device or _get_object_coll_device(group) + # Serialize object_list elements to tensors on src rank. + tensor_list, size_list = zip( + *[_object_to_tensor(obj, current_device, group) for obj in object_list] + ) + object_sizes_tensor = torch.cat(size_list) + + # Send object sizes + send(object_sizes_tensor, group_dst=group_dst, group=group) + + # Concatenate and send serialized object tensors + # Note: torch.cat will do an extra memory copy to the current device, if the tensor_list + # has only one element, we can skip the copy. + if len(tensor_list) == 1: # type: ignore[possibly-undefined] + object_tensor = tensor_list[0] + else: + object_tensor = torch.cat(tensor_list) + + send(object_tensor, group_dst=group_dst, group=group) + + +@_exception_logger +def recv_object_list( + object_list: List[Any], + src: Optional[int] = None, + group: Optional[ProcessGroup] = None, + device: Optional[torch.device] = None, + group_src: Optional[int] = None, +): + """ + Receives picklable objects in ``object_list`` synchronously. + + Similar to :func:`recv`, but can receive Python objects. + + Args: + object_list (List[Any]): List of objects to receive into. + Must provide a list of sizes equal to the size of the list being sent. + src (int, optional): Source rank from which to recv ``object_list``. + Source rank is based on global process group (regardless of ``group`` argument) + Will receive from any rank if set to None. Default is ``None``. + group: (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. Default is ``None``. + device (``torch.device``, optional): If not None, receives on this device. + Default is ``None``. + group_src (int, optional): Destination rank on ``group``. Invalid to specify both ``src`` and ``group_src``. + + Returns: + Sender rank. -1 if rank is not part of the group. If rank is part of the group, + ``object_list`` will contain the sent objects from ``src`` rank. + + .. note:: For NCCL-based process groups, internal tensor representations + of objects must be moved to the GPU device before communication takes + place. In this case, the device used is given by + ``torch.cuda.current_device()`` and it is the user's responsibility to + ensure that this is set so that each rank has an individual GPU, via + ``torch.cuda.set_device()``. + + .. warning:: + :func:`recv_object_list` uses ``pickle`` module implicitly, which + is known to be insecure. It is possible to construct malicious pickle + data which will execute arbitrary code during unpickling. Only call this + function with data you trust. + + .. warning:: + Calling :func:`recv_object_list` with GPU tensors is not well supported + and inefficient as it incurs GPU -> CPU transfer since tensors would be + pickled. Please consider using :func:`recv` instead. + + Example:: + >>> # xdoctest: +SKIP("need process group init") + >>> # Note: Process group initialization omitted on each rank. + >>> import torch.distributed as dist + >>> # Assumes backend is not NCCL + >>> device = torch.device("cpu") + >>> if dist.get_rank() == 0: + >>> # Assumes world_size of 2. + >>> objects = ["foo", 12, {1: 2}] # any picklable object + >>> dist.send_object_list(objects, dst=1, device=device) + >>> else: + >>> objects = [None, None, None] + >>> dist.recv_object_list(objects, src=0, device=device) + >>> objects + ['foo', 12, {1: 2}] + """ + if _rank_not_in_group(group): + _warn_not_in_group("recv_object_list") + return -1 + + # Current device selection. + # To preserve backwards compatibility, ``device`` is default to ``None`` + # in which case we run current logic of device selection, i.e. + # ``current_device`` is CUDA if backend is NCCL otherwise CPU device. In the + # case it is not ``None`` we move the size and object tensors to be + # received to this device. + current_device = device or _get_object_coll_device(group) + object_sizes_tensor = torch.empty( + len(object_list), dtype=torch.long, device=current_device + ) + + # Receive object sizes + rank_sizes = recv(object_sizes_tensor, src=src, group=group, group_src=group_src) + + # Tensor to receive serialized objects into. + object_tensor = torch.empty( # type: ignore[call-overload] + torch.sum(object_sizes_tensor).item(), # type: ignore[arg-type] + dtype=torch.uint8, + device=current_device, + ) + + rank_objects = recv(object_tensor, src=src, group=group, group_src=group_src) + assert ( + rank_sizes == rank_objects + ), "Mismatch in return ranks for object sizes and objects." + # Deserialize objects using their stored sizes. + offset = 0 + for i, obj_size in enumerate(object_sizes_tensor): + obj_view = object_tensor[offset : offset + obj_size] + obj_view = obj_view.type(torch.uint8) + offset += obj_size + object_list[i] = _tensor_to_object(obj_view, obj_size, group) + return rank_objects + + +@_exception_logger +def broadcast_object_list( + object_list: List[Any], + src: Optional[int] = None, + group: Optional[ProcessGroup] = None, + device: Optional[torch.device] = None, + group_src: Optional[int] = None, +): + """ + Broadcasts picklable objects in ``object_list`` to the whole group. + + Similar to :func:`broadcast`, but Python objects can be passed in. + Note that all objects in ``object_list`` must be picklable in order to be + broadcasted. + + Args: + object_list (List[Any]): List of input objects to broadcast. + Each object must be picklable. Only objects on the ``src`` rank will + be broadcast, but each rank must provide lists of equal sizes. + src (int): Source rank from which to broadcast ``object_list``. + Source rank is based on global process group (regardless of ``group`` argument) + group: (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. Default is ``None``. + device (``torch.device``, optional): If not None, the objects are + serialized and converted to tensors which are moved to the + ``device`` before broadcasting. Default is ``None``. + group_src (int): Source rank on ``group``. Must not specify one of ``group_src`` + and ``src`` but not both. + + Returns: + ``None``. If rank is part of the group, ``object_list`` will contain the + broadcasted objects from ``src`` rank. + + .. note:: For NCCL-based process groups, internal tensor representations + of objects must be moved to the GPU device before communication takes + place. In this case, the device used is given by + ``torch.cuda.current_device()`` and it is the user's responsibility to + ensure that this is set so that each rank has an individual GPU, via + ``torch.cuda.set_device()``. + + .. note:: Note that this API differs slightly from the :func:`broadcast` + collective since it does not provide an ``async_op`` handle and thus + will be a blocking call. + + .. warning:: + :func:`broadcast_object_list` uses ``pickle`` module implicitly, which + is known to be insecure. It is possible to construct malicious pickle + data which will execute arbitrary code during unpickling. Only call this + function with data you trust. + + .. warning:: + Calling :func:`broadcast_object_list` with GPU tensors is not well supported + and inefficient as it incurs GPU -> CPU transfer since tensors would be + pickled. Please consider using :func:`broadcast` instead. + + Example:: + >>> # xdoctest: +SKIP("need process group init") + >>> # Note: Process group initialization omitted on each rank. + >>> import torch.distributed as dist + >>> if dist.get_rank() == 0: + >>> # Assumes world_size of 3. + >>> objects = ["foo", 12, {1: 2}] # any picklable object + >>> else: + >>> objects = [None, None, None] + >>> # Assumes backend is not NCCL + >>> device = torch.device("cpu") + >>> dist.broadcast_object_list(objects, src=0, device=device) + >>> objects + ['foo', 12, {1: 2}] + """ + group = _group_or_default_group(group) + if src is None and group_src is None: + src = 0 + global_src = _canonicalize_group_rank(group, src, group_src, return_global=True) + if _rank_not_in_group(group): + _warn_not_in_group("broadcast_object_list") + return + + # Current device selection. + # To preserve backwards compatibility, ``device`` is default to ``None`` + # in which case we run current logic of device selection, i.e. + # ``current_device`` is CUDA if backend is NCCL otherwise CPU device. In the + # case it is not ``None`` we move the size and object tensors to be + # broadcasted to this device. + current_device = device or _get_object_coll_device(group) + my_global_rank = get_rank() + # Serialize object_list elements to tensors on src rank. + if my_global_rank == global_src: + tensor_list, size_list = zip( + *[_object_to_tensor(obj, current_device, group) for obj in object_list] + ) + object_sizes_tensor = torch.cat(size_list) + else: + object_sizes_tensor = torch.empty( + len(object_list), dtype=torch.long, device=current_device + ) + + # Broadcast object sizes + broadcast(object_sizes_tensor, src=global_src, group=group) + + # Concatenate and broadcast serialized object tensors + # Note: torch.cat will do an extra memory copy to the current device, if the tensor_list + # has only one element, we can skip the copy. + if my_global_rank == global_src: + if len(tensor_list) == 1: # type: ignore[possibly-undefined] + object_tensor = tensor_list[0] + else: + object_tensor = torch.cat(tensor_list) + else: + object_tensor = torch.empty( # type: ignore[call-overload] + torch.sum(object_sizes_tensor).item(), # type: ignore[arg-type] + dtype=torch.uint8, + device=current_device, + ) + + broadcast(object_tensor, src=global_src, group=group) + # Deserialize objects using their stored sizes. + offset = 0 + if my_global_rank != global_src: + for i, obj_size in enumerate(object_sizes_tensor): + obj_view = object_tensor[offset : offset + obj_size] + obj_view = obj_view.type(torch.uint8) + offset += obj_size + object_list[i] = _tensor_to_object(obj_view, obj_size, group) + + +@_exception_logger +def scatter_object_list( + scatter_object_output_list: List[Any], + scatter_object_input_list: Optional[List[Any]] = None, + src: Optional[int] = None, + group: Optional[ProcessGroup] = None, + group_src: Optional[int] = None, +): + """ + Scatters picklable objects in ``scatter_object_input_list`` to the whole group. + + Similar to :func:`scatter`, but Python objects can be passed in. On + each rank, the scattered object will be stored as the first element of + ``scatter_object_output_list``. Note that all objects in + ``scatter_object_input_list`` must be picklable in order to be scattered. + + Args: + scatter_object_output_list (List[Any]): Non-empty list whose first + element will store the object scattered to this rank. + scatter_object_input_list (List[Any], optional): List of input objects to scatter. + Each object must be picklable. Only objects on the ``src`` rank will + be scattered, and the argument can be ``None`` for non-src ranks. + src (int): Source rank from which to scatter ``scatter_object_input_list``. + Source rank is based on global process group (regardless of ``group`` argument). + (If both ``src`` and ``group_src`` are None, default is global rank 0) + group: (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. Default is ``None``. + group_src (int, optional): Source rank on ``group``. Invalid to specify both ``src`` and ``group_src`` + + Returns: + ``None``. If rank is part of the group, ``scatter_object_output_list`` + will have its first element set to the scattered object for this rank. + + .. note:: Note that this API differs slightly from the scatter collective + since it does not provide an ``async_op`` handle and thus will be a + blocking call. + + .. warning:: + :func:`scatter_object_list` uses ``pickle`` module implicitly, which + is known to be insecure. It is possible to construct malicious pickle + data which will execute arbitrary code during unpickling. Only call this + function with data you trust. + + .. warning:: + Calling :func:`scatter_object_list` with GPU tensors is not well supported + and inefficient as it incurs GPU -> CPU transfer since tensors would be + pickled. Please consider using :func:`scatter` instead. + + Example:: + >>> # xdoctest: +SKIP("need process group init") + >>> # Note: Process group initialization omitted on each rank. + >>> import torch.distributed as dist + >>> if dist.get_rank() == 0: + >>> # Assumes world_size of 3. + >>> objects = ["foo", 12, {1: 2}] # any picklable object + >>> else: + >>> # Can be any list on non-src ranks, elements are not used. + >>> objects = [None, None, None] + >>> output_list = [None] + >>> dist.scatter_object_list(output_list, objects, src=0) + >>> # Rank i gets objects[i]. For example, on rank 2: + >>> output_list + [{1: 2}] + """ + group = _group_or_default_group(group) + if src is None and group_src is None: + src = 0 + global_src = _canonicalize_group_rank(group, src, group_src, return_global=True) + if _rank_not_in_group(group): + _warn_not_in_group("scatter_object_list") + return + + if ( + not isinstance(scatter_object_output_list, list) + or len(scatter_object_output_list) < 1 + ): + raise ValueError( + "Expected argument scatter_object_output_list to be a list of size at least 1." + ) + + my_global_rank = get_rank() + pg_device = _get_object_coll_device(group) + if my_global_rank == global_src: + if scatter_object_input_list is None: + raise ValueError( + "source rank must provide non-None scatter_object_input_list" + ) + tensor_list, tensor_sizes = zip( + *[ + _object_to_tensor(obj, pg_device, group) + for obj in scatter_object_input_list + ] + ) + tensor_list, tensor_sizes = list(tensor_list), list(tensor_sizes) + + # Src rank broadcasts the maximum tensor size. This is because all ranks are + # expected to call into scatter() with equal-sized tensors. + max_tensor_size = max(tensor_sizes) # type: ignore[possibly-undefined] + for tensor in tensor_list: # type: ignore[possibly-undefined] + tensor.resize_(max_tensor_size) + else: + max_tensor_size = torch.tensor([0], dtype=torch.long, device=pg_device) + broadcast(max_tensor_size, src=global_src, group=group) + + # Scatter actual serialized objects + output_tensor = torch.empty( + max_tensor_size.item(), dtype=torch.uint8, device=pg_device + ) + scatter( + output_tensor, + scatter_list=None if my_global_rank != global_src else tensor_list, # type: ignore[possibly-undefined] + src=global_src, + group=group, + ) + + # Scatter per-object sizes to trim tensors when deserializing back to object + obj_tensor_size = torch.tensor([0], dtype=torch.long, device=pg_device) + scatter( + obj_tensor_size, + scatter_list=None if my_global_rank != global_src else tensor_sizes, # type: ignore[possibly-undefined] + src=global_src, + group=group, + ) + + # Deserialize back to object + scatter_object_output_list[0] = _tensor_to_object( + output_tensor, obj_tensor_size, group + ) + + +@_exception_logger +def all_gather(tensor_list, tensor, group=None, async_op=False): + """ + Gathers tensors from the whole group in a list. + + Complex and uneven sized tensors are supported. + + Args: + tensor_list (list[Tensor]): Output list. It should contain + correctly-sized tensors to be used for output of the collective. + Uneven sized tensors are supported. + tensor (Tensor): Tensor to be broadcast from current process. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + Examples: + >>> # xdoctest: +SKIP("need process group init") + >>> # All tensors below are of torch.int64 dtype. + >>> # We have 2 process groups, 2 ranks. + >>> device = torch.device(f'cuda:{rank}') + >>> tensor_list = [torch.zeros(2, dtype=torch.int64, device=device) for _ in range(2)] + >>> tensor_list + [tensor([0, 0], device='cuda:0'), tensor([0, 0], device='cuda:0')] # Rank 0 + [tensor([0, 0], device='cuda:1'), tensor([0, 0], device='cuda:1')] # Rank 1 + >>> tensor = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank + >>> tensor + tensor([1, 2], device='cuda:0') # Rank 0 + tensor([3, 4], device='cuda:1') # Rank 1 + >>> dist.all_gather(tensor_list, tensor) + >>> tensor_list + [tensor([1, 2], device='cuda:0'), tensor([3, 4], device='cuda:0')] # Rank 0 + [tensor([1, 2], device='cuda:1'), tensor([3, 4], device='cuda:1')] # Rank 1 + + >>> # All tensors below are of torch.cfloat dtype. + >>> # We have 2 process groups, 2 ranks. + >>> tensor_list = [torch.zeros(2, dtype=torch.cfloat, device=device) for _ in range(2)] + >>> tensor_list + [tensor([0.+0.j, 0.+0.j], device='cuda:0'), tensor([0.+0.j, 0.+0.j], device='cuda:0')] # Rank 0 + [tensor([0.+0.j, 0.+0.j], device='cuda:1'), tensor([0.+0.j, 0.+0.j], device='cuda:1')] # Rank 1 + >>> tensor = torch.tensor([1+1j, 2+2j], dtype=torch.cfloat, device=device) + 2 * rank * (1+1j) + >>> tensor + tensor([1.+1.j, 2.+2.j], device='cuda:0') # Rank 0 + tensor([3.+3.j, 4.+4.j], device='cuda:1') # Rank 1 + >>> dist.all_gather(tensor_list, tensor) + >>> tensor_list + [tensor([1.+1.j, 2.+2.j], device='cuda:0'), tensor([3.+3.j, 4.+4.j], device='cuda:0')] # Rank 0 + [tensor([1.+1.j, 2.+2.j], device='cuda:1'), tensor([3.+3.j, 4.+4.j], device='cuda:1')] # Rank 1 + + """ + _check_tensor_list(tensor_list, "tensor_list") + _check_single_tensor(tensor, "tensor") + _ensure_all_tensors_same_dtype(tensor_list, tensor) + if _rank_not_in_group(group): + _warn_not_in_group("all_gather") + return + + tensor_list = [ + t if not t.is_complex() else torch.view_as_real(t) for t in tensor_list + ] + tensor = tensor if not tensor.is_complex() else torch.view_as_real(tensor) + + group = group or _get_default_group() + work = group.allgather([tensor_list], [tensor]) + + if async_op: + return work + else: + work.wait() + + +@_exception_logger +def all_gather_into_tensor(output_tensor, input_tensor, group=None, async_op=False): + """ + Gather tensors from all ranks and put them in a single output tensor. + + This function requires all tensors to be the same size on each process. + + Args: + output_tensor (Tensor): Output tensor to accommodate tensor elements + from all ranks. It must be correctly sized to have one of the + following forms: + (i) a concatenation of all the input tensors along the primary + dimension; for definition of "concatenation", see ``torch.cat()``; + (ii) a stack of all the input tensors along the primary dimension; + for definition of "stack", see ``torch.stack()``. + Examples below may better explain the supported output forms. + input_tensor (Tensor): Tensor to be gathered from current rank. + Different from the ``all_gather`` API, the input tensors in this + API must have the same size across all ranks. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + Examples: + >>> # xdoctest: +SKIP("need process group init") + >>> # All tensors below are of torch.int64 dtype and on CUDA devices. + >>> # We have two ranks. + >>> device = torch.device(f'cuda:{rank}') + >>> tensor_in = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank + >>> tensor_in + tensor([1, 2], device='cuda:0') # Rank 0 + tensor([3, 4], device='cuda:1') # Rank 1 + >>> # Output in concatenation form + >>> tensor_out = torch.zeros(world_size * 2, dtype=torch.int64, device=device) + >>> dist.all_gather_into_tensor(tensor_out, tensor_in) + >>> tensor_out + tensor([1, 2, 3, 4], device='cuda:0') # Rank 0 + tensor([1, 2, 3, 4], device='cuda:1') # Rank 1 + >>> # Output in stack form + >>> tensor_out2 = torch.zeros(world_size, 2, dtype=torch.int64, device=device) + >>> dist.all_gather_into_tensor(tensor_out2, tensor_in) + >>> tensor_out2 + tensor([[1, 2], + [3, 4]], device='cuda:0') # Rank 0 + tensor([[1, 2], + [3, 4]], device='cuda:1') # Rank 1 + + .. warning:: + The Gloo backend does not support this API. + + """ + _check_single_tensor(input_tensor, "input_tensor") + _check_single_tensor(output_tensor, "output_tensor") + if _rank_not_in_group(group): + _warn_not_in_group("all_gather_into_tensor") + return + + output_tensor = ( + output_tensor + if not output_tensor.is_complex() + else torch.view_as_real(output_tensor) + ) + input_tensor = ( + input_tensor + if not input_tensor.is_complex() + else torch.view_as_real(input_tensor) + ) + + opts = AllgatherOptions() + opts.asyncOp = async_op + + group = group or _get_default_group() + + if group in _world.pg_coalesce_state.keys(): + # We are in coalescing context, do not issue single operation, just append a collective representation + coll = _CollOp(all_gather_into_tensor, input_tensor, output_tensor) + _world.pg_coalesce_state[group].append(coll) + if async_op: + return _IllegalWork() + else: + return None + + work = group._allgather_base(output_tensor, input_tensor, opts) + + if async_op: + return work + else: + work.wait() + + +@_exception_logger +@deprecated( + "`torch.distributed._all_gather_base` is a private function and will be deprecated. " + "Please use `torch.distributed.all_gather_into_tensor` instead.", + category=FutureWarning, +) +def _all_gather_base(output_tensor, input_tensor, group=None, async_op=False): + """ + Single tensor all gather. Gathers a single tensor from all ranks, and puts them in a single output tensor. + + Args: + output_tensor (Tensor): Output tensor. It should contain + correctly-sized tensors to be used for output of the collective. + input_tensor (Tensor): Tensor to be broadcast from current process. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + .. warning:: + `_all_gather_base` is a private function. Users should use + `all_gather_into_tensor` instead. + + """ + return all_gather_into_tensor(output_tensor, input_tensor, group, async_op) + + +@_exception_logger +@deprecated( + "`torch.distributed.all_gather_coalesced` will be deprecated. If you must use it, " + "please revisit our documentation later at " + "https://pytorch.org/docs/main/distributed.html#collective-functions", + category=FutureWarning, +) +def all_gather_coalesced( + output_tensor_lists, input_tensor_list, group=None, async_op=False +): + """ + Gathers input tensors from the whole group in a list in a coalesced manner. + + Complex tensors are supported. + + Args: + output_tensor_lists (list[list[Tensor]]): Output list. It should contain + correctly-sized tensors to be used for output of the collective. + input_tensor_list (list[Tensor]): Tensors to be broadcast from + current process. At least one tensor has to be non empty. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + Example: + we have 2 process groups, 2 ranks. + rank 0 passes: + input_tensor_list = [[[1, 1], [1, 1]], [2], [3, 3]] + output_tensor_lists = + [[[[-1, -1], [-1, -1]], [-1], [-1, -1]], + [[[-1, -1], [-1, -1]], [-1], [-1, -1]]] + rank 1 passes: + input_tensor_list = [[[3, 3], [3, 3]], [5], [1, 1]] + output_tensor_lists = + [[[[-1, -1], [-1, -1]], [-1], [-1, -1]], + [[[-1, -1], [-1, -1]], [-1], [-1, -1]]] + both rank 0 and 1 get: + output_tensor_lists = + [[[1, 1], [1, 1]], [2], [3, 3]], + [[3, 3], [3, 3]], [5], [1, 1]]]. + + WARNING: at this time individual shape checking is not implemented across nodes. + For example, if the rank 0 node passes [torch.rand(4), torch.rand(2)] and the + rank 1 node passes [torch.rand(2), torch.rand(2), torch.rand(2)], the + all_gather_coalesced operation will proceed without complaint and return + erroneous outputs. This lack of shape checking results in significant + performance improvements but users of this function should take extra care + to ensure that each node passes in tensors whose shapes match across nodes. + """ + # We only check basic compatibility with C++ params here, C++ code will + # do shape and type checking. + if _rank_not_in_group(group): + _warn_not_in_group("all_gather_coalesced") + return + _check_tensor_list(input_tensor_list, "input_tensor_list") + _ensure_all_tensors_same_dtype(input_tensor_list) + if not isinstance(output_tensor_lists, list): + raise TypeError( + "Invalid function argument: output_tensor_lists should be a list" + ) + for output_tensor_list in output_tensor_lists: + _check_tensor_list(output_tensor_list, "output_tensor_lists") + _ensure_all_tensors_same_dtype(output_tensor_list) + + output_tensor_lists = [ + [t if not t.is_complex() else torch.view_as_real(t) for t in l] + for l in output_tensor_lists + ] + input_tensor_list = [ + t if not t.is_complex() else torch.view_as_real(t) for t in input_tensor_list + ] + + group = group or _get_default_group() + work = group.allgather_coalesced(output_tensor_lists, input_tensor_list) + + if async_op: + return work.get_future() + else: + work.wait() + + +def _validate_output_list_for_rank(my_rank, dst, gather_list): + if dst == my_rank: + if not gather_list: + raise ValueError( + "Argument ``gather_list`` must be specified on destination rank." + ) + elif gather_list: + raise ValueError( + "Argument ``gather_list`` must NOT be specified " + "on non-destination ranks." + ) + + +@_exception_logger +def gather( + tensor: torch.Tensor, + gather_list: Optional[List[torch.Tensor]] = None, + dst: Optional[int] = None, + group: Optional[ProcessGroup] = None, + async_op: bool = False, + group_dst: Optional[int] = None, +): + """ + Gathers a list of tensors in a single process. + + This function requires all tensors to be the same size on each process. + + Args: + tensor (Tensor): Input tensor. + gather_list (list[Tensor], optional): List of appropriately, + same-sized tensors to use for gathered data + (default is None, must be specified on the destination rank) + dst (int, optional): Destination rank on global process group (regardless of ``group`` argument). + (If both ``dst`` and ``group_dst`` are None, default is global rank 0) + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op + group_dst (int, optional): Destination rank on ``group``. Invalid to specify both ``dst`` and ``group_dst`` + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + .. note:: Note that all Tensors in gather_list must have the same size. + + Example:: + >>> # xdoctest: +SKIP("no rank") + >>> # We have 2 process groups, 2 ranks. + >>> tensor_size = 2 + >>> device = torch.device(f'cuda:{rank}') + >>> tensor = torch.ones(tensor_size, device=device) + rank + >>> if dist.get_rank() == 0: + >>> gather_list = [torch.zeros_like(tensor, device=device) for i in range(2)] + >>> else: + >>> gather_list = None + >>> dist.gather(tensor, gather_list, dst=0) + >>> # Rank 0 gets gathered data. + >>> gather_list + [tensor([1., 1.], device='cuda:0'), tensor([2., 2.], device='cuda:0')] # Rank 0 + None # Rank 1 + + """ + _check_single_tensor(tensor, "tensor") + + # Parameter ``gather_list`` may be left unspecified on non-dst ranks. + if gather_list: + _check_tensor_list(gather_list, "gather_list") + else: + gather_list = [] + _ensure_all_tensors_same_dtype(tensor, gather_list) + group = _group_or_default_group(group) + if _rank_not_in_group(group): + _warn_not_in_group("gather") + return + if dst is None and group_dst is None: + dst = 0 + global_dst = _canonicalize_group_rank(group, dst, group_dst, return_global=True) + group_dst = _canonicalize_group_rank(group, dst, group_dst, return_global=False) + my_global_rank = get_rank() + _validate_output_list_for_rank(my_global_rank, global_dst, gather_list) + output_tensors = [gather_list] if global_dst == my_global_rank else [] + input_tensors = [tensor] + + opts = GatherOptions() + opts.rootRank = group_dst + work = group.gather(output_tensors, input_tensors, opts) + + if async_op: + return work + else: + work.wait() + + +@_exception_logger +def scatter( + tensor: torch.Tensor, + scatter_list: Optional[List[torch.Tensor]] = None, + src: Optional[int] = None, + group: Optional[ProcessGroup] = None, + async_op: bool = False, + group_src: Optional[int] = None, +): + """ + Scatters a list of tensors to all processes in a group. + + Each process will receive exactly one tensor and store its data in the + ``tensor`` argument. + + Complex tensors are supported. + + Args: + tensor (Tensor): Output tensor. + scatter_list (list[Tensor]): List of tensors to scatter (default is + None, must be specified on the source rank) + src (int): Source rank on global process group (regardless of ``group`` argument). + (If both ``src`` and ``group_src`` are None, default is global rank 0) + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op + group_src (int, optional): Source rank on ``group``. Invalid to specify both ``src`` and ``group_src`` + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + .. note:: Note that all Tensors in scatter_list must have the same size. + + Example:: + >>> # xdoctest: +SKIP("need process group init") + >>> # Note: Process group initialization omitted on each rank. + >>> import torch.distributed as dist + >>> tensor_size = 2 + >>> device = torch.device(f'cuda:{rank}') + >>> output_tensor = torch.zeros(tensor_size, device=device) + >>> if dist.get_rank() == 0: + >>> # Assumes world_size of 2. + >>> # Only tensors, all of which must be the same size. + >>> t_ones = torch.ones(tensor_size, device=device) + >>> t_fives = torch.ones(tensor_size, device=device) * 5 + >>> scatter_list = [t_ones, t_fives] + >>> else: + >>> scatter_list = None + >>> dist.scatter(output_tensor, scatter_list, src=0) + >>> # Rank i gets scatter_list[i]. + >>> output_tensor + tensor([1., 1.], device='cuda:0') # Rank 0 + tensor([5., 5.], device='cuda:1') # Rank 1 + + """ + _check_single_tensor(tensor, "tensor") + # Parameter ``scatter_list`` may be left unspecified on non-src ranks. + if scatter_list: + _check_tensor_list(scatter_list, "scatter_list") + else: + scatter_list = [] + _ensure_all_tensors_same_dtype(tensor, scatter_list) + group = _group_or_default_group(group) + if src is None and group_src is None: + src = 0 + global_src = _canonicalize_group_rank(group, src, group_src, return_global=True) + group_src = _canonicalize_group_rank(group, src, group_src, return_global=False) + if _rank_not_in_group(group): + _warn_not_in_group("scatter") + return + scatter_list = [ + t if not t.is_complex() else torch.view_as_real(t) for t in scatter_list + ] + tensor = tensor if not tensor.is_complex() else torch.view_as_real(tensor) + + my_global_rank = get_rank() + if global_src == my_global_rank: + if not scatter_list: + raise ValueError( + "Argument ``scatter_list`` must be specified on source rank." + ) + input_tensors = [scatter_list] + output_tensors = [tensor] + else: + if scatter_list: + raise ValueError( + "Argument ``scatter_list`` must NOT be specified " + "on non-source ranks." + ) + input_tensors = [] + output_tensors = [tensor] + + opts = ScatterOptions() + opts.rootRank = group_src + opts.asyncOp = async_op + work = group.scatter(output_tensors, input_tensors, opts) + + if async_op: + return work + else: + work.wait() + + +@_exception_logger +def reduce_scatter(output, input_list, op=ReduceOp.SUM, group=None, async_op=False): + """ + Reduces, then scatters a list of tensors to all processes in a group. + + Args: + output (Tensor): Output tensor. + input_list (list[Tensor]): List of tensors to reduce and scatter. + op (optional): One of the values from + ``torch.distributed.ReduceOp`` + enum. Specifies an operation used for element-wise reductions. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group. + + """ + _check_single_tensor(output, "output") + _check_tensor_list(input_list, "input_list") + _ensure_all_tensors_same_dtype(output, input_list) + if _rank_not_in_group(group): + _warn_not_in_group("reduce_scatter") + return + + opts = ReduceScatterOptions() + opts.reduceOp = op + + group = group or _get_default_group() + work = group.reduce_scatter([output], [input_list], opts) + + if async_op: + return work + else: + work.wait() + + +@_exception_logger +def reduce_scatter_tensor(output, input, op=ReduceOp.SUM, group=None, async_op=False): + """ + Reduces, then scatters a tensor to all ranks in a group. + + Args: + output (Tensor): Output tensor. It should have the same size across all + ranks. + input (Tensor): Input tensor to be reduced and scattered. Its size + should be output tensor size times the world size. The input tensor + can have one of the following shapes: + (i) a concatenation of the output tensors along the primary + dimension, or + (ii) a stack of the output tensors along the primary dimension. + For definition of "concatenation", see ``torch.cat()``. + For definition of "stack", see ``torch.stack()``. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group. + + Examples: + >>> # xdoctest: +SKIP("need process group init") + >>> # All tensors below are of torch.int64 dtype and on CUDA devices. + >>> # We have two ranks. + >>> device = torch.device(f'cuda:{rank}') + >>> tensor_out = torch.zeros(2, dtype=torch.int64, device=device) + >>> # Input in concatenation form + >>> tensor_in = torch.arange(world_size * 2, dtype=torch.int64, device=device) + >>> tensor_in + tensor([0, 1, 2, 3], device='cuda:0') # Rank 0 + tensor([0, 1, 2, 3], device='cuda:1') # Rank 1 + >>> dist.reduce_scatter_tensor(tensor_out, tensor_in) + >>> tensor_out + tensor([0, 2], device='cuda:0') # Rank 0 + tensor([4, 6], device='cuda:1') # Rank 1 + >>> # Input in stack form + >>> tensor_in = torch.reshape(tensor_in, (world_size, 2)) + >>> tensor_in + tensor([[0, 1], + [2, 3]], device='cuda:0') # Rank 0 + tensor([[0, 1], + [2, 3]], device='cuda:1') # Rank 1 + >>> dist.reduce_scatter_tensor(tensor_out, tensor_in) + >>> tensor_out + tensor([0, 2], device='cuda:0') # Rank 0 + tensor([4, 6], device='cuda:1') # Rank 1 + + .. warning:: + The Gloo backend does not support this API. + + """ + _check_single_tensor(output, "output") + _check_single_tensor(input, "input") + + if _rank_not_in_group(group): + _warn_not_in_group("reduce_scatter_tensor") + return + + opts = ReduceScatterOptions() + opts.reduceOp = op + opts.asyncOp = async_op + + group = group or _get_default_group() + + # Check if we are in coalescing context + # If we are, do not issue single operation, just append a collective representation + if group in _world.pg_coalesce_state.keys(): + coll = _CollOp(reduce_scatter_tensor, input, output, op, None) + _world.pg_coalesce_state[group].append(coll) + if async_op: + return _IllegalWork() + else: + return None + + work = group._reduce_scatter_base(output, input, opts) + + if async_op: + return work + else: + work.wait() + + +@deprecated( + "`torch.distributed._reduce_scatter_base` is a private function and will be deprecated. " + "Please use `torch.distributed.reduce_scatter_tensor` instead.", + category=FutureWarning, +) +def _reduce_scatter_base(output, input, op=ReduceOp.SUM, group=None, async_op=False): + """ + Reduces, then scatters a flattened tensor to all processes in a group. + + Args: + output (Tensor): Output tensor. + input (Tensor): Input tensor that is of size output tensor size times world size + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group. + + .. warning:: + `_reduce_scatter_base` is a private function. Users should use + `reduce_scatter_tensor` instead. + + """ + return reduce_scatter_tensor(output, input, op, group, async_op) + + +@_exception_logger +def all_to_all_single( + output, + input, + output_split_sizes=None, + input_split_sizes=None, + group=None, + async_op=False, +): + """ + Split input tensor and then scatter the split list to all processes in a group. + + Later the received tensors are concatenated from all the processes in the group + and returned as a single output tensor. + + Complex tensors are supported. + + Args: + output (Tensor): Gathered concatenated output tensor. + input (Tensor): Input tensor to scatter. + output_split_sizes: (list[Int], optional): Output split sizes for dim 0 + if specified None or empty, dim 0 of ``output`` tensor must divide + equally by ``world_size``. + input_split_sizes: (list[Int], optional): Input split sizes for dim 0 + if specified None or empty, dim 0 of ``input`` tensor must divide + equally by ``world_size``. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group. + + .. warning:: + `all_to_all_single` is experimental and subject to change. + + Examples: + >>> # xdoctest: +SKIP("Undefined rank") + >>> input = torch.arange(4) + rank * 4 + >>> input + tensor([0, 1, 2, 3]) # Rank 0 + tensor([4, 5, 6, 7]) # Rank 1 + tensor([8, 9, 10, 11]) # Rank 2 + tensor([12, 13, 14, 15]) # Rank 3 + >>> output = torch.empty([4], dtype=torch.int64) + >>> dist.all_to_all_single(output, input) + >>> output + tensor([0, 4, 8, 12]) # Rank 0 + tensor([1, 5, 9, 13]) # Rank 1 + tensor([2, 6, 10, 14]) # Rank 2 + tensor([3, 7, 11, 15]) # Rank 3 + + >>> # Essentially, it is similar to following operation: + >>> scatter_list = list(input.chunk(world_size)) + >>> gather_list = list(output.chunk(world_size)) + >>> for i in range(world_size): + >>> dist.scatter(gather_list[i], scatter_list if i == rank else [], src = i) + + >>> # Another example with uneven split + >>> input + tensor([0, 1, 2, 3, 4, 5]) # Rank 0 + tensor([10, 11, 12, 13, 14, 15, 16, 17, 18]) # Rank 1 + tensor([20, 21, 22, 23, 24]) # Rank 2 + tensor([30, 31, 32, 33, 34, 35, 36]) # Rank 3 + >>> input_splits + [2, 2, 1, 1] # Rank 0 + [3, 2, 2, 2] # Rank 1 + [2, 1, 1, 1] # Rank 2 + [2, 2, 2, 1] # Rank 3 + >>> output_splits + [2, 3, 2, 2] # Rank 0 + [2, 2, 1, 2] # Rank 1 + [1, 2, 1, 2] # Rank 2 + [1, 2, 1, 1] # Rank 3 + >>> output = ... + >>> dist.all_to_all_single(output, input, output_splits, input_splits) + >>> output + tensor([ 0, 1, 10, 11, 12, 20, 21, 30, 31]) # Rank 0 + tensor([ 2, 3, 13, 14, 22, 32, 33]) # Rank 1 + tensor([ 4, 15, 16, 23, 34, 35]) # Rank 2 + tensor([ 5, 17, 18, 24, 36]) # Rank 3 + + + >>> # Another example with tensors of torch.cfloat type. + >>> input = torch.tensor([1+1j, 2+2j, 3+3j, 4+4j], dtype=torch.cfloat) + 4 * rank * (1+1j) + >>> input + tensor([1+1j, 2+2j, 3+3j, 4+4j]) # Rank 0 + tensor([5+5j, 6+6j, 7+7j, 8+8j]) # Rank 1 + tensor([9+9j, 10+10j, 11+11j, 12+12j]) # Rank 2 + tensor([13+13j, 14+14j, 15+15j, 16+16j]) # Rank 3 + >>> output = torch.empty([4], dtype=torch.int64) + >>> dist.all_to_all_single(output, input) + >>> output + tensor([1+1j, 5+5j, 9+9j, 13+13j]) # Rank 0 + tensor([2+2j, 6+6j, 10+10j, 14+14j]) # Rank 1 + tensor([3+3j, 7+7j, 11+11j, 15+15j]) # Rank 2 + tensor([4+4j, 8+8j, 12+12j, 16+16j]) # Rank 3 + """ + if _rank_not_in_group(group): + _warn_not_in_group("all_to_all_single") + return + + opts = AllToAllOptions() + _check_single_tensor(output, "output") + _check_single_tensor(input, "input") + _ensure_all_tensors_same_dtype(output, input) + + if input.is_complex(): + input = torch.view_as_real(input) + if output.is_complex(): + output = torch.view_as_real(output) + + output_split_sizes = [] if output_split_sizes is None else output_split_sizes + input_split_sizes = [] if input_split_sizes is None else input_split_sizes + + group = group or _get_default_group() + work = group.alltoall_base( + output, input, output_split_sizes, input_split_sizes, opts + ) + + if async_op: + return work + else: + work.wait() + + +@_exception_logger +def all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False): + """ + Scatters list of input tensors to all processes in a group and return gathered list of tensors in output list. + + Complex tensors are supported. + + Args: + output_tensor_list (list[Tensor]): List of tensors to be gathered one + per rank. + input_tensor_list (list[Tensor]): List of tensors to scatter one per rank. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group. + + .. warning:: + `all_to_all` is experimental and subject to change. + + Examples: + >>> # xdoctest: +SKIP("Undefined rank") + >>> input = torch.arange(4) + rank * 4 + >>> input = list(input.chunk(4)) + >>> input + [tensor([0]), tensor([1]), tensor([2]), tensor([3])] # Rank 0 + [tensor([4]), tensor([5]), tensor([6]), tensor([7])] # Rank 1 + [tensor([8]), tensor([9]), tensor([10]), tensor([11])] # Rank 2 + [tensor([12]), tensor([13]), tensor([14]), tensor([15])] # Rank 3 + >>> output = list(torch.empty([4], dtype=torch.int64).chunk(4)) + >>> dist.all_to_all(output, input) + >>> output + [tensor([0]), tensor([4]), tensor([8]), tensor([12])] # Rank 0 + [tensor([1]), tensor([5]), tensor([9]), tensor([13])] # Rank 1 + [tensor([2]), tensor([6]), tensor([10]), tensor([14])] # Rank 2 + [tensor([3]), tensor([7]), tensor([11]), tensor([15])] # Rank 3 + + >>> # Essentially, it is similar to following operation: + >>> scatter_list = input + >>> gather_list = output + >>> for i in range(world_size): + >>> dist.scatter(gather_list[i], scatter_list if i == rank else [], src=i) + + >>> input + tensor([0, 1, 2, 3, 4, 5]) # Rank 0 + tensor([10, 11, 12, 13, 14, 15, 16, 17, 18]) # Rank 1 + tensor([20, 21, 22, 23, 24]) # Rank 2 + tensor([30, 31, 32, 33, 34, 35, 36]) # Rank 3 + >>> input_splits + [2, 2, 1, 1] # Rank 0 + [3, 2, 2, 2] # Rank 1 + [2, 1, 1, 1] # Rank 2 + [2, 2, 2, 1] # Rank 3 + >>> output_splits + [2, 3, 2, 2] # Rank 0 + [2, 2, 1, 2] # Rank 1 + [1, 2, 1, 2] # Rank 2 + [1, 2, 1, 1] # Rank 3 + >>> input = list(input.split(input_splits)) + >>> input + [tensor([0, 1]), tensor([2, 3]), tensor([4]), tensor([5])] # Rank 0 + [tensor([10, 11, 12]), tensor([13, 14]), tensor([15, 16]), tensor([17, 18])] # Rank 1 + [tensor([20, 21]), tensor([22]), tensor([23]), tensor([24])] # Rank 2 + [tensor([30, 31]), tensor([32, 33]), tensor([34, 35]), tensor([36])] # Rank 3 + >>> output = ... + >>> dist.all_to_all(output, input) + >>> output + [tensor([0, 1]), tensor([10, 11, 12]), tensor([20, 21]), tensor([30, 31])] # Rank 0 + [tensor([2, 3]), tensor([13, 14]), tensor([22]), tensor([32, 33])] # Rank 1 + [tensor([4]), tensor([15, 16]), tensor([23]), tensor([34, 35])] # Rank 2 + [tensor([5]), tensor([17, 18]), tensor([24]), tensor([36])] # Rank 3 + + >>> # Another example with tensors of torch.cfloat type. + >>> input = torch.tensor([1+1j, 2+2j, 3+3j, 4+4j], dtype=torch.cfloat) + 4 * rank * (1+1j) + >>> input = list(input.chunk(4)) + >>> input + [tensor([1+1j]), tensor([2+2j]), tensor([3+3j]), tensor([4+4j])] # Rank 0 + [tensor([5+5j]), tensor([6+6j]), tensor([7+7j]), tensor([8+8j])] # Rank 1 + [tensor([9+9j]), tensor([10+10j]), tensor([11+11j]), tensor([12+12j])] # Rank 2 + [tensor([13+13j]), tensor([14+14j]), tensor([15+15j]), tensor([16+16j])] # Rank 3 + >>> output = list(torch.empty([4], dtype=torch.int64).chunk(4)) + >>> dist.all_to_all(output, input) + >>> output + [tensor([1+1j]), tensor([5+5j]), tensor([9+9j]), tensor([13+13j])] # Rank 0 + [tensor([2+2j]), tensor([6+6j]), tensor([10+10j]), tensor([14+14j])] # Rank 1 + [tensor([3+3j]), tensor([7+7j]), tensor([11+11j]), tensor([15+15j])] # Rank 2 + [tensor([4+4j]), tensor([8+8j]), tensor([12+12j]), tensor([16+16j])] # Rank 3 + + """ + if _rank_not_in_group(group): + _warn_not_in_group("all_to_all") + return + + opts = AllToAllOptions() + _check_tensor_list(output_tensor_list, "output_tensor_list") + _check_tensor_list(input_tensor_list, "input_tensor_list") + _ensure_all_tensors_same_dtype(output_tensor_list, input_tensor_list) + + input_tensor_list = [ + t if not t.is_complex() else torch.view_as_real(t) for t in input_tensor_list + ] + output_tensor_list = [ + t if not t.is_complex() else torch.view_as_real(t) for t in output_tensor_list + ] + + group = group or _get_default_group() + work = group.alltoall(output_tensor_list, input_tensor_list, opts) + + if async_op: + return work + else: + work.wait() + + +@_exception_logger +def barrier( + group: Optional[ProcessGroup] = GroupMember.WORLD, async_op=False, device_ids=None +): + """ + Synchronize all processes. + + This collective blocks processes until the whole group enters this function, + if async_op is False, or if async work handle is called on wait(). + + Args: + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + async_op (bool, optional): Whether this op should be an async op + device_ids ([int], optional): List of device/GPU ids. + + Returns: + Async work handle, if async_op is set to True. + None, if not async_op or if not part of the group + + .. note:: `ProcessGroupNCCL` now blocks the cpu thread till the completion of the barrier collective. + """ + if _rank_not_in_group(group): + _warn_not_in_group("barrier") + return + + opts = BarrierOptions() + opts.device = torch.device(_get_object_coll_device(group)) + if device_ids is not None: + if isinstance(device_ids, list): + opts.device_ids = device_ids + else: + raise TypeError( + "Invalid function argument: device_ids type should be List[int]" + ) + + group = group or _get_default_group() + work = group.barrier(opts=opts) + + if async_op: + return work + else: + work.wait() + + +def monitored_barrier( + group: Optional[ProcessGroup] = GroupMember.WORLD, + timeout=None, + wait_all_ranks=False, +): + """ + Synchronize processes similar to ``torch.distributed.barrier``, but consider a configurable timeout. + + It is able to report ranks that did not pass this barrier within the provided timeout. + Specifically, for non-zero ranks, will block until a send/recv is processed from rank 0. + Rank 0 will block until all send /recv from other ranks are processed, and will report + failures for ranks that failed to respond in time. Note that if one rank does not reach the + monitored_barrier (for example due to a hang), all other ranks would fail in monitored_barrier. + + This collective will block all processes/ranks in the group, until the + whole group exits the function successfully, making it useful for debugging + and synchronizing. However, it can have a performance impact and should only + be used for debugging or scenarios that require full synchronization points + on the host-side. For debugging purposes, this barrier can be inserted + before the application's collective calls to check if any ranks are + desynchronized. + + .. note:: Note that this collective is only supported with the GLOO backend. + + Args: + group (ProcessGroup, optional): The process group to work on. If + ``None``, the default process group will be used. + timeout (datetime.timedelta, optional): Timeout for monitored_barrier. + If ``None``, the default process group timeout will be used. + wait_all_ranks (bool, optional): Whether to collect all failed ranks or + not. By default, this is ``False`` and ``monitored_barrier`` on rank 0 + will throw on the first failed rank it encounters in order to fail + fast. By setting ``wait_all_ranks=True`` ``monitored_barrier`` will + collect all failed ranks and throw an error containing information + about all failed ranks. + + Returns: + ``None``. + + Example:: + >>> # xdoctest: +SKIP("need process group init") + >>> # Note: Process group initialization omitted on each rank. + >>> import torch.distributed as dist + >>> if dist.get_rank() != 1: + >>> dist.monitored_barrier() # Raises exception indicating that + >>> # rank 1 did not call into monitored_barrier. + >>> # Example with wait_all_ranks=True + >>> if dist.get_rank() == 0: + >>> dist.monitored_barrier(wait_all_ranks=True) # Raises exception + >>> # indicating that ranks 1, 2, ... world_size - 1 did not call into + >>> # monitored_barrier. + """ + # Need to call rank not in group before using the group, otherwise + # "Invalid process group" error is raised. + if _rank_not_in_group(group): + _warn_not_in_group("monitored_barrier") + return + + if get_backend(group) != Backend.GLOO: + raise ValueError("monitored_barrier is only implemented for GLOO backend.") + + if timeout is None: + timeout = _get_default_timeout(get_backend(group)) + elif isinstance(timeout, float): + # TODO(whc) aparently some existing test case for monitored_barrier passes in a timeout in float format? + warnings.warn( + "Please specify timeout arg as a timedelta. " + f"Converting current value of {timeout} assuming it represents seconds", + ) + timeout = timedelta(seconds=timeout) + + _check_valid_timeout(timeout) + + group_to_use = _get_default_group() if group is None else group + return group_to_use.monitored_barrier( # type:ignore[attr-defined] + timeout, wait_all_ranks=wait_all_ranks + ) + + +def _create_process_group_wrapper( + wrapped_pg: torch._C._distributed_c10d.Backend, + store_prefix: str, + store: Store, + rank: int, + world_size: int, + timeout: timedelta = default_pg_timeout, +): + assert _GLOO_AVAILABLE, "ProcessGroupWrapper unsupported without GLOO backend." + + # (whc) this appears to be just for the gloo backend? if so, `default_pg_timeout` is appropriate... + + # Create a separate prefix store for the helper process group. + prefix = f"{PG_WRAPPER_STORE_PREFIX}:{store_prefix}" + store = PrefixStore(prefix, store) + helper_pg = ProcessGroupGloo(store, rank, world_size, timeout=timeout) + # Wrap the underlying pg with ProcessGroupWrapper. + wrapped_pg = _ProcessGroupWrapper(wrapped_pg, helper_pg) + return wrapped_pg + + +# helper function for deterministically hashing a list of ranks to a unique +# string +def _hash_ranks_to_str(ranks: List[int]) -> str: + rank_join: str = "_".join(map(str, ranks)) + # In case there is already a PG with the same rank composition + unique_str = "_".join([rank_join, str(len(_world.pg_names))]) + return hashlib.sha1(bytes(unique_str, "utf-8")).hexdigest() + + +# Takes a list of ranks and computes an integer color +def _process_group_color(ranks: List[int]) -> int: + # Convert list to tuple to make it hashable + ranks = tuple(ranks) + hash_value = hash(ranks) + # Split color must be: + # - a non-negative integer; + # - a type compatible with C's int because we are pybinding to the latter. + # Thus, we limit the hash value within c_int's max value. + max_c_int = 2 ** (ctypes.sizeof(ctypes.c_int) * 8 - 1) + color = abs(hash_value) % max_c_int + return color + + +def _process_group_name(ranks, use_hashed_name): + # Create name for a process group. + global _world + if use_hashed_name: + pg_name = _hash_ranks_to_str(ranks) + else: + pg_name = str(_world.group_count) + _world.group_count += 1 + # TODO: why is group count incremented only in the else path? + return pg_name + + +def _get_backend_from_str(backend: Optional[str] = None) -> Backend: + # Default to the same backend as the global process group + # if backend is not specified. + if not backend: + backend = get_backend(_get_default_group()) + return Backend(backend) + + +def _is_safe_to_split() -> bool: + """ + Checks if it is safe to split the any process group in the world. + This is only safe if the default pg has a bound device id, otherwise + users must be aware that a pg is only splittable after the first collective is + issued. + """ + return False if _get_default_group().bound_device_id is None else True + + +@_time_logger +def split_group( + parent_pg: Optional[ProcessGroup] = None, + split_ranks: Optional[list] = None, + timeout: Optional[timedelta] = None, + pg_options: Optional[Any] = None, + group_desc: Optional[str] = None, +) -> Optional[ProcessGroup]: + """ + Create a new process group splitted from the given parent process group. + + warning:: This is an experimental API and only the ``NCCL`` backend supports this API. + Other backends will raise an error. + Users of this API must gurantee that all ranks in the parent group enter this API call, + and the split of the sub groups is the same accross all ranks in the parent group. + + Args: + parent_pg (ProcessGroup, optional): The parent process group. If None, + the default process group will be used. Users need to gurantee that + the parent group is fully initialized (e.g, communicators are initialized) + split_ranks (list[list[int]]): the split ranks, which is a list of list of ranks. + Users need to make sure the validity of the split ranks such that one + split (represented by one inner list of ints) does not overlap with any other split. + Note that the ranks in each split is the group rank (instead of global rank) + in the parent pg. For example, if the parent group has 4 ranks, and split_ranks can be + [[0, 1], [2, 3]]. Note [[0,1]] is also a valid split, in which case ranks 2, 3 would + return a non-group member. + timeout (timedelta, optional): see `init_process_group` for details and default value. + pg_options (ProcessGroupOptions, optional): only ProcessGroupNCCLOptions is supported now. + specifying what additional options need to be passed in during + the construction of specific process groups. i.e.``is_high_priority_stream`` + can be specified so that process group can pick up high priority cuda streams. + For other availble options to config nccl, + See https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t + group_desc (str, optional): a string to describe the process group. + + Returns: + ProcessGroup if the current rank is within one split/subgroup given by split_ranks, + or None if the current rank is not part of any split_ranks`. + + """ + # check inputs + if split_ranks is None: + raise ValueError("split_ranks cannot be None") + + global _world + default_pg = _get_default_group() + device_id = default_pg.bound_device_id + if not device_id: + raise RuntimeError( + "No device associated with the default pg, not safe to split any process groups" + ) + _default_backend, default_store = _world.pg_map[default_pg] + global_rank = default_pg.rank() + global_world_size = default_pg.size() + + if not parent_pg: + parent_pg = default_pg + if parent_pg not in _world.pg_group_ranks: + raise ValueError(f"Group {parent_pg} is not registered") + + parent_global_to_group_ranks = _world.pg_group_ranks[parent_pg] + parent_group_to_global_ranks = { + group_rank: global_rank + for global_rank, group_rank in parent_global_to_group_ranks.items() + } + + if global_rank not in parent_global_to_group_ranks: + raise ValueError( + f"Global rank {global_rank} is not part of the parent group {parent_pg}" + ) + + parent_group_rank = parent_global_to_group_ranks[global_rank] + parent_backend = parent_pg._get_backend(torch.device("cuda")) + + # if the parent backend does not support splitting, raise error + # currently this API only support NCCL backend + if ( + not parent_backend + or not parent_backend.supports_splitting + or not isinstance(parent_backend, ProcessGroupNCCL) + ): + raise RuntimeError( + "No backend for the parent process group or its backend does not support splitting" + ) + + # set the group_desc before the color or no_cloor split + group_desc = ( + f"{parent_pg.group_desc}:split:{parent_backend.comm_split_count()}" + if group_desc is None + else group_desc + ) + + parent_backend_str, _ = _world.pg_map[parent_pg] + # same type of backend as the parent process group + backend = Backend(parent_backend_str) + backend_config = BackendConfig(backend) + + if pg_options is not None: + assert isinstance( + pg_options, ProcessGroupNCCL.Options + ), "Expected pg_options argument to be of type ProcessGroupNCCL.Options" + else: + # default pg_options same as the parent process group + pg_options = parent_backend.options + + # this timeout defaulting/validation is used for all the new_groups/new_subgroups variants, + # which may just pass their timeout value (or None) + if timeout is None: + timeout = _get_default_timeout(backend) + _check_valid_timeout(timeout) + + # find my group of ranks and my group local rank in split_ranks + my_group = None + group_rank = -1 + + for split_group in split_ranks: + if len(split_group) == 0: + raise ValueError("the split group cannot be empty") + if len(split_group) > global_world_size: + raise ValueError( + "the split group's size should be less or equal to the world_size set by init_process_group" + ) + if len(split_group) != len(set(split_group)): + raise ValueError("the split group cannot have duplicate ranks") + split_group = sorted(split_group) + if parent_group_rank in split_group: + my_group = split_group + group_rank = split_group.index(parent_group_rank) + break + # if my rank does not belong to any sub group, + # no_color split should be called + if my_group is None or group_rank == -1: + parent_backend.perform_nocolor_split(device_id) + return None + + group_name = _process_group_name(my_group, use_hashed_name=False) + global_ranks_in_my_group = [parent_group_to_global_ranks[rank] for rank in my_group] + + prefix_store = PrefixStore(f"{group_name}/", default_store) + # We register the backend after initializing and timeout is set in pg_options. + pg: ProcessGroup = ProcessGroup( + prefix_store, + group_rank, + len(my_group), + ) + backend_type = ProcessGroup.BackendType.NCCL + pg.bound_device_id = device_id + pg._set_default_backend(backend_type) + + pg_options._timeout = timeout + pg_options.split_from = parent_backend + pg_options.split_color = _process_group_color(my_group) + pg_options.global_ranks_in_group = global_ranks_in_my_group + pg_options.group_name = group_name + backend_class = ProcessGroupNCCL( + prefix_store, group_rank, len(my_group), pg_options + ) + backend_class._set_sequence_number_for_group() + + pg._register_backend(torch.device("cuda"), backend_type, backend_class) + + # set group_name and group_desc to backend + assert group_name is not None + assert group_desc is not None + pg._set_group_name(group_name) + pg._set_group_desc(group_desc) + + # always eagerly initialize the backend in split_group + eager_backend = pg._get_backend(device_id) + eager_backend.eager_connect_single_device(device_id) + + # update global state + _world.pg_map[pg] = (backend, prefix_store) + _world.pg_names[pg] = group_name + _register_process_group(group_name, pg) + _world.pg_backend_config[pg] = str(backend_config) + pg_tag = f"ptd:{group_name}" + _world.tags_to_pg.setdefault(pg_tag, []).append(pg) + _world.pg_to_tag[pg] = pg_tag + + # Create the global rank to group rank mapping + _world.pg_group_ranks[pg] = { + global_rank: group_rank + for group_rank, global_rank in enumerate(global_ranks_in_my_group) + } + + return pg + + +@_time_logger +def new_group( + ranks=None, + timeout=None, + backend=None, + pg_options=None, + use_local_synchronization=False, + group_desc=None, + device_id: Optional[torch.device] = None, +): + """ + Create a new distributed group. + + This function requires that all processes in the main group (i.e. all + processes that are part of the distributed job) enter this function, even + if they are not going to be members of the group. Additionally, groups + should be created in the same order in all processes. + + .. warning:: + Safe concurrent usage: + When using multiple process groups with the ``NCCL`` backend, the user + must ensure a globally consistent execution order of collectives across + ranks. + + If multiple threads within a process issue collectives, explicit + synchronization is necessary to ensure consistent ordering. + + When using async variants of torch.distributed communication APIs, + a work object is returned and the communication kernel is + enqueued on a separate CUDA stream, allowing overlap of communication + and computation. Once one or more async ops have been issued on one process + group, they must be synchronized with other cuda streams by calling `work.wait()` + before using another process group. + + See `Using multiple NCCL communicators concurrently `_ for more details. + + Args: + ranks (list[int]): List of ranks of group members. If ``None``, will be + set to all ranks. Default is ``None``. + timeout (timedelta, optional): see `init_process_group` for details and default value. + backend (str or Backend, optional): The backend to use. Depending on + build-time configurations, valid values are ``gloo`` and ``nccl``. + By default uses the same backend as the global group. This field + should be given as a lowercase string (e.g., ``"gloo"``), which can + also be accessed via :class:`Backend` attributes (e.g., + ``Backend.GLOO``). If ``None`` is passed in, the backend + corresponding to the default process group will be used. Default is + ``None``. + pg_options (ProcessGroupOptions, optional): process group options + specifying what additional options need to be passed in during + the construction of specific process groups. i.e. for the ``nccl`` + backend, ``is_high_priority_stream`` can be specified so that + process group can pick up high priority cuda streams. For other availble options to config nccl, + See https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t + use_local_synchronization (bool, optional): perform a group-local + barrier at the end of the process group creation. This is different + in that non-member ranks don't need to call into API and don't + join the barrier. + group_desc (str, optional): a string to describe the process group. + device_id (torch.device, optional): a single, specific device + to "bind" this process to, The `new_group` call will try to initialize + a communication backend immediately for the device if this field is given. + + Returns: + A handle of distributed group that can be given to collective calls or + GroupMember.NON_GROUP_MEMBER if the rank is not part of ``ranks``. + + N.B. use_local_synchronization doesn't work with MPI. + + N.B. While use_local_synchronization=True can be significantly faster with larger + clusters and small process groups, care must be taken since it changes cluster behavior + as non-member ranks don't join the group barrier(). + + N.B. use_local_synchronization=True can lead to deadlocks when each rank creates + multiple overlaping process groups. To avoid that, make sure all ranks follow the + same global creation order. + """ + return _new_group_with_tag( + ranks, + timeout, + backend, + pg_options, + None, + use_local_synchronization=use_local_synchronization, + group_desc=group_desc, + device_id=device_id, + ) + + +def _new_group_with_tag( + ranks=None, + timeout=None, + backend=None, + backend_options=None, + pg_tag=None, + use_local_synchronization=False, + group_desc=None, + device_id: Optional[torch.device] = None, +): + """ + Variant of ``new_group`` that exposes tag creation. + + :: N.B. The mechanism is experimental and tied to the functional collectives effort, see + ``torch.distributed._functional_collectives`` for reference on how to use it. + """ + global _world + + default_pg = _get_default_group() + if device_id is None: + device_id = default_pg.bound_device_id + elif default_pg.bound_device_id is not None: + assert ( + device_id == default_pg.bound_device_id + ), "Mismatched bound device between new pg and the default pg." + default_backend, default_store = _world.pg_map[default_pg] + global_rank = default_pg.rank() + global_world_size = default_pg.size() + + # Default to the same backend as the global process group + # if the backend is not specified. + if not backend: + backend = default_backend + backend = Backend(backend) + + # this timeout defaulting/validation is used for all the new_groups/new_subgroups variants, + # which may just pass their timeout value (or None) + if timeout is None: + timeout = _get_default_timeout(backend) + _check_valid_timeout(timeout) + + if use_local_synchronization: + # MPI backend doesn't have have a way for us to perform a partial sync + if backend == Backend.MPI: + raise ValueError( + "MPI backend doesn't support use_local_synchronization=True" + ) + if ranks is not None and get_rank() not in ranks: + return None + + # checks the input ranks + if ranks is not None: + ranks = sorted(ranks) + group_world_size = len(ranks) + if group_world_size > global_world_size: + raise ValueError( + "the new group's world size should be less or " + "equal to the world size set by " + "init_process_group" + ) + # check ranks' sanity + for rank in ranks: + if rank < 0 or rank >= global_world_size: + raise ValueError( + "The new group's rank should be within " + "the world_size set by init_process_group" + ) + if global_rank in ranks: + group_rank = ranks.index(global_rank) + else: + group_rank = None + else: + ranks = list(range(global_world_size)) + group_world_size = global_world_size + group_rank = global_rank + + group_name = _process_group_name(ranks, use_hashed_name=use_local_synchronization) + + pg, pg_store = _new_process_group_helper( + group_world_size, + group_rank, + ranks, + backend, + default_store, + group_name, + backend_options=backend_options, + timeout=timeout, + pg_tag=pg_tag, + device_id=device_id, + group_desc=group_desc, + ) + + # Create the global rank to group rank mapping + _world.pg_group_ranks[pg] = { + global_rank: group_rank for group_rank, global_rank in enumerate(ranks) + } + + if _is_barrier_after_init() == 1: + # barrier at the end to ensure that once we return from this method, all + # process groups including global variables (if any) are updated + # correctly on all ranks. + # Update 04/2023: for large-scale runs, this barrier (esp. store-based + # barrier) may be costly and/or unscalable. Also, in a lot of cases, + # these barriers may be unnecessary, as proven by a green CI after + # removal. An environment variable `TORCH_DIST_INIT_BARRIER` has been + # added which enables this barrier only when set to 1. + logger.info( + "Performing barrier after ProcessGroup initialization since " + "TORCH_DIST_INIT_BARRIER = 1" + ) + if backend == Backend.MPI: + # MPI doesn't have store. + barrier() + else: + barrier_store = pg_store if use_local_synchronization else default_store + world_size = len(ranks) if use_local_synchronization else get_world_size() + # Use store based barrier here since barrier() used a bunch of + # default devices and messes up NCCL internal state. + _store_based_barrier( + global_rank, barrier_store, group_name, world_size, timeout + ) + + return pg + + +def new_subgroups( + group_size=None, + group=None, + timeout=None, + backend=None, + pg_options=None, + group_desc=None, +): + """ + Create subgroups of equal size. + + By default, it creates intra-machine subgroups, + where each of which contains all the ranks of a machine, based on the assumption + that each machine has the same number of devices. + + This is a convenience API that calls ``new_group`` to generate multiple subgroups. + It requires that all processes in the main group (i.e. all + processes that are part of the distributed job) enter this function, even + if they are not going to be members of the group. + + .. warning:: + If ``group_size`` is passed in, the world size must be divisible by ``group_size``. + If no ``group_size`` is passed in, it believe that you are creating a group based + on CUDA and determining the group size by number of CUDA devices, and if not all + the machines have the same number of devices, the subgroup division will be + different across nodes and can cause unexpected behaviors. Therefore, if you are + creating a subgroup that does not depend on CUDA (such as Gloo on CPU), please + pass in ``group_size`` correctly. + + .. warning:: + See warning `Safe concurrent usage` for `new_group` API for important details about + using multiple process groups concurrently in a safe manner. + + Args: + group_size (int, optional): The size of each subgroup. If ``None``, + the default subgroup size is equal to the number of devices on each machine, + based on the assumption that each machine has exactly the same + number of devices. Default is ``None``. + timeout (timedelta, optional): see `init_process_group` for details and default value. + backend (str or Backend, optional): The backend to use. Depending on + build-time configurations, valid values are ``gloo`` and ``nccl``. + By default uses the same backend as the global group. This field + should be given as a lowercase string (e.g., ``"gloo"``), which can + also be accessed via :class:`Backend` attributes (e.g., + ``Backend.GLOO``). If ``None`` is passed in, the backend + corresponding to the default process group will be used. Default is + ``None``. + pg_options (ProcessGroupOptions, optional): process group options + specifying what additional options need to be passed in during + the construction of specific process groups. i.e. for the ``nccl`` + backend, ``is_high_priority_stream`` can be specified so that + process group can pick up high priority cuda streams. + group_desc (str, optional): A string describing the group. Each subgroup will + inherit its group_desc + + Returns: + The subgroup containing the current rank, and all the subgroups used for cleanup. + + Examples: + >>> # Create intra-machine subgroups. + >>> # xdoctest: +SKIP("need process group init") + >>> cur_subgroup, subgroups = dist.new_subgroups() + >>> # Allreduce within the machine. + >>> rank = dist.get_rank() + >>> tensor = torch.ones(1, device=rank) * rank + >>> dist.all_reduce(tensor, group=cur_subgroup) + >>> tensor + tensor([28]) # Assume 8 CUDA devices per machine. 28 is sum(range(8)). + >>> # Cleanup. + >>> for subgroup in subgroups: + >>> dist.destroy_process_group(subgroup) + """ + if group_size is None: + if not torch.cuda.is_available(): + raise ValueError( + "Default group size only takes effect when CUDA is available." + "If your subgroup using a backend that does not depend on CUDA," + "please pass in 'group_size' correctly." + ) + group_size = torch.cuda.device_count() + if group_size <= 0: + raise ValueError(f"The arg 'group_size' ({group_size}) must be positive") + + world_size = get_world_size() + if world_size < group_size: + raise ValueError( + f"The arg 'group_size' ({group_size}) must not exceed the world size ({world_size})" + ) + if world_size % group_size != 0: + raise ValueError("The world size must be divisible by 'group_size'") + + subgroups = [] + cur_subgroup = None + + for subgroup_id in range(world_size // group_size): + start_rank = subgroup_id * group_size + end_rank = start_rank + group_size + ranks_in_subgroup = list(range(start_rank, end_rank)) + subgroup = new_group( + ranks=ranks_in_subgroup, + timeout=timeout, + backend=backend, + pg_options=pg_options, + group_desc=group_desc, + ) + subgroups.append(subgroup) + + rank = get_rank() + if rank in ranks_in_subgroup: + cur_subgroup = subgroup + logger.info("Rank %s is assigned to subgroup %s", rank, ranks_in_subgroup) + + return cur_subgroup, subgroups + + +def new_subgroups_by_enumeration( + ranks_per_subgroup_list, + timeout=None, + backend=None, + pg_options=None, + group_desc=None, +): + """ + Create subgroups by dividing the global world. + + The division is specified by a nested list of ranks. The subgroups cannot have + overlap, and some ranks may not have to be in any subgroup. + + This is a convenience API that calls ``new_group`` to generate multiple subgroups. + It requires that all processes in the main group (i.e. all + processes that are part of the distributed job) enter this function, even + if they are not going to be members of the group. + + .. warning:: + See warning `Safe concurrent usage` for `new_group` API for important details about + using multiple process groups concurrently in a safe manner. + + Args: + ranks_per_subgroup_list (list[list[int]]): A nested list of ranks of + group members. + timeout (timedelta, optional): see `init_process_group` for details and default value. + backend (str or Backend, optional): The backend to use. Depending on + build-time configurations, valid values are ``gloo`` and ``nccl``. + By default uses the same backend as the global group. This field + should be given as a lowercase string (e.g., ``"gloo"``), which can + also be accessed via :class:`Backend` attributes (e.g., + ``Backend.GLOO``). If ``None`` is passed in, the backend + corresponding to the default process group will be used. Default is + ``None``. + pg_options (ProcessGroupOptions, optional): process group options + specifying what additional options need to be passed in during + the construction of specific process groups. i.e. for the ``nccl`` + backend, ``is_high_priority_stream`` can be specified so that + process group can pick up high priority cuda streams. + group_desc (str, optional): A string describing the group. Each subgroup will + inherit its group_desc. + + Returns: + The subgroup containing the current rank, and all the subgroups used for cleanup. + + Examples: + >>> # Create two subgroups, where each has 2 processes. + >>> # xdoctest: +SKIP("need process group init") + >>> cur_subgroup, subgroups = dist.new_subgroups(ranks=[[0, 2], [1, 3]]) + >>> rank = dist.get_rank() + >>> tensor = torch.ones(1, device=rank) * rank + >>> dist.all_reduce(tensor, group=cur_subgroup) + >>> tensor + tensor([2]) # Subgroup 0: ranks 0 and 2 + tensor([4]) # Subgroup 1: ranks 1 and 3 + """ + if ranks_per_subgroup_list is None or len(ranks_per_subgroup_list) == 0: + raise ValueError("The arg 'ranks_per_subgroup_list' cannot be empty") + + subgroups = [] + cur_subgroup = None + # Create a mapping from rank to subgroup to check if there is any subgroup overlap. + rank_to_ranks_dict = {} # type: ignore[var-annotated] + for ranks in ranks_per_subgroup_list: + subgroup = new_group( + ranks=ranks, + timeout=timeout, + backend=backend, + pg_options=pg_options, + group_desc=group_desc, + ) + subgroups.append(subgroup) + my_rank = get_rank() + for rank in ranks: + if rank in rank_to_ranks_dict: + raise ValueError( + f"Rank {rank} has appeared in both subgroup {rank_to_ranks_dict[rank]} and {ranks}" + ) + rank_to_ranks_dict[rank] = ranks + if my_rank == rank: + cur_subgroup = subgroup + logger.info("Rank %s is assigned to subgroup %s", rank, ranks) + + return cur_subgroup, subgroups + + +def _find_pg_by_ranks_and_tag(tag: str, ranks: List[int]) -> Optional[ProcessGroup]: + if len(tag) > 0 and not tag.startswith("ptd:") and not tag.startswith("user:"): + tag = f"user:{tag}" + + for group in _world.tags_to_pg.get(tag, []): + if group.size() != len(ranks): + continue + + group_ranks = get_process_group_ranks(group) + good = all(r in group_ranks for r in ranks) + if good: + return group + return None + + +def _find_or_create_pg_by_ranks_and_tag( + tag: str, ranks: List[int], stride: int +) -> ProcessGroup: + assert ( + len(ranks) % stride == 0 + ), f"Ranks length ({len(ranks)}) must be divisible by stride ({stride})" + + my_rank = get_rank() + my_ranks = None + + if stride == len(ranks): + my_ranks = ranks.copy() + assert my_rank in my_ranks, "rankset doesn't include the current node" + else: + for i in range(0, len(ranks), stride): + rank_set = ranks[i : i + stride] + if my_rank in rank_set: + my_ranks = rank_set + assert my_ranks is not None, "rankset doesn't include the current node" + + my_ranks = sorted(my_ranks) + + pg = _find_pg_by_ranks_and_tag(tag, my_ranks) + if pg is not None: + return pg + if tag == "": + raise ValueError("Cannot automatically create PG with empty tag") + # TODO copy settings and timeout from default PG + return _new_group_with_tag(my_ranks, pg_tag=tag) + + +def _get_group_tag(pg: ProcessGroup) -> str: + """Return the tag associated with ``pg``.""" + tag = _world.pg_to_tag[pg] + if tag.startswith("user:"): + tag = tag[5:] + return tag + + +def _get_process_group_name(pg: ProcessGroup) -> str: + return _world.pg_names.get(pg, "None") + + +def _get_process_group_store(pg: ProcessGroup) -> Store: + return _world.pg_map[pg][1] diff --git a/lib/python3.12/site-packages/torch/distributed/launch.py b/lib/python3.12/site-packages/torch/distributed/launch.py new file mode 100644 index 0000000000000000000000000000000000000000..a9e35c36db7fbc585351128b9bbd712924e809af --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/launch.py @@ -0,0 +1,208 @@ +# mypy: allow-untyped-defs +r""" +Module ``torch.distributed.launch``. + +``torch.distributed.launch`` is a module that spawns up multiple distributed +training processes on each of the training nodes. + +.. warning:: + + This module is going to be deprecated in favor of :ref:`torchrun `. + +The utility can be used for single-node distributed training, in which one or +more processes per node will be spawned. The utility can be used for either +CPU training or GPU training. If the utility is used for GPU training, +each distributed process will be operating on a single GPU. This can achieve +well-improved single-node training performance. It can also be used in +multi-node distributed training, by spawning up multiple processes on each node +for well-improved multi-node distributed training performance as well. +This will especially be beneficial for systems with multiple Infiniband +interfaces that have direct-GPU support, since all of them can be utilized for +aggregated communication bandwidth. + +In both cases of single-node distributed training or multi-node distributed +training, this utility will launch the given number of processes per node +(``--nproc-per-node``). If used for GPU training, this number needs to be less +or equal to the number of GPUs on the current system (``nproc_per_node``), +and each process will be operating on a single GPU from *GPU 0 to +GPU (nproc_per_node - 1)*. + +**How to use this module:** + +1. Single-Node multi-process distributed training + +:: + + python -m torch.distributed.launch --nproc-per-node=NUM_GPUS_YOU_HAVE + YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other + arguments of your training script) + +2. Multi-Node multi-process distributed training: (e.g. two nodes) + + +Node 1: *(IP: 192.168.1.1, and has a free port: 1234)* + +:: + + python -m torch.distributed.launch --nproc-per-node=NUM_GPUS_YOU_HAVE + --nnodes=2 --node-rank=0 --master-addr="192.168.1.1" + --master-port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 + and all other arguments of your training script) + +Node 2: + +:: + + python -m torch.distributed.launch --nproc-per-node=NUM_GPUS_YOU_HAVE + --nnodes=2 --node-rank=1 --master-addr="192.168.1.1" + --master-port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 + and all other arguments of your training script) + +3. To look up what optional arguments this module offers: + +:: + + python -m torch.distributed.launch --help + + +**Important Notices:** + +1. This utility and multi-process distributed (single-node or +multi-node) GPU training currently only achieves the best performance using +the NCCL distributed backend. Thus NCCL backend is the recommended backend to +use for GPU training. + +2. In your training program, you must parse the command-line argument: +``--local-rank=LOCAL_PROCESS_RANK``, which will be provided by this module. +If your training program uses GPUs, you should ensure that your code only +runs on the GPU device of LOCAL_PROCESS_RANK. This can be done by: + +Parsing the local_rank argument + +:: + + >>> # xdoctest: +SKIP + >>> import argparse + >>> parser = argparse.ArgumentParser() + >>> parser.add_argument("--local-rank", "--local_rank", type=int) + >>> args = parser.parse_args() + +Set your device to local rank using either + +:: + + >>> torch.cuda.set_device(args.local_rank) # before your code runs + +or + +:: + + >>> with torch.cuda.device(args.local_rank): + >>> # your code to run + >>> ... + +.. versionchanged:: 2.0.0 + + The launcher will passes the ``--local-rank=`` argument to your script. + From PyTorch 2.0.0 onwards, the dashed ``--local-rank`` is preferred over the + previously used underscored ``--local_rank``. + + For backward compatibility, it may be necessary for users to handle both + cases in their argument parsing code. This means including both ``"--local-rank"`` + and ``"--local_rank"`` in the argument parser. If only ``"--local_rank"`` is + provided, the launcher will trigger an error: "error: unrecognized arguments: + --local-rank=". For training code that only supports PyTorch 2.0.0+, + including ``"--local-rank"`` should be sufficient. + +3. In your training program, you are supposed to call the following function +at the beginning to start the distributed backend. It is strongly recommended +that ``init_method=env://``. Other init methods (e.g. ``tcp://``) may work, +but ``env://`` is the one that is officially supported by this module. + +:: + + >>> torch.distributed.init_process_group(backend='YOUR BACKEND', + >>> init_method='env://') + +4. In your training program, you can either use regular distributed functions +or use :func:`torch.nn.parallel.DistributedDataParallel` module. If your +training program uses GPUs for training and you would like to use +:func:`torch.nn.parallel.DistributedDataParallel` module, +here is how to configure it. + +:: + + >>> model = torch.nn.parallel.DistributedDataParallel(model, + >>> device_ids=[args.local_rank], + >>> output_device=args.local_rank) + +Please ensure that ``device_ids`` argument is set to be the only GPU device id +that your code will be operating on. This is generally the local rank of the +process. In other words, the ``device_ids`` needs to be ``[args.local_rank]``, +and ``output_device`` needs to be ``args.local_rank`` in order to use this +utility + +5. Another way to pass ``local_rank`` to the subprocesses via environment variable +``LOCAL_RANK``. This behavior is enabled when you launch the script with +``--use-env=True``. You must adjust the subprocess example above to replace +``args.local_rank`` with ``os.environ['LOCAL_RANK']``; the launcher +will not pass ``--local-rank`` when you specify this flag. + +.. warning:: + + ``local_rank`` is NOT globally unique: it is only unique per process + on a machine. Thus, don't use it to decide if you should, e.g., + write to a networked filesystem. See + https://github.com/pytorch/pytorch/issues/12042 for an example of + how things can go wrong if you don't do this correctly. + + + +""" + +from typing_extensions import deprecated as _deprecated + +from torch.distributed.run import get_args_parser, run + + +def parse_args(args): + parser = get_args_parser() + parser.add_argument( + "--use-env", + "--use_env", + default=False, + action="store_true", + help="Use environment variable to pass " + "'local rank'. For legacy reasons, the default value is False. " + "If set to True, the script will not pass " + "--local-rank as argument, and will instead set LOCAL_RANK.", + ) + return parser.parse_args(args) + + +def launch(args): + if args.no_python and not args.use_env: + raise ValueError( + "When using the '--no-python' flag," + " you must also set the '--use-env' flag." + ) + run(args) + + +@_deprecated( + "The module torch.distributed.launch is deprecated\n" + "and will be removed in future. Use torchrun.\n" + "Note that --use-env is set by default in torchrun.\n" + "If your script expects `--local-rank` argument to be set, please\n" + "change it to read from `os.environ['LOCAL_RANK']` instead. See \n" + "https://pytorch.org/docs/stable/distributed.html#launch-utility for \n" + "further instructions\n", + category=FutureWarning, +) +def main(args=None): + args = parse_args(args) + launch(args) + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/torch/distributed/logging_handlers.py b/lib/python3.12/site-packages/torch/distributed/logging_handlers.py new file mode 100644 index 0000000000000000000000000000000000000000..021ad100f06a89fa944d7e6dba18f1a9558d88f5 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/logging_handlers.py @@ -0,0 +1,17 @@ +#!/usr/bin/env python3 + +# Copyright (c) Facebook, Inc. and its affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from typing import Dict, List + + +__all__: List[str] = [] + +_log_handlers: Dict[str, logging.Handler] = { + "default": logging.NullHandler(), +} diff --git a/lib/python3.12/site-packages/torch/distributed/nn/__init__.py b/lib/python3.12/site-packages/torch/distributed/nn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e15fb517052e4aefeb7377d1f0ca63cf2b2da753 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/nn/__init__.py @@ -0,0 +1,7 @@ +import torch + +from .functional import * # noqa: F403 + + +if torch.distributed.rpc.is_available(): + from .api.remote_module import RemoteModule diff --git a/lib/python3.12/site-packages/torch/distributed/nn/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/nn/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c5aee08b4f0866e310127da9c73f40e6f05cc627 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/nn/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/nn/__pycache__/functional.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/nn/__pycache__/functional.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..216160b960f942ae9d5e33021853fb6477a988ec Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/nn/__pycache__/functional.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/nn/api/__init__.py b/lib/python3.12/site-packages/torch/distributed/nn/api/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/distributed/nn/api/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/nn/api/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64dc5b3f4bb6ba57b1dd1361f70a74df894288e9 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/nn/api/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/nn/api/__pycache__/remote_module.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/nn/api/__pycache__/remote_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..611053a1ed56a628e21c528a3f6a4273a891f741 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/nn/api/__pycache__/remote_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/nn/api/remote_module.py b/lib/python3.12/site-packages/torch/distributed/nn/api/remote_module.py new file mode 100644 index 0000000000000000000000000000000000000000..6ce293c883c5ca65ca5af5623d4708d52c3fd6be --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/nn/api/remote_module.py @@ -0,0 +1,762 @@ +#!/usr/bin/python3 +# mypy: allow-untyped-defs +import collections +import io +import sys +import types +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Mapping, + Optional, + Set, + Tuple, + Type, + TypeVar, + Union, +) + +import torch +import torch.distributed.rpc as rpc +from torch import device, dtype, nn, Tensor +from torch.distributed import _remote_device +from torch.distributed.nn.jit import instantiator +from torch.distributed.rpc.internal import _internal_rpc_pickler +from torch.nn import Module +from torch.nn.parameter import Parameter +from torch.utils.hooks import RemovableHandle + + +__all__ = ["RemoteModule"] + +_grad_t = Union[Tuple[Tensor, ...], Tensor] +# See https://mypy.readthedocs.io/en/latest/generics.html#generic-methods-and-generic-self for the use +# of `T` to annotate `self`. Many methods of `Module` return `self` and we want those return values to be +# the type of the subclass, not the looser type of `Module`. +T = TypeVar("T", bound="Module") + +_NON_SCRIPTABLE_REMOTE_MODULE_MODULE = ( + instantiator.instantiate_non_scriptable_remote_module_template() +) + +_REMOTE_MODULE_PICKLED_ATTRIBUTES = ( + "on", + "device", + "is_device_map_set", + "is_scriptable", + "generated_methods", + "module_rref", +) + +_SerializedRemoteModule = collections.namedtuple("_SerializedRemoteModule", _REMOTE_MODULE_PICKLED_ATTRIBUTES) # type: ignore[misc] + +# These attributes are mostly from RemoteModule's parent class and are intentionally not pickled. +# A new attribute of RemoteModule should be either in _REMOTE_MODULE_PICKLED_ATTRIBUTES +# or _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING. +# Otherwise, it will not be pickled. +_REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING = ( + "training", + "_parameters", + "_buffers", + "_non_persistent_buffers_set", + "_backward_hooks", + "_backward_pre_hooks", + "_is_full_backward_hook", + "_forward_hooks", + "_forward_hooks_with_kwargs", + "_forward_hooks_always_called", + "_forward_pre_hooks", + "_forward_pre_hooks_with_kwargs", + "_state_dict_hooks", + "_state_dict_pre_hooks", + "_load_state_dict_pre_hooks", + "_load_state_dict_post_hooks", + "_state_dict_pre_hooks", + "_modules", + # The two attributes below are generated methods, not available at pickling time. + "forward_async", + "forward", +) + + +# RPC handler. +def _instantiate_template(module_interface_cls, enable_moving_cpu_tensors_to_cuda): + instantiator.instantiate_scriptable_remote_module_template( + module_interface_cls, enable_moving_cpu_tensors_to_cuda + ) + + +def _create_module(module_cls, args, kwargs, device): + module = module_cls(*args, **kwargs) + if not isinstance(module, nn.Module): + raise ValueError( + "Expect `module_cls(*args, **kwargs)` returns an instance of , " + f"but it returns an instance of {type(module)}." + ) + module.to(device) + return module + + +def _create_module_with_interface( + module_cls, args, kwargs, device, module_interface_cls +): + module = _create_module(module_cls, args, kwargs, device) + if module_interface_cls is not None: + module = torch.jit.script(module) + return rpc.RRef(module, module_interface_cls) + + +def _param_rrefs(module_rref, recurse) -> List[rpc.RRef[Parameter]]: + ret: List[rpc.RRef[Parameter]] = [ + rpc.RRef(param) for param in module_rref.local_value().parameters(recurse) + ] + return ret + + +def _raise_not_supported(name: str) -> None: + raise ValueError(f"Method ``{name}`` not supported for RemoteModule") + + +class _RemoteModule(nn.Module): + def __new__(cls, *args, **kwargs): + # Use __new__ for logging purposes. + torch._C._log_api_usage_once("torch.distributed.nn.api.remote_module") + return super().__new__(cls) + + def __init__( + self, + remote_device: str, + module_cls: Type[nn.Module], + args: Optional[Tuple] = None, + kwargs: Optional[Dict[str, Any]] = None, + _module_interface_cls: Any = None, + ): + """ + RemoteModule instance can only be created after RPC initialization. + + It creates a user-specified module on a specified remote node. + It behaves like a regular ``nn.Module`` except that the ``forward`` method is + executed on the remote node. + It takes care of autograd recording to ensure the backward pass propagates + gradients back to the corresponding remote module. + It can be shared across processors using `RPC framework `__, + without incurring any overheads of copying the actual module, + which is equivalent to an :class:`~torch.distributed.rpc.RRef` + pointing to the remote module. + + The arguments of ``forward_async`` and ``forward`` are the same as + the ``forward`` method of the module returned by the ``module_cls``. + + Apart from ``forward_async`` and ``forward``, no other methods are supported from nn.Module for now. + + Particularly, to create a hybrid model, typically the local modules should be + created outside of remote modules, rather than as submodules of any remote module (by calling ``add_module``). + Hybrid Example: + >>> class HybridModel(nn.Module): + >>> def __init__(self) -> None: + >>> nn.Module.__init__(self) + >>> self.remote_embedding = RemoteModule(...) + >>> self.local_linear = nn.Linear(...) + + For example, if ``module_cls`` returns an instance of ``nn.Linear``, + that has ``forward`` method signature, ``def forward(input: Tensor) -> Tensor:``, + the generated ``RemoteModule`` will have 2 methods in signature of + ``def forward(input: Tensor) -> Tensor:`` and + ``def forward_async(input: Tensor) -> Future[Tensor]:``. + + .. note:: + If the remote module is placed on a cuda device, + any input CPU tensors will be automatically moved to the same cuda device, + and GPU tensors are returned over the wire according to the device map of the remote worker on TensorPipe RPC backend. + + Args: + remote_device (str): Device on the destination worker where we'd like to place this module. + The device can be a local device or a remote device specified by one of the following remote + formats: + + 1. "rank:/" (ex: "rank:0/cuda:0"). + 2. "/" (ex: "trainer0/cuda:0"). + + In addition, the device field can be optional and the default value is "cpu". + module_cls (nn.Module): For example, + >>> class MyModule(nn.Module): + >>> def forward(input): + >>> return input + 1 + >>> + >>> module_cls = MyModule + args (Sequence, optional): args to be passed to ``module_cls``. + kwargs (Dict, optional): kwargs to be passed to ``module_cls``. + _module_interface_cls (type, optional): The TorchScript interface type for the module + to be created. The type object should be decorated by @torch.jit.interface. + If not provided, the generated RemoteModule is not torchscript-able. + Warning, this is an experimental API and susceptible to frequent changes. + + Returns: + A remote module instance which wraps the :class:`~nn.Module` created by the + user-provided ``module_cls``, it has a blocking ``forward`` method and an + asynchronous ``forward_async`` method that returns a future of the ``forward`` call + on the user-provided module on the remote side. + + Example:: + Run the following code in two different processes: + + >>> # xdoctest: +SKIP("distributed") + >>> # On worker 0: + >>> import torch + >>> import torch.distributed.rpc as rpc + >>> from torch import nn, Tensor + >>> from torch.distributed.nn.api.remote_module import RemoteModule + >>> + >>> rpc.init_rpc("worker0", rank=0, world_size=2) + >>> remote_linear_module = RemoteModule( + >>> "worker1/cpu", nn.Linear, args=(20, 30), + >>> ) + >>> input = torch.randn(128, 20) + >>> ret_fut = remote_linear_module.forward_async(input) + >>> ret = ret_fut.wait() + >>> rpc.shutdown() + + >>> # On worker 1: + >>> import torch + >>> import torch.distributed.rpc as rpc + >>> + >>> rpc.init_rpc("worker1", rank=1, world_size=2) + >>> rpc.shutdown() + """ + super().__init__() + + enable_moving_cpu_tensors_to_cuda = self._prepare_init(remote_device) + + # Default arguments preparation. + args = args if args is not None else () + kwargs = kwargs if kwargs is not None else {} + + if _module_interface_cls is not None: + # Users reply on this field to know if this generated RemoteModule is TorchScript-able. + self.is_scriptable = True + + # Instantiate template on remote side. + fut = rpc.rpc_async( + self.on, + _instantiate_template, + (_module_interface_cls, enable_moving_cpu_tensors_to_cuda), + ) + + self._init_template( + _module_interface_cls, enable_moving_cpu_tensors_to_cuda + ) + + # Instantiate template on remote side. + fut = rpc.rpc_async( + self.on, + _instantiate_template, + (_module_interface_cls, enable_moving_cpu_tensors_to_cuda), + ) + + # Create the module on the remote side. + fut.wait() # Ensure remote_module_cls is available on remote side. + + # TODO: We need to change this to rpc.remote, and make it async (see the else branch below). + # For that we need to be able to apply _module_interface_cls to the RRef returned by rpc.remote + # See https://github.com/pytorch/pytorch/issues/58098 for more context. + self.module_rref = rpc.rpc_sync( + self.on, + _create_module_with_interface, + (module_cls, args, kwargs, self.device, _module_interface_cls), + ) + else: + self.is_scriptable = False + self.generated_methods = ( + _NON_SCRIPTABLE_REMOTE_MODULE_MODULE._generated_methods + ) + # Create the module on the remote side. + self.module_rref = rpc.remote( + self.on, + _create_module, + (module_cls, args, kwargs, self.device), + ) + + self._install_generated_methods() + self._check_attribute_picklability() + + def remote_parameters(self, recurse: bool = True) -> List[rpc.RRef[Parameter]]: + """ + Return a list of :class:`~torch.distributed.rpc.RRef` pointing to the remote module's parameters. + + This can typically be used in conjunction + with :class:`~torch.distributed.optim.DistributedOptimizer`. + + Args: + recurse (bool): if True, then returns parameters of the remote + module and all submodules of the remote module. Otherwise, + returns only parameters that are direct members of the + remote module. + + Returns: + A list of :class:`~torch.distributed.rpc.RRef` (``List[RRef[nn.Parameter]]``) + to remote module's parameters. + """ + return rpc.rpc_sync(self.on, _param_rrefs, args=(self.module_rref, recurse)) + + def get_module_rref(self) -> rpc.RRef[nn.Module]: + """Return an :class:`~torch.distributed.rpc.RRef` (``RRef[nn.Module]``) pointing to the remote module.""" + return self.module_rref + + @torch.jit.export + def __getstate__(self): + raise RuntimeError( + "Cannot pickle RemoteModule in python pickler. RemoteModule can only be pickled when using RPC" + ) + + @torch.jit.export + def __setstate__(self, state): + raise RuntimeError( + "Cannot unpickle RemoteModule in python pickler. RemoteModule can only be unpickled when using RPC" + ) + + def register_buffer( + self, name: str, tensor: Optional[Tensor], persistent: bool = True + ) -> None: + _raise_not_supported(self.register_buffer.__name__) + + def register_parameter(self, name: str, param: Optional[Parameter]) -> None: + _raise_not_supported(self.register_parameter.__name__) + + def add_module(self, name: str, module: Optional[Module]) -> None: + _raise_not_supported(self.add_module.__name__) + + def apply(self: T, fn: Callable[[Module], None]) -> T: # type: ignore[return] + _raise_not_supported(self.apply.__name__) + + def cuda(self: T, device: Optional[Union[int, device]] = None) -> T: # type: ignore[return] + _raise_not_supported(self.cuda.__name__) + + def ipu(self: T, device: Optional[Union[int, device]] = None) -> T: # type: ignore[return] + _raise_not_supported(self.ipu.__name__) + + def xpu(self: T, device: Optional[Union[int, device]] = None) -> T: # type: ignore[return] + _raise_not_supported(self.xpu.__name__) + + def cpu(self: T) -> T: # type: ignore[return] + _raise_not_supported(self.cpu.__name__) + + def type(self: T, dst_type: Union[dtype, str]) -> T: # type: ignore[return] + _raise_not_supported(self.type.__name__) + + def float(self: T) -> T: # type: ignore[return] + _raise_not_supported(self.float.__name__) + + def double(self: T) -> T: # type: ignore[return] + _raise_not_supported(self.double.__name__) + + def half(self: T) -> T: # type: ignore[return] + _raise_not_supported(self.half.__name__) + + def bfloat16(self: T) -> T: # type: ignore[return] + _raise_not_supported(self.bfloat16.__name__) + + def to(self, *args, **kwargs) -> T: # type: ignore[misc, return, type-var] + _raise_not_supported(self.to.__name__) + + def register_backward_hook( # type: ignore[return] + self, hook: Callable[[Module, _grad_t, _grad_t], Union[None, _grad_t]] + ) -> RemovableHandle: + _raise_not_supported(self.register_backward_hook.__name__) + + def register_forward_pre_hook( # type: ignore[return] + self, + hook: Union[ + Callable[[T, Tuple[Any, ...]], Optional[Any]], + Callable[ + [T, Tuple[Any, ...], Dict[str, Any]], + Optional[Tuple[Any, Dict[str, Any]]], + ], + ], + prepend: bool = False, + with_kwargs: bool = False, + ) -> RemovableHandle: + _raise_not_supported(self.register_forward_pre_hook.__name__) + + def register_forward_hook( # type: ignore[return, override] + self, + hook: Union[ + Callable[[T, Tuple[Any, ...], Any], Optional[Any]], + Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]], + ], + prepend: bool = False, + with_kwargs: bool = False, + ) -> RemovableHandle: + _raise_not_supported(self.register_forward_hook.__name__) + + def state_dict(self, *args, **kwargs): + _raise_not_supported(self.state_dict.__name__) + + def load_state_dict( + self, + state_dict: Mapping[str, Any], + strict: bool = True, + assign: bool = False, + ): + _raise_not_supported(self.load_state_dict.__name__) + + def parameters(self, recurse: bool = True) -> Iterator[Parameter]: + raise ValueError( + "Method ``parameters`` not supported for RemoteModule. Please use ``remote_parameters`` instead." + ) + + def named_parameters( # type: ignore[return] + self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True + ) -> Iterator[Tuple[str, Parameter]]: + _raise_not_supported(self.named_parameters.__name__) + + def buffers(self, recurse: bool = True) -> Iterator[Tensor]: # type: ignore[return] + _raise_not_supported(self.buffers.__name__) + + def named_buffers( # type: ignore[return] + self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True + ) -> Iterator[Tuple[str, Tensor]]: + _raise_not_supported(self.named_buffers.__name__) + + def children(self) -> Iterator[Module]: # type: ignore[return] + _raise_not_supported(self.children.__name__) + + def named_children(self) -> Iterator[Tuple[str, Module]]: # type: ignore[return] + _raise_not_supported(self.named_children.__name__) + + def modules(self) -> Iterator[Module]: # type: ignore[return] + _raise_not_supported(self.modules.__name__) + + def named_modules( + self, + memo: Optional[Set[Module]] = None, + prefix: str = "", + remove_duplicate: bool = True, + ): + _raise_not_supported(self.named_modules.__name__) + + def train(self: T, mode: bool = True) -> T: + return self.module_rref.rpc_sync().train() # type: ignore[operator, union-attr] + + def eval(self: T) -> T: + return self.module_rref.rpc_sync().eval() # type: ignore[operator, union-attr] + + def requires_grad_(self: T, requires_grad: bool = True) -> T: # type: ignore[return] + _raise_not_supported(self.requires_grad_.__name__) + + def zero_grad(self, set_to_none: bool = True) -> None: + _raise_not_supported(self.zero_grad.__name__) + + def share_memory(self: T) -> T: # type: ignore[return] + _raise_not_supported(self.share_memory.__name__) + + def extra_repr(self) -> str: # type: ignore[return] + _raise_not_supported(self.extra_repr.__name__) + + def _prepare_init(self, remote_device_str: str) -> bool: + """Prepare the initialization and returns whether to enable automatically moving CPU tensors to CUDA devices.""" + # Sanity check. + assert rpc._is_current_rpc_agent_set(), "RemoteModule only works in RPC." + + remote_device = _remote_device(remote_device_str) + self.on = ( + remote_device.worker_name() + if remote_device.worker_name() is not None + else remote_device.rank() + ) + self.device = str(remote_device.device()) + agent = rpc._get_current_rpc_agent() + # If the device map of the remote worker is set, + # then enable moving any input CPU tensors to the same cuda device. + self.is_device_map_set = bool( + agent._get_device_map(agent.get_worker_info(self.on)) # type: ignore[arg-type] + ) + # ``enable_moving_cpu_tensors_to_cuda`` is less strict than ``is_device_map_set``: + # If ``enable_moving_cpu_tensors_to_cuda`` is true, but the device map is not set, + # then any CPU tensors can still be moved to a cuda device to run forward, + # but the output must be moved back to CPU before being sent over the wire. + enable_moving_cpu_tensors_to_cuda = torch.device(self.device).type == "cuda" + return enable_moving_cpu_tensors_to_cuda + + def _init_template(self, module_interface_cls, enable_moving_cpu_tensors_to_cuda): + """Instantiate template on local side.""" + generated_module = instantiator.instantiate_scriptable_remote_module_template( + module_interface_cls, enable_moving_cpu_tensors_to_cuda + ) + self.generated_methods = generated_module._generated_methods + + def _check_attribute_picklability(self): + """Check if all the attribute has explicitly defined whether to be pickled (i.e., picklability).""" + for k in self.__dict__.keys(): + if ( + k not in _REMOTE_MODULE_PICKLED_ATTRIBUTES + and k not in _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING + ): + raise AttributeError( + f"Attribute {k} must be either in ``_REMOTE_MODULE_PICKLED_ATTRIBUTES`` or " + "``_REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING``." + ) + + def _install_generated_methods(self): + for method in self.generated_methods: + method_name = method.__name__ + method = torch.jit.export(method) + setattr(self, method_name, types.MethodType(method, self)) + + @staticmethod + def init_from_module_rref( + remote_device: str, + module_rref: rpc.RRef[nn.Module], + _module_interface_cls: Any = None, + ): + """ + Besides the constructor, a RemoteModule instance can also be initialized given a module RRef. + + This alternate initialization method can be particularly useful if we want to create multiple + RemoteModule instances that share the same underlying module and reduce memory consumption. + + Moreover, this also provides a workaround for passing script RemoteModule over RPC, + which is not supported. The recommended way is as follows: + + 1. the sender creates a RemoteModule; + 2. the sender sends its ``module_rref`` over RPC; + 3. the receiver calls this method to initialize another RemoteModule using the same ``module_rref``. + + Example:: + Run the following code in two different processes: + + >>> # xdoctest: +SKIP("distributed") + >>> # On worker 0: + >>> import torch + >>> import torch.distributed.rpc as rpc + >>> from torch import nn, Tensor + >>> from torch.distributed.nn.api.remote_module import RemoteModule + >>> + >>> rpc.init_rpc("worker0", rank=0, world_size=2) + >>> remote_module = RemoteModule( + >>> "worker1/cpu", nn.Linear, args=(20, 30), + >>> ) + >>> + >>> remote_module1 = rpc.rpc_sync( + >>> "worker1/cpu", + >>> RemoteModule.init_from_module_rref, + >>> ("worker1/cpu", remote_module1.get_module_rref()), + >>> ) + >>> rpc.shutdown() + + >>> # On worker 1: + >>> import torch + >>> import torch.distributed.rpc as rpc + >>> + >>> rpc.init_rpc("worker1", rank=1, world_size=2) + >>> rpc.shutdown() + + Args: + remote_device (str): Device on the destination worker where we'd like to place this module. + The device can be a local device or a remote device specified by one of the following remote + formats: + + 1. "rank:/" (ex: "rank:0/cuda:0"). + 2. "/" (ex: "trainer0/cuda:0"). + + In addition, the device field can be optional and the default value is "cpu". + module_rref (RRef[nn.Module]): The module reference shared by both the caller and + the created remote module. + _module_interface_cls (type, optional): The TorchScript interface type for the module + to be created. The type object should be decorated by @torch.jit.interface. + If not provided, the generated RemoteModule is not torchscript-able. + Warning, this is an experimental API and susceptible to frequent changes. + + Returns: + A remote module instance which wraps the :class:`~nn.Module` created by the + user-provided ``module_rref``, it has a blocking ``forward`` method and an + asynchronous ``forward_async`` method that returns a future of the ``forward`` call + on the user-provided module on the remote side. + """ + # NOTE: if a new attribute is added to this class, also need to add it + # to ``_REMOTE_MODULE_PICKLED_ATTRIBUTES`` for pickling/unpickling. + + remote_module = object.__new__(RemoteModule) + + enable_moving_cpu_tensors_to_cuda = remote_module._prepare_init(remote_device) + + if _module_interface_cls is not None: + # Users reply on this field to know if this generated RemoteModule is TorchScript-able. + remote_module.is_scriptable = True + + remote_module._init_template( + _module_interface_cls, enable_moving_cpu_tensors_to_cuda + ) + else: + remote_module.is_scriptable = False + remote_module.generated_methods = ( + _NON_SCRIPTABLE_REMOTE_MODULE_MODULE._generated_methods + ) + remote_module.module_rref = module_rref + + remote_module._install_generated_methods() + remote_module._check_attribute_picklability() + + return remote_module + + +class RemoteModule(_RemoteModule): + """ + A RemoteModule instance can only be created after RPC initialization. + + It creates a user-specified module on a specified remote node. + It behaves like a regular ``nn.Module`` except that the ``forward`` method is + executed on the remote node. + It takes care of autograd recording to ensure the backward pass propagates + gradients back to the corresponding remote module. + + It generates two methods ``forward_async`` and ``forward`` based on the + signature of the ``forward`` method of ``module_cls``. ``forward_async`` + runs asynchronously and returns a Future. The arguments of ``forward_async`` + and ``forward`` are the same as the ``forward`` method of the module + returned by the ``module_cls``. + + For example, if ``module_cls`` returns an instance of ``nn.Linear``, + that has ``forward`` method signature: ``def forward(input: Tensor) -> Tensor:``, + the generated ``RemoteModule`` will have 2 methods with the signatures: + + | ``def forward(input: Tensor) -> Tensor:`` + | ``def forward_async(input: Tensor) -> Future[Tensor]:`` + + Args: + remote_device (str): Device on the destination worker where we'd like to place this module. + The format should be "/", where the device field can be parsed as torch.device type. + E.g., "trainer0/cpu", "trainer0", "ps0/cuda:0". + In addition, the device field can be optional and the default value is "cpu". + module_cls (nn.Module): Class for the module to be created remotely. For example, + + >>> class MyModule(nn.Module): + >>> def forward(input): + >>> return input + 1 + >>> + >>> module_cls = MyModule + + args (Sequence, optional): args to be passed to ``module_cls``. + kwargs (Dict, optional): kwargs to be passed to ``module_cls``. + + Returns: + A remote module instance which wraps the :class:`~nn.Module` created by the + user-provided ``module_cls``, it has a blocking ``forward`` method and an + asynchronous ``forward_async`` method that returns a future of the ``forward`` call + on the user-provided module on the remote side. + + Example:: + Run the following code in two different processes: + + >>> # xdoctest: +SKIP("distributed") + >>> # On worker 0: + >>> import torch + >>> import torch.distributed.rpc as rpc + >>> from torch import nn, Tensor + >>> from torch.distributed.nn.api.remote_module import RemoteModule + >>> + >>> rpc.init_rpc("worker0", rank=0, world_size=2) + >>> remote_linear_module = RemoteModule( + >>> "worker1/cpu", nn.Linear, args=(20, 30), + >>> ) + >>> input = torch.randn(128, 20) + >>> ret_fut = remote_linear_module.forward_async(input) + >>> ret = ret_fut.wait() + >>> rpc.shutdown() + + >>> # On worker 1: + >>> import torch + >>> import torch.distributed.rpc as rpc + >>> + >>> rpc.init_rpc("worker1", rank=1, world_size=2) + >>> rpc.shutdown() + + Furthermore, a more practical example that is combined with + `DistributedDataParallel `__ (DDP) + can be found in this `tutorial `__. + """ + + def __init__( + self, + remote_device: str, + module_cls: Type[nn.Module], + args: Optional[Tuple] = None, + kwargs: Optional[Dict[str, Any]] = None, + ): + super().__init__(remote_device, module_cls, args, kwargs) + + +def _remote_module_receiver( + *remote_module_pickled_attrs, +): + """Deserializes a RemoteModule.""" + serialized_remote_module = _SerializedRemoteModule._make( + remote_module_pickled_attrs + ) + m = object.__new__(RemoteModule) + m.__dict__.update(serialized_remote_module._asdict()) + + # Unpickling the attribute `module_rref` must invoke RRef's `_deserialize()` method. + m.module_rref = rpc.PyRRef._deserialize(m.module_rref) + + # Install generated methods when unpickled. + for method in m.generated_methods: + method_name = method.__name__ + method = torch.jit.export(method) + setattr(m, method_name, types.MethodType(method, m)) + + return m + + +def _remote_module_reducer(remote_module): + """Serialize a RemoteModule.""" + pickled_attrs = {} + for k, v in remote_module.__dict__.items(): + # Pickling the attribute `module_rref` must invoke RRef's `_serialize()` method. + if k == "module_rref": + pickled_attrs[k] = v._serialize() + elif k in _REMOTE_MODULE_PICKLED_ATTRIBUTES: + pickled_attrs[k] = v + # Check if unpickled attributes are all in _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING. + elif k not in _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING: + print( + f"The new attribute ``{k}`` of RemoteModule is ignored during RPC pickling. " + "To pickle this attribute, please add it to ``_REMOTE_MODULE_PICKLED_ATTRIBUTES``. " + "Otherwise, please explicitly add it to ``_REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING``.", + file=sys.stderr, + ) + + return ( + _remote_module_receiver, + tuple(pickled_attrs.values()), + ) + + +def _recursive_script_module_receiver( + recursive_script_module_serialized, +): + """Deserializes a RecursiveScriptModule that does not contain a script RemoteModule.""" + f = io.BytesIO(recursive_script_module_serialized) + m = torch.jit.load(f) + return m + + +def _recursive_script_module_reducer(recursive_script_module): + """Serialize a RecursiveScriptModule that does not contain a script RemoteModule, and raises an error otherwise.""" + if hasattr(recursive_script_module._c, "module_rref"): + raise RuntimeError( + "Passing a script RemoteModule over RPC is not supported. Please create a RemoteModule in the sender, " + "send the `module_rref` to the receiver, and create a new instance on the receiver end by passing this `module_rref`." + ) + + f = io.BytesIO() + torch.jit.save(recursive_script_module, f) + return (_recursive_script_module_receiver, (f.getvalue(),)) + + +_internal_rpc_pickler._register_reducer(RemoteModule, _remote_module_reducer) +_internal_rpc_pickler._register_reducer( + torch.jit.RecursiveScriptModule, _recursive_script_module_reducer +) diff --git a/lib/python3.12/site-packages/torch/distributed/nn/functional.py b/lib/python3.12/site-packages/torch/distributed/nn/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..110df578552a59ab95f622ad5a306bc2a57f6457 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/nn/functional.py @@ -0,0 +1,452 @@ +# mypy: allow-untyped-defs +import torch +import torch.distributed as dist +from torch.autograd import Function + +# The two imports below are not always available depending on the +# USE_DISTRIBUTED compile flag. Make sure they raise import error +# if we're trying to use them. +from torch.distributed import group, ReduceOp + + +def broadcast(tensor, src, group=group.WORLD): + """ + Broadcasts the tensor to the whole group. + + ``tensor`` must have the same number of elements in all processes + participating in the collective. + + Arguments: + tensor (Tensor): Data to be sent if ``src`` is the rank of current + process. + src (int): Source rank. + group (ProcessGroup, optional): The process group to work on. + + Returns: + Tensor: Received tensor from the broadcast op. + + """ + return _Broadcast.apply(src, group, tensor) + + +def gather(tensor, dst=0, group=group.WORLD): + """ + Gathers a list of tensors in a single process. + + Arguments: + tensor (Tensor): Input tensor. + dst (int, optional): Destination rank (default is 0). + group (ProcessGroup, optional): The process group to work on. + + Returns: + tuple[Tensor]: List of appropriately-sized tensors with the gathered data. + """ + return _Gather.apply(dst, group, tensor) + + +def scatter(tensors, src=0, group=group.WORLD): + """ + Scatters a list of tensors to all processes in a group. + + Each process will receive exactly one tensor and store its data in the + ``tensor`` argument. + + Arguments: + tensors (list[Tensor]): List of tensors to scatter on the source rank. + Receivers must pass ``None`. + src (int, optional): Source rank (default is 0). + group (ProcessGroup, optional): The process group to work on. + + Returns: + Tensor: Output tensor from the scatter operation. + + """ + return _Scatter.apply(src, group, *tensors) + + +def reduce(tensor, dst, op=ReduceOp.SUM, group=group.WORLD): + """ + Reduces the tensor data across all machines. + + Only the process with rank ``dst`` is going to receive the final result. + + Arguments: + tensor (Tensor): Input of the collective. + dst (int): Destination rank. + op (optional): One of the values from + ``torch.distributed.ReduceOp`` + enum. Specifies an operation used for element-wise reductions. + group (ProcessGroup, optional): The process group to work on. + + Returns: + Tensor: Output of the collective. + + """ + return _Reduce.apply(dst, op, group, tensor) + + +def reduce_scatter(output, input_list, op=ReduceOp.SUM, group=group.WORLD): + """ + Reduces, then scatters a list of tensors to all processes in a group. + + Arguments: + output (Tensor): Output tensor. + input_list (list[Tensor]): List of tensors to reduce and scatter. + op (optional): One of the values from + ``torch.distributed.ReduceOp`` + enum. Specifies an operation used for element-wise reductions. + group (ProcessGroup, optional): The process group to work on. + + Returns: + Tensor: Output of the collective. + + """ + return _Reduce_Scatter.apply(op, group, output, *input_list) + + +def all_gather(tensor, group=group.WORLD): + """ + Gathers tensors from the whole group in a list. + + Arguments: + tensor (Tensor): Tensor to be broadcast from current process. + group (ProcessGroup, optional): The process group to work on. + + Returns: + tuple([Tensor]): Output of the collective. + + """ + return _AllGather.apply(group, tensor) + + +def _all_gather_base(output_tensor, input_tensor, group=group.WORLD): + """ + Single tensor all gather. Gathers a single tensor from all ranks, and puts them in a single output tensor. + + Args: + output_tensor (Tensor): Output tensor. It should contain + correctly-sized tensors to be used for output of the collective. + input_tensor (Tensor): Tensor to be broadcast from current process. + group (ProcessGroup, optional): The process group to work on. If None, + the default process group will be used. + + Examples: + >>> # All tensors below are of torch.int64 dtype. + >>> # We have 2 process groups, 2 ranks. + >>> # xdoctest: +SKIP("incorrect want text") + >>> output_tensor = torch.zeros(2, dtype=torch.int64) + >>> output_tensor + [tensor([0, 0])] # Rank 0 and 1 + >>> tensor = torch.arange(1, dtype=torch.int64) + 1 + rank + >>> tensor + tensor([1]) # Rank 0 + tensor([2]) # Rank 1 + >>> dist.all_gather_base(output_tensor, tensor) + >>> output_tensor + tensor([1,2]) # Rank 0 + tensor([1,2]) # Rank 1 + + .. warning:: + `_all_gather_base` is experimental and subject to change. + It is the caller's responsibility to ensure the output_tensor + is correctly sized. + + """ + return _AllGatherBase.apply(output_tensor, input_tensor, group) + + +def all_to_all(output_tensor_list, input_tensor_list, group=group.WORLD): + """ + Each process scatters list of input tensors to all processes in a group and return gathered list of tensors in output list. + + Arguments: + output_tensor_list (list[Tensor]): list of tensors to gather one per rank. + input_tensor_list (list[Tensor]): List of tensors to scatter one per rank. + group (ProcessGroup, optional): The process group to work on. + + Returns: + tuple([Tensor]): Output of the collective. + + """ + return _AlltoAll.apply(group, output_tensor_list, *input_tensor_list) + + +def all_to_all_single( + output, + input, + output_split_sizes=None, + input_split_sizes=None, + group=group.WORLD, +): + """ + Each process splits input tensor and then scatters the split list to all processes in a group. + + Then concatenate the received tensors from all the processes in the group and return single output tensor. + + Arguments: + output (Tensor): Gathered concatenated output tensor. + input (Tensor): Input tensor to scatter. + output_split_sizes: (list[Int], optional): Output split sizes for dim 0 + if specified None or empty, dim 0 of ``output`` tensor must divide + equally by ``world_size``. + input_split_sizes: (list[Int], optional): Input split sizes for dim 0 + if specified None or empty, dim 0 of ``input`` tensor must divide + equally by ``world_size``. + + Returns: + Tensor: Output of the collective. + + """ + return _AlltoAllSingle.apply( + group, output, output_split_sizes, input_split_sizes, input + ) + + +def all_reduce(tensor, op=ReduceOp.SUM, group=group.WORLD): + """ + Reduces the tensor data across all machines in such a way that all get the final result. + + After the call the returned tensor is going to be bitwise + identical in all processes. + + Arguments: + tensor (Tensor): Input of the collective. + op (optional): One of the values from + ``torch.distributed.ReduceOp`` + enum. Specifies an operation used for element-wise reductions. + group (ProcessGroup, optional): The process group to work on. + + Returns: + Tensor: Output of the collective + + """ + return _AllReduce.apply(op, group, tensor) + + +class _Broadcast(Function): + @staticmethod + def forward(ctx, src, group, tensor): + ctx.src = src + ctx.group = group + ctx.rank = dist.get_rank(group=group) + # torch.distributed makes all the calls in place + # we allocate new tensors to avoid this + tensor = tensor.clone() + dist.broadcast(tensor, src, group=group) + return tensor + + @staticmethod + def backward(ctx, grad_output): + gx = _Reduce.apply(ctx.src, ReduceOp.SUM, ctx.group, grad_output) + if ctx.src != ctx.rank: + gx.zero_() + return (None, None, gx) + + +class _Gather(Function): + @staticmethod + def forward(ctx, dst, group, tensor): + ctx.dst = dst + ctx.group = group + # Need to create a list of tensors here to do the + # aggregation, get it from the group size + # tensor should be correctly sized for the method + # gathering + tensor_list = [ + torch.zeros_like(tensor) for i in range(dist.get_world_size(group=group)) + ] + + tensor = tensor.contiguous() + if dist.get_rank(group=group) == dst: + dist.gather(tensor, tensor_list, dst, group=group) + else: + dist.gather(tensor, None, dst, group=group) + return tuple(tensor_list) + + @staticmethod + def backward(ctx, *grad_outputs): + return (None, None) + (_Scatter.apply(ctx.dst, ctx.group, *grad_outputs),) + + +class _Scatter(Function): + @staticmethod + def forward(ctx, src, group, *tensors): + ctx.src = src + ctx.group = group + assert all(t.size() == tensors[0].size() for t in tensors) + output = torch.zeros_like(tensors[0]) + if dist.get_rank(group=group) == src: + dist.scatter(output, list(tensors), src, group=group) + else: + dist.scatter(output, None, src, group=group) + return output + + @staticmethod + def backward(ctx, grad_output): + return (None, None) + _Gather.apply(ctx.src, ctx.group, grad_output) + + +class _Reduce(Function): + @staticmethod + def forward(ctx, src, op, group, tensor): + ctx.src = src + ctx.group = group + tensor = tensor.clone() + dist.reduce(tensor, src, op=op, group=group) + return tensor + + @staticmethod + def backward(ctx, grad_output): + return (None, None, None) + (_Broadcast.apply(ctx.src, ctx.group, grad_output),) + + +class _Reduce_Scatter(Function): + @staticmethod + def forward(ctx, op, group, tensor, *input_tensor_list): + ctx.group = group + # Need contiguous tensors for collectives. + tensor = tensor.contiguous() + input_tensor_list = tuple(t.contiguous() for t in input_tensor_list) + dist.reduce_scatter(tensor, list(input_tensor_list), op=op, group=group) + return tensor + + @staticmethod + def backward(ctx, grad_output): + return (None, None, None) + _AllGather.apply(ctx.group, grad_output) + + +class _AllGather(Function): + @staticmethod + def forward(ctx, group, tensor): + # Need contiguous tensors for collectives. + tensor = tensor.contiguous() + + ctx.group = group + out_tensor_list = [ + torch.empty_like(tensor) for _ in range(dist.get_world_size(group=group)) + ] + + dist.all_gather(out_tensor_list, tensor, group=group) + return tuple(out_tensor_list) + + @staticmethod + def backward(ctx, *grad_outputs): + if dist.get_backend(group=ctx.group) is dist.Backend.NCCL: + rank = dist.get_rank(group=ctx.group) + gx = torch.empty_like(grad_outputs[rank]) + gx = _Reduce_Scatter.apply(ReduceOp.SUM, ctx.group, gx, *grad_outputs) + else: + # As many backends doesn't support ReduceScatter, we use AlltoAll with .sum() + # to emulate the ReduceScatter behavior + tensor_list = [torch.empty_like(tensor) for tensor in grad_outputs] + gxs = _AlltoAll.apply(ctx.group, tensor_list, *grad_outputs) + gx = torch.sum(torch.stack(gxs), dim=0) + return (None, gx) + + +class _AllGatherBase(Function): + @staticmethod + def forward(ctx, output_tensor, input_tensor, group): + ctx.group = group + dist._all_gather_base(output_tensor, input_tensor.contiguous(), group=group) + return output_tensor + + @staticmethod + def backward(ctx, grad_output): + if dist.get_backend(group=ctx.group) is dist.Backend.NCCL: + world_size = dist.get_world_size(group=ctx.group) + out_size = list(grad_output.size()) + if out_size[0] % world_size != 0: + raise RuntimeError( + f"Tensor with dimensions: {out_size} does " + f"not have first dimension divisible by world_size: {world_size}" + ) + out_size[0] = out_size[0] // dist.get_world_size(group=ctx.group) + gx = torch.empty( + out_size, device=grad_output.device, dtype=grad_output.dtype + ) + dist._reduce_scatter_base(gx, grad_output, ReduceOp.SUM, ctx.group) + else: + raise RuntimeError("Backend not supported!") + return (None, gx, None) + + +class _AlltoAll(Function): + @staticmethod + def forward(ctx, group, out_tensor_list, *tensors): + ctx.group = group + ctx.input_tensor_size_list = [ + tensors[i].size() for i in range(dist.get_world_size(group=group)) + ] + my_rank = dist.get_rank(group=group) + tensors = tuple(t.contiguous() for t in tensors) + # Implement it on means of scatter/gather, send/recv async operations have issues + if dist.get_backend(group=group) is dist.Backend.GLOO: + for i in range(dist.get_world_size(group=group)): + to_send = None + if i == my_rank: + to_send = list(tensors) + dist.scatter(out_tensor_list[i], to_send, i, group=group) + else: + dist.all_to_all( + out_tensor_list, + list(tensors), + group=group, + ) + return tuple(out_tensor_list) + + @staticmethod + def backward(ctx, *grad_outputs): + tensor_list = [ + torch.empty( + size, device=grad_outputs[0].device, dtype=grad_outputs[0].dtype + ) + for size in ctx.input_tensor_size_list + ] + return (None, None) + _AlltoAll.apply(ctx.group, tensor_list, *grad_outputs) + + +class _AlltoAllSingle(Function): + @staticmethod + def forward(ctx, group, output, output_split_sizes, input_split_sizes, input): + ctx.group = group + ctx.input_size = input.size() + ctx.output_split_sizes = input_split_sizes + ctx.input_split_sizes = output_split_sizes + dist.all_to_all_single( + output, + input, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + ) + return output + + @staticmethod + def backward(ctx, grad_output): + tensor = torch.empty( + ctx.input_size, device=grad_output.device, dtype=grad_output.dtype + ) + return (None, None, None, None) + ( + _AlltoAllSingle.apply( + ctx.group, + tensor, + ctx.output_split_sizes, + ctx.input_split_sizes, + grad_output.contiguous(), + ), + ) + + +class _AllReduce(Function): + @staticmethod + def forward(ctx, op, group, tensor): + ctx.group = group + ctx.op = op + tensor = tensor.clone() + dist.all_reduce(tensor, op=op, group=group) + return tensor + + @staticmethod + def backward(ctx, grad_output): + return (None, None) + (_AllReduce.apply(ctx.op, ctx.group, grad_output),) diff --git a/lib/python3.12/site-packages/torch/distributed/nn/jit/__init__.py b/lib/python3.12/site-packages/torch/distributed/nn/jit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/distributed/nn/jit/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/nn/jit/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..827ff5f8c0717e2df59d88d85e30a75a444ceef7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/nn/jit/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/nn/jit/__pycache__/instantiator.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/nn/jit/__pycache__/instantiator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f66600283ecd3fa29229dee349fb868e7bc907dd Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/nn/jit/__pycache__/instantiator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/nn/jit/instantiator.py b/lib/python3.12/site-packages/torch/distributed/nn/jit/instantiator.py new file mode 100644 index 0000000000000000000000000000000000000000..d529fc740945a058c11bcc87ddf766c198a3ebfe --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/nn/jit/instantiator.py @@ -0,0 +1,154 @@ +#!/usr/bin/python3 +# mypy: allow-untyped-defs +import importlib +import logging +import os +import sys +import tempfile +from typing import Optional + +import torch +from torch.distributed.nn.jit.templates.remote_module_template import ( + get_remote_module_template, +) + + +logger = logging.getLogger(__name__) + + +_FILE_PREFIX = "_remote_module_" +_TEMP_DIR = tempfile.TemporaryDirectory() +INSTANTIATED_TEMPLATE_DIR_PATH = _TEMP_DIR.name +logger.info("Created a temporary directory at %s", INSTANTIATED_TEMPLATE_DIR_PATH) +sys.path.append(INSTANTIATED_TEMPLATE_DIR_PATH) + + +def get_arg_return_types_from_interface(module_interface): + assert getattr( + module_interface, "__torch_script_interface__", False + ), "Expect a TorchScript class interface decorated by @torch.jit.interface." + qualified_name = torch._jit_internal._qualified_name(module_interface) + cu = torch.jit._state._python_cu + module_interface_c = cu.get_interface(qualified_name) + assert ( + "forward" in module_interface_c.getMethodNames() + ), f"Expect forward in interface methods, while it has {module_interface_c.getMethodNames()}" + method_schema = module_interface_c.getMethod("forward") + + arg_str_list = [] + arg_type_str_list = [] + assert method_schema is not None + for argument in method_schema.arguments: + arg_str_list.append(argument.name) + + if argument.has_default_value(): + default_value_str = f" = {argument.default_value}" + else: + default_value_str = "" + arg_type_str = f"{argument.name}: {argument.type}{default_value_str}" + arg_type_str_list.append(arg_type_str) + + arg_str_list = arg_str_list[1:] # Remove "self". + args_str = ", ".join(arg_str_list) + + arg_type_str_list = arg_type_str_list[1:] # Remove "self". + arg_types_str = ", ".join(arg_type_str_list) + + assert len(method_schema.returns) == 1 + argument = method_schema.returns[0] + return_type_str = str(argument.type) + + return args_str, arg_types_str, return_type_str + + +def _write(out_path, text): + old_text: Optional[str] + try: + with open(out_path) as f: + old_text = f.read() + except OSError: + old_text = None + if old_text != text: + with open(out_path, "w") as f: + logger.info("Writing %s", out_path) + f.write(text) + else: + logger.info("Skipped writing %s", out_path) + + +def _do_instantiate_remote_module_template( + generated_module_name, str_dict, enable_moving_cpu_tensors_to_cuda +): + generated_code_text = get_remote_module_template( + enable_moving_cpu_tensors_to_cuda + ).format(**str_dict) + out_path = os.path.join( + INSTANTIATED_TEMPLATE_DIR_PATH, f"{generated_module_name}.py" + ) + _write(out_path, generated_code_text) + + # From importlib doc, + # > If you are dynamically importing a module that was created since + # the interpreter began execution (e.g., created a Python source file), + # you may need to call invalidate_caches() in order for the new module + # to be noticed by the import system. + importlib.invalidate_caches() + generated_module = importlib.import_module(f"{generated_module_name}") + return generated_module + + +def instantiate_scriptable_remote_module_template( + module_interface_cls, enable_moving_cpu_tensors_to_cuda=True +): + if not getattr(module_interface_cls, "__torch_script_interface__", False): + raise ValueError( + f"module_interface_cls {module_interface_cls} must be a type object decorated by " + "@torch.jit.interface" + ) + + # Generate the template instance name. + module_interface_cls_name = torch._jit_internal._qualified_name( + module_interface_cls + ).replace(".", "_") + generated_module_name = f"{_FILE_PREFIX}{module_interface_cls_name}" + + # Generate type annotation strs. + assign_module_interface_cls_str = ( + f"from {module_interface_cls.__module__} import " + f"{module_interface_cls.__name__} as module_interface_cls" + ) + args_str, arg_types_str, return_type_str = get_arg_return_types_from_interface( + module_interface_cls + ) + kwargs_str = "" + arrow_and_return_type_str = f" -> {return_type_str}" + arrow_and_future_return_type_str = f" -> Future[{return_type_str}]" + + str_dict = dict( + assign_module_interface_cls=assign_module_interface_cls_str, + arg_types=arg_types_str, + arrow_and_return_type=arrow_and_return_type_str, + arrow_and_future_return_type=arrow_and_future_return_type_str, + args=args_str, + kwargs=kwargs_str, + jit_script_decorator="@torch.jit.script", + ) + return _do_instantiate_remote_module_template( + generated_module_name, str_dict, enable_moving_cpu_tensors_to_cuda + ) + + +def instantiate_non_scriptable_remote_module_template(): + generated_module_name = f"{_FILE_PREFIX}non_scriptable" + str_dict = dict( + assign_module_interface_cls="module_interface_cls = None", + args="*args", + kwargs="**kwargs", + arg_types="*args, **kwargs", + arrow_and_return_type="", + arrow_and_future_return_type="", + jit_script_decorator="", + ) + # For a non-scriptable template, always enable moving CPU tensors to a cuda device, + # because there is no syntax limitation on the extra handling caused by the script. + return _do_instantiate_remote_module_template(generated_module_name, str_dict, True) diff --git a/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/__init__.py b/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..674eb9650359ca1aad144c2929964138509a6937 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/__pycache__/remote_module_template.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/__pycache__/remote_module_template.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..764301d708715cda6f73760f3c84c64f82c663a3 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/__pycache__/remote_module_template.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/remote_module_template.py b/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/remote_module_template.py new file mode 100644 index 0000000000000000000000000000000000000000..07b055774b36af4835e308c8a1f85afd0ab35f0f --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/nn/jit/templates/remote_module_template.py @@ -0,0 +1,108 @@ +#!/usr/bin/python3 +# mypy: allow-untyped-defs + + +def get_remote_module_template(enable_moving_cpu_tensors_to_cuda: bool): + return _TEMPLATE_PREFIX + ( + _REMOTE_FORWARD_TEMPLATE_ENABLE_MOVING_CPU_TENSORS_TO_CUDA + if enable_moving_cpu_tensors_to_cuda + else _REMOTE_FORWARD_TEMPLATE + ) + + +_TEMPLATE_PREFIX = """from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +{assign_module_interface_cls} + + +def forward_async(self, {arg_types}){arrow_and_future_return_type}: + args = (self.module_rref, self.device, self.is_device_map_set, {args}) + kwargs = {{{kwargs}}} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, {arg_types}){arrow_and_return_type}: + args = (self.module_rref, self.device, self.is_device_map_set, {args}) + kwargs = {{{kwargs}}} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + +{jit_script_decorator} +""" + +# This template may cause typing error (the mismatch between ``Tuple[()]`` and ``Tuple[Any]``) +# even if the code is only used for instantiation but not execution. +# Therefore, only include handling moving CPU tensors to a cuda device if necessary. +# TODO: Merge these two templates together in the future once TorchScript syntax is improved. +_REMOTE_FORWARD_TEMPLATE_ENABLE_MOVING_CPU_TENSORS_TO_CUDA = """ +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, {arg_types}){arrow_and_return_type}: + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward({args}, {kwargs}) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = ({args},) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {{{kwargs}}} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, {kwargs}) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, {kwargs}))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, {kwargs}): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret +""" + +_REMOTE_FORWARD_TEMPLATE = """ +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, {arg_types}){arrow_and_return_type}: + module = module_rref.local_value() + + return module.forward({args}, {kwargs}) +""" diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__init__.py b/lib/python3.12/site-packages/torch/distributed/optim/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d8fee468504de67bc02601f9775623acc8935f55 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/__init__.py @@ -0,0 +1,43 @@ +""" +:mod:`torch.distributed.optim` exposes DistributedOptimizer, which takes a list +of remote parameters (:class:`~torch.distributed.rpc.RRef`) and runs the +optimizer locally on the workers where the parameters live. The distributed +optimizer can use any of the local optimizer :ref:`optimizer-algorithms` to +apply the gradients on each worker. +""" +import warnings + +import torch +from torch import optim + +from .apply_optimizer_in_backward import ( + _apply_optimizer_in_backward, + _get_in_backward_optimizers, +) +from .functional_adadelta import _FunctionalAdadelta +from .functional_adagrad import _FunctionalAdagrad +from .functional_adam import _FunctionalAdam +from .functional_adamax import _FunctionalAdamax +from .functional_adamw import _FunctionalAdamW +from .functional_rmsprop import _FunctionalRMSprop +from .functional_rprop import _FunctionalRprop +from .functional_sgd import _FunctionalSGD +from .named_optimizer import _NamedOptimizer +from .utils import as_functional_optim + + +# DistributedOptimizer imports torch.distributed.rpc names, so gate availability +# based on RPC being available. +if hasattr(torch._C, "_rpc_init"): + from .optimizer import DistributedOptimizer + +from .post_localSGD_optimizer import PostLocalSGDOptimizer +from .zero_redundancy_optimizer import ZeroRedundancyOptimizer + + +__all__ = [ + "as_functional_optim", + "DistributedOptimizer", + "PostLocalSGDOptimizer", + "ZeroRedundancyOptimizer", +] diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7dcce4c6224152408e05d89e73d5c5011e2e1483 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/_deprecation_warning.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/_deprecation_warning.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..696199a0e85e87d90bcae68bb7a1954f6e1a626a Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/_deprecation_warning.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/apply_optimizer_in_backward.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/apply_optimizer_in_backward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b98b7f824ecf8fdd05e34fc7c5630dc3b00c99e1 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/apply_optimizer_in_backward.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adadelta.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adadelta.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..18b0142ee146bbdc214b04839a035904c135488d Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adadelta.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adagrad.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adagrad.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c2efc1e6ace51c7d234710ef96cefae87e1b607 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adagrad.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adam.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f171b4718506d6546a3a04ecfbcb35c897533016 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adamax.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adamax.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..31d12d094c4816877f648cf3dc777eb51389ecc2 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adamax.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adamw.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adamw.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..31ebcfd51dfbeada84c2dc158f63117d5fcef140 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_adamw.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_rmsprop.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_rmsprop.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ad448e7e5fc5baf50e3cccfb2d671e5d2a327b5 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_rmsprop.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_rprop.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_rprop.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42bcf8a598bb6c54946ed1bae3c7d10e78fa4dc2 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_rprop.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_sgd.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_sgd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f297864b3197380e3489c2a7511c4d7f3b43a4e Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/functional_sgd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/named_optimizer.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/named_optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fab72eb6c4cde92c03b6a198b2f0ee75b574cd77 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/named_optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/optimizer.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9cb9a2635180e883f554ff0235dbd1fb1abd54a6 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/post_localSGD_optimizer.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/post_localSGD_optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b225702b7397b742629f2c7f3380960cbf2f0eb8 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/post_localSGD_optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eba4d670325e9fd82d83669b6fdbea680273e0ca Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/zero_redundancy_optimizer.cpython-312.pyc b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/zero_redundancy_optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b2f4ae4316a12e3a547f773ca05b01e4896b2f4 Binary files /dev/null and b/lib/python3.12/site-packages/torch/distributed/optim/__pycache__/zero_redundancy_optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/distributed/optim/_deprecation_warning.py b/lib/python3.12/site-packages/torch/distributed/optim/_deprecation_warning.py new file mode 100644 index 0000000000000000000000000000000000000000..c3434a4cd4f081843295e488c18a67a5c297fcbf --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/_deprecation_warning.py @@ -0,0 +1,16 @@ +import warnings + +import torch + + +@torch.jit.ignore # type: ignore[misc] +def _scripted_functional_optimizer_deprecation_warning(stacklevel: int = 0) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("always") + warnings.warn( + "`TorchScript` support for functional optimizers is deprecated " + "and will be removed in a future PyTorch release. " + "Consider using the `torch.compile` optimizer instead.", + DeprecationWarning, + stacklevel=stacklevel + 2, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/apply_optimizer_in_backward.py b/lib/python3.12/site-packages/torch/distributed/optim/apply_optimizer_in_backward.py new file mode 100644 index 0000000000000000000000000000000000000000..36f679f4eba49b378036bae2f8dc89e7278c91a4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/apply_optimizer_in_backward.py @@ -0,0 +1,120 @@ +from typing import Any, Dict, Iterable, List, no_type_check, Type + +import torch + + +__all__: List[str] = [] + +# WeakTensorKeyDictionary to store relevant meta-data for the Tensor/Parameter +# without changing it's life-time. +# NOTE: Alternative is to add the meta-data as an attribute to the tensor, +# but that will serialize the meta-data if Tensor is serialized. +param_to_optim_hook_handle_map = torch.utils.weak.WeakTensorKeyDictionary() +param_to_acc_grad_map = torch.utils.weak.WeakTensorKeyDictionary() + + +@no_type_check +def _apply_optimizer_in_backward( + optimizer_class: Type[torch.optim.Optimizer], + params: Iterable[torch.nn.Parameter], + optimizer_kwargs: Dict[str, Any], + register_hook: bool = True, +) -> None: + """ + Upon ``backward()``, the optimizer specified for each parameter will fire after + the gradient has been accumulated into the parameter. + + Note - gradients for these parameters will be set to None after ``backward()``. + This means that any other optimizer not specified via `_apply_optimizer_in_backward` + over this parameter will be a no-op. + + Args: + optimizer_class: (Type[torch.optim.Optimizer]): Optimizer to apply to parameter + params: (Iterator[nn.Parameter]): parameters to apply optimizer state to + optimizer_kwargs: (Dict[str, Any]): kwargs to pass to optimizer constructor + register_hook: (bool): whether to register a hook that runs the optimizer + after gradient for this parameter is accumulated. This is the default + way that optimizer in backward is implemented, but specific use cases + (such as DDP) may wish to override this to implement custom behavior. + (Default = True) + + Example:: + params_generator = model.parameters() + param_1 = next(params_generator) + remainder_params = list(params_generator) + + apply_optimizer_in_backward(torch.optim.SGD, [param_1], {"lr": .02}) + apply_optimizer_in_backward(torch.optim.Adam, remainder_params, {"lr": .04}) + + model(...).sum().backward() # after backward, parameters will already + # have their registered optimizer(s) applied. + + """ + torch._C._log_api_usage_once("torch.distributed.optim.apply_optimizer_in_backward") + + @no_type_check + def _apply_optimizer_in_backward_to_param(param: torch.nn.Parameter) -> None: + # view_as creates a node in autograd graph that allows us access to the + # parameter's AccumulateGrad autograd function object. We register a + # hook on this object to fire the optimizer when the gradient for + # this parameter is ready (has been accumulated into .grad field) + + # Don't create a new acc_grad if we already have one + # i.e. for shared parameters or attaching multiple optimizers to a param. + if param not in param_to_acc_grad_map: + param_to_acc_grad_map[param] = param.view_as(param).grad_fn.next_functions[ + 0 + ][0] + + optimizer = optimizer_class([param], **optimizer_kwargs) + + if not hasattr(param, "_in_backward_optimizers"): + param._in_backward_optimizers = [] # type: ignore[attr-defined] + # TODO: Remove these attributes once we have a better way of accessing + # optimizer classes and kwargs for a parameter. + param._optimizer_classes = [] # type: ignore[attr-defined] + param._optimizer_kwargs = [] # type: ignore[attr-defined] + + param._in_backward_optimizers.append(optimizer) # type: ignore[attr-defined] + param._optimizer_classes.append(optimizer_class) # type: ignore[attr-defined] + param._optimizer_kwargs.append(optimizer_kwargs) # type: ignore[attr-defined] + + if not register_hook: + return + + def optimizer_hook(*_unused) -> None: + for opt in param._in_backward_optimizers: # type: ignore[attr-defined] + opt.step() + + param.grad = None + + handle = param_to_acc_grad_map[param].register_hook(optimizer_hook) # type: ignore[attr-defined] + if param not in param_to_optim_hook_handle_map: + param_to_optim_hook_handle_map[param] = [] + param_to_optim_hook_handle_map[param].append(handle) + + for param in params: + _apply_optimizer_in_backward_to_param(param) + + +def _get_in_backward_optimizers(module: torch.nn.Module) -> List[torch.optim.Optimizer]: + """ + Return a list of in-backward optimizers applied to ``module``'s parameters. Note that these + optimizers are not intended to directly have their ``step`` or ``zero_grad`` methods called + by the user and are intended to be used for things like checkpointing. + + Args: + module: (torch.nn.Module): model to retrieve in-backward optimizers for + + Returns: + List[torch.optim.Optimizer]: the in-backward optimizers. + + Example:: + _apply_optimizer_in_backward(torch.optim.SGD, model.parameters(), {'lr': 0.01}) + optims = _get_optimizers_in_backward(model) + """ + optims: List[torch.optim.Optimizer] = [] + for param in module.parameters(): + optims.extend(getattr(param, "_in_backward_optimizers", [])) + + return optims diff --git a/lib/python3.12/site-packages/torch/distributed/optim/functional_adadelta.py b/lib/python3.12/site-packages/torch/distributed/optim/functional_adadelta.py new file mode 100644 index 0000000000000000000000000000000000000000..81bff38773a90244d59094ef261503545309c2a2 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/functional_adadelta.py @@ -0,0 +1,111 @@ +# mypy: allow-untyped-defs +from typing import Dict, List, Optional + +import torch +import torch.optim._functional as F +from torch import Tensor +from torch.distributed.optim._deprecation_warning import ( + _scripted_functional_optimizer_deprecation_warning, +) + + +__all__: List[str] = [] + + +# Define a TorchScript compatible Functional Adadelta Optimizer +# where we use these optimizer in a functional way. +# Instead of using the `param.grad` when updating parameters, +# we explicitly allow the distributed optimizer pass gradients to +# the `step` function. In this way, we could separate the gradients +# and parameters and allow multithreaded trainer to update the +# parameters without data traces on accumulating to the same .grad. +# NOTE: This should be only used by distributed optimizer internals +# and not meant to expose to the user. +@torch.jit.script +class _FunctionalAdadelta: + def __init__( + self, + params: List[Tensor], + lr: float = 1.0, + rho: float = 0.9, + eps: float = 1e-6, + weight_decay: float = 0.0, + foreach: bool = False, + maximize: bool = False, + _allow_empty_param_list: bool = False, + ): + _scripted_functional_optimizer_deprecation_warning(stacklevel=2) + self.defaults = { + "lr": lr, + "rho": rho, + "eps": eps, + "weight_decay": weight_decay, + } + self.foreach = foreach + self.maximize = maximize + + if len(params) == 0 and not _allow_empty_param_list: + raise ValueError("optimizer got an empty parameter list") + + # NOTE: we only have one param_group and don't allow user to add additional + # param group as it's not a common use case. + self.param_group = {"params": params} + + self.state = torch.jit.annotate(Dict[torch.Tensor, Dict[str, torch.Tensor]], {}) + + def step(self, gradients: List[Optional[Tensor]]): + params = self.param_group["params"] + params_with_grad = [] + grads = [] + square_avgs = [] + acc_deltas = [] + state_steps = [] + lr = self.defaults["lr"] + rho = self.defaults["rho"] + eps = self.defaults["eps"] + weight_decay = self.defaults["weight_decay"] + + if len(params) != len(gradients): + raise ValueError( + "the gradients passed in does not equal to the size of the parameters!" + + f"Params length: {len(params)}. " + + f"Gradients length: {len(gradients)}" + ) + has_complex = False + for param, gradient in zip(params, gradients): + if gradient is not None: + has_complex |= torch.is_complex(param) + params_with_grad.append(param) + grads.append(gradient) + # Lazy state initialization + if param not in self.state: + self.state[param] = {} + state = self.state[param] + state["step"] = torch.tensor(0.0) + state["square_avg"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + state["acc_delta"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + + state = self.state[param] + square_avgs.append(state["square_avg"]) + acc_deltas.append(state["acc_delta"]) + state_steps.append(state["step"]) + + with torch.no_grad(): + F.adadelta( + params_with_grad, + grads, + square_avgs, + acc_deltas, + state_steps, + lr=lr, + rho=rho, + eps=eps, + weight_decay=weight_decay, + foreach=self.foreach, + maximize=self.maximize, + has_complex=has_complex, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/functional_adagrad.py b/lib/python3.12/site-packages/torch/distributed/optim/functional_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..2c8ac898a45f26312e04dfdba6eb2a8c3f15f6ad --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/functional_adagrad.py @@ -0,0 +1,115 @@ +# mypy: allow-untyped-defs +from typing import Dict, List, Optional + +import torch +import torch.optim._functional as F +from torch import Tensor +from torch.distributed.optim._deprecation_warning import ( + _scripted_functional_optimizer_deprecation_warning, +) + + +__all__: List[str] = [] + + +# Define a TorchScript compatible Functional Adagrad Optimizer +# where we use these optimizer in a functional way. +# Instead of using the `param.grad` when updating parameters, +# we explicitly let the user pass gradients to the `step` function +# this is so that we could separate the gradients and parameters +# and allow multithreaded trainer to update the parameters +# without data traces on accumulating to the same .grad. +# NOTE: This should be only used by distributed optimizer internals +# and not meant to expose to the user. +@torch.jit.script +class _FunctionalAdagrad: + def __init__( + self, + params: List[Tensor], + lr: float = 1e-2, + lr_decay: float = 0.0, + weight_decay: float = 0.0, + initial_accumulator_value: float = 0.0, + warmup_lr_multiplier: float = 1.0, + warmup_num_iters: float = 0.0, + eps: float = 1e-10, + coalesce_grad: bool = True, + foreach: bool = False, + fused: bool = False, + maximize: bool = False, + _allow_empty_param_list: bool = False, + ): + _scripted_functional_optimizer_deprecation_warning(stacklevel=2) + self.defaults = { + "lr": lr, + "lr_decay": lr_decay, + "eps": eps, + "weight_decay": weight_decay, + "initial_accumulator_value": initial_accumulator_value, + "warmup_lr_multiplier": warmup_lr_multiplier, + "warmup_num_iters": warmup_num_iters, + } + self.coalesce_grad = coalesce_grad + self.foreach = foreach + self.fused = fused + self.maximize = maximize + self.state = torch.jit.annotate(Dict[torch.Tensor, Dict[str, torch.Tensor]], {}) + + if len(params) == 0 and not _allow_empty_param_list: + raise ValueError("optimizer got an empty parameter list") + + # NOTE: we only have one param_group and don't allow user to add additional + # param group as it's not a common use case. + self.param_group = {"params": params} + + # TODO: no union or any types in TorchScript, make step a scalar tensor instead + # This is also needed by if we want to share_memory on the step across processes + for p in self.param_group["params"]: + self.state[p] = { + "sum": torch.full_like(p.data, initial_accumulator_value), + "step": torch.tensor(0.0), + } + + def step(self, gradients: List[Optional[Tensor]]): + params = self.param_group["params"] + params_with_grad = [] + grads = [] + state_sums = [] + state_steps: List[Tensor] = [] + + if len(params) != len(gradients): + raise ValueError( + "the gradients passed in does not equal to the size of the parameters!" + + f"Params length: {len(params)}. " + + f"Gradients length: {len(gradients)}" + ) + + has_sparse_grad, has_complex = False, False + for param, gradient in zip(self.param_group["params"], gradients): + if gradient is not None: + has_sparse_grad |= gradient.is_sparse + has_complex |= torch.is_complex(param) + params_with_grad.append(param) + grads.append(gradient) + state = self.state[param] + state_sums.append(state["sum"]) + state_steps.append(state["step"]) + + with torch.no_grad(): + F.adagrad( + params, + grads, + state_sums, + state_steps, + lr=self.defaults["lr"], + weight_decay=self.defaults["weight_decay"], + lr_decay=self.defaults["lr_decay"], + eps=self.defaults["eps"], + has_sparse_grad=has_sparse_grad, + foreach=self.foreach, + maximize=self.maximize, + has_complex=has_complex, + fused=self.fused, + grad_scale=None, + found_inf=None, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/functional_adam.py b/lib/python3.12/site-packages/torch/distributed/optim/functional_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..d52d2c90260e52cfe180b3d5a0b0e1d14704fe3b --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/functional_adam.py @@ -0,0 +1,202 @@ +# mypy: allow-untyped-defs +from typing import Dict, List, Optional, Tuple + +import torch +import torch.optim._functional as F +from torch import Tensor +from torch.distributed.optim._deprecation_warning import ( + _scripted_functional_optimizer_deprecation_warning, +) + + +__all__: List[str] = [] + + +# Define a TorchScript compatible Functional Adam Optimizer +# where we use these optimizer in a functional way. +# Instead of using the `param.grad` when updating parameters, +# we explicitly allow the distributed optimizer pass gradients to +# the `step` function. In this way, we could separate the gradients +# and parameters and allow multithreaded trainer to update the +# parameters without data traces on accumulating to the same .grad. +# NOTE: This should be only used by distributed optimizer internals +# and not meant to expose to the user. +@torch.jit.script +class _FunctionalAdam: + def __init__( + self, + params: List[Tensor], + lr: float = 1e-3, + betas: Tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0.0, + amsgrad: bool = False, + maximize: bool = False, + foreach: bool = False, + fused: bool = False, + _allow_empty_param_list: bool = False, + ): + _scripted_functional_optimizer_deprecation_warning(stacklevel=2) + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + + self.defaults = { + "lr": lr, + "eps": eps, + "beta1": betas[0], + "beta2": betas[1], + "weight_decay": weight_decay, + } + self.amsgrad = amsgrad + self.maximize = maximize + self.foreach = foreach + self.fused = fused + self.state = torch.jit.annotate(Dict[torch.Tensor, Dict[str, torch.Tensor]], {}) + + if len(params) == 0 and not _allow_empty_param_list: + raise ValueError("optimizer got an empty parameter list") + + # NOTE: we only have one param_group and don't allow user to add additional + # param group as it's not a common use case. + self.param_group = {"params": params} + + def step_param(self, param: Tensor, grad: Optional[Tensor]): + """ + Similar to step, but operates on a single parameter and optionally a + gradient tensor. + """ + params_with_grad = [] + grads = [] + exp_avgs = [] + exp_avg_sqs = [] + max_exp_avg_sqs = [] + state_steps: List[Tensor] = [] + has_complex = torch.is_complex(param) + if grad is not None: + params_with_grad.append(param) + grads.append(grad) + if param not in self.state: + self.state[param] = {} + state = self.state[param] + state["step"] = torch.tensor(0.0) + state["exp_avg"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + state["exp_avg_sq"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + if self.amsgrad: + state["max_exp_avg_sq"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + + state = self.state[param] + exp_avgs.append(state["exp_avg"]) + exp_avg_sqs.append(state["exp_avg_sq"]) + + if self.amsgrad: + max_exp_avg_sqs.append(state["max_exp_avg_sq"]) + + state_steps.append(state["step"]) + with torch.no_grad(): + F.adam( + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + amsgrad=self.amsgrad, + has_complex=has_complex, + maximize=self.maximize, + beta1=self.defaults["beta1"], + beta2=self.defaults["beta2"], + lr=self.defaults["lr"], + weight_decay=self.defaults["weight_decay"], + eps=self.defaults["eps"], + foreach=self.foreach, + fused=self.fused, + grad_scale=None, + found_inf=None, + ) + + def step(self, gradients: List[Optional[Tensor]]): + params = self.param_group["params"] + params_with_grad = [] + grads = [] + exp_avgs = [] + exp_avg_sqs = [] + max_exp_avg_sqs = [] + state_steps: List[Tensor] = [] + has_complex = False + + if len(params) != len(gradients): + raise ValueError( + "the gradients passed in does not equal to the size of the parameters!" + + f"Params length: {len(params)}. " + + f"Gradients length: {len(gradients)}" + ) + + for param, gradient in zip(self.param_group["params"], gradients): + if gradient is not None: + has_complex |= torch.is_complex(param) + params_with_grad.append(param) + grads.append(gradient) + # Lazy state initialization + if param not in self.state: + self.state[param] = {} + state = self.state[param] + state["step"] = torch.tensor(0.0) + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + if self.amsgrad: + # Maintains max of all exp. moving avg. of sq. grad. values + state["max_exp_avg_sq"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + + state = self.state[param] + + exp_avgs.append(state["exp_avg"]) + exp_avg_sqs.append(state["exp_avg_sq"]) + + if self.amsgrad: + max_exp_avg_sqs.append(state["max_exp_avg_sq"]) + + state_steps.append(state["step"]) + + with torch.no_grad(): + F.adam( + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + amsgrad=self.amsgrad, + has_complex=has_complex, + maximize=self.maximize, + beta1=self.defaults["beta1"], + beta2=self.defaults["beta2"], + lr=self.defaults["lr"], + weight_decay=self.defaults["weight_decay"], + eps=self.defaults["eps"], + foreach=self.foreach, + fused=self.fused, + grad_scale=None, + found_inf=None, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/functional_adamax.py b/lib/python3.12/site-packages/torch/distributed/optim/functional_adamax.py new file mode 100644 index 0000000000000000000000000000000000000000..6691bc3b724447e2fb8a6ba0c766c7aa08163af2 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/functional_adamax.py @@ -0,0 +1,123 @@ +# mypy: allow-untyped-defs +from typing import Dict, List, Optional, Tuple + +import torch +import torch.optim._functional as F +from torch import Tensor +from torch.distributed.optim._deprecation_warning import ( + _scripted_functional_optimizer_deprecation_warning, +) + + +__all__: List[str] = [] + + +# Define a TorchScript compatible Functional Adamax Optimizer +# where we use these optimizer in a functional way. +# Instead of using the `param.grad` when updating parameters, +# we explicitly allow the distributed optimizer pass gradients to +# the `step` function. In this way, we could separate the gradients +# and parameters and allow multithreaded trainer to update the +# parameters without data traces on accumulating to the same .grad. +# NOTE: This should be only used by distributed optimizer internals +# and not meant to expose to the user. +@torch.jit.script +class _FunctionalAdamax: + def __init__( + self, + params: List[Tensor], + lr: float = 1e-3, + betas: Tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0.0, + foreach: bool = False, + maximize: bool = False, + _allow_empty_param_list: bool = False, + ): + _scripted_functional_optimizer_deprecation_warning(stacklevel=2) + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + + self.defaults = { + "lr": lr, + "eps": eps, + "beta1": betas[0], + "beta2": betas[1], + "weight_decay": weight_decay, + } + self.foreach = foreach + self.maximize = maximize + self.state = torch.jit.annotate(Dict[torch.Tensor, Dict[str, torch.Tensor]], {}) + + if len(params) == 0 and not _allow_empty_param_list: + raise ValueError("optimizer got an empty parameter list") + + # NOTE: we only have one param_group and don't allow user to add additional + # param group as it's not a common use case. + self.param_group = {"params": params} + + def step(self, gradients: List[Optional[Tensor]]): + params = self.param_group["params"] + params_with_grad = [] + grads = [] + exp_avgs = [] + exp_infs = [] + state_steps: List[Tensor] = [] + + if len(params) != len(gradients): + raise ValueError( + "the gradients passed in does not equal to the size of the parameters!" + + f"Params length: {len(params)}. " + + f"Gradients length: {len(gradients)}" + ) + + has_complex = False + for param, gradient in zip(self.param_group["params"], gradients): + if gradient is not None: + has_complex |= torch.is_complex(param) + params_with_grad.append(param) + grads.append(gradient) + # Lazy state initialization + if param not in self.state: + self.state[param] = {} + state = self.state[param] + state["step"] = torch.tensor(0.0) + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_inf"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + + state = self.state[param] + + exp_avgs.append(state["exp_avg"]) + exp_infs.append(state["exp_inf"]) + state_steps.append(state["step"]) + + with torch.no_grad(): + F.adamax( + params_with_grad, + grads, + exp_avgs, + exp_infs, + state_steps, + eps=self.defaults["eps"], + beta1=self.defaults["beta1"], + beta2=self.defaults["beta2"], + lr=self.defaults["lr"], + weight_decay=self.defaults["weight_decay"], + foreach=self.foreach, + maximize=self.maximize, + has_complex=has_complex, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/functional_adamw.py b/lib/python3.12/site-packages/torch/distributed/optim/functional_adamw.py new file mode 100644 index 0000000000000000000000000000000000000000..cde0591a174fb03b83c0d2fdb926b4f95666276b --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/functional_adamw.py @@ -0,0 +1,203 @@ +# mypy: allow-untyped-defs +from typing import Dict, List, Optional, Tuple + +import torch +import torch.optim._functional as F +from torch import Tensor +from torch.distributed.optim._deprecation_warning import ( + _scripted_functional_optimizer_deprecation_warning, +) + + +__all__: List[str] = [] + + +# Define a TorchScript compatible Functional AdamW Optimizer +# where we use these optimizer in a functional way. +# Instead of using the `param.grad` when updating parameters, +# we explicitly allow the distributed optimizer pass gradients to +# the `step` function. In this way, we could separate the gradients +# and parameters and allow multithreaded trainer to update the +# parameters without data traces on accumulating to the same .grad. +# NOTE: This should be only used by distributed optimizer internals +# and not meant to expose to the user. +@torch.jit.script +class _FunctionalAdamW: + def __init__( + self, + params: List[Tensor], + lr: float = 1e-3, + betas: Tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 1e-2, + amsgrad: bool = False, + maximize: bool = False, + foreach: bool = False, + fused: bool = False, + _allow_empty_param_list: bool = False, + ): + _scripted_functional_optimizer_deprecation_warning(stacklevel=2) + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + + self.defaults = { + "lr": lr, + "eps": eps, + "beta1": betas[0], + "beta2": betas[1], + "weight_decay": weight_decay, + } + self.amsgrad = amsgrad + self.maximize = maximize + self.foreach = foreach + self.fused = fused + self.state = torch.jit.annotate(Dict[torch.Tensor, Dict[str, torch.Tensor]], {}) + + if len(params) == 0 and not _allow_empty_param_list: + raise ValueError("optimizer got an empty parameter list") + + # NOTE: we only have one param_group and don't allow user to add additional + # param group as it's not a common use case. + self.param_group = {"params": params} + + def step_param(self, param: Tensor, grad: Optional[Tensor]): + params_with_grad = [] + grads = [] + exp_avgs = [] + exp_avg_sqs = [] + max_exp_avg_sqs = [] + state_steps: List[Tensor] = [] + has_complex = torch.is_complex(param) + if grad is not None: + params_with_grad.append(param) + grads.append(grad) + # Lazy state initialization + if param not in self.state: + self.state[param] = {} + state = self.state[param] + state["step"] = torch.tensor(0.0) + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + if self.amsgrad: + # Maintains max of all exp. moving avg. of sq. grad. values + state["max_exp_avg_sq"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + + state = self.state[param] + + exp_avgs.append(state["exp_avg"]) + exp_avg_sqs.append(state["exp_avg_sq"]) + + if self.amsgrad: + max_exp_avg_sqs.append(state["max_exp_avg_sq"]) + + state_steps.append(state["step"]) + with torch.no_grad(): + F.adamw( + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + amsgrad=self.amsgrad, + maximize=self.maximize, + beta1=self.defaults["beta1"], + beta2=self.defaults["beta2"], + lr=self.defaults["lr"], + weight_decay=self.defaults["weight_decay"], + eps=self.defaults["eps"], + foreach=self.foreach, + fused=self.fused, + grad_scale=None, + found_inf=None, + has_complex=has_complex, + ) + + def step(self, gradients: List[Optional[Tensor]]): + params = self.param_group["params"] + params_with_grad = [] + grads = [] + exp_avgs = [] + exp_avg_sqs = [] + max_exp_avg_sqs = [] + state_steps: List[Tensor] = [] + + if len(params) != len(gradients): + raise ValueError( + "the gradients passed in does not equal to the size of the parameters!" + + f"Params length: {len(params)}. " + + f"Gradients length: {len(gradients)}" + ) + + has_complex = False + for param, gradient in zip(self.param_group["params"], gradients): + if gradient is not None: + has_complex |= torch.is_complex(param) + params_with_grad.append(param) + grads.append(gradient) + # Lazy state initialization + if param not in self.state: + self.state[param] = {} + state = self.state[param] + state["step"] = torch.tensor(0.0) + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + if self.amsgrad: + # Maintains max of all exp. moving avg. of sq. grad. values + state["max_exp_avg_sq"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + + state = self.state[param] + + exp_avgs.append(state["exp_avg"]) + exp_avg_sqs.append(state["exp_avg_sq"]) + + if self.amsgrad: + max_exp_avg_sqs.append(state["max_exp_avg_sq"]) + + state_steps.append(state["step"]) + + with torch.no_grad(): + F.adamw( + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + amsgrad=self.amsgrad, + maximize=self.maximize, + beta1=self.defaults["beta1"], + beta2=self.defaults["beta2"], + lr=self.defaults["lr"], + weight_decay=self.defaults["weight_decay"], + eps=self.defaults["eps"], + foreach=self.foreach, + fused=self.fused, + grad_scale=None, + found_inf=None, + has_complex=has_complex, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/functional_rmsprop.py b/lib/python3.12/site-packages/torch/distributed/optim/functional_rmsprop.py new file mode 100644 index 0000000000000000000000000000000000000000..1362c1e635fcb14b08bacc65c904494a192bf5a0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/functional_rmsprop.py @@ -0,0 +1,130 @@ +# mypy: allow-untyped-defs +from typing import Dict, List, Optional + +import torch +import torch.optim._functional as F +from torch import Tensor +from torch.distributed.optim._deprecation_warning import ( + _scripted_functional_optimizer_deprecation_warning, +) + + +__all__: List[str] = [] + + +# Define a TorchScript compatible Functional RMSprop Optimizer +# where we use these optimizer in a functional way. +# Instead of using the `param.grad` when updating parameters, +# we explicitly allow the distributed optimizer pass gradients to +# the `step` function. In this way, we could separate the gradients +# and parameters and allow multithreaded trainer to update the +# parameters without data traces on accumulating to the same .grad. +# NOTE: This should be only used by distributed optimizer internals +# and not meant to expose to the user. +@torch.jit.script +class _FunctionalRMSprop: + def __init__( + self, + params: List[Tensor], + lr: float = 1e-2, + alpha: float = 0.99, + eps: float = 1e-8, + weight_decay: float = 0.0, + momentum: float = 0.0, + centered: bool = False, + foreach: bool = False, + maximize: bool = False, + _allow_empty_param_list: bool = False, + ): + _scripted_functional_optimizer_deprecation_warning(stacklevel=2) + self.defaults = { + "lr": lr, + "alpha": alpha, + "eps": eps, + "weight_decay": weight_decay, + "momentum": momentum, + } + self.centered = centered + self.foreach = foreach + self.maximize = maximize + + if len(params) == 0 and not _allow_empty_param_list: + raise ValueError("optimizer got an empty parameter list") + + # NOTE: we only have one param_group and don't allow user to add additional + # param group as it's not a common use case. + self.param_group = {"params": params} + + self.state = torch.jit.annotate(Dict[torch.Tensor, Dict[str, torch.Tensor]], {}) + + def step(self, gradients: List[Optional[Tensor]]): + params = self.param_group["params"] + params_with_grad = [] + grads = [] + square_avgs = [] + grad_avgs = [] + momentum_buffer_list = [] + state_steps = [] + lr = self.defaults["lr"] + alpha = self.defaults["alpha"] + eps = self.defaults["eps"] + momentum = self.defaults["momentum"] + weight_decay = self.defaults["weight_decay"] + + if len(params) != len(gradients): + raise ValueError( + "the gradients passed in does not equal to the size of the parameters!" + + f"Params length: {len(params)}. " + + f"Gradients length: {len(gradients)}" + ) + + has_complex = False + for param, gradient in zip(params, gradients): + if gradient is not None: + has_complex |= torch.is_complex(param) + params_with_grad.append(param) + grads.append(gradient) + # Lazy state initialization + if param not in self.state: + self.state[param] = {} + state = self.state[param] + state["step"] = torch.tensor(0.0) + state["square_avg"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + if momentum > 0: + state["momentum_buffer"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + if self.centered: + state["grad_avg"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + + state = self.state[param] + square_avgs.append(state["square_avg"]) + if momentum > 0: + momentum_buffer_list.append(state["momentum_buffer"]) + if self.centered: + grad_avgs.append(state["grad_avg"]) + + state_steps.append(state["step"]) + + with torch.no_grad(): + F.rmsprop( + params_with_grad, + grads, + square_avgs, + grad_avgs, + momentum_buffer_list, + state_steps, + lr=lr, + alpha=alpha, + eps=eps, + weight_decay=weight_decay, + momentum=momentum, + centered=self.centered, + foreach=self.foreach, + maximize=self.maximize, + has_complex=has_complex, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/functional_rprop.py b/lib/python3.12/site-packages/torch/distributed/optim/functional_rprop.py new file mode 100644 index 0000000000000000000000000000000000000000..e00ce9006737d6e4d52430dc705fd010b51852da --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/functional_rprop.py @@ -0,0 +1,107 @@ +# mypy: allow-untyped-defs +from typing import Dict, List, Optional, Tuple + +import torch +import torch.optim._functional as F +from torch import Tensor +from torch.distributed.optim._deprecation_warning import ( + _scripted_functional_optimizer_deprecation_warning, +) + + +__all__: List[str] = [] + + +# Define a TorchScript compatible Functional Rprop Optimizer +# where we use these optimizer in a functional way. +# Instead of using the `param.grad` when updating parameters, +# we explicitly allow the distributed optimizer pass gradients to +# the `step` function. In this way, we could separate the gradients +# and parameters and allow multithreaded trainer to update the +# parameters without data traces on accumulating to the same .grad. +# NOTE: This should be only used by distributed optimizer internals +# and not meant to expose to the user. +@torch.jit.script +class _FunctionalRprop: + def __init__( + self, + params: List[Tensor], + lr: float = 1e-2, + etas: Tuple[float, float] = (0.5, 1.2), + step_sizes: Tuple[float, float] = (1e-6, 50), + foreach: bool = False, + maximize: bool = False, + _allow_empty_param_list: bool = False, + ): + _scripted_functional_optimizer_deprecation_warning(stacklevel=2) + self.defaults = { + "lr": lr, + } + self.etas = etas + self.step_sizes = step_sizes + self.foreach = foreach + self.maximize = maximize + + if len(params) == 0 and not _allow_empty_param_list: + raise ValueError("optimizer got an empty parameter list") + + # NOTE: we only have one param_group and don't allow user to add additional + # param group as it's not a common use case. + self.param_group = {"params": params} + + self.state = torch.jit.annotate(Dict[torch.Tensor, Dict[str, torch.Tensor]], {}) + + def step(self, gradients: List[Optional[Tensor]]): + params = self.param_group["params"] + params_with_grad = [] + grads = [] + prevs = [] + step_sizes = [] + state_steps = [] + lr = self.defaults["lr"] + etaminus, etaplus = self.etas + step_size_min, step_size_max = self.step_sizes + + if len(params) != len(gradients): + raise ValueError( + "the gradients passed in does not equal to the size of the parameters!" + + f"Params length: {len(params)}. " + + f"Gradients length: {len(gradients)}" + ) + + has_complex = False + for param, gradient in zip(params, gradients): + if gradient is not None: + has_complex |= torch.is_complex(param) + params_with_grad.append(param) + grads.append(gradient) + # Lazy state initialization + if param not in self.state: + self.state[param] = {} + state = self.state[param] + state["step"] = torch.tensor(0.0) + state["prev"] = torch.zeros_like( + param, memory_format=torch.preserve_format + ) + state["step_size"] = torch.full_like(gradient, lr) + + state = self.state[param] + prevs.append(state["prev"]) + step_sizes.append(state["step_size"]) + state_steps.append(state["step"]) + + with torch.no_grad(): + F.rprop( + params_with_grad, + grads, + prevs, + step_sizes, + state_steps, + step_size_min=step_size_min, + step_size_max=step_size_max, + etaminus=etaminus, + etaplus=etaplus, + foreach=self.foreach, + maximize=self.maximize, + has_complex=has_complex, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/functional_sgd.py b/lib/python3.12/site-packages/torch/distributed/optim/functional_sgd.py new file mode 100644 index 0000000000000000000000000000000000000000..4644b6a4030edc1b650032a6836f6484affefd66 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/functional_sgd.py @@ -0,0 +1,166 @@ +# mypy: allow-untyped-defs +from typing import Dict, List, Optional + +import torch +import torch.optim._functional as F +from torch import Tensor +from torch.distributed.optim._deprecation_warning import ( + _scripted_functional_optimizer_deprecation_warning, +) + + +__all__: List[str] = [] + + +# Define a TorchScript compatible Functional SGD Optimizer +# where we use these optimizer in a functional way. +# Instead of using the `param.grad` when updating parameters, +# we explicitly allow the distributed optimizer pass gradients to +# the `step` function. In this way, we could separate the gradients +# and parameters and allow multithreaded trainer to update the +# parameters without data traces on accumulating to the same .grad. +# NOTE: This should be only used by distributed optimizer internals +# and not meant to expose to the user. +@torch.jit.script +class _FunctionalSGD: + def __init__( + self, + params: List[Tensor], + lr: float = 1e-2, + momentum: float = 0.0, + dampening: float = 0.0, + weight_decay: float = 0.0, + nesterov: bool = False, + maximize: bool = False, + foreach: bool = False, + fused: bool = False, + _allow_empty_param_list: bool = False, + ): + _scripted_functional_optimizer_deprecation_warning(stacklevel=2) + self.defaults = { + "lr": lr, + "momentum": momentum, + "dampening": dampening, + "weight_decay": weight_decay, + } + self.nesterov = nesterov + self.maximize = maximize + self.foreach = foreach + self.fused = fused + self.state = torch.jit.annotate(Dict[torch.Tensor, Dict[str, torch.Tensor]], {}) + + if len(params) == 0 and not _allow_empty_param_list: + raise ValueError("optimizer got an empty parameter list") + + # NOTE: we only have one param_group and don't allow user to add additional + # param group as it's not a common use case. + self.param_group = {"params": params} + + def step_param(self, param: Tensor, grad: Optional[Tensor]): + """Similar to self.step, but operates on a single parameter and + its gradient. + """ + # TODO: Once step_param interface is robust, refactor step to call + # step param on each param. + weight_decay = self.defaults["weight_decay"] + momentum = self.defaults["momentum"] + dampening = self.defaults["dampening"] + lr = self.defaults["lr"] + params = [param] + momentum_buffer_list: List[Optional[Tensor]] = [] + grads = [] + + has_sparse_grad = False + if grad is not None: + grads.append(grad) + if grad.is_sparse: + has_sparse_grad = True + if param not in self.state: + self.state[param] = {} + state = self.state[param] + if "momentum_buffer" not in state: + momentum_buffer_list.append(None) + else: + momentum_buffer_list.append(state["momentum_buffer"]) + + with torch.no_grad(): + F.sgd( + params, + grads, + momentum_buffer_list, + weight_decay=weight_decay, + momentum=momentum, + lr=lr, + dampening=dampening, + nesterov=self.nesterov, + maximize=self.maximize, + has_sparse_grad=has_sparse_grad, + foreach=self.foreach, + fused=self.fused, + grad_scale=None, + found_inf=None, + ) + # update momentum_buffer in state + state = self.state[param] + momentum_buffer = momentum_buffer_list[0] + if momentum_buffer is not None: + state["momentum_buffer"] = momentum_buffer + + def step(self, gradients: List[Optional[Tensor]]): + params = self.param_group["params"] + params_with_grad = [] + grads = [] + momentum_buffer_list: List[Optional[Tensor]] = [] + lr = self.defaults["lr"] + weight_decay = self.defaults["weight_decay"] + momentum = self.defaults["momentum"] + dampening = self.defaults["dampening"] + + if len(params) != len(gradients): + raise ValueError( + "the gradients passed in does not equal to the size of the parameters!" + + f"Params length: {len(params)}. " + + f"Gradients length: {len(gradients)}" + ) + + has_sparse_grad = False + for param, gradient in zip(params, gradients): + if gradient is not None: + params_with_grad.append(param) + grads.append(gradient) + if gradient.is_sparse: + has_sparse_grad = True + + if param not in self.state: + self.state[param] = {} + + state = self.state[param] + if "momentum_buffer" not in state: + momentum_buffer_list.append(None) + else: + momentum_buffer_list.append(state["momentum_buffer"]) + + with torch.no_grad(): + F.sgd( + params_with_grad, + grads, + momentum_buffer_list, + weight_decay=weight_decay, + momentum=momentum, + lr=lr, + dampening=dampening, + nesterov=self.nesterov, + maximize=self.maximize, + has_sparse_grad=has_sparse_grad, + foreach=self.foreach, + fused=self.fused, + grad_scale=None, + found_inf=None, + ) + + # update momentum_buffers in state + for i, p in enumerate(params_with_grad): + state = self.state[p] + momentum_buffer = momentum_buffer_list[i] + if momentum_buffer is not None: + state["momentum_buffer"] = momentum_buffer diff --git a/lib/python3.12/site-packages/torch/distributed/optim/named_optimizer.py b/lib/python3.12/site-packages/torch/distributed/optim/named_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..4130189967a0a7912cd9aee675da352697d98cbb --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/named_optimizer.py @@ -0,0 +1,339 @@ +# mypy: allow-untyped-defs +import logging +import warnings +from copy import deepcopy +from typing import ( + Any, + Callable, + Collection, + Dict, + List, + Mapping, + Optional, + overload, + Union, +) + +import torch +import torch.nn as nn +from torch import optim +from torch.distributed._shard.sharded_tensor import ShardedTensor +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + + +__all__: List[str] = [] + +logger = logging.getLogger(__name__) + + +class _NamedOptimizer(optim.Optimizer): + """ + ``_NamedOptimizer`` takes a dict of parameters and exposes ``state_dict`` by parameter key. + + We replace the original key (number) in an optim to the + fully qualified name (FQN) string. User can initialize the optim as they + initialize a PyTorch optim, the only difference is that they also need to + pass in the FQN of each parameters. + + Args: + named_parameters (Mapping[str, Union[torch.Tensor, ShardedTensor]]): + Mapping from FQN to parameter. + optimizer_class (optim.Optimizer): + The class of optimizer to instantiate. + param_groups (Collection[Mapping[str, Any]]): + `param_groups` to pass to optimizer if specified. + The key of the inner map needs to be FQNs. + Default: None + module (nn.Module): the module whose parameters to updated + by the optimizer. + args: arguments to pass to the optimizer constructor. + kwargs: arguments to pass to the optimizer constructor. + + Example:: + >>> # xdoctest: +SKIP("distributed") + >>> from torch import optim + >>> from torch.distributed.optim import _NamedOptimizer + >>> + >>> # Define the named optimizer. + >>> m = Model(...) + >>> named_optim = _NamedOptimizer(m.named_parameters(), optim.SGD) + >>> # Forward pass + backward pass. + >>> named_optim.step() + >>> ... + >>> # Call state_dict for the named optimizer returns a FQN state_dict. + >>> named_optim.state_dict() + + Warning: This API is still in development and subject to change. + + TODO: Add tutorial for _NamedOptimizer. + TODO: Add documentation in the docstring for the public attributes + like self.param_groups and self.named_parameters. + """ + + def __init__( + self, + named_parameters: Mapping[str, Union[torch.Tensor, ShardedTensor]], + optimizer_class: optim.Optimizer, + param_groups: Optional[Collection[Mapping[str, Any]]] = None, + module: Optional[nn.Module] = None, + *args, + **kwargs, + ) -> None: + torch._C._log_api_usage_once("torch.distributed.optim._NamedOptimizer") + self.param_groups: Collection[Mapping[str, Any]] = param_groups # type: ignore[assignment] + self._param_groups_check() + self.named_parameters = dict(named_parameters) + params_for_optimizer = ( + self.named_parameters.values() if param_groups is None else param_groups + ) + self._optimizer = optimizer_class( # type: ignore[operator] + params_for_optimizer, + *args, + **kwargs, + ) + self.module = module + if param_groups is None: + self.ordered_param_keys = list(self.named_parameters.keys()) + else: + warnings.warn( + "Since we pass in param_groups, we will use param_groups to " + "initialize the optimizer, not all parameters of the module." + ) + param_to_key = {param: key for key, param in self.named_parameters.items()} # type: ignore[misc, has-type] + ordered_param_keys = [] + for group in param_groups: + for param in group["params"]: + if param not in param_to_key: + raise ValueError( + f"Expect param name {param} found in param group but is missing." + ) + ordered_param_keys.append(param_to_key[param]) + self.ordered_param_keys = ordered_param_keys + # Update param_groups from optimizer. + self.param_groups = self._optimizer.param_groups + + def _param_groups_check(self): + if self.param_groups is not None: + for param_group in self.param_groups: + assert isinstance(param_group, dict), "param group must be a dict" + assert "params" in param_group, "param group must contain key params" + params = param_group["params"] + if isinstance(params, torch.Tensor): + params = [params] + params = list(params) + for param in params: + if not isinstance(param, torch.Tensor): + raise TypeError( + "optimizer can only optimize Tensors, " + "but one of the params is " + torch.typename(param) + ) + param_group["params"] = params + + def state_dict(self) -> Dict[str, Any]: + """ + Return the ``state_dict`` of the optimizer. + + Instead of using number to index + parameters, we will use module fully qualified name (FQN) as the key. + """ + state_dict = self._optimizer.state_dict() + param_groups = state_dict["param_groups"] + + ret_state = { + self.ordered_param_keys[st_key]: state_val + for st_key, state_val in state_dict["state"].items() + } + + ret_groups = [] + for group in param_groups: + param_keys = [self.ordered_param_keys[param] for param in group["params"]] + ret_group = {"params": sorted(param_keys)} + for k, v in group.items(): + if k != "params": + ret_group[k] = deepcopy(v) + ret_groups.append(ret_group) + + return self._post_state_dict({"state": ret_state, "param_groups": ret_groups}) + + @overload + def step(self, closure: None = ...) -> None: + ... + + @overload + def step(self, closure: Callable[[], float]) -> float: + ... + + def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]: + """ + Perform a single optimization step. + + This will call :meth:`torch.optim.Optimizer.step` on the wrapped + optimizer. + """ + return self._optimizer.step(closure=closure) + + @property + def state(self) -> Mapping[torch.Tensor, Any]: # type: ignore[override] + return self._optimizer.state + + def load_state_dict(self, state_dict: Mapping[str, Any]) -> None: + """ + Define the default behavior to load a state_dict for ``_NamedOptimizer``. + + Sample Code + ``` + my_model = MyModule() + optimizer = _NamedOptimizer(my_model.named_parameters(), Adagrad) + ... + + optim_state_dict = optimizer.state_dict() + ... + ... + + optimizer.load_state_dict(optim_state_dict) + ... + ``` + Args: + state_dict (Dict[str, Any]) : A ``state_dict`` to load into the optimizer. + Note that this state dict update is performed in place. + + .. note:: PyTorch is using lazy init to initialize the optim states. + So it is possible that there is no optim state when user call + ``load_state_dict`` and for ``_NamedOptimizer`` we make it stricter + that users can only call ``load_state_dict`` after the state is initialized. + By doing this, we can validate the optim ``state_dict`` to be loaded. + """ + new_state_dict = self._optimizer.state_dict() + state_dict = self._pre_load_state_dict(state_dict) + state = state_dict["state"] + new_state = new_state_dict["state"] + if len(new_state) == 0: + raise ValueError( + "Expects the optim to be initialized before load but found not initialized." + ) + + for idx, param_key in enumerate(self.ordered_param_keys): + # When the conditional training is performed, not all parameters are updated in the optim. + if param_key not in state.keys(): + continue + if len(state[param_key]) != len(new_state[idx]): + raise ValueError( + f"Expects equal length as {len(new_state[idx])} for parameter {param_key} but found: {len(state[param_key])}" + ) + # Iterate through all optimizer states. + for state_key, state_val in new_state[idx].items(): + if state_key not in state[param_key]: + raise ValueError( + f"Expects state {state_key} for parameter {param_key} but not found." + ) + + src_state_val = state[param_key][state_key] + if isinstance(state_val, ShardedTensor): + assert isinstance(src_state_val, ShardedTensor) + num_shards = len(state_val.local_shards()) + num_new_shards = len(src_state_val.local_shards()) + if num_shards != num_new_shards: + raise ValueError( + f"Expects equal number of shards as {num_new_shards} but found {num_shards} for {param_key}/{state_key}" + ) + for shard, src_shard in zip( + state_val.local_shards(), src_state_val.local_shards() + ): + shard.tensor.detach().copy_(src_shard.tensor) + elif isinstance(state_val, torch.Tensor): + assert isinstance(src_state_val, torch.Tensor) + state_val.detach().copy_(src_state_val) + else: + new_state[idx][state_key] = deepcopy(src_state_val) + + # Load param_groups of state_dict + src_param_groups = state_dict["param_groups"] + new_param_groups = new_state_dict["param_groups"] + + src_group_map = {} + for group in src_param_groups: + param_keys = list(group["params"]) + src_group_map[_gen_param_group_key(param_keys)] = group + new_group_map = {} + for new_group in new_param_groups: + param_keys = [] + for param_key in new_group["params"]: + param_keys.append(self.ordered_param_keys[param_key]) # type: ignore[call-overload] + new_group_map[_gen_param_group_key(param_keys)] = new_group + for group_key, new_group in new_group_map.items(): + # When not all parameters are used in training or receive gradient, aka., not all parameters + # would be in the param_group. Thus we skip the group_key here. + if group_key not in src_group_map: + continue + src_group = src_group_map[group_key] + if len(src_group) != len(new_group): + raise ValueError( + f"Expects equal param_group size as {len(new_group)} for group {group_key} but found {len(src_group)}." + ) + for k in src_group: + if k not in new_group: + raise ValueError( + f"Expects group key {k} to be in group {group_key} in `state_dict` but is missing." + ) + if k != "params": + new_group[k] = deepcopy(src_group[k]) + + self._optimizer.load_state_dict(new_state_dict) + + def add_param_group(self, param_group: Mapping[str, Any]) -> None: + """ + Add a param group to the :class:`_NamedOptimizer` s `param_groups`. + + Warning: This API is still in development and subject to change. + """ + assert isinstance(param_group, dict), "param group must be a dict" + + params = param_group["params"] + if isinstance(params, torch.Tensor): + param_group["params"] = [params] + else: + param_group["params"] = list(params) + + param_to_key = {param: key for key, param in self.named_parameters.items()} # type: ignore[misc, has-type] + for param in param_group["params"]: + if param not in param_to_key: + raise ValueError("some parameters are not in the module") + self.ordered_param_keys.append(param_to_key[param]) + + self._optimizer.add_param_group(param_group) + # Update param_groups from optimizer. + self.param_groups = self._optimizer.param_groups + + def init_state(self) -> None: + """ + Run a dummy optimizer step, which allows to initialize optimizer state because we do lazy init for most optimizers. + + This allows doing in-place loading of optimizer state from a checkpoint. + """ + for param in self.named_parameters.values(): + if param.requires_grad: + t = torch.zeros_like(param) + param.grad = torch.autograd.Variable(t) + # Calling ``step`` will load the initial state for optimizer states. + self.step(closure=None) + + def _pre_load_state_dict(self, state_dict) -> Dict[str, Any]: + # TODO(chienchin): This API should be FSDP agnostic and should support + # general user hooks. + if isinstance(self.module, FSDP): + return FSDP.optim_state_dict_to_load( + self.module, self._optimizer, state_dict, is_named_optimizer=True + ) + return state_dict + + def _post_state_dict(self, state_dict) -> Dict[str, Any]: + # TODO(chienchin): This API should be FSDP agnostic and should support + # general user hooks. + if isinstance(self.module, FSDP): + FSDP.optim_state_dict(self.module, self._optimizer, state_dict) + return state_dict + + +def _gen_param_group_key(param_keys: List[str]) -> str: + """Concatenate all param keys as a unique indentifier for one param group.""" + return "/".join(sorted(param_keys)) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/optimizer.py b/lib/python3.12/site-packages/torch/distributed/optim/optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..a32e5505ffeb88c9d4f966c77503d7240322f84f --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/optimizer.py @@ -0,0 +1,256 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +import logging +from collections import defaultdict +from threading import Lock +from typing import List, Optional + +import torch +import torch.distributed.autograd as dist_autograd +import torch.distributed.rpc as rpc +import torch.jit as jit +import torch.nn as nn +from torch import Tensor +from torch.distributed.rpc import RRef + +from .utils import functional_optim_map + + +__all__ = ["DistributedOptimizer"] + +logger = logging.getLogger(__name__) + + +# XXX: we define a _ScriptModuleOptimizer here to explicitly +# compile the FunctionalOptimizer class into TorchScript +# This is because ScriptClass instance still lives in +# python unless you explicitly compile it as an attribute +# in ScriptModule or pass it to a ScriptFunction +# _ScriptLocalOptimizerInterface serves as a common +# interface type for Optimizer ScriptModules. +# +# TODO (wanchaol): remove this once we added TorchScript +# class reference semantics +@jit.interface +class _ScriptLocalOptimizerInterface: + def step(self, autograd_ctx_id: int) -> None: + pass + + +class _ScriptLocalOptimizer(nn.Module): + # TorchScript does not support multithread concurrent compiling. + # request_callback might invoke concurrent compiling, so we + # serialize the compiling with a lock + compile_lock = Lock() + + def __init__(self, optim_cls, local_params_rref, *args, **kwargs): + super().__init__() + self._local_params = [rref.local_value() for rref in local_params_rref] + self.optim = optim_cls(self._local_params, *args, **kwargs) + + @jit.export + def step(self, autograd_ctx_id: int): + all_local_grads = dist_autograd.get_gradients(autograd_ctx_id) + # apply functional optimizer step with a list of gradients + grads: List[Optional[Tensor]] = [ + all_local_grads[p] if p in all_local_grads else None + for p in self._local_params + ] + + self.optim.step(grads) + + +# TODO (wanchaol): remove/merge this with ScriptLocalOptimizer once +# we have converted all to functional optimizer in distributed.optim +class _LocalOptimizer: + # Ideally we would only need to share a lock for instances of + # _LocalOptimizer that deal with the same parameters. We are + # making a simplifying assumption here that if there is more + # than one instance of _LocalOptimizer per worker, they will + # be optimizing the same parameters (e.g. each data parallel + # trainer will create its own instance of _LocalOptimizer but + # they will all optimize the same parameters on each worker) + global_lock = Lock() + + def __init__(self, optim_cls, local_params_rref, *args, **kwargs): + self._local_params = [rref.local_value() for rref in local_params_rref] + self.optim = optim_cls(self._local_params, *args, **kwargs) + + def step(self, autograd_ctx_id): + all_local_grads = dist_autograd.get_gradients(autograd_ctx_id) + + with _LocalOptimizer.global_lock: + for param, grad in all_local_grads.items(): + param.grad = grad + self.optim.step() + + +def _new_local_optimizer(optim_cls, local_params_rref, *args, **kwargs): + return rpc.RRef(_LocalOptimizer(optim_cls, local_params_rref, *args, **kwargs)) + + +def _local_optimizer_step(local_optim_rref, autograd_ctx_id): + local_optim = local_optim_rref.local_value() + local_optim.step(autograd_ctx_id) + + +# new/step functions combined with _ScriptLocalOptimizer to provide GIL-free optimizer +def _new_script_local_optimizer(optim_cls, local_params_rref, *args, **kwargs): + optim = _ScriptLocalOptimizer(optim_cls, local_params_rref, *args, **kwargs) + + with _ScriptLocalOptimizer.compile_lock: + script_optim = jit.script(optim) + return rpc.RRef(script_optim, _ScriptLocalOptimizerInterface) + + +@jit.script +def _script_local_optimizer_step( + local_optim_rref: RRef[_ScriptLocalOptimizerInterface], autograd_ctx_id: int +) -> None: + local_optim = local_optim_rref.local_value() + local_optim.step(autograd_ctx_id) + + +def _wait_for_all(rpc_futs): + # TODO: improve error propagation + exception = None + results = [] + for fut in rpc_futs: + try: + results.append(fut.wait()) + except Exception as e: + results.append(e) + exception = e + if exception is not None: + raise exception + return results + + +class DistributedOptimizer: + """ + DistributedOptimizer takes remote references to parameters scattered + across workers and applies the given optimizer locally for each parameter. + + This class uses :meth:`~torch.distributed.autograd.get_gradients` in order + to retrieve the gradients for specific parameters. + + Concurrent calls to + :meth:`~torch.distributed.optim.DistributedOptimizer.step`, + either from the same or different clients, will + be serialized on each worker -- as each worker's optimizer can only work + on one set of gradients at a time. However, there is no guarantee that + the full forward-backward-optimizer sequence will execute for one client + at a time. This means that the gradients being applied may not correspond + to the latest forward pass executed on a given worker. Also, there is no + guaranteed ordering across workers. + + `DistributedOptimizer` creates the local optimizer with TorchScript enabled + by default, so that optimizer updates are not blocked by the Python Global + Interpreter Lock (GIL) in the case of multithreaded training (e.g. Distributed + Model Parallel). This feature is currently enabled for most optimizers. You + can also follow `the recipe`__ in PyTorch tutorials to enable TorchScript support + for your own custom optimizers. + + Args: + optimizer_class (optim.Optimizer): the class of optimizer to + instantiate on each worker. + params_rref (list[RRef]): list of RRefs to local or remote parameters + to optimize. + args: arguments to pass to the optimizer constructor on each worker. + kwargs: arguments to pass to the optimizer constructor on each worker. + + Example:: + >>> # xdoctest: +SKIP("distributed") + >>> import torch.distributed.autograd as dist_autograd + >>> import torch.distributed.rpc as rpc + >>> from torch import optim + >>> from torch.distributed.optim import DistributedOptimizer + >>> + >>> with dist_autograd.context() as context_id: + >>> # Forward pass. + >>> rref1 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3)) + >>> rref2 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 1)) + >>> loss = rref1.to_here() + rref2.to_here() + >>> + >>> # Backward pass. + >>> dist_autograd.backward(context_id, [loss.sum()]) + >>> + >>> # Optimizer. + >>> dist_optim = DistributedOptimizer( + >>> optim.SGD, + >>> [rref1, rref2], + >>> lr=0.05, + >>> ) + >>> dist_optim.step(context_id) + + __ https://github.com/pytorch/tutorials/pull/1465 + """ + + def __init__(self, optimizer_class, params_rref, *args, **kwargs): + torch._C._log_api_usage_once("torch.distributed.optim.DistributedOptimizer") + per_worker_params_rref = defaultdict(list) + for param in params_rref: + per_worker_params_rref[param.owner()].append(param) + + if optimizer_class in functional_optim_map and jit._state._enabled: + optim_ctor = functional_optim_map.get(optimizer_class) + else: + optim_ctor = optimizer_class + self.is_functional_optim = optim_ctor != optimizer_class + + if self.is_functional_optim: + optimizer_new_func = _new_script_local_optimizer + else: + logger.warning( + "Creating the optimizer %s without TorchScript support, " + "this might result in slow computation time in multithreading environment" + "(i.e. Distributed Model Parallel training on CPU) due to the Python's " + "Global Interpreter Lock (GIL). Please file an issue if you need this " + "optimizer in TorchScript. ", + optimizer_class, + ) + optimizer_new_func = _new_local_optimizer + + remote_optim_futs = [] + for worker, param_rrefs in per_worker_params_rref.items(): + remote_optim_rref_fut = rpc.rpc_async( + worker, + optimizer_new_func, + args=(optim_ctor, param_rrefs) + args, + kwargs=kwargs, + ) + remote_optim_futs.append(remote_optim_rref_fut) + + self.remote_optimizers = _wait_for_all(remote_optim_futs) + + def step(self, context_id): + """ + Performs a single optimization step. + + This will call :meth:`torch.optim.Optimizer.step` on each worker + containing parameters to be optimized, and will block until all workers + return. The provided ``context_id`` will be used to retrieve the + corresponding :class:`~torch.distributed.autograd.context` that + contains the gradients that should be applied to the parameters. + + Args: + context_id: the autograd context id for which we should run the + optimizer step. + """ + dist_autograd._is_valid_context(context_id) + + optimizer_step_func = ( + _script_local_optimizer_step + if self.is_functional_optim + else _local_optimizer_step + ) + + rpc_futs = [ + rpc.rpc_async( + optimizer.owner(), + optimizer_step_func, + args=(optimizer, context_id), + ) + for optimizer in self.remote_optimizers + ] + _wait_for_all(rpc_futs) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/post_localSGD_optimizer.py b/lib/python3.12/site-packages/torch/distributed/optim/post_localSGD_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..3c0027d112407396c86a6f9966ac3f4e7a5660a3 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/post_localSGD_optimizer.py @@ -0,0 +1,110 @@ +# mypy: allow-untyped-defs +import warnings + +import torch +import torch.distributed.algorithms.model_averaging.averagers as averagers + + +class PostLocalSGDOptimizer(torch.optim.Optimizer): + r""" + Wraps an arbitrary :class:`torch.optim.Optimizer` and runs `post-local SGD `_, + This optimizer runs local optimizer at every step. + After the warm-up stage, it averages parameters periodically afer the local optimizer is applied. + + Args: + optim: The local optimizer. + averager: A model averager instance to run post-localSGD algorithm. + + Example:: + + >>> # xdoctest: +SKIP("undefined variables") + >>> import torch + >>> import torch.distributed as dist + >>> import torch.distributed.algorithms.model_averaging.averagers as averagers + >>> import torch.nn as nn + >>> from torch.distributed.optim import PostLocalSGDOptimizer + >>> from torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook import ( + >>> PostLocalSGDState, + >>> post_localSGD_hook, + >>> ) + >>> + >>> model = nn.parallel.DistributedDataParallel( + >>> module, device_ids=[rank], output_device=rank + >>> ) + >>> + >>> # Register a post-localSGD communication hook. + >>> state = PostLocalSGDState(process_group=None, subgroup=None, start_localSGD_iter=100) + >>> model.register_comm_hook(state, post_localSGD_hook) + >>> + >>> # Create a post-localSGD optimizer that wraps a local optimizer. + >>> # Note that ``warmup_steps`` used in ``PostLocalSGDOptimizer`` must be the same as + >>> # ``start_localSGD_iter`` used in ``PostLocalSGDState``. + >>> local_optim = torch.optim.SGD(params=model.parameters(), lr=0.01) + >>> opt = PostLocalSGDOptimizer( + >>> optim=local_optim, + >>> averager=averagers.PeriodicModelAverager(period=4, warmup_steps=100) + >>> ) + >>> + >>> # In the first 100 steps, DDP runs global gradient averaging at every step. + >>> # After 100 steps, DDP runs gradient averaging within each subgroup (intra-node by default), + >>> # and post-localSGD optimizer runs global model averaging every 4 steps after applying the local optimizer. + >>> for step in range(0, 200): + >>> opt.zero_grad() + >>> loss = loss_fn(output, labels) + >>> loss.backward() + >>> opt.step() + """ + + def __init__(self, optim: torch.optim.Optimizer, averager: averagers.ModelAverager): + self.optim = optim + self.param_groups = self.optim.param_groups + self.averager = averager + + @property + def state(self): + return self.optim.state + + def __repr__(self): + return self.optim.__repr__() + + def state_dict(self): + r""" + This is the same as :class:`torch.optim.Optimizer` :meth:`state_dict`, + but adds an extra entry to record model averager's step to the checkpoint + to ensure reload does not cause unnecessary warm up again. + """ + optim_state_dict = self.optim.state_dict() + optim_state_dict["step"] = self.averager.step + return optim_state_dict + + def load_state_dict(self, state_dict): + r""" + This is the same as :class:`torch.optim.Optimizer` :meth:`load_state_dict`, + but also restores model averager's step value to the one + saved in the provided ``state_dict``. + + If there is no ``"step"`` entry in ``state_dict``, + it will raise a warning and initialize the model averager's step to 0. + """ + self.optim.load_state_dict(state_dict) + if "step" in state_dict: + self.averager.step = state_dict["step"] + else: + warnings.warn( + "Loaded state dict does not contain a step counter for an averager. " + "Setting step counter to 0." + ) + self.averager.step = 0 + + def step(self): # type: ignore[override] + r""" + Performs a single optimization step (parameter update). + """ + self.optim.step() + self.averager.average_parameters(params=self.param_groups) + + def zero_grad(self, set_to_none: bool = True): # type: ignore[override] + self.optim.zero_grad(set_to_none=set_to_none) + + def add_param_group(self, param_group): + self.optim.add_param_group(param_group) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/utils.py b/lib/python3.12/site-packages/torch/distributed/optim/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d2c75eee7e39bc1e079c871e7a0e40ba47af0026 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/utils.py @@ -0,0 +1,66 @@ +# mypy: allow-untyped-defs +from typing import Type + +from torch import optim + +from .functional_adadelta import _FunctionalAdadelta +from .functional_adagrad import _FunctionalAdagrad +from .functional_adam import _FunctionalAdam +from .functional_adamax import _FunctionalAdamax +from .functional_adamw import _FunctionalAdamW +from .functional_rmsprop import _FunctionalRMSprop +from .functional_rprop import _FunctionalRprop +from .functional_sgd import _FunctionalSGD + + +# dict to map a user passed in optimizer_class to a functional +# optimizer class if we have already defined inside the +# distributed.optim package, this is so that we hide the +# functional optimizer to user and still provide the same API. +functional_optim_map = { + optim.Adagrad: _FunctionalAdagrad, + optim.Adam: _FunctionalAdam, + optim.AdamW: _FunctionalAdamW, + optim.SGD: _FunctionalSGD, + optim.Adadelta: _FunctionalAdadelta, + optim.RMSprop: _FunctionalRMSprop, + optim.Rprop: _FunctionalRprop, + optim.Adamax: _FunctionalAdamax, +} + + +def register_functional_optim(key, optim): + """ + Interface to insert a new functional optimizer to functional_optim_map + ``fn_optim_key`` and ``fn_optimizer`` are user defined. The optimizer and key + need not be of :class:`torch.optim.Optimizer` (e.g. for custom optimizers) + Example:: + >>> # import the new functional optimizer + >>> # xdoctest: +SKIP + >>> from xyz import fn_optimizer + >>> from torch.distributed.optim.utils import register_functional_optim + >>> fn_optim_key = "XYZ_optim" + >>> register_functional_optim(fn_optim_key, fn_optimizer) + """ + if key not in functional_optim_map: + functional_optim_map[key] = optim + + +def as_functional_optim(optim_cls: Type, *args, **kwargs): + try: + functional_cls = functional_optim_map[optim_cls] + except KeyError as e: + raise ValueError( + f"Optimizer {optim_cls} does not have a functional " f"counterpart!" + ) from e + + return _create_functional_optim(functional_cls, *args, **kwargs) + + +def _create_functional_optim(functional_optim_cls: Type, *args, **kwargs): + return functional_optim_cls( + [], + *args, + **kwargs, + _allow_empty_param_list=True, + ) diff --git a/lib/python3.12/site-packages/torch/distributed/optim/zero_redundancy_optimizer.py b/lib/python3.12/site-packages/torch/distributed/optim/zero_redundancy_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..7eed2dc387607305bd2aa8f61a1bfa7d405c0427 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/optim/zero_redundancy_optimizer.py @@ -0,0 +1,1652 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. +# +# This source code is licensed under the BSD license found in the +# LICENSE file in the root directory of this source tree. + +r"""Zero Redundancy Optimizer.""" +import collections +import copy +import enum +import inspect +import io +import logging +from itertools import chain +from typing import Any, Callable, Dict, List, Optional, Set, Type, Union + +import torch +import torch.distributed as dist +from torch.distributed.algorithms.join import Join, Joinable, JoinHook +from torch.distributed.optim.utils import functional_optim_map +from torch.optim import Optimizer + + +__all__ = ["ZeroRedundancyOptimizer"] + + +logger = logging.getLogger(__name__) + + +# Credits: classy_vision/generic/distributed_util.py +def _recursive_copy_to_device( + value: Any, + non_blocking: bool, + device: torch.device, +) -> Any: + r""" + Recursively searches lists, tuples, dicts and copies tensors to device if possible. + + Non-tensor values are passed as-is in the result. + + .. note: These are all copies, so if there are two objects that reference + the same object, then after this call, there will be two different objects + referenced on the device. + """ + if isinstance(value, torch.Tensor): + return value.to(device, non_blocking=non_blocking) + + if isinstance(value, (list, tuple)): + values = [ + _recursive_copy_to_device(val, non_blocking=non_blocking, device=device) + for val in value + ] + return values if isinstance(value, list) else tuple(values) + + if isinstance(value, collections.abc.Mapping): + return { + key: _recursive_copy_to_device( + val, non_blocking=non_blocking, device=device + ) + for key, val in value.items() + } + + return value + + +def _is_trainable(param: torch.Tensor) -> bool: + r"""Return if a parameter is trainable, where trainability is equivalent to requiring a gradient.""" + return param.requires_grad + + +def _broadcast_object( + obj: Any, + src_rank: int, + group: object = dist.group.WORLD, + device: torch.device = torch.device("cpu"), +) -> Any: + r""" + Broadcasts an object to the given group. + + It will be sending the object if called from the source rank and receiving + the object otherwise. + + Arguments: + obj: object to broadcast; only used if called on the source rank. + src_rank (int): source rank. + group (``ProcessGroup``, optional): group used for the broadcast + (default: ``dist.group.WORLD``). + device (``torch.device``, optional): device to send from or receive + to (default: ``torch.device("cpu")``). + + Returns: + The broadcasted object. + """ + if dist.get_rank() == src_rank: + # Send the object + buffer = io.BytesIO() + torch.save(obj, buffer) + data = bytearray(buffer.getbuffer()) + length_tensor = torch.LongTensor([len(data)]).to(device) + data_send_tensor = torch.ByteTensor(data).to(device) + dist.broadcast(length_tensor, src=src_rank, group=group, async_op=False) + dist.broadcast(data_send_tensor, src=src_rank, group=group, async_op=False) + else: + # Receive the object + length_tensor = torch.LongTensor([0]).to(device) + dist.broadcast(length_tensor, src=src_rank, group=group, async_op=False) + data_recv_tensor = torch.empty( + [int(length_tensor.item())], dtype=torch.uint8, device=device + ) + dist.broadcast(data_recv_tensor, src=src_rank, group=group, async_op=False) + buffer = io.BytesIO(data_recv_tensor.cpu().numpy()) + obj = torch.load(buffer, map_location=device, weights_only=False) + return obj + + +class _ZeROJoinHook(JoinHook): + def __init__(self, zero): + assert isinstance(zero, ZeroRedundancyOptimizer), ( + "ZeRO join hook requires passing in a ZeroRedundancyOptimizer " + "instance as the state" + ) + self.zero = zero + super().__init__() + + def main_hook(self): + """ + Perform an optimizer step. + + This step updates the joined process's shard of + the parameters and broadcasts those parameters. + """ + self.zero.step() + + +class _DDPBucketAssignment: + r""" + Represent a :class:`DistributedDataParallel` bucket assignment. + + This means that a (possibly non-strict) subset of the parameters corresponding to + a DDP bucket assigned to a rank to update. + + Attributes: + bucket_index (int): index of the bucket determined by the DDP gradient + bucket all-reduce order. + parameters (List[torch.Tensor]): model parameters in the bucket + assigned to this rank. + offset (int): offset into the :class:`GradBucket` 's :meth:`parameters` + giving the index of the first element in the passed-in + ``parameters``; this equivalently indexes into the + :class:`GradBucket` 's :meth:`gradients`. + device (torch.device): device on which the parameters are stored. + tensor (torch.Tensor): flattened tensor giving the data of the + parameter subset assigned to the rank. + """ + + def __init__( + self, + bucket_index: int, + parameters: List[torch.Tensor], + offset: int, + ): + self.bucket_index = bucket_index + self.parameters = parameters + self.offset = offset + if len(self.parameters) == 0: + raise ValueError("Empty bucket assignment") + # DDP guarantees all parameters in the bucket have the same device + self.device: torch.device = self.parameters[0].device + self.tensor: Optional[torch.Tensor] = None + + +class _OverlapStatus(enum.IntEnum): + r""" + Define possible statuses that :class:`ZeroRedundancyOptimizer` can be in when overlapping with :class:`DistributedDataParallel`. + + Attributes: + ``UNINITIALIZED``: The ZeRO instance is effectively uninitialized and + is waiting for DDP to finalize its bucketing. + ``DDP_HAS_REBUILT_BUCKETS``: DDP has rebuilt its buckets, meaning that + its bucketing is finalized. The ZeRO instance can now collect the + necessary information about the DDP bucketing. + ``INITIALIZED``: The ZeRO instance is fully initialized and can now + optimize parameters. + """ + + UNINITIALIZED = 0 + DDP_HAS_REBUILT_BUCKETS = 1 + INITIALIZED = 2 + + +class _OverlapInfo: + r""" + Information needed by :class:`ZeroRedundancyOptimizer` to overlap with :class:`DistributedDataParallel`. + + Arguments: + world_size (int): world size of the process group being used. + + Attributes: + shard_buckets (bool): if ``True``, then the assignment of each + :class:`DistributedDataParallel` bucket is partitioned across + possibly multiple :class:`ZeroRedundancyOptimizer` instances (i.e. + across possibly multiple ranks) to approximate uniformity following + a threshold given by the total parameter size divided by the world + size; if ``False``, then each bucket is wholly assigned to a single + :class:`ZeroRedundancyOptimizer` instance (i.e. to a single rank); + this should be set to the value passed into the hook constructor. + status (_OverlapStatus): current status; see :class:`_OverlapStatus` + for more information. + params_per_bucket (List[List[torch.Tensor]]): ``params_per_bucket[i]`` + gives the model parameters in the ``i``th bucket. + params_per_rank (List[List[torch.Tensor]]): ``params_per_rank[i]`` + gives the model parameters assigned to the ``i``th rank, where the + parameters are grouped by increasing bucket indices. + offsets (Dict[int, int]): maps from bucket index to the offset in + ``self.params_per_rank[rank]`` giving the index of the first + parameter in that bucket, where ``rank`` is this process's own + rank; the keys of this :class:`dict` are the bucket indices + assigned to this rank. + num_bucket_assignments (int): total number of bucket assignments across + all ranks; this is equal to the number of + :class:`DistributedDataParallel` gradient buckets if + ``shard_buckets=False`` and possibly greater otherwise. + total_size (int, optional): total size of all buckets (i.e. sum of + ``param.numel()`` for all ``param`` across all buckets) if + ``shard_buckets=True``; otherwise, ``None``. + broadcast_handles (List[Work]): :class:`list` of async work handles for + the parameter broadcasts. + bucket_index_to_future (Dict[int, torch.futures.Future]): + :class:`dict` mapping bucket index to the corresponding all-reduce + future. + bucket_index_to_bucket (Dict[int, dist.GradBucket]): :class:`dict` + mapping bucket index to the corresponding bucket. + bucket_indices_seen (List[int]): :class:`list` of the bucket indices + seen on this iteration. + """ + + def __init__(self, world_size) -> None: + self.status: _OverlapStatus = _OverlapStatus.UNINITIALIZED + self.shard_buckets: bool = False + + # Modified per bucket reconstruction + self.params_per_bucket: List[List[torch.Tensor]] = [] + self.params_per_rank: List[List[torch.Tensor]] = [[] for _ in range(world_size)] + self.offsets: Dict[int, int] = {} + # Group Ranks + self.assigned_ranks_per_bucket: List[Set[int]] = [] + self.num_bucket_assignments: int = 0 + self.total_size: Optional[int] = None + + # Modified per iteration + self.broadcast_handles: List[Any] = [] + self.bucket_indices_seen: List[int] = [] + # Used by `hook_with_zero_step()` + self.bucket_index_to_future: Dict[int, torch.futures.Future] = {} + self.bucket_index_to_bucket: Dict[int, dist.GradBucket] = {} + + def wait_for_broadcasts(self) -> None: + r""" + Wait for all parameter broadcasts. + + This function should be called once all broadcasts have been scheduled, + meaning ``self.broadcast_handles`` is filled. This clears ``self.broadcast_handles`` + in preparation for the next iteration. + """ + assert ( + len(self.broadcast_handles) == self.num_bucket_assignments + ), f"Missing at least one broadcast handle on rank {dist.get_rank()}" + _ = [x.wait() for x in self.broadcast_handles] + self.broadcast_handles.clear() + + def clear_per_iter_info(self) -> None: + r""" + Clear the data structures that are modified per-iteration. + + This function should be called at the end of an iteration. + """ + self.bucket_indices_seen.clear() + self.bucket_index_to_future.clear() + self.bucket_index_to_bucket.clear() + + +class ZeroRedundancyOptimizer(Optimizer, Joinable): + r""" + Wrap an arbitrary :class:`optim.Optimizer ` and shards its states across ranks in the group. + + The sharing is done as described by ZeRO_. + + The local optimizer instance in each rank is only + responsible for updating approximately ``1 / world_size`` parameters and + hence only needs to keep ``1 / world_size`` optimizer states. After + parameters are updated locally, each rank will broadcast its parameters to + all other peers to keep all model replicas in the same state. + ``ZeroRedundancyOptimizer`` can be used in conjunction with + :class:`torch.nn.parallel.DistributedDataParallel` to reduce per-rank peak + memory consumption. + + ``ZeroRedundancyOptimizer`` uses a sorted-greedy algorithm to pack a number + of parameters at each rank. Each parameter belongs to a single rank and is + not divided among ranks. The partition is arbitrary and might not match the + the parameter registration or usage order. + + Arguments: + params (``Iterable``): an ``Iterable`` of :class:`torch.Tensor` s + or :class:`dict` s giving all parameters, which will be sharded + across ranks. + + Keyword Args: + optimizer_class (:class:`torch.nn.Optimizer`): the class of the local + optimizer. + process_group (``ProcessGroup``, optional): ``torch.distributed`` + ``ProcessGroup`` (default: ``dist.group.WORLD`` initialized by + :meth:`torch.distributed.init_process_group`). + parameters_as_bucket_view (bool, optional): if ``True``, parameters are + packed into buckets to speed up communication, and ``param.data`` + fields point to bucket views at different offsets; if ``False``, + each individual parameter is communicated separately, and each + ``params.data`` stays intact (default: ``False``). + overlap_with_ddp (bool, optional): if ``True``, :meth:`step` is + overlapped with :class:`DistributedDataParallel` 's gradient + synchronization; this requires (1) either a functional optimizer + for the ``optimizer_class`` argument or one with a functional + equivalent and (2) registering a DDP communication hook + constructed from one of the functions in ``ddp_zero_hook.py``; + parameters are packed into buckets matching those in + :class:`DistributedDataParallel`, meaning that the + ``parameters_as_bucket_view`` argument is ignored. + If ``False``, :meth:`step` runs disjointly after the backward pass + (per normal). + (default: ``False``) + **defaults: any trailing arguments, which are forwarded to the local + optimizer. + + Example:: + + >>> # xdoctest: +SKIP + >>> import torch.nn as nn + >>> from torch.distributed.optim import ZeroRedundancyOptimizer + >>> from torch.nn.parallel import DistributedDataParallel as DDP + >>> model = nn.Sequential(*[nn.Linear(2000, 2000).to(rank) for _ in range(20)]) + >>> ddp = DDP(model, device_ids=[rank]) + >>> opt = ZeroRedundancyOptimizer( + >>> ddp.parameters(), + >>> optimizer_class=torch.optim.Adam, + >>> lr=0.01 + >>> ) + >>> ddp(inputs).sum().backward() + >>> opt.step() + + .. warning:: + Currently, ``ZeroRedundancyOptimizer`` requires that all of the + passed-in parameters are the same dense type. + + .. warning:: + If you pass ``overlap_with_ddp=True``, be wary of the following: Given + the way that overlapping :class:`DistributedDataParallel` with + :class:`ZeroRedundancyOptimizer` is currently implemented, the first + two or three training iterations do not perform parameter updates in + the optimizer step, depending on if ``static_graph=False`` or + ``static_graph=True``, respectively. This is because it needs + information about the gradient bucketing strategy used by + :class:`DistributedDataParallel`, which is not finalized until the + second forward pass if ``static_graph=False`` or until the third + forward pass if ``static_graph=True``. To adjust for this, one option + is to prepend dummy inputs. + + .. warning:: ZeroRedundancyOptimizer is experimental and subject to change. + + .. _ZeRO: https://arxiv.org/abs/1910.02054 + + """ + + def __init__( + self, + params, + optimizer_class: Type[Optimizer], + process_group: Optional[Any] = None, + parameters_as_bucket_view: bool = False, + overlap_with_ddp: bool = False, + **defaults: Any, + ): + r"""Init.""" + # Perform type and assumption checks on the input parameters + params = self._verify_and_init_params(params) + self._verify_same_dense_param_type() + + # NOTE: The parent constructor uses `add_param_group()` which is + # partially overloaded in ZeroRedundancyOptimizer, so we use the + # `initialized` flag to dissociate the behaviour of `add_param_group()` + # between the parent and child. + self.initialized = False + + Optimizer.__init__(self, params, defaults) + Joinable.__init__(self) + # Now, all parameters are held in both `self._all_params` and + # `self.param_groups` + + # Internal data structures (`_cache` indicates lazily evaluated) + self._param_to_rank_cache: Dict[torch.Tensor, int] = {} + self._param_to_index_cache: Dict[torch.Tensor, int] = {} + self._partition_parameters_cache: List[List[Dict]] = [] + self._index_to_param_cache: List[torch.Tensor] = [] + self._device_to_params_per_rank_cache: Dict[ + torch.device, List[List[torch.Tensor]] + ] = {} + self._bucket_assignments_per_rank_cache: List[ + Dict[int, _DDPBucketAssignment] + ] = [] + self._is_trainable_mask = self._get_is_trainable_mask() + + # Default device for collective communication and buckets + self._default_device = self._all_params[0].device + + self.process_group = ( + process_group if process_group is not None else dist.group.WORLD + ) + self.world_size: int = dist.get_world_size(self.process_group) + self.rank: int = dist.get_rank(self.process_group) + self.global_rank: int = dist.distributed_c10d.get_global_rank( + self.process_group, self.rank + ) + + self._overlap_with_ddp: bool = overlap_with_ddp + self._optim_defaults = defaults + self._optim_constructor = self._get_optimizer_constructor(optimizer_class) + + # If `overlap_with_ddp=True`, local optimizer initialization is delayed + # to run time after the necessary information has been collected + if not overlap_with_ddp: + self._init_local_optimizer() + else: + self._overlap_info: _OverlapInfo = _OverlapInfo(self.world_size) + if parameters_as_bucket_view: + logger.warning( + "`parameters_as_bucket_view=True` will be ignored since " + "`overlap_with_ddp=True`; instead, a different bucketing " + "strategy will be used" + ) + + # `self._buckets` is used if `parameters_as_bucket_view=True`, in + # which case parameter data is flattened into contiguous bucket tensors + self.parameters_as_bucket_view = parameters_as_bucket_view + self._buckets: List[List[torch.Tensor]] = [] + self._build_param_buckets() + + # Optional consolidated optimizer state, only populated if this rank + # is the target in `consolidate_state_dict()` + self._all_state_dicts: List[Dict[str, Any]] = [] + + self.initialized = True + + def _clear_cache(self) -> None: + r"""Clear the cached data structures giving partition information.""" + self._partition_parameters_cache.clear() + self._param_to_rank_cache.clear() + self._index_to_param_cache.clear() + self._param_to_index_cache.clear() + self._device_to_params_per_rank_cache.clear() + self._bucket_assignments_per_rank_cache.clear() + + def add_param_group(self, param_group: Dict[str, Any]) -> None: + r""" + Add a parameter group to the :class:`Optimizer` 's ``param_groups``. + + This can be useful when fine tuning a pre-trained network, as frozen + layers can be made trainable and added to the :class:`Optimizer` as + training progresses. + + Arguments: + param_group (dict): specifies the parameters to be optimized and + group-specific optimization options. + + .. warning:: This method handles updating the shards on all partitions + but needs to be called on all ranks. Calling this on a subset of + the ranks will cause the training to hang because communication + primitives are called depending on the managed parameters and + expect all the ranks to participate on the same set of parameters. + """ + if self.initialized and self._overlap_with_ddp: + raise RuntimeError( + "ZeroRedundancyOptimizer with `overlap_with_ddp=True` only " + "supports a single parameter group" + ) + + super().add_param_group(param_group) + # NOTE: The rest of the method assumes that the call to the parent's + # `add_param_group()` appends the new parameter group and preserves + # the previous parameter-group ordering + + if self.initialized: + # Force a re-partitioning of the parameters + self._clear_cache() + param_groups = self._partition_parameters()[self.rank] + # NOTE: All parameters in the old parameter groups should be + # assigned to the same ranks so that the local optimizers do not + # need to be reinitialized + + # Add the parameters assigned to this rank from the new parameter + # group to the local optimizer, if any + if len(param_groups) == len(self.optim.param_groups) + 1: + self.optim.add_param_group(param_groups[-1]) + + # Update the bucketing strategy accordingly + if self.parameters_as_bucket_view: + self._build_param_buckets() + + def consolidate_state_dict(self, to: int = 0) -> None: + r""" + Consolidate a list of ``state_dict`` s (one per rank) on the target rank. + + Arguments: + to (int): the rank that receives the optimizer states (default: 0). + + Raises: + RuntimeError: if ``overlap_with_ddp=True`` and this method is + called before this :class:`ZeroRedundancyOptimizer` instance + has been fully initialized, which happens once + :class:`DistributedDataParallel` gradient buckets have been + rebuilt. + + .. warning:: This needs to be called on all ranks. + """ + self._check_overlap_initialized() + + # Sync the exposed `param_groups` attributes to the local optimizer in + # case they have been updated + self._sync_param_groups(self.param_groups, self.optim.param_groups) + + # Pull the sharded state from all ranks and store them in rank order + empty_messenger = torch.tensor( + [0], dtype=torch.uint8, device=self._default_device + ) + + # NOTE: We wastefully use `broadcast()` (e.g. instead of `gather()`) + # due to compatibility issues with NCCL backend; a possible follow-up + # is to move all sharded state management to RPC RRef + self._all_state_dicts = [] + for rank in range(self.world_size): + global_rank = dist.distributed_c10d.get_global_rank( + self.process_group, rank + ) + if self.rank == to: + # Consolidate all local `state_dict`s on this rank, storing on + # CPU to save GPU memory + if rank == self.rank: + # Directly append own optimizer state + self._all_state_dicts.append( + _recursive_copy_to_device( + self.optim.state_dict(), + non_blocking=True, + device=torch.device("cpu"), + ) + ) + else: + # Receive the optimizer state from the source rank + local_state_dict = _broadcast_object( + empty_messenger, + src_rank=global_rank, + group=self.process_group, + device=self._default_device, + ) + self._all_state_dicts.append( + _recursive_copy_to_device( + local_state_dict, + non_blocking=True, + device=torch.device("cpu"), + ) + ) + else: + if rank == self.rank: + # Send the optimizer state to the target rank + _ = _broadcast_object( + self.optim.state_dict(), + src_rank=self.global_rank, + group=self.process_group, + device=self._default_device, + ) + elif rank != to: + # Discard the received object; `broadcast()` is used for + # compatibility reasons + _ = _broadcast_object( + empty_messenger, + src_rank=global_rank, + group=self.process_group, + device=self._default_device, + ) + + def _verify_params_per_rank( + self, + params_per_rank: List[List[torch.Tensor]], + ) -> None: + r""" + Verify ``params_per_rank`` for :meth:`_partition_parameters`. + + The verification is done by checking that ``params_per_rank`` has length equal + to the world size and that it does not contain any parameters not passed into the + :class:`ZeroRedundancyOptimizer` constructor. + + The parameters in ``params_per_rank`` being a strict subset of those + passed into the constructor is valid since some parameters may be + frozen. + + Raises: + ValueError: if ``params_per_rank`` does not have length equal to + the world size or if it contains a parameter that was not + passed into the :class:`ZeroRedundancyOptimizer` constructor. + """ + if len(params_per_rank) != self.world_size: + raise ValueError( + "`params_per_rank` must have length equal to the world size" + ) + all_params_set = set(self._all_params) + for params in params_per_rank: + for param in params: + if param not in all_params_set: + raise ValueError( + "Passing a new parameter in `params_per_rank` that " + "was not passed into the ZeroRedundancyOptimizer " + "constructor" + ) + + def _partition_param_group( + self, param_group: Dict[str, Any], params_per_rank: List[List[torch.Tensor]] + ) -> None: + r""" + Partition the parameter group ``param_group`` according to ``params_per_rank``. + + The partition will modify the ``self._partition_parameters_cache``. This method should + only be used as a subroutine for :meth:`_partition_parameters`. + + Arguments: + param_group (dict[str, Any]): a parameter group as normally defined + in an optimizer state. + params_per_rank (list[list[torch.Tensor]]): a :class:`list` of + length world size containing :class:`list` s of parameters to + assign to each rank. + """ + for rank, params in enumerate(params_per_rank): + rank_param_group = copy.copy(param_group) + rank_param_group["params"] = params + self._partition_parameters_cache[rank].append(rank_param_group) + + def _partition_parameters( + self, + params_per_rank: Optional[List[List[torch.Tensor]]] = None, + ) -> List[List[Dict]]: + r""" + Partitions parameters across distributed data parallel ranks. + + Arguments: + params_per_rank (list[list[torch.Tensor]], optional): a + :class:`list` of length world size containing :class:`list` s + of parameters to assign to each rank; this provides a way to + specify a partition manually. + If ``None``, the parameters are partitioned according to an + internal algorithm. + (default: ``None``) + + Returns: + A :class:`list` where each element of the list contains the + ``param_groups`` for a rank (which itself is a :class:`list` of + :class:`dict`); element 0 corresponds to rank 0, etc.; each rank + stores the ``param_groups`` for all ranks for the collective + communication in :meth:`step`. + + Raises: + ValueError: see :meth:`_validate_params_per_rank`. + RuntimeError: if ``params_per_rank`` is not ``None`` and this + :class:`ZeroRedundancyOptimizer` instance is using more than + one parameter group. + """ + if params_per_rank is None: + # Partition the parameters optimizing for uniformity + if len(self._partition_parameters_cache) == 0: + self._partition_parameters_cache = [[] for _ in range(self.world_size)] + sizes = [0] * self.world_size + for param_group in self.param_groups: + param_group_params_per_rank: List[List] = [ + [] for _ in range(self.world_size) + ] + # Sort the parameters by size (largest first) + params_sorted = sorted( + param_group["params"], key=lambda t: t.numel(), reverse=True + ) + for param in params_sorted: + # Greedily add the parameter to rank with smallest size so far + rank = self._get_min_index(sizes) + param_group_params_per_rank[rank].append(param) + sizes[rank] += param.numel() + # Apply the constructed partition of the parameter group + self._partition_param_group( + param_group, param_group_params_per_rank + ) + + return self._partition_parameters_cache + + # Partition the parameters according to `params_per_rank` + assert len(self._partition_parameters_cache) == 0, ( + "Specifying `params_per_rank` should only be done when the " + "parameters have not been partitioned yet" + ) + if len(self.param_groups) != 1: + raise RuntimeError( + "Specifying `params_per_rank` only supports a single parameter group" + ) + self._verify_params_per_rank(params_per_rank) + self._partition_parameters_cache = [[] for _ in range(self.world_size)] + + # Apply the passed-in partition of the parameter group + param_group = self.param_groups[0] + self._partition_param_group(param_group, params_per_rank) + + return self._partition_parameters_cache + + @property + def _param_to_rank(self) -> Dict[torch.Tensor, int]: + r""":class:`dict` mapping parameters to their assigned data parallel rank in the partition.""" + if len(self._param_to_rank_cache) == 0: + for rank, param_groups in enumerate(self._partition_parameters()): + for param_group in param_groups: + for param in param_group["params"]: + self._param_to_rank_cache[param] = rank + return self._param_to_rank_cache + + @property + def _param_to_index(self) -> Dict[torch.Tensor, int]: + r""" + :class:`dict` mapping parameters to their indices in the global optimizer state. + + NOTE: This assumes that the global optimizer state's indexing (in + ``state_dict``) follows a linear ordering over the parameter groups. + """ + if len(self._param_to_index_cache) == 0: + self._param_to_index_cache = { + p: i + for i, p in enumerate(chain(*(g["params"] for g in self.param_groups))) + } + return self._param_to_index_cache + + @property + def _index_to_param(self) -> List[torch.Tensor]: + r"""List mapping parameter indices in the global optimizer scheme to the actual params.""" + if len(self._index_to_param_cache) == 0: + self._index_to_param_cache = list( + chain(*(g["params"] for g in self.param_groups)) + ) + return self._index_to_param_cache + + def _broadcast_params_from_rank(self, rank: int): + r""" + Broadcast the shard of parameters from a given rank to all other ranks asynchronously. + + Arguments: + rank (int): the source rank. + + Returns: + A :class:`list` of async work handles for the ``broadcast()`` s + performed to synchronize the parameters. + """ + assert not self._overlap_with_ddp, ( + "`_broadcast_params_from_rank()` should not be used if " + "`overlap_with_ddp=True`; instead, the broadcasting should " + "happen in the DDP communication hook" + ) + handles = [] + if self.parameters_as_bucket_view: + for dev_i_buckets in self._buckets: + bucket = dev_i_buckets[rank] + global_rank = dist.distributed_c10d.get_global_rank( + self.process_group, rank + ) + handles.append( + dist.broadcast( + tensor=bucket, + src=global_rank, + group=self.process_group, + async_op=True, + ) + ) + else: + param_groups = self._partition_parameters()[rank] + global_rank = dist.distributed_c10d.get_global_rank( + self.process_group, rank + ) + for param_group in param_groups: + handles.extend( + dist.broadcast( + tensor=param.data, + src=global_rank, + group=self.process_group, + async_op=True, + ) + for param in param_group["params"] + ) + return handles + + def _sync_params(self): + r""" + Sync all parameter shards across the ranks. + + This rank sends its shard of the parameters to all other ranks and + receives a shard from each other rank. This is done using + ``broadcast()``. Parameters are sent bucket-by-bucket if + ``parameters_as_bucket_view=True``and sent parameter-by-parameter + otherwise. + """ + handles = [] + for rank in range(self.world_size): + handles.extend(self._broadcast_params_from_rank(rank)) + _ = [x.wait() for x in handles] + + @property + def _device_to_params_per_rank( + self, + ) -> Dict[torch.device, List[List[torch.Tensor]]]: + r""" + Return device parameters assigned per rank. + + :class:`dict` mapping each device to a :class:`list` of the per-rank parameter + lists filtered to only include the parameters stored on that device. + Each per-rank parameter list gives the parameters assigned to that rank + to update. + + This is used for constructing the parameter buckets if + ``parameters_as_bucket_view=True``. + + Let ``dev_i`` denote the ``i``th device for this rank. Then: + ``dev_0`` maps to a list containing: + rank 0's assigned parameters stored on ``dev_0``, + rank 1's assigned parameters stored on ``dev_0``, + ... + ``dev_1`` maps to a list containing: + rank 0's assigned parameters stored on ``dev_1``, + rank 1's assigned parameters stored on ``dev_1``, + ... + ... + """ + assert self.parameters_as_bucket_view, ( + "`_device_to_params_per_rank` should only be used if " + "`parameters_as_bucket_view=True`" + ) + if len(self._device_to_params_per_rank_cache) == 0: + for rank, param_groups in enumerate(self._partition_parameters()): + for param_group in param_groups: + for param in param_group["params"]: + device = param.device + if device not in self._device_to_params_per_rank_cache: + self._device_to_params_per_rank_cache[device] = [ + [] for _ in range(self.world_size) + ] + self._device_to_params_per_rank_cache[device][rank].append( + param + ) + return self._device_to_params_per_rank_cache + + def _get_min_index( + self, + values: List[int], + disallowed_indices: Optional[Set[int]] = None, + ) -> int: + r""" + Return ``values.index(min(values))``, except only uses one pass. + + It also excludes any indices in ``disallowed_indices`` if provided. + + Arguments: + values: (List[int]): :class:`list` of values. + disallowed_indices (Optional[Set[int]]): indices that are + disallowed from being the returned min index. + """ + min_index = -1 + min_value = float("inf") + for i, value in enumerate(values): + if disallowed_indices and i in disallowed_indices: + continue + if value < min_value: + min_value = value + min_index = i + assert min_index >= 0, "All indices are disallowed" + return min_index + + def _assign_bucket_subset_to_rank( + self, + bucket_index: int, + bucket_params: List[torch.Tensor], + bucket_offset: int, + assigned_rank: int, + assigned_ranks_per_bucket: List[Set[int]], + ) -> None: + r""" + Assign ``bucket_params`` to the rank with the least size assigned so far and collects relevant information. + + The model parameters given by ``bucket_params`` represents a (possibly non-strict) + subset of the parameters corresponding to a :class:`DistributedDataParallel` bucket. + + Arguments: + bucket_index (int): index of the :class:`DistributedDataParallel` + gradient bucket. + bucket_params (List[torch.Tensor]): subset of the parameters + corresponding to the bucket to assign. + bucket_offset (int): offset giving the index of the first element + in ``bucket_params`` in the bucket's full parameter list. + assigned_rank (int): group rank to assign to. + assigned_ranks_per_bucket (List[Set[int]]): :class:`set` of group ranks + assigned to each bucket. + """ + overlap_info = self._overlap_info + if len(bucket_params) == 0: + raise ValueError("Empty bucket assignment") + params_per_rank = overlap_info.params_per_rank + offsets = overlap_info.offsets + + self._bucket_assignments_per_rank_cache[assigned_rank][ + bucket_index + ] = _DDPBucketAssignment(bucket_index, bucket_params, bucket_offset) + if self.global_rank == assigned_rank: + offsets[bucket_index] = len(params_per_rank[assigned_rank]) + params_per_rank[assigned_rank].extend(bucket_params) + assigned_ranks_per_bucket[bucket_index].add(assigned_rank) + self._overlap_info.num_bucket_assignments += 1 + + @property + def _bucket_assignments_per_rank(self) -> List[Dict[int, _DDPBucketAssignment]]: + r""" + Return DDP bucket parameters assigned per rank. + + :class:`list` of length world size consisting of :class:`dict` s + mapping bucket indices to :class:`_DDPBucketAssignment` s for each + rank. + """ + assert ( + self._overlap_with_ddp + ), "`_bucket_assignments_per_rank` only be used if `overlap_with_ddp=True`" + if len(self._bucket_assignments_per_rank_cache) > 0: + return self._bucket_assignments_per_rank_cache + + overlap_info = self._overlap_info + assert overlap_info.status == _OverlapStatus.INITIALIZED + + self._bucket_assignments_per_rank_cache = [{} for _ in range(self.world_size)] + params_per_bucket = overlap_info.params_per_bucket + + if overlap_info.shard_buckets: + # Define the assignment threshold to approximate uniformity + assert overlap_info.total_size is not None, "`total_size` was not computed" + threshold = overlap_info.total_size / self.world_size # type: ignore[operator] + size_per_rank = [0 for _ in range(self.world_size)] + + num_buckets = len(params_per_bucket) + overlap_info.assigned_ranks_per_bucket = [set() for _ in range(num_buckets)] + assigned_ranks_per_bucket = overlap_info.assigned_ranks_per_bucket + if not overlap_info.shard_buckets: + # Assign each DDP bucket entirely to a single rank + for bucket_index, bucket_params in enumerate(params_per_bucket): + assert len(bucket_params) > 0, "Empty bucket" + assigned_rank = self._get_assigned_rank(bucket_index) + self._assign_bucket_subset_to_rank( + bucket_index, + bucket_params, + 0, + assigned_rank, + assigned_ranks_per_bucket, + ) + else: + # Assign each DDP bucket to possibly multiple ranks + # Specifically, sort the DDP buckets by increasing size, and for + # each bucket, iteratively assign the maximal unassigned subset + # with size less than `threshold` to the rank with the least total + # size so far -- each such assignment is represented by a + # `_DDPBucketAssignment` instance and only contains parameters from + # a single DDP bucket + params_per_bucket_enum = sorted( + enumerate(params_per_bucket), key=lambda x: sum(p.numel() for p in x[1]) + ) + for bucket_index, bucket_params in params_per_bucket_enum: + assert len(bucket_params) > 0, "Empty bucket" + bucket_offset = 0 + assignment_size = 0 + for param_index, param in enumerate(bucket_params): + param_numel = param.numel() + if ( + assignment_size + param_numel >= threshold + and param_index > bucket_offset + ): + assigned_rank = self._get_min_index( + size_per_rank, assigned_ranks_per_bucket[bucket_index] + ) + # Include up to but not including the parameter that + # exceeded the threshold + self._assign_bucket_subset_to_rank( + bucket_index, + bucket_params[bucket_offset:param_index], + bucket_offset, + assigned_rank, + assigned_ranks_per_bucket, + ) + size_per_rank[assigned_rank] += assignment_size + bucket_offset = param_index + assignment_size = 0 + assignment_size += param_numel + # Assign the remainder of the bucket so that no assignment + # spans across two buckets + assigned_rank = self._get_min_index( + size_per_rank, assigned_ranks_per_bucket[bucket_index] + ) + self._assign_bucket_subset_to_rank( + bucket_index, + bucket_params[bucket_offset:], + bucket_offset, + assigned_rank, + assigned_ranks_per_bucket, + ) + size_per_rank[assigned_rank] += assignment_size + + return self._bucket_assignments_per_rank_cache + + def _local_step( + self, + gradients: Optional[List[Optional[torch.Tensor]]] = None, + closure: Optional[Callable[[], float]] = None, + **kwargs: Any, + ) -> Optional[float]: + r""" + Perform a single optimizer step without syncing parameters across ranks. + + Arguments: + gradients (list[Optional[torch.Tensor]], optional): a :class:`list` + of length equal to the number of parameters assigned to this + rank containing gradient tensors or ``None`` as its elements; + a ``None`` in the :class:`list` indicates that the + corresponding parameter should not be updated. + If the argument itself is ``None``, then all parameters are + updated, and the gradients are assumed to be already populated. + (default: ``None``) + closure (Callable): a closure that re-evaluates the model and + returns the loss; optional for most optimizers and should be + ``None`` if ``gradients`` is not ``None``; (default: ``None``) + Returns: + Optional loss depending on the underlying local optimizer. + + .. warning:: + The argument ``gradients`` should only be specified (i.e. not + ``None``) if ``overlap_with_ddp=True``, in which case + :class:`ZeroRedundancyOptimizer` wraps a functional optimizer. + """ + Join.notify_join_context(self) + # Check if the model trainability has changed + is_trainable_mask = self._get_is_trainable_mask() + if is_trainable_mask != self._is_trainable_mask: + if self._overlap_with_ddp: + raise RuntimeError( + "ZeroRedundancyOptimizer with `overlap_with_ddp=True` " + "does not support changing parameter trainability at run " + "time" + ) + logger.warning( + "ZeroRedundancyOptimizer detected that the trainable " + "parameters changed; rebuilding the parameter buckets if " + "enabled" + ) + self._build_param_buckets() + self._is_trainable_mask = is_trainable_mask + + # Sync the exposed `param_groups` attributes to the local optimizer in + # case they have been updated + self._sync_param_groups(self.param_groups, self.optim.param_groups) + + # Run the optimizer step on this shard only + if gradients is None: + loss = ( + self.optim.step(**kwargs) + if closure is None + else self.optim.step(closure=closure, **kwargs) + ) + else: + assert self._overlap_with_ddp, ( + "Specifying `gradients` should not " + "be used when `overlap_with_ddp=False`" + ) + assert ( + closure is None + ), "`closure` is not supported when using a local functional optimizer" + loss = self.optim.step(gradients=gradients) + + # Sync any updated attributes in the local optimizer to the exposed + # `param_groups` + self._sync_param_groups(self.optim.param_groups, self.param_groups) + + return loss + + def step( + self, + closure: Optional[Callable[[], float]] = None, + **kwargs: Any, + ) -> Optional[float]: + r""" + Perform a single optimizer step and syncs parameters across all ranks. + + Arguments: + closure (Callable): a closure that re-evaluates the model and + returns the loss; optional for most optimizers. + Returns: + Optional loss depending on the underlying local optimizer. + + .. note: Any extra parameters are passed to the base optimizer as-is. + """ + if self._overlap_with_ddp: + logger.warning( + "`step()` should not be included in the training loop when " + "`overlap_with_ddp=True`" + ) + return None + + # Perform the local optimizer step + loss = self._local_step(closure=closure, **kwargs) + + # Sync all of the updated parameter shards across the ranks + self._sync_params() + + return loss + + def join_hook(self, **kwargs): + r""" + Return the ZeRO join hook. + + It enables training on uneven inputs by + shadowing the collective communications in the optimizer step. + + Gradients must be properly set before this hook is called. + + Arguments: + kwargs (dict): a :class:`dict` containing any keyword arguments + to modify the behavior of the join hook at run time; all + :class:`Joinable` instances sharing the same join context + manager are forwarded the same value for ``kwargs``. + + This hook does not support any keyword arguments; i.e. ``kwargs`` is + unused. + """ + return _ZeROJoinHook(self) + + @property + def join_device(self) -> torch.device: + r"""Return default device.""" + return self._default_device + + @property + def join_process_group(self) -> Any: + r"""Return process group.""" + return self.process_group + + def load_state_dict(self, state_dict: Dict[str, Any]) -> None: + r""" + Load the state pertaining to the given rank from the input ``state_dict``, updating the local optimizer as needed. + + Arguments: + state_dict (dict): optimizer state; should be an object returned + from a call to :meth:`state_dict`. + + Raises: + RuntimeError: if ``overlap_with_ddp=True`` and this method is + called before this :class:`ZeroRedundancyOptimizer` instance + has been fully initialized, which happens once + :class:`DistributedDataParallel` gradient buckets have been + rebuilt. + """ + self._check_overlap_initialized() + + for index, value in state_dict["state"].items(): + param = self._index_to_param[index] + if self._param_to_rank[param] != self.rank: + # Clear any state irrelevant to this rank + state_dict["state"][index] = None + else: + # Load the parameter state to the local optimizer + self.optim.state[param] = _recursive_copy_to_device( + value, non_blocking=True, device=param.device + ) + # Force zero-dimensional tensors (like Adam "step") on CPU + for state_name, state_value in self.optim.state[param].items(): + if torch.is_tensor(state_value) and state_value.dim() == 0: + self.optim.state[param][state_name] = state_value.cpu() + + super().load_state_dict(state_dict) + + # Sync the input state with the exposed and local optimizer states + self._sync_param_groups(state_dict["param_groups"], self.param_groups) + self._sync_param_groups(self.param_groups, self.optim.param_groups) + + def state_dict(self) -> Dict[str, Any]: + r""" + Return the last global optimizer state known to this rank. + + .. warning: + If the state has not been consolidated to this rank, this raises a + runtime error, and even if it has, the state may not be up-to-date, + depending on when :meth:`consolidate_state_dict` was last called. + + Raises: + RuntimeError: if ``overlap_with_ddp=True`` and this method is + called before this :class:`ZeroRedundancyOptimizer` instance + has been fully initialized, which happens once + :class:`DistributedDataParallel` gradient buckets have been + rebuilt; or if this method is called without a preceding call + to :meth:`consolidate_state_dict`. + """ + self._check_overlap_initialized() + + if len(self._all_state_dicts) == 0: + raise RuntimeError( + "Optimizer state has not been consolidated on this rank. " + f"Please call `consolidate_state_dict(to={self.rank})` on " + "all ranks beforehand if you meant to save the global state." + ) + + # Get the possibly-stale global optimizer state that uses global + # parameter indexing + state_dict = super().state_dict() + + # Update the global optimizer state with local state information, + # factoring in the translation from local to global indexing + for rank, local_state_dict in enumerate(self._all_state_dicts): + local_param_groups = local_state_dict["param_groups"] + global_param_groups = self._partition_parameters()[rank] + assert len(local_param_groups) == len( + global_param_groups + ), "Mismatch between number of local and global parameter groups" + + for local_param_group, global_param_group in zip( + local_param_groups, global_param_groups + ): + # `local_param_group` stores local indices, while + # `global_param_group` stores the tensors directly + local_param_indices = local_param_group["params"] + global_params = global_param_group["params"] + + assert len(local_param_indices) == len( + global_params + ), "Mismatch between number of local and global parameters in parameter group" + for local_param_index, global_param in zip( + local_param_indices, global_params + ): + # Update the global parameter state, if any + if local_param_index in local_state_dict["state"]: + global_param_index = self._param_to_index[global_param] + state_dict["state"][global_param_index] = local_state_dict[ + "state" + ][local_param_index] + + # Sort the parameters in the state + state_dict["state"] = dict(sorted(state_dict["state"].items())) + return state_dict + + @staticmethod + def _sync_param_groups( + src_param_groups: List[Dict[Any, Any]], + dst_param_groups: List[Dict[Any, Any]], + ) -> None: + r""" + Sync the attributes from the source parameter groups to the destination parameter groups. + + Example attributes include learning rate or scheduler attributes. The + two parameter groups should have the same length (i.e. same number of + parameter groups). + + Arguments: + src_param_groups (list[dict]): parameter groups giving the + attribute settings to copy. + dst_param_groups (list[dict]): parameter groups giving the + attribute settings to set. + """ + assert len(src_param_groups) == len( + dst_param_groups + ), "Mismatch between number of source and destination parameter groups" + for src_param_group, dst_param_group in zip(src_param_groups, dst_param_groups): + # Sync all attributes except the parameters + for attr in filter(lambda x: x != "params", src_param_group.keys()): + dst_param_group[attr] = src_param_group[attr] + + def _build_param_buckets(self) -> None: + r""" + Build parameter buckets if ``parameters_as_bucket_view=True``. + + For each device that stores this rank's parameters, there is a + bucket (represented as a tensor) containing all of the parameters on + that device that are assigned to a given rank in the parameter update + partition. + + This method is called in the constructor and any time parameter + trainability is changed. + + .. warning:: + The current implementation assumes that all of the parameters in a + bucket are of the same dense type when allocating the bucket's + tensor. + + .. warning:: + If the model parameters are stored across more than one device, + then the storage partitioning must be the same across all + processes in order for parameter synchronization to work. + """ + if not self.parameters_as_bucket_view or self._overlap_with_ddp: + return + + # `self._buckets[i][j]` are the parameters stored on device i and + # assigned to rank j + num_devices = len(self._device_to_params_per_rank) + self._buckets = [[] for _ in range(num_devices)] # type: ignore[assignment] + + for dev_i, (device, params_per_rank) in enumerate( + self._device_to_params_per_rank.items() + ): + for params in params_per_rank: + bucket_size = 0 + dtype = None + trainable_params = [] + for param in params: + if not _is_trainable(param): + # Clone in case the parameter was previously part of + # a bucket to avoid the data from being destroyed + param.data = param.data.detach().clone() + else: + bucket_size += param.numel() + trainable_params.append(param) + dtype = param.dtype # assumes all same dtype + + if bucket_size == 0: + # Create a dummy bucket if there are no parameters + bucket = torch.zeros(1, device=device) + else: + # Construct the bucket (assuming all dense and same dtype) + bucket = torch.empty(bucket_size, dtype=dtype, device=device) + offset = 0 + for param in trainable_params: + offset_next = offset + param.numel() + bucket[offset:offset_next].copy_(param.data.flatten()) + param.data = bucket[offset:offset_next].view_as(param.data) + offset = offset_next + self._buckets[dev_i].append(bucket) # type: ignore[arg-type] + + def _build_ddp_param_buckets(self) -> None: + r""" + Build the DDP bucket with parameters assigned to this rank. + + For each DDP bucket with parameters assigned to this rank, flattens the + data of those parameters into a single tensor and saves the tensor to + the ``tensor`` attribute in the corresponding + :class:`_DDPBucketAssignment` instance stored in + ``self._bucket_assignments_per_rank``. + + :class:`DistributedDataParallel` guarantees that the parameters + corresponding to a gradient bucket have the same device and the same + dtype. + """ + for bucket_assignments in self._bucket_assignments_per_rank: + for bucket_assignment in bucket_assignments.values(): + params = bucket_assignment.parameters + bucket_size = 0 + dtype = None + for param in params: + assert _is_trainable(param), ( + "Model parameter " + "corresponding to a gradient in a DDP bucket should " + "require a gradient" + ) + bucket_size += param.numel() + dtype = param.dtype # assumes all same dtype + assert bucket_size > 0, "Empty bucket" + + # Construct the bucket tensor (assuming all dense and same dtype) + tensor = torch.empty( + bucket_size, dtype=dtype, device=bucket_assignment.device + ) + offset = 0 + for param in params: + offset_next = offset + param.numel() + tensor[offset:offset_next].copy_(param.data.flatten()) + param.data = tensor[offset:offset_next].view_as(param.data) + offset = offset_next + bucket_assignment.tensor = tensor + + def _verify_and_init_params( + self, + params: Any, + ) -> Union[List[torch.Tensor], List[dict]]: + r""" + Verify the type of ``params`` and initializes ``self._all_params`` as a :class:`list` of all parameters. + + The initializagtion will first make sure that provided ``params`` is valid. + + Arguments: + params (Any): Candidate parameter list or parameter groups to verify. + + Raises: + TypeError: ``params`` has an invalid type. + ValueError: ``params`` is empty. + + Returns: + The persistent form of ``params`` to be passed into the parent + :class:`Optimizer` constructor -- i.e. returns ``params`` as a + :class:`list` to ensure that it can be iterated over again. + """ + if isinstance(params, torch.Tensor): + raise TypeError( + "`params` argument should be an iterable of " + f"Tensors, but got {torch.typename(params)}" + ) + try: + all_params = list(params) + except TypeError as e: + raise TypeError( + "`params` argument should be an iterable of Tensors" + f" or dicts, but got {torch.typename(params)}" + ) from e + if len(all_params) == 0: + raise ValueError("ZeroRedundancyOptimizer got an empty parameter list") + all_tensors = True + all_dicts = True + for param in all_params: + all_tensors &= isinstance(param, torch.Tensor) + all_dicts &= isinstance(param, dict) + if not all_tensors and not all_dicts: + raise TypeError( + "`params` argument should be an iterable of Tensors or dicts" + ) + # Ensure that `self._all_params` contains a list of all parameters + if all_tensors: + self._all_params = all_params + elif all_dicts: + self._all_params = [] + # `all_params` contains parameter groups (not parameters) + for param_group in all_params: + if "params" not in param_group: + raise ValueError( + "Each parameter group passed-in via `params` must " + "have a 'params' key mapping to the parameters in " + "the group" + ) + self._all_params.extend(param_group["params"]) + return all_params + + def _verify_same_dense_param_type(self) -> None: + r""" + Verify that all parameters are of the same dense type. + + The method assumes that ``self._all_params`` has been initialized + and is non-empty. + + Raises: + ValueError: ``params`` contains sparse parameters or parameters + of varying dense types. + + NOTE: This method can be removed once support for sparse parameters + and varying parameter types is added. + """ + typename = torch.typename(self._all_params[0]) + if self._all_params[0].is_sparse: + raise ValueError( + "ZeroRedundancyOptimizer only supports using " + "the same dense type for all parameters but got " + f"{typename}" + ) + for param in self._all_params[1:]: + other_typename = torch.typename(param) + if other_typename != typename: + raise ValueError( + "ZeroRedundancyOptimizer only supports " + "using the same dense type for all " + f"parameters but got both {typename} and " + f"{other_typename}" + ) + + def _get_is_trainable_mask(self) -> List[bool]: + r"""Return a boolean mask indicating if each parameter is trainable (``requires_grad``) or not.""" + return list(map(_is_trainable, self._all_params)) + + def _init_local_optimizer(self) -> None: + r""" + Initialize this rank's local optimizer, responsible for its subset of the parameters. + + The local optimizer is saved in ``self.optim``. + """ + assert ( + self._optim_constructor is not None + ), "The local optimizer class has not been set" + + param_groups = self._partition_parameters()[self.rank] + # `overlap_with_ddp=True` requires a local functional optimizer + if self._overlap_with_ddp: + # Functional optimizers only support a single parameter group and + # require passing in the parameters as a list + assert len(param_groups) == 1, ( + "Initializing the local " + "functional optimizer with more than one parameter group" + ) + params = param_groups[0]["params"] + # Try to pass `_allow_empty_param_list=True` to avoid erroring + if ( + "_allow_empty_param_list" + in inspect.signature(self._optim_constructor).parameters + ): + self.optim: Any = self._optim_constructor( + params, **self._optim_defaults, _allow_empty_param_list=True + ) + else: + logger.warning( + "%s does not support the argument " + "`_allow_empty_param_list`; ZeroRedundancyOptimizer may " + "error due to an empty parameter list", + self._optim_constructor, + ) + self.optim: Any = self._optim_constructor(params, **self._optim_defaults) # type: ignore[no-redef] + + # Log information about the DDP and ZeRO bucketing + if dist.get_debug_level() != dist.DebugLevel.OFF: + local_numel = sum(p.numel() for p in params) + num_assigned_buckets = len( + self._bucket_assignments_per_rank[self.global_rank] + ) + logger.info( + "rank %s with %s parameters " "across %s buckets", + self.global_rank, + local_numel, + num_assigned_buckets, + ) + if self.global_rank == 0: + logger.info( + "%s DDP " "buckets and " "%s bucket " "assignments", + len(self._overlap_info.params_per_bucket), + self._overlap_info.num_bucket_assignments, + ) + else: + # NOTE: Passing `param_groups` into the local optimizer constructor + # bypasses the empty parameter list check + self.optim: Optimizer = self._optim_constructor(param_groups, **self._optim_defaults) # type: ignore[no-redef] + + # TODO: Manually add `self.param_groups` if using a functional + # optimizer; remove this if/when the functional optimizers support + # multiple parameter groups + if self._overlap_with_ddp and not hasattr(self.optim, "param_groups"): + assert hasattr(self.optim, "param_group"), ( + "The functional optimizer should set at least one of the " + "attributes `param_group` or `param_groups`" + ) + self.optim.param_groups = [self.optim.param_group] # type: ignore[attr-defined] + + self._sync_param_groups(self.optim.param_groups, self.param_groups) + + def _init_zero_for_overlap(self) -> None: + r"""Perform a delayed initialization of the local optimizer and the supporting data structures.""" + assert self._overlap_with_ddp, ( + "`_init_zero_for_overlap()` should only be called when " + "`overlap_with_ddp=True`" + ) + self._overlap_info.status = _OverlapStatus.INITIALIZED + self._clear_cache() + self._partition_parameters(self._overlap_info.params_per_rank) + self._build_ddp_param_buckets() + self._init_local_optimizer() + + def _get_assigned_rank(self, bucket_index: int) -> int: + r""" + Return the single rank assigned to a :class:`DistributedDataParallel` gradient bucket. + + Arguments: + bucket_index (int): index of the :class:`DistributedDataParallel` + bucket for which to get the assigned rank. + """ + assert not self._overlap_info.shard_buckets, ( + "The bucket assignment requires global bucket information and " + "will be computed later; there should be no need to use this " + "method" + ) + return bucket_index % self.world_size + + def _check_overlap_initialized(self): + r""" + Check the delayed initialization depending on the value of ``overlap_with_ddp``. + + The delayed initialization has occurred (see + :meth:`_init_zero_for_overlap`) if ``overlap_with_ddp=True``, and + raises a ``RuntimeError`` if not. This should preface methods that + should not be run before that delayed initialization. + + Raises: + RuntimeError: if ``overlap_with_ddp=True`` and + :meth:`_init_zero_for_overlap` has not been called. + """ + if ( + self._overlap_with_ddp + and self._overlap_info.status != _OverlapStatus.INITIALIZED + ): + raise RuntimeError( + "This method should not be called until this " + "ZeroRedundancyOptimizer instance has been fully " + "initialized" + ) + + def _get_optimizer_constructor(self, optimizer_class: Any) -> Any: + r""" + Return the optimizer constructor using validation and transformation depending on ``overlap_with_ddp``. + + Returns: + - ``optimizer_class`` if ``overlap_with_ddp=False`` and + ``optimizer_class`` is not a functional optimizer. + - ``optimizer_class`` if ``overlap_with_ddp=True`` and + ``optimizer_class`` is already a functional optimizer. + - The functional equivalent of ``optimizer_class`` if + ``overlap_with_ddp=True`` and ``optimizer_class`` is not + already a functional optimizer (assuming the equivalent + exists). + + Raises: + ValueError: + + - if ``overlap_with_ddp=True`` but ``optimizer_class`` is + neither a functional optimizer nor translatable to a + functional optimizer. + - if ``overlap_with_ddp=False`` and ``optimizer_class`` is a + functional optimizer. + """ + functional_optims = functional_optim_map.values() + if not self._overlap_with_ddp: + if optimizer_class in functional_optims: + # Using a functional optimizer is only supported when + # `overlap_with_ddp=True` + raise ValueError( + f"Passing in a functional optimizer {optimizer_class} " + "when `overlap_with_ddp=False`" + ) + else: + return optimizer_class + else: + if optimizer_class in functional_optims: + # Already a functional optimizer + return optimizer_class + elif optimizer_class in functional_optim_map: + # Translate the passed-in optimizer class to its functional + # equivalent if `overlap_with_ddp=True` + optim_constructor = functional_optim_map[optimizer_class] + logger.info( + "Using the functional optimizer %s " + "instead of %s since " + "`overlap_with_ddp=True`", + optim_constructor, + optimizer_class, + ) + return optim_constructor + else: + raise ValueError( + "Using `ddp_with_overlap=True` requires using a " + "functional optimizer, but there is no supported functional " + f"optimizer equivalent for {optimizer_class}" + ) diff --git a/lib/python3.12/site-packages/torch/distributed/remote_device.py b/lib/python3.12/site-packages/torch/distributed/remote_device.py new file mode 100644 index 0000000000000000000000000000000000000000..ab5215e2f83a71e56c6f638bd1453341da671d37 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/remote_device.py @@ -0,0 +1,120 @@ +# mypy: allow-untyped-defs +from typing import Optional, Union + +import torch + + +class _remote_device: + """ + Represents a device on a remote worker. + + Args: + remote_device (str or torch.device): Represents a device on a remote worker. + The string format should be one of the following: + + 1. "/", where the device field can be parsed as torch.device type. + E.g., "trainer0/cpu", "trainer0", "ps0/cuda:0". + In addition, the device field can be optional and the default value is "cpu". + 2. "rank:/", where is the rank of the + process and device can be parsed as torch.device type. + E.g., "rank:0/cpu", "rank:0", "rank:0/cuda:0" + 3. and are optional and formats like "cpu" + and "cuda:1", just represent local devices. + """ + + def __init__(self, remote_device: Union[str, torch.device]): + PARSE_ERROR = ( + f"Could not parse remote_device: {remote_device}. The valid format is " + "'/' or 'rank:/' or ''" + ) + self._worker_name = None + self._rank = None + self._device: Optional[Union[str, int, torch.device]] = None + + if isinstance(remote_device, torch.device): + self._device = remote_device + elif isinstance(remote_device, str): + fields = remote_device.split("/") + if len(fields) == 2: + self._worker_name, self._device = fields + elif len(fields) == 1: + # Check if this is a valid device. + if _remote_device._is_valid_local_device(fields[0]): + self._device = fields[0] + else: + self._worker_name = fields[0] + self._device = "cpu" + else: + raise ValueError(PARSE_ERROR) + else: + raise TypeError(f"Invalid type for remote_device: {type(remote_device)}") + + # Do some basic sanity check (no empty string) + if self._worker_name is not None and not self._worker_name: + raise ValueError(PARSE_ERROR) + + # Validate the device. + self._device = torch.device(self._device) + + # Check for rank based format. + if self._worker_name is not None: + fields = self._worker_name.split(":") + if len(fields) == 2: + # rank:/device format, extract rank + if fields[0] == "rank" and fields[1].isdigit(): + self._rank = int(fields[1]) # type: ignore[assignment] + self._worker_name = None + else: + raise ValueError(PARSE_ERROR) + elif len(fields) > 2: + raise ValueError(PARSE_ERROR) + + @staticmethod + def _is_valid_local_device(device): + # Check for torch.device + try: + torch.device(device) + return True + except Exception: + return False + + def worker_name(self) -> Optional[str]: + """Return the name of remote worker representing the remote device and ``None`` if no worker name is available.""" + return self._worker_name + + def rank(self) -> Optional[int]: + """ + Returns the rank of remote worker representing the remote device. + Returns ``None`` if no rank is available. + """ + return self._rank + + def device(self) -> torch.device: + """Return the local device on the remote worker.""" + return self._device # type: ignore[return-value] + + def __repr__(self): + if self._device is not None: + if self._worker_name is not None: + return f"{self._worker_name}/{self._device}" + elif self._rank is not None: + return f"rank:{self._rank}/{self._device}" + else: + return str(self._device) + else: + if self._worker_name is not None: + return f"{self._worker_name}" + elif self._rank is not None: + return f"{self._rank}" + else: + raise RuntimeError("Invalid state!") + + def __eq__(self, other): + return isinstance(other, _remote_device) and ( + self._worker_name == other._worker_name + and self._device == other._device + and self._rank == other._rank + ) + + def __hash__(self): + return hash(self._worker_name) ^ hash(self._device) ^ hash(self._rank) diff --git a/lib/python3.12/site-packages/torch/distributed/rendezvous.py b/lib/python3.12/site-packages/torch/distributed/rendezvous.py new file mode 100644 index 0000000000000000000000000000000000000000..4686f3ce96869fdb24ce609fe7a6d48e0f37a97b --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/rendezvous.py @@ -0,0 +1,286 @@ +# mypy: allow-untyped-defs +try: + from urllib.parse import urlparse, urlunparse +except ImportError as e: + raise ImportError( + "urllib cannot be found, urlparse from python2 is no longer supported." + ) from e + +import numbers +import os +import sys +from datetime import timedelta +from typing import Callable, Dict, Iterator, Optional, Tuple + +from torch.distributed import FileStore, Store, TCPStore + +from .constants import default_pg_timeout + + +_rendezvous_handlers: Dict[str, Callable[..., Iterator[Tuple[Store, int, int]]]] = {} + +__all__ = ["register_rendezvous_handler", "rendezvous"] + + +def register_rendezvous_handler(scheme, handler): + """ + Register a new rendezvous handler. + + Before we can run collective algorithms, participating processes + need to find each other and exchange information to be able to + communicate. We call this process rendezvous. + + The outcome of the rendezvous process is a triplet containing a + shared key/value store, the rank of the process, and the total + number of participating processes. + + If none of the bundled rendezvous methods apply to your execution + environment you can opt to register your own rendezvous handler. + Pick a unique name and use the URL scheme to identify it when + calling the `rendezvous()` function. + + Args: + scheme (str): URL scheme to identify your rendezvous handler. + handler (function): Handler that is invoked when the + `rendezvous()` function is called with a URL that uses + the corresponding scheme. It must be a generator function + that yields the triplet. + """ + global _rendezvous_handlers + if scheme in _rendezvous_handlers: + raise RuntimeError(f"Rendezvous handler for {scheme}:// already registered") + _rendezvous_handlers[scheme] = handler + + +# Query will have format "rank=0&world_size=1" and is +# converted into {"rank": 0, "world_size": 1} +def _query_to_dict(query: str) -> Dict[str, str]: + return { + pair[0]: pair[1] + for pair in (pair.split("=") for pair in filter(None, query.split("&"))) + } + + +def _get_use_libuv_from_query_dict(query_dict: Dict[str, str]) -> bool: + # libuv is the default backend for TCPStore. To enable the non-libuv backend, + # user can explicitly specify ``use_libuv=0`` in the URL parameter. + return query_dict.get("use_libuv", os.environ.get("USE_LIBUV", "1")) == "1" + + +def _rendezvous_helper(url: str, rank: int, world_size_opt: Optional[int], **kwargs): + result = urlparse(url) + if world_size_opt is None: + world_size = -1 + if result.scheme == "env": + rank = int(os.environ.get("RANK", rank)) + # If the world_size env variable is not present then it is a dynamic group + world_size = int(os.environ.get("WORLD_SIZE", world_size)) + else: + world_size = world_size_opt + if rank != -1 or world_size != -1 or world_size_opt is None: + query_dict = _query_to_dict(result.query) + assert ( + "rank" not in query_dict and "world_size" not in query_dict + ), f"The url: {url} has node-specific arguments(rank, world_size) already." + if rank != -1: + query_dict["rank"] = str(rank) + if world_size != -1 or world_size_opt is None: + query_dict["world_size"] = str(world_size) + result = result._replace( + query=f"{'&'.join([f'{k}={v}' for k, v in query_dict.items()])}" + ) + url = urlunparse(result) + + if result.scheme not in _rendezvous_handlers: + raise RuntimeError(f"No rendezvous handler for {result.scheme}://") + return _rendezvous_handlers[result.scheme](url, **kwargs) + + +def rendezvous(url: str, rank: int = -1, world_size: int = -1, **kwargs): + if not isinstance(url, (str, bytes)): + raise RuntimeError(f"`url` must be a string. {type(url)}: {url}") + + if not isinstance(rank, numbers.Integral): + raise RuntimeError(f"`rank` must be an integer. {rank}") + + if not isinstance(world_size, numbers.Integral): + raise RuntimeError(f"`world_size` must be an integer. {world_size}") + + return _rendezvous_helper(url, rank, world_size, **kwargs) + + +def _create_store_from_options(backend_options, rank): + store, _, _ = next(_rendezvous_helper(backend_options.init_method, rank, None)) + return store + + +def _rendezvous_error(msg): + return ValueError("Error initializing torch.distributed using " + msg) + + +def _file_rendezvous_handler(url: str, **kwargs): + def _error(msg): + return _rendezvous_error("file:// rendezvous: " + msg) + + result = urlparse(url) + path = result.path + if sys.platform == "win32": + import urllib.request + + full_path = result.netloc + result.path + path = urllib.request.url2pathname(full_path) + if path: + # Normalizing an empty string produces ".", which is not expected. + path = os.path.normpath(path) + + if not path: + raise _error("path missing") + query_dict = _query_to_dict(result.query) + if "rank" not in query_dict: + raise _error("rank parameter missing") + if "world_size" not in query_dict: + raise _error("world size parameter missing") + + rank = int(query_dict["rank"]) + world_size = int(query_dict["world_size"]) + store = FileStore(path, world_size) + yield (store, rank, world_size) + + # If this configuration is invalidated, there is nothing we can do about it + raise RuntimeError("Unable to perform rerendezvous using file:// method") + + +def _torchelastic_use_agent_store() -> bool: + return os.environ.get("TORCHELASTIC_USE_AGENT_STORE", None) == str(True) + + +def _create_c10d_store( + hostname, port, rank, world_size, timeout, use_libuv=True +) -> Store: + """ + Smartly creates a c10d Store object on ``rank`` based on whether we need to re-use agent store. + + The TCPStore server is assumed to be hosted + on ``hostname:port``. + + By default, the TCPStore server uses the asynchronous implementation + ``LibUVStoreDaemon`` which utilizes libuv. + + If ``torchelastic_use_agent_store()`` is ``True``, then it is assumed that + the agent leader (node rank 0) hosts the TCPStore server (for which the + endpoint is specified by the given ``hostname:port``). Hence + ALL ranks will create and return a TCPStore client (e.g. ``start_daemon=False``). + + If ``torchelastic_use_agent_store()`` is ``False``, then rank 0 will host + the TCPStore (with multi-tenancy) and it is assumed that rank 0's hostname + and port are correctly passed via ``hostname`` and ``port``. All + non-zero ranks will create and return a TCPStore client. + """ + # check if port is uint16_t + if not 0 <= port < 2**16: + raise ValueError(f"port must have value from 0 to 65535 but was {port}.") + + if _torchelastic_use_agent_store(): + # We create a new TCPStore for every retry so no need to add prefix for each attempt. + return TCPStore( + host_name=hostname, + port=port, + world_size=world_size, + is_master=False, + timeout=timeout, + ) + else: + start_daemon = rank == 0 + return TCPStore( + host_name=hostname, + port=port, + world_size=world_size, + is_master=start_daemon, + timeout=timeout, + multi_tenant=True, + use_libuv=use_libuv, + ) + + +def _tcp_rendezvous_handler( + url: str, timeout: timedelta = default_pg_timeout, **kwargs +): + def _error(msg): + return _rendezvous_error("tcp:// rendezvous: " + msg) + + result = urlparse(url) + if not result.port: + raise _error("port number missing") + query_dict = _query_to_dict(result.query) + if "rank" not in query_dict: + raise _error("rank parameter missing") + if "world_size" not in query_dict: + raise _error("world size parameter missing") + + rank = int(query_dict["rank"]) + world_size = int(query_dict["world_size"]) + use_libuv = _get_use_libuv_from_query_dict(query_dict) + + assert result.hostname is not None + + store = _create_c10d_store( + result.hostname, result.port, rank, world_size, timeout, use_libuv + ) + + yield (store, rank, world_size) + + # If this configuration is invalidated, there is nothing we can do about it + raise RuntimeError("Unable to perform re-rendezvous using tcp:// method") + + +def _env_rendezvous_handler( + url: str, timeout: timedelta = default_pg_timeout, **kwargs +): + def _error(msg): + return _rendezvous_error("env:// rendezvous: " + msg) + + def _env_error(var): + return _error(f"environment variable {var} expected, but not set") + + def _get_env_or_raise(env_var: str) -> str: + env_val = os.environ.get(env_var, None) + if not env_val: + raise _env_error(env_var) + else: + return env_val + + result = urlparse(url) + query_dict = _query_to_dict(result.query) + + rank: int + world_size: int + master_port: int + master_addr: str + + if "rank" in query_dict: + rank = int(query_dict["rank"]) + else: + rank = int(_get_env_or_raise("RANK")) + + if "world_size" in query_dict: + world_size = int(query_dict["world_size"]) + else: + world_size = int(_get_env_or_raise("WORLD_SIZE")) + + master_addr = _get_env_or_raise("MASTER_ADDR") + master_port = int(_get_env_or_raise("MASTER_PORT")) + use_libuv = _get_use_libuv_from_query_dict(query_dict) + + store = _create_c10d_store( + master_addr, master_port, rank, world_size, timeout, use_libuv + ) + + yield (store, rank, world_size) + + # If this configuration is invalidated, there is nothing we can do about it + raise RuntimeError("Unable to perform re-rendezvous using env:// method") + + +register_rendezvous_handler("tcp", _tcp_rendezvous_handler) +register_rendezvous_handler("env", _env_rendezvous_handler) +register_rendezvous_handler("file", _file_rendezvous_handler) diff --git a/lib/python3.12/site-packages/torch/distributed/run.py b/lib/python3.12/site-packages/torch/distributed/run.py new file mode 100644 index 0000000000000000000000000000000000000000..088bcf367fcf847d56bfffe0ea57a1bc75810cb9 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/run.py @@ -0,0 +1,922 @@ +#!/usr/bin/env python3 +# mypy: allow-untyped-defs + +# Copyright (c) Facebook, Inc. and its affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Superset of ``torch.distributed.launch``. + +``torchrun`` provides a superset of the functionality as ``torch.distributed.launch`` +with the following additional functionalities: + +1. Worker failures are handled gracefully by restarting all workers. + +2. Worker ``RANK`` and ``WORLD_SIZE`` are assigned automatically. + +3. Number of nodes is allowed to change between minimum and maximum sizes (elasticity). + +.. note:: ``torchrun`` is a python + `console script `_ + to the main module + `torch.distributed.run `_ + declared in the ``entry_points`` configuration in + `setup.py `_. + It is equivalent to invoking ``python -m torch.distributed.run``. + + +Transitioning from torch.distributed.launch to torchrun +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +``torchrun`` supports the same arguments as ``torch.distributed.launch`` **except** +for ``--use-env`` which is now deprecated. To migrate from ``torch.distributed.launch`` +to ``torchrun`` follow these steps: + +1. If your training script is already reading ``local_rank`` from the ``LOCAL_RANK`` environment variable. + Then you need simply omit the ``--use-env`` flag, e.g.: + + +--------------------------------------------------------------------+--------------------------------------------+ + | ``torch.distributed.launch`` | ``torchrun`` | + +====================================================================+============================================+ + | | | + | .. code-block:: shell-session | .. code-block:: shell-session | + | | | + | $ python -m torch.distributed.launch --use-env train_script.py | $ torchrun train_script.py | + | | | + +--------------------------------------------------------------------+--------------------------------------------+ + +2. If your training script reads local rank from a ``--local-rank`` cmd argument. + Change your training script to read from the ``LOCAL_RANK`` environment variable as + demonstrated by the following code snippet: + + +-------------------------------------------------------+----------------------------------------------------+ + | ``torch.distributed.launch`` | ``torchrun`` | + +=======================================================+====================================================+ + | | | + | .. code-block:: python | .. code-block:: python | + | | | + | | | + | import argparse | import os | + | parser = argparse.ArgumentParser() | local_rank = int(os.environ["LOCAL_RANK"]) | + | parser.add_argument("--local-rank", type=int) | | + | args = parser.parse_args() | | + | | | + | local_rank = args.local_rank | | + | | | + +-------------------------------------------------------+----------------------------------------------------+ + +.. versionchanged:: 2.0.0 + + The launcher will pass the ``--local-rank=`` argument to your script. + From PyTorch 2.0.0 onwards, the dashed ``--local-rank`` is preferred over the + previously used underscored ``--local_rank``. + + For backward compatibility, it may be necessary for users to handle both + cases in their argument parsing code. This means including both ``"--local-rank"`` + and ``"--local_rank"`` in the argument parser. If only ``"--local_rank"`` is + provided, the launcher will trigger an error: "error: unrecognized arguments: + --local-rank=". For training code that only supports PyTorch 2.0.0+, + including ``"--local-rank"`` should be sufficient. + + :: + + >>> # xdoctest: +SKIP + >>> import argparse + >>> parser = argparse.ArgumentParser() + >>> parser.add_argument("--local-rank", "--local_rank", type=int) + >>> args = parser.parse_args() + +The aformentioned changes suffice to migrate from ``torch.distributed.launch`` to ``torchrun``. +To take advantage of new features such as elasticity, fault-tolerance, and error reporting of ``torchrun`` +please refer to: + +* :ref:`elastic_train_script` for more information on authoring training scripts that are ``torchrun`` compliant. +* the rest of this page for more information on the features of ``torchrun``. + + +Usage +-------- + +Single-node multi-worker +++++++++++++++++++++++++++++++ + +:: + + torchrun + --standalone + --nnodes=1 + --nproc-per-node=$NUM_TRAINERS + YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...) + +Stacked single-node multi-worker ++++++++++++++++++++++++++++++++++++ + +To run multiple instances (separate jobs) of single-node, multi-worker on the +same host, we need to make sure that each instance (job) is +setup on different ports to avoid port conflicts (or worse, two jobs being merged +as a single job). To do this you have to run with ``--rdzv-backend=c10d`` +and specify a different port by setting ``--rdzv-endpoint=localhost:$PORT_k``. +For ``--nodes=1``, its often convenient to let ``torchrun`` pick a free random +port automatically instead of manually assigning different ports for each run. + +:: + + torchrun + --rdzv-backend=c10d + --rdzv-endpoint=localhost:0 + --nnodes=1 + --nproc-per-node=$NUM_TRAINERS + YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...) + + +Fault tolerant (fixed sized number of workers, no elasticity, tolerates 3 failures) +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + +:: + + torchrun + --nnodes=$NUM_NODES + --nproc-per-node=$NUM_TRAINERS + --max-restarts=3 + --rdzv-id=$JOB_ID + --rdzv-backend=c10d + --rdzv-endpoint=$HOST_NODE_ADDR + YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...) + +``HOST_NODE_ADDR``, in form [:] (e.g. node1.example.com:29400), specifies the node and +the port on which the C10d rendezvous backend should be instantiated and hosted. It can be any +node in your training cluster, but ideally you should pick a node that has a high bandwidth. + +.. note:: + If no port number is specified ``HOST_NODE_ADDR`` defaults to 29400. + +Elastic (``min=1``, ``max=4``, tolerates up to 3 membership changes or failures) ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + +:: + + torchrun + --nnodes=1:4 + --nproc-per-node=$NUM_TRAINERS + --max-restarts=3 + --rdzv-id=$JOB_ID + --rdzv-backend=c10d + --rdzv-endpoint=$HOST_NODE_ADDR + YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...) + +``HOST_NODE_ADDR``, in form [:] (e.g. node1.example.com:29400), specifies the node and +the port on which the C10d rendezvous backend should be instantiated and hosted. It can be any +node in your training cluster, but ideally you should pick a node that has a high bandwidth. + +.. note:: + If no port number is specified ``HOST_NODE_ADDR`` defaults to 29400. + +Note on rendezvous backend +------------------------------ + +For multi-node training you need to specify: + +1. ``--rdzv-id``: A unique job id (shared by all nodes participating in the job) +2. ``--rdzv-backend``: An implementation of + :py:class:`torch.distributed.elastic.rendezvous.RendezvousHandler` +3. ``--rdzv-endpoint``: The endpoint where the rendezvous backend is running; usually in form + ``host:port``. + +Currently ``c10d`` (recommended), ``etcd-v2``, and ``etcd`` (legacy) rendezvous backends are +supported out of the box. To use ``etcd-v2`` or ``etcd``, setup an etcd server with the ``v2`` api +enabled (e.g. ``--enable-v2``). + +.. warning:: + ``etcd-v2`` and ``etcd`` rendezvous use etcd API v2. You MUST enable the v2 API on the etcd + server. Our tests use etcd v3.4.3. + +.. warning:: + For etcd-based rendezvous we recommend using ``etcd-v2`` over ``etcd`` which is functionally + equivalent, but uses a revised implementation. ``etcd`` is in maintenance mode and will be + removed in a future version. + +Definitions +-------------- + +1. ``Node`` - A physical instance or a container; maps to the unit that the job manager works with. + +2. ``Worker`` - A worker in the context of distributed training. + +3. ``WorkerGroup`` - The set of workers that execute the same function (e.g. trainers). + +4. ``LocalWorkerGroup`` - A subset of the workers in the worker group running on the same node. + +5. ``RANK`` - The rank of the worker within a worker group. + +6. ``WORLD_SIZE`` - The total number of workers in a worker group. + +7. ``LOCAL_RANK`` - The rank of the worker within a local worker group. + +8. ``LOCAL_WORLD_SIZE`` - The size of the local worker group. + +9. ``rdzv_id`` - A user-defined id that uniquely identifies the worker group for a job. This id is + used by each node to join as a member of a particular worker group. + +9. ``rdzv_backend`` - The backend of the rendezvous (e.g. ``c10d``). This is typically a strongly + consistent key-value store. + +10. ``rdzv_endpoint`` - The rendezvous backend endpoint; usually in form ``:``. + +A ``Node`` runs ``LOCAL_WORLD_SIZE`` workers which comprise a ``LocalWorkerGroup``. The union of +all ``LocalWorkerGroups`` in the nodes in the job comprise the ``WorkerGroup``. + +Environment Variables +---------------------- + +The following environment variables are made available to you in your script: + +1. ``LOCAL_RANK`` - The local rank. + +2. ``RANK`` - The global rank. + +3. ``GROUP_RANK`` - The rank of the worker group. A number between 0 and ``max_nnodes``. When + running a single worker group per node, this is the rank of the node. + +4. ``ROLE_RANK`` - The rank of the worker across all the workers that have the same role. The role + of the worker is specified in the ``WorkerSpec``. + +5. ``LOCAL_WORLD_SIZE`` - The local world size (e.g. number of workers running locally); equals to + ``--nproc-per-node`` specified on ``torchrun``. + +6. ``WORLD_SIZE`` - The world size (total number of workers in the job). + +7. ``ROLE_WORLD_SIZE`` - The total number of workers that was launched with the same role specified + in ``WorkerSpec``. + +8. ``MASTER_ADDR`` - The FQDN of the host that is running worker with rank 0; used to initialize + the Torch Distributed backend. + +9. ``MASTER_PORT`` - The port on the ``MASTER_ADDR`` that can be used to host the C10d TCP store. + +10. ``TORCHELASTIC_RESTART_COUNT`` - The number of worker group restarts so far. + +11. ``TORCHELASTIC_MAX_RESTARTS`` - The configured maximum number of restarts. + +12. ``TORCHELASTIC_RUN_ID`` - Equal to the rendezvous ``run_id`` (e.g. unique job id). + +13. ``PYTHON_EXEC`` - System executable override. If provided, the python user script will + use the value of ``PYTHON_EXEC`` as executable. The `sys.executable` is used by default. + +Deployment +------------ + +1. (Not needed for the C10d backend) Start the rendezvous backend server and get the endpoint (to be + passed as ``--rdzv-endpoint`` to the launcher script) + +2. Single-node multi-worker: Start the launcher on the host to start the agent process which + creates and monitors a local worker group. + +3. Multi-node multi-worker: Start the launcher with the same arguments on all the nodes + participating in training. + +When using a job/cluster manager the entry point command to the multi-node job should be this +launcher. + +Failure Modes +--------------- + +1. Worker failure: For a training job with ``n`` workers, if ``k<=n`` workers fail all workers + are stopped and restarted up to ``max_restarts``. + +2. Agent failure: An agent failure results in a local worker group failure. It is up to the job + manager to fail the entire job (gang semantics) or attempt to replace the node. Both behaviors + are supported by the agent. + +3. Node failure: Same as agent failure. + +Membership Changes +-------------------- + +1. Node departure (scale-down): The agent is notified of the departure, all existing workers are + stopped, a new ``WorkerGroup`` is formed, and all workers are started with a new ``RANK`` and + ``WORLD_SIZE``. + +2. Node arrival (scale-up): The new node is admitted to the job, all existing workers are stopped, + a new ``WorkerGroup`` is formed, and all workers are started with a new ``RANK`` and + ``WORLD_SIZE``. + +Important Notices +-------------------- + +1. This utility and multi-process distributed (single-node or + multi-node) GPU training currently only achieves the best performance using + the NCCL distributed backend. Thus NCCL backend is the recommended backend to + use for GPU training. + +2. The environment variables necessary to initialize a Torch process group are provided to you by + this module, no need for you to pass ``RANK`` manually. To initialize a process group in your + training script, simply run: + +:: + + >>> # xdoctest: +SKIP("stub") + >>> import torch.distributed as dist + >>> dist.init_process_group(backend="gloo|nccl") + +3. In your training program, you can either use regular distributed functions + or use :func:`torch.nn.parallel.DistributedDataParallel` module. If your + training program uses GPUs for training and you would like to use + :func:`torch.nn.parallel.DistributedDataParallel` module, + here is how to configure it. + +:: + + local_rank = int(os.environ["LOCAL_RANK"]) + model = torch.nn.parallel.DistributedDataParallel(model, + device_ids=[local_rank], + output_device=local_rank) + +Please ensure that ``device_ids`` argument is set to be the only GPU device id +that your code will be operating on. This is generally the local rank of the +process. In other words, the ``device_ids`` needs to be ``[int(os.environ("LOCAL_RANK"))]``, +and ``output_device`` needs to be ``int(os.environ("LOCAL_RANK"))`` in order to use this +utility + + +4. On failures or membership changes ALL surviving workers are killed immediately. Make sure to + checkpoint your progress. The frequency of checkpoints should depend on your job's tolerance + for lost work. + +5. This module only supports homogeneous ``LOCAL_WORLD_SIZE``. That is, it is assumed that all + nodes run the same number of local workers (per role). + +6. ``RANK`` is NOT stable. Between restarts, the local workers on a node can be assigned a + different range of ranks than before. NEVER hard code any assumptions about the stable-ness of + ranks or some correlation between ``RANK`` and ``LOCAL_RANK``. + +7. When using elasticity (``min_size!=max_size``) DO NOT hard code assumptions about + ``WORLD_SIZE`` as the world size can change as nodes are allowed to leave and join. + +8. It is recommended for your script to have the following structure: + +:: + + def main(): + load_checkpoint(checkpoint_path) + initialize() + train() + + def train(): + for batch in iter(dataset): + train_step(batch) + + if should_checkpoint: + save_checkpoint(checkpoint_path) + +9. (Recommended) On worker errors, this tool will summarize the details of the error + (e.g. time, rank, host, pid, traceback, etc). On each node, the first error (by timestamp) + is heuristically reported as the "Root Cause" error. To get tracebacks as part of this + error summary print out, you must decorate your main entrypoint function in your + training script as shown in the example below. If not decorated, then the summary + will not include the traceback of the exception and will only contain the exitcode. + For details on torchelastic error handling see: https://pytorch.org/docs/stable/elastic/errors.html + +:: + + from torch.distributed.elastic.multiprocessing.errors import record + + @record + def main(): + # do train + pass + + if __name__ == "__main__": + main() + +""" +import os +import sys +import uuid +from argparse import ArgumentParser, REMAINDER +from importlib import metadata +from typing import Callable, List, Optional, Set, Tuple, Type, Union + +import torch +from torch.distributed.argparse_util import check_env, env +from torch.distributed.elastic.multiprocessing import DefaultLogsSpecs, LogsSpecs, Std +from torch.distributed.elastic.multiprocessing.errors import record +from torch.distributed.elastic.rendezvous.utils import _parse_rendezvous_config +from torch.distributed.elastic.utils import macros +from torch.distributed.elastic.utils.logging import get_logger +from torch.distributed.launcher.api import elastic_launch, LaunchConfig +from torch.utils.backend_registration import _get_custom_mod_func + + +logger = get_logger(__name__) + + +def get_args_parser() -> ArgumentParser: + """Parse the command line options.""" + parser = ArgumentParser(description="Torch Distributed Elastic Training Launcher") + + # + # Worker/node size related arguments. + # + + parser.add_argument( + "--nnodes", + action=env, + type=str, + default="1:1", + help="Number of nodes, or the range of nodes in form :.", + ) + parser.add_argument( + "--nproc-per-node", + "--nproc_per_node", + action=env, + type=str, + default="1", + help="Number of workers per node; supported values: [auto, cpu, gpu, int].", + ) + + # + # Rendezvous related arguments + # + + parser.add_argument( + "--rdzv-backend", + "--rdzv_backend", + action=env, + type=str, + default="static", + help="Rendezvous backend.", + ) + parser.add_argument( + "--rdzv-endpoint", + "--rdzv_endpoint", + action=env, + type=str, + default="", + help="Rendezvous backend endpoint; usually in form :.", + ) + parser.add_argument( + "--rdzv-id", + "--rdzv_id", + action=env, + type=str, + default="none", + help="User-defined group id.", + ) + parser.add_argument( + "--rdzv-conf", + "--rdzv_conf", + action=env, + type=str, + default="", + help="Additional rendezvous configuration (=,=,...).", + ) + parser.add_argument( + "--standalone", + action=check_env, + help="Start a local standalone rendezvous backend that is represented by a C10d TCP store " + "on a free port. Useful when launching single-node, multi-worker job. If specified " + "--rdzv-backend, --rdzv-endpoint, --rdzv-id are auto-assigned and any explicitly set values " + "are ignored.", + ) + + # + # User-code launch related arguments. + # + + parser.add_argument( + "--max-restarts", + "--max_restarts", + action=env, + type=int, + default=0, + help="Maximum number of worker group restarts before failing.", + ) + parser.add_argument( + "--monitor-interval", + "--monitor_interval", + action=env, + type=float, + default=0.1, + help="Interval, in seconds, to monitor the state of workers.", + ) + parser.add_argument( + "--start-method", + "--start_method", + action=env, + type=str, + default="spawn", + choices=["spawn", "fork", "forkserver"], + help="Multiprocessing start method to use when creating workers.", + ) + parser.add_argument( + "--role", + action=env, + type=str, + default="default", + help="User-defined role for the workers.", + ) + parser.add_argument( + "-m", + "--module", + action=check_env, + help="Change each process to interpret the launch script as a Python module, executing " + "with the same behavior as 'python -m'.", + ) + parser.add_argument( + "--no-python", + "--no_python", + action=check_env, + help="Skip prepending the training script with 'python' - just execute it directly. Useful " + "when the script is not a Python script.", + ) + + parser.add_argument( + "--run-path", + "--run_path", + action=check_env, + help="Run the training script with runpy.run_path in the same interpreter." + " Script must be provided as an abs path (e.g. /abs/path/script.py)." + " Takes precedence over --no-python.", + ) + parser.add_argument( + "--log-dir", + "--log_dir", + action=env, + type=str, + default=None, + help="Base directory to use for log files (e.g. /var/log/torch/elastic). The same " + "directory is re-used for multiple runs (a unique job-level sub-directory is created with " + "rdzv_id as the prefix).", + ) + parser.add_argument( + "-r", + "--redirects", + action=env, + type=str, + default="0", + help="Redirect std streams into a log file in the log directory (e.g. [-r 3] redirects " + "both stdout+stderr for all workers, [-r 0:1,1:2] redirects stdout for local rank 0 and " + "stderr for local rank 1).", + ) + parser.add_argument( + "-t", + "--tee", + action=env, + type=str, + default="0", + help="Tee std streams into a log file and also to console (see --redirects for format).", + ) + + parser.add_argument( + "--local-ranks-filter", + "--local_ranks_filter", + action=env, + type=str, + default="", + help="Only show logs from specified ranks in console (e.g. [--local_ranks_filter=0,1,2] will " + "only show logs from rank 0, 1 and 2). This will only apply to stdout and stderr, not to" + "log files saved via --redirect or --tee", + ) + + # + # Backwards compatible parameters with caffe2.distributed.launch. + # + + parser.add_argument( + "--node-rank", + "--node_rank", + type=int, + action=env, + default=0, + help="Rank of the node for multi-node distributed training.", + ) + parser.add_argument( + "--master-addr", + "--master_addr", + default="127.0.0.1", + type=str, + action=env, + help="Address of the master node (rank 0) that only used for static rendezvous. It should " + "be either the IP address or the hostname of rank 0. For single node multi-proc training " + "the --master-addr can simply be 127.0.0.1; IPv6 should have the pattern " + "`[0:0:0:0:0:0:0:1]`.", + ) + parser.add_argument( + "--master-port", + "--master_port", + default=29500, + type=int, + action=env, + help="Port on the master node (rank 0) to be used for communication during distributed " + "training. It is only used for static rendezvous.", + ) + parser.add_argument( + "--local-addr", + "--local_addr", + default=None, + type=str, + action=env, + help="Address of the local node. If specified, will use the given address for connection. " + "Else, will look up the local node address instead. Else, it will be default to local " + "machine's FQDN.", + ) + + parser.add_argument( + "--logs-specs", + "--logs_specs", + default=None, + type=str, + help="torchrun.logs_specs group entrypoint name, value must be type of LogsSpecs. " + "Can be used to override custom logging behavior.", + ) + + # + # Positional arguments. + # + + parser.add_argument( + "training_script", + type=str, + help="Full path to the (single GPU) training program/script to be launched in parallel, " + "followed by all the arguments for the training script.", + ) + + # Rest from the training program. + parser.add_argument("training_script_args", nargs=REMAINDER) + + return parser + + +def parse_args(args): + parser = get_args_parser() + return parser.parse_args(args) + + +def parse_min_max_nnodes(nnodes: str): + arr = nnodes.split(":") + + if len(arr) == 1: + min_nodes = max_nodes = int(arr[0]) + elif len(arr) == 2: + min_nodes = int(arr[0]) + max_nodes = int(arr[1]) + else: + raise RuntimeError(f'nnodes={nnodes} is not in "MIN:MAX" format') # noqa: E231 + + return min_nodes, max_nodes + + +def determine_local_world_size(nproc_per_node: str): + try: + logger.info("Using nproc_per_node=%s.", nproc_per_node) + return int(nproc_per_node) + except ValueError as e: + if nproc_per_node == "cpu": + num_proc = os.cpu_count() + device_type = "cpu" + elif nproc_per_node == "gpu": + if not torch.cuda.is_available(): + raise ValueError("Cuda is not available.") from e + device_type = "gpu" + num_proc = torch.cuda.device_count() + elif nproc_per_node == torch._C._get_privateuse1_backend_name(): + if not _get_custom_mod_func("is_available")(): + raise ValueError(f"{nproc_per_node} is not available.") from e + device_type = nproc_per_node + num_proc = _get_custom_mod_func("device_count")() + elif nproc_per_node == "auto": + if torch.cuda.is_available(): + num_proc = torch.cuda.device_count() + device_type = "gpu" + elif ( + hasattr(torch, torch._C._get_privateuse1_backend_name()) + and _get_custom_mod_func("is_available")() + ): + num_proc = _get_custom_mod_func("device_count")() + device_type = torch._C._get_privateuse1_backend_name() + else: + num_proc = os.cpu_count() + device_type = "cpu" + else: + raise ValueError( + f"Unsupported nproc_per_node value: {nproc_per_node}" + ) from e + + logger.info( + "Using nproc_per_node=%s, setting nproc_per_node to %s since the instance has %s %s", + nproc_per_node, + num_proc, + num_proc, + device_type, + ) + return num_proc + + +def get_rdzv_endpoint(args): + if args.rdzv_backend == "static" and not args.rdzv_endpoint: + return f"{args.master_addr}:{args.master_port}" # noqa: E231 + return args.rdzv_endpoint + + +def get_use_env(args) -> bool: + """ + Retrieve ``use_env`` from the args. + + ``use_env`` is a legacy argument, if ``use_env`` is False, the + ``--node-rank`` argument will be transferred to all worker processes. + ``use_env`` is only used by the ``torch.distributed.launch`` and will + be deprecated in future releases. + """ + if not hasattr(args, "use_env"): + return True + return args.use_env + + +def _get_logs_specs_class(logs_specs_name: Optional[str]) -> Type[LogsSpecs]: + """ + Attemps to load `torchrun.logs_spec` entrypoint with key of `logs_specs_name` param. + Provides plugin mechanism to provide custom implementation of LogsSpecs. + + Returns `DefaultLogsSpecs` when logs_spec_name is None. + Raises ValueError when entrypoint for `logs_spec_name` can't be found in entrypoints. + """ + logs_specs_cls = None + if logs_specs_name is not None: + eps = metadata.entry_points() + if hasattr(eps, "select"): # >= 3.10 + group = eps.select(group="torchrun.logs_specs") + if group.select(name=logs_specs_name): + logs_specs_cls = group[logs_specs_name].load() + + elif specs := eps.get("torchrun.logs_specs"): # < 3.10 + if entrypoint_list := [ep for ep in specs if ep.name == logs_specs_name]: + logs_specs_cls = entrypoint_list[0].load() + + if logs_specs_cls is None: + raise ValueError( + f"Could not find entrypoint under 'torchrun.logs_specs[{logs_specs_name}]' key" + ) + + logger.info( + "Using logs_spec '%s' mapped to %s", logs_specs_name, str(logs_specs_cls) + ) + else: + logs_specs_cls = DefaultLogsSpecs + + return logs_specs_cls + + +def config_from_args(args) -> Tuple[LaunchConfig, Union[Callable, str], List[str]]: + # If ``args`` not passed, defaults to ``sys.argv[:1]`` + min_nodes, max_nodes = parse_min_max_nnodes(args.nnodes) + assert 0 < min_nodes <= max_nodes + assert args.max_restarts >= 0 + + if ( + hasattr(args, "master_addr") + and args.rdzv_backend != "static" + and not args.rdzv_endpoint + ): + logger.warning( + "master_addr is only used for static rdzv_backend and when rdzv_endpoint " + "is not specified." + ) + + nproc_per_node = determine_local_world_size(args.nproc_per_node) + if "OMP_NUM_THREADS" not in os.environ and nproc_per_node > 1: + omp_num_threads = 1 + logger.warning( + "\n*****************************************\n" + "Setting OMP_NUM_THREADS environment variable for each process to be " + "%s in default, to avoid your system being overloaded, " + "please further tune the variable for optimal performance in " + "your application as needed. \n" + "*****************************************", + omp_num_threads, + ) + # This env variable will be passed down to the subprocesses + os.environ["OMP_NUM_THREADS"] = str(omp_num_threads) + + log_line_prefix_template = os.getenv("TORCHELASTIC_LOG_LINE_PREFIX_TEMPLATE") + + rdzv_configs = _parse_rendezvous_config(args.rdzv_conf) + + if args.rdzv_backend == "static": + rdzv_configs["rank"] = args.node_rank + + rdzv_endpoint = get_rdzv_endpoint(args) + + ranks: Optional[Set[int]] = None + if args.local_ranks_filter: + try: + ranks = set(map(int, args.local_ranks_filter.split(","))) + assert ranks + except Exception as e: + raise ValueError( + "--local_ranks_filter must be a comma-separated list of integers e.g. --local_ranks_filter=0,1,2" + ) from e + + logs_specs_cls: Type[LogsSpecs] = _get_logs_specs_class(args.logs_specs) + logs_specs = logs_specs_cls( + log_dir=args.log_dir, + redirects=Std.from_str(args.redirects), + tee=Std.from_str(args.tee), + local_ranks_filter=ranks, + ) + + config = LaunchConfig( + min_nodes=min_nodes, + max_nodes=max_nodes, + nproc_per_node=nproc_per_node, + run_id=args.rdzv_id, + role=args.role, + rdzv_endpoint=rdzv_endpoint, + rdzv_backend=args.rdzv_backend, + rdzv_configs=rdzv_configs, + max_restarts=args.max_restarts, + monitor_interval=args.monitor_interval, + start_method=args.start_method, + log_line_prefix_template=log_line_prefix_template, + local_addr=args.local_addr, + logs_specs=logs_specs, + ) + + with_python = not args.no_python + cmd: Union[Callable, str] + cmd_args = [] + use_env = get_use_env(args) + if args.run_path: + cmd = run_script_path + cmd_args.append(args.training_script) + else: + if with_python: + cmd = os.getenv("PYTHON_EXEC", sys.executable) + cmd_args.append("-u") + if args.module: + cmd_args.append("-m") + cmd_args.append(args.training_script) + else: + if args.module: + raise ValueError( + "Don't use both the '--no-python' flag" + " and the '--module' flag at the same time." + ) + cmd = args.training_script + if not use_env: + cmd_args.append(f"--local-rank={macros.local_rank}") + cmd_args.extend(args.training_script_args) + + return config, cmd, cmd_args + + +def run_script_path(training_script: str, *training_script_args: str): + """ + Run the provided `training_script` from within this interpreter. + + Usage: `script_as_function("/abs/path/to/script.py", "--arg1", "val1")` + """ + import runpy + import sys + + sys.argv = [training_script] + [*training_script_args] + runpy.run_path(sys.argv[0], run_name="__main__") + + +def run(args): + torch.multiprocessing._set_thread_name("pt_elastic") + + if args.standalone: + args.rdzv_backend = "c10d" + args.rdzv_endpoint = "localhost:0" + args.rdzv_id = str(uuid.uuid4()) + logger.info( + "\n**************************************\n" + "Rendezvous info:\n" + "--rdzv-backend=%s " + "--rdzv-endpoint=%s " + "--rdzv-id=%s\n" + "**************************************\n", + args.rdzv_backend, + args.rdzv_endpoint, + args.rdzv_id, + ) + + config, cmd, cmd_args = config_from_args(args) + elastic_launch( + config=config, + entrypoint=cmd, + )(*cmd_args) + + +@record +def main(args=None): + args = parse_args(args) + run(args) + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/torch/distributed/utils.py b/lib/python3.12/site-packages/torch/distributed/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..af2372af19f345feddc768b84b817fd225290c98 --- /dev/null +++ b/lib/python3.12/site-packages/torch/distributed/utils.py @@ -0,0 +1,389 @@ +# mypy: allow-untyped-defs +import dataclasses +import traceback +from typing import ( + Any, + Callable, + Container, + Dict, + List, + Optional, + OrderedDict, + overload, + Set, + Tuple, + TypeVar, +) + +import torch +import torch.distributed as dist +from torch import nn +from torch.nn.utils.rnn import PackedSequence + + +__all__ = [] # type: ignore[var-annotated] + + +def _pack_kwargs(*args: Any, **kwargs: Any) -> Tuple[Tuple[Any, ...], Tuple[str, ...]]: + """ + Turn argument list into separate key list and value list (unpack_kwargs does the opposite). + + Inspiration: https://github.com/facebookresearch/fairscale/blob/eeb6684/fairscale/internal/containers.py#L70 + Usage:: + + kwarg_keys, flat_args = pack_kwargs(1, 2, a=3, b=4) + assert kwarg_keys == ("a", "b") + assert flat_args == (1, 2, 3, 4) + args, kwargs = unpack_kwargs(kwarg_keys, flat_args) + assert args == (1, 2) + assert kwargs == {"a": 3, "b": 4} + Returns: + Tuple[Tuple[Any, ...], Tuple[str, ...]]: The first tuple element gives + gives both positional args and kwarg values, where the positional args + proceed kwarg values and kwarg values are ordered consistently with the + kwarg keys. The second tuple element gives the kwarg keys. + The second tuple element's length is at most the first tuple element's length. + """ + kwarg_keys: List[str] = [] + flat_args: List[Any] = list(args) + for k, v in kwargs.items(): + kwarg_keys.append(k) + flat_args.append(v) + + return tuple(flat_args), tuple(kwarg_keys) + + +def _cast_forward_inputs( + dtype: Optional[torch.dtype], + *args: Any, + **kwargs: Any, +) -> Tuple[Any, Any]: + """ + Cast floating point tensors in ``args`` and ``kwargs`` to ``input_dtype``. + + This respects the existing ``requires_grad`` on the tensors. + """ + if dtype is None: + return args, kwargs + + def cast_fn(x: torch.Tensor) -> torch.Tensor: + if not torch.is_floating_point(x) or x.dtype == dtype: + return x + return x.to(dtype) + + return (_apply_to_tensors(cast_fn, args), _apply_to_tensors(cast_fn, kwargs)) + + +def _unpack_kwargs( + flat_args: Tuple[Any, ...], kwarg_keys: Tuple[str, ...] +) -> Tuple[Tuple[Any, ...], Dict[str, Any]]: + """See _pack_kwargs.""" + assert len(kwarg_keys) <= len( + flat_args + ), f"too many keys {len(kwarg_keys)} vs. {len(flat_args)}" + if len(kwarg_keys) == 0: + return flat_args, {} + args = flat_args[: -len(kwarg_keys)] + kwargs = dict(zip(kwarg_keys, flat_args[-len(kwarg_keys) :])) + return args, kwargs + + +S = TypeVar("S", dict, list, tuple) +T = TypeVar("T", torch.Tensor, PackedSequence) + + +@overload +def _recursive_to( + inputs: S, target_device: torch.device, use_side_stream_for_tensor_copies: bool +) -> List[S]: + ... + + +@overload +def _recursive_to( + inputs: T, target_device: torch.device, use_side_stream_for_tensor_copies: bool +) -> Tuple[T]: + ... + + +def _recursive_to(inputs, target_device, use_side_stream_for_tensor_copies): + r"""Recursively moves input to the target_device.""" + + def to_map(obj): + if isinstance(obj, (torch.Tensor, PackedSequence)): + device = obj.data.device if isinstance(obj, PackedSequence) else obj.device + if device == target_device: + return (obj,) + if not use_side_stream_for_tensor_copies: + return (obj.to(target_device),) + else: + # If the custom module is not registered to torch, stream is not used for acceleration + device_mod = getattr(torch, device.type, None) + if device.type == "cpu" or device_mod is None: + return (obj.to(target_device),) + + from torch.nn.parallel._functions import _get_stream + + # Perform CPU -> target_device copies in a background stream. This code is + # motivated from similar logic in torch/nn/parallel/_functions.py + stream = _get_stream(target_device) + with device_mod.stream(stream): + output = obj.to(target_device) + # synchronize with the copy stream + with device_mod.device(target_device.index): + current_stream = device_mod.current_stream() + # Sync the current stream with the copy stream + current_stream.wait_stream(stream) + # Ensure tensor memory is not reused until work on + # main stream is complete + if isinstance(obj, PackedSequence): + output.data.record_stream(current_stream) # type: ignore[arg-type] + else: + assert isinstance(output, torch.Tensor) + output.record_stream(current_stream) # type: ignore[arg-type] + return (output,) + + from torch.nn.parallel.scatter_gather import _is_namedtuple + + if _is_namedtuple(obj): + return [type(obj)(*args) for args in zip(*map(to_map, obj))] + if isinstance(obj, tuple) and len(obj) > 0: + return list(zip(*map(to_map, obj))) + if isinstance(obj, list) and len(obj) > 0: + return [list(i) for i in zip(*map(to_map, obj))] + if isinstance(obj, dict) and len(obj) > 0: + return [type(obj)(i) for i in zip(*map(to_map, obj.items()))] + return [obj] + + # Avoid reference cycle + try: + res = to_map(inputs) + finally: + to_map = None # type: ignore[assignment] + return res + + +def _p_assert(cond: Any, s: str, raise_assertion_error: bool = True) -> None: + """Alternate to ``assert`` when in the backward context to print the error message ``s`` since otherwise, it is swallowed.""" + if not cond: + print(s) + traceback.print_stack() + if raise_assertion_error: + raise AssertionError(s) + + +def _alloc_storage(tensor: torch.Tensor, size: torch.Size) -> None: + """ + Allocate storage for ``tensor`` with the given size. + + Returns: + bool: ``True`` if this method allocated storage and ``False`` if the + storage was already allocated. + """ + with torch.no_grad(): + if not torch.distributed._functional_collectives.is_torchdynamo_compiling(): + already_allocated = tensor._typed_storage()._size() == size.numel() + if not already_allocated: + tensor_storage_size = tensor._typed_storage()._size() + _p_assert( + tensor_storage_size == 0, + "Tensor storage should have been resized to be 0 but got PLACEHOLDEr", + ) + tensor._typed_storage()._resize_(size.numel()) + + +def _free_storage(tensor: torch.Tensor): + """ + Frees the underlying storage of ``tensor``. + + Returns: + bool: ``True`` if the method freed the storage and ``False`` if the + storage was already freed. + """ + with torch.no_grad(): + if not torch.distributed._functional_collectives.is_torchdynamo_compiling(): + already_freed = tensor._typed_storage()._size() == 0 + if not already_freed: + _p_assert( + tensor.storage_offset() == 0, + "Freeing a tensor's storage is unsafe when it is not the sole occupant\n" + f"storage offset: {tensor.storage_offset()}\n" + f"storage size: {tensor._typed_storage()._size()}\n" + f"tensor shape: {tensor.shape}", + ) + tensor._typed_storage()._resize_(0) + + +Q = TypeVar("Q") +R = TypeVar("R", dict, list, tuple, set, OrderedDict, PackedSequence, Any) + + +@overload +def _apply_to_tensors(fn: Callable[[torch.Tensor], Q], container: torch.Tensor) -> Q: + ... + + +@overload +def _apply_to_tensors(fn: Callable[[torch.Tensor], Any], container: R) -> R: + ... + + +def _apply_to_tensors(fn, container): + """Recursively apply to all tensor in different kinds of container types.""" + + def apply(x): + from torch.nn.parallel.scatter_gather import _is_namedtuple + + if isinstance(x, torch.Tensor): + return fn(x) + elif hasattr(x, "__dataclass_fields__"): + dc = dataclasses.replace(x) + changes = { + f.name: apply(getattr(dc, f.name)) for f in dataclasses.fields(dc) + } + return dataclasses.replace(dc, **changes) + elif isinstance(x, OrderedDict): + od = x.__class__() + for key, value in x.items(): + od[key] = apply(value) + return od + elif isinstance(x, PackedSequence): + apply(x.data) + return x + elif isinstance(x, dict): + return {key: apply(value) for key, value in x.items()} + elif _is_namedtuple(x): + res = (apply(el) for el in x) + return type(x)(*res) + elif isinstance(x, (list, tuple, set)): + return type(x)(apply(el) for el in x) + else: + return x + + return apply(container) + + +def _to_kwargs( + inputs: Tuple[Any, ...], + kwargs: Optional[Dict[str, Any]], + target_device: torch.device, + use_side_stream_for_tensor_copies: bool, +) -> Tuple[Tuple[Any, ...], Tuple[Dict[str, Any], ...]]: + moved_inputs = ( + _recursive_to(inputs, target_device, use_side_stream_for_tensor_copies) + if inputs + else [] + ) + moved_kwargs = ( + _recursive_to(kwargs, target_device, use_side_stream_for_tensor_copies) + if kwargs + else [] + ) + if len(moved_inputs) < len(moved_kwargs): + moved_inputs.extend([() for _ in range(len(moved_kwargs) - len(inputs))]) + elif len(moved_kwargs) < len(moved_inputs): + moved_kwargs.extend([{} for _ in range(len(moved_inputs) - len(moved_kwargs))]) + return tuple(moved_inputs), tuple(moved_kwargs) + + +def _verify_param_shape_across_processes( + process_group: dist.ProcessGroup, + tensors: List[torch.Tensor], + logger: Optional["dist.Logger"] = None, +): + return dist._verify_params_across_processes(process_group, tensors, logger) + + +def _sync_module_states( + module: nn.Module, + process_group: dist.ProcessGroup, + broadcast_bucket_size: int, + src: int, + params_and_buffers_to_ignore: Container[str], + broadcast_buffers: bool = True, +) -> None: + """ + Sync ``module``'s parameters and buffers state. + + Syncs ``module``'s parameters and buffers state so that all ranks contain + the same module state across all ranks. Note that this API assumes that all + parameter shapes are consistent before running the synchronization. This can + be checked with ``_verify_param_shape_across_processes``. + """ + module_states: List[torch.Tensor] = [] + for name, param in module.named_parameters(): + if name not in params_and_buffers_to_ignore: + module_states.append(param.detach()) + + if broadcast_buffers: + for name, buffer in module.named_buffers(): + if name not in params_and_buffers_to_ignore: + module_states.append(buffer.detach()) + + _sync_params_and_buffers(process_group, module_states, broadcast_bucket_size, src) + + +def _sync_params_and_buffers( + process_group: dist.ProcessGroup, + module_states: List[torch.Tensor], + broadcast_bucket_size: int, + src: int, +) -> None: + """Synchronize ``module_states`` (list of tensors) across all processes by broadcasting them from rank 0.""" + if len(module_states) > 0: + dist._broadcast_coalesced( + process_group, module_states, broadcast_bucket_size, src + ) + + +def _replace_by_prefix( + state_dict: Dict[str, Any], + old_prefix: str, + new_prefix: str, +) -> None: + """ + Replace all keys that match a given old_prefix with a new_prefix (in-place). + + Usage:: + + state_dict = {"layer.xyz": torch.tensor(1)} + replace_by_prefix_(state_dict, "layer.", "module.layer.") + assert state_dict == {"module.layer.xyz": torch.tensor(1)} + """ + if old_prefix == new_prefix: + raise ValueError("old_prefix and new_prefix must be distinct") + for key in list(state_dict.keys()): + if not key.startswith(old_prefix): + continue + new_key = new_prefix + key[len(old_prefix) :] + state_dict[new_key] = state_dict[key] + del state_dict[key] + + +def _data_ptr_allocated(tensor: torch.Tensor) -> bool: + return tensor.untyped_storage().data_ptr() > 0 + + +def _get_root_modules(modules: List[nn.Module]) -> List[nn.Module]: + """ + Returns the modules in ``modules`` that are root modules (i.e. + parent-less) with respect to the set ``modules``. In other words, these + are the modules in ``modules`` that are the not child of any other + module in ``modules``. + """ + root_modules: List[nn.Module] = [] + module_to_modules: Dict[nn.Module, Set[nn.Module]] = { + module: set(module.modules()) for module in modules + } + for candidate_module in modules: + is_root_module = True + for module, _modules in module_to_modules.items(): + is_child_module = ( + candidate_module is not module and candidate_module in _modules + ) + if is_child_module: + is_root_module = False + break + if is_root_module: + root_modules.append(candidate_module) + return root_modules diff --git a/lib/python3.12/site-packages/torch/linalg/__init__.py b/lib/python3.12/site-packages/torch/linalg/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..00aaf323c8831c593f892eb9552466517ebd61d4 --- /dev/null +++ b/lib/python3.12/site-packages/torch/linalg/__init__.py @@ -0,0 +1,2854 @@ +import torch +from torch._C import _add_docstr, _linalg # type: ignore[attr-defined] + +LinAlgError = torch._C._LinAlgError # type: ignore[attr-defined] + +Tensor = torch.Tensor + +common_notes = { + "experimental_warning": """This function is "experimental" and it may change in a future PyTorch release.""", + "sync_note": "When inputs are on a CUDA device, this function synchronizes that device with the CPU.", + "sync_note_ex": r"When the inputs are on a CUDA device, this function synchronizes only when :attr:`check_errors`\ `= True`.", + "sync_note_has_ex": ("When inputs are on a CUDA device, this function synchronizes that device with the CPU. " + "For a version of this function that does not synchronize, see :func:`{}`.") +} + + +# Note: This not only adds doc strings for functions in the linalg namespace, but +# also connects the torch.linalg Python namespace to the torch._C._linalg builtins. + +cross = _add_docstr(_linalg.linalg_cross, r""" +linalg.cross(input, other, *, dim=-1, out=None) -> Tensor + + +Computes the cross product of two 3-dimensional vectors. + +Supports input of float, double, cfloat and cdouble dtypes. Also supports batches +of vectors, for which it computes the product along the dimension :attr:`dim`. +It broadcasts over the batch dimensions. + +Args: + input (Tensor): the first input tensor. + other (Tensor): the second input tensor. + dim (int, optional): the dimension along which to take the cross-product. Default: `-1`. + +Keyword args: + out (Tensor, optional): the output tensor. Ignored if `None`. Default: `None`. + +Example: + >>> a = torch.randn(4, 3) + >>> a + tensor([[-0.3956, 1.1455, 1.6895], + [-0.5849, 1.3672, 0.3599], + [-1.1626, 0.7180, -0.0521], + [-0.1339, 0.9902, -2.0225]]) + >>> b = torch.randn(4, 3) + >>> b + tensor([[-0.0257, -1.4725, -1.2251], + [-1.1479, -0.7005, -1.9757], + [-1.3904, 0.3726, -1.1836], + [-0.9688, -0.7153, 0.2159]]) + >>> torch.linalg.cross(a, b) + tensor([[ 1.0844, -0.5281, 0.6120], + [-2.4490, -1.5687, 1.9792], + [-0.8304, -1.3037, 0.5650], + [-1.2329, 1.9883, 1.0551]]) + >>> a = torch.randn(1, 3) # a is broadcast to match shape of b + >>> a + tensor([[-0.9941, -0.5132, 0.5681]]) + >>> torch.linalg.cross(a, b) + tensor([[ 1.4653, -1.2325, 1.4507], + [ 1.4119, -2.6163, 0.1073], + [ 0.3957, -1.9666, -1.0840], + [ 0.2956, -0.3357, 0.2139]]) +""") + +cholesky = _add_docstr(_linalg.linalg_cholesky, r""" +linalg.cholesky(A, *, upper=False, out=None) -> Tensor + +Computes the Cholesky decomposition of a complex Hermitian or real symmetric positive-definite matrix. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **Cholesky decomposition** of a complex Hermitian or real symmetric positive-definite matrix +:math:`A \in \mathbb{K}^{n \times n}` is defined as + +.. math:: + + A = LL^{\text{H}}\mathrlap{\qquad L \in \mathbb{K}^{n \times n}} + +where :math:`L` is a lower triangular matrix with real positive diagonal (even in the complex case) and +:math:`L^{\text{H}}` is the conjugate transpose when :math:`L` is complex, and the transpose when :math:`L` is real-valued. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +""" + fr""" +.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.cholesky_ex")} +""" + r""" + +.. seealso:: + + :func:`torch.linalg.cholesky_ex` for a version of this operation that + skips the (slow) error checking by default and instead returns the debug + information. This makes it a faster way to check if a matrix is + positive-definite. + + :func:`torch.linalg.eigh` for a different decomposition of a Hermitian matrix. + The eigenvalue decomposition gives more information about the matrix but it + slower to compute than the Cholesky decomposition. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of symmetric or Hermitian positive-definite matrices. + +Keyword args: + upper (bool, optional): whether to return an upper triangular matrix. + The tensor returned with upper=True is the conjugate transpose of the tensor + returned with upper=False. + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Raises: + RuntimeError: if the :attr:`A` matrix or any matrix in a batched :attr:`A` is not Hermitian + (resp. symmetric) positive-definite. If :attr:`A` is a batch of matrices, + the error message will include the batch index of the first matrix that fails + to meet this condition. + +Examples:: + + >>> A = torch.randn(2, 2, dtype=torch.complex128) + >>> A = A @ A.T.conj() + torch.eye(2) # creates a Hermitian positive-definite matrix + >>> A + tensor([[2.5266+0.0000j, 1.9586-2.0626j], + [1.9586+2.0626j, 9.4160+0.0000j]], dtype=torch.complex128) + >>> L = torch.linalg.cholesky(A) + >>> L + tensor([[1.5895+0.0000j, 0.0000+0.0000j], + [1.2322+1.2976j, 2.4928+0.0000j]], dtype=torch.complex128) + >>> torch.dist(L @ L.T.conj(), A) + tensor(4.4692e-16, dtype=torch.float64) + + >>> A = torch.randn(3, 2, 2, dtype=torch.float64) + >>> A = A @ A.mT + torch.eye(2) # batch of symmetric positive-definite matrices + >>> L = torch.linalg.cholesky(A) + >>> torch.dist(L @ L.mT, A) + tensor(5.8747e-16, dtype=torch.float64) +""") + +cholesky_ex = _add_docstr(_linalg.linalg_cholesky_ex, r""" +linalg.cholesky_ex(A, *, upper=False, check_errors=False, out=None) -> (Tensor, Tensor) + +Computes the Cholesky decomposition of a complex Hermitian or real +symmetric positive-definite matrix. + +This function skips the (slow) error checking and error message construction +of :func:`torch.linalg.cholesky`, instead directly returning the LAPACK +error codes as part of a named tuple ``(L, info)``. This makes this function +a faster way to check if a matrix is positive-definite, and it provides an +opportunity to handle decomposition errors more gracefully or performantly +than :func:`torch.linalg.cholesky` does. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +If :attr:`A` is not a Hermitian positive-definite matrix, or if it's a batch of matrices +and one or more of them is not a Hermitian positive-definite matrix, +then ``info`` stores a positive integer for the corresponding matrix. +The positive integer indicates the order of the leading minor that is not positive-definite, +and the decomposition could not be completed. +``info`` filled with zeros indicates that the decomposition was successful. +If ``check_errors=True`` and ``info`` contains positive integers, then a RuntimeError is thrown. + +""" + fr""" +.. note:: {common_notes["sync_note_ex"]} + +.. warning:: {common_notes["experimental_warning"]} +""" + r""" + +.. seealso:: + :func:`torch.linalg.cholesky` is a NumPy compatible variant that always checks for errors. + +Args: + A (Tensor): the Hermitian `n \times n` matrix or the batch of such matrices of size + `(*, n, n)` where `*` is one or more batch dimensions. + +Keyword args: + upper (bool, optional): whether to return an upper triangular matrix. + The tensor returned with upper=True is the conjugate transpose of the tensor + returned with upper=False. + check_errors (bool, optional): controls whether to check the content of ``infos``. Default: `False`. + out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`. + +Examples:: + + >>> A = torch.randn(2, 2, dtype=torch.complex128) + >>> A = A @ A.t().conj() # creates a Hermitian positive-definite matrix + >>> L, info = torch.linalg.cholesky_ex(A) + >>> A + tensor([[ 2.3792+0.0000j, -0.9023+0.9831j], + [-0.9023-0.9831j, 0.8757+0.0000j]], dtype=torch.complex128) + >>> L + tensor([[ 1.5425+0.0000j, 0.0000+0.0000j], + [-0.5850-0.6374j, 0.3567+0.0000j]], dtype=torch.complex128) + >>> info + tensor(0, dtype=torch.int32) + +""") + +inv = _add_docstr(_linalg.linalg_inv, r""" +linalg.inv(A, *, out=None) -> Tensor + +Computes the inverse of a square matrix if it exists. +Throws a `RuntimeError` if the matrix is not invertible. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +for a matrix :math:`A \in \mathbb{K}^{n \times n}`, +its **inverse matrix** :math:`A^{-1} \in \mathbb{K}^{n \times n}` (if it exists) is defined as + +.. math:: + + A^{-1}A = AA^{-1} = \mathrm{I}_n + +where :math:`\mathrm{I}_n` is the `n`-dimensional identity matrix. + +The inverse matrix exists if and only if :math:`A` is `invertible`_. In this case, +the inverse is unique. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices +then the output has the same batch dimensions. + +""" + fr""" +.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.inv_ex")} +""" + r""" + +.. note:: + Consider using :func:`torch.linalg.solve` if possible for multiplying a matrix on the left by + the inverse, as:: + + linalg.solve(A, B) == linalg.inv(A) @ B # When B is a matrix + + It is always preferred to use :func:`~solve` when possible, as it is faster and more + numerically stable than computing the inverse explicitly. + +.. seealso:: + + :func:`torch.linalg.pinv` computes the pseudoinverse (Moore-Penrose inverse) of matrices + of any shape. + + :func:`torch.linalg.solve` computes :attr:`A`\ `.inv() @ \ `:attr:`B` with a + numerically stable algorithm. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of invertible matrices. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Raises: + RuntimeError: if the matrix :attr:`A` or any matrix in the batch of matrices :attr:`A` is not invertible. + +Examples:: + + >>> A = torch.randn(4, 4) + >>> Ainv = torch.linalg.inv(A) + >>> torch.dist(A @ Ainv, torch.eye(4)) + tensor(1.1921e-07) + + >>> A = torch.randn(2, 3, 4, 4) # Batch of matrices + >>> Ainv = torch.linalg.inv(A) + >>> torch.dist(A @ Ainv, torch.eye(4)) + tensor(1.9073e-06) + + >>> A = torch.randn(4, 4, dtype=torch.complex128) # Complex matrix + >>> Ainv = torch.linalg.inv(A) + >>> torch.dist(A @ Ainv, torch.eye(4)) + tensor(7.5107e-16, dtype=torch.float64) + +.. _invertible: + https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem +""") + +solve_ex = _add_docstr(_linalg.linalg_solve_ex, r""" +linalg.solve_ex(A, B, *, left=True, check_errors=False, out=None) -> (Tensor, Tensor) + +A version of :func:`~solve` that does not perform error checks unless :attr:`check_errors`\ `= True`. +It also returns the :attr:`info` tensor returned by `LAPACK's getrf`_. + +""" + fr""" +.. note:: {common_notes["sync_note_ex"]} + +.. warning:: {common_notes["experimental_warning"]} +""" + r""" + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions. + +Keyword args: + left (bool, optional): whether to solve the system :math:`AX=B` or :math:`XA = B`. Default: `True`. + check_errors (bool, optional): controls whether to check the content of ``infos`` and raise + an error if it is non-zero. Default: `False`. + out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(result, info)`. + +Examples:: + + >>> A = torch.randn(3, 3) + >>> Ainv, info = torch.linalg.solve_ex(A) + >>> torch.dist(torch.linalg.inv(A), Ainv) + tensor(0.) + >>> info + tensor(0, dtype=torch.int32) + +.. _LAPACK's getrf: + https://www.netlib.org/lapack/explore-html/dd/d9a/group__double_g_ecomputational_ga0019443faea08275ca60a734d0593e60.html +""") + +inv_ex = _add_docstr(_linalg.linalg_inv_ex, r""" +linalg.inv_ex(A, *, check_errors=False, out=None) -> (Tensor, Tensor) + +Computes the inverse of a square matrix if it is invertible. + +Returns a namedtuple ``(inverse, info)``. ``inverse`` contains the result of +inverting :attr:`A` and ``info`` stores the LAPACK error codes. + +If :attr:`A` is not an invertible matrix, or if it's a batch of matrices +and one or more of them is not an invertible matrix, +then ``info`` stores a positive integer for the corresponding matrix. +The positive integer indicates the diagonal element of the LU decomposition of +the input matrix that is exactly zero. +``info`` filled with zeros indicates that the inversion was successful. +If ``check_errors=True`` and ``info`` contains positive integers, then a RuntimeError is thrown. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +""" + fr""" +.. note:: {common_notes["sync_note_ex"]} + +.. warning:: {common_notes["experimental_warning"]} +""" + r""" + +.. seealso:: + + :func:`torch.linalg.inv` is a NumPy compatible variant that always checks for errors. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of square matrices. + check_errors (bool, optional): controls whether to check the content of ``info``. Default: `False`. + +Keyword args: + out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`. + +Examples:: + + >>> A = torch.randn(3, 3) + >>> Ainv, info = torch.linalg.inv_ex(A) + >>> torch.dist(torch.linalg.inv(A), Ainv) + tensor(0.) + >>> info + tensor(0, dtype=torch.int32) + +""") + +det = _add_docstr(_linalg.linalg_det, r""" +linalg.det(A, *, out=None) -> Tensor + +Computes the determinant of a square matrix. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +.. seealso:: + + :func:`torch.linalg.slogdet` computes the sign and natural logarithm of the absolute + value of the determinant of square matrices. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Examples:: + + >>> A = torch.randn(3, 3) + >>> torch.linalg.det(A) + tensor(0.0934) + + >>> A = torch.randn(3, 2, 2) + >>> torch.linalg.det(A) + tensor([1.1990, 0.4099, 0.7386]) +""") + +slogdet = _add_docstr(_linalg.linalg_slogdet, r""" +linalg.slogdet(A, *, out=None) -> (Tensor, Tensor) + +Computes the sign and natural logarithm of the absolute value of the determinant of a square matrix. + +For complex :attr:`A`, it returns the sign and the natural logarithm of the modulus of the +determinant, that is, a logarithmic polar decomposition of the determinant. + +The determinant can be recovered as `sign * exp(logabsdet)`. +When a matrix has a determinant of zero, it returns `(0, -inf)`. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +.. seealso:: + + :func:`torch.linalg.det` computes the determinant of square matrices. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions. + +Keyword args: + out (tuple, optional): output tuple of two tensors. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(sign, logabsdet)`. + + `sign` will have the same dtype as :attr:`A`. + + `logabsdet` will always be real-valued, even when :attr:`A` is complex. + +Examples:: + + >>> A = torch.randn(3, 3) + >>> A + tensor([[ 0.0032, -0.2239, -1.1219], + [-0.6690, 0.1161, 0.4053], + [-1.6218, -0.9273, -0.0082]]) + >>> torch.linalg.det(A) + tensor(-0.7576) + >>> torch.logdet(A) + tensor(nan) + >>> torch.linalg.slogdet(A) + torch.return_types.linalg_slogdet(sign=tensor(-1.), logabsdet=tensor(-0.2776)) +""") + +eig = _add_docstr(_linalg.linalg_eig, r""" +linalg.eig(A, *, out=None) -> (Tensor, Tensor) + +Computes the eigenvalue decomposition of a square matrix if it exists. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **eigenvalue decomposition** of a square matrix +:math:`A \in \mathbb{K}^{n \times n}` (if it exists) is defined as + +.. math:: + + A = V \operatorname{diag}(\Lambda) V^{-1}\mathrlap{\qquad V \in \mathbb{C}^{n \times n}, \Lambda \in \mathbb{C}^n} + +This decomposition exists if and only if :math:`A` is `diagonalizable`_. +This is the case when all its eigenvalues are different. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +The returned eigenvalues are not guaranteed to be in any specific order. + +.. note:: The eigenvalues and eigenvectors of a real matrix may be complex. + +""" + fr""" +.. note:: {common_notes["sync_note"]} +""" + r""" + +.. warning:: This function assumes that :attr:`A` is `diagonalizable`_ (for example, when all the + eigenvalues are different). If it is not diagonalizable, the returned + eigenvalues will be correct but :math:`A \neq V \operatorname{diag}(\Lambda)V^{-1}`. + +.. warning:: The returned eigenvectors are normalized to have norm `1`. + Even then, the eigenvectors of a matrix are not unique, nor are they continuous with respect to + :attr:`A`. Due to this lack of uniqueness, different hardware and software may compute + different eigenvectors. + + This non-uniqueness is caused by the fact that multiplying an eigenvector by + by :math:`e^{i \phi}, \phi \in \mathbb{R}` produces another set of valid eigenvectors + of the matrix. For this reason, the loss function shall not depend on the phase of the + eigenvectors, as this quantity is not well-defined. + This is checked when computing the gradients of this function. As such, + when inputs are on a CUDA device, the computation of the gradients + of this function synchronizes that device with the CPU. + + +.. warning:: Gradients computed using the `eigenvectors` tensor will only be finite when + :attr:`A` has distinct eigenvalues. + Furthermore, if the distance between any two eigenvalues is close to zero, + the gradient will be numerically unstable, as it depends on the eigenvalues + :math:`\lambda_i` through the computation of + :math:`\frac{1}{\min_{i \neq j} \lambda_i - \lambda_j}`. + +.. seealso:: + + :func:`torch.linalg.eigvals` computes only the eigenvalues. + Unlike :func:`torch.linalg.eig`, the gradients of :func:`~eigvals` are always + numerically stable. + + :func:`torch.linalg.eigh` for a (faster) function that computes the eigenvalue decomposition + for Hermitian and symmetric matrices. + + :func:`torch.linalg.svd` for a function that computes another type of spectral + decomposition that works on matrices of any shape. + + :func:`torch.linalg.qr` for another (much faster) decomposition that works on matrices of + any shape. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of diagonalizable matrices. + +Keyword args: + out (tuple, optional): output tuple of two tensors. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(eigenvalues, eigenvectors)` which corresponds to :math:`\Lambda` and :math:`V` above. + + `eigenvalues` and `eigenvectors` will always be complex-valued, even when :attr:`A` is real. The eigenvectors + will be given by the columns of `eigenvectors`. + +Examples:: + + >>> A = torch.randn(2, 2, dtype=torch.complex128) + >>> A + tensor([[ 0.9828+0.3889j, -0.4617+0.3010j], + [ 0.1662-0.7435j, -0.6139+0.0562j]], dtype=torch.complex128) + >>> L, V = torch.linalg.eig(A) + >>> L + tensor([ 1.1226+0.5738j, -0.7537-0.1286j], dtype=torch.complex128) + >>> V + tensor([[ 0.9218+0.0000j, 0.1882-0.2220j], + [-0.0270-0.3867j, 0.9567+0.0000j]], dtype=torch.complex128) + >>> torch.dist(V @ torch.diag(L) @ torch.linalg.inv(V), A) + tensor(7.7119e-16, dtype=torch.float64) + + >>> A = torch.randn(3, 2, 2, dtype=torch.float64) + >>> L, V = torch.linalg.eig(A) + >>> torch.dist(V @ torch.diag_embed(L) @ torch.linalg.inv(V), A) + tensor(3.2841e-16, dtype=torch.float64) + +.. _diagonalizable: + https://en.wikipedia.org/wiki/Diagonalizable_matrix#Definition +""") + +eigvals = _add_docstr(_linalg.linalg_eigvals, r""" +linalg.eigvals(A, *, out=None) -> Tensor + +Computes the eigenvalues of a square matrix. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **eigenvalues** of a square matrix :math:`A \in \mathbb{K}^{n \times n}` are defined +as the roots (counted with multiplicity) of the polynomial `p` of degree `n` given by + +.. math:: + + p(\lambda) = \operatorname{det}(A - \lambda \mathrm{I}_n)\mathrlap{\qquad \lambda \in \mathbb{C}} + +where :math:`\mathrm{I}_n` is the `n`-dimensional identity matrix. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +The returned eigenvalues are not guaranteed to be in any specific order. + +.. note:: The eigenvalues of a real matrix may be complex, as the roots of a real polynomial may be complex. + + The eigenvalues of a matrix are always well-defined, even when the matrix is not diagonalizable. + +""" + fr""" +.. note:: {common_notes["sync_note"]} +""" + r""" + +.. seealso:: + + :func:`torch.linalg.eig` computes the full eigenvalue decomposition. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Returns: + A complex-valued tensor containing the eigenvalues even when :attr:`A` is real. + +Examples:: + + >>> A = torch.randn(2, 2, dtype=torch.complex128) + >>> L = torch.linalg.eigvals(A) + >>> L + tensor([ 1.1226+0.5738j, -0.7537-0.1286j], dtype=torch.complex128) + + >>> torch.dist(L, torch.linalg.eig(A).eigenvalues) + tensor(2.4576e-07) +""") + +eigh = _add_docstr(_linalg.linalg_eigh, r""" +linalg.eigh(A, UPLO='L', *, out=None) -> (Tensor, Tensor) + +Computes the eigenvalue decomposition of a complex Hermitian or real symmetric matrix. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **eigenvalue decomposition** of a complex Hermitian or real symmetric matrix +:math:`A \in \mathbb{K}^{n \times n}` is defined as + +.. math:: + + A = Q \operatorname{diag}(\Lambda) Q^{\text{H}}\mathrlap{\qquad Q \in \mathbb{K}^{n \times n}, \Lambda \in \mathbb{R}^n} + +where :math:`Q^{\text{H}}` is the conjugate transpose when :math:`Q` is complex, and the transpose when :math:`Q` is real-valued. +:math:`Q` is orthogonal in the real case and unitary in the complex case. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +:attr:`A` is assumed to be Hermitian (resp. symmetric), but this is not checked internally, instead: + +- If :attr:`UPLO`\ `= 'L'` (default), only the lower triangular part of the matrix is used in the computation. +- If :attr:`UPLO`\ `= 'U'`, only the upper triangular part of the matrix is used. + +The eigenvalues are returned in ascending order. + +""" + fr""" +.. note:: {common_notes["sync_note"]} +""" + r""" + +.. note:: The eigenvalues of real symmetric or complex Hermitian matrices are always real. + +.. warning:: The eigenvectors of a symmetric matrix are not unique, nor are they continuous with + respect to :attr:`A`. Due to this lack of uniqueness, different hardware and + software may compute different eigenvectors. + + This non-uniqueness is caused by the fact that multiplying an eigenvector by + `-1` in the real case or by :math:`e^{i \phi}, \phi \in \mathbb{R}` in the complex + case produces another set of valid eigenvectors of the matrix. + For this reason, the loss function shall not depend on the phase of the eigenvectors, as + this quantity is not well-defined. + This is checked for complex inputs when computing the gradients of this function. As such, + when inputs are complex and are on a CUDA device, the computation of the gradients + of this function synchronizes that device with the CPU. + +.. warning:: Gradients computed using the `eigenvectors` tensor will only be finite when + :attr:`A` has distinct eigenvalues. + Furthermore, if the distance between any two eigenvalues is close to zero, + the gradient will be numerically unstable, as it depends on the eigenvalues + :math:`\lambda_i` through the computation of + :math:`\frac{1}{\min_{i \neq j} \lambda_i - \lambda_j}`. + +.. warning:: User may see pytorch crashes if running `eigh` on CUDA devices with CUDA versions before 12.1 update 1 + with large ill-conditioned matrices as inputs. + Refer to :ref:`Linear Algebra Numerical Stability` for more details. + If this is the case, user may (1) tune their matrix inputs to be less ill-conditioned, + or (2) use :func:`torch.backends.cuda.preferred_linalg_library` to + try other supported backends. + +.. seealso:: + + :func:`torch.linalg.eigvalsh` computes only the eigenvalues of a Hermitian matrix. + Unlike :func:`torch.linalg.eigh`, the gradients of :func:`~eigvalsh` are always + numerically stable. + + :func:`torch.linalg.cholesky` for a different decomposition of a Hermitian matrix. + The Cholesky decomposition gives less information about the matrix but is much faster + to compute than the eigenvalue decomposition. + + :func:`torch.linalg.eig` for a (slower) function that computes the eigenvalue decomposition + of a not necessarily Hermitian square matrix. + + :func:`torch.linalg.svd` for a (slower) function that computes the more general SVD + decomposition of matrices of any shape. + + :func:`torch.linalg.qr` for another (much faster) decomposition that works on general + matrices. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of symmetric or Hermitian matrices. + UPLO ('L', 'U', optional): controls whether to use the upper or lower triangular part + of :attr:`A` in the computations. Default: `'L'`. + +Keyword args: + out (tuple, optional): output tuple of two tensors. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(eigenvalues, eigenvectors)` which corresponds to :math:`\Lambda` and :math:`Q` above. + + `eigenvalues` will always be real-valued, even when :attr:`A` is complex. + It will also be ordered in ascending order. + + `eigenvectors` will have the same dtype as :attr:`A` and will contain the eigenvectors as its columns. + +Examples:: + >>> A = torch.randn(2, 2, dtype=torch.complex128) + >>> A = A + A.T.conj() # creates a Hermitian matrix + >>> A + tensor([[2.9228+0.0000j, 0.2029-0.0862j], + [0.2029+0.0862j, 0.3464+0.0000j]], dtype=torch.complex128) + >>> L, Q = torch.linalg.eigh(A) + >>> L + tensor([0.3277, 2.9415], dtype=torch.float64) + >>> Q + tensor([[-0.0846+-0.0000j, -0.9964+0.0000j], + [ 0.9170+0.3898j, -0.0779-0.0331j]], dtype=torch.complex128) + >>> torch.dist(Q @ torch.diag(L.cdouble()) @ Q.T.conj(), A) + tensor(6.1062e-16, dtype=torch.float64) + + >>> A = torch.randn(3, 2, 2, dtype=torch.float64) + >>> A = A + A.mT # creates a batch of symmetric matrices + >>> L, Q = torch.linalg.eigh(A) + >>> torch.dist(Q @ torch.diag_embed(L) @ Q.mH, A) + tensor(1.5423e-15, dtype=torch.float64) +""") + +eigvalsh = _add_docstr(_linalg.linalg_eigvalsh, r""" +linalg.eigvalsh(A, UPLO='L', *, out=None) -> Tensor + +Computes the eigenvalues of a complex Hermitian or real symmetric matrix. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **eigenvalues** of a complex Hermitian or real symmetric matrix :math:`A \in \mathbb{K}^{n \times n}` +are defined as the roots (counted with multiplicity) of the polynomial `p` of degree `n` given by + +.. math:: + + p(\lambda) = \operatorname{det}(A - \lambda \mathrm{I}_n)\mathrlap{\qquad \lambda \in \mathbb{R}} + +where :math:`\mathrm{I}_n` is the `n`-dimensional identity matrix. +The eigenvalues of a real symmetric or complex Hermitian matrix are always real. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +The eigenvalues are returned in ascending order. + +:attr:`A` is assumed to be Hermitian (resp. symmetric), but this is not checked internally, instead: + +- If :attr:`UPLO`\ `= 'L'` (default), only the lower triangular part of the matrix is used in the computation. +- If :attr:`UPLO`\ `= 'U'`, only the upper triangular part of the matrix is used. + +""" + fr""" +.. note:: {common_notes["sync_note"]} +""" + r""" + +.. seealso:: + + :func:`torch.linalg.eigh` computes the full eigenvalue decomposition. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of symmetric or Hermitian matrices. + UPLO ('L', 'U', optional): controls whether to use the upper or lower triangular part + of :attr:`A` in the computations. Default: `'L'`. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Returns: + A real-valued tensor containing the eigenvalues even when :attr:`A` is complex. + The eigenvalues are returned in ascending order. + +Examples:: + + >>> A = torch.randn(2, 2, dtype=torch.complex128) + >>> A = A + A.T.conj() # creates a Hermitian matrix + >>> A + tensor([[2.9228+0.0000j, 0.2029-0.0862j], + [0.2029+0.0862j, 0.3464+0.0000j]], dtype=torch.complex128) + >>> torch.linalg.eigvalsh(A) + tensor([0.3277, 2.9415], dtype=torch.float64) + + >>> A = torch.randn(3, 2, 2, dtype=torch.float64) + >>> A = A + A.mT # creates a batch of symmetric matrices + >>> torch.linalg.eigvalsh(A) + tensor([[ 2.5797, 3.4629], + [-4.1605, 1.3780], + [-3.1113, 2.7381]], dtype=torch.float64) +""") + +householder_product = _add_docstr(_linalg.linalg_householder_product, r""" +householder_product(A, tau, *, out=None) -> Tensor + +Computes the first `n` columns of a product of Householder matrices. + +Let :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, and +let :math:`A \in \mathbb{K}^{m \times n}` be a matrix with columns :math:`a_i \in \mathbb{K}^m` +for :math:`i=1,\ldots,m` with :math:`m \geq n`. Denote by :math:`b_i` the vector resulting from +zeroing out the first :math:`i-1` components of :math:`a_i` and setting to `1` the :math:`i`-th. +For a vector :math:`\tau \in \mathbb{K}^k` with :math:`k \leq n`, this function computes the +first :math:`n` columns of the matrix + +.. math:: + + H_1H_2 ... H_k \qquad\text{with}\qquad H_i = \mathrm{I}_m - \tau_i b_i b_i^{\text{H}} + +where :math:`\mathrm{I}_m` is the `m`-dimensional identity matrix and :math:`b^{\text{H}}` is the +conjugate transpose when :math:`b` is complex, and the transpose when :math:`b` is real-valued. +The output matrix is the same size as the input matrix :attr:`A`. + +See `Representation of Orthogonal or Unitary Matrices`_ for further details. + +Supports inputs of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if the inputs are batches of matrices then +the output has the same batch dimensions. + +.. seealso:: + + :func:`torch.geqrf` can be used together with this function to form the `Q` from the + :func:`~qr` decomposition. + + :func:`torch.ormqr` is a related function that computes the matrix multiplication + of a product of Householder matrices with another matrix. + However, that function is not supported by autograd. + +.. warning:: + Gradient computations are only well-defined if :math:`\tau_i \neq \frac{1}{||a_i||^2}`. + If this condition is not met, no error will be thrown, but the gradient produced may contain `NaN`. + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + tau (Tensor): tensor of shape `(*, k)` where `*` is zero or more batch dimensions. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Raises: + RuntimeError: if :attr:`A` doesn't satisfy the requirement `m >= n`, + or :attr:`tau` doesn't satisfy the requirement `n >= k`. + +Examples:: + + >>> A = torch.randn(2, 2) + >>> h, tau = torch.geqrf(A) + >>> Q = torch.linalg.householder_product(h, tau) + >>> torch.dist(Q, torch.linalg.qr(A).Q) + tensor(0.) + + >>> h = torch.randn(3, 2, 2, dtype=torch.complex128) + >>> tau = torch.randn(3, 1, dtype=torch.complex128) + >>> Q = torch.linalg.householder_product(h, tau) + >>> Q + tensor([[[ 1.8034+0.4184j, 0.2588-1.0174j], + [-0.6853+0.7953j, 2.0790+0.5620j]], + + [[ 1.4581+1.6989j, -1.5360+0.1193j], + [ 1.3877-0.6691j, 1.3512+1.3024j]], + + [[ 1.4766+0.5783j, 0.0361+0.6587j], + [ 0.6396+0.1612j, 1.3693+0.4481j]]], dtype=torch.complex128) + +.. _Representation of Orthogonal or Unitary Matrices: + https://www.netlib.org/lapack/lug/node128.html +""") + +ldl_factor = _add_docstr(_linalg.linalg_ldl_factor, r""" +linalg.ldl_factor(A, *, hermitian=False, out=None) -> (Tensor, Tensor) + +Computes a compact representation of the LDL factorization of a Hermitian or symmetric (possibly indefinite) matrix. + +When :attr:`A` is complex valued it can be Hermitian (:attr:`hermitian`\ `= True`) +or symmetric (:attr:`hermitian`\ `= False`). + +The factorization is of the form the form :math:`A = L D L^T`. +If :attr:`hermitian` is `True` then transpose operation is the conjugate transpose. + +:math:`L` (or :math:`U`) and :math:`D` are stored in compact form in ``LD``. +They follow the format specified by `LAPACK's sytrf`_ function. +These tensors may be used in :func:`torch.linalg.ldl_solve` to solve linear systems. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +""" + fr""" +.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.ldl_factor_ex")} +""" + r""" + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of symmetric or Hermitian matrices. + +Keyword args: + hermitian (bool, optional): whether to consider the input to be Hermitian or symmetric. + For real-valued matrices, this switch has no effect. Default: `False`. + out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(LD, pivots)`. + +Examples:: + + >>> A = torch.randn(3, 3) + >>> A = A @ A.mT # make symmetric + >>> A + tensor([[7.2079, 4.2414, 1.9428], + [4.2414, 3.4554, 0.3264], + [1.9428, 0.3264, 1.3823]]) + >>> LD, pivots = torch.linalg.ldl_factor(A) + >>> LD + tensor([[ 7.2079, 0.0000, 0.0000], + [ 0.5884, 0.9595, 0.0000], + [ 0.2695, -0.8513, 0.1633]]) + >>> pivots + tensor([1, 2, 3], dtype=torch.int32) + +.. _LAPACK's sytrf: + https://www.netlib.org/lapack/explore-html/d3/db6/group__double_s_ycomputational_gad91bde1212277b3e909eb6af7f64858a.html +""") + +ldl_factor_ex = _add_docstr(_linalg.linalg_ldl_factor_ex, r""" +linalg.ldl_factor_ex(A, *, hermitian=False, check_errors=False, out=None) -> (Tensor, Tensor, Tensor) + +This is a version of :func:`~ldl_factor` that does not perform error checks unless :attr:`check_errors`\ `= True`. +It also returns the :attr:`info` tensor returned by `LAPACK's sytrf`_. +``info`` stores integer error codes from the backend library. +A positive integer indicates the diagonal element of :math:`D` that is zero. +Division by 0 will occur if the result is used for solving a system of linear equations. +``info`` filled with zeros indicates that the factorization was successful. +If ``check_errors=True`` and ``info`` contains positive integers, then a `RuntimeError` is thrown. + +""" + fr""" +.. note:: {common_notes["sync_note_ex"]} + +.. warning:: {common_notes["experimental_warning"]} +""" + r""" + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of symmetric or Hermitian matrices. + +Keyword args: + hermitian (bool, optional): whether to consider the input to be Hermitian or symmetric. + For real-valued matrices, this switch has no effect. Default: `False`. + check_errors (bool, optional): controls whether to check the content of ``info`` and raise + an error if it is non-zero. Default: `False`. + out (tuple, optional): tuple of three tensors to write the output to. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(LD, pivots, info)`. + +Examples:: + + >>> A = torch.randn(3, 3) + >>> A = A @ A.mT # make symmetric + >>> A + tensor([[7.2079, 4.2414, 1.9428], + [4.2414, 3.4554, 0.3264], + [1.9428, 0.3264, 1.3823]]) + >>> LD, pivots, info = torch.linalg.ldl_factor_ex(A) + >>> LD + tensor([[ 7.2079, 0.0000, 0.0000], + [ 0.5884, 0.9595, 0.0000], + [ 0.2695, -0.8513, 0.1633]]) + >>> pivots + tensor([1, 2, 3], dtype=torch.int32) + >>> info + tensor(0, dtype=torch.int32) + +.. _LAPACK's sytrf: + https://www.netlib.org/lapack/explore-html/d3/db6/group__double_s_ycomputational_gad91bde1212277b3e909eb6af7f64858a.html +""") + +ldl_solve = _add_docstr(_linalg.linalg_ldl_solve, r""" +linalg.ldl_solve(LD, pivots, B, *, hermitian=False, out=None) -> Tensor + +Computes the solution of a system of linear equations using the LDL factorization. + +:attr:`LD` and :attr:`pivots` are the compact representation of the LDL factorization and +are expected to be computed by :func:`torch.linalg.ldl_factor_ex`. +:attr:`hermitian` argument to this function should be the same +as the corresponding arguments in :func:`torch.linalg.ldl_factor_ex`. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +""" + fr""" +.. warning:: {common_notes["experimental_warning"]} +""" + r""" + +Args: + LD (Tensor): the `n \times n` matrix or the batch of such matrices of size + `(*, n, n)` where `*` is one or more batch dimensions. + pivots (Tensor): the pivots corresponding to the LDL factorization of :attr:`LD`. + B (Tensor): right-hand side tensor of shape `(*, n, k)`. + +Keyword args: + hermitian (bool, optional): whether to consider the decomposed matrix to be Hermitian or symmetric. + For real-valued matrices, this switch has no effect. Default: `False`. + out (tuple, optional): output tensor. `B` may be passed as `out` and the result is computed in-place on `B`. + Ignored if `None`. Default: `None`. + +Examples:: + + >>> A = torch.randn(2, 3, 3) + >>> A = A @ A.mT # make symmetric + >>> LD, pivots, info = torch.linalg.ldl_factor_ex(A) + >>> B = torch.randn(2, 3, 4) + >>> X = torch.linalg.ldl_solve(LD, pivots, B) + >>> torch.linalg.norm(A @ X - B) + >>> tensor(0.0001) +""") + +lstsq = _add_docstr(_linalg.linalg_lstsq, r""" +torch.linalg.lstsq(A, B, rcond=None, *, driver=None) -> (Tensor, Tensor, Tensor, Tensor) + +Computes a solution to the least squares problem of a system of linear equations. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **least squares problem** for a linear system :math:`AX = B` with +:math:`A \in \mathbb{K}^{m \times n}, B \in \mathbb{K}^{m \times k}` is defined as + +.. math:: + + \min_{X \in \mathbb{K}^{n \times k}} \|AX - B\|_F + +where :math:`\|-\|_F` denotes the Frobenius norm. + +Supports inputs of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if the inputs are batches of matrices then +the output has the same batch dimensions. + +:attr:`driver` chooses the backend function that will be used. +For CPU inputs the valid values are `'gels'`, `'gelsy'`, `'gelsd`, `'gelss'`. +To choose the best driver on CPU consider: + +- If :attr:`A` is well-conditioned (its `condition number`_ is not too large), or you do not mind some precision loss. + + - For a general matrix: `'gelsy'` (QR with pivoting) (default) + - If :attr:`A` is full-rank: `'gels'` (QR) + +- If :attr:`A` is not well-conditioned. + + - `'gelsd'` (tridiagonal reduction and SVD) + - But if you run into memory issues: `'gelss'` (full SVD). + +For CUDA input, the only valid driver is `'gels'`, which assumes that :attr:`A` is full-rank. + +See also the `full description of these drivers`_ + +:attr:`rcond` is used to determine the effective rank of the matrices in :attr:`A` +when :attr:`driver` is one of (`'gelsy'`, `'gelsd'`, `'gelss'`). +In this case, if :math:`\sigma_i` are the singular values of `A` in decreasing order, +:math:`\sigma_i` will be rounded down to zero if :math:`\sigma_i \leq \text{rcond} \cdot \sigma_1`. +If :attr:`rcond`\ `= None` (default), :attr:`rcond` is set to the machine precision of the dtype of :attr:`A` times `max(m, n)`. + +This function returns the solution to the problem and some extra information in a named tuple of +four tensors `(solution, residuals, rank, singular_values)`. For inputs :attr:`A`, :attr:`B` +of shape `(*, m, n)`, `(*, m, k)` respectively, it contains + +- `solution`: the least squares solution. It has shape `(*, n, k)`. +- `residuals`: the squared residuals of the solutions, that is, :math:`\|AX - B\|_F^2`. + It has shape equal to the batch dimensions of :attr:`A`. + It is computed when `m > n` and every matrix in :attr:`A` is full-rank, + otherwise, it is an empty tensor. + If :attr:`A` is a batch of matrices and any matrix in the batch is not full rank, + then an empty tensor is returned. This behavior may change in a future PyTorch release. +- `rank`: tensor of ranks of the matrices in :attr:`A`. + It has shape equal to the batch dimensions of :attr:`A`. + It is computed when :attr:`driver` is one of (`'gelsy'`, `'gelsd'`, `'gelss'`), + otherwise it is an empty tensor. +- `singular_values`: tensor of singular values of the matrices in :attr:`A`. + It has shape `(*, min(m, n))`. + It is computed when :attr:`driver` is one of (`'gelsd'`, `'gelss'`), + otherwise it is an empty tensor. + +.. note:: + This function computes `X = \ `:attr:`A`\ `.pinverse() @ \ `:attr:`B` in a faster and + more numerically stable way than performing the computations separately. + +.. warning:: + The default value of :attr:`rcond` may change in a future PyTorch release. + It is therefore recommended to use a fixed value to avoid potential + breaking changes. + +Args: + A (Tensor): lhs tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + B (Tensor): rhs tensor of shape `(*, m, k)` where `*` is zero or more batch dimensions. + rcond (float, optional): used to determine the effective rank of :attr:`A`. + If :attr:`rcond`\ `= None`, :attr:`rcond` is set to the machine + precision of the dtype of :attr:`A` times `max(m, n)`. Default: `None`. + +Keyword args: + driver (str, optional): name of the LAPACK/MAGMA method to be used. + If `None`, `'gelsy'` is used for CPU inputs and `'gels'` for CUDA inputs. + Default: `None`. + +Returns: + A named tuple `(solution, residuals, rank, singular_values)`. + +Examples:: + + >>> A = torch.randn(1,3,3) + >>> A + tensor([[[-1.0838, 0.0225, 0.2275], + [ 0.2438, 0.3844, 0.5499], + [ 0.1175, -0.9102, 2.0870]]]) + >>> B = torch.randn(2,3,3) + >>> B + tensor([[[-0.6772, 0.7758, 0.5109], + [-1.4382, 1.3769, 1.1818], + [-0.3450, 0.0806, 0.3967]], + [[-1.3994, -0.1521, -0.1473], + [ 1.9194, 1.0458, 0.6705], + [-1.1802, -0.9796, 1.4086]]]) + >>> X = torch.linalg.lstsq(A, B).solution # A is broadcasted to shape (2, 3, 3) + >>> torch.dist(X, torch.linalg.pinv(A) @ B) + tensor(1.5152e-06) + + >>> S = torch.linalg.lstsq(A, B, driver='gelsd').singular_values + >>> torch.dist(S, torch.linalg.svdvals(A)) + tensor(2.3842e-07) + + >>> A[:, 0].zero_() # Decrease the rank of A + >>> rank = torch.linalg.lstsq(A, B).rank + >>> rank + tensor([2]) + +.. _condition number: + https://pytorch.org/docs/main/linalg.html#torch.linalg.cond +.. _full description of these drivers: + https://www.netlib.org/lapack/lug/node27.html +""") + +matrix_power = _add_docstr(_linalg.linalg_matrix_power, r""" +matrix_power(A, n, *, out=None) -> Tensor + +Computes the `n`-th power of a square matrix for an integer `n`. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +If :attr:`n`\ `= 0`, it returns the identity matrix (or batch) of the same shape +as :attr:`A`. If :attr:`n` is negative, it returns the inverse of each matrix +(if invertible) raised to the power of `abs(n)`. + +.. note:: + Consider using :func:`torch.linalg.solve` if possible for multiplying a matrix on the left by + a negative power as, if :attr:`n`\ `> 0`:: + + torch.linalg.solve(matrix_power(A, n), B) == matrix_power(A, -n) @ B + + It is always preferred to use :func:`~solve` when possible, as it is faster and more + numerically stable than computing :math:`A^{-n}` explicitly. + +.. seealso:: + + :func:`torch.linalg.solve` computes :attr:`A`\ `.inverse() @ \ `:attr:`B` with a + numerically stable algorithm. + +Args: + A (Tensor): tensor of shape `(*, m, m)` where `*` is zero or more batch dimensions. + n (int): the exponent. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Raises: + RuntimeError: if :attr:`n`\ `< 0` and the matrix :attr:`A` or any matrix in the + batch of matrices :attr:`A` is not invertible. + +Examples:: + + >>> A = torch.randn(3, 3) + >>> torch.linalg.matrix_power(A, 0) + tensor([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]) + >>> torch.linalg.matrix_power(A, 3) + tensor([[ 1.0756, 0.4980, 0.0100], + [-1.6617, 1.4994, -1.9980], + [-0.4509, 0.2731, 0.8001]]) + >>> torch.linalg.matrix_power(A.expand(2, -1, -1), -2) + tensor([[[ 0.2640, 0.4571, -0.5511], + [-1.0163, 0.3491, -1.5292], + [-0.4899, 0.0822, 0.2773]], + [[ 0.2640, 0.4571, -0.5511], + [-1.0163, 0.3491, -1.5292], + [-0.4899, 0.0822, 0.2773]]]) +""") + +matrix_rank = _add_docstr(_linalg.linalg_matrix_rank, r""" +linalg.matrix_rank(A, *, atol=None, rtol=None, hermitian=False, out=None) -> Tensor + +Computes the numerical rank of a matrix. + +The matrix rank is computed as the number of singular values +(or eigenvalues in absolute value when :attr:`hermitian`\ `= True`) +that are greater than :math:`\max(\text{atol}, \sigma_1 * \text{rtol})` threshold, +where :math:`\sigma_1` is the largest singular value (or eigenvalue). + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +If :attr:`hermitian`\ `= True`, :attr:`A` is assumed to be Hermitian if complex or +symmetric if real, but this is not checked internally. Instead, just the lower +triangular part of the matrix is used in the computations. + +If :attr:`rtol` is not specified and :attr:`A` is a matrix of dimensions `(m, n)`, +the relative tolerance is set to be :math:`\text{rtol} = \max(m, n) \varepsilon` +and :math:`\varepsilon` is the epsilon value for the dtype of :attr:`A` (see :class:`.finfo`). +If :attr:`rtol` is not specified and :attr:`atol` is specified to be larger than zero then +:attr:`rtol` is set to zero. + +If :attr:`atol` or :attr:`rtol` is a :class:`torch.Tensor`, its shape must be broadcastable to that +of the singular values of :attr:`A` as returned by :func:`torch.linalg.svdvals`. + +.. note:: + This function has NumPy compatible variant `linalg.matrix_rank(A, tol, hermitian=False)`. + However, use of the positional argument :attr:`tol` is deprecated in favor of :attr:`atol` and :attr:`rtol`. + +""" + fr""" +.. note:: The matrix rank is computed using a singular value decomposition + :func:`torch.linalg.svdvals` if :attr:`hermitian`\ `= False` (default) and the eigenvalue + decomposition :func:`torch.linalg.eigvalsh` when :attr:`hermitian`\ `= True`. + {common_notes["sync_note"]} +""" + r""" + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + tol (float, Tensor, optional): [NumPy Compat] Alias for :attr:`atol`. Default: `None`. + +Keyword args: + atol (float, Tensor, optional): the absolute tolerance value. When `None` it's considered to be zero. + Default: `None`. + rtol (float, Tensor, optional): the relative tolerance value. See above for the value it takes when `None`. + Default: `None`. + hermitian(bool): indicates whether :attr:`A` is Hermitian if complex + or symmetric if real. Default: `False`. + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Examples:: + + >>> A = torch.eye(10) + >>> torch.linalg.matrix_rank(A) + tensor(10) + >>> B = torch.eye(10) + >>> B[0, 0] = 0 + >>> torch.linalg.matrix_rank(B) + tensor(9) + + >>> A = torch.randn(4, 3, 2) + >>> torch.linalg.matrix_rank(A) + tensor([2, 2, 2, 2]) + + >>> A = torch.randn(2, 4, 2, 3) + >>> torch.linalg.matrix_rank(A) + tensor([[2, 2, 2, 2], + [2, 2, 2, 2]]) + + >>> A = torch.randn(2, 4, 3, 3, dtype=torch.complex64) + >>> torch.linalg.matrix_rank(A) + tensor([[3, 3, 3, 3], + [3, 3, 3, 3]]) + >>> torch.linalg.matrix_rank(A, hermitian=True) + tensor([[3, 3, 3, 3], + [3, 3, 3, 3]]) + >>> torch.linalg.matrix_rank(A, atol=1.0, rtol=0.0) + tensor([[3, 2, 2, 2], + [1, 2, 1, 2]]) + >>> torch.linalg.matrix_rank(A, atol=1.0, rtol=0.0, hermitian=True) + tensor([[2, 2, 2, 1], + [1, 2, 2, 2]]) +""") + +norm = _add_docstr(_linalg.linalg_norm, r""" +linalg.norm(A, ord=None, dim=None, keepdim=False, *, out=None, dtype=None) -> Tensor + +Computes a vector or matrix norm. + +Supports input of float, double, cfloat and cdouble dtypes. + +Whether this function computes a vector or matrix norm is determined as follows: + +- If :attr:`dim` is an `int`, the vector norm will be computed. +- If :attr:`dim` is a `2`-`tuple`, the matrix norm will be computed. +- If :attr:`dim`\ `= None` and :attr:`ord`\ `= None`, + :attr:`A` will be flattened to 1D and the `2`-norm of the resulting vector will be computed. +- If :attr:`dim`\ `= None` and :attr:`ord` `!= None`, :attr:`A` must be 1D or 2D. + +:attr:`ord` defines the norm that is computed. The following norms are supported: + +====================== ========================= ======================================================== +:attr:`ord` norm for matrices norm for vectors +====================== ========================= ======================================================== +`None` (default) Frobenius norm `2`-norm (see below) +`'fro'` Frobenius norm -- not supported -- +`'nuc'` nuclear norm -- not supported -- +`inf` `max(sum(abs(x), dim=1))` `max(abs(x))` +`-inf` `min(sum(abs(x), dim=1))` `min(abs(x))` +`0` -- not supported -- `sum(x != 0)` +`1` `max(sum(abs(x), dim=0))` as below +`-1` `min(sum(abs(x), dim=0))` as below +`2` largest singular value as below +`-2` smallest singular value as below +other `int` or `float` -- not supported -- `sum(abs(x)^{ord})^{(1 / ord)}` +====================== ========================= ======================================================== + +where `inf` refers to `float('inf')`, NumPy's `inf` object, or any equivalent object. + +.. seealso:: + + :func:`torch.linalg.vector_norm` computes a vector norm. + + :func:`torch.linalg.matrix_norm` computes a matrix norm. + + The above functions are often clearer and more flexible than using :func:`torch.linalg.norm`. + For example, `torch.linalg.norm(A, ord=1, dim=(0, 1))` always + computes a matrix norm, but with `torch.linalg.vector_norm(A, ord=1, dim=(0, 1))` it is possible + to compute a vector norm over the two dimensions. + +Args: + A (Tensor): tensor of shape `(*, n)` or `(*, m, n)` where `*` is zero or more batch dimensions + ord (int, float, inf, -inf, 'fro', 'nuc', optional): order of norm. Default: `None` + dim (int, Tuple[int], optional): dimensions over which to compute + the vector or matrix norm. See above for the behavior when :attr:`dim`\ `= None`. + Default: `None` + keepdim (bool, optional): If set to `True`, the reduced dimensions are retained + in the result as dimensions with size one. Default: `False` + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + dtype (:class:`torch.dtype`, optional): If specified, the input tensor is cast to + :attr:`dtype` before performing the operation, and the returned tensor's type + will be :attr:`dtype`. Default: `None` + +Returns: + A real-valued tensor, even when :attr:`A` is complex. + +Examples:: + + >>> from torch import linalg as LA + >>> a = torch.arange(9, dtype=torch.float) - 4 + >>> a + tensor([-4., -3., -2., -1., 0., 1., 2., 3., 4.]) + >>> B = a.reshape((3, 3)) + >>> B + tensor([[-4., -3., -2.], + [-1., 0., 1.], + [ 2., 3., 4.]]) + + >>> LA.norm(a) + tensor(7.7460) + >>> LA.norm(B) + tensor(7.7460) + >>> LA.norm(B, 'fro') + tensor(7.7460) + >>> LA.norm(a, float('inf')) + tensor(4.) + >>> LA.norm(B, float('inf')) + tensor(9.) + >>> LA.norm(a, -float('inf')) + tensor(0.) + >>> LA.norm(B, -float('inf')) + tensor(2.) + + >>> LA.norm(a, 1) + tensor(20.) + >>> LA.norm(B, 1) + tensor(7.) + >>> LA.norm(a, -1) + tensor(0.) + >>> LA.norm(B, -1) + tensor(6.) + >>> LA.norm(a, 2) + tensor(7.7460) + >>> LA.norm(B, 2) + tensor(7.3485) + + >>> LA.norm(a, -2) + tensor(0.) + >>> LA.norm(B.double(), -2) + tensor(1.8570e-16, dtype=torch.float64) + >>> LA.norm(a, 3) + tensor(5.8480) + >>> LA.norm(a, -3) + tensor(0.) + +Using the :attr:`dim` argument to compute vector norms:: + + >>> c = torch.tensor([[1., 2., 3.], + ... [-1, 1, 4]]) + >>> LA.norm(c, dim=0) + tensor([1.4142, 2.2361, 5.0000]) + >>> LA.norm(c, dim=1) + tensor([3.7417, 4.2426]) + >>> LA.norm(c, ord=1, dim=1) + tensor([6., 6.]) + +Using the :attr:`dim` argument to compute matrix norms:: + + >>> A = torch.arange(8, dtype=torch.float).reshape(2, 2, 2) + >>> LA.norm(A, dim=(1,2)) + tensor([ 3.7417, 11.2250]) + >>> LA.norm(A[0, :, :]), LA.norm(A[1, :, :]) + (tensor(3.7417), tensor(11.2250)) +""") + +vector_norm = _add_docstr(_linalg.linalg_vector_norm, r""" +linalg.vector_norm(x, ord=2, dim=None, keepdim=False, *, dtype=None, out=None) -> Tensor + +Computes a vector norm. + +If :attr:`x` is complex valued, it computes the norm of :attr:`x`\ `.abs()` + +Supports input of float, double, cfloat and cdouble dtypes. + +This function does not necessarily treat multidimensional :attr:`x` as a batch of +vectors, instead: + +- If :attr:`dim`\ `= None`, :attr:`x` will be flattened before the norm is computed. +- If :attr:`dim` is an `int` or a `tuple`, the norm will be computed over these dimensions + and the other dimensions will be treated as batch dimensions. + +This behavior is for consistency with :func:`torch.linalg.norm`. + +:attr:`ord` defines the vector norm that is computed. The following norms are supported: + +====================== =============================== +:attr:`ord` vector norm +====================== =============================== +`2` (default) `2`-norm (see below) +`inf` `max(abs(x))` +`-inf` `min(abs(x))` +`0` `sum(x != 0)` +other `int` or `float` `sum(abs(x)^{ord})^{(1 / ord)}` +====================== =============================== + +where `inf` refers to `float('inf')`, NumPy's `inf` object, or any equivalent object. + +:attr:`dtype` may be used to perform the computation in a more precise dtype. +It is semantically equivalent to calling ``linalg.vector_norm(x.to(dtype))`` +but it is faster in some cases. + +.. seealso:: + + :func:`torch.linalg.matrix_norm` computes a matrix norm. + +Args: + x (Tensor): tensor, flattened by default, but this behavior can be + controlled using :attr:`dim`. (Note: the keyword argument + `input` can also be used as an alias for `x`.) + ord (int, float, inf, -inf, 'fro', 'nuc', optional): order of norm. Default: `2` + dim (int, Tuple[int], optional): dimensions over which to compute + the norm. See above for the behavior when :attr:`dim`\ `= None`. + Default: `None` + keepdim (bool, optional): If set to `True`, the reduced dimensions are retained + in the result as dimensions with size one. Default: `False` + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + dtype (:class:`torch.dtype`, optional): type used to perform the accumulation and the return. + If specified, :attr:`x` is cast to :attr:`dtype` before performing the operation, + and the returned tensor's type will be :attr:`dtype` if real and of its real counterpart if complex. + :attr:`dtype` may be complex if :attr:`x` is complex, otherwise it must be real. + :attr:`x` should be convertible without narrowing to :attr:`dtype`. Default: None + +Returns: + A real-valued tensor, even when :attr:`x` is complex. + +Examples:: + + >>> from torch import linalg as LA + >>> a = torch.arange(9, dtype=torch.float) - 4 + >>> a + tensor([-4., -3., -2., -1., 0., 1., 2., 3., 4.]) + >>> B = a.reshape((3, 3)) + >>> B + tensor([[-4., -3., -2.], + [-1., 0., 1.], + [ 2., 3., 4.]]) + >>> LA.vector_norm(a, ord=3.5) + tensor(5.4345) + >>> LA.vector_norm(B, ord=3.5) + tensor(5.4345) +""") + +matrix_norm = _add_docstr(_linalg.linalg_matrix_norm, r""" +linalg.matrix_norm(A, ord='fro', dim=(-2, -1), keepdim=False, *, dtype=None, out=None) -> Tensor + +Computes a matrix norm. + +If :attr:`A` is complex valued, it computes the norm of :attr:`A`\ `.abs()` + +Support input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices: the norm will be computed over the +dimensions specified by the 2-tuple :attr:`dim` and the other dimensions will +be treated as batch dimensions. The output will have the same batch dimensions. + +:attr:`ord` defines the matrix norm that is computed. The following norms are supported: + +====================== ======================================================== +:attr:`ord` matrix norm +====================== ======================================================== +`'fro'` (default) Frobenius norm +`'nuc'` nuclear norm +`inf` `max(sum(abs(x), dim=1))` +`-inf` `min(sum(abs(x), dim=1))` +`1` `max(sum(abs(x), dim=0))` +`-1` `min(sum(abs(x), dim=0))` +`2` largest singular value +`-2` smallest singular value +====================== ======================================================== + +where `inf` refers to `float('inf')`, NumPy's `inf` object, or any equivalent object. + +Args: + A (Tensor): tensor with two or more dimensions. By default its + shape is interpreted as `(*, m, n)` where `*` is zero or more + batch dimensions, but this behavior can be controlled using :attr:`dim`. + ord (int, inf, -inf, 'fro', 'nuc', optional): order of norm. Default: `'fro'` + dim (Tuple[int, int], optional): dimensions over which to compute the norm. Default: `(-2, -1)` + keepdim (bool, optional): If set to `True`, the reduced dimensions are retained + in the result as dimensions with size one. Default: `False` + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + dtype (:class:`torch.dtype`, optional): If specified, the input tensor is cast to + :attr:`dtype` before performing the operation, and the returned tensor's type + will be :attr:`dtype`. Default: `None` + +Returns: + A real-valued tensor, even when :attr:`A` is complex. + +Examples:: + + >>> from torch import linalg as LA + >>> A = torch.arange(9, dtype=torch.float).reshape(3, 3) + >>> A + tensor([[0., 1., 2.], + [3., 4., 5.], + [6., 7., 8.]]) + >>> LA.matrix_norm(A) + tensor(14.2829) + >>> LA.matrix_norm(A, ord=-1) + tensor(9.) + >>> B = A.expand(2, -1, -1) + >>> B + tensor([[[0., 1., 2.], + [3., 4., 5.], + [6., 7., 8.]], + + [[0., 1., 2.], + [3., 4., 5.], + [6., 7., 8.]]]) + >>> LA.matrix_norm(B) + tensor([14.2829, 14.2829]) + >>> LA.matrix_norm(B, dim=(0, 2)) + tensor([ 3.1623, 10.0000, 17.2627]) +""") + +matmul = _add_docstr(_linalg.linalg_matmul, r""" +linalg.matmul(input, other, *, out=None) -> Tensor + +Alias for :func:`torch.matmul` +""") + +diagonal = _add_docstr(_linalg.linalg_diagonal, r""" +linalg.diagonal(A, *, offset=0, dim1=-2, dim2=-1) -> Tensor + +Alias for :func:`torch.diagonal` with defaults :attr:`dim1`\ `= -2`, :attr:`dim2`\ `= -1`. +""") + +multi_dot = _add_docstr(_linalg.linalg_multi_dot, r""" +linalg.multi_dot(tensors, *, out=None) + +Efficiently multiplies two or more matrices by reordering the multiplications so that +the fewest arithmetic operations are performed. + +Supports inputs of float, double, cfloat and cdouble dtypes. +This function does not support batched inputs. + +Every tensor in :attr:`tensors` must be 2D, except for the first and last which +may be 1D. If the first tensor is a 1D vector of shape `(n,)` it is treated as a row vector +of shape `(1, n)`, similarly if the last tensor is a 1D vector of shape `(n,)` it is treated +as a column vector of shape `(n, 1)`. + +If the first and last tensors are matrices, the output will be a matrix. +However, if either is a 1D vector, then the output will be a 1D vector. + +Differences with `numpy.linalg.multi_dot`: + +- Unlike `numpy.linalg.multi_dot`, the first and last tensors must either be 1D or 2D + whereas NumPy allows them to be nD + +.. warning:: This function does not broadcast. + +.. note:: This function is implemented by chaining :func:`torch.mm` calls after + computing the optimal matrix multiplication order. + +.. note:: The cost of multiplying two matrices with shapes `(a, b)` and `(b, c)` is + `a * b * c`. Given matrices `A`, `B`, `C` with shapes `(10, 100)`, + `(100, 5)`, `(5, 50)` respectively, we can calculate the cost of different + multiplication orders as follows: + + .. math:: + + \begin{align*} + \operatorname{cost}((AB)C) &= 10 \times 100 \times 5 + 10 \times 5 \times 50 = 7500 \\ + \operatorname{cost}(A(BC)) &= 10 \times 100 \times 50 + 100 \times 5 \times 50 = 75000 + \end{align*} + + In this case, multiplying `A` and `B` first followed by `C` is 10 times faster. + +Args: + tensors (Sequence[Tensor]): two or more tensors to multiply. The first and last + tensors may be 1D or 2D. Every other tensor must be 2D. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Examples:: + + >>> from torch.linalg import multi_dot + + >>> multi_dot([torch.tensor([1, 2]), torch.tensor([2, 3])]) + tensor(8) + >>> multi_dot([torch.tensor([[1, 2]]), torch.tensor([2, 3])]) + tensor([8]) + >>> multi_dot([torch.tensor([[1, 2]]), torch.tensor([[2], [3]])]) + tensor([[8]]) + + >>> A = torch.arange(2 * 3).view(2, 3) + >>> B = torch.arange(3 * 2).view(3, 2) + >>> C = torch.arange(2 * 2).view(2, 2) + >>> multi_dot((A, B, C)) + tensor([[ 26, 49], + [ 80, 148]]) +""") + +svd = _add_docstr(_linalg.linalg_svd, r""" +linalg.svd(A, full_matrices=True, *, driver=None, out=None) -> (Tensor, Tensor, Tensor) + +Computes the singular value decomposition (SVD) of a matrix. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **full SVD** of a matrix +:math:`A \in \mathbb{K}^{m \times n}`, if `k = min(m,n)`, is defined as + +.. math:: + + A = U \operatorname{diag}(S) V^{\text{H}} + \mathrlap{\qquad U \in \mathbb{K}^{m \times m}, S \in \mathbb{R}^k, V \in \mathbb{K}^{n \times n}} + +where :math:`\operatorname{diag}(S) \in \mathbb{K}^{m \times n}`, +:math:`V^{\text{H}}` is the conjugate transpose when :math:`V` is complex, and the transpose when :math:`V` is real-valued. +The matrices :math:`U`, :math:`V` (and thus :math:`V^{\text{H}}`) are orthogonal in the real case, and unitary in the complex case. + +When `m > n` (resp. `m < n`) we can drop the last `m - n` (resp. `n - m`) columns of `U` (resp. `V`) to form the **reduced SVD**: + +.. math:: + + A = U \operatorname{diag}(S) V^{\text{H}} + \mathrlap{\qquad U \in \mathbb{K}^{m \times k}, S \in \mathbb{R}^k, V \in \mathbb{K}^{n \times k}} + +where :math:`\operatorname{diag}(S) \in \mathbb{K}^{k \times k}`. +In this case, :math:`U` and :math:`V` also have orthonormal columns. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +The returned decomposition is a named tuple `(U, S, Vh)` +which corresponds to :math:`U`, :math:`S`, :math:`V^{\text{H}}` above. + +The singular values are returned in descending order. + +The parameter :attr:`full_matrices` chooses between the full (default) and reduced SVD. + +The :attr:`driver` kwarg may be used in CUDA with a cuSOLVER backend to choose the algorithm used to compute the SVD. +The choice of a driver is a trade-off between accuracy and speed. + +- If :attr:`A` is well-conditioned (its `condition number`_ is not too large), or you do not mind some precision loss. + + - For a general matrix: `'gesvdj'` (Jacobi method) + - If :attr:`A` is tall or wide (`m >> n` or `m << n`): `'gesvda'` (Approximate method) + +- If :attr:`A` is not well-conditioned or precision is relevant: `'gesvd'` (QR based) + +By default (:attr:`driver`\ `= None`), we call `'gesvdj'` and, if it fails, we fallback to `'gesvd'`. + +Differences with `numpy.linalg.svd`: + +- Unlike `numpy.linalg.svd`, this function always returns a tuple of three tensors + and it doesn't support `compute_uv` argument. + Please use :func:`torch.linalg.svdvals`, which computes only the singular values, + instead of `compute_uv=False`. + +.. note:: When :attr:`full_matrices`\ `= True`, the gradients with respect to `U[..., :, min(m, n):]` + and `Vh[..., min(m, n):, :]` will be ignored, as those vectors can be arbitrary bases + of the corresponding subspaces. + +.. warning:: The returned tensors `U` and `V` are not unique, nor are they continuous with + respect to :attr:`A`. + Due to this lack of uniqueness, different hardware and software may compute + different singular vectors. + + This non-uniqueness is caused by the fact that multiplying any pair of singular + vectors :math:`u_k, v_k` by `-1` in the real case or by + :math:`e^{i \phi}, \phi \in \mathbb{R}` in the complex case produces another two + valid singular vectors of the matrix. + For this reason, the loss function shall not depend on this :math:`e^{i \phi}` quantity, + as it is not well-defined. + This is checked for complex inputs when computing the gradients of this function. As such, + when inputs are complex and are on a CUDA device, the computation of the gradients + of this function synchronizes that device with the CPU. + +.. warning:: Gradients computed using `U` or `Vh` will only be finite when + :attr:`A` does not have repeated singular values. If :attr:`A` is rectangular, + additionally, zero must also not be one of its singular values. + Furthermore, if the distance between any two singular values is close to zero, + the gradient will be numerically unstable, as it depends on the singular values + :math:`\sigma_i` through the computation of + :math:`\frac{1}{\min_{i \neq j} \sigma_i^2 - \sigma_j^2}`. + In the rectangular case, the gradient will also be numerically unstable when + :attr:`A` has small singular values, as it also depends on the computation of + :math:`\frac{1}{\sigma_i}`. + +.. seealso:: + + :func:`torch.linalg.svdvals` computes only the singular values. + Unlike :func:`torch.linalg.svd`, the gradients of :func:`~svdvals` are always + numerically stable. + + :func:`torch.linalg.eig` for a function that computes another type of spectral + decomposition of a matrix. The eigendecomposition works just on square matrices. + + :func:`torch.linalg.eigh` for a (faster) function that computes the eigenvalue decomposition + for Hermitian and symmetric matrices. + + :func:`torch.linalg.qr` for another (much faster) decomposition that works on general + matrices. + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + full_matrices (bool, optional): controls whether to compute the full or reduced + SVD, and consequently, + the shape of the returned tensors + `U` and `Vh`. Default: `True`. + +Keyword args: + driver (str, optional): name of the cuSOLVER method to be used. This keyword argument only works on CUDA inputs. + Available options are: `None`, `gesvd`, `gesvdj`, and `gesvda`. + Default: `None`. + out (tuple, optional): output tuple of three tensors. Ignored if `None`. + +Returns: + A named tuple `(U, S, Vh)` which corresponds to :math:`U`, :math:`S`, :math:`V^{\text{H}}` above. + + `S` will always be real-valued, even when :attr:`A` is complex. + It will also be ordered in descending order. + + `U` and `Vh` will have the same dtype as :attr:`A`. The left / right singular vectors will be given by + the columns of `U` and the rows of `Vh` respectively. + +Examples:: + + >>> A = torch.randn(5, 3) + >>> U, S, Vh = torch.linalg.svd(A, full_matrices=False) + >>> U.shape, S.shape, Vh.shape + (torch.Size([5, 3]), torch.Size([3]), torch.Size([3, 3])) + >>> torch.dist(A, U @ torch.diag(S) @ Vh) + tensor(1.0486e-06) + + >>> U, S, Vh = torch.linalg.svd(A) + >>> U.shape, S.shape, Vh.shape + (torch.Size([5, 5]), torch.Size([3]), torch.Size([3, 3])) + >>> torch.dist(A, U[:, :3] @ torch.diag(S) @ Vh) + tensor(1.0486e-06) + + >>> A = torch.randn(7, 5, 3) + >>> U, S, Vh = torch.linalg.svd(A, full_matrices=False) + >>> torch.dist(A, U @ torch.diag_embed(S) @ Vh) + tensor(3.0957e-06) + +.. _condition number: + https://pytorch.org/docs/main/linalg.html#torch.linalg.cond +.. _the resulting vectors will span the same subspace: + https://en.wikipedia.org/wiki/Singular_value_decomposition#Singular_values,_singular_vectors,_and_their_relation_to_the_SVD +""") + +svdvals = _add_docstr(_linalg.linalg_svdvals, r""" +linalg.svdvals(A, *, driver=None, out=None) -> Tensor + +Computes the singular values of a matrix. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +The singular values are returned in descending order. + +.. note:: This function is equivalent to NumPy's `linalg.svd(A, compute_uv=False)`. + +""" + fr""" +.. note:: {common_notes["sync_note"]} +""" + r""" + +.. seealso:: + + :func:`torch.linalg.svd` computes the full singular value decomposition. + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + +Keyword args: + driver (str, optional): name of the cuSOLVER method to be used. This keyword argument only works on CUDA inputs. + Available options are: `None`, `gesvd`, `gesvdj`, and `gesvda`. + Check :func:`torch.linalg.svd` for details. + Default: `None`. + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Returns: + A real-valued tensor, even when :attr:`A` is complex. + +Examples:: + + >>> A = torch.randn(5, 3) + >>> S = torch.linalg.svdvals(A) + >>> S + tensor([2.5139, 2.1087, 1.1066]) + + >>> torch.dist(S, torch.linalg.svd(A, full_matrices=False).S) + tensor(2.4576e-07) +""") + +cond = _add_docstr(_linalg.linalg_cond, r""" +linalg.cond(A, p=None, *, out=None) -> Tensor + +Computes the condition number of a matrix with respect to a matrix norm. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **condition number** :math:`\kappa` of a matrix +:math:`A \in \mathbb{K}^{n \times n}` is defined as + +.. math:: + + \kappa(A) = \|A\|_p\|A^{-1}\|_p + +The condition number of :attr:`A` measures the numerical stability of the linear system `AX = B` +with respect to a matrix norm. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +:attr:`p` defines the matrix norm that is computed. The following norms are supported: + +========= ================================= +:attr:`p` matrix norm +========= ================================= +`None` `2`-norm (largest singular value) +`'fro'` Frobenius norm +`'nuc'` nuclear norm +`inf` `max(sum(abs(x), dim=1))` +`-inf` `min(sum(abs(x), dim=1))` +`1` `max(sum(abs(x), dim=0))` +`-1` `min(sum(abs(x), dim=0))` +`2` largest singular value +`-2` smallest singular value +========= ================================= + +where `inf` refers to `float('inf')`, NumPy's `inf` object, or any equivalent object. + +For :attr:`p` is one of `('fro', 'nuc', inf, -inf, 1, -1)`, this function uses +:func:`torch.linalg.norm` and :func:`torch.linalg.inv`. +As such, in this case, the matrix (or every matrix in the batch) :attr:`A` has to be square +and invertible. + +For :attr:`p` in `(2, -2)`, this function can be computed in terms of the singular values +:math:`\sigma_1 \geq \ldots \geq \sigma_n` + +.. math:: + + \kappa_2(A) = \frac{\sigma_1}{\sigma_n}\qquad \kappa_{-2}(A) = \frac{\sigma_n}{\sigma_1} + +In these cases, it is computed using :func:`torch.linalg.svdvals`. For these norms, the matrix +(or every matrix in the batch) :attr:`A` may have any shape. + +.. note :: When inputs are on a CUDA device, this function synchronizes that device with the CPU + if :attr:`p` is one of `('fro', 'nuc', inf, -inf, 1, -1)`. + +.. seealso:: + + :func:`torch.linalg.solve` for a function that solves linear systems of square matrices. + + :func:`torch.linalg.lstsq` for a function that solves linear systems of general matrices. + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions + for :attr:`p` in `(2, -2)`, and of shape `(*, n, n)` where every matrix + is invertible for :attr:`p` in `('fro', 'nuc', inf, -inf, 1, -1)`. + p (int, inf, -inf, 'fro', 'nuc', optional): + the type of the matrix norm to use in the computations (see above). Default: `None` + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Returns: + A real-valued tensor, even when :attr:`A` is complex. + +Raises: + RuntimeError: + if :attr:`p` is one of `('fro', 'nuc', inf, -inf, 1, -1)` + and the :attr:`A` matrix or any matrix in the batch :attr:`A` is not square + or invertible. + +Examples:: + + >>> A = torch.randn(3, 4, 4, dtype=torch.complex64) + >>> torch.linalg.cond(A) + >>> A = torch.tensor([[1., 0, -1], [0, 1, 0], [1, 0, 1]]) + >>> torch.linalg.cond(A) + tensor([1.4142]) + >>> torch.linalg.cond(A, 'fro') + tensor(3.1623) + >>> torch.linalg.cond(A, 'nuc') + tensor(9.2426) + >>> torch.linalg.cond(A, float('inf')) + tensor(2.) + >>> torch.linalg.cond(A, float('-inf')) + tensor(1.) + >>> torch.linalg.cond(A, 1) + tensor(2.) + >>> torch.linalg.cond(A, -1) + tensor(1.) + >>> torch.linalg.cond(A, 2) + tensor([1.4142]) + >>> torch.linalg.cond(A, -2) + tensor([0.7071]) + + >>> A = torch.randn(2, 3, 3) + >>> torch.linalg.cond(A) + tensor([[9.5917], + [3.2538]]) + >>> A = torch.randn(2, 3, 3, dtype=torch.complex64) + >>> torch.linalg.cond(A) + tensor([[4.6245], + [4.5671]]) +""") + +pinv = _add_docstr(_linalg.linalg_pinv, r""" +linalg.pinv(A, *, atol=None, rtol=None, hermitian=False, out=None) -> Tensor + +Computes the pseudoinverse (Moore-Penrose inverse) of a matrix. + +The pseudoinverse may be `defined algebraically`_ +but it is more computationally convenient to understand it `through the SVD`_ + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +If :attr:`hermitian`\ `= True`, :attr:`A` is assumed to be Hermitian if complex or +symmetric if real, but this is not checked internally. Instead, just the lower +triangular part of the matrix is used in the computations. + +The singular values (or the norm of the eigenvalues when :attr:`hermitian`\ `= True`) +that are below :math:`\max(\text{atol}, \sigma_1 \cdot \text{rtol})` threshold are +treated as zero and discarded in the computation, +where :math:`\sigma_1` is the largest singular value (or eigenvalue). + +If :attr:`rtol` is not specified and :attr:`A` is a matrix of dimensions `(m, n)`, +the relative tolerance is set to be :math:`\text{rtol} = \max(m, n) \varepsilon` +and :math:`\varepsilon` is the epsilon value for the dtype of :attr:`A` (see :class:`.finfo`). +If :attr:`rtol` is not specified and :attr:`atol` is specified to be larger than zero then +:attr:`rtol` is set to zero. + +If :attr:`atol` or :attr:`rtol` is a :class:`torch.Tensor`, its shape must be broadcastable to that +of the singular values of :attr:`A` as returned by :func:`torch.linalg.svd`. + +.. note:: This function uses :func:`torch.linalg.svd` if :attr:`hermitian`\ `= False` and + :func:`torch.linalg.eigh` if :attr:`hermitian`\ `= True`. + For CUDA inputs, this function synchronizes that device with the CPU. + +.. note:: + Consider using :func:`torch.linalg.lstsq` if possible for multiplying a matrix on the left by + the pseudoinverse, as:: + + torch.linalg.lstsq(A, B).solution == A.pinv() @ B + + It is always preferred to use :func:`~lstsq` when possible, as it is faster and more + numerically stable than computing the pseudoinverse explicitly. + +.. note:: + This function has NumPy compatible variant `linalg.pinv(A, rcond, hermitian=False)`. + However, use of the positional argument :attr:`rcond` is deprecated in favor of :attr:`rtol`. + +.. warning:: + This function uses internally :func:`torch.linalg.svd` (or :func:`torch.linalg.eigh` + when :attr:`hermitian`\ `= True`), so its derivative has the same problems as those of these + functions. See the warnings in :func:`torch.linalg.svd` and :func:`torch.linalg.eigh` for + more details. + +.. seealso:: + + :func:`torch.linalg.inv` computes the inverse of a square matrix. + + :func:`torch.linalg.lstsq` computes :attr:`A`\ `.pinv() @ \ `:attr:`B` with a + numerically stable algorithm. + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + rcond (float, Tensor, optional): [NumPy Compat]. Alias for :attr:`rtol`. Default: `None`. + +Keyword args: + atol (float, Tensor, optional): the absolute tolerance value. When `None` it's considered to be zero. + Default: `None`. + rtol (float, Tensor, optional): the relative tolerance value. See above for the value it takes when `None`. + Default: `None`. + hermitian(bool, optional): indicates whether :attr:`A` is Hermitian if complex + or symmetric if real. Default: `False`. + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Examples:: + + >>> A = torch.randn(3, 5) + >>> A + tensor([[ 0.5495, 0.0979, -1.4092, -0.1128, 0.4132], + [-1.1143, -0.3662, 0.3042, 1.6374, -0.9294], + [-0.3269, -0.5745, -0.0382, -0.5922, -0.6759]]) + >>> torch.linalg.pinv(A) + tensor([[ 0.0600, -0.1933, -0.2090], + [-0.0903, -0.0817, -0.4752], + [-0.7124, -0.1631, -0.2272], + [ 0.1356, 0.3933, -0.5023], + [-0.0308, -0.1725, -0.5216]]) + + >>> A = torch.randn(2, 6, 3) + >>> Apinv = torch.linalg.pinv(A) + >>> torch.dist(Apinv @ A, torch.eye(3)) + tensor(8.5633e-07) + + >>> A = torch.randn(3, 3, dtype=torch.complex64) + >>> A = A + A.T.conj() # creates a Hermitian matrix + >>> Apinv = torch.linalg.pinv(A, hermitian=True) + >>> torch.dist(Apinv @ A, torch.eye(3)) + tensor(1.0830e-06) + +.. _defined algebraically: + https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse#Existence_and_uniqueness +.. _through the SVD: + https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse#Singular_value_decomposition_(SVD) +""") + +matrix_exp = _add_docstr(_linalg.linalg_matrix_exp, r""" +linalg.matrix_exp(A) -> Tensor + +Computes the matrix exponential of a square matrix. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +this function computes the **matrix exponential** of :math:`A \in \mathbb{K}^{n \times n}`, which is defined as + +.. math:: + \mathrm{matrix\_exp}(A) = \sum_{k=0}^\infty \frac{1}{k!}A^k \in \mathbb{K}^{n \times n}. + +If the matrix :math:`A` has eigenvalues :math:`\lambda_i \in \mathbb{C}`, +the matrix :math:`\mathrm{matrix\_exp}(A)` has eigenvalues :math:`e^{\lambda_i} \in \mathbb{C}`. + +Supports input of bfloat16, float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions. + +Example:: + + >>> A = torch.empty(2, 2, 2) + >>> A[0, :, :] = torch.eye(2, 2) + >>> A[1, :, :] = 2 * torch.eye(2, 2) + >>> A + tensor([[[1., 0.], + [0., 1.]], + + [[2., 0.], + [0., 2.]]]) + >>> torch.linalg.matrix_exp(A) + tensor([[[2.7183, 0.0000], + [0.0000, 2.7183]], + + [[7.3891, 0.0000], + [0.0000, 7.3891]]]) + + >>> import math + >>> A = torch.tensor([[0, math.pi/3], [-math.pi/3, 0]]) # A is skew-symmetric + >>> torch.linalg.matrix_exp(A) # matrix_exp(A) = [[cos(pi/3), sin(pi/3)], [-sin(pi/3), cos(pi/3)]] + tensor([[ 0.5000, 0.8660], + [-0.8660, 0.5000]]) +""") + + +solve = _add_docstr(_linalg.linalg_solve, r""" +linalg.solve(A, B, *, left=True, out=None) -> Tensor + +Computes the solution of a square system of linear equations with a unique solution. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +this function computes the solution :math:`X \in \mathbb{K}^{n \times k}` of the **linear system** associated to +:math:`A \in \mathbb{K}^{n \times n}, B \in \mathbb{K}^{n \times k}`, which is defined as + +.. math:: AX = B + +If :attr:`left`\ `= False`, this function returns the matrix :math:`X \in \mathbb{K}^{n \times k}` that solves the system + +.. math:: + + XA = B\mathrlap{\qquad A \in \mathbb{K}^{k \times k}, B \in \mathbb{K}^{n \times k}.} + +This system of linear equations has one solution if and only if :math:`A` is `invertible`_. +This function assumes that :math:`A` is invertible. + +Supports inputs of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if the inputs are batches of matrices then +the output has the same batch dimensions. + +Letting `*` be zero or more batch dimensions, + +- If :attr:`A` has shape `(*, n, n)` and :attr:`B` has shape `(*, n)` (a batch of vectors) or shape + `(*, n, k)` (a batch of matrices or "multiple right-hand sides"), this function returns `X` of shape + `(*, n)` or `(*, n, k)` respectively. +- Otherwise, if :attr:`A` has shape `(*, n, n)` and :attr:`B` has shape `(n,)` or `(n, k)`, :attr:`B` + is broadcasted to have shape `(*, n)` or `(*, n, k)` respectively. + This function then returns the solution of the resulting batch of systems of linear equations. + +.. note:: + This function computes `X = \ `:attr:`A`\ `.inverse() @ \ `:attr:`B` in a faster and + more numerically stable way than performing the computations separately. + +.. note:: + It is possible to compute the solution of the system :math:`XA = B` by passing the inputs + :attr:`A` and :attr:`B` transposed and transposing the output returned by this function. + +.. note:: + :attr:`A` is allowed to be a non-batched `torch.sparse_csr_tensor`, but only with `left=True`. + +""" + fr""" +.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.solve_ex")} +""" + r""" + +.. seealso:: + + :func:`torch.linalg.solve_triangular` computes the solution of a triangular system of linear + equations with a unique solution. + +Args: + A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions. + B (Tensor): right-hand side tensor of shape `(*, n)` or `(*, n, k)` or `(n,)` or `(n, k)` + according to the rules described above + +Keyword args: + left (bool, optional): whether to solve the system :math:`AX=B` or :math:`XA = B`. Default: `True`. + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Raises: + RuntimeError: if the :attr:`A` matrix is not invertible or any matrix in a batched :attr:`A` + is not invertible. + +Examples:: + + >>> A = torch.randn(3, 3) + >>> b = torch.randn(3) + >>> x = torch.linalg.solve(A, b) + >>> torch.allclose(A @ x, b) + True + >>> A = torch.randn(2, 3, 3) + >>> B = torch.randn(2, 3, 4) + >>> X = torch.linalg.solve(A, B) + >>> X.shape + torch.Size([2, 3, 4]) + >>> torch.allclose(A @ X, B) + True + + >>> A = torch.randn(2, 3, 3) + >>> b = torch.randn(3, 1) + >>> x = torch.linalg.solve(A, b) # b is broadcasted to size (2, 3, 1) + >>> x.shape + torch.Size([2, 3, 1]) + >>> torch.allclose(A @ x, b) + True + >>> b = torch.randn(3) + >>> x = torch.linalg.solve(A, b) # b is broadcasted to size (2, 3) + >>> x.shape + torch.Size([2, 3]) + >>> Ax = A @ x.unsqueeze(-1) + >>> torch.allclose(Ax, b.unsqueeze(-1).expand_as(Ax)) + True + +.. _invertible: + https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem +""") + +solve_triangular = _add_docstr(_linalg.linalg_solve_triangular, r""" +linalg.solve_triangular(A, B, *, upper, left=True, unitriangular=False, out=None) -> Tensor + +Computes the solution of a triangular system of linear equations with a unique solution. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +this function computes the solution :math:`X \in \mathbb{K}^{n \times k}` of the **linear system** +associated to the triangular matrix :math:`A \in \mathbb{K}^{n \times n}` without zeros on the diagonal +(that is, it is `invertible`_) and the rectangular matrix , :math:`B \in \mathbb{K}^{n \times k}`, +which is defined as + +.. math:: AX = B + +The argument :attr:`upper` signals whether :math:`A` is upper or lower triangular. + +If :attr:`left`\ `= False`, this function returns the matrix :math:`X \in \mathbb{K}^{n \times k}` that +solves the system + +.. math:: + + XA = B\mathrlap{\qquad A \in \mathbb{K}^{k \times k}, B \in \mathbb{K}^{n \times k}.} + +If :attr:`upper`\ `= True` (resp. `False`) just the upper (resp. lower) triangular half of :attr:`A` +will be accessed. The elements below the main diagonal will be considered to be zero and will not be accessed. + +If :attr:`unitriangular`\ `= True`, the diagonal of :attr:`A` is assumed to be ones and will not be accessed. + +The result may contain `NaN` s if the diagonal of :attr:`A` contains zeros or elements that +are very close to zero and :attr:`unitriangular`\ `= False` (default) or if the input matrix +has very small eigenvalues. + +Supports inputs of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if the inputs are batches of matrices then +the output has the same batch dimensions. + +.. seealso:: + + :func:`torch.linalg.solve` computes the solution of a general square system of linear + equations with a unique solution. + +Args: + A (Tensor): tensor of shape `(*, n, n)` (or `(*, k, k)` if :attr:`left`\ `= False`) + where `*` is zero or more batch dimensions. + B (Tensor): right-hand side tensor of shape `(*, n, k)`. + +Keyword args: + upper (bool): whether :attr:`A` is an upper or lower triangular matrix. + left (bool, optional): whether to solve the system :math:`AX=B` or :math:`XA = B`. Default: `True`. + unitriangular (bool, optional): if `True`, the diagonal elements of :attr:`A` are assumed to be + all equal to `1`. Default: `False`. + out (Tensor, optional): output tensor. `B` may be passed as `out` and the result is computed in-place on `B`. + Ignored if `None`. Default: `None`. + +Examples:: + + >>> A = torch.randn(3, 3).triu_() + >>> B = torch.randn(3, 4) + >>> X = torch.linalg.solve_triangular(A, B, upper=True) + >>> torch.allclose(A @ X, B) + True + + >>> A = torch.randn(2, 3, 3).tril_() + >>> B = torch.randn(2, 3, 4) + >>> X = torch.linalg.solve_triangular(A, B, upper=False) + >>> torch.allclose(A @ X, B) + True + + >>> A = torch.randn(2, 4, 4).tril_() + >>> B = torch.randn(2, 3, 4) + >>> X = torch.linalg.solve_triangular(A, B, upper=False, left=False) + >>> torch.allclose(X @ A, B) + True + +.. _invertible: + https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem +""") + +lu_factor = _add_docstr(_linalg.linalg_lu_factor, r""" +linalg.lu_factor(A, *, bool pivot=True, out=None) -> (Tensor, Tensor) + +Computes a compact representation of the LU factorization with partial pivoting of a matrix. + +This function computes a compact representation of the decomposition given by :func:`torch.linalg.lu`. +If the matrix is square, this representation may be used in :func:`torch.linalg.lu_solve` +to solve system of linear equations that share the matrix :attr:`A`. + +The returned decomposition is represented as a named tuple `(LU, pivots)`. +The ``LU`` matrix has the same shape as the input matrix ``A``. Its upper and lower triangular +parts encode the non-constant elements of ``L`` and ``U`` of the LU decomposition of ``A``. + +The returned permutation matrix is represented by a 1-indexed vector. `pivots[i] == j` represents +that in the `i`-th step of the algorithm, the `i`-th row was permuted with the `j-1`-th row. + +On CUDA, one may use :attr:`pivot`\ `= False`. In this case, this function returns the LU +decomposition without pivoting if it exists. + +Supports inputs of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if the inputs are batches of matrices then +the output has the same batch dimensions. + +""" + fr""" +.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.lu_factor_ex")} +""" + r""" +.. warning:: The LU decomposition is almost never unique, as often there are different permutation + matrices that can yield different LU decompositions. + As such, different platforms, like SciPy, or inputs on different devices, + may produce different valid decompositions. + + Gradient computations are only supported if the input matrix is full-rank. + If this condition is not met, no error will be thrown, but the gradient may not be finite. + This is because the LU decomposition with pivoting is not differentiable at these points. + +.. seealso:: + + :func:`torch.linalg.lu_solve` solves a system of linear equations given the output of this + function provided the input matrix was square and invertible. + + :func:`torch.lu_unpack` unpacks the tensors returned by :func:`~lu_factor` into the three + matrices `P, L, U` that form the decomposition. + + :func:`torch.linalg.lu` computes the LU decomposition with partial pivoting of a possibly + non-square matrix. It is a composition of :func:`~lu_factor` and :func:`torch.lu_unpack`. + + :func:`torch.linalg.solve` solves a system of linear equations. It is a composition + of :func:`~lu_factor` and :func:`~lu_solve`. + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + +Keyword args: + pivot (bool, optional): Whether to compute the LU decomposition with partial pivoting, or the regular LU + decomposition. :attr:`pivot`\ `= False` not supported on CPU. Default: `True`. + out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(LU, pivots)`. + +Raises: + RuntimeError: if the :attr:`A` matrix is not invertible or any matrix in a batched :attr:`A` + is not invertible. + +Examples:: + + >>> A = torch.randn(2, 3, 3) + >>> B1 = torch.randn(2, 3, 4) + >>> B2 = torch.randn(2, 3, 7) + >>> LU, pivots = torch.linalg.lu_factor(A) + >>> X1 = torch.linalg.lu_solve(LU, pivots, B1) + >>> X2 = torch.linalg.lu_solve(LU, pivots, B2) + >>> torch.allclose(A @ X1, B1) + True + >>> torch.allclose(A @ X2, B2) + True + +.. _invertible: + https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem +""") + +lu_factor_ex = _add_docstr(_linalg.linalg_lu_factor_ex, r""" +linalg.lu_factor_ex(A, *, pivot=True, check_errors=False, out=None) -> (Tensor, Tensor, Tensor) + +This is a version of :func:`~lu_factor` that does not perform error checks unless :attr:`check_errors`\ `= True`. +It also returns the :attr:`info` tensor returned by `LAPACK's getrf`_. + +""" + fr""" +.. note:: {common_notes["sync_note_ex"]} + +.. warning:: {common_notes["experimental_warning"]} +""" + r""" + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + +Keyword args: + pivot (bool, optional): Whether to compute the LU decomposition with partial pivoting, or the regular LU + decomposition. :attr:`pivot`\ `= False` not supported on CPU. Default: `True`. + check_errors (bool, optional): controls whether to check the content of ``infos`` and raise + an error if it is non-zero. Default: `False`. + out (tuple, optional): tuple of three tensors to write the output to. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(LU, pivots, info)`. + +.. _LAPACK's getrf: + https://www.netlib.org/lapack/explore-html/dd/d9a/group__double_g_ecomputational_ga0019443faea08275ca60a734d0593e60.html +""") + +lu_solve = _add_docstr(_linalg.linalg_lu_solve, r""" +linalg.lu_solve(LU, pivots, B, *, left=True, adjoint=False, out=None) -> Tensor + +Computes the solution of a square system of linear equations with a unique solution given an LU decomposition. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +this function computes the solution :math:`X \in \mathbb{K}^{n \times k}` of the **linear system** associated to +:math:`A \in \mathbb{K}^{n \times n}, B \in \mathbb{K}^{n \times k}`, which is defined as + +.. math:: AX = B + +where :math:`A` is given factorized as returned by :func:`~lu_factor`. + +If :attr:`left`\ `= False`, this function returns the matrix :math:`X \in \mathbb{K}^{n \times k}` that solves the system + +.. math:: + + XA = B\mathrlap{\qquad A \in \mathbb{K}^{k \times k}, B \in \mathbb{K}^{n \times k}.} + +If :attr:`adjoint`\ `= True` (and :attr:`left`\ `= True`), given an LU factorization of :math:`A` +this function function returns the :math:`X \in \mathbb{K}^{n \times k}` that solves the system + +.. math:: + + A^{\text{H}}X = B\mathrlap{\qquad A \in \mathbb{K}^{k \times k}, B \in \mathbb{K}^{n \times k}.} + +where :math:`A^{\text{H}}` is the conjugate transpose when :math:`A` is complex, and the +transpose when :math:`A` is real-valued. The :attr:`left`\ `= False` case is analogous. + +Supports inputs of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if the inputs are batches of matrices then +the output has the same batch dimensions. + +Args: + LU (Tensor): tensor of shape `(*, n, n)` (or `(*, k, k)` if :attr:`left`\ `= True`) + where `*` is zero or more batch dimensions as returned by :func:`~lu_factor`. + pivots (Tensor): tensor of shape `(*, n)` (or `(*, k)` if :attr:`left`\ `= True`) + where `*` is zero or more batch dimensions as returned by :func:`~lu_factor`. + B (Tensor): right-hand side tensor of shape `(*, n, k)`. + +Keyword args: + left (bool, optional): whether to solve the system :math:`AX=B` or :math:`XA = B`. Default: `True`. + adjoint (bool, optional): whether to solve the system :math:`AX=B` or :math:`A^{\text{H}}X = B`. Default: `False`. + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Examples:: + + >>> A = torch.randn(3, 3) + >>> LU, pivots = torch.linalg.lu_factor(A) + >>> B = torch.randn(3, 2) + >>> X = torch.linalg.lu_solve(LU, pivots, B) + >>> torch.allclose(A @ X, B) + True + + >>> B = torch.randn(3, 3, 2) # Broadcasting rules apply: A is broadcasted + >>> X = torch.linalg.lu_solve(LU, pivots, B) + >>> torch.allclose(A @ X, B) + True + + >>> B = torch.randn(3, 5, 3) + >>> X = torch.linalg.lu_solve(LU, pivots, B, left=False) + >>> torch.allclose(X @ A, B) + True + + >>> B = torch.randn(3, 3, 4) # Now solve for A^T + >>> X = torch.linalg.lu_solve(LU, pivots, B, adjoint=True) + >>> torch.allclose(A.mT @ X, B) + True + +.. _invertible: + https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem +""") + +lu = _add_docstr(_linalg.linalg_lu, r""" +lu(A, *, pivot=True, out=None) -> (Tensor, Tensor, Tensor) + +Computes the LU decomposition with partial pivoting of a matrix. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **LU decomposition with partial pivoting** of a matrix +:math:`A \in \mathbb{K}^{m \times n}` is defined as + +.. math:: + + A = PLU\mathrlap{\qquad P \in \mathbb{K}^{m \times m}, L \in \mathbb{K}^{m \times k}, U \in \mathbb{K}^{k \times n}} + +where `k = min(m,n)`, :math:`P` is a `permutation matrix`_, :math:`L` is lower triangular with ones on the diagonal +and :math:`U` is upper triangular. + +If :attr:`pivot`\ `= False` and :attr:`A` is on GPU, then the **LU decomposition without pivoting** is computed + +.. math:: + + A = LU\mathrlap{\qquad L \in \mathbb{K}^{m \times k}, U \in \mathbb{K}^{k \times n}} + +When :attr:`pivot`\ `= False`, the returned matrix :attr:`P` will be empty. +The LU decomposition without pivoting `may not exist`_ if any of the principal minors of :attr:`A` is singular. +In this case, the output matrix may contain `inf` or `NaN`. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +.. seealso:: + + :func:`torch.linalg.solve` solves a system of linear equations using the LU decomposition + with partial pivoting. + +.. warning:: The LU decomposition is almost never unique, as often there are different permutation + matrices that can yield different LU decompositions. + As such, different platforms, like SciPy, or inputs on different devices, + may produce different valid decompositions. + +.. warning:: Gradient computations are only supported if the input matrix is full-rank. + If this condition is not met, no error will be thrown, but the gradient + may not be finite. + This is because the LU decomposition with pivoting is not differentiable at these points. + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + pivot (bool, optional): Controls whether to compute the LU decomposition with partial pivoting or + no pivoting. Default: `True`. + +Keyword args: + out (tuple, optional): output tuple of three tensors. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(P, L, U)`. + +Examples:: + + >>> A = torch.randn(3, 2) + >>> P, L, U = torch.linalg.lu(A) + >>> P + tensor([[0., 1., 0.], + [0., 0., 1.], + [1., 0., 0.]]) + >>> L + tensor([[1.0000, 0.0000], + [0.5007, 1.0000], + [0.0633, 0.9755]]) + >>> U + tensor([[0.3771, 0.0489], + [0.0000, 0.9644]]) + >>> torch.dist(A, P @ L @ U) + tensor(5.9605e-08) + + >>> A = torch.randn(2, 5, 7, device="cuda") + >>> P, L, U = torch.linalg.lu(A, pivot=False) + >>> P + tensor([], device='cuda:0') + >>> torch.dist(A, L @ U) + tensor(1.0376e-06, device='cuda:0') + +.. _permutation matrix: + https://en.wikipedia.org/wiki/Permutation_matrix +.. _may not exist: + https://en.wikipedia.org/wiki/LU_decomposition#Definitions +""") + +tensorinv = _add_docstr(_linalg.linalg_tensorinv, r""" +linalg.tensorinv(A, ind=2, *, out=None) -> Tensor + +Computes the multiplicative inverse of :func:`torch.tensordot`. + +If `m` is the product of the first :attr:`ind` dimensions of :attr:`A` and `n` is the product of +the rest of the dimensions, this function expects `m` and `n` to be equal. +If this is the case, it computes a tensor `X` such that +`tensordot(\ `:attr:`A`\ `, X, \ `:attr:`ind`\ `)` is the identity matrix in dimension `m`. +`X` will have the shape of :attr:`A` but with the first :attr:`ind` dimensions pushed back to the end + +.. code:: text + + X.shape == A.shape[ind:] + A.shape[:ind] + +Supports input of float, double, cfloat and cdouble dtypes. + +.. note:: When :attr:`A` is a `2`-dimensional tensor and :attr:`ind`\ `= 1`, + this function computes the (multiplicative) inverse of :attr:`A` + (see :func:`torch.linalg.inv`). + +.. note:: + Consider using :func:`torch.linalg.tensorsolve` if possible for multiplying a tensor on the left + by the tensor inverse, as:: + + linalg.tensorsolve(A, B) == torch.tensordot(linalg.tensorinv(A), B) # When B is a tensor with shape A.shape[:B.ndim] + + It is always preferred to use :func:`~tensorsolve` when possible, as it is faster and more + numerically stable than computing the pseudoinverse explicitly. + +.. seealso:: + + :func:`torch.linalg.tensorsolve` computes + `torch.tensordot(tensorinv(\ `:attr:`A`\ `), \ `:attr:`B`\ `)`. + +Args: + A (Tensor): tensor to invert. Its shape must satisfy + `prod(\ `:attr:`A`\ `.shape[:\ `:attr:`ind`\ `]) == + prod(\ `:attr:`A`\ `.shape[\ `:attr:`ind`\ `:])`. + ind (int): index at which to compute the inverse of :func:`torch.tensordot`. Default: `2`. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Raises: + RuntimeError: if the reshaped :attr:`A` is not invertible or the product of the first + :attr:`ind` dimensions is not equal to the product of the rest. + +Examples:: + + >>> A = torch.eye(4 * 6).reshape((4, 6, 8, 3)) + >>> Ainv = torch.linalg.tensorinv(A, ind=2) + >>> Ainv.shape + torch.Size([8, 3, 4, 6]) + >>> B = torch.randn(4, 6) + >>> torch.allclose(torch.tensordot(Ainv, B), torch.linalg.tensorsolve(A, B)) + True + + >>> A = torch.randn(4, 4) + >>> Atensorinv = torch.linalg.tensorinv(A, ind=1) + >>> Ainv = torch.linalg.inv(A) + >>> torch.allclose(Atensorinv, Ainv) + True +""") + +tensorsolve = _add_docstr(_linalg.linalg_tensorsolve, r""" +linalg.tensorsolve(A, B, dims=None, *, out=None) -> Tensor + +Computes the solution `X` to the system `torch.tensordot(A, X) = B`. + +If `m` is the product of the first :attr:`B`\ `.ndim` dimensions of :attr:`A` and +`n` is the product of the rest of the dimensions, this function expects `m` and `n` to be equal. + +The returned tensor `x` satisfies +`tensordot(\ `:attr:`A`\ `, x, dims=x.ndim) == \ `:attr:`B`. +`x` has shape :attr:`A`\ `[B.ndim:]`. + +If :attr:`dims` is specified, :attr:`A` will be reshaped as + +.. code:: text + + A = movedim(A, dims, range(len(dims) - A.ndim + 1, 0)) + +Supports inputs of float, double, cfloat and cdouble dtypes. + +.. seealso:: + + :func:`torch.linalg.tensorinv` computes the multiplicative inverse of + :func:`torch.tensordot`. + +Args: + A (Tensor): tensor to solve for. Its shape must satisfy + `prod(\ `:attr:`A`\ `.shape[:\ `:attr:`B`\ `.ndim]) == + prod(\ `:attr:`A`\ `.shape[\ `:attr:`B`\ `.ndim:])`. + B (Tensor): tensor of shape :attr:`A`\ `.shape[:\ `:attr:`B`\ `.ndim]`. + dims (Tuple[int], optional): dimensions of :attr:`A` to be moved. + If `None`, no dimensions are moved. Default: `None`. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Raises: + RuntimeError: if the reshaped :attr:`A`\ `.view(m, m)` with `m` as above is not + invertible or the product of the first :attr:`ind` dimensions is not equal + to the product of the rest of the dimensions. + +Examples:: + + >>> A = torch.eye(2 * 3 * 4).reshape((2 * 3, 4, 2, 3, 4)) + >>> B = torch.randn(2 * 3, 4) + >>> X = torch.linalg.tensorsolve(A, B) + >>> X.shape + torch.Size([2, 3, 4]) + >>> torch.allclose(torch.tensordot(A, X, dims=X.ndim), B) + True + + >>> A = torch.randn(6, 4, 4, 3, 2) + >>> B = torch.randn(4, 3, 2) + >>> X = torch.linalg.tensorsolve(A, B, dims=(0, 2)) + >>> X.shape + torch.Size([6, 4]) + >>> A = A.permute(1, 3, 4, 0, 2) + >>> A.shape[B.ndim:] + torch.Size([6, 4]) + >>> torch.allclose(torch.tensordot(A, X, dims=X.ndim), B, atol=1e-6) + True +""") + +qr = _add_docstr(_linalg.linalg_qr, r""" +qr(A, mode='reduced', *, out=None) -> (Tensor, Tensor) + +Computes the QR decomposition of a matrix. + +Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, +the **full QR decomposition** of a matrix +:math:`A \in \mathbb{K}^{m \times n}` is defined as + +.. math:: + + A = QR\mathrlap{\qquad Q \in \mathbb{K}^{m \times m}, R \in \mathbb{K}^{m \times n}} + +where :math:`Q` is orthogonal in the real case and unitary in the complex case, +and :math:`R` is upper triangular with real diagonal (even in the complex case). + +When `m > n` (tall matrix), as `R` is upper triangular, its last `m - n` rows are zero. +In this case, we can drop the last `m - n` columns of `Q` to form the +**reduced QR decomposition**: + +.. math:: + + A = QR\mathrlap{\qquad Q \in \mathbb{K}^{m \times n}, R \in \mathbb{K}^{n \times n}} + +The reduced QR decomposition agrees with the full QR decomposition when `n >= m` (wide matrix). + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :attr:`A` is a batch of matrices then +the output has the same batch dimensions. + +The parameter :attr:`mode` chooses between the full and reduced QR decomposition. +If :attr:`A` has shape `(*, m, n)`, denoting `k = min(m, n)` + +- :attr:`mode`\ `= 'reduced'` (default): Returns `(Q, R)` of shapes `(*, m, k)`, `(*, k, n)` respectively. + It is always differentiable. +- :attr:`mode`\ `= 'complete'`: Returns `(Q, R)` of shapes `(*, m, m)`, `(*, m, n)` respectively. + It is differentiable for `m <= n`. +- :attr:`mode`\ `= 'r'`: Computes only the reduced `R`. Returns `(Q, R)` with `Q` empty and `R` of shape `(*, k, n)`. + It is never differentiable. + +Differences with `numpy.linalg.qr`: + +- :attr:`mode`\ `= 'raw'` is not implemented. +- Unlike `numpy.linalg.qr`, this function always returns a tuple of two tensors. + When :attr:`mode`\ `= 'r'`, the `Q` tensor is an empty tensor. + +.. warning:: The elements in the diagonal of `R` are not necessarily positive. + As such, the returned QR decomposition is only unique up to the sign of the diagonal of `R`. + Therefore, different platforms, like NumPy, or inputs on different devices, + may produce different valid decompositions. + +.. warning:: The QR decomposition is only well-defined if the first `k = min(m, n)` columns + of every matrix in :attr:`A` are linearly independent. + If this condition is not met, no error will be thrown, but the QR produced + may be incorrect and its autodiff may fail or produce incorrect results. + +Args: + A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + mode (str, optional): one of `'reduced'`, `'complete'`, `'r'`. + Controls the shape of the returned tensors. Default: `'reduced'`. + +Keyword args: + out (tuple, optional): output tuple of two tensors. Ignored if `None`. Default: `None`. + +Returns: + A named tuple `(Q, R)`. + +Examples:: + + >>> A = torch.tensor([[12., -51, 4], [6, 167, -68], [-4, 24, -41]]) + >>> Q, R = torch.linalg.qr(A) + >>> Q + tensor([[-0.8571, 0.3943, 0.3314], + [-0.4286, -0.9029, -0.0343], + [ 0.2857, -0.1714, 0.9429]]) + >>> R + tensor([[ -14.0000, -21.0000, 14.0000], + [ 0.0000, -175.0000, 70.0000], + [ 0.0000, 0.0000, -35.0000]]) + >>> (Q @ R).round() + tensor([[ 12., -51., 4.], + [ 6., 167., -68.], + [ -4., 24., -41.]]) + >>> (Q.T @ Q).round() + tensor([[ 1., 0., 0.], + [ 0., 1., -0.], + [ 0., -0., 1.]]) + >>> Q2, R2 = torch.linalg.qr(A, mode='r') + >>> Q2 + tensor([]) + >>> torch.equal(R, R2) + True + >>> A = torch.randn(3, 4, 5) + >>> Q, R = torch.linalg.qr(A, mode='complete') + >>> torch.dist(Q @ R, A) + tensor(1.6099e-06) + >>> torch.dist(Q.mT @ Q, torch.eye(4)) + tensor(6.2158e-07) +""") + +vander = _add_docstr(_linalg.linalg_vander, r""" +vander(x, N=None) -> Tensor + +Generates a Vandermonde matrix. + +Returns the Vandermonde matrix :math:`V` + +.. math:: + + V = \begin{pmatrix} + 1 & x_1 & x_1^2 & \dots & x_1^{N-1}\\ + 1 & x_2 & x_2^2 & \dots & x_2^{N-1}\\ + 1 & x_3 & x_3^2 & \dots & x_3^{N-1}\\ + \vdots & \vdots & \vdots & \ddots &\vdots \\ + 1 & x_n & x_n^2 & \dots & x_n^{N-1} + \end{pmatrix}. + +for `N > 1`. +If :attr:`N`\ `= None`, then `N = x.size(-1)` so that the output is a square matrix. + +Supports inputs of float, double, cfloat, cdouble, and integral dtypes. +Also supports batches of vectors, and if :attr:`x` is a batch of vectors then +the output has the same batch dimensions. + +Differences with `numpy.vander`: + +- Unlike `numpy.vander`, this function returns the powers of :attr:`x` in ascending order. + To get them in the reverse order call ``linalg.vander(x, N).flip(-1)``. + +Args: + x (Tensor): tensor of shape `(*, n)` where `*` is zero or more batch dimensions + consisting of vectors. + +Keyword args: + N (int, optional): Number of columns in the output. Default: `x.size(-1)` + +Example:: + + >>> x = torch.tensor([1, 2, 3, 5]) + >>> linalg.vander(x) + tensor([[ 1, 1, 1, 1], + [ 1, 2, 4, 8], + [ 1, 3, 9, 27], + [ 1, 5, 25, 125]]) + >>> linalg.vander(x, N=3) + tensor([[ 1, 1, 1], + [ 1, 2, 4], + [ 1, 3, 9], + [ 1, 5, 25]]) +""") + +vecdot = _add_docstr(_linalg.linalg_vecdot, r""" +linalg.vecdot(x, y, *, dim=-1, out=None) -> Tensor + +Computes the dot product of two batches of vectors along a dimension. + +In symbols, this function computes + +.. math:: + + \sum_{i=1}^n \overline{x_i}y_i. + +over the dimension :attr:`dim` where :math:`\overline{x_i}` denotes the conjugate for complex +vectors, and it is the identity for real vectors. + +Supports input of half, bfloat16, float, double, cfloat, cdouble and integral dtypes. +It also supports broadcasting. + +Args: + x (Tensor): first batch of vectors of shape `(*, n)`. + y (Tensor): second batch of vectors of shape `(*, n)`. + +Keyword args: + dim (int): Dimension along which to compute the dot product. Default: `-1`. + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Examples:: + + >>> v1 = torch.randn(3, 2) + >>> v2 = torch.randn(3, 2) + >>> linalg.vecdot(v1, v2) + tensor([ 0.3223, 0.2815, -0.1944]) + >>> torch.vdot(v1[0], v2[0]) + tensor(0.3223) +""") diff --git a/lib/python3.12/site-packages/torch/xpu/__init__.py b/lib/python3.12/site-packages/torch/xpu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..485820cd4d2f7e92aed4877287540c2e79fac2b5 --- /dev/null +++ b/lib/python3.12/site-packages/torch/xpu/__init__.py @@ -0,0 +1,531 @@ +# mypy: allow-untyped-defs +r""" +This package introduces support for the XPU backend, specifically tailored for +Intel GPU optimization. + +This package is lazily initialized, so you can always import it, and use +:func:`is_available()` to determine if your system supports XPU. +""" +import threading +import traceback +from functools import lru_cache +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import torch +import torch._C +from torch import device as _device +from torch._utils import _dummy_type, _LazySeedTracker + +from ._utils import _get_device_index +from .streams import Event, Stream + + +_initialized = False +_tls = threading.local() +_initialization_lock = threading.Lock() +_queued_calls: List[ + Tuple[Callable[[], None], List[str]] +] = [] # don't invoke these until initialization occurs +_is_in_bad_fork = getattr(torch._C, "_xpu_isInBadFork", lambda: False) +_device_t = Union[_device, str, int, None] +_lazy_seed_tracker = _LazySeedTracker() +default_generators: Tuple[torch._C.Generator] = () # type: ignore[assignment] + + +def _is_compiled() -> bool: + r"""Return true if compile with XPU support.""" + return torch._C._has_xpu + + +if _is_compiled(): + _XpuDeviceProperties = torch._C._XpuDeviceProperties + _exchange_device = torch._C._xpu_exchangeDevice + _maybe_exchange_device = torch._C._xpu_maybeExchangeDevice +else: + # Define dummy if PyTorch was compiled without XPU + _XpuDeviceProperties = _dummy_type("_XpuDeviceProperties") # type: ignore[assignment, misc] + + def _exchange_device(device: int) -> int: + raise NotImplementedError("PyTorch was compiled without XPU support") + + def _maybe_exchange_device(device: int) -> int: + raise NotImplementedError("PyTorch was compiled without XPU support") + + +@lru_cache(maxsize=1) +def device_count() -> int: + r"""Return the number of XPU device available.""" + if not _is_compiled(): + return 0 + return torch._C._xpu_getDeviceCount() + + +def is_available() -> bool: + r"""Return a bool indicating if XPU is currently available.""" + # This function nerver throws. + return device_count() > 0 + + +def is_bf16_supported(): + r"""Return a bool indicating if the current XPU device supports dtype bfloat16.""" + return True + + +def is_initialized(): + r"""Return whether PyTorch's XPU state has been initialized.""" + return _initialized and not _is_in_bad_fork() + + +def _lazy_call(callable, **kwargs): + if is_initialized(): + callable() + else: + global _lazy_seed_tracker + if kwargs.get("seed_all", False): + _lazy_seed_tracker.queue_seed_all(callable, traceback.format_stack()) + elif kwargs.get("seed", False): + _lazy_seed_tracker.queue_seed(callable, traceback.format_stack()) + else: + # Don't store the actual traceback to avoid memory cycle + _queued_calls.append((callable, traceback.format_stack())) + + +def init(): + r"""Initialize PyTorch's XPU state. + This is a Python API about lazy initialization that avoids initializing + XPU until the first time it is accessed. Does nothing if the XPU state is + already initialized. + """ + _lazy_init() + + +def _lazy_init(): + global _initialized, _queued_calls + if is_initialized() or hasattr(_tls, "is_initializing"): + return + with _initialization_lock: + # This test was was protected via GIL. Double-check whether XPU has + # already been initialized. + if is_initialized(): + return + # Stop promptly upon encountering a bad fork error. + if _is_in_bad_fork(): + raise RuntimeError( + "Cannot re-initialize XPU in forked subprocess. To use XPU with " + "multiprocessing, you must use the 'spawn' start method" + ) + if not _is_compiled(): + raise AssertionError("Torch not compiled with XPU enabled") + # This function inits XPU backend and detects bad fork processing. + torch._C._xpu_init() + # Some of the queued calls may reentrantly call _lazy_init(); We need to + # just return without initializing in that case. + _tls.is_initializing = True + + _queued_calls.extend(calls for calls in _lazy_seed_tracker.get_calls() if calls) + + try: + for queued_call, orig_traceback in _queued_calls: + try: + queued_call() + except Exception as e: + msg = ( + f"XPU call failed lazily at initialization with error: {str(e)}\n\n" + f"XPU call was originally invoked at:\n\n{''.join(orig_traceback)}" + ) + raise Exception(msg) from e # noqa: TRY002 + finally: + delattr(_tls, "is_initializing") + _initialized = True + + +class _DeviceGuard: + def __init__(self, index: int): + self.idx = index + self.prev_idx = -1 + + def __enter__(self): + self.prev_idx = torch.xpu._exchange_device(self.idx) + + def __exit__(self, type: Any, value: Any, traceback: Any): + self.idx = torch.xpu._maybe_exchange_device(self.prev_idx) + return False + + +class device: + r"""Context-manager that changes the selected device. + + Args: + device (torch.device or int or str): device index to select. It's a no-op if + this argument is a negative integer or ``None``. + """ + + def __init__(self, device: Any): + self.idx = _get_device_index(device, optional=True) + self.prev_idx = -1 + + def __enter__(self): + self.prev_idx = torch.xpu._exchange_device(self.idx) + + def __exit__(self, type: Any, value: Any, traceback: Any): + self.idx = torch.xpu._maybe_exchange_device(self.prev_idx) + return False + + +class device_of(device): + r"""Context-manager that changes the current device to that of given object. + + You can use both tensors and storages as arguments. If a given object is + not allocated on a XPU, this is a no-op. + + Args: + obj (Tensor or Storage): object allocated on the selected device. + """ + + def __init__(self, obj): + idx = obj.get_device() if obj.is_xpu else -1 + super().__init__(idx) + + +def set_device(device: _device_t) -> None: + r"""Set the current device. + + Args: + device (torch.device or int or str): selected device. This function is a + no-op if this argument is negative. + """ + _lazy_init() + device = _get_device_index(device) + if device >= 0: + torch._C._xpu_setDevice(device) + + +def get_device_name(device: Optional[_device_t] = None) -> str: + r"""Get the name of a device. + + Args: + device (torch.device or int or str, optional): device for which to + return the name. This function is a no-op if this argument is a + negative integer. It uses the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + + Returns: + str: the name of the device + """ + return get_device_properties(device).name + + +@lru_cache(None) +def get_device_capability(device: Optional[_device_t] = None) -> Dict[str, Any]: + r"""Get the xpu capability of a device. + + Args: + device (torch.device or int or str, optional): device for which to + return the device capability. This function is a no-op if this + argument is a negative integer. It uses the current device, given by + :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` + (default). + + Returns: + Dict[str, Any]: the xpu capability dictionary of the device + """ + props = get_device_properties(device) + # pybind service attributes are no longer needed and their presence breaks + # the further logic related to the serialization of the created dictionary. + # In particular it filters out `` + # to fix Triton tests. + # This field appears after updating pybind to 2.13.6. + return { + prop: getattr(props, prop) + for prop in dir(props) + if not prop.startswith(("__", "_pybind11_")) + } + + +def get_device_properties(device: Optional[_device_t] = None) -> _XpuDeviceProperties: + r"""Get the properties of a device. + + Args: + device (torch.device or int or str): device for which to return the + properties of the device. + + Returns: + _XpuDeviceProperties: the properties of the device + """ + _lazy_init() + device = _get_device_index(device, optional=True) + if device < 0 or device >= device_count(): + raise AssertionError("Invalid device index") + return _get_device_properties(device) # type: ignore[name-defined] # noqa: F821 + + +def current_device() -> int: + r"""Return the index of a currently selected device.""" + _lazy_init() + return torch._C._xpu_getDevice() + + +def _get_device(device: Union[int, str, torch.device]) -> torch.device: + r"""Return the torch.device type object from the passed in device. + + Args: + device (torch.device or int or str): selected device. + """ + if isinstance(device, str): + device = torch.device(device) + elif isinstance(device, int): + device = torch.device("xpu", device) + return device + + +class StreamContext: + r"""Context-manager that selects a given stream. + + All XPU kernels queued within its context will be enqueued on a selected + stream. + + Args: + Stream (Stream): selected stream. This manager is a no-op if it's + ``None``. + .. note:: Streams are per-device. + """ + cur_stream: Optional["torch.xpu.Stream"] + + def __init__(self, stream: Optional["torch.xpu.Stream"]): + self.stream = stream + self.idx = _get_device_index(None, True) + if self.idx is None: + self.idx = -1 + + def __enter__(self): + cur_stream = self.stream + if cur_stream is None or self.idx == -1: + return + self.src_prev_stream = torch.xpu.current_stream(None) + + # If the stream is not on the current device, then set the current stream on the device + if self.src_prev_stream.device != cur_stream.device: + with device(cur_stream.device): + self.dst_prev_stream = torch.xpu.current_stream(cur_stream.device) + torch.xpu.set_stream(cur_stream) + + def __exit__(self, type: Any, value: Any, traceback: Any): + cur_stream = self.stream + if cur_stream is None or self.idx == -1: + return + + # Reset the stream on the original device and destination device + if self.src_prev_stream.device != cur_stream.device: + torch.xpu.set_stream(self.dst_prev_stream) + torch.xpu.set_stream(self.src_prev_stream) + + +def stream(stream: Optional["torch.xpu.Stream"]) -> StreamContext: + r"""Wrap around the Context-manager StreamContext that selects a given stream. + + Arguments: + stream (Stream): selected stream. This manager is a no-op if it's ``None``. + """ + return StreamContext(stream) + + +def _set_stream_by_id(stream_id, device_index, device_type): + r"""set stream specified by the stream id, device index and device type + + Args: stream_id (int): not visible to the user, used to assigned to the specific stream. + device_index (int): selected device index. + device_type (int): selected device type. + """ + torch._C._xpu_setStream( + stream_id=stream_id, + device_index=device_index, + device_type=device_type, + ) + + +def set_stream(stream: Stream): + r"""Set the current stream.This is a wrapper API to set the stream. + Usage of this function is discouraged in favor of the ``stream`` + context manager. + + Args: + stream (Stream): selected stream. This function is a no-op + if this argument is ``None``. + """ + if stream is None: + return + _lazy_init() + _set_stream_by_id( + stream_id=stream.stream_id, + device_index=stream.device_index, + device_type=stream.device_type, + ) + + +def current_stream(device: Optional[_device_t] = None) -> Stream: + r"""Return the currently selected :class:`Stream` for a given device. + + Args: + device (torch.device or int, optional): selected device. Returns + the currently selected :class:`Stream` for the current device, given + by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` + (default). + """ + _lazy_init() + streamdata = torch._C._xpu_getCurrentStream( + _get_device_index(device, optional=True) + ) + return Stream( + stream_id=streamdata[0], device_index=streamdata[1], device_type=streamdata[2] + ) + + +def synchronize(device: _device_t = None) -> None: + r"""Wait for all kernels in all streams on a XPU device to complete. + + Args: + device (torch.device or int, optional): device for which to synchronize. + It uses the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + _lazy_init() + device = _get_device_index(device, optional=True) + return torch._C._xpu_synchronize(device) + + +def get_arch_list() -> List[str]: + r"""Return list XPU architectures this library was compiled for.""" + if not is_available(): + return [] + arch_flags = torch._C._xpu_getArchFlags() + if arch_flags is None: + return [] + return arch_flags.split() + + +def get_gencode_flags() -> str: + r"""Return XPU AOT(ahead-of-time) build flags this library was compiled with.""" + arch_list = get_arch_list() + if len(arch_list) == 0: + return "" + return f'-device {",".join(arch for arch in arch_list)}' + + +def _get_generator(device: torch.device) -> torch._C.Generator: + r"""Return the XPU Generator object for the given device. + + Args: + device (torch.device): selected device. + """ + idx = device.index + if idx is None: + idx = current_device() + return torch.xpu.default_generators[idx] + + +def _set_rng_state_offset( + offset: int, device: Union[int, str, torch.device] = "xpu" +) -> None: + r"""Set the random number generator state offset of the specified GPU. + + Args: + offset (int): The desired offset + device (torch.device or int, optional): The device to set the RNG state. + Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). + """ + final_device = _get_device(device) + + def cb(): + default_generator = _get_generator(final_device) + default_generator.set_offset(offset) + + _lazy_call(cb) + + +def _get_rng_state_offset(device: Union[int, str, torch.device] = "xpu") -> int: + r"""Return the random number generator state offset of the specified GPU. + + Args: + device (torch.device or int, optional): The device to return the RNG state offset of. + Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). + + .. warning:: + This function eagerly initializes XPU. + """ + _lazy_init() + final_device = _get_device(device) + default_generator = _get_generator(final_device) + return default_generator.get_offset() + + +# import here to avoid circular import +from .memory import ( + empty_cache, + max_memory_allocated, + max_memory_reserved, + mem_get_info, + memory_allocated, + memory_reserved, + memory_stats, + memory_stats_as_nested_dict, + reset_accumulated_memory_stats, + reset_peak_memory_stats, +) +from .random import ( + get_rng_state, + get_rng_state_all, + initial_seed, + manual_seed, + manual_seed_all, + seed, + seed_all, + set_rng_state, + set_rng_state_all, +) + + +__all__ = [ + "Event", + "Stream", + "StreamContext", + "current_device", + "current_stream", + "default_generators", + "device", + "device_of", + "device_count", + "empty_cache", + "get_arch_list", + "get_device_capability", + "get_device_name", + "get_device_properties", + "get_gencode_flags", + "get_rng_state", + "get_rng_state_all", + "get_stream", + "init", + "initial_seed", + "is_available", + "is_bf16_supported", + "is_initialized", + "manual_seed", + "manual_seed_all", + "max_memory_allocated", + "max_memory_reserved", + "mem_get_info", + "memory_allocated", + "memory_reserved", + "memory_stats", + "memory_stats_as_nested_dict", + "reset_accumulated_memory_stats", + "reset_peak_memory_stats", + "seed", + "seed_all", + "set_device", + "set_rng_state", + "set_rng_state_all", + "set_stream", + "stream", + "streams", + "synchronize", +] diff --git a/lib/python3.12/site-packages/torch/xpu/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/xpu/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c6d5a017be45442e4dbcaf2ee5e84fbeca1e77f Binary files /dev/null and b/lib/python3.12/site-packages/torch/xpu/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/xpu/__pycache__/_gpu_trace.cpython-312.pyc b/lib/python3.12/site-packages/torch/xpu/__pycache__/_gpu_trace.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9444f31bad45f70fb0202821593eab13d0de49c8 Binary files /dev/null and b/lib/python3.12/site-packages/torch/xpu/__pycache__/_gpu_trace.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/xpu/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/xpu/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d1cf405a0d589d74dd7b466cd9a4444a442862f7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/xpu/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/xpu/__pycache__/memory.cpython-312.pyc b/lib/python3.12/site-packages/torch/xpu/__pycache__/memory.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35cc9266f569b9d433b9582de4d60a2e54459c16 Binary files /dev/null and b/lib/python3.12/site-packages/torch/xpu/__pycache__/memory.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/xpu/__pycache__/random.cpython-312.pyc b/lib/python3.12/site-packages/torch/xpu/__pycache__/random.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a3e39d10688184c8a8bd15d1700703da3d743d75 Binary files /dev/null and b/lib/python3.12/site-packages/torch/xpu/__pycache__/random.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/xpu/__pycache__/streams.cpython-312.pyc b/lib/python3.12/site-packages/torch/xpu/__pycache__/streams.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..515c1e0c9dd65819cceead1c983088decdc699f8 Binary files /dev/null and b/lib/python3.12/site-packages/torch/xpu/__pycache__/streams.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/xpu/_gpu_trace.py b/lib/python3.12/site-packages/torch/xpu/_gpu_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..0407abbf24959346ed78572e4be183c877d8a855 --- /dev/null +++ b/lib/python3.12/site-packages/torch/xpu/_gpu_trace.py @@ -0,0 +1,75 @@ +from typing import Callable + +from torch._utils import CallbackRegistry + + +EventCreationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU event creation" +) +EventDeletionCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU event deletion" +) +EventRecordCallbacks: "CallbackRegistry[int, int]" = CallbackRegistry( + "XPU event record" +) +EventWaitCallbacks: "CallbackRegistry[int, int]" = CallbackRegistry( + "XPU event wait" +) +MemoryAllocationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU memory allocation" +) +MemoryDeallocationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU memory deallocation" +) +StreamCreationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU stream creation" +) +DeviceSynchronizationCallbacks: "CallbackRegistry[[]]" = CallbackRegistry( + "XPU device synchronization" +) +StreamSynchronizationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU stream synchronization" +) +EventSynchronizationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU event synchronization" +) + + +def register_callback_for_event_creation(cb: Callable[[int], None]) -> None: + EventCreationCallbacks.add_callback(cb) + + +def register_callback_for_event_deletion(cb: Callable[[int], None]) -> None: + EventDeletionCallbacks.add_callback(cb) + + +def register_callback_for_event_record(cb: Callable[[int, int], None]) -> None: + EventRecordCallbacks.add_callback(cb) + + +def register_callback_for_event_wait(cb: Callable[[int, int], None]) -> None: + EventWaitCallbacks.add_callback(cb) + + +def register_callback_for_memory_allocation(cb: Callable[[int], None]) -> None: + MemoryAllocationCallbacks.add_callback(cb) + + +def register_callback_for_memory_deallocation(cb: Callable[[int], None]) -> None: + MemoryDeallocationCallbacks.add_callback(cb) + + +def register_callback_for_stream_creation(cb: Callable[[int], None]) -> None: + StreamCreationCallbacks.add_callback(cb) + + +def register_callback_for_device_synchronization(cb: Callable[[], None]) -> None: + DeviceSynchronizationCallbacks.add_callback(cb) + + +def register_callback_for_stream_synchronization(cb: Callable[[int], None]) -> None: + StreamSynchronizationCallbacks.add_callback(cb) + + +def register_callback_for_event_synchronization(cb: Callable[[int], None]) -> None: + EventSynchronizationCallbacks.add_callback(cb) diff --git a/lib/python3.12/site-packages/torch/xpu/_utils.py b/lib/python3.12/site-packages/torch/xpu/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8f738267459a2791a4a33ca4bec74800a58f0b9a --- /dev/null +++ b/lib/python3.12/site-packages/torch/xpu/_utils.py @@ -0,0 +1,39 @@ +from typing import Any + +import torch + +# The _get_device_index has been moved to torch.utils._get_device_index +from torch._utils import _get_device_index as _torch_get_device_index + + +def _get_device_index( + device: Any, optional: bool = False, allow_cpu: bool = False +) -> int: + r"""Get the device index from :attr:`device`, which can be a torch.device + object, a Python integer, or ``None``. + + If :attr:`device` is a torch.device object, returns the device index if it + is a XPU device. Note that for a XPU device without a specified index, + i.e., ``torch.device('xpu')``, this will return the current default XPU + device if :attr:`optional` is ``True``. If :attr:`allow_cpu` is ``True``, + CPU devices will be accepted and ``-1`` will be returned in this case. + + If :attr:`device` is a Python integer, it is returned as is. + + If :attr:`device` is ``None``, this will return the current default XPU + device if :attr:`optional` is ``True``. + """ + if isinstance(device, int): + return device + if isinstance(device, str): + device = torch.device(device) + if isinstance(device, torch.device): + if allow_cpu: + if device.type not in ["xpu", "cpu"]: + raise ValueError(f"Expected a xpu or cpu device, but got: {device}") + elif device.type != "xpu": + raise ValueError(f"Expected a xpu device, but got: {device}") + if not torch.jit.is_scripting(): + if isinstance(device, torch.xpu.device): + return device.idx + return _torch_get_device_index(device, optional, allow_cpu) diff --git a/lib/python3.12/site-packages/torch/xpu/memory.py b/lib/python3.12/site-packages/torch/xpu/memory.py new file mode 100644 index 0000000000000000000000000000000000000000..2d8bd296dff0fb923bce6bf5743cd34abbb87034 --- /dev/null +++ b/lib/python3.12/site-packages/torch/xpu/memory.py @@ -0,0 +1,208 @@ +import collections +from typing import Any, Dict, Tuple, Union + +import torch +from torch.types import Device + +from . import _get_device_index, is_initialized + + +_device_t = Union[Device, str, int, None] + + +def empty_cache() -> None: + r"""Release all unoccupied cached memory currently held by the caching + allocator so that those can be used in other XPU application. + + .. note:: + :func:`~torch.xpu.empty_cache` doesn't increase the amount of XPU + memory available for PyTorch. However, it may help reduce fragmentation + of XPU memory in certain cases. + """ + if is_initialized(): + torch._C._xpu_emptyCache() + + +def reset_peak_memory_stats(device: _device_t = None) -> None: + r"""Reset the "peak" stats tracked by the XPU memory allocator. + + See :func:`~torch.xpu.memory_stats` for details. Peak stats correspond to the + `"peak"` key in each individual stat dict. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + device = _get_device_index(device, optional=True) + return torch._C._xpu_resetPeakMemoryStats(device) + + +def reset_accumulated_memory_stats(device: _device_t = None) -> None: + r"""Reset the "accumulated" (historical) stats tracked by the XPU memory allocator. + + See :func:`~torch.xpu.memory_stats` for details. Accumulated stats correspond to + the `"allocated"` and `"freed"` keys in each individual stat dict. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + device = _get_device_index(device, optional=True) + return torch._C._xpu_resetAccumulatedMemoryStats(device) + + +def memory_stats_as_nested_dict(device: _device_t = None) -> Dict[str, Any]: + r"""Return the result of :func:`~torch.xpu.memory_stats` as a nested dictionary.""" + if not is_initialized(): + return {} + device = _get_device_index(device, optional=True) + return torch._C._xpu_memoryStats(device) + + +def memory_stats(device: _device_t = None) -> Dict[str, Any]: + r"""Return a dictionary of XPU memory allocator statistics for a given device. + + The return value of this function is a dictionary of statistics, each of + which is a non-negative integer. + + Core statistics: + + - ``"allocated_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``: + amount of allocated memory. + - ``"reserved_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``: + amount of reserved memory. + - ``"active_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``: + amount of active memory. + - ``"requested_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``: + memory requested by client code, compare this with allocated_bytes to check if + allocation rounding adds too much overhead. + + For these core statistics, values are broken down as follows. + + Pool type: + + - ``all``: combined statistics across all memory pools. + - ``large_pool``: statistics for the large allocation pool (for size >= 1MB allocations). + - ``small_pool``: statistics for the small allocation pool (for size < 1MB allocations). + + Metric type: + + - ``current``: current value of this metric. + - ``peak``: maximum value of this metric. + - ``allocated``: historical total increase in this metric. + - ``freed``: historical total decrease in this metric. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistics for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + result = [] + + def _recurse_add_to_result(prefix: str, obj: Any) -> None: + if isinstance(obj, dict): + if len(prefix) > 0: + prefix += "." + for k, v in obj.items(): + _recurse_add_to_result(prefix + k, v) + else: + result.append((prefix, obj)) + + stats = memory_stats_as_nested_dict(device=device) + _recurse_add_to_result("", stats) + result.sort() + + return collections.OrderedDict(result) + + +def memory_allocated(device: _device_t = None) -> int: + r"""Return the current GPU memory occupied by tensors in bytes for a given device. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + + .. note:: + This is likely less than the amount shown in `xpu-smi` since some + unused memory can be held by the caching allocator and some context + needs to be created on GPU. + """ + return memory_stats(device=device).get("allocated_bytes.all.current", 0) + + +def max_memory_allocated(device: _device_t = None) -> int: + r"""Return the maximum GPU memory occupied by tensors in bytes for a given device. + + By default, this returns the peak allocated memory since the beginning of + this program. :func:`~torch.xpu.reset_peak_memory_stats` can be used to + reset the starting point in tracking this metric. For example, these two + functions can measure the peak allocated memory usage of each iteration in a + training loop. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + return memory_stats(device=device).get("allocated_bytes.all.peak", 0) + + +def memory_reserved(device: _device_t = None) -> int: + r"""Return the current GPU memory managed by the caching allocator in bytes for a given device. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + return memory_stats(device=device).get("reserved_bytes.all.current", 0) + + +def max_memory_reserved(device: _device_t = None) -> int: + r"""Return the maximum GPU memory managed by the caching allocator in bytes for a given device. + + By default, this returns the peak cached memory since the beginning of this + program. :func:`~torch.xpu.reset_peak_memory_stats` can be used to reset + the starting point in tracking this metric. For example, these two functions + can measure the peak cached memory amount of each iteration in a training + loop. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + return memory_stats(device=device).get("reserved_bytes.all.peak", 0) + + +def mem_get_info(device: _device_t = None) -> Tuple[int, int]: + r"""Return the global free and total GPU memory for a given device. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + + Returns: + int: the memory available on the device in units of bytes. + int: the total memory on the device in units of bytes + """ + device = _get_device_index(device, optional=True) + return torch._C._xpu_getMemoryInfo(device) + + +__all__ = [ + "empty_cache", + "max_memory_allocated", + "max_memory_reserved", + "mem_get_info", + "memory_allocated", + "memory_reserved", + "memory_stats", + "memory_stats_as_nested_dict", + "reset_accumulated_memory_stats", + "reset_peak_memory_stats", +] diff --git a/lib/python3.12/site-packages/torch/xpu/random.py b/lib/python3.12/site-packages/torch/xpu/random.py new file mode 100644 index 0000000000000000000000000000000000000000..5bc142418637cbb69dd5602fc0f9b9eafa654252 --- /dev/null +++ b/lib/python3.12/site-packages/torch/xpu/random.py @@ -0,0 +1,176 @@ +# mypy: allow-untyped-defs +from typing import Iterable, List, Union + +import torch +from torch import Tensor + +from . import _lazy_call, _lazy_init, current_device, device_count + + +def get_rng_state(device: Union[int, str, torch.device] = "xpu") -> Tensor: + r"""Return the random number generator state of the specified GPU as a ByteTensor. + + Args: + device (torch.device or int, optional): The device to return the RNG state of. + Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). + + .. warning:: + This function eagerly initializes XPU. + """ + _lazy_init() + if isinstance(device, str): + device = torch.device(device) + elif isinstance(device, int): + device = torch.device("xpu", device) + idx = device.index + if idx is None: + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + return default_generator.get_state() + + +def get_rng_state_all() -> List[Tensor]: + r"""Return a list of ByteTensor representing the random number states of all devices.""" + results = [get_rng_state(i) for i in range(device_count())] + return results + + +def set_rng_state( + new_state: Tensor, device: Union[int, str, torch.device] = "xpu" +) -> None: + r"""Set the random number generator state of the specified GPU. + + Args: + new_state (torch.ByteTensor): The desired state + device (torch.device or int, optional): The device to set the RNG state. + Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). + """ + with torch._C._DisableFuncTorch(): + new_state_copy = new_state.clone(memory_format=torch.contiguous_format) + if isinstance(device, str): + device = torch.device(device) + elif isinstance(device, int): + device = torch.device("xpu", device) + + def cb(): + idx = device.index + if idx is None: + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + default_generator.set_state(new_state_copy) + + _lazy_call(cb) + + +def set_rng_state_all(new_states: Iterable[Tensor]) -> None: + r"""Set the random number generator state of all devices. + + Args: + new_states (Iterable of torch.ByteTensor): The desired state for each device. + """ + for i, state in enumerate(new_states): + set_rng_state(state, i) + + +def manual_seed(seed: int) -> None: + r"""Set the seed for generating random numbers for the current GPU. + + It's safe to call this function if XPU is not available; in that case, it is silently ignored. + + Args: + seed (int): The desired seed. + + .. warning:: + If you are working with a multi-GPU model, this function is insufficient + to get determinism. To seed all GPUs, use :func:`manual_seed_all`. + """ + seed = int(seed) + + def cb(): + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + default_generator.manual_seed(seed) + + _lazy_call(cb, seed=True) + + +def manual_seed_all(seed: int) -> None: + r"""Set the seed for generating random numbers on all GPUs. + + It's safe to call this function if XPU is not available; in that case, it is silently ignored. + + Args: + seed (int): The desired seed. + """ + seed = int(seed) + + def cb(): + for i in range(device_count()): + default_generator = torch.xpu.default_generators[i] + default_generator.manual_seed(seed) + + _lazy_call(cb, seed_all=True) + + +def seed() -> None: + r"""Set the seed for generating random numbers to a random number for the current GPU. + + It's safe to call this function if XPU is not available; in that case, it is silently ignored. + + .. warning:: + If you are working with a multi-GPU model, this function will only initialize + the seed on one GPU. To initialize all GPUs, use :func:`seed_all`. + """ + + def cb(): + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + default_generator.seed() + + _lazy_call(cb) + + +def seed_all() -> None: + r"""Set the seed for generating random numbers to a random number on all GPUs. + + It's safe to call this function if XPU is not available; in that case, it is silently ignored. + """ + + def cb(): + random_seed = 0 + seeded = False + for i in range(device_count()): + default_generator = torch.xpu.default_generators[i] + if not seeded: + default_generator.seed() + random_seed = default_generator.initial_seed() + seeded = True + else: + default_generator.manual_seed(random_seed) + + _lazy_call(cb) + + +def initial_seed() -> int: + r"""Return the current random seed of the current GPU. + + .. warning:: + This function eagerly initializes XPU. + """ + _lazy_init() + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + return default_generator.initial_seed() + + +__all__ = [ + "get_rng_state", + "get_rng_state_all", + "set_rng_state", + "set_rng_state_all", + "manual_seed", + "manual_seed_all", + "seed", + "seed_all", + "initial_seed", +] diff --git a/lib/python3.12/site-packages/torch/xpu/streams.py b/lib/python3.12/site-packages/torch/xpu/streams.py new file mode 100644 index 0000000000000000000000000000000000000000..beb438be466d998ad57b49e0327ecc532df1efd3 --- /dev/null +++ b/lib/python3.12/site-packages/torch/xpu/streams.py @@ -0,0 +1,169 @@ +# mypy: allow-untyped-defs +import ctypes + +import torch +from torch._utils import _dummy_type + + +if not hasattr(torch._C, "_XpuStreamBase"): + # Define dummy base classes + torch._C.__dict__["_XpuStreamBase"] = _dummy_type("_XpuStreamBase") + torch._C.__dict__["_XpuEventBase"] = _dummy_type("_XpuEventBase") + + +class Stream(torch._C._XpuStreamBase): + r"""Wrapper around a XPU stream. + + A XPU stream is a linear sequence of execution that belongs to a specific + device, independent from other streams. + + Args: + device(torch.device or int, optional): a device on which to allocate + the stream. If :attr:`device` is ``None`` (default) or a negative + integer, this will use the current device. + priority(int, optional): priority of the stream, should be 0 or + negative, where negative numbers indicate higher priority. By default, + streams have priority 0. + """ + + def __new__(cls, device=None, priority=0, **kwargs): + # setting device manager is expensive, so we avoid it unless necessary + if device is None or ("stream_id" in kwargs and "device_index" in kwargs): + return super().__new__(cls, priority=priority, **kwargs) + else: + with torch.xpu.device(device): + return super().__new__(cls, priority=priority, **kwargs) + + def wait_event(self, event) -> None: + r"""Make all future work submitted to the stream wait for an event. + + Args: + event (torch.xpu.Event): an event to wait for. + """ + event.wait(self) + + def wait_stream(self, stream) -> None: + r"""Synchronize with another stream. + + All future work submitted to this stream will wait until all kernels + submitted to a given stream at the time of call complete. + + Args: + stream (Stream): a stream to synchronize. + """ + self.wait_event(stream.record_event()) + + def record_event(self, event=None): + r"""Record an event. + + Args: + event (torch.xpu.Event, optional): event to record. If not given, a new one + will be allocated. + + Returns: + Recorded event. + """ + if event is None: + event = Event() + event.record(self) + return event + + def query(self) -> bool: + r"""Check if all the work submitted has been completed. + + Returns: + A boolean indicating if all kernels in this stream are completed. + """ + return super().query() + + def synchronize(self) -> None: + r"""Wait for all the kernels in this stream to complete.""" + super().synchronize() + + @property + def _as_parameter_(self): + return ctypes.c_void_p(self.sycl_queue) + + def __eq__(self, o): + if isinstance(o, Stream): + return super().__eq__(o) + return False + + def __hash__(self): + return hash((self.sycl_queue, self.device)) + + def __repr__(self): + return f"torch.xpu.Stream(device={self.device} sycl_queue={self.sycl_queue:#x})" + + +class Event(torch._C._XpuEventBase): + r"""Wrapper around a XPU event. + + XPU events are synchronization markers that can be used to monitor the + device's progress, and to synchronize XPU streams. + + The underlying XPU events are lazily initialized when the event is first + recorded. After creation, only streams on the same device may record the + event. However, streams on any device can wait on the event. + + Args: + enable_timing (bool, optional): indicates if the event should measure time + (default: ``False``) + """ + + def __new__(cls, enable_timing=False): + return super().__new__(cls, enable_timing=enable_timing) + + def record(self, stream=None) -> None: + r"""Record the event in a given stream. + + Uses ``torch.xpu.current_stream()`` if no stream is specified. The + stream's device must match the event's device. + """ + if stream is None: + stream = torch.xpu.current_stream() + super().record(stream) + + def wait(self, stream=None) -> None: + r"""Make all future work submitted to the given stream wait for this event. + + Use ``torch.xpu.current_stream()`` if no stream is specified. + """ + if stream is None: + stream = torch.xpu.current_stream() + super().wait(stream) + + def query(self) -> bool: + r"""Check if all work currently captured by event has completed. + + Returns: + A boolean indicating if all work currently captured by event has + completed. + """ + return super().query() + + def elapsed_time(self, end_event): + r"""Return the time elapsed. + + Time reported in milliseconds after the event was recorded and + before the end_event was recorded. + """ + return super().elapsed_time(end_event) + + def synchronize(self) -> None: + r"""Wait for the event to complete. + + Waits until the completion of all work currently captured in this event. + This prevents the CPU thread from proceeding until the event completes. + """ + super().synchronize() + + @property + def _as_parameter_(self): + return ctypes.c_void_p(self.sycl_event) + + def __repr__(self): + if self.sycl_event: + return f"torch.xpu.Event(sycl_event={self.sycl_event:#x})" + else: + return "torch.xpu.Event(uninitialized)"