已合并
Optimize the patch for DTensor #34610
yc_233创建于 4月28日
Optimize the patch for DTensor #34610
已合并
共 3 个文件变更+127-121
| @@ -0,0 +1,127 @@ | |||
| 1 | +# Regression tests: DTensor kwargs strategy path should work with native PyTorch 2.11+ support. | ||
| 2 | +import torch | ||
| 3 | +from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta | ||
| 4 | +from torch.distributed.tensor._op_schema import ( | ||
| 5 | + OpSchema, | ||
| 6 | + OpSpec, | ||
| 7 | + OpStrategy, | ||
| 8 | + RuntimeSchemaInfo, | ||
| 9 | + TupleStrategy, | ||
| 10 | +) | ||
| 11 | +from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy | ||
| 12 | +from torch.distributed.tensor.device_mesh import DeviceMesh | ||
| 13 | +from torch.distributed.tensor.placement_types import Replicate | ||
| 14 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +def _cpu_mesh(size: int = 4) -> DeviceMesh: | ||
| 18 | + return DeviceMesh("cpu", torch.arange(size)) | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +def _replicate_strategy(mesh: DeviceMesh, shape: tuple[int, ...]) -> OpStrategy: | ||
| 22 | + placements = tuple(Replicate() for _ in range(mesh.ndim)) | ||
| 23 | + strides: list[int] = [] | ||
| 24 | + st = 1 | ||
| 25 | + for s in reversed(shape): | ||
| 26 | + strides.append(st) | ||
| 27 | + st *= s | ||
| 28 | + strides_t = tuple(reversed(strides)) | ||
| 29 | + meta = TensorMeta(shape=torch.Size(shape), stride=strides_t, dtype=torch.float32) | ||
| 30 | + spec = DTensorSpec(mesh=mesh, placements=placements, tensor_meta=meta) | ||
| 31 | + op_spec = OpSpec(output_specs=spec, input_specs=(spec,), redistribute_cost=[[0.0]]) | ||
| 32 | + return OpStrategy([op_spec]) | ||
| 33 | + | ||
| 34 | + | ||
| 35 | +class TestNativeKwargsStrategyAndExpand(TestCase): | ||
| 36 | + """Mimics torch_npu _matrix_ops / _math_ops usage: args_strategy + kwargs_strategy length drives expand.""" | ||
| 37 | + | ||
| 38 | + def setUp(self): | ||
| 39 | + from torch.testing._internal.distributed.fake_pg import FakeStore | ||
| 40 | + | ||
| 41 | + super().setUp() | ||
| 42 | + self.ws = 4 | ||
| 43 | + torch.distributed.init_process_group( | ||
| 44 | + backend="fake", rank=0, world_size=self.ws, store=FakeStore() | ||
| 45 | + ) | ||
| 46 | + | ||
| 47 | + def tearDown(self): | ||
| 48 | + torch.distributed.destroy_process_group() | ||
| 49 | + super().tearDown() | ||
| 50 | + | ||
| 51 | + def test_op_schema_kwargs_strategy_flat(self): | ||
| 52 | + mesh = _cpu_mesh(self.ws) | ||
| 53 | + a = _replicate_strategy(mesh, (2, 3)) | ||
| 54 | + b = _replicate_strategy(mesh, (2, 3)) | ||
| 55 | + schema = OpSchema( | ||
| 56 | + torch.ops.aten.add.Tensor, | ||
| 57 | + (a,), | ||
| 58 | + {"other": b, "dim": 0}, | ||
| 59 | + RuntimeSchemaInfo(needs_pytree=False), | ||
| 60 | + ) | ||
| 61 | + ks = schema.kwargs_strategy | ||
| 62 | + self.assertEqual(len(ks), 1) | ||
| 63 | + self.assertIs(ks[0], b) | ||
| 64 | + | ||
| 65 | + def test_op_schema_kwargs_strategy_pytree_tuple_in_kwarg(self): | ||
| 66 | + mesh = _cpu_mesh(self.ws) | ||
| 67 | + c0 = _replicate_strategy(mesh, (1, 1)) | ||
| 68 | + c1 = _replicate_strategy(mesh, (2, 2)) | ||
| 69 | + ts = TupleStrategy([c0, c1]) | ||
| 70 | + schema = OpSchema( | ||
| 71 | + torch.ops.aten.add.Tensor, | ||
| 72 | + (_replicate_strategy(mesh, (3, 3)),), | ||
| 73 | + {"tensors": ts, "alpha": 1.0}, | ||
| 74 | + RuntimeSchemaInfo(needs_pytree=True), | ||
| 75 | + ) | ||
| 76 | + ks = schema.kwargs_strategy | ||
| 77 | + self.assertEqual(len(ks), 2) | ||
| 78 | + self.assertCountEqual(ks, (c0, c1)) | ||
| 79 | + | ||
| 80 | + def test_expand_one_output_plus_args_and_kwarg_tensor_strategies(self): | ||
| 81 | + """Placement list: [output] + args + kwargs OpStrategies (same invariant as npu_grouped_matmul-style code).""" | ||
| 82 | + mesh = _cpu_mesh(self.ws) | ||
| 83 | + x_st = _replicate_strategy(mesh, (4, 4)) | ||
| 84 | + bias_st = _replicate_strategy(mesh, (4, 4)) | ||
| 85 | + # Positional tensors only in args_schema; bias only in kwargs (typical fused-op pattern). | ||
| 86 | + op_schema = OpSchema( | ||
| 87 | + torch.ops.aten.add.Tensor, | ||
| 88 | + (x_st,), | ||
| 89 | + {"bias": bias_st}, | ||
| 90 | + RuntimeSchemaInfo(needs_pytree=False), | ||
| 91 | + ) | ||
| 92 | + self.assertEqual(len(op_schema.args_strategy) + len(op_schema.kwargs_strategy), 2) | ||
| 93 | + # [y_out, x_in, bias_kw] | ||
| 94 | + single = [[Replicate(), Replicate(), Replicate()]] | ||
| 95 | + strat = expand_to_full_mesh_op_strategy(mesh, op_schema, single, input_index=1) | ||
| 96 | + self.assertGreater(len(strat.strategies), 0) | ||
| 97 | + self.assertEqual(len(strat.strategies[0].input_specs), 2) | ||
| 98 | + | ||
| 99 | + def test_expand_multi_output_index_with_kwargs(self): | ||
| 100 | + """Multiple leading outputs in placement list; kwargs still counted in input_args_strategy.""" | ||
| 101 | + mesh = _cpu_mesh(self.ws) | ||
| 102 | + in0 = _replicate_strategy(mesh, (2, 8, 64, 128)) | ||
| 103 | + in1 = _replicate_strategy(mesh, (2, 8, 64, 192)) | ||
| 104 | + kw_st = _replicate_strategy(mesh, (1,)) | ||
| 105 | + op_schema = OpSchema( | ||
| 106 | + torch.ops.aten.add.Tensor, | ||
| 107 | + (in0, in1), | ||
| 108 | + {"aux": kw_st}, | ||
| 109 | + RuntimeSchemaInfo(needs_pytree=False), | ||
| 110 | + ) | ||
| 111 | + # 2 outputs + 2 args + 1 kw = 5 placements; input_index=2 -> two output slots. | ||
| 112 | + single = [ | ||
| 113 | + [ | ||
| 114 | + Replicate(), | ||
| 115 | + Replicate(), | ||
| 116 | + Replicate(), | ||
| 117 | + Replicate(), | ||
| 118 | + Replicate(), | ||
| 119 | + ] | ||
| 120 | + ] | ||
| 121 | + strat = expand_to_full_mesh_op_strategy(mesh, op_schema, single, input_index=2) | ||
| 122 | + self.assertGreater(len(strat.strategies), 0) | ||
| 123 | + self.assertEqual(len(strat.strategies[0].input_specs), 3) | ||
| 124 | + | ||
| 125 | + | ||
| 126 | +if __name__ == "__main__": | ||
| 127 | + run_tests() | ||
| @@ -1,4 +1,3 @@ | |||
| 1 | -import torch_npu.distributed.tensor._dtensor_patch # patch before register strategy | ||
| 2 | import torch_npu.distributed.tensor._pointwise_ops | 1 | import torch_npu.distributed.tensor._pointwise_ops |
| 3 | import torch_npu.distributed.tensor._matrix_ops | 2 | import torch_npu.distributed.tensor._matrix_ops |
| 4 | import torch_npu.distributed.tensor._attention | 3 | import torch_npu.distributed.tensor._attention |
Dtorch_npu/distributed/tensor/_dtensor_patch.py+0-120
| @@ -1,120 +0,0 @@ | |||
| 1 | -import itertools | ||
| 2 | -from typing import Callable, Optional | ||
| 3 | - | ||
| 4 | -import torch | ||
| 5 | -from torch.distributed.tensor._dtensor_spec import DTensorSpec | ||
| 6 | -from torch.distributed.tensor._op_schema import ( | ||
| 7 | - OpSchema, | ||
| 8 | - OpSpec, | ||
| 9 | - OpStrategy, | ||
| 10 | - PlacementList, | ||
| 11 | -) | ||
| 12 | -from torch.distributed.tensor._ops.utils import ( | ||
| 13 | - generate_redistribute_costs, | ||
| 14 | - is_tensor_shardable | ||
| 15 | -) | ||
| 16 | -from torch.distributed.tensor.device_mesh import DeviceMesh | ||
| 17 | - | ||
| 18 | -try: | ||
| 19 | - from torch.utils._cxx_pytree import tree_leaves | ||
| 20 | -except ImportError: | ||
| 21 | - from torch.utils._pytree import tree_leaves | ||
| 22 | - | ||
| 23 | - | ||
| 24 | -def _patched_kwargs_strategy(self) -> tuple[OpStrategy, ...]: | ||
| 25 | - kwargs_vals = ( | ||
| 26 | - tree_leaves(self.kwargs_schema) | ||
| 27 | - if self.schema_info is not None and self.schema_info.needs_pytree | ||
| 28 | - else self.kwargs_schema.values() | ||
| 29 | - ) | ||
| 30 | - return tuple(item for item in kwargs_vals if isinstance(item, OpStrategy)) | ||
| 31 | - | ||
| 32 | - | ||
| 33 | -def _patched_expand_to_full_mesh_op_strategy( | ||
| 34 | - mesh: DeviceMesh, | ||
| 35 | - op_schema: OpSchema, | ||
| 36 | - single_mesh_dim_strategies: list[PlacementList], | ||
| 37 | - *, | ||
| 38 | - input_index: int = 1, | ||
| 39 | - inplace_op: bool = False, | ||
| 40 | - is_valid_strategy_cb: Optional[ | ||
| 41 | - Callable[[list[DTensorSpec], tuple[Optional[DTensorSpec], ...]], bool] | ||
| 42 | - ] = None, | ||
| 43 | -) -> OpStrategy: | ||
| 44 | - # Expand the single_mesh_dim_strategies to full mesh dim strategies. | ||
| 45 | - all_mesh_dim_strategies = [single_mesh_dim_strategies] * mesh.ndim | ||
| 46 | - | ||
| 47 | - strategy_combs = itertools.product(*all_mesh_dim_strategies) | ||
| 48 | - | ||
| 49 | - all_strategies = [] | ||
| 50 | - for strategy_comb in strategy_combs: | ||
| 51 | - spec_list: list[Optional[DTensorSpec]] = [] | ||
| 52 | - for specs in zip(*strategy_comb): | ||
| 53 | - if specs[0] is not None: | ||
| 54 | - spec_list.append(DTensorSpec(mesh, specs)) | ||
| 55 | - else: | ||
| 56 | - spec_list.append(None) | ||
| 57 | - | ||
| 58 | - input_specs: list[DTensorSpec] = [s for s in spec_list[input_index:] if isinstance(s, DTensorSpec)] | ||
| 59 | - | ||
| 60 | - args_strategy = op_schema.args_strategy | ||
| 61 | - kwargs_strategy = op_schema.kwargs_strategy | ||
| 62 | - input_args_strategy = args_strategy + kwargs_strategy | ||
| 63 | - | ||
| 64 | - if len(input_specs) != len(input_args_strategy): | ||
| 65 | - raise AssertionError( | ||
| 66 | - f"input_specs({len(input_specs)}) != strategies({len(input_args_strategy)}: " | ||
| 67 | - f"{len(args_strategy)} args + {len(kwargs_strategy)} kwargs)" | ||
| 68 | - ) | ||
| 69 | - self_spec = input_args_strategy[0].strategies[0].output_spec | ||
| 70 | - | ||
| 71 | - if inplace_op and self_spec.placements != input_specs[0].placements: | ||
| 72 | - # if it's inplace op, we would only allow the OpSpec to be added when the | ||
| 73 | - # input_spec matches the first argument's runtime sharding, otherwise we skip | ||
| 74 | - continue | ||
| 75 | - | ||
| 76 | - output_specs: tuple[Optional[DTensorSpec], ...] | ||
| 77 | - if input_index > 1: | ||
| 78 | - output_specs = tuple(spec_list[:input_index]) | ||
| 79 | - else: | ||
| 80 | - if spec_list[0] is not None: | ||
| 81 | - output_specs = spec_list[0] # type: ignore[assignment] | ||
| 82 | - else: | ||
| 83 | - raise RuntimeError("output spec is None") | ||
| 84 | - | ||
| 85 | - # check all inputs are shardable | ||
| 86 | - if not all( | ||
| 87 | - is_tensor_shardable(inp.shape, s) | ||
| 88 | - for inp, s in zip(input_args_strategy, input_specs) | ||
| 89 | - ): | ||
| 90 | - continue | ||
| 91 | - | ||
| 92 | - # perform additional op-specific filtering | ||
| 93 | - if is_valid_strategy_cb is not None: | ||
| 94 | - if not is_valid_strategy_cb(input_specs, output_specs): | ||
| 95 | - continue | ||
| 96 | - | ||
| 97 | - redistribute_cost = [ | ||
| 98 | - generate_redistribute_costs(input_strategy, input_spec) | ||
| 99 | - for input_strategy, input_spec in zip(input_args_strategy, input_specs) | ||
| 100 | - ] | ||
| 101 | - | ||
| 102 | - strategy = OpSpec( | ||
| 103 | - output_specs=output_specs, | ||
| 104 | - input_specs=input_specs, | ||
| 105 | - redistribute_cost=redistribute_cost, | ||
| 106 | - ) | ||
| 107 | - all_strategies.append(strategy) | ||
| 108 | - return OpStrategy(all_strategies) | ||
| 109 | - | ||
| 110 | - | ||
| 111 | -def _apply_dtensor_patch(): | ||
| 112 | - # adding kwarg inputs handling in register sharding for previous pytorch version | ||
| 113 | - # See pytorch/pytorch/pull/168249 | ||
| 114 | - if torch.__version__ < "2.10": | ||
| 115 | - if not hasattr(OpSchema, "kwargs_strategy"): | ||
| 116 | - OpSchema.kwargs_strategy = property(_patched_kwargs_strategy) | ||
| 117 | - torch.distributed.tensor._ops.utils.expand_to_full_mesh_op_strategy = _patched_expand_to_full_mesh_op_strategy | ||
| 118 | - | ||
| 119 | - | ||
| 120 | -_apply_dtensor_patch() | ||