已合并
refactor dtensor rules #34096
kisnwang创建于 4月21日
refactor dtensor rules #34096
已合并
共 10 个文件变更+138-74
| @@ -1,3 +1,5 @@ | |||
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates | ||
| 2 | +# Owner(s): ["oncall: distributed"] | ||
| 1 | import itertools | 3 | import itertools |
| 2 | from typing import cast, List, Optional | 4 | from typing import cast, List, Optional |
| 3 | from unittest import skip | 5 | from unittest import skip |
| @@ -6,7 +8,7 @@ import torch | |||
| 6 | from torch.distributed._tensor import DeviceMesh, distribute_tensor | 8 | from torch.distributed._tensor import DeviceMesh, distribute_tensor |
| 7 | from torch.distributed._tensor.api import DTensor | 9 | from torch.distributed._tensor.api import DTensor |
| 8 | from torch.distributed._tensor.placement_types import ( | 10 | from torch.distributed._tensor.placement_types import ( |
| 9 | - _Partial, | 11 | + Partial, |
| 10 | Placement, | 12 | Placement, |
| 11 | Replicate, | 13 | Replicate, |
| 12 | Shard, | 14 | Shard, |
| @@ -16,6 +18,7 @@ from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTe | |||
| 16 | 18 | ||
| 17 | import torch_npu | 19 | import torch_npu |
| 18 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 20 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU |
| 21 | +npu = torch.ops.npu | ||
| 19 | 22 | ||
| 20 | 23 | ||
| 21 | class DistMatrixOpsTest(DTensorTestBase): | 24 | class DistMatrixOpsTest(DTensorTestBase): |
| @@ -60,7 +63,7 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 60 | 63 | ||
| 61 | # test if addmm output is a partial | 64 | # test if addmm output is a partial |
| 62 | self.assertIsInstance(dist_res, DTensor) | 65 | self.assertIsInstance(dist_res, DTensor) |
| 63 | - self.assertIsInstance(dist_res.placements[0], _Partial) | 66 | + self.assertIsInstance(dist_res.placements[0], Partial) |
| 64 | 67 | ||
| 65 | # test if result is the same as tensor | 68 | # test if result is the same as tensor |
| 66 | replica_res = dist_res.redistribute(device_mesh, replica_spec) | 69 | replica_res = dist_res.redistribute(device_mesh, replica_spec) |
| @@ -130,10 +133,10 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 130 | da = distribute_tensor(a, device_mesh, [Shard(1)]) | 133 | da = distribute_tensor(a, device_mesh, [Shard(1)]) |
| 131 | db = distribute_tensor(b, device_mesh, [Shard(0)]) | 134 | db = distribute_tensor(b, device_mesh, [Shard(0)]) |
| 132 | 135 | ||
| 133 | - # mm(da, db) should return a _Partial tensor. | 136 | + # mm(da, db) should return a Partial tensor. |
| 134 | - # transposing it should keep it _Partial | 137 | + # transposing it should keep it Partial |
| 135 | dc = torch.mm(da, db).t() | 138 | dc = torch.mm(da, db).t() |
| 136 | - self.assertTrue(isinstance(dc.placements[0], _Partial)) | 139 | + self.assertTrue(isinstance(dc.placements[0], Partial)) |
| 137 | # check that the local and distributed op results match | 140 | # check that the local and distributed op results match |
| 138 | self.assertEqual( | 141 | self.assertEqual( |
| 139 | c, | 142 | c, |
| @@ -272,5 +275,49 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 272 | test_placement_comb([spec[0]], [spec[1]]) | 275 | test_placement_comb([spec[0]], [spec[1]]) |
| 273 | 276 | ||
| 274 | 277 | ||
| 278 | + | ||
| 279 | + | ||
| 280 | + def test_npu_bmmV2(self): | ||
| 281 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 282 | + mat1 = torch.rand(4, 8, 4, device=self.device_type, requires_grad=True) | ||
| 283 | + mat2 = torch.rand(4, 4, 8, device=self.device_type, requires_grad=True) | ||
| 284 | + local_result = npu.npu_bmmV2(mat1, mat2) | ||
| 285 | + grad_local_res = torch.ones_like(local_result) | ||
| 286 | + local_result.backward(grad_local_res) | ||
| 287 | + | ||
| 288 | + def test_placement_comb( | ||
| 289 | + placements1: List[Placement], | ||
| 290 | + placements2: List[Placement], | ||
| 291 | + ) -> None: | ||
| 292 | + mat1_dt = distribute_tensor(mat1, device_mesh, placements1) | ||
| 293 | + mat2_dt = distribute_tensor(mat2, device_mesh, placements2) | ||
| 294 | + dist_res = cast(DTensor, npu.npu_bmmV2(mat1_dt, mat2_dt)).redistribute( | ||
| 295 | + device_mesh, [Replicate()] | ||
| 296 | + ) | ||
| 297 | + dist_local_res = dist_res.to_local() | ||
| 298 | + self.assertEqual(dist_local_res, local_result) | ||
| 299 | + | ||
| 300 | + # test backward | ||
| 301 | + # it generates a different grad shape | ||
| 302 | + grad_dist_res = torch.ones_like(dist_res) | ||
| 303 | + dist_res.backward(grad_dist_res) | ||
| 304 | + self.assertIsNotNone(mat1_dt.grad) | ||
| 305 | + mat1_dt_grad = cast(DTensor, mat1_dt.grad) | ||
| 306 | + mat1_grad_local = mat1_dt_grad.redistribute( | ||
| 307 | + device_mesh, [Replicate()] | ||
| 308 | + ).to_local() | ||
| 309 | + self.assertEqual(mat1_grad_local, mat1.grad) | ||
| 310 | + | ||
| 311 | + shard0_spec = Shard(0) | ||
| 312 | + shard1_spec = Shard(1) | ||
| 313 | + shard2_spec = Shard(2) | ||
| 314 | + replica_spec = Replicate() | ||
| 315 | + placement_specs = [shard0_spec, shard1_spec, shard2_spec, replica_spec] | ||
| 316 | + shard_specs_comb = list(itertools.product(placement_specs, placement_specs)) | ||
| 317 | + | ||
| 318 | + # tests that currently pass | ||
| 319 | + for spec in shard_specs_comb: | ||
| 320 | + test_placement_comb([spec[0]], [spec[1]]) | ||
| 321 | + | ||
| 275 | if __name__ == "__main__": | 322 | if __name__ == "__main__": |
| 276 | run_tests() | 323 | run_tests() |
| @@ -11,7 +11,7 @@ from torch import Tensor | |||
| 11 | 11 | ||
| 12 | from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor | 12 | from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor |
| 13 | from torch.distributed._tensor.placement_types import ( | 13 | from torch.distributed._tensor.placement_types import ( |
| 14 | - _Partial, | 14 | + Partial, |
| 15 | Placement, | 15 | Placement, |
| 16 | Replicate, | 16 | Replicate, |
| 17 | Shard, | 17 | Shard, |
| @@ -22,6 +22,7 @@ from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTe | |||
| 22 | 22 | ||
| 23 | import torch_npu | 23 | import torch_npu |
| 24 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 24 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU |
| 25 | +npu = torch.ops.npu | ||
| 25 | 26 | ||
| 26 | 27 | ||
| 27 | def no_op(): | 28 | def no_op(): |
| @@ -140,8 +141,8 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 140 | 141 | ||
| 141 | def test_partial_add(self): | 142 | def test_partial_add(self): |
| 142 | device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | 143 | device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) |
| 143 | - d_1 = DTensor.from_local(torch.rand(2, 2), device_mesh, [_Partial()]) | 144 | + d_1 = DTensor.from_local(torch.rand(2, 2), device_mesh, [Partial()]) |
| 144 | - d_2 = DTensor.from_local(torch.rand(2, 2), device_mesh, [_Partial()]) | 145 | + d_2 = DTensor.from_local(torch.rand(2, 2), device_mesh, [Partial()]) |
| 145 | d_3 = d_1 + d_2 | 146 | d_3 = d_1 + d_2 |
| 146 | self.assertEqual(d_3._spec.placements[0].is_partial(), True) | 147 | self.assertEqual(d_3._spec.placements[0].is_partial(), True) |
| 147 | 148 | ||
| @@ -161,6 +162,30 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 161 | input_size=(8, 5), | 162 | input_size=(8, 5), |
| 162 | op=torch.nn.functional.gelu, | 163 | op=torch.nn.functional.gelu, |
| 163 | ) | 164 | ) |
| 165 | + self._run_sharded_elementwise_ops( | ||
| 166 | + device_mesh=device_mesh, | ||
| 167 | + placements=[Shard(0)], | ||
| 168 | + input_size=(8, 5), | ||
| 169 | + op=npu.fast_gelu, | ||
| 170 | + ) | ||
| 171 | + self._run_sharded_elementwise_ops( | ||
| 172 | + device_mesh=device_mesh, | ||
| 173 | + placements=[Replicate()], | ||
| 174 | + input_size=(8, 5), | ||
| 175 | + op=npu.fast_gelu, | ||
| 176 | + ) | ||
| 177 | + self._run_sharded_elementwise_ops( | ||
| 178 | + device_mesh=device_mesh, | ||
| 179 | + placements=[Shard(0)], | ||
| 180 | + input_size=(8, 5), | ||
| 181 | + op=npu.npu_fast_gelu, | ||
| 182 | + ) | ||
| 183 | + self._run_sharded_elementwise_ops( | ||
| 184 | + device_mesh=device_mesh, | ||
| 185 | + placements=[Replicate()], | ||
| 186 | + input_size=(8, 5), | ||
| 187 | + op=npu.npu_fast_gelu, | ||
| 188 | + ) | ||
| 164 | self._run_sharded_elementwise_ops( | 189 | self._run_sharded_elementwise_ops( |
| 165 | device_mesh=device_mesh, | 190 | device_mesh=device_mesh, |
| 166 | placements=[Shard(1)], | 191 | placements=[Shard(1)], |
| @@ -221,7 +246,7 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 221 | with self.assertRaisesRegex(RuntimeError, "supported"): | 246 | with self.assertRaisesRegex(RuntimeError, "supported"): |
| 222 | self._run_sharded_elementwise_ops( | 247 | self._run_sharded_elementwise_ops( |
| 223 | device_mesh=device_mesh, | 248 | device_mesh=device_mesh, |
| 224 | - placements=[_Partial(ReduceOp.SUM)], | 249 | + placements=[Partial(ReduceOp.SUM)], |
| 225 | input_size=(8, 5), | 250 | input_size=(8, 5), |
| 226 | op=torch.nn.functional.dropout, | 251 | op=torch.nn.functional.dropout, |
| 227 | ) | 252 | ) |
| @@ -1,12 +1,14 @@ | |||
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates | ||
| 2 | +# Owner(s): ["oncall: distributed"] | ||
| 1 | import itertools | 3 | import itertools |
| 2 | from typing import cast, List | 4 | from typing import cast, List |
| 3 | 5 | ||
| 4 | import torch | 6 | import torch |
| 5 | import torch.distributed as dist | 7 | import torch.distributed as dist |
| 6 | from torch import rand, randn, Tensor | 8 | from torch import rand, randn, Tensor |
| 7 | -from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate, Shard | 9 | +from torch.distributed.tensor import DeviceMesh, distribute_tensor, Replicate, Shard |
| 8 | -from torch.distributed._tensor.debug import CommDebugMode | 10 | +from torch.distributed.tensor.debug import CommDebugMode |
| 9 | -from torch.distributed._tensor.ops._view_ops import ( | 11 | +from torch.distributed.tensor._ops._view_ops import ( |
| 10 | Broadcast, | 12 | Broadcast, |
| 11 | dim_maps, | 13 | dim_maps, |
| 12 | Flatten, | 14 | Flatten, |
| @@ -16,13 +18,14 @@ from torch.distributed._tensor.ops._view_ops import ( | |||
| 16 | Split, | 18 | Split, |
| 17 | view_groups, | 19 | view_groups, |
| 18 | ) | 20 | ) |
| 19 | -from torch.distributed._tensor.placement_types import Placement | 21 | +from torch.distributed.tensor.placement_types import Placement |
| 20 | from torch.testing._internal.common_utils import run_tests | 22 | from torch.testing._internal.common_utils import run_tests |
| 21 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase | 23 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase |
| 22 | from torch.utils import _pytree as pytree | 24 | from torch.utils import _pytree as pytree |
| 23 | 25 | ||
| 24 | import torch_npu | 26 | import torch_npu |
| 25 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 27 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU |
| 28 | +npu = torch.ops.npu | ||
| 26 | 29 | ||
| 27 | 30 | ||
| 28 | class TestViewOps(DTensorTestBase): | 31 | class TestViewOps(DTensorTestBase): |
| @@ -358,6 +361,16 @@ class TestViewOps(DTensorTestBase): | |||
| 358 | (randn(24, 36, 28), (-1, -3, -2)), | 361 | (randn(24, 36, 28), (-1, -3, -2)), |
| 359 | (InputDim(2), InputDim(0), InputDim(1)), | 362 | (InputDim(2), InputDim(0), InputDim(1)), |
| 360 | ) | 363 | ) |
| 364 | + self.dimmap_test( | ||
| 365 | + npu.npu_transpose, | ||
| 366 | + (randn(24, 36, 28), (2, 0, 1)), | ||
| 367 | + (InputDim(2), InputDim(0), InputDim(1)), | ||
| 368 | + ) | ||
| 369 | + self.dimmap_test( | ||
| 370 | + npu.npu_transpose, | ||
| 371 | + (randn(24, 36, 28), (-1, -3, -2)), | ||
| 372 | + (InputDim(2), InputDim(0), InputDim(1)), | ||
| 373 | + ) | ||
| 361 | self.dimmap_test( | 374 | self.dimmap_test( |
| 362 | torch.ravel, | 375 | torch.ravel, |
| 363 | (randn(24, 36),), | 376 | (randn(24, 36),), |
| @@ -82,7 +82,6 @@ from torch_npu.profiler._add_mstx_patch import _apply_mstx_patch | |||
| 82 | from torch_npu.distributed.fsdp._add_fsdp_patch import _apply_fsdp_patch | 82 | from torch_npu.distributed.fsdp._add_fsdp_patch import _apply_fsdp_patch |
| 83 | from torch_npu.distributed.rpc.backend_registry import _rpc_backend_registry | 83 | from torch_npu.distributed.rpc.backend_registry import _rpc_backend_registry |
| 84 | from torch_npu.utils import _cann_package_check, _add_intercept_methods | 84 | from torch_npu.utils import _cann_package_check, _add_intercept_methods |
| 85 | -from torch_npu.utils import _register_ops_under_dtensor_rules | ||
| 86 | from torch_npu.utils.exposed_api import public_npu_functions | 85 | from torch_npu.utils.exposed_api import public_npu_functions |
| 87 | from torch_npu.multiprocessing.reductions import _add_reductions_methods | 86 | from torch_npu.multiprocessing.reductions import _add_reductions_methods |
| 88 | from torch_npu.npu.utils import _erase_stream as erase_stream | 87 | from torch_npu.npu.utils import _erase_stream as erase_stream |
| @@ -298,9 +297,6 @@ atexit.register(_npu_shutdown) | |||
| 298 | # init and register rpc npu backend | 297 | # init and register rpc npu backend |
| 299 | _rpc_backend_registry() | 298 | _rpc_backend_registry() |
| 300 | 299 | ||
| 301 | -# register rules for ops in dtensor | ||
| 302 | -_register_ops_under_dtensor_rules() | ||
| 303 | - | ||
| 304 | # Enable NPU Sanitizer | 300 | # Enable NPU Sanitizer |
| 305 | if 'TORCH_NPU_SANITIZER' in os.environ: | 301 | if 'TORCH_NPU_SANITIZER' in os.environ: |
| 306 | import torch_npu.npu._sanitizer as csan | 302 | import torch_npu.npu._sanitizer as csan |
| @@ -4,4 +4,5 @@ import torch_npu.distributed.tensor._matrix_ops | |||
| 4 | import torch_npu.distributed.tensor._attention | 4 | import torch_npu.distributed.tensor._attention |
| 5 | import torch_npu.distributed.tensor._math_ops | 5 | import torch_npu.distributed.tensor._math_ops |
| 6 | import torch_npu.distributed.tensor._moe_ops | 6 | import torch_npu.distributed.tensor._moe_ops |
| 7 | +import torch_npu.distributed.tensor._view_ops | ||
| 7 | import torch_npu.distributed.tensor._sharded_tensor_patch | 8 | import torch_npu.distributed.tensor._sharded_tensor_patch |
| @@ -5,6 +5,7 @@ import torch | |||
| 5 | from torch.distributed._tensor.experimental import register_sharding | 5 | from torch.distributed._tensor.experimental import register_sharding |
| 6 | from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta | 6 | from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta |
| 7 | from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy | 7 | from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy |
| 8 | +from torch.distributed.tensor._ops._matrix_ops import _mm_like_strategy | ||
| 8 | from torch_npu._compat.distributed import register_op_strategy | 9 | from torch_npu._compat.distributed import register_op_strategy |
| 9 | from torch.distributed.tensor import DTensor, Partial, Replicate, Shard | 10 | from torch.distributed.tensor import DTensor, Partial, Replicate, Shard |
| 10 | from torch.distributed.tensor._op_schema import ( | 11 | from torch.distributed.tensor._op_schema import ( |
| @@ -822,6 +823,12 @@ def custom_dropout_backward_sharding(op_schema: OpSchema) -> OpStrategy: | |||
| 822 | return output_strategy | 823 | return output_strategy |
| 823 | 824 | ||
| 824 | 825 | ||
| 826 | + | ||
| 827 | +def custom_bmm_strategy(op_schema: OpSchema): | ||
| 828 | + mesh = op_schema.get_mesh_from_args() | ||
| 829 | + return _mm_like_strategy("bmk,bkn->bmn", mesh, op_schema) | ||
| 830 | + | ||
| 831 | + | ||
| 825 | customized_ops = { | 832 | customized_ops = { |
| 826 | npu.npu_grouped_matmul.default: _npu_grouped_matmul_handler, | 833 | npu.npu_grouped_matmul.default: _npu_grouped_matmul_handler, |
| 827 | npu.npu_grouped_matmul.List: _npu_grouped_matmul_handler, | 834 | npu.npu_grouped_matmul.List: _npu_grouped_matmul_handler, |
| @@ -4,11 +4,11 @@ from torch.distributed.tensor._op_schema import OpSchema, RuntimeSchemaInfo | |||
| 4 | from torch_npu._compat.distributed import register_op_strategy | 4 | from torch_npu._compat.distributed import register_op_strategy |
| 5 | from torch.distributed.tensor._ops._pointwise_ops import pointwise_strategy | 5 | from torch.distributed.tensor._ops._pointwise_ops import pointwise_strategy |
| 6 | 6 | ||
| 7 | - | 7 | +aten = torch.ops.aten |
| 8 | npu = torch.ops.npu | 8 | npu = torch.ops.npu |
| 9 | 9 | ||
| 10 | 10 | ||
| 11 | -custom_pointwise_ops = { | 11 | +custom_linear_pointwise_ops = { |
| 12 | npu.npu_dtype_cast.default: 0, | 12 | npu.npu_dtype_cast.default: 0, |
| 13 | npu._npu_dtype_cast.default: 0, | 13 | npu._npu_dtype_cast.default: 0, |
| 14 | npu.npu_dtype_cast_backward.default: 0, | 14 | npu.npu_dtype_cast_backward.default: 0, |
| @@ -16,12 +16,33 @@ custom_pointwise_ops = { | |||
| 16 | } | 16 | } |
| 17 | 17 | ||
| 18 | 18 | ||
| 19 | -def custom_pointwise_strategy(op_schema: OpSchema): | 19 | +def custom_linear_pointwise_strategy(op_schema: OpSchema): |
| 20 | - op_type = custom_pointwise_ops.get(op_schema.op, -1) | 20 | + op_type = custom_linear_pointwise_ops.get(op_schema.op, -1) |
| 21 | return pointwise_strategy(op_schema, linearity=op_type) | 21 | return pointwise_strategy(op_schema, linearity=op_type) |
| 22 | 22 | ||
| 23 | 23 | ||
| 24 | +for op in custom_linear_pointwise_ops: | ||
| 25 | + register_op_strategy( | ||
| 26 | + op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"]) | ||
| 27 | + )(custom_linear_pointwise_strategy) | ||
| 28 | + | ||
| 29 | + | ||
| 30 | +custom_pointwise_ops = [ | ||
| 31 | + # please keep the entries below alphabetically sorted | ||
| 32 | + # native ops | ||
| 33 | + aten.isclose.default, | ||
| 34 | + aten.isfinite.default, | ||
| 35 | + # custom ops | ||
| 36 | + npu.fast_gelu.default, | ||
| 37 | + npu.npu_fast_gelu.default, | ||
| 38 | + npu.npu_layer_norm_eval.default, | ||
| 39 | + # backward point-wise ops | ||
| 40 | + # please keep the entries below alphabetically sorted | ||
| 41 | + npu.npu_fast_gelu_backward.default | ||
| 42 | +] | ||
| 43 | + | ||
| 44 | + | ||
| 24 | for op in custom_pointwise_ops: | 45 | for op in custom_pointwise_ops: |
| 25 | register_op_strategy( | 46 | register_op_strategy( |
| 26 | op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"]) | 47 | op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"]) |
| 27 | - )(custom_pointwise_strategy) | 48 | + )(pointwise_strategy) |
| @@ -0,0 +1,6 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.distributed.tensor._ops._view_ops import register_op_strategy_map | ||
| 3 | +npu = torch.ops.npu | ||
| 4 | + | ||
| 5 | + | ||
| 6 | +register_op_strategy_map(npu.npu_transpose.default, torch.permute) | ||
| @@ -9,7 +9,6 @@ from .storage import _add_storage_methods | |||
| 9 | from .combine_tensors import npu_combine_tensors, get_part_combined_tensor, is_combined_tensor_valid | 9 | from .combine_tensors import npu_combine_tensors, get_part_combined_tensor, is_combined_tensor_valid |
| 10 | from .serialization import _add_serialization_methods, save_async | 10 | from .serialization import _add_serialization_methods, save_async |
| 11 | from .npu_intercept import _cann_package_check, _add_intercept_methods | 11 | from .npu_intercept import _cann_package_check, _add_intercept_methods |
| 12 | -from .dtensor import _register_ops_under_dtensor_rules | ||
| 13 | from .collect_env import _add_collect_env_methods | 12 | from .collect_env import _add_collect_env_methods |
| 14 | from ._dynamo import add_dynamo_methods | 13 | from ._dynamo import add_dynamo_methods |
| 15 | from ._inductor import _inductor_register_device_op_overrides | 14 | from ._inductor import _inductor_register_device_op_overrides |
| @@ -1,51 +0,0 @@ | |||
| 1 | -import torch | ||
| 2 | -from torch.distributed.tensor._ops._common_rules import pointwise_rule | ||
| 3 | -from torch.distributed.tensor._ops.utils import normalize_dims | ||
| 4 | -from torch_npu._compat.distributed import register_prop_rule | ||
| 5 | -from torch.distributed.tensor._ops._matrix_ops import bmm_strategy | ||
| 6 | -from torch.distributed.tensor._ops._view_ops import ( | ||
| 7 | - register_op_strategy_map, | ||
| 8 | - dim_maps, | ||
| 9 | - InputDim | ||
| 10 | -) | ||
| 11 | -import torch_npu | ||
| 12 | - | ||
| 13 | -__all__ = [] | ||
| 14 | - | ||
| 15 | - | ||
| 16 | -def _register_ops_under_dtensor_rules(): | ||
| 17 | - npu = torch.ops.npu | ||
| 18 | - aten = torch.ops.aten | ||
| 19 | - | ||
| 20 | - pointwise_ops = [ | ||
| 21 | - # please keep the entries below alphabetically sorted | ||
| 22 | - # native ops | ||
| 23 | - aten.isclose.default, | ||
| 24 | - aten.isfinite.default, | ||
| 25 | - # custom ops | ||
| 26 | - npu.fast_gelu.default, | ||
| 27 | - npu.npu_fast_gelu.default, | ||
| 28 | - npu.npu_layer_norm_eval.default, | ||
| 29 | - # backward point-wise ops | ||
| 30 | - # please keep the entries below alphabetically sorted | ||
| 31 | - npu.npu_fast_gelu_backward.default | ||
| 32 | - ] | ||
| 33 | - | ||
| 34 | - matrix_ops = [ | ||
| 35 | - npu.npu_bmmV2.default | ||
| 36 | - ] | ||
| 37 | - # pointwise rule | ||
| 38 | - for op in pointwise_ops: | ||
| 39 | - register_prop_rule(op)(pointwise_rule) | ||
| 40 | - | ||
| 41 | - # bmm rules | ||
| 42 | - for op in matrix_ops: | ||
| 43 | - register_prop_rule(op)(bmm_strategy) | ||
| 44 | - | ||
| 45 | - # reshape_prop under view_ops | ||
| 46 | - dim_maps.update({ | ||
| 47 | - torch_npu.npu_transpose: lambda input, dims: tuple( | ||
| 48 | - InputDim(i) for i in normalize_dims(dims, input.ndim) | ||
| 49 | - ) | ||
| 50 | - }) | ||
| 51 | - register_op_strategy_map(npu.npu_transpose.default, torch_npu.npu_transpose) | ||