# mypy: allow-untyped-defs # Copyright (c) Meta Platforms, Inc. and affiliates import math from collections.abc import Sequence from dataclasses import dataclass from enum import Enum from typing import Any, cast import torch from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta from torch.distributed.tensor._op_schema import ( OpSchema, OpSpec, OpStrategy, PlacementList, RuntimeSchemaInfo, TupleStrategy, ) from torch.distributed.tensor._ops.single_dim_strategy import ( _ShardingPlaceholder, register_single_dim_strategy, ) from torch.distributed.tensor._ops.utils import ( as_list, expand_to_full_mesh_op_strategy, generate_redistribute_costs, is_tensor_evenly_shardable, is_tensor_evenly_shardable_on_dim, normalize_dim, normalize_dims, register_op_strategy, ) from torch.distributed.tensor._utils import normalize_to_torch_size from torch.distributed.tensor.placement_types import ( _is_shard_like, _StridedShard, Partial, Placement, Replicate, Shard, ) aten = torch.ops.aten prims = torch.ops.prims class Reduction(Enum): NONE = 0 MEAN = 1 SUM = 2 @dataclass(frozen=True) class NormReduction: norm_type: int | float ReductionOpType = NormReduction | str @dataclass(frozen=True) class _NormPartial(Partial): """ This placement is used for partial p-norm (p not in {inf, -inf, 0}). For p-norms, the p-norm over n elements computes (sum_i x_i^p)^(1/p). For example, consider 2 ranks, a (4,) tensor sharded on dim-0, and 2-norm: Rank 0: [t1, t2] | Rank 1: [t3, t4] After computing 2-norm per gradient (partial placement): Rank 0: [sqrt(t1^2 + t2^2)] | Rank 1: [sqrt(t3^2 + t4^2)] Converting from partial to replicate wants to ultimately get: Rank 0/1: [sqrt(t1^2 + t2^2 + t3^2 + t4^2)] This is achieved by: x^p -> allreduce sum -> x^(1/p). """ norm_type: int | float = 2 def __init__(self, norm_type: int | float = 2): super().__init__("sum") object.__setattr__(self, "norm_type", norm_type) def _partition_value( self, tensor: torch.Tensor, mesh: DeviceMesh, mesh_dim: int ) -> torch.Tensor: return tensor / math.pow(mesh.size(mesh_dim), 1 / self.norm_type) def _reduce_shard_value( self, tensor: torch.Tensor, mesh: DeviceMesh, mesh_dim: int, shard_spec: Placement, ) -> torch.Tensor: if not isinstance(shard_spec, Shard): raise AssertionError(f"Expected Shard, got {type(shard_spec)}") tensor = self._pre_reduce_transform(tensor) reduced_tensor = super()._reduce_shard_value(tensor, mesh, mesh_dim, shard_spec) return self._post_reduce_transform(reduced_tensor) def _reduce_value( self, tensor: torch.Tensor, mesh: DeviceMesh, mesh_dim: int ) -> torch.Tensor: tensor = self._pre_reduce_transform(tensor) reduced_tensor = super()._reduce_value(tensor, mesh, mesh_dim) return self._post_reduce_transform(reduced_tensor) def _pre_reduce_transform(self, tensor: torch.Tensor) -> torch.Tensor: return tensor**self.norm_type def _post_reduce_transform(self, tensor: torch.Tensor) -> torch.Tensor: return tensor ** (1.0 / self.norm_type) def __eq__(self, other: object) -> bool: if not isinstance(other, _NormPartial): return False return self.norm_type == other.norm_type def __hash__(self) -> int: return 1 + hash(self.norm_type) def __repr__(self) -> str: return f"_NormPartial({self.norm_type})" def __str__(self) -> str: return f"_NormP({self.norm_type})" def _infer_reduction_dims(dims_arg: object, ndim: int) -> list[int] | None: if dims_arg is None: return None dims = cast(list[int], as_list(dims_arg)) dims = cast(list[int], normalize_dims(dims, ndim)) empty_dims = [[0], [-1], []] if ndim == 0 and dims_arg in empty_dims: return None return dims def _infer_reduce_dims_map( reduction_dims: list[int], input_ndim: int, keep_dim=False ) -> list[int]: reduction_dims_map = [] new_dim_count = 0 for input_dim in range(input_ndim): if input_dim in reduction_dims and not keep_dim: # if input dim in reduction dims, mark it as -1 reduction_dims_map.append(-1) else: # otherwise mark it as the new dim reduction_dims_map.append(new_dim_count) new_dim_count += 1 return reduction_dims_map def _replicate_dims_start_at( placements: Sequence[Placement], start_dim: int = 0 ) -> tuple[Placement, ...]: new_placements: list[Placement] = [] for p in placements: if p.is_partial() or (_is_shard_like(p) and p.dim >= start_dim): new_placements.append(Replicate()) # make it replicate else: new_placements.append(p) # keep the placement return tuple(new_placements) # return new_placements which align with placements but skip the skipped_dim # Precondition: no shard-like placement on skipped_dim (callers must # replicate it first via replicate_reduction_dims). def _skip_dim( placements: tuple[Placement, ...], skipped_dim: int ) -> tuple[Placement, ...]: new_placements: list[Placement] = [] for p in placements: if isinstance(p, _StridedShard) and p.dim >= skipped_dim: new_placements.append(_StridedShard(p.dim - 1, split_factor=p.split_factor)) elif isinstance(p, Shard) and p.dim >= skipped_dim: new_placements.append(Shard(p.dim - 1)) else: new_placements.append(p) return tuple(new_placements) def replicate_reduction_dims( placements: tuple[Placement, ...], reduction_dims: list[int] ) -> tuple[Placement, ...]: # replicate the reduction dims if not reduction_linear new_placements: list[Placement] = [] for p in placements: if p.is_partial(): new_placements.append(Replicate()) elif _is_shard_like(p) and p.dim in reduction_dims: new_placements.append(Replicate()) else: new_placements.append(p) return tuple(new_placements) def map_placements_after_reduction( placements: tuple[Placement, ...], reduction_dims: list[int], reduction_dims_map: list[int], reduction_op: ReductionOpType, ) -> tuple[Placement, ...]: """ Map each placement based on the output shape after reduction. """ new_placements: list[Placement] = [] for placement in placements: if isinstance(placement, (Replicate, Partial)): new_placements.append(placement) else: if not _is_shard_like(placement): raise AssertionError( f"Expected Shard/_StridedShard, got {type(placement)}" ) shard_dim = placement.dim new_shard_dim = reduction_dims_map[shard_dim] if new_shard_dim == -1 or shard_dim in reduction_dims: # if new_shard_dim collapsed or its in the reduction dims # (i.e. for the case where keepdims=True), we generate partial new_placements.append(get_placement_from_reduction_op(reduction_op)) else: if isinstance(placement, _StridedShard): new_placements.append( _StridedShard( new_shard_dim, split_factor=placement.split_factor ) ) elif isinstance(placement, Shard): new_placements.append(Shard(new_shard_dim)) return tuple(new_placements) def get_placement_from_reduction_op(reduction_op: ReductionOpType) -> Placement: if isinstance(reduction_op, NormReduction): if reduction_op.norm_type == 0: # return P(sum) for easier reduction_linear handling. return Partial("sum") return _NormPartial(norm_type=reduction_op.norm_type) return Partial(reduction_op) def common_reduction_strategy( input_strategy: OpStrategy, reduce_dims: list[int], keep_dim: bool = False, reduction_linear: bool = True, reduction_op: ReductionOpType = "sum", ) -> OpStrategy: """ reduction_linear means that the reduction `f` follows this rule: f([f(a), f(b)]) = f([a, b]) reduction linear should be super set of linearity. """ # by default follow reduction input strategy reduction_strategy = OpStrategy([]) for op_spec in input_strategy.strategies: if reduction_op == "avg": output_spec = op_spec.output_spec local_shape = list(output_spec.tensor_meta.shape) # type:ignore[union-attr] for dim in reduce_dims: if not is_tensor_evenly_shardable_on_dim(local_shape, output_spec, dim): # reduce(avg) is not linear for unevenly sharded tensors reduction_linear = False break for p in op_spec.output_spec.placements: # when the partial reduction op matches the global reduction op, # we can delay redistribution (i.e max, max) if isinstance(p, Partial) and p.reduce_op != reduction_op: reduction_linear = False break if not reduction_linear: # input placements for this strategy should clear out pending sum and sharding # on the reduction dimension input_placements = replicate_reduction_dims( op_spec.output_spec.placements, reduce_dims ) else: input_placements = op_spec.output_spec.placements input_spec = DTensorSpec( mesh=input_strategy.mesh, placements=input_placements, tensor_meta=op_spec.output_spec.tensor_meta, ) reduce_dims_map = _infer_reduce_dims_map(reduce_dims, input_spec.ndim, keep_dim) out_placements = map_placements_after_reduction( input_spec.placements, reduce_dims, reduce_dims_map, reduction_op ) redistribute_cost = [generate_redistribute_costs(input_strategy, input_spec)] reduction_strategy.strategies.append( OpSpec( output_specs=DTensorSpec( mesh=input_strategy.mesh, placements=out_placements, ), input_specs=(input_spec,), redistribute_cost=redistribute_cost, ) ) return reduction_strategy LINEAR_REDUCTION_OP_MAP = { aten.all.default: "product", aten.all.dim: "product", aten.sum.default: "sum", aten.sum.dim_IntList: "sum", prims.sum.default: "sum", aten.any.default: "sum", aten.any.dim: "sum", aten.any.dims: "sum", aten.any.out: "sum", # These are only valid when there is no padding aten.prod.default: "product", aten.prod.dim_int: "product", aten.prod.int_out: "product", prims.prod.default: "product", # avg is only linear when there is no padding aten.mean.default: "avg", aten.mean.dim: "avg", aten.mean.out: "avg", aten.max.default: "max", aten.max.out: "max", aten.min.default: "min", aten.min.out: "min", aten.amax.default: "max", aten.amax.out: "max", aten.amin.default: "min", aten.amin.out: "min", aten.nansum.default: "sum", } # argmax/argmin return indices which cannot be combined with P(max/min). # They need special handling that forces redistribution on reduction dims. ARGMAX_ARGMIN_OPS = { aten.argmax.default: "max", aten.argmin.default: "min", } @register_op_strategy( list(LINEAR_REDUCTION_OP_MAP.keys()), schema_info=RuntimeSchemaInfo(1) ) def linear_reduction_strategy(op_schema: OpSchema) -> OpStrategy: args_schema = op_schema.args_schema input_strategy = args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") dims = None if len(op_schema.args_schema) > 1: dims = _infer_reduction_dims(args_schema[1], input_strategy.ndim) reduce_dims = list(range(input_strategy.ndim)) if dims is None else dims keep_dim = len(op_schema.args_schema) > 2 and bool(op_schema.args_schema[2]) reduction_op = LINEAR_REDUCTION_OP_MAP[op_schema.op] return common_reduction_strategy( input_strategy, reduce_dims, keep_dim=keep_dim, reduction_linear=True, reduction_op=reduction_op, ) # max.dim/min.dim return (values, indices). Indices are local to each shard # and cannot be combined across ranks, so we force Replicate on reduction dims # (same approach as argmax/argmin). @register_single_dim_strategy( [aten.max.dim, aten.min.dim], schema_info=RuntimeSchemaInfo(1) ) def max_min_dim_single_dim_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: input_meta = args_schema[0] if not isinstance(input_meta, TensorMeta): raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}") ndim = len(input_meta.shape) dim = normalize_dim(cast(int, args_schema[1]), ndim) keep_dim = len(args_schema) > 2 and bool(args_schema[2]) strategies: list[list[Placement | _ShardingPlaceholder]] = [] for d in range(ndim): if d == dim: continue out_d = d if keep_dim or d < dim else d - 1 # [values, indices, input]: shard on non-reduction dim strategies.append( [ _ShardingPlaceholder(out_d), _ShardingPlaceholder(out_d), _ShardingPlaceholder(d), ] ) return strategies @register_op_strategy(list(ARGMAX_ARGMIN_OPS.keys()), schema_info=RuntimeSchemaInfo(1)) def argmax_argmin_strategy(op_schema: OpSchema) -> OpStrategy: """ Strategy for argmax/argmin. These return indices, not values, so they cannot use P(max/min) output placements. The indices are local to each shard and cannot be meaningfully combined across ranks with a max/min reduction. Force redistribution on reduction dimensions by using reduction_linear=False. """ args_schema = op_schema.args_schema input_strategy = args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") dims = None if len(op_schema.args_schema) > 1: dims = _infer_reduction_dims(args_schema[1], input_strategy.ndim) reduce_dims = list(range(input_strategy.ndim)) if dims is None else dims keep_dim = len(op_schema.args_schema) > 2 and bool(op_schema.args_schema[2]) reduction_op = ARGMAX_ARGMIN_OPS[op_schema.op] return common_reduction_strategy( input_strategy, reduce_dims, keep_dim=keep_dim, reduction_linear=False, # Force redistribution - indices can't use P(max/min) # reduction_op is effectively unused here: reduction_linear=False # forces all reduction-dim Shard placements to Replicate before # map_placements_after_reduction, so no Shard-on-reduction-dim # remains to convert to Partial. Passed for consistency. reduction_op=reduction_op, ) @register_op_strategy( [aten.cumsum.default, aten.cumprod.default, aten.logcumsumexp.default], schema_info=RuntimeSchemaInfo(1), ) def scan_strategy(op_schema: OpSchema) -> OpStrategy: args_schema = op_schema.args_schema input_strategy = args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") dim = args_schema[1] if not isinstance(dim, int): raise AssertionError(f"Expected int, got {type(dim)}") return common_reduction_strategy( input_strategy, [dim], keep_dim=True, reduction_linear=False ) @register_op_strategy( [aten.median.default, aten.nanmedian.default], schema_info=RuntimeSchemaInfo(1), ) def global_median_strategy(op_schema: OpSchema) -> OpStrategy: input_strategy = cast(OpStrategy, op_schema.args_schema[0]) reduce_dims = list(range(input_strategy.ndim)) return common_reduction_strategy( input_strategy, reduce_dims, reduction_linear=False ) @register_single_dim_strategy( [aten.median.dim, aten.nanmedian.dim, aten.mode.default], schema_info=RuntimeSchemaInfo(1), ) def dim_reduction_with_indices_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: input_meta = args_schema[0] if not isinstance(input_meta, TensorMeta): raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}") ndim = len(input_meta.shape) dim = normalize_dim(cast(int, args_schema[1]) if len(args_schema) > 1 else -1, ndim) keep_dim = len(args_schema) > 2 and bool(args_schema[2]) strategies: list[list[Placement | _ShardingPlaceholder]] = [] for d in range(ndim): if d == dim: continue out_d = d if keep_dim or d < dim else d - 1 strategies.append( [ _ShardingPlaceholder(out_d), _ShardingPlaceholder(out_d), _ShardingPlaceholder(d), ] ) return strategies @register_single_dim_strategy( [aten.kthvalue.default], schema_info=RuntimeSchemaInfo(2), ) def kthvalue_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: input_meta = args_schema[0] if not isinstance(input_meta, TensorMeta): raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}") ndim = len(input_meta.shape) dim = normalize_dim(cast(int, args_schema[2]) if len(args_schema) > 2 else -1, ndim) keep_dim = len(args_schema) > 3 and bool(args_schema[3]) strategies: list[list[Placement | _ShardingPlaceholder]] = [] for d in range(ndim): if d == dim: continue out_d = d if keep_dim or d < dim else d - 1 strategies.append( [ _ShardingPlaceholder(out_d), _ShardingPlaceholder(out_d), _ShardingPlaceholder(d), ] ) return strategies @register_op_strategy( [aten.cummax.default, aten.cummin.default], schema_info=RuntimeSchemaInfo(1), ) def cummax_cummin_strategy(op_schema: OpSchema) -> OpStrategy: dim = cast(int, op_schema.args_schema[1]) return sort_strategy(op_schema, dim) @register_op_strategy( [ aten.std.correction, aten.std.correction_out, aten.var.correction, aten.var.correction_out, aten.var_mean.correction, aten.var_mean.correction_out, prims.var.default, ], schema_info=RuntimeSchemaInfo(1, ["keepdim"]), ) def std_var_reduction_strategy(op_schema: OpSchema) -> OpStrategy: args_schema = op_schema.args_schema input_strategy = args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") dims = None if len(op_schema.args_schema) > 1: dims = _infer_reduction_dims(args_schema[1], input_strategy.ndim) reduce_dims = list(range(input_strategy.ndim)) if dims is None else dims keep_dim = cast(bool, op_schema.kwargs_schema.get("keepdim", False)) return common_reduction_strategy( input_strategy, reduce_dims, keep_dim=keep_dim, reduction_linear=False ) def _get_norm_reduction_op(norm_type: int | float | str) -> ReductionOpType: """Get the reduction op for vector/foreach norm based on norm_type. For inf/-inf norms, returns simple reduction ops ("max", "min"). For other norms (including 0), returns NormReduction which produces the appropriate Partial placement via get_placement_from_reduction_op. """ if norm_type in (float("inf"), "inf"): return "max" elif norm_type in (float("-inf"), "-inf"): return "min" else: if not isinstance(norm_type, (int, float)): raise AssertionError return NormReduction(norm_type) @register_op_strategy( [aten.linalg_vector_norm.default, aten.norm.Scalar], schema_info=RuntimeSchemaInfo(1), ) def vector_norm_strategy(op_schema: OpSchema) -> OpStrategy: args_schema = op_schema.args_schema input_strategy = args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") norm_type = args_schema[1] if len(args_schema) > 1 else 2 if not isinstance(norm_type, (int, float, str)): raise AssertionError(f"Expected int, float, or str, got {type(norm_type)}") dim = args_schema[2] if len(args_schema) > 2 else None keepdim = args_schema[3] if len(args_schema) > 3 else False dims = _infer_reduction_dims(dim, input_strategy.ndim) reduce_dims = list(range(input_strategy.ndim)) if dims is None else dims return common_reduction_strategy( input_strategy, reduce_dims, keep_dim=cast(bool, keepdim), reduction_op=_get_norm_reduction_op(norm_type), ) @register_op_strategy( [aten._foreach_norm.Scalar], schema_info=RuntimeSchemaInfo(1, needs_pytree=True) ) def foreach_norm_strategy(op_schema: OpSchema) -> TupleStrategy: args_schema = op_schema.args_schema input_tuple_strategy = args_schema[0] if not isinstance(input_tuple_strategy, TupleStrategy): raise AssertionError( f"Expected TupleStrategy, got {type(input_tuple_strategy)}" ) norm_type = args_schema[1] if len(args_schema) > 1 else 2 if not isinstance(norm_type, (int, float, str)): raise AssertionError(f"Expected int, float, or str, got {type(norm_type)}") output_tuple_strategy_children: list[OpStrategy] = [] for op_strategy in input_tuple_strategy.children: if not isinstance(op_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(op_strategy)}") reduce_dims = list(range(op_strategy.ndim)) output_strategy = common_reduction_strategy( op_strategy, reduce_dims, reduction_op=_get_norm_reduction_op(norm_type), ) output_tuple_strategy_children.append(output_strategy) return TupleStrategy(output_tuple_strategy_children) @register_op_strategy([aten.linalg__powsum.default], schema_info=RuntimeSchemaInfo(1)) def powsum_strategy(op_schema: OpSchema) -> OpStrategy: """ Strategy for linalg__powsum: computes sum(|x|^ord) without the final root. Output is always reducible with Partial("sum"). """ args_schema = op_schema.args_schema input_strategy = args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") dim = args_schema[2] if len(args_schema) > 2 else None keepdim = args_schema[3] if len(args_schema) > 3 else False dims = _infer_reduction_dims(dim, input_strategy.ndim) reduce_dims = list(range(input_strategy.ndim)) if dims is None else dims return common_reduction_strategy( input_strategy, reduce_dims, keep_dim=cast(bool, keepdim), reduction_linear=True, reduction_op="sum", ) @register_op_strategy( [aten._foreach_powsum.Scalar], schema_info=RuntimeSchemaInfo(1, needs_pytree=True) ) def foreach_powsum_strategy(op_schema: OpSchema) -> TupleStrategy: """ Strategy for _foreach_powsum: computes sum(|x|^ord) for each tensor. Output is always reducible with Partial("sum"). """ args_schema = op_schema.args_schema input_tuple_strategy = args_schema[0] if not isinstance(input_tuple_strategy, TupleStrategy): raise AssertionError( f"Expected TupleStrategy, got {type(input_tuple_strategy)}" ) output_tuple_strategy_children: list[OpStrategy] = [] for op_strategy in input_tuple_strategy.children: if not isinstance(op_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(op_strategy)}") reduce_dims = list(range(op_strategy.ndim)) output_strategy = common_reduction_strategy( op_strategy, reduce_dims, reduction_linear=True, reduction_op="sum", ) output_tuple_strategy_children.append(output_strategy) return TupleStrategy(output_tuple_strategy_children) @register_op_strategy( [aten._foreach_max.default], schema_info=RuntimeSchemaInfo(1, needs_pytree=True) ) def foreach_max_strategy(op_schema: OpSchema) -> TupleStrategy: """ Strategy for _foreach_max, which reduces each tensor in a list to its maximum value. """ args_schema = op_schema.args_schema input_tuple_strategy = args_schema[0] if not isinstance(input_tuple_strategy, TupleStrategy): raise AssertionError( f"Expected TupleStrategy, got {type(input_tuple_strategy)}" ) output_tuple_strategy_children: list[OpStrategy] = [] for op_strategy in input_tuple_strategy.children: if not isinstance(op_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(op_strategy)}") # Reduce all dimensions to get a scalar reduce_dims = list(range(op_strategy.ndim)) output_strategy = common_reduction_strategy( op_strategy, reduce_dims, reduction_linear=True, reduction_op="max", ) output_tuple_strategy_children.append(output_strategy) return TupleStrategy(output_tuple_strategy_children) @register_op_strategy( [ aten._linalg_svd.default, aten.linalg_qr.default, # TODO: The diagonal ops can have an improved sharding strategy for # shard placements that does not require redistributing to replicate. aten.diagonal_copy.default, aten.diag_embed.default, aten.diag.default, aten.diagonal.default, aten.tril.default, aten.triu.default, aten._linalg_eigh.default, ], schema_info=RuntimeSchemaInfo(1), ) def linalg_replicate_strategy(op_schema: OpSchema) -> OpStrategy: """ Since we do not have a simple way to compute some linear algebra operations like SVD or QR decomposition, always fall back to replicate. """ args_schema = op_schema.args_schema input_strategy = args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") mesh = input_strategy.mesh output_strategies: list[OpSpec] = [] for placement_strategy in input_strategy.strategies: replicate_placements = tuple(Replicate() for _ in range(mesh.ndim)) replicate_spec = DTensorSpec( mesh=mesh, placements=replicate_placements, tensor_meta=placement_strategy.output_spec.tensor_meta, ) redistribute_cost = [ generate_redistribute_costs(input_strategy, replicate_spec) ] replicate_strategy = OpSpec( output_specs=replicate_spec, input_specs=(replicate_spec,), redistribute_cost=redistribute_cost, ) output_strategies.append(replicate_strategy) return OpStrategy(output_strategies) # Maps each pooling op to its spatial rank (number of spatial dimensions). # Batched inputs have layout (N, C, *spatial) with ndim = spatial_rank + 2; # unbatched inputs drop the batch dim giving ndim = spatial_rank + 1. POOL_SPATIAL_RANK: dict[torch._ops.OpOverload, int] = { aten.avg_pool1d.default: 1, aten.avg_pool2d.default: 2, aten.avg_pool3d.default: 3, aten.adaptive_avg_pool1d.default: 1, aten._adaptive_avg_pool2d.default: 2, aten._adaptive_avg_pool3d.default: 3, aten.adaptive_max_pool1d.default: 1, aten.adaptive_max_pool2d.default: 2, aten.adaptive_max_pool3d.default: 3, aten.fractional_max_pool2d.default: 2, aten.fractional_max_pool3d.default: 3, aten.max_pool1d_with_indices.default: 1, aten.max_pool2d_with_indices.default: 2, aten.max_pool3d_with_indices.default: 3, } AVG_POOL_OPS = [ aten.avg_pool1d.default, aten.avg_pool2d.default, aten.avg_pool3d.default, aten.adaptive_avg_pool1d.default, aten._adaptive_avg_pool2d.default, aten._adaptive_avg_pool3d.default, ] MAX_POOL_OPS = [ aten.adaptive_max_pool1d.default, aten.adaptive_max_pool2d.default, aten.adaptive_max_pool3d.default, aten.fractional_max_pool2d.default, aten.fractional_max_pool3d.default, aten.max_pool1d_with_indices.default, aten.max_pool2d_with_indices.default, aten.max_pool3d_with_indices.default, ] @register_op_strategy( AVG_POOL_OPS + MAX_POOL_OPS, schema_info=RuntimeSchemaInfo(1), ) def pooling_strategy(op_schema: OpSchema) -> OpStrategy: input_strategy = cast(OpStrategy, op_schema.args_schema[0]) mesh = input_strategy.mesh num_outputs = 2 if op_schema.op in MAX_POOL_OPS else 1 num_inputs = len(op_schema.args_strategy) + len(op_schema.kwargs_strategy) n = num_outputs + num_inputs single_mesh_dim_strategies: list[PlacementList] = [ [Replicate()] * n, [Shard(0)] * n, ] # avg_pool is linear: Partial(sum) and Partial(avg) pass through unchanged. if op_schema.op in AVG_POOL_OPS: single_mesh_dim_strategies.append([Partial("sum")] * n) single_mesh_dim_strategies.append([Partial("avg")] * n) # S(1) is safe when dim 1 is the channel dim (pooling never touches it). # Batched inputs have layout (N, C, *spatial) with ndim = spatial_rank + 2. spatial_rank = POOL_SPATIAL_RANK[op_schema.op] is_batched = input_strategy.ndim >= spatial_rank + 2 if is_batched: single_mesh_dim_strategies.append([Shard(1)] * n) return expand_to_full_mesh_op_strategy( mesh, op_schema, single_mesh_dim_strategies, input_index=num_outputs ) @register_op_strategy( [aten._log_softmax.default, aten._softmax.default, aten._safe_softmax.default], schema_info=RuntimeSchemaInfo(1), ) def softmax_strategy(op_schema: OpSchema) -> OpStrategy: input_strategy, softmax_dim, *_ = op_schema.args_schema input_strategy = cast(OpStrategy, input_strategy) softmax_dim = cast(int, softmax_dim) softmax_dim = normalize_dim(softmax_dim, input_strategy.ndim) output_strategy = OpStrategy([]) for input_placement_strategy in input_strategy.strategies: redistribute_costs = [] input_src_spec = input_placement_strategy.output_spec # make sure input is replicated along the softmax dim input_target_spec = DTensorSpec( mesh=input_strategy.mesh, placements=replicate_reduction_dims( input_src_spec.placements, [softmax_dim] ), tensor_meta=input_src_spec.tensor_meta, ) redistribute_costs.append( generate_redistribute_costs(input_strategy, input_target_spec) ) output_target_spec = input_target_spec output_strategy.strategies.append( OpSpec( output_specs=output_target_spec, input_specs=[input_target_spec], redistribute_cost=redistribute_costs, ) ) return output_strategy @register_op_strategy( [ aten._log_softmax_backward_data.default, aten._softmax_backward_data.default, ], schema_info=RuntimeSchemaInfo(2), ) def softmax_backward_strategy(op_schema: OpSchema) -> OpStrategy: grad_out_strategy, out_strategy, softmax_dim, _ = op_schema.args_schema grad_out_strategy = cast(OpStrategy, grad_out_strategy) out_strategy = cast(OpStrategy, out_strategy) softmax_dim = cast(int, softmax_dim) softmax_dim = normalize_dim(softmax_dim, grad_out_strategy.ndim) grad_in_strategy = OpStrategy([]) for grad_out_placement_strat, out_placement_strat in zip( grad_out_strategy.strategies, out_strategy.strategies ): # follow the sharding of the grad_out or out depending on which has more shards grad_out_src_spec = grad_out_placement_strat.output_spec out_src_spec = out_placement_strat.output_spec src_spec = ( grad_out_src_spec if grad_out_src_spec.num_shards >= out_src_spec.num_shards else out_src_spec ) # make sure inputs are replicated along the softmax dim tgt_spec = DTensorSpec( mesh=grad_out_strategy.mesh, placements=replicate_reduction_dims(src_spec.placements, [softmax_dim]), ) new_grad_out_spec = DTensorSpec( mesh=tgt_spec.mesh, placements=tgt_spec.placements, tensor_meta=grad_out_src_spec.tensor_meta, ) new_out_spec = DTensorSpec( mesh=tgt_spec.mesh, placements=tgt_spec.placements, tensor_meta=out_src_spec.tensor_meta, ) redist_grad_out_cost = generate_redistribute_costs(grad_out_strategy, tgt_spec) redist_out_cost = generate_redistribute_costs(out_strategy, tgt_spec) grad_in_strategy.strategies.append( OpSpec( output_specs=tgt_spec, input_specs=(new_grad_out_spec, new_out_spec), redistribute_cost=[redist_grad_out_cost, redist_out_cost], ) ) return grad_in_strategy @register_op_strategy( [aten.nll_loss_forward.default, aten.nll_loss2d_forward.default], schema_info=RuntimeSchemaInfo(3), ) def nll_loss_forward_strategy(op_schema: OpSchema) -> OpStrategy: mesh = op_schema.get_mesh_from_args() if not len(op_schema.args_schema) == 5: raise AssertionError(f"Expected 5 args, got {len(op_schema.args_schema)}") ( input_strategy, target_strategy, weight_strategy, reduction, _, ) = op_schema.args_schema input_strategy = cast(OpStrategy, input_strategy) target_strategy = cast(OpStrategy, target_strategy) reduction = cast(int, reduction) input_shape = input_strategy.shape channel_dim = 1 if len(input_shape) >= 2 else 0 output_strategy = OpStrategy([]) for idx, input_placement_strategy in enumerate(input_strategy.strategies): op_args_target_specs = [] redistribute_costs = [] # make sure input is replicated along the channel dim input_src_spec = input_placement_strategy.output_spec input_expected_spec = DTensorSpec( mesh=mesh, placements=replicate_reduction_dims( input_src_spec.placements, [channel_dim] ), tensor_meta=input_src_spec.tensor_meta, ) op_args_target_specs.append(input_expected_spec) redistribute_costs.append( generate_redistribute_costs(input_strategy, input_expected_spec) ) # target doesn't have channel dim, and it follows input on other dims target_src_spec = target_strategy.strategies[idx].output_spec target_expected_spec = DTensorSpec( mesh=mesh, placements=_skip_dim(input_expected_spec.placements, channel_dim), tensor_meta=target_src_spec.tensor_meta, ) op_args_target_specs.append(target_expected_spec) redistribute_costs.append( generate_redistribute_costs(target_strategy, target_expected_spec) ) # weight tensor, if given, has to be a Tensor of size input_shape[channel_dim] # make sure it is replicated if weight_strategy is not None: if not isinstance(weight_strategy, OpStrategy): raise AssertionError( f"Expected OpStrategy, got {type(weight_strategy)}" ) weight_src_spec = weight_strategy.strategies[idx].output_spec weight_expected_spec = DTensorSpec( mesh=mesh, placements=_replicate_dims_start_at(weight_src_spec.placements), tensor_meta=weight_src_spec.tensor_meta, ) op_args_target_specs.append(weight_expected_spec) redistribute_costs.append( generate_redistribute_costs(weight_strategy, weight_expected_spec) ) if reduction == Reduction.NONE.value: output_expected_spec = target_expected_spec total_weight_expected_spec = DTensorSpec( mesh=mesh, placements=tuple([Replicate()] * mesh.ndim) ) else: if reduction == Reduction.MEAN.value: reduction_op = "avg" if not is_tensor_evenly_shardable( target_expected_spec.shape, target_expected_spec ): raise ValueError( "The intermediate results of nll_loss cannot be evenly sharded, \ resulting in biased mean result." ) else: # reduction == Reduction.SUM.value: reduction_op = "sum" reduce_dims = list(range(target_expected_spec.ndim)) reduce_dims_map = _infer_reduce_dims_map( reduce_dims, target_expected_spec.ndim, keep_dim=False ) out_placements = map_placements_after_reduction( target_expected_spec.placements, reduce_dims, reduce_dims_map, reduction_op, ) output_expected_spec = DTensorSpec( mesh=mesh, placements=out_placements, ) # whether reduction is sum or mean, the total weight has to be summed up if not replicated total_weight_placements = map_placements_after_reduction( target_expected_spec.placements, reduce_dims, reduce_dims_map, "sum", ) total_weight_expected_spec = DTensorSpec( mesh=mesh, placements=total_weight_placements, ) output_strategy.strategies.append( OpSpec( output_specs=(output_expected_spec, total_weight_expected_spec), input_specs=op_args_target_specs, redistribute_cost=redistribute_costs, ) ) return output_strategy @register_op_strategy( [aten.nll_loss_backward.default, aten.nll_loss2d_backward.default], schema_info=RuntimeSchemaInfo(4), ) def nll_loss_backward_strategy(op_schema: OpSchema) -> OpStrategy: # backward op does not need to validate the mesh since forward op has already done it mesh = op_schema.get_mesh_from_args(validate=False) if not len(op_schema.args_schema) == 7: raise AssertionError(f"Expected 7 args, got {len(op_schema.args_schema)}") ( grad_out_strategy, input_strategy, target_strategy, weight_strategy, reduction, _, total_weight_strategy, ) = op_schema.args_schema grad_out_strategy = cast(OpStrategy, grad_out_strategy) input_strategy = cast(OpStrategy, input_strategy) target_strategy = cast(OpStrategy, target_strategy) reduction = cast(int, reduction) total_weight_strategy = cast(OpStrategy, total_weight_strategy) input_shape = input_strategy.shape channel_dim = 1 if len(input_shape) >= 2 else 0 grad_in_strategy = OpStrategy([]) for idx, input_placement_strategy in enumerate(input_strategy.strategies): op_args_target_specs = [] redistribute_costs = [] # make sure input is replicated along the channel dim input_src_spec = input_placement_strategy.output_spec input_expected_spec = DTensorSpec( mesh=mesh, placements=replicate_reduction_dims( input_src_spec.placements, [channel_dim] ), tensor_meta=input_src_spec.tensor_meta, ) op_args_target_specs.append(input_expected_spec) redistribute_costs.append( generate_redistribute_costs(input_strategy, input_expected_spec) ) # target doesn't have channel dim, and it follows input on other dims target_src_spec = target_strategy.strategies[idx].output_spec target_expected_spec = DTensorSpec( mesh=mesh, placements=_skip_dim(input_expected_spec.placements, channel_dim), tensor_meta=target_src_spec.tensor_meta, ) op_args_target_specs.append(target_expected_spec) redistribute_costs.append( generate_redistribute_costs(target_strategy, target_expected_spec) ) # grad_out follows target if there is no reduction; # otherwise, it should be a replicated scalar. grad_out_src_spec = grad_out_strategy.strategies[idx].output_spec if reduction == Reduction.NONE.value: grad_out_expected_spec = target_expected_spec else: grad_out_expected_spec = DTensorSpec( mesh=mesh, placements=_replicate_dims_start_at(grad_out_src_spec.placements), tensor_meta=grad_out_src_spec.tensor_meta, ) op_args_target_specs.insert(0, grad_out_expected_spec) redistribute_costs.insert( 0, generate_redistribute_costs(grad_out_strategy, grad_out_expected_spec) ) # weight tensor, if given, has to be a Tensor of size input_shape[channel_dim] # make sure it is replicated if weight_strategy is not None: if not isinstance(weight_strategy, OpStrategy): raise AssertionError( f"Expected OpStrategy, got {type(weight_strategy)}" ) weight_src_spec = weight_strategy.strategies[idx].output_spec weight_expected_spec = DTensorSpec( mesh=mesh, placements=_replicate_dims_start_at(weight_src_spec.placements), tensor_meta=weight_src_spec.tensor_meta, ) op_args_target_specs.append(weight_expected_spec) redistribute_costs.append( generate_redistribute_costs(weight_strategy, weight_expected_spec) ) # total_weight is only used by the backward kernel for reduction='mean'. # For reduction='sum' or 'none', it is unused, so no redistribution needed. total_weight_src_spec = total_weight_strategy.strategies[idx].output_spec if reduction == Reduction.MEAN.value: total_weight_expected_spec = DTensorSpec( mesh=mesh, placements=_replicate_dims_start_at(total_weight_src_spec.placements), tensor_meta=total_weight_src_spec.tensor_meta, ) else: total_weight_expected_spec = total_weight_src_spec op_args_target_specs.append(total_weight_expected_spec) redistribute_costs.append( generate_redistribute_costs( total_weight_strategy, total_weight_expected_spec ) ) grad_in_expected_spec = input_expected_spec grad_in_strategy.strategies.append( OpSpec( output_specs=grad_in_expected_spec, input_specs=op_args_target_specs, redistribute_cost=redistribute_costs, ) ) return grad_in_strategy @register_single_dim_strategy( [aten.native_layer_norm.default], schema_info=RuntimeSchemaInfo(1), ) def layer_norm_single_dim_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: input_meta = args_schema[0] normalized_shape = args_schema[1] weight_meta = args_schema[2] bias_meta = args_schema[3] axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape)) strategies: list[list[Placement | _ShardingPlaceholder]] = [] for dim in range(axis): # [out, mean, rstd, input, weight?, bias?] rule: list[Placement | _ShardingPlaceholder] = [ _ShardingPlaceholder(dim), # out _ShardingPlaceholder(dim), # mean _ShardingPlaceholder(dim), # rstd _ShardingPlaceholder(dim), # input ] if weight_meta is not None: rule.append(Replicate()) if bias_meta is not None: rule.append(Replicate()) strategies.append(rule) return strategies @register_single_dim_strategy( [aten._fused_rms_norm.default], schema_info=RuntimeSchemaInfo(1), ) def rms_norm_single_dim_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: input_meta = args_schema[0] normalized_shape = args_schema[1] weight_meta = args_schema[2] axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape)) strategies: list[list[Placement | _ShardingPlaceholder]] = [] for dim in range(axis): # [out, rrms, input, weight?] rule: list[Placement | _ShardingPlaceholder] = [ _ShardingPlaceholder(dim), # out _ShardingPlaceholder(dim), # rrms _ShardingPlaceholder(dim), # input ] if weight_meta is not None: rule.append(Replicate()) strategies.append(rule) return strategies @register_single_dim_strategy( [aten.native_layer_norm_backward.default], schema_info=RuntimeSchemaInfo(2), ) def layer_norm_bwd_single_dim_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder | None]]: input_meta = args_schema[1] normalized_shape = args_schema[2] # mean = args_schema[3], rstd = args_schema[4] weight_meta = args_schema[5] bias_meta = args_schema[6] axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape)) strategies: list[list[Placement | _ShardingPlaceholder | None]] = [] for dim in range(axis): # outputs: [d_input, d_weight, d_bias] — always 3 per schema # d_weight/d_bias use None when weight/bias are None rule: list[Placement | _ShardingPlaceholder | None] = [ _ShardingPlaceholder(dim), # d_input Partial("sum") if weight_meta is not None else None, # d_weight Partial("sum") if bias_meta is not None else None, # d_bias ] # inputs: [grad_out, input, mean, rstd, weight?, bias?] rule.extend( [ _ShardingPlaceholder(dim), # grad_out _ShardingPlaceholder(dim), # input _ShardingPlaceholder(dim), # mean _ShardingPlaceholder(dim), # rstd ] ) if weight_meta is not None: rule.append(Replicate()) if bias_meta is not None: rule.append(Replicate()) strategies.append(rule) return strategies @register_single_dim_strategy( [aten._fused_rms_norm_backward.default], schema_info=RuntimeSchemaInfo(2), ) def rms_norm_bwd_single_dim_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder | None]]: input_meta = args_schema[1] normalized_shape = args_schema[2] # rstd = args_schema[3] weight_meta = args_schema[4] axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape)) strategies: list[list[Placement | _ShardingPlaceholder | None]] = [] for dim in range(axis): # outputs: [d_input, d_weight] — always 2 per schema # d_weight uses None when weight is None # inputs: [grad_out, input, rstd, weight?] rule: list[Placement | _ShardingPlaceholder | None] = [ _ShardingPlaceholder(dim), # d_input Partial("sum") if weight_meta is not None else None, # d_weight _ShardingPlaceholder(dim), # grad_out _ShardingPlaceholder(dim), # input _ShardingPlaceholder(dim), # rstd ] if weight_meta is not None: rule.append(Replicate()) strategies.append(rule) return strategies def sort_strategy(op_schema: OpSchema, sort_dim: int) -> OpStrategy: input_strategy = cast(OpStrategy, op_schema.args_schema[0]) sort_dim = normalize_dim(sort_dim, input_strategy.ndim) single_mesh_dim_strategies = [] all_replicate: PlacementList = [Replicate()] * 3 single_mesh_dim_strategies.append(all_replicate) for dim in range(input_strategy.ndim): if dim != sort_dim: dim_shardings: PlacementList = [Shard(dim)] * 3 single_mesh_dim_strategies.append(dim_shardings) return expand_to_full_mesh_op_strategy( input_strategy.mesh, op_schema, single_mesh_dim_strategies, input_index=2 ) @register_op_strategy( [aten.topk.default], schema_info=RuntimeSchemaInfo(2), ) def topk_strategy(op_schema: OpSchema) -> OpStrategy: topk_dim = ( cast(int, op_schema.args_schema[2]) if len(op_schema.args_schema) > 2 else -1 ) return sort_strategy(op_schema, topk_dim) @register_op_strategy( aten.sort.default, schema_info=RuntimeSchemaInfo( 1, ), ) def sort_default_strategy(op_schema: OpSchema) -> OpStrategy: # mostly copy paste from topk_strategy input_strategy = op_schema.args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") sort_dim = -1 if len(op_schema.args_schema) > 1: sort_dim = cast(int, op_schema.args_schema[1]) return sort_strategy(op_schema, sort_dim) @register_op_strategy( aten.sort.stable, schema_info=RuntimeSchemaInfo( 1, static_kwargkey=["dim", "descending", "stable"], ), ) def sort_stable_strategy(op_schema: OpSchema) -> OpStrategy: # mostly copy paste from topk_strategy input_strategy = op_schema.args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") sort_dim = -1 if "dim" in op_schema.kwargs_schema: sort_dim = cast(int, op_schema.kwargs_schema["dim"]) return sort_strategy(op_schema, sort_dim) @register_op_strategy( [aten.histc.default], # strategy choice depends on the value of 'min' and 'max' kwargs, which are position 2 and 3 schema_info=RuntimeSchemaInfo(2), ) def histc_strategy(op_schema: OpSchema) -> OpStrategy: input_strategy = cast(OpStrategy, op_schema.args_schema[0]) single_mesh_dim_strategies: list[PlacementList] = [] single_mesh_dim_strategies.append([Replicate(), Replicate()]) # histc can support sharded input and partial output on any input dim, provided the min and max # values are user-specified. If not user-specified, the true min and max of the data in each local # tensor will be used to compute bin boundaries, which will not be the same across ranks, leading to # an incorrect final result if len(op_schema.args_schema) == 4: for dim in range(input_strategy.ndim): dim_shardings: PlacementList = [Partial(), Shard(dim)] single_mesh_dim_strategies.append(dim_shardings) return expand_to_full_mesh_op_strategy( input_strategy.mesh, op_schema, single_mesh_dim_strategies ) @register_op_strategy( [aten.logsumexp.default], schema_info=RuntimeSchemaInfo( # static_argnum is the position where non-Tensor args beings. static_argnum=1, # static_kwargkey is the name of kwargs to hash (which determines # whether sharding prop can be cached). static_kwargkey=["keepdim"], ), ) def logsumexp_strategy(op_schema: OpSchema) -> OpStrategy: """Implements the sharding propagation strategy for logsumexp.""" # args_schema contains all but the DTensor args (e.g., dim, keepdim). args_schema = op_schema.args_schema if not len(args_schema) > 1: raise AssertionError( f"Expected more than 1 arg (input and dim are required), got {len(args_schema)}" ) input_strategy = args_schema[0] if not isinstance(input_strategy, OpStrategy): raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") dims_arg = args_schema[1] reduce_dims = _infer_reduction_dims(dims_arg, input_strategy.ndim) if reduce_dims is None: raise AssertionError("Expected reduce_dims to not be None") keep_dim = cast(bool, op_schema.kwargs_schema.get("keepdim", False)) return common_reduction_strategy( input_strategy, reduce_dims, keep_dim=keep_dim, reduction_linear=False, ) _LINALG_NUM_PLACEMENTS = { # 1 in 1 out aten.cholesky.default: 2, aten.cholesky_inverse.default: 2, aten.linalg_matrix_exp.default: 2, # 2 in 1 out aten.cholesky_solve.default: 3, aten.linalg_householder_product.default: 3, aten.linalg_solve_triangular.default: 3, # 3 in 1 out aten.linalg_ldl_solve.default: 4, aten.linalg_lu_solve.default: 4, aten.ormqr.default: 4, # 1 in 2 out aten.geqrf.default: 3, aten.linalg_cholesky_ex.default: 3, aten.linalg_eig.default: 3, aten.linalg_inv_ex.default: 3, # 2 in 2 out aten.triangular_solve.default: 4, # 1 in 3 out aten._linalg_det.default: 4, aten.linalg_ldl_factor_ex.default: 4, aten.linalg_lu.default: 4, aten.linalg_lu_factor_ex.default: 4, # 2 in 3 out aten.lu_unpack.default: 5, # 1 in 4 out aten._linalg_slogdet.default: 5, # 2 in 4 out aten._linalg_solve_ex.default: 6, # 1 in aten._linalg_check_errors.default: 1, } def _linalg_batch_dim_strategies( ndim: int, n_placements: int ) -> list[list[Placement | _ShardingPlaceholder]]: """Build single-dim strategies for linalg ops that operate on the last 1-2 dims. Returns sharding on each batch dim (all dims except the last 2), with all outputs and inputs sharded on the same dim. """ strategies: list[list[Placement | _ShardingPlaceholder]] = [] for dim in range(ndim - 2): strategies.append([_ShardingPlaceholder(dim)] * n_placements) return strategies def _get_ndim(tensor_meta: Any) -> int: if not isinstance(tensor_meta, TensorMeta): raise AssertionError(f"Expected TensorMeta, got {type(tensor_meta)}") return len(tensor_meta.shape) @register_single_dim_strategy( [ aten.cholesky.default, aten.cholesky_inverse.default, aten.linalg_matrix_exp.default, aten.cholesky_solve.default, aten.linalg_householder_product.default, aten.linalg_solve_triangular.default, aten.linalg_ldl_solve.default, aten.linalg_lu_solve.default, aten.ormqr.default, aten.geqrf.default, aten.linalg_cholesky_ex.default, aten.linalg_eig.default, aten.linalg_inv_ex.default, aten.triangular_solve.default, aten._linalg_det.default, aten.linalg_ldl_factor_ex.default, aten.linalg_lu.default, aten.linalg_lu_factor_ex.default, aten.lu_unpack.default, aten._linalg_slogdet.default, aten._linalg_solve_ex.default, aten._linalg_check_errors.default, ], schema_info=RuntimeSchemaInfo(1), ) def linalg_batch_dim_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: ndim = _get_ndim(args_schema[0]) if op not in _LINALG_NUM_PLACEMENTS: raise AssertionError(f"Expected op in _LINALG_NUM_PLACEMENTS, got {op}") n_placements = _LINALG_NUM_PLACEMENTS[op] strategies = _linalg_batch_dim_strategies(ndim, n_placements=n_placements) if op == aten.linalg_solve_triangular.default: # solve_triangular(A, B) -> result: linear in B strategies.append([Partial(), Replicate(), Partial()]) strategies.append([Partial("avg"), Replicate(), Partial("avg")]) # A replicated, B sharded on batch dims (B may have more batch dims than A) ndim_b = _get_ndim(args_schema[1]) for dim in range(ndim_b - 2): strategies.append( [_ShardingPlaceholder(dim), Replicate(), _ShardingPlaceholder(dim)] ) elif op == aten.cholesky_solve.default: # cholesky_solve(B, A) -> result (B is arg0) strategies.append([Partial(), Partial(), Replicate()]) elif op == aten.linalg_lu_solve.default: # linalg_lu_solve(LU, pivots, B) -> result strategies.append([Partial(), Replicate(), Replicate(), Partial()]) elif op == aten.linalg_ldl_solve.default: # linalg_ldl_solve(LD, pivots, B) -> result strategies.append([Partial(), Replicate(), Replicate(), Partial()]) elif op == aten.ormqr.default: # ormqr(a, tau, C) -> result (linear in C) strategies.append([Partial(), Replicate(), Replicate(), Partial()]) elif op == aten._linalg_solve_ex.default: # _linalg_solve_ex(A, B) -> (result, LU, pivots, info) strategies.append( [Partial(), Replicate(), Replicate(), Replicate(), Replicate(), Partial()] ) return strategies # linalg_pinv has optional tensor kwargs atol, rtol (scalar tensors when present). # Schema: (Tensor self, *, Tensor? atol=None, Tensor? rtol=None, bool hermitian) -> Tensor # When atol/rtol are None, num_inputs=1; when present, they add to num_inputs. @register_single_dim_strategy( [aten.linalg_pinv.atol_rtol_tensor], schema_info=RuntimeSchemaInfo(1), ) def linalg_pinv_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: ndim = _get_ndim(args_schema[0]) # Count optional tensor kwargs that are actually present extra_tensors = sum( isinstance(kwargs_schema.get(k), TensorMeta) for k in ("atol", "rtol") ) strategies: list[list[Placement | _ShardingPlaceholder]] = [] for dim in range(ndim - 2): s: list[Placement | _ShardingPlaceholder] = [ _ShardingPlaceholder(dim), _ShardingPlaceholder(dim), ] # atol, rtol are scalar tensors — always Replicate s.extend([Replicate()] * extra_tensors) strategies.append(s) return strategies # linalg_cross is pointwise on every dim except the cross-product dim (which # must be size 3). Shard on any other dim. @register_single_dim_strategy( [aten.linalg_cross.default], schema_info=RuntimeSchemaInfo(1), ) def linalg_cross_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: ndim = _get_ndim(args_schema[0]) cross_dim = kwargs_schema.get("dim", -1) % ndim strategies: list[list[Placement | _ShardingPlaceholder]] = [] for dim in range(ndim): if dim == cross_dim: continue strategies.append([_ShardingPlaceholder(dim)] * 3) return strategies # --------------------------------------------------------------------------- # Interpolation / upsample / pooling ops # # These ops operate on spatial dims and are safely shardable on batch (dim 0) # and channel (dim 1). grid_sampler is batch-only because the grid tensor has # no channel dimension. # --------------------------------------------------------------------------- @register_single_dim_strategy( [ aten.upsample_nearest1d.default, aten.upsample_nearest2d.default, aten.upsample_nearest3d.default, aten._upsample_nearest_exact1d.default, aten._upsample_nearest_exact2d.default, aten._upsample_nearest_exact3d.default, aten._upsample_bilinear2d_aa.default, aten.upsample_bicubic2d.default, aten.upsample_bilinear2d.default, aten.upsample_linear1d.default, aten.upsample_trilinear3d.default, ], schema_info=RuntimeSchemaInfo(1), ) def interp_upsample_1out_1in_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: # 1 output + 1 input = 2 placements; shard on batch (0) and channel (1) # Upsample is a linear transformation so Partial(sum/avg) is valid. return [ [_ShardingPlaceholder(0)] * 2, [_ShardingPlaceholder(1)] * 2, [Partial("sum"), Partial("sum")], [Partial("avg"), Partial("avg")], ] @register_single_dim_strategy( [ aten.max_unpool2d.default, aten.max_unpool3d.default, aten._adaptive_avg_pool2d_backward.default, ], schema_info=RuntimeSchemaInfo(1), ) def interp_pool_1out_2in_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: # 1 output + 2 inputs = 3 placements; shard on batch (0) and channel (1) return [ [_ShardingPlaceholder(0)] * 3, [_ShardingPlaceholder(1)] * 3, ] @register_single_dim_strategy( [aten.max_pool2d_with_indices_backward.default], schema_info=RuntimeSchemaInfo(1), ) def pool_backward_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: # max_pool2d_with_indices_backward(grad_output, self, ..., indices) -> grad_input # 1 output + 3 tensor inputs = 4 placements # Order: [output, grad_output, self, indices] input_meta = cast(TensorMeta, args_schema[0]) strategies: list[list[Placement | _ShardingPlaceholder]] = [ [_ShardingPlaceholder(0)] * 4, ] if len(input_meta.shape) >= 4: # batched: (N, C, H, W) strategies.append([_ShardingPlaceholder(1)] * 4) # The backward is linear in grad_output, so P(sum/avg) pass through. # indices must be replicated (integer positions, not reducible). # self is only used for shape, so replicate it too. r = Replicate() for reduce_op in ("sum", "avg"): p = Partial(reduce_op) strategies.append([p, p, r, r]) return strategies @register_single_dim_strategy( [aten.grid_sampler_2d.default, aten.grid_sampler_3d.default], schema_info=RuntimeSchemaInfo(1), ) def grid_sampler_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: # grid_sampler_{2,3}d(input[N,C,...], grid[N,...,{2,3}]) -> output[N,C,...] # grid has no channel dim, so only batch sharding applies to both inputs. # Linear in input: P(sum/avg) on input with replicated grid is valid. return [ [_ShardingPlaceholder(0)] * 3, [Partial("sum"), Partial("sum"), Replicate()], [Partial("avg"), Partial("avg"), Replicate()], ] @register_single_dim_strategy( [aten.grid_sampler_2d_backward.default, aten.grid_sampler_3d_backward.default], schema_info=RuntimeSchemaInfo(1), ) def grid_sampler_backward_strategy( op: torch._ops.OpOverload, args_schema: tuple[Any, ...], kwargs_schema: dict[str, Any], ) -> list[list[Placement | _ShardingPlaceholder]]: # grid_sampler_{2,3}d_backward: 2 outputs (grad_input, grad_grid) + 3 inputs = 5 placements, batch-only return [[_ShardingPlaceholder(0)] * 5]