# Copyright (c) Meta Platforms, Inc. and affiliates import functools import logging from collections.abc import Callable, Sequence from typing import Any, cast, Optional, TypeAlias, TypeVar, Union import torch from torch._ops import OpOverload from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta from torch.distributed.tensor._op_schema import ( ArgsType, KwargsType, OpSchema, OpStrategy, PlacementList, RuntimeSchemaInfo, StrategyType, TupleStrategy, ) from torch.distributed.tensor.device_mesh import DeviceMesh from torch.distributed.tensor.placement_types import ( _StridedShard, Placement, Replicate, Shard, ) from torch.utils._pytree import tree_map_only logger = logging.getLogger(__name__) class _ShardingPlaceholder: """ A placeholder for a sharding placement that has a specified tensor dim, but the other metadata (e.g. split factor if it's a StridedShard) will be filled in later. """ dim: int def __init__(self, dim: int): self.dim = dim def __repr__(self) -> str: return f"_ShardingPlaceholder(dim={self.dim})" _StrategyTypeT = TypeVar("_StrategyTypeT", bound=StrategyType) _PlacementT = TypeVar("_PlacementT", bound=Placement) _ShardingPlaceholderT = TypeVar("_ShardingPlaceholderT", bound=_ShardingPlaceholder) _SingleDimStrategyFunc: TypeAlias = Callable[ [OpOverload, ArgsType, KwargsType], list[list[_PlacementT | _ShardingPlaceholderT]] ] _ExpandedSingleDimStrategyFunc: TypeAlias = Callable[ [OpOverload, ArgsType, KwargsType], _StrategyTypeT ] def _insert_single_dim_replication_strategy( single_dim_strategies_with_placeholders: list[ list[Placement | _ShardingPlaceholder] ], num_input_tensors: int, ) -> list[list[Placement | _ShardingPlaceholder]]: """ Inserts the [Replicate(), Replicate(), ...] strategy after asserting that such strategy does not yet exist. """ for strategy in single_dim_strategies_with_placeholders: assert not all(isinstance(p, Replicate) for p in strategy) single_dim_strategies_with_placeholders.append( [Replicate()] * (1 + num_input_tensors) ) return single_dim_strategies_with_placeholders def _fill_single_dim_strategy_placeholders( unique_input_placements: set[Placement], single_dim_strategies_with_placeholders: list[ list[Placement | _ShardingPlaceholder] ], ) -> list[list[Placement]]: """ Replace any _ShardingPlaceholder with the specific Sharding types used by the inputs in op_schema. Supports implicit replication. Example: single_dim_strategies_with_placeholders = [[Partial(), _ShardingPlaceholder(1), _ShardingPlaceholder(0)]] input0: Shard(0) input1: StridedShard(1, split_factor=2) returns: [ [Partial(), Shard(1), Shard(0)], [Partial(), StridedShard(1, split_factor=2), StridedShard(0, split_factor=2)], [Replicate(), Replicate(), Replicate()] ] """ shard_builders: dict[str, Callable[[int], Placement]] = {} for placement in unique_input_placements: if isinstance(placement, _StridedShard): key = f"StridedShard(sf={placement.split_factor})" if key not in shard_builders: shard_builders[key] = functools.partial( _StridedShard, split_factor=placement.split_factor ) elif isinstance(placement, Shard): key = "Shard()" if key not in shard_builders: shard_builders[key] = lambda tensor_dim: Shard(tensor_dim) # if any of the placements is a placeholder, we need to expand the strategy # to all possible combinations of placements expanded_strategies_over_one_mesh_dim: list[list[Placement]] = [] for s in single_dim_strategies_with_placeholders: if any(isinstance(p, _ShardingPlaceholder) for p in s): for shard_builder in shard_builders.values(): expanded_strategy: list[Placement] = [] for maybe_placeholder in s: if isinstance(maybe_placeholder, _ShardingPlaceholder): # we combine the tensor dim to shard from the placeholder # with other metadata (e.g. split_factor) from the sharding class expanded_strategy.append(shard_builder(maybe_placeholder.dim)) else: assert isinstance(maybe_placeholder, Placement) expanded_strategy.append(maybe_placeholder) expanded_strategies_over_one_mesh_dim.append(expanded_strategy) else: assert all(isinstance(p, Placement) for p in s) expanded_strategies_over_one_mesh_dim.append(cast(list[Placement], (s))) return expanded_strategies_over_one_mesh_dim def _get_unique_placements(op_schema: OpSchema) -> set[Placement]: unique_placements = set() def _update_placements(obj: Any): if isinstance(obj, DTensorSpec): unique_placements.update(obj.placements) elif isinstance(obj, OpStrategy): assert len(obj.strategies) == 1 unique_placements.update(obj.strategies[0].output_spec.placements) elif isinstance(obj, TupleStrategy): for child in obj.children: _update_placements(child) for obj in op_schema.args_schema: _update_placements(obj) return unique_placements def _get_num_tensor_inputs(op_schema: OpSchema) -> int: num_inputs = 0 for obj in op_schema.args_schema: if isinstance(obj, OpStrategy): num_inputs += 1 elif isinstance(obj, TupleStrategy): num_inputs += len(obj.children) return num_inputs def _expand_single_dim_strategy_to_mesh( mesh: DeviceMesh, op_schema: OpSchema, single_dim_strategy: _SingleDimStrategyFunc, output_tensor_meta: TensorMeta | Sequence[TensorMeta | None], ) -> _ExpandedSingleDimStrategyFunc: """ Expands the single_mesh_dim impl across all mesh dims, and expands ShardingPlacholder into all sharding types used by inputs. This supports functional correctness but will generate all possible combinations, which is prohibitively expensive for larger numbers of mesh dimensions. The expanded_strategy function accesses both the args_schema/kwargs_schema, which contains TensorMeta in place of tensor arguments, but also the op_schema which contains OpStrategy in place of Tensor args. Args: output_tensor_meta: tensor metadata for the output(s), precomputed during sharding prop """ # Note: circular import, failed to untangle with #168221, reverted from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy @functools.lru_cache def _create_expanded_strategy( op_schema: OpSchema, output_tensor_meta: TensorMeta | Sequence[TensorMeta | None], ) -> Callable[[OpOverload, ArgsType, KwargsType], StrategyType]: def expanded_strategy( op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType ) -> StrategyType: # Note: op_schema vs [args_schema, kwargs_schema] # ----------------------------------------------- # Inside `expanded_strategy function we purposefully have access to 2 similar structures. # 1) (op, args_schema, kwargs_schema): This is all the single_dim_strategy is allowed to see. # importantly, it does not contain information about input placements or meshes - just TensorMeta. # 2) op_schema - captured from the parent scope, this contains the input placement and mesh info, needed # to actually perform expansion. unique_input_placements = _get_unique_placements(op_schema) num_inputs = _get_num_tensor_inputs(op_schema) # Note: Trees vs Flat Lists # ------------------------- # op_schema.args_schema may contain a TupleStrategy with child strategies for List[Tensor] inputs. # args_schema has corresponding TupleStrategy, but with TensorSpec in place of child strategies. # CURRENTLY: single_dim_strategy will return a flat list of Placements for each strategy, where any # input tuple strategies have been inlined. I'm not sure if we want to keep doing this, or preserve a pytree # structure here. I'm following the convention in the current DTensor sharding strategies for now. # Inside expanded_strategy, we need to carefully align the OpStrategies / Specs from op_schema which are _not_ # flattened, with the flat Placement list returned from single_dim strategy. strategies_over_one_mesh_dim = single_dim_strategy( op, args_schema, kwargs_schema ) strategies_over_one_mesh_dim = _insert_single_dim_replication_strategy( strategies_over_one_mesh_dim, num_inputs ) expanded_strategies_over_one_mesh_dim = ( _fill_single_dim_strategy_placeholders( unique_input_placements, strategies_over_one_mesh_dim ) ) # Note: does not support `allow_unbacked_sharding` which is needed by matmul rules for some compile test # currently, we should probably change that test though, since it seems wrong to me to allow sharding unbacked # dims return expand_to_full_mesh_op_strategy( mesh, op_schema, cast(list[PlacementList], expanded_strategies_over_one_mesh_dim), output_tensor_meta=output_tensor_meta, ) return expanded_strategy def _translate_foreach_op_schema( op_schema: OpSchema, output_tensor_meta: Sequence[TensorMeta], index: int ) -> tuple[OpSchema, TensorMeta]: """Translate foreach op to per-element version of schema.""" op_parts = str(op_schema.op).split(".") base_op_name = op_parts[-2].replace("_foreach_", "") foreach_variant = op_parts[-1] # select per-element inputs, outputs target_args, target_kwargs = tree_map_only( TupleStrategy, lambda x: x.children[index], (op_schema.args_schema, op_schema.kwargs_schema), is_leaf=lambda x: isinstance(x, TupleStrategy), ) target_output_meta = output_tensor_meta[index] # figure out target op variant variant_map = { "List": "Tensor", "ScalarList": "Scalar", "Scalar": "Scalar", "Tensor": "Tensor", "default": "default", } target_variant = ( "default" if len(target_args) == 1 else variant_map.get(foreach_variant, "default") ) # this seems a bit messy base_op = getattr(torch.ops.aten, base_op_name) target_op = ( getattr(base_op, target_variant) if target_variant in base_op.overloads() else base_op.default ) op_schema = OpSchema( target_op, # type: ignore[arg-type] args_schema=tuple(target_args), kwargs_schema=op_schema.kwargs_schema, ) return op_schema, target_output_meta def expanded_foreach_strategy( op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType ) -> StrategyType: tensorlist_len: int | None = None for i, obj in enumerate(op_schema.args_schema): if isinstance(obj, TupleStrategy): if tensorlist_len is None: tensorlist_len = len(obj.children) elif len(obj.children) != tensorlist_len: raise AssertionError( f"Expected {tensorlist_len} children in index {i}, but found {len(obj.children)}." ) if tensorlist_len is None: raise AssertionError("Must have at least one tuple input to a foreach op") child_strategies: list[StrategyType] = [] for tensorlist_i in range(tensorlist_len): per_index_schema, per_index_output_meta = _translate_foreach_op_schema( op_schema, output_tensor_meta, # type: ignore[arg-type] tensorlist_i, ) per_index_strategy = _create_expanded_strategy( per_index_schema, per_index_output_meta ) child_strategies.append( per_index_strategy( op, per_index_schema.args_meta, per_index_schema.kwargs_meta ) ) return TupleStrategy(children=child_strategies) # TODO maybe this could be helped by adding a new 'tag' to the OpOverload? if op_schema.op.name().startswith("aten::_foreach_"): return expanded_foreach_strategy return _create_expanded_strategy(op_schema, output_tensor_meta) def register_single_dim_strategy( op: Union[torch._ops.OpOverload, list[torch._ops.OpOverload]], schema_info: Optional[RuntimeSchemaInfo] = None, ) -> Callable[[_SingleDimStrategyFunc], _SingleDimStrategyFunc]: """ Registers a single_dim_strategy function for the given op. A single_dim_strategy enumerates all the non-trivial sharding specifications for the operator over a single mesh dim. Since it generates the full set of valid shardings regardless of the given input placements, tensor inputs are represented via TensorMetas instead of OpStrategies like in op_strategy functions. TensorMeta inputs, along with all other op inputs (int, float, bool, etc) inform which sharding placements are valid. Single-dim strategies are fed into infra that expands them over multiple mesh dims and fills sharding placeholders with concrete sharding types. Single-dim-strategies should not list the trivial "all Replicate" rule, which is assumed for all ops. Sharding placeholders should be used to represent generic sharding rules in single_dim_strategies. For example, most operators that support sharded inputs and outputs do not care how the shards are laid out globally, as long as they are consistent. It is just as valid to run a matmul on (Shard(0), Replicate -> Shard(0)) as on (StridedShard(0), Replicate -> StridedShard(0)). For this reason, we use a 'ShardingPlaceholder' to represent all generic sharding types. Placeholders will be filled by concrete sharding types seen in runtime inputs, not all types known to DTensor. """ # Note: circular import, failed to untangle with #168221, reverted from torch.distributed.tensor._api import DTensor from torch.distributed.tensor._ops.utils import _get_registration_wrapper # For every ATen op that accepts any args in this list, # the arg itself can impact the strides (and potentially the sharding strategy) # of the output tensor. # thus, we will detect ATen schemas with any of these args and ensure # that they get specialized here. arg_names_that_require_specializing_cache_strategy = [ "memory_format", ] return _get_registration_wrapper( DTensor._op_dispatcher.sharding_propagator.register_single_dim_op_strategy, op, schema_info, arg_names_that_require_specializing_cache_strategy, )