已合并
refactor dtensor rules #34586
kisnwang创建于 4月28日
refactor dtensor rules #34586
已合并
共 11 个文件变更+780-416
| @@ -1,12 +1,16 @@ | |||
| 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 |
| 3 | from unittest import skip | 5 | from unittest import skip |
| 4 | 6 | ||
| 7 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU, with_comms | ||
| 8 | + | ||
| 5 | import torch | 9 | import torch |
| 6 | from torch.distributed._tensor import DeviceMesh, distribute_tensor | 10 | from torch.distributed._tensor import DeviceMesh, distribute_tensor |
| 7 | from torch.distributed._tensor.api import DTensor | 11 | from torch.distributed._tensor.api import DTensor |
| 8 | from torch.distributed._tensor.placement_types import ( | 12 | from torch.distributed._tensor.placement_types import ( |
| 9 | - _Partial, | 13 | + Partial, |
| 10 | Placement, | 14 | Placement, |
| 11 | Replicate, | 15 | Replicate, |
| 12 | Shard, | 16 | Shard, |
| @@ -14,8 +18,8 @@ from torch.distributed._tensor.placement_types import ( | |||
| 14 | from torch.testing._internal.common_utils import run_tests | 18 | from torch.testing._internal.common_utils import run_tests |
| 15 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase | 19 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase |
| 16 | 20 | ||
| 17 | -import torch_npu | 21 | + |
| 18 | -from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 22 | +npu = torch.ops.npu |
| 19 | 23 | ||
| 20 | 24 | ||
| 21 | class DistMatrixOpsTest(DTensorTestBase): | 25 | class DistMatrixOpsTest(DTensorTestBase): |
| @@ -48,9 +52,13 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 48 | shard1_spec = [Shard(1)] | 52 | shard1_spec = [Shard(1)] |
| 49 | replica_spec = [Replicate()] | 53 | replica_spec = [Replicate()] |
| 50 | 54 | ||
| 51 | - tensor_to_shard1 = torch.randn(12, 8, requires_grad=True, device=self.device_type) | 55 | + tensor_to_shard1 = torch.randn( |
| 56 | + 12, 8, requires_grad=True, device=self.device_type | ||
| 57 | + ) | ||
| 52 | mat1 = distribute_tensor(tensor_to_shard1, device_mesh, shard1_spec) | 58 | mat1 = distribute_tensor(tensor_to_shard1, device_mesh, shard1_spec) |
| 53 | - tensor_to_shard0 = torch.randn(8, 4, requires_grad=True, device=self.device_type) | 59 | + tensor_to_shard0 = torch.randn( |
| 60 | + 8, 4, requires_grad=True, device=self.device_type | ||
| 61 | + ) | ||
| 54 | mat2 = distribute_tensor(tensor_to_shard0, device_mesh, shard0_spec) | 62 | mat2 = distribute_tensor(tensor_to_shard0, device_mesh, shard0_spec) |
| 55 | input_tensor = torch.randn(4, requires_grad=True, device=self.device_type) | 63 | input_tensor = torch.randn(4, requires_grad=True, device=self.device_type) |
| 56 | input1 = distribute_tensor(input_tensor, device_mesh, replica_spec) | 64 | input1 = distribute_tensor(input_tensor, device_mesh, replica_spec) |
| @@ -60,7 +68,7 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 60 | 68 | ||
| 61 | # test if addmm output is a partial | 69 | # test if addmm output is a partial |
| 62 | self.assertIsInstance(dist_res, DTensor) | 70 | self.assertIsInstance(dist_res, DTensor) |
| 63 | - self.assertIsInstance(dist_res.placements[0], _Partial) | 71 | + self.assertIsInstance(dist_res.placements[0], Partial) |
| 64 | 72 | ||
| 65 | # test if result is the same as tensor | 73 | # test if result is the same as tensor |
| 66 | replica_res = dist_res.redistribute(device_mesh, replica_spec) | 74 | replica_res = dist_res.redistribute(device_mesh, replica_spec) |
| @@ -87,7 +95,7 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 87 | local_res = torch.mm(t1, t2) | 95 | local_res = torch.mm(t1, t2) |
| 88 | 96 | ||
| 89 | def test_placement_comb( | 97 | def test_placement_comb( |
| 90 | - placements1: List[Placement], placements2: List[Placement] | 98 | + placements1: list[Placement], placements2: list[Placement] |
| 91 | ) -> None: | 99 | ) -> None: |
| 92 | dt1 = distribute_tensor(t1, device_mesh, placements1) | 100 | dt1 = distribute_tensor(t1, device_mesh, placements1) |
| 93 | dt2 = distribute_tensor(t2, device_mesh, placements2) | 101 | dt2 = distribute_tensor(t2, device_mesh, placements2) |
| @@ -130,10 +138,10 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 130 | da = distribute_tensor(a, device_mesh, [Shard(1)]) | 138 | da = distribute_tensor(a, device_mesh, [Shard(1)]) |
| 131 | db = distribute_tensor(b, device_mesh, [Shard(0)]) | 139 | db = distribute_tensor(b, device_mesh, [Shard(0)]) |
| 132 | 140 | ||
| 133 | - # mm(da, db) should return a _Partial tensor. | 141 | + # mm(da, db) should return a Partial tensor. |
| 134 | - # transposing it should keep it _Partial | 142 | + # transposing it should keep it Partial |
| 135 | dc = torch.mm(da, db).t() | 143 | dc = torch.mm(da, db).t() |
| 136 | - self.assertTrue(isinstance(dc.placements[0], _Partial)) | 144 | + self.assertTrue(isinstance(dc.placements[0], Partial)) |
| 137 | # check that the local and distributed op results match | 145 | # check that the local and distributed op results match |
| 138 | self.assertEqual( | 146 | self.assertEqual( |
| 139 | c, | 147 | c, |
| @@ -150,12 +158,12 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 150 | batch_2 = torch.rand(4, 8, 8, requires_grad=True) | 158 | batch_2 = torch.rand(4, 8, 8, requires_grad=True) |
| 151 | 159 | ||
| 152 | def test_placement_comb( | 160 | def test_placement_comb( |
| 153 | - tensor_placements: List[Placement], | 161 | + tensor_placements: list[Placement], |
| 154 | - batch_1_placements: List[Placement], | 162 | + batch_1_placements: list[Placement], |
| 155 | - batch_2_placements: List[Placement], | 163 | + batch_2_placements: list[Placement], |
| 156 | beta: int, | 164 | beta: int, |
| 157 | alpha: int, | 165 | alpha: int, |
| 158 | - batch_1_grad: Optional[torch.Tensor], | 166 | + batch_1_grad: torch.Tensor | None, |
| 159 | ) -> None: | 167 | ) -> None: |
| 160 | tensor_dt = distribute_tensor(tensor, device_mesh, tensor_placements) | 168 | tensor_dt = distribute_tensor(tensor, device_mesh, tensor_placements) |
| 161 | batch_1_dt = distribute_tensor(batch_1, device_mesh, batch_1_placements) | 169 | batch_1_dt = distribute_tensor(batch_1, device_mesh, batch_1_placements) |
| @@ -238,8 +246,8 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 238 | local_result.backward(grad_local_res) | 246 | local_result.backward(grad_local_res) |
| 239 | 247 | ||
| 240 | def test_placement_comb( | 248 | def test_placement_comb( |
| 241 | - placements1: List[Placement], | 249 | + placements1: list[Placement], |
| 242 | - placements2: List[Placement], | 250 | + placements2: list[Placement], |
| 243 | ) -> None: | 251 | ) -> None: |
| 244 | mat1_dt = distribute_tensor(mat1, device_mesh, placements1) | 252 | mat1_dt = distribute_tensor(mat1, device_mesh, placements1) |
| 245 | mat2_dt = distribute_tensor(mat2, device_mesh, placements2) | 253 | mat2_dt = distribute_tensor(mat2, device_mesh, placements2) |
| @@ -271,6 +279,50 @@ class DistMatrixOpsTest(DTensorTestBase): | |||
| 271 | for spec in shard_specs_comb: | 279 | for spec in shard_specs_comb: |
| 272 | test_placement_comb([spec[0]], [spec[1]]) | 280 | test_placement_comb([spec[0]], [spec[1]]) |
| 273 | 281 | ||
| 282 | + | ||
| 283 | + | ||
| 284 | + def test_npu_bmmV2(self): | ||
| 285 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 286 | + mat1 = torch.rand(4, 8, 4, device=self.device_type, requires_grad=True) | ||
| 287 | + mat2 = torch.rand(4, 4, 8, device=self.device_type, requires_grad=True) | ||
| 288 | + local_result = npu.npu_bmmV2(mat1, mat2) | ||
| 289 | + grad_local_res = torch.ones_like(local_result) | ||
| 290 | + local_result.backward(grad_local_res) | ||
| 291 | + | ||
| 292 | + def test_placement_comb( | ||
| 293 | + placements1: list[Placement], | ||
| 294 | + placements2: list[Placement], | ||
| 295 | + ) -> None: | ||
| 296 | + mat1_dt = distribute_tensor(mat1, device_mesh, placements1) | ||
| 297 | + mat2_dt = distribute_tensor(mat2, device_mesh, placements2) | ||
| 298 | + dist_res = cast(DTensor, npu.npu_bmmV2(mat1_dt, mat2_dt)).redistribute( | ||
| 299 | + device_mesh, [Replicate()] | ||
| 300 | + ) | ||
| 301 | + dist_local_res = dist_res.to_local() | ||
| 302 | + self.assertEqual(dist_local_res, local_result) | ||
| 303 | + | ||
| 304 | + # test backward | ||
| 305 | + # it generates a different grad shape | ||
| 306 | + grad_dist_res = torch.ones_like(dist_res) | ||
| 307 | + dist_res.backward(grad_dist_res) | ||
| 308 | + self.assertIsNotNone(mat1_dt.grad) | ||
| 309 | + mat1_dt_grad = cast(DTensor, mat1_dt.grad) | ||
| 310 | + mat1_grad_local = mat1_dt_grad.redistribute( | ||
| 311 | + device_mesh, [Replicate()] | ||
| 312 | + ).to_local() | ||
| 313 | + self.assertEqual(mat1_grad_local, mat1.grad) | ||
| 314 | + | ||
| 315 | + shard0_spec = Shard(0) | ||
| 316 | + shard1_spec = Shard(1) | ||
| 317 | + shard2_spec = Shard(2) | ||
| 318 | + replica_spec = Replicate() | ||
| 319 | + placement_specs = [shard0_spec, shard1_spec, shard2_spec, replica_spec] | ||
| 320 | + shard_specs_comb = list(itertools.product(placement_specs, placement_specs)) | ||
| 321 | + | ||
| 322 | + # tests that currently pass | ||
| 323 | + for spec in shard_specs_comb: | ||
| 324 | + test_placement_comb([spec[0]], [spec[1]]) | ||
| 325 | + | ||
| 274 | 326 | ||
| 275 | if __name__ == "__main__": | 327 | if __name__ == "__main__": |
| 276 | run_tests() | 328 | run_tests() |
| @@ -1,17 +1,19 @@ | |||
| 1 | # Copyright (c) Meta Platforms, Inc. and affiliates | 1 | # Copyright (c) Meta Platforms, Inc. and affiliates |
| 2 | # Owner(s): ["oncall: distributed"] | 2 | # Owner(s): ["oncall: distributed"] |
| 3 | 3 | ||
| 4 | -from typing import Any, Callable, Dict, Optional, Sequence | 4 | +from collections.abc import Callable, Sequence |
| 5 | +from typing import Any | ||
| 5 | from unittest import skip | 6 | from unittest import skip |
| 6 | 7 | ||
| 7 | -import torch | 8 | +import torch_npu |
| 9 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU, with_comms | ||
| 8 | 10 | ||
| 11 | +import torch | ||
| 9 | import torch.utils._pytree as pytree | 12 | import torch.utils._pytree as pytree |
| 10 | from torch import Tensor | 13 | from torch import Tensor |
| 11 | - | ||
| 12 | from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor | 14 | from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor |
| 13 | from torch.distributed._tensor.placement_types import ( | 15 | from torch.distributed._tensor.placement_types import ( |
| 14 | - _Partial, | 16 | + Partial, |
| 15 | Placement, | 17 | Placement, |
| 16 | Replicate, | 18 | Replicate, |
| 17 | Shard, | 19 | Shard, |
| @@ -20,8 +22,8 @@ from torch.distributed.distributed_c10d import ReduceOp | |||
| 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 | 24 | ||
| 23 | -import torch_npu | 25 | + |
| 24 | -from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 26 | +npu = torch.ops.npu |
| 25 | 27 | ||
| 26 | 28 | ||
| 27 | def no_op(): | 29 | def no_op(): |
| @@ -79,9 +81,9 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 79 | device_mesh: DeviceMesh, | 81 | device_mesh: DeviceMesh, |
| 80 | placements: Sequence[Placement], | 82 | placements: Sequence[Placement], |
| 81 | op: Callable, | 83 | op: Callable, |
| 82 | - pre_op_fn: Optional[Callable] = None, | 84 | + pre_op_fn: Callable | None = None, |
| 83 | args: Sequence[Any] = tuple(), | 85 | args: Sequence[Any] = tuple(), |
| 84 | - kwargs: Optional[Dict[str, Any]] = None, | 86 | + kwargs: dict[str, Any] | None = None, |
| 85 | ): | 87 | ): |
| 86 | if pre_op_fn is None: | 88 | if pre_op_fn is None: |
| 87 | pre_op_fn = no_op | 89 | pre_op_fn = no_op |
| @@ -114,7 +116,7 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 114 | *, | 116 | *, |
| 115 | device_mesh: DeviceMesh, | 117 | device_mesh: DeviceMesh, |
| 116 | placements: Sequence[Placement], | 118 | placements: Sequence[Placement], |
| 117 | - pre_op_fn: Optional[Callable] = None, | 119 | + pre_op_fn: Callable | None = None, |
| 118 | input_size: Sequence[int], | 120 | input_size: Sequence[int], |
| 119 | op: Callable, | 121 | op: Callable, |
| 120 | **kwargs, | 122 | **kwargs, |
| @@ -140,8 +142,8 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 140 | 142 | ||
| 141 | def test_partial_add(self): | 143 | def test_partial_add(self): |
| 142 | device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | 144 | device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) |
| 143 | - d_1 = DTensor.from_local(torch.rand(2, 2), device_mesh, [_Partial()]) | 145 | + 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()]) | 146 | + d_2 = DTensor.from_local(torch.rand(2, 2), device_mesh, [Partial()]) |
| 145 | d_3 = d_1 + d_2 | 147 | d_3 = d_1 + d_2 |
| 146 | self.assertEqual(d_3._spec.placements[0].is_partial(), True) | 148 | self.assertEqual(d_3._spec.placements[0].is_partial(), True) |
| 147 | 149 | ||
| @@ -161,6 +163,30 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 161 | input_size=(8, 5), | 163 | input_size=(8, 5), |
| 162 | op=torch.nn.functional.gelu, | 164 | op=torch.nn.functional.gelu, |
| 163 | ) | 165 | ) |
| 166 | + self._run_sharded_elementwise_ops( | ||
| 167 | + device_mesh=device_mesh, | ||
| 168 | + placements=[Shard(0)], | ||
| 169 | + input_size=(8, 5), | ||
| 170 | + op=npu.fast_gelu, | ||
| 171 | + ) | ||
| 172 | + self._run_sharded_elementwise_ops( | ||
| 173 | + device_mesh=device_mesh, | ||
| 174 | + placements=[Replicate()], | ||
| 175 | + input_size=(8, 5), | ||
| 176 | + op=npu.fast_gelu, | ||
| 177 | + ) | ||
| 178 | + self._run_sharded_elementwise_ops( | ||
| 179 | + device_mesh=device_mesh, | ||
| 180 | + placements=[Shard(0)], | ||
| 181 | + input_size=(8, 5), | ||
| 182 | + op=npu.npu_fast_gelu, | ||
| 183 | + ) | ||
| 184 | + self._run_sharded_elementwise_ops( | ||
| 185 | + device_mesh=device_mesh, | ||
| 186 | + placements=[Replicate()], | ||
| 187 | + input_size=(8, 5), | ||
| 188 | + op=npu.npu_fast_gelu, | ||
| 189 | + ) | ||
| 164 | self._run_sharded_elementwise_ops( | 190 | self._run_sharded_elementwise_ops( |
| 165 | device_mesh=device_mesh, | 191 | device_mesh=device_mesh, |
| 166 | placements=[Shard(1)], | 192 | placements=[Shard(1)], |
| @@ -221,7 +247,7 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 221 | with self.assertRaisesRegex(RuntimeError, "supported"): | 247 | with self.assertRaisesRegex(RuntimeError, "supported"): |
| 222 | self._run_sharded_elementwise_ops( | 248 | self._run_sharded_elementwise_ops( |
| 223 | device_mesh=device_mesh, | 249 | device_mesh=device_mesh, |
| 224 | - placements=[_Partial(ReduceOp.SUM)], | 250 | + placements=[Partial(ReduceOp.SUM)], |
| 225 | input_size=(8, 5), | 251 | input_size=(8, 5), |
| 226 | op=torch.nn.functional.dropout, | 252 | op=torch.nn.functional.dropout, |
| 227 | ) | 253 | ) |
| @@ -257,16 +283,21 @@ class DistElementwiseOpsTest(DTensorTestBase): | |||
| 257 | local_result = torch_npu.npu_dtype_cast(npu_input, dst_dtype) | 283 | local_result = torch_npu.npu_dtype_cast(npu_input, dst_dtype) |
| 258 | 284 | ||
| 259 | # distributed tensor | 285 | # distributed tensor |
| 260 | - device_mesh = init_device_mesh(self.device_type, [1, 4], mesh_dim_names=["dp", "tp"]) | 286 | + device_mesh = init_device_mesh( |
| 287 | + self.device_type, [1, 4], mesh_dim_names=["dp", "tp"] | ||
| 288 | + ) | ||
| 261 | shard0_spec = Shard(0) | 289 | shard0_spec = Shard(0) |
| 262 | strided_shard0_spec = _StridedShard(0, split_factor=4) | 290 | strided_shard0_spec = _StridedShard(0, split_factor=4) |
| 263 | replica_spec = Replicate() | 291 | replica_spec = Replicate() |
| 264 | - dt_input = distribute_tensor(npu_input, device_mesh, [strided_shard0_spec, shard0_spec]) | 292 | + dt_input = distribute_tensor( |
| 293 | + npu_input, device_mesh, [strided_shard0_spec, shard0_spec] | ||
| 294 | + ) | ||
| 265 | dist_res: DTensor = torch_npu.npu_dtype_cast(dt_input, dst_dtype).redistribute( | 295 | dist_res: DTensor = torch_npu.npu_dtype_cast(dt_input, dst_dtype).redistribute( |
| 266 | device_mesh, [replica_spec, replica_spec] | 296 | device_mesh, [replica_spec, replica_spec] |
| 267 | ) | 297 | ) |
| 268 | self.assertEqual(dist_res.to_local().dtype, dst_dtype) | 298 | self.assertEqual(dist_res.to_local().dtype, dst_dtype) |
| 269 | self.assertEqual(dist_res.to_local(), local_result) | 299 | self.assertEqual(dist_res.to_local(), local_result) |
| 270 | 300 | ||
| 301 | + | ||
| 271 | if __name__ == "__main__": | 302 | if __name__ == "__main__": |
| 272 | run_tests() | 303 | run_tests() |
| @@ -1,12 +1,15 @@ | |||
| 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 |
| 5 | + | ||
| 6 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU, with_comms | ||
| 3 | 7 | ||
| 4 | import torch | 8 | import torch |
| 5 | import torch.distributed as dist | 9 | import torch.distributed as dist |
| 6 | from torch import rand, randn, Tensor | 10 | from torch import rand, randn, Tensor |
| 7 | -from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate, Shard | 11 | +from torch.distributed.tensor import DeviceMesh, distribute_tensor, Replicate, Shard |
| 8 | -from torch.distributed._tensor.debug import CommDebugMode | 12 | +from torch.distributed.tensor._ops._view_ops import ( |
| 9 | -from torch.distributed._tensor.ops._view_ops import ( | ||
| 10 | Broadcast, | 13 | Broadcast, |
| 11 | dim_maps, | 14 | dim_maps, |
| 12 | Flatten, | 15 | Flatten, |
| @@ -16,13 +19,14 @@ from torch.distributed._tensor.ops._view_ops import ( | |||
| 16 | Split, | 19 | Split, |
| 17 | view_groups, | 20 | view_groups, |
| 18 | ) | 21 | ) |
| 19 | -from torch.distributed._tensor.placement_types import Placement | 22 | +from torch.distributed.tensor.debug import CommDebugMode |
| 23 | +from torch.distributed.tensor.placement_types import Placement | ||
| 20 | from torch.testing._internal.common_utils import run_tests | 24 | from torch.testing._internal.common_utils import run_tests |
| 21 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase | 25 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase |
| 22 | from torch.utils import _pytree as pytree | 26 | from torch.utils import _pytree as pytree |
| 23 | 27 | ||
| 24 | -import torch_npu | 28 | + |
| 25 | -from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 29 | +npu = torch.ops.npu |
| 26 | 30 | ||
| 27 | 31 | ||
| 28 | class TestViewOps(DTensorTestBase): | 32 | class TestViewOps(DTensorTestBase): |
| @@ -152,7 +156,7 @@ class TestViewOps(DTensorTestBase): | |||
| 152 | if op == torch.unbind: | 156 | if op == torch.unbind: |
| 153 | no_shard_dims.add(kwargs.get("dim", 0)) | 157 | no_shard_dims.add(kwargs.get("dim", 0)) |
| 154 | 158 | ||
| 155 | - sharding_choices = cast(List[Placement], [Replicate()]) + [ | 159 | + sharding_choices = cast(list[Placement], [Replicate()]) + [ |
| 156 | Shard(i) for i, s in enumerate(in_shape) if s > 1 and i not in no_shard_dims | 160 | Shard(i) for i, s in enumerate(in_shape) if s > 1 and i not in no_shard_dims |
| 157 | ] | 161 | ] |
| 158 | 162 | ||
| @@ -358,6 +362,16 @@ class TestViewOps(DTensorTestBase): | |||
| 358 | (randn(24, 36, 28), (-1, -3, -2)), | 362 | (randn(24, 36, 28), (-1, -3, -2)), |
| 359 | (InputDim(2), InputDim(0), InputDim(1)), | 363 | (InputDim(2), InputDim(0), InputDim(1)), |
| 360 | ) | 364 | ) |
| 365 | + self.dimmap_test( | ||
| 366 | + npu.npu_transpose, | ||
| 367 | + (randn(24, 36, 28), (2, 0, 1)), | ||
| 368 | + (InputDim(2), InputDim(0), InputDim(1)), | ||
| 369 | + ) | ||
| 370 | + self.dimmap_test( | ||
| 371 | + npu.npu_transpose, | ||
| 372 | + (randn(24, 36, 28), (-1, -3, -2)), | ||
| 373 | + (InputDim(2), InputDim(0), InputDim(1)), | ||
| 374 | + ) | ||
| 361 | self.dimmap_test( | 375 | self.dimmap_test( |
| 362 | torch.ravel, | 376 | torch.ravel, |
| 363 | (randn(24, 36),), | 377 | (randn(24, 36),), |
| @@ -11,6 +11,7 @@ bStores | |||
| 11 | BU | 11 | BU |
| 12 | CANN | 12 | CANN |
| 13 | cann | 13 | cann |
| 14 | +childs | ||
| 14 | contiguities | 15 | contiguities |
| 15 | contiguity | 16 | contiguity |
| 16 | coo | 17 | coo |
| @@ -1,23 +1,24 @@ | |||
| 1 | __all__ = ["erase_stream", "matmul_checksum", "HiFloat8Tensor"] | 1 | __all__ = ["erase_stream", "matmul_checksum", "HiFloat8Tensor"] |
| 2 | 2 | ||
| 3 | +import atexit | ||
| 4 | +import ctypes | ||
| 3 | import os | 5 | import os |
| 4 | import sys | 6 | import sys |
| 5 | -import types | ||
| 6 | -import atexit | ||
| 7 | import traceback | 7 | import traceback |
| 8 | -import ctypes | 8 | +import types |
| 9 | import warnings | 9 | import warnings |
| 10 | - | ||
| 11 | from functools import wraps | 10 | from functools import wraps |
| 12 | 11 | ||
| 12 | + | ||
| 13 | # Disable autoloading before running 'import torch' to avoid circular dependencies | 13 | # Disable autoloading before running 'import torch' to avoid circular dependencies |
| 14 | ORG_AUTOLOAD = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1") | 14 | ORG_AUTOLOAD = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1") |
| 15 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" | 15 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" |
| 16 | 16 | ||
| 17 | import torch | 17 | import torch |
| 18 | +import torch_npu | ||
| 18 | from torch.distributed.fsdp import sharded_grad_scaler | 19 | from torch.distributed.fsdp import sharded_grad_scaler |
| 19 | from torch.utils.checkpoint import DefaultDeviceType | 20 | from torch.utils.checkpoint import DefaultDeviceType |
| 20 | -import torch_npu | 21 | + |
| 21 | 22 | ||
| 22 | acc = torch._C._get_accelerator() | 23 | acc = torch._C._get_accelerator() |
| 23 | if acc.type != "cpu": | 24 | if acc.type != "cpu": |
| @@ -27,15 +28,16 @@ if acc.type != "cpu": | |||
| 27 | error_code = "ERR00007" | 28 | error_code = "ERR00007" |
| 28 | error_code_msg = "feature not supported" | 29 | error_code_msg = "feature not supported" |
| 29 | submodule_name = "PTA" | 30 | submodule_name = "PTA" |
| 30 | - raise RuntimeError(f"Two accelerators cannot be used at the same time " | 31 | + raise RuntimeError( |
| 31 | - f"in PyTorch: npu and {acc.type}. You can install " | 32 | + f"Two accelerators cannot be used at the same time " |
| 32 | - f"the cpu version of PyTorch to use your npu device, " | 33 | + f"in PyTorch: npu and {acc.type}. You can install " |
| 33 | - f"or use the {acc.type} device with " | 34 | + f"the cpu version of PyTorch to use your npu device, " |
| 34 | - f"'export TORCH_DEVICE_BACKEND_AUTOLOAD=0'.\n" | 35 | + f"or use the {acc.type} device with " |
| 35 | - f"[ERROR] {time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime())} " | 36 | + f"'export TORCH_DEVICE_BACKEND_AUTOLOAD=0'.\n" |
| 36 | - f"(PID:{os.getpid()}, Device:-1, RankID:-1) " | 37 | + f"[ERROR] {time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime())} " |
| 37 | - f"{error_code} {submodule_name} {error_code_msg}" | 38 | + f"(PID:{os.getpid()}, Device:-1, RankID:-1) " |
| 38 | - ) | 39 | + f"{error_code} {submodule_name} {error_code_msg}" |
| 40 | + ) | ||
| 39 | 41 | ||
| 40 | try: | 42 | try: |
| 41 | import torch_npu.npu | 43 | import torch_npu.npu |
| @@ -43,58 +45,76 @@ except ImportError as e: | |||
| 43 | if "libhccl.so" in str(e): | 45 | if "libhccl.so" in str(e): |
| 44 | if "ASCEND_OPP_PATH" in os.environ: | 46 | if "ASCEND_OPP_PATH" in os.environ: |
| 45 | # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!! | 47 | # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!! |
| 46 | - e.msg += ". Please check that the compiler package is installed. "\ | 48 | + e.msg += ( |
| 47 | - "Please run 'source set_env.sh' in the CANN installation path." | 49 | + ". Please check that the compiler package is installed. " |
| 50 | + "Please run 'source set_env.sh' in the CANN installation path." | ||
| 51 | + ) | ||
| 48 | else: | 52 | else: |
| 49 | # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!! | 53 | # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!! |
| 50 | - e.msg += ". Please check that the cann package is installed. "\ | 54 | + e.msg += ( |
| 51 | - "Please run 'source set_env.sh' in the CANN installation path." | 55 | + ". Please check that the cann package is installed. " |
| 56 | + "Please run 'source set_env.sh' in the CANN installation path." | ||
| 57 | + ) | ||
| 52 | elif "libascendcl.so" in str(e): | 58 | elif "libascendcl.so" in str(e): |
| 53 | # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!! | 59 | # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!! |
| 54 | - e.msg += ". Please check that the runtime package is installed. "\ | 60 | + e.msg += ( |
| 55 | - "Please run 'source set_env.sh' in the CANN installation path." | 61 | + ". Please check that the runtime package is installed. " |
| 62 | + "Please run 'source set_env.sh' in the CANN installation path." | ||
| 63 | + ) | ||
| 56 | raise | 64 | raise |
| 57 | 65 | ||
| 58 | -import torch_npu.npu.amp | 66 | +import torch_npu._afd |
| 59 | -import torch_npu.npu.aclnn | ||
| 60 | -import torch_npu.optim | ||
| 61 | -import torch_npu.dynamo | ||
| 62 | import torch_npu._C | 67 | import torch_npu._C |
| 63 | import torch_npu._logging | 68 | import torch_npu._logging |
| 64 | -from torch_npu.utils import patch_getenv | 69 | +import torch_npu.distributed.rpc |
| 65 | -from torch_npu.utils.utils import _is_interactive_command_line | 70 | +import torch_npu.dynamo |
| 66 | -import torch_npu._afd | 71 | +import torch_npu.npu.aclnn |
| 67 | -from torch_npu import profiler | 72 | +import torch_npu.npu.amp |
| 68 | -from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler | 73 | +import torch_npu.op_plugin |
| 69 | -from torch_npu.contrib.function import npu_functional | 74 | +import torch_npu.optim |
| 70 | -from torch_npu.contrib.module import npu_modules | ||
| 71 | -from torch_npu.utils import _apply_module_patch, _add_tensor_methods, _add_collect_env_methods, \ | ||
| 72 | - _add_storage_methods, _add_serialization_methods, add_dynamo_methods, add_perf_dump_patch, \ | ||
| 73 | - add_optim_method, _inductor_register_device_op_overrides, \ | ||
| 74 | - _apply_npu_show_warning, _apply_npugraph_tree_methods, _apply_dlpack_patch, npu_patch_meta | ||
| 75 | -from torch_npu.utils._dynamo_device import _dynamo_register_interface_for_device | ||
| 76 | -from torch_npu.npu._format import _apply_npu_format_patch | ||
| 77 | import torch_npu.utils._afd_ops | 75 | import torch_npu.utils._afd_ops |
| 78 | import torch_npu.utils.custom_ops | 76 | import torch_npu.utils.custom_ops |
| 79 | -import torch_npu.distributed.rpc | 77 | +from torch_npu import _op_plugin_docs, profiler |
| 80 | -import torch_npu.op_plugin | 78 | +from torch_npu._C._distributed_c10d import ParallelStore |
| 81 | -from torch_npu.profiler._add_mstx_patch import _apply_mstx_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 | ||
| 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 | ||
| 87 | -from torch_npu.multiprocessing.reductions import _add_reductions_methods | ||
| 88 | -from torch_npu.npu.utils import _erase_stream as erase_stream | ||
| 89 | -from torch_npu.utils.hif8_tensor import _HiFloat8Tensor as HiFloat8Tensor | ||
| 90 | -from torch_npu.utils._error_code import ErrCode, pta_error, _except_handler | ||
| 91 | from torch_npu.asd.asd import _asd_patch | 79 | from torch_npu.asd.asd import _asd_patch |
| 92 | from torch_npu.asd.checksum import _matmul_checksum as matmul_checksum | 80 | from torch_npu.asd.checksum import _matmul_checksum as matmul_checksum |
| 81 | +from torch_npu.contrib.function import npu_functional | ||
| 82 | +from torch_npu.contrib.module import npu_modules | ||
| 83 | +from torch_npu.distributed.fsdp._add_fsdp_patch import _apply_fsdp_patch | ||
| 84 | +from torch_npu.distributed.rpc.backend_registry import _rpc_backend_registry | ||
| 93 | from torch_npu.dynamo import _patch_npu_trace_rules | 85 | from torch_npu.dynamo import _patch_npu_trace_rules |
| 94 | -from torch_npu._C._distributed_c10d import ParallelStore | 86 | +from torch_npu.multiprocessing.reductions import _add_reductions_methods |
| 87 | +from torch_npu.npu._format import _apply_npu_format_patch | ||
| 88 | +from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler | ||
| 89 | +from torch_npu.npu.utils import _erase_stream as erase_stream | ||
| 95 | from torch_npu.op_plugin.meta import _meta_registrations | 90 | from torch_npu.op_plugin.meta import _meta_registrations |
| 91 | +from torch_npu.profiler._add_mstx_patch import _apply_mstx_patch | ||
| 92 | +from torch_npu.utils import ( | ||
| 93 | + _add_collect_env_methods, | ||
| 94 | + _add_intercept_methods, | ||
| 95 | + _add_serialization_methods, | ||
| 96 | + _add_storage_methods, | ||
| 97 | + _add_tensor_methods, | ||
| 98 | + _apply_dlpack_patch, | ||
| 99 | + _apply_module_patch, | ||
| 100 | + _apply_npu_show_warning, | ||
| 101 | + _apply_npugraph_tree_methods, | ||
| 102 | + _cann_package_check, | ||
| 103 | + _inductor_register_device_op_overrides, | ||
| 104 | + add_dynamo_methods, | ||
| 105 | + add_optim_method, | ||
| 106 | + add_perf_dump_patch, | ||
| 107 | + npu_patch_meta, | ||
| 108 | + patch_getenv, | ||
| 109 | +) | ||
| 110 | +from torch_npu.utils._dynamo_device import _dynamo_register_interface_for_device | ||
| 111 | +from torch_npu.utils._error_code import _except_handler, ErrCode, pta_error | ||
| 112 | +from torch_npu.utils.exposed_api import public_npu_functions | ||
| 113 | +from torch_npu.utils.hif8_tensor import _HiFloat8Tensor as HiFloat8Tensor | ||
| 114 | +from torch_npu.utils.utils import _is_interactive_command_line | ||
| 96 | from torch_npu.version import __version__ as __version__ | 115 | from torch_npu.version import __version__ as __version__ |
| 97 | -from torch_npu import _op_plugin_docs | 116 | + |
| 117 | + | ||
| 98 | del _op_plugin_docs | 118 | del _op_plugin_docs |
| 99 | 119 | ||
| 100 | _cann_package_check() | 120 | _cann_package_check() |
| @@ -103,14 +123,16 @@ _cann_package_check() | |||
| 103 | def _wrap_torch_error_func(func): | 123 | def _wrap_torch_error_func(func): |
| 104 | 124 | ||
| 105 | def wrapper(*args, **kwargs): | 125 | def wrapper(*args, **kwargs): |
| 106 | - raise RuntimeError(f"torch.{func.__name__} is deprecated and will be removed in future version. " | 126 | + raise RuntimeError( |
| 107 | - f"Use torch_npu.{func.__name__} instead." + pta_error(ErrCode.NOT_SUPPORT)) | 127 | + f"torch.{func.__name__} is deprecated and will be removed in future version. " |
| 128 | + f"Use torch_npu.{func.__name__} instead." + pta_error(ErrCode.NOT_SUPPORT) | ||
| 129 | + ) | ||
| 108 | 130 | ||
| 109 | return wrapper | 131 | return wrapper |
| 110 | 132 | ||
| 111 | 133 | ||
| 112 | for name in dir(torch.ops.npu): | 134 | for name in dir(torch.ops.npu): |
| 113 | - if name.startswith('__') or name in ['_dir', 'name']: | 135 | + if name.startswith("__") or name in ["_dir", "name"]: |
| 114 | continue | 136 | continue |
| 115 | globals()[name] = getattr(torch.ops.npu, name) | 137 | globals()[name] = getattr(torch.ops.npu, name) |
| 116 | if name in public_npu_functions: | 138 | if name in public_npu_functions: |
| @@ -118,7 +140,7 @@ for name in dir(torch.ops.npu): | |||
| 118 | setattr(torch, name, _wrap_torch_error_func(getattr(torch.ops.npu, name))) | 140 | setattr(torch, name, _wrap_torch_error_func(getattr(torch.ops.npu, name))) |
| 119 | 141 | ||
| 120 | for name in dir(torch_npu._C._cd.DType): | 142 | for name in dir(torch_npu._C._cd.DType): |
| 121 | - if name.startswith('__') or name in ['_dir', 'name']: | 143 | + if name.startswith("__") or name in ["_dir", "name"]: |
| 122 | continue | 144 | continue |
| 123 | setattr(torch_npu, name, getattr(torch_npu._C._cd.DType, name)) | 145 | setattr(torch_npu, name, getattr(torch_npu._C._cd.DType, name)) |
| 124 | 146 | ||
| @@ -136,26 +158,29 @@ def _apply_patches(monkey_patches): | |||
| 136 | if hasattr(root_module, module_list[0]): | 158 | if hasattr(root_module, module_list[0]): |
| 137 | return _getattr(module_list[1:], getattr(root_module, module_list[0])) | 159 | return _getattr(module_list[1:], getattr(root_module, module_list[0])) |
| 138 | else: | 160 | else: |
| 139 | - empty_module_name = f'{root_module.__name__}.{module_list[0]}' | 161 | + empty_module_name = f"{root_module.__name__}.{module_list[0]}" |
| 140 | sys.modules[empty_module_name] = types.ModuleType(empty_module_name) | 162 | sys.modules[empty_module_name] = types.ModuleType(empty_module_name) |
| 141 | setattr(root_module, module_list[0], sys.modules.get(empty_module_name)) | 163 | setattr(root_module, module_list[0], sys.modules.get(empty_module_name)) |
| 142 | return _getattr(module_list[1:], getattr(root_module, module_list[0])) | 164 | return _getattr(module_list[1:], getattr(root_module, module_list[0])) |
| 143 | 165 | ||
| 144 | for patch_pair in monkey_patches: | 166 | for patch_pair in monkey_patches: |
| 145 | dest, patch = patch_pair | 167 | dest, patch = patch_pair |
| 146 | - dest_module = _getattr(dest.split('.'), root_module=torch) | 168 | + dest_module = _getattr(dest.split("."), root_module=torch) |
| 147 | last_module_level = dest.split(".")[-1] | 169 | last_module_level = dest.split(".")[-1] |
| 148 | if not isinstance(patch, types.ModuleType): | 170 | if not isinstance(patch, types.ModuleType): |
| 149 | setattr(dest_module, last_module_level, patch) | 171 | setattr(dest_module, last_module_level, patch) |
| 150 | continue | 172 | continue |
| 151 | 173 | ||
| 152 | - if not hasattr(dest_module, last_module_level) or not hasattr(patch, '__all__'): | 174 | + if not hasattr(dest_module, last_module_level) or not hasattr(patch, "__all__"): |
| 153 | setattr(dest_module, last_module_level, patch) | 175 | setattr(dest_module, last_module_level, patch) |
| 154 | - sys.modules[f'{dest_module.__name__}.{last_module_level}'] = patch | 176 | + sys.modules[f"{dest_module.__name__}.{last_module_level}"] = patch |
| 155 | continue | 177 | continue |
| 156 | 178 | ||
| 157 | - if not hasattr(patch, '__all__'): | 179 | + if not hasattr(patch, "__all__"): |
| 158 | - raise NotImplementedError("Patch module must have __all__ definition." + pta_error(ErrCode.NOT_SUPPORT)) | 180 | + raise NotImplementedError( |
| 181 | + "Patch module must have __all__ definition." | ||
| 182 | + + pta_error(ErrCode.NOT_SUPPORT) | ||
| 183 | + ) | ||
| 159 | dest_module = getattr(dest_module, last_module_level) | 184 | dest_module = getattr(dest_module, last_module_level) |
| 160 | for attr in patch.__all__: | 185 | for attr in patch.__all__: |
| 161 | setattr(dest_module, attr, getattr(patch, attr)) | 186 | setattr(dest_module, attr, getattr(patch, attr)) |
| @@ -188,29 +213,64 @@ def _apply_class_patches(): | |||
| 188 | 213 | ||
| 189 | 214 | ||
| 190 | def _apply_distributed_methods_patch(): | 215 | def _apply_distributed_methods_patch(): |
| 191 | - torch._C._distributed_c10d._verify_params_across_processes = torch_npu.distributed._verify_params_across_processes | 216 | + torch._C._distributed_c10d._verify_params_across_processes = ( |
| 192 | - torch.distributed.batch_isend_irecv = torch_npu.distributed.distributed_c10d._batch_isend_irecv | 217 | + torch_npu.distributed._verify_params_across_processes |
| 193 | - torch.distributed.distributed_c10d.batch_isend_irecv = torch_npu.distributed.distributed_c10d._batch_isend_irecv | 218 | + ) |
| 219 | + torch.distributed.batch_isend_irecv = ( | ||
| 220 | + torch_npu.distributed.distributed_c10d._batch_isend_irecv | ||
| 221 | + ) | ||
| 222 | + torch.distributed.distributed_c10d.batch_isend_irecv = ( | ||
| 223 | + torch_npu.distributed.distributed_c10d._batch_isend_irecv | ||
| 224 | + ) | ||
| 194 | torch.distributed.gather = torch_npu.distributed.distributed_c10d._gather | 225 | torch.distributed.gather = torch_npu.distributed.distributed_c10d._gather |
| 195 | - torch.distributed.distributed_c10d.gather = torch_npu.distributed.distributed_c10d._gather | 226 | + torch.distributed.distributed_c10d.gather = ( |
| 196 | - torch.distributed.gather_object = torch_npu.distributed.distributed_c10d._gather_object | 227 | + torch_npu.distributed.distributed_c10d._gather |
| 197 | - torch.distributed.distributed_c10d.gather_object = torch_npu.distributed.distributed_c10d._gather_object | 228 | + ) |
| 229 | + torch.distributed.gather_object = ( | ||
| 230 | + torch_npu.distributed.distributed_c10d._gather_object | ||
| 231 | + ) | ||
| 232 | + torch.distributed.distributed_c10d.gather_object = ( | ||
| 233 | + torch_npu.distributed.distributed_c10d._gather_object | ||
| 234 | + ) | ||
| 198 | torch.distributed.is_hccl_available = torch_npu.distributed.is_hccl_available | 235 | torch.distributed.is_hccl_available = torch_npu.distributed.is_hccl_available |
| 199 | torch.distributed.reinit_process_group = torch_npu.distributed.reinit_process_group | 236 | torch.distributed.reinit_process_group = torch_npu.distributed.reinit_process_group |
| 200 | - torch.distributed.distributed_c10d.rendezvous = torch_npu.distributed.distributed_c10d._trigger_rendezvous_decorator(torch.distributed.distributed_c10d.rendezvous) | 237 | + torch.distributed.distributed_c10d.rendezvous = ( |
| 201 | - torch.distributed.launcher.api._get_addr_and_port = torch_npu.distributed.distributed_c10d._trigger__get_addr_and_port_decorator(torch.distributed.launcher.api._get_addr_and_port) | 238 | + torch_npu.distributed.distributed_c10d._trigger_rendezvous_decorator( |
| 239 | + torch.distributed.distributed_c10d.rendezvous | ||
| 240 | + ) | ||
| 241 | + ) | ||
| 242 | + torch.distributed.launcher.api._get_addr_and_port = ( | ||
| 243 | + torch_npu.distributed.distributed_c10d._trigger__get_addr_and_port_decorator( | ||
| 244 | + torch.distributed.launcher.api._get_addr_and_port | ||
| 245 | + ) | ||
| 246 | + ) | ||
| 202 | torch._C._distributed_c10d.ProcessGroup._get_sequence_number_for_group = ( | 247 | torch._C._distributed_c10d.ProcessGroup._get_sequence_number_for_group = ( |
| 203 | - torch_npu.distributed.distributed_c10d._hccl_get_sequence_number_for_group) | 248 | + torch_npu.distributed.distributed_c10d._hccl_get_sequence_number_for_group |
| 204 | - torch.distributed.nn.functional._AllGatherBase.backward = torch_npu.distributed.nn.functional._allgather_base_backward_hccl | 249 | + ) |
| 205 | - torch.distributed.distributed_c10d._add_ephemeral_timeout_for_all_pgs = torch_npu.distributed.distributed_c10d._hccl_add_ephemeral_timeout_for_all_pgs | 250 | + torch.distributed.nn.functional._AllGatherBase.backward = ( |
| 251 | + torch_npu.distributed.nn.functional._allgather_base_backward_hccl | ||
| 252 | + ) | ||
| 253 | + torch.distributed.distributed_c10d._add_ephemeral_timeout_for_all_pgs = ( | ||
| 254 | + torch_npu.distributed.distributed_c10d._hccl_add_ephemeral_timeout_for_all_pgs | ||
| 255 | + ) | ||
| 206 | 256 | ||
| 207 | 257 | ||
| 208 | torch.utils.rename_privateuse1_backend("npu") | 258 | torch.utils.rename_privateuse1_backend("npu") |
| 209 | # rename device name to 'npu' and register funcs | 259 | # rename device name to 'npu' and register funcs |
| 210 | -torch._register_device_module('npu', torch_npu.npu) | 260 | +torch._register_device_module("npu", torch_npu.npu) |
| 211 | -unsupported_dtype = [torch.quint8, torch.quint4x2, torch.quint2x4, torch.qint32, torch.qint8] | 261 | +unsupported_dtype = [ |
| 212 | -torch.utils.generate_methods_for_privateuse1_backend(for_tensor=True, for_module=True, for_storage=True, | 262 | + torch.quint8, |
| 213 | - unsupported_dtype=unsupported_dtype) | 263 | + torch.quint4x2, |
| 264 | + torch.quint2x4, | ||
| 265 | + torch.qint32, | ||
| 266 | + torch.qint8, | ||
| 267 | +] | ||
| 268 | +torch.utils.generate_methods_for_privateuse1_backend( | ||
| 269 | + for_tensor=True, | ||
| 270 | + for_module=True, | ||
| 271 | + for_storage=True, | ||
| 272 | + unsupported_dtype=unsupported_dtype, | ||
| 273 | +) | ||
| 214 | torch.nn.parameter.UninitializedTensorMixin._allowed_methods.append(torch.Tensor.npu) | 274 | torch.nn.parameter.UninitializedTensorMixin._allowed_methods.append(torch.Tensor.npu) |
| 215 | 275 | ||
| 216 | # register npu device interface for dynamo | 276 | # register npu device interface for dynamo |
| @@ -223,8 +283,14 @@ _asd_patch() | |||
| 223 | _except_handler.patch_excepthook() | 283 | _except_handler.patch_excepthook() |
| 224 | 284 | ||
| 225 | _warn_msg = { | 285 | _warn_msg = { |
| 226 | - "DropoutWithByteMask" : "torch.nn.DropoutWithByteMask is deprecated and will be removed in future version. Use torch_npu.contrib.module.DropoutWithByteMask instead.", | 286 | + "DropoutWithByteMask": ( |
| 227 | - "dropout_with_byte_mask" : "torch.nn.functional.dropout_with_byte_mask is deprecated and will be removed in future version. Use torch_npu.contrib.function.dropout_with_byte_mask instead.", | 287 | + "torch.nn.DropoutWithByteMask is deprecated and will be removed in future version. " |
| 288 | + "Use torch_npu.contrib.module.DropoutWithByteMask instead." | ||
| 289 | + ), | ||
| 290 | + "dropout_with_byte_mask": ( | ||
| 291 | + "torch.nn.functional.dropout_with_byte_mask is deprecated and will be removed in future version. " | ||
| 292 | + "Use torch_npu.contrib.function.dropout_with_byte_mask instead." | ||
| 293 | + ), | ||
| 228 | } | 294 | } |
| 229 | 295 | ||
| 230 | 296 | ||
| @@ -233,9 +299,16 @@ def _wrap_torch_patch_warning_func(func): | |||
| 233 | def wrapper(*args, **kwargs): | 299 | def wrapper(*args, **kwargs): |
| 234 | warnings.warn(_warn_msg[func.__name__]) | 300 | warnings.warn(_warn_msg[func.__name__]) |
| 235 | return func(*args, **kwargs) | 301 | return func(*args, **kwargs) |
| 302 | + | ||
| 236 | return wrapper | 303 | return wrapper |
| 237 | -setattr(torch.nn, "DropoutWithByteMask", _wrap_torch_patch_warning_func(torch.nn.DropoutWithByteMask)) | 304 | + |
| 238 | -setattr(torch.nn.functional, "dropout_with_byte_mask", _wrap_torch_patch_warning_func(torch.nn.functional.dropout_with_byte_mask)) | 305 | + |
| 306 | +torch.nn.DropoutWithByteMask = _wrap_torch_patch_warning_func( | ||
| 307 | + torch.nn.DropoutWithByteMask | ||
| 308 | +) | ||
| 309 | +torch.nn.functional.dropout_with_byte_mask = _wrap_torch_patch_warning_func( | ||
| 310 | + torch.nn.functional.dropout_with_byte_mask | ||
| 311 | +) | ||
| 239 | # this must be placed at the end | 312 | # this must be placed at the end |
| 240 | torch_npu._C._initExtension() | 313 | torch_npu._C._initExtension() |
| 241 | 314 | ||
| @@ -244,32 +317,49 @@ def _new_process_group_hccl_helper(dist_backend_opts, pg_options): | |||
| 244 | store = dist_backend_opts.store | 317 | store = dist_backend_opts.store |
| 245 | group_rank = dist_backend_opts.group_rank | 318 | group_rank = dist_backend_opts.group_rank |
| 246 | group_size = dist_backend_opts.group_size | 319 | group_size = dist_backend_opts.group_size |
| 247 | - if pg_options is None or not isinstance(pg_options, torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options): | 320 | + if pg_options is None or not isinstance( |
| 321 | + pg_options, torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options | ||
| 322 | + ): | ||
| 248 | pg_options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options() | 323 | pg_options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options() |
| 249 | pg_options.is_high_priority_stream = False | 324 | pg_options.is_high_priority_stream = False |
| 250 | pg_options._timeout = dist_backend_opts.timeout | 325 | pg_options._timeout = dist_backend_opts.timeout |
| 251 | pg_options.global_ranks_in_group = dist_backend_opts.global_ranks_in_group | 326 | pg_options.global_ranks_in_group = dist_backend_opts.global_ranks_in_group |
| 252 | pg_options.group_id = dist_backend_opts.group_id | 327 | pg_options.group_id = dist_backend_opts.group_id |
| 253 | - return torch_npu._C._distributed_c10d.ProcessGroupHCCL(store, group_rank, group_size, pg_options) | 328 | + return torch_npu._C._distributed_c10d.ProcessGroupHCCL( |
| 329 | + store, group_rank, group_size, pg_options | ||
| 330 | + ) | ||
| 254 | 331 | ||
| 255 | 332 | ||
| 256 | def _new_process_group_lccl_helper(dist_backend_opts, pg_options): | 333 | def _new_process_group_lccl_helper(dist_backend_opts, pg_options): |
| 257 | store = dist_backend_opts.store | 334 | store = dist_backend_opts.store |
| 258 | group_rank = dist_backend_opts.group_rank | 335 | group_rank = dist_backend_opts.group_rank |
| 259 | group_size = dist_backend_opts.group_size | 336 | group_size = dist_backend_opts.group_size |
| 260 | - return torch_npu._C._distributed_c10d.ProcessGroupLCCL(store, group_rank, group_size) | 337 | + return torch_npu._C._distributed_c10d.ProcessGroupLCCL( |
| 338 | + store, group_rank, group_size | ||
| 339 | + ) | ||
| 261 | 340 | ||
| 262 | 341 | ||
| 263 | def _register_distributed_backend_for_npu(): | 342 | def _register_distributed_backend_for_npu(): |
| 264 | # init and register hccl backend | 343 | # init and register hccl backend |
| 265 | # Note: Since torch 2.8, the hccl backend must be registered at first to keep a right default_device_backend_map | 344 | # Note: Since torch 2.8, the hccl backend must be registered at first to keep a right default_device_backend_map |
| 266 | - torch.distributed.Backend.register_backend("hccl", lambda dist_backend_opts, pg_options: | 345 | + torch.distributed.Backend.register_backend( |
| 267 | - _new_process_group_hccl_helper(dist_backend_opts, pg_options), extended_api=True, devices=["npu"]) | 346 | + "hccl", |
| 347 | + lambda dist_backend_opts, pg_options: _new_process_group_hccl_helper( | ||
| 348 | + dist_backend_opts, pg_options | ||
| 349 | + ), | ||
| 350 | + extended_api=True, | ||
| 351 | + devices=["npu"], | ||
| 352 | + ) | ||
| 268 | 353 | ||
| 269 | # init and register lccl backend | 354 | # init and register lccl backend |
| 270 | - torch.distributed.Backend.register_backend("lccl", lambda dist_backend_opts, pg_options: | 355 | + torch.distributed.Backend.register_backend( |
| 271 | - _new_process_group_lccl_helper(dist_backend_opts, pg_options), extended_api=True, devices=["npu"]) | 356 | + "lccl", |
| 272 | - | 357 | + lambda dist_backend_opts, pg_options: _new_process_group_lccl_helper( |
| 358 | + dist_backend_opts, pg_options | ||
| 359 | + ), | ||
| 360 | + extended_api=True, | ||
| 361 | + devices=["npu"], | ||
| 362 | + ) | ||
| 273 | 363 | ||
| 274 | 364 | ||
| 275 | # init and register distributed backend | 365 | # init and register distributed backend |
| @@ -290,6 +380,7 @@ def _npu_shutdown(): | |||
| 290 | torch_npu.asd.asd.matmul_check._cleanup() | 380 | torch_npu.asd.asd.matmul_check._cleanup() |
| 291 | if torch_npu.npu.aclnn._use_static_aclnn_kernel: | 381 | if torch_npu.npu.aclnn._use_static_aclnn_kernel: |
| 292 | from torch_npu._inductor.npu_static_kernel import uninstall_static_kernel | 382 | from torch_npu._inductor.npu_static_kernel import uninstall_static_kernel |
| 383 | + | ||
| 293 | uninstall_static_kernel() | 384 | uninstall_static_kernel() |
| 294 | 385 | ||
| 295 | 386 | ||
| @@ -299,11 +390,8 @@ atexit.register(_npu_shutdown) | |||
| 299 | # init and register rpc npu backend | 390 | # init and register rpc npu backend |
| 300 | _rpc_backend_registry() | 391 | _rpc_backend_registry() |
| 301 | 392 | ||
| 302 | -# register rules for ops in dtensor | ||
| 303 | -_register_ops_under_dtensor_rules() | ||
| 304 | - | ||
| 305 | # Enable NPU Sanitizer | 393 | # Enable NPU Sanitizer |
| 306 | -if 'TORCH_NPU_SANITIZER' in os.environ: | 394 | +if "TORCH_NPU_SANITIZER" in os.environ: |
| 307 | import torch_npu.npu._sanitizer as csan | 395 | import torch_npu.npu._sanitizer as csan |
| 308 | 396 | ||
| 309 | csan.enable_npu_sanitizer() | 397 | csan.enable_npu_sanitizer() |
| @@ -315,20 +403,24 @@ _inductor_register_device_op_overrides() | |||
| 315 | _patch_npu_trace_rules() | 403 | _patch_npu_trace_rules() |
| 316 | 404 | ||
| 317 | if _is_interactive_command_line(): | 405 | if _is_interactive_command_line(): |
| 318 | - os.environ["TASK_QUEUE_ENABLE"] = '0' | 406 | + os.environ["TASK_QUEUE_ENABLE"] = "0" |
| 319 | - warnings.warn("On the interactive interface, the value of TASK_QUEUE_ENABLE is set to 0 by default. \ | 407 | + warnings.warn( |
| 320 | - Do not set it to 1 to prevent some unknown errors") | 408 | + "On the interactive interface, the value of TASK_QUEUE_ENABLE is set to 0 by default. \ |
| 409 | + Do not set it to 1 to prevent some unknown errors" | ||
| 410 | + ) | ||
| 321 | 411 | ||
| 322 | # Enable transfer_to_npu via environment variable | 412 | # Enable transfer_to_npu via environment variable |
| 323 | _transfer_to_npu_env = os.getenv("TORCH_TRANSFER_TO_NPU", "0") | 413 | _transfer_to_npu_env = os.getenv("TORCH_TRANSFER_TO_NPU", "0") |
| 324 | if _transfer_to_npu_env == "1": | 414 | if _transfer_to_npu_env == "1": |
| 325 | from torch_npu.contrib import transfer_to_npu | 415 | from torch_npu.contrib import transfer_to_npu |
| 326 | elif _transfer_to_npu_env != "0": | 416 | elif _transfer_to_npu_env != "0": |
| 327 | - raise ValueError(f"Invalid value for TORCH_TRANSFER_TO_NPU: {_transfer_to_npu_env}. Only '0' or '1' is supported.") | 417 | + raise ValueError( |
| 418 | + f"Invalid value for TORCH_TRANSFER_TO_NPU: {_transfer_to_npu_env}. Only '0' or '1' is supported." | ||
| 419 | + ) | ||
| 328 | 420 | ||
| 329 | 421 | ||
| 330 | # This function is an entrypoint called by PyTorch | 422 | # This function is an entrypoint called by PyTorch |
| 331 | # when running 'import torch'. There is no need to do anything. | 423 | # when running 'import torch'. There is no need to do anything. |
| 332 | def _autoload(): | 424 | def _autoload(): |
| 333 | # We should restore this switch as sub processes need to inherit its value | 425 | # We should restore this switch as sub processes need to inherit its value |
| 334 | - os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD | 426 | + os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD |
| @@ -1,7 +1,8 @@ | |||
| 1 | -import torch_npu.distributed.tensor._dtensor_patch # patch before register strategy | ||
| 2 | -import torch_npu.distributed.tensor._pointwise_ops | ||
| 3 | -import torch_npu.distributed.tensor._matrix_ops | ||
| 4 | import torch_npu.distributed.tensor._attention | 1 | import torch_npu.distributed.tensor._attention |
| 2 | +import torch_npu.distributed.tensor._dtensor_patch # patch before register strategy | ||
| 5 | import torch_npu.distributed.tensor._math_ops | 3 | import torch_npu.distributed.tensor._math_ops |
| 4 | +import torch_npu.distributed.tensor._matrix_ops | ||
| 6 | import torch_npu.distributed.tensor._moe_ops | 5 | import torch_npu.distributed.tensor._moe_ops |
| 6 | +import torch_npu.distributed.tensor._pointwise_ops | ||
| 7 | import torch_npu.distributed.tensor._sharded_tensor_patch | 7 | import torch_npu.distributed.tensor._sharded_tensor_patch |
| 8 | +import torch_npu.distributed.tensor._view_ops | ||
| @@ -1,23 +1,24 @@ | |||
| 1 | -from typing import cast, Dict, List, Optional, Tuple, Union | ||
| 2 | import os | 1 | import os |
| 2 | +from typing import cast | ||
| 3 | 3 | ||
| 4 | import torch | 4 | import torch |
| 5 | +import torch_npu | ||
| 5 | from torch.distributed._tensor.experimental import register_sharding | 6 | from torch.distributed._tensor.experimental import register_sharding |
| 6 | -from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta | ||
| 7 | -from torch.distributed.tensor._ops.registration import register_op_strategy | ||
| 8 | -from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy | ||
| 9 | from torch.distributed.tensor import DTensor, Partial, Replicate, Shard | 7 | from torch.distributed.tensor import DTensor, Partial, Replicate, Shard |
| 8 | +from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta | ||
| 10 | from torch.distributed.tensor._op_schema import ( | 9 | from torch.distributed.tensor._op_schema import ( |
| 11 | OpInfo, | 10 | OpInfo, |
| 12 | OpSchema, | 11 | OpSchema, |
| 13 | - OpStrategy, | ||
| 14 | OpSpec, | 12 | OpSpec, |
| 13 | + OpStrategy, | ||
| 15 | OutputSharding, | 14 | OutputSharding, |
| 16 | RuntimeSchemaInfo, | 15 | RuntimeSchemaInfo, |
| 17 | - TupleStrategy | 16 | + TupleStrategy, |
| 18 | ) | 17 | ) |
| 18 | +from torch.distributed.tensor._ops._matrix_ops import _mm_like_strategy | ||
| 19 | +from torch.distributed.tensor._ops.registration import register_op_strategy | ||
| 20 | +from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy | ||
| 19 | 21 | ||
| 20 | -import torch_npu | ||
| 21 | 22 | ||
| 22 | try: | 23 | try: |
| 23 | from torch.utils import _cxx_pytree as pytree | 24 | from torch.utils import _cxx_pytree as pytree |
| @@ -25,11 +26,12 @@ except ImportError: | |||
| 25 | from torch.utils import _pytree as pytree | 26 | from torch.utils import _pytree as pytree |
| 26 | 27 | ||
| 27 | from ._common import ( | 28 | from ._common import ( |
| 29 | + get_empty_local_results, | ||
| 28 | get_redistributed_local_args, | 30 | get_redistributed_local_args, |
| 29 | get_redistributed_local_kwargs, | 31 | get_redistributed_local_kwargs, |
| 30 | - get_empty_local_results | ||
| 31 | ) | 32 | ) |
| 32 | 33 | ||
| 34 | + | ||
| 33 | aten = torch.ops.aten | 35 | aten = torch.ops.aten |
| 34 | npu = torch.ops.npu | 36 | npu = torch.ops.npu |
| 35 | 37 | ||
| @@ -37,7 +39,9 @@ npu = torch.ops.npu | |||
| 37 | def _get_max_shardable_dim(tensor): | 39 | def _get_max_shardable_dim(tensor): |
| 38 | shape = tensor.shape | 40 | shape = tensor.shape |
| 39 | world_size = torch.distributed.get_world_size() | 41 | world_size = torch.distributed.get_world_size() |
| 40 | - divisible_dims = [(idx, dim) for idx, dim in enumerate(shape) if dim % world_size == 0] | 42 | + divisible_dims = [ |
| 43 | + (idx, dim) for idx, dim in enumerate(shape) if dim % world_size == 0 | ||
| 44 | + ] | ||
| 41 | if divisible_dims: | 45 | if divisible_dims: |
| 42 | idx, _ = max(divisible_dims, key=lambda x: x[1]) | 46 | idx, _ = max(divisible_dims, key=lambda x: x[1]) |
| 43 | return idx | 47 | return idx |
| @@ -45,7 +49,7 @@ def _get_max_shardable_dim(tensor): | |||
| 45 | return -1 | 49 | return -1 |
| 46 | 50 | ||
| 47 | 51 | ||
| 48 | -def _handle_tensor_list_in_kwargs(kwargs: Dict[str, object], op_info: OpInfo) -> None: | 52 | +def _handle_tensor_list_in_kwargs(kwargs: dict[str, object], op_info: OpInfo) -> None: |
| 49 | for key, value in kwargs.items(): | 53 | for key, value in kwargs.items(): |
| 50 | if isinstance(value, list) and all(isinstance(e, DTensor) for e in value): | 54 | if isinstance(value, list) and all(isinstance(e, DTensor) for e in value): |
| 51 | new_schema = [] | 55 | new_schema = [] |
| @@ -53,12 +57,16 @@ def _handle_tensor_list_in_kwargs(kwargs: Dict[str, object], op_info: OpInfo) -> | |||
| 53 | for dtensor in value: | 57 | for dtensor in value: |
| 54 | new_schema.append(dtensor._spec) | 58 | new_schema.append(dtensor._spec) |
| 55 | new_local_tensors.append(dtensor._local_tensor) | 59 | new_local_tensors.append(dtensor._local_tensor) |
| 56 | - op_info.schema.kwargs_schema[key] = tuple(new_schema) # list is not hashable for cache | 60 | + op_info.schema.kwargs_schema[key] = tuple( |
| 61 | + new_schema | ||
| 62 | + ) # list is not hashable for cache | ||
| 57 | op_info.local_kwargs[key] = new_local_tensors | 63 | op_info.local_kwargs[key] = new_local_tensors |
| 58 | 64 | ||
| 59 | op_info.schema._recompute_comparison_key() | 65 | op_info.schema._recompute_comparison_key() |
| 60 | 66 | ||
| 67 | + | ||
| 61 | if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | 68 | if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": |
| 69 | + | ||
| 62 | 70 | ||
| 63 | def custom_matmul_strategy( | 71 | def custom_matmul_strategy( |
| 64 | tensor1: DTensorSpec, | 72 | tensor1: DTensorSpec, |
| @@ -111,13 +119,24 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 111 | diff = abs(len1 - len2) | 119 | diff = abs(len1 - len2) |
| 112 | is_shape1_longer = len1 > len2 | 120 | is_shape1_longer = len1 > len2 |
| 113 | 121 | ||
| 114 | - for i in range(min(len1, len2) - 3, -1, -1): | 122 | + for i in range(min(len1, len2) - 3, -1, -1): |
| 115 | - shape1_shardable = shape1[i + diff] % tensor1.mesh.size(0) == 0 if is_shape1_longer else shape1[i] % tensor1.mesh.size(0) == 0 | 123 | + shape1_shardable = ( |
| 116 | - shape2_shardable = shape2[i] % tensor2.mesh.size(0) == 0 if is_shape1_longer else shape2[i + diff] % tensor2.mesh.size(0) == 0 | 124 | + shape1[i + diff] % tensor1.mesh.size(0) == 0 |
| 125 | + if is_shape1_longer | ||
| 126 | + else shape1[i] % tensor1.mesh.size(0) == 0 | ||
| 127 | + ) | ||
| 128 | + shape2_shardable = ( | ||
| 129 | + shape2[i] % tensor2.mesh.size(0) == 0 | ||
| 130 | + if is_shape1_longer | ||
| 131 | + else shape2[i + diff] % tensor2.mesh.size(0) == 0 | ||
| 132 | + ) | ||
| 117 | 133 | ||
| 118 | if shape1_shardable and shape2_shardable: | 134 | if shape1_shardable and shape2_shardable: |
| 119 | - strategy_batch = ([Shard(i + diff)], [Shard(i + diff), Shard(i)]) if is_shape1_longer \ | 135 | + strategy_batch = ( |
| 136 | + ([Shard(i + diff)], [Shard(i + diff), Shard(i)]) | ||
| 137 | + if is_shape1_longer | ||
| 120 | else ([Shard(i + diff)], [Shard(i), Shard(i + diff)]) | 138 | else ([Shard(i + diff)], [Shard(i), Shard(i + diff)]) |
| 139 | + ) | ||
| 121 | acceptable_shardings.append(strategy_batch) | 140 | acceptable_shardings.append(strategy_batch) |
| 122 | 141 | ||
| 123 | for i in range(diff - 1, -1, -1): | 142 | for i in range(diff - 1, -1, -1): |
| @@ -129,26 +148,34 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 129 | acceptable_shardings.append(strategy_batch) | 148 | acceptable_shardings.append(strategy_batch) |
| 130 | # Shard tensor1 | 149 | # Shard tensor1 |
| 131 | if shape1[len(shape1) - 2] % tensor1.mesh.size(0) == 0: | 150 | if shape1[len(shape1) - 2] % tensor1.mesh.size(0) == 0: |
| 132 | - strategy_tensor1 = ([Shard(len(output_shape) - 2)], [Shard(len(shape1) - 2), Replicate()]) | 151 | + strategy_tensor1 = ( |
| 152 | + [Shard(len(output_shape) - 2)], | ||
| 153 | + [Shard(len(shape1) - 2), Replicate()], | ||
| 154 | + ) | ||
| 133 | acceptable_shardings.append(strategy_tensor1) | 155 | acceptable_shardings.append(strategy_tensor1) |
| 134 | # Shard tensor2 | 156 | # Shard tensor2 |
| 135 | if shape2[len(shape2) - 1] % tensor2.mesh.size(0) == 0: | 157 | if shape2[len(shape2) - 1] % tensor2.mesh.size(0) == 0: |
| 136 | - strategy_tensor2 = ([Shard(len(output_shape) - 1)], [Replicate(), Shard(len(shape2) - 1)]) | 158 | + strategy_tensor2 = ( |
| 159 | + [Shard(len(output_shape) - 1)], | ||
| 160 | + [Replicate(), Shard(len(shape2) - 1)], | ||
| 161 | + ) | ||
| 137 | acceptable_shardings.append(strategy_tensor2) | 162 | acceptable_shardings.append(strategy_tensor2) |
| 138 | # Shard tensor1 and tensor2 | 163 | # Shard tensor1 and tensor2 |
| 139 | if shape1[len(shape1) - 1] % tensor1.mesh.size(0) == 0: | 164 | if shape1[len(shape1) - 1] % tensor1.mesh.size(0) == 0: |
| 140 | - strategy_3 = ([Partial()], [Shard(len(shape1) - 1), Shard(len(shape2) - 2)]) | 165 | + strategy_3 = ( |
| 166 | + [Partial()], | ||
| 167 | + [Shard(len(shape1) - 1), Shard(len(shape2) - 2)], | ||
| 168 | + ) | ||
| 141 | acceptable_shardings.append(strategy_3) | 169 | acceptable_shardings.append(strategy_3) |
| 142 | 170 | ||
| 143 | return acceptable_shardings | 171 | return acceptable_shardings |
| 144 | 172 | ||
| 145 | - | ||
| 146 | 173 | ||
| 147 | def custom_matmul_backward_strategy( | 174 | def custom_matmul_backward_strategy( |
| 148 | grad: DTensorSpec, | 175 | grad: DTensorSpec, |
| 149 | tensor1: DTensorSpec, | 176 | tensor1: DTensorSpec, |
| 150 | tensor2: DTensorSpec, | 177 | tensor2: DTensorSpec, |
| 151 | - mask: List[bool], | 178 | + mask: list[bool], |
| 152 | ): | 179 | ): |
| 153 | grad_dim = len(grad.shape) | 180 | grad_dim = len(grad.shape) |
| 154 | tensor1_dim = len(tensor1.shape) | 181 | tensor1_dim = len(tensor1.shape) |
| @@ -157,8 +184,8 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 157 | acceptable_shardings = [] | 184 | acceptable_shardings = [] |
| 158 | 185 | ||
| 159 | replicate_strategy = ( | 186 | replicate_strategy = ( |
| 160 | - [Replicate(), Replicate()], | 187 | + [Replicate(), Replicate()], |
| 161 | - [Replicate(), Replicate(), Replicate(), None] | 188 | + [Replicate(), Replicate(), Replicate(), None], |
| 162 | ) | 189 | ) |
| 163 | acceptable_shardings.append(replicate_strategy) | 190 | acceptable_shardings.append(replicate_strategy) |
| 164 | 191 | ||
| @@ -167,21 +194,21 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 167 | elif tensor1_dim >= 2 and (tensor2_dim == 1 or tensor2_dim == 2): | 194 | elif tensor1_dim >= 2 and (tensor2_dim == 1 or tensor2_dim == 2): |
| 168 | if tensor2.shape[0] % tensor2.mesh.size(0) == 0: | 195 | if tensor2.shape[0] % tensor2.mesh.size(0) == 0: |
| 169 | strategy_1 = ( | 196 | strategy_1 = ( |
| 170 | - [Shard(tensor1_dim - 1), Shard(0)], | 197 | + [Shard(tensor1_dim - 1), Shard(0)], |
| 171 | - [Replicate(), Shard(tensor1_dim - 1), Shard(0), None] | 198 | + [Replicate(), Shard(tensor1_dim - 1), Shard(0), None], |
| 172 | ) | 199 | ) |
| 173 | acceptable_shardings.append(strategy_1) | 200 | acceptable_shardings.append(strategy_1) |
| 174 | for i in range(tensor1_dim - 1): | 201 | for i in range(tensor1_dim - 1): |
| 175 | if tensor1.shape[i] % tensor1.mesh.size(0) == 0: | 202 | if tensor1.shape[i] % tensor1.mesh.size(0) == 0: |
| 176 | strategy_2 = ( | 203 | strategy_2 = ( |
| 177 | - [Shard(i), Partial()], | 204 | + [Shard(i), Partial()], |
| 178 | - [Shard(i), Shard(i), Replicate(), None] | 205 | + [Shard(i), Shard(i), Replicate(), None], |
| 179 | ) | 206 | ) |
| 180 | acceptable_shardings.append(strategy_2) | 207 | acceptable_shardings.append(strategy_2) |
| 181 | if tensor2_dim == 2 and tensor2.shape[1] % tensor2.mesh.size(0) == 0: | 208 | if tensor2_dim == 2 and tensor2.shape[1] % tensor2.mesh.size(0) == 0: |
| 182 | strategy_3 = ( | 209 | strategy_3 = ( |
| 183 | - [Partial(), Shard(1)], | 210 | + [Partial(), Shard(1)], |
| 184 | - [Shard(grad_dim - 1), Replicate(), Shard(1), None] | 211 | + [Shard(grad_dim - 1), Replicate(), Shard(1), None], |
| 185 | ) | 212 | ) |
| 186 | acceptable_shardings.append(strategy_3) | 213 | acceptable_shardings.append(strategy_3) |
| 187 | return acceptable_shardings | 214 | return acceptable_shardings |
| @@ -189,100 +216,135 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 189 | is_special = tensor2_dim == 2 and tensor1_dim == 1 | 216 | is_special = tensor2_dim == 2 and tensor1_dim == 1 |
| 190 | if tensor1.shape[-1] % tensor1.mesh.size(0) == 0: | 217 | if tensor1.shape[-1] % tensor1.mesh.size(0) == 0: |
| 191 | strategy_1 = ( | 218 | strategy_1 = ( |
| 192 | - [Shard(tensor1_dim if is_special else tensor1_dim - 1), Shard(tensor2_dim - 2)], | 219 | + [ |
| 193 | - [Replicate(), Shard(tensor1_dim - 1), Shard(tensor2_dim - 2), None] | 220 | + Shard(tensor1_dim if is_special else tensor1_dim - 1), |
| 221 | + Shard(tensor2_dim - 2), | ||
| 222 | + ], | ||
| 223 | + [Replicate(), Shard(tensor1_dim - 1), Shard(tensor2_dim - 2), None], | ||
| 194 | ) | 224 | ) |
| 195 | acceptable_shardings.append(strategy_1) | 225 | acceptable_shardings.append(strategy_1) |
| 196 | if tensor2.shape[-1] % tensor2.mesh.size(0) == 0: | 226 | if tensor2.shape[-1] % tensor2.mesh.size(0) == 0: |
| 197 | strategy_2 = ( | 227 | strategy_2 = ( |
| 198 | - [Partial(), Shard(tensor2_dim - 1)], | 228 | + [Partial(), Shard(tensor2_dim - 1)], |
| 199 | - [Shard(grad_dim - 1), Replicate(), Shard(tensor2_dim - 1), None] | 229 | + [Shard(grad_dim - 1), Replicate(), Shard(tensor2_dim - 1), None], |
| 200 | ) | 230 | ) |
| 201 | acceptable_shardings.append(strategy_2) | 231 | acceptable_shardings.append(strategy_2) |
| 202 | for i in range(tensor2_dim - 2): | 232 | for i in range(tensor2_dim - 2): |
| 203 | if tensor2.shape[i] % tensor2.mesh.size(0) == 0: | 233 | if tensor2.shape[i] % tensor2.mesh.size(0) == 0: |
| 204 | strategy_3 = ( | 234 | strategy_3 = ( |
| 205 | - [Partial(), Shard(i)], | 235 | + [Partial(), Shard(i)], |
| 206 | - [Shard(i), Replicate(), Shard(i)] | 236 | + [Shard(i), Replicate(), Shard(i)], |
| 207 | ) | 237 | ) |
| 208 | acceptable_shardings.append(strategy_3) | 238 | acceptable_shardings.append(strategy_3) |
| 209 | if tensor1_dim == 2 and tensor1.shape[0] % tensor1.mesh.size(0) == 0: | 239 | if tensor1_dim == 2 and tensor1.shape[0] % tensor1.mesh.size(0) == 0: |
| 210 | strategy_4 = ( | 240 | strategy_4 = ( |
| 211 | - [Shard(0), Partial()], | 241 | + [Shard(0), Partial()], |
| 212 | - [Shard(grad_dim - 2), Shard(0), Replicate(), None] | 242 | + [Shard(grad_dim - 2), Shard(0), Replicate(), None], |
| 213 | ) | 243 | ) |
| 214 | acceptable_shardings.append(strategy_4) | 244 | acceptable_shardings.append(strategy_4) |
| 215 | return acceptable_shardings | 245 | return acceptable_shardings |
| 216 | else: | 246 | else: |
| 217 | if grad.shape[-1] % grad.mesh.size(0) == 0: | 247 | if grad.shape[-1] % grad.mesh.size(0) == 0: |
| 218 | strategy_1 = ( | 248 | strategy_1 = ( |
| 219 | - [Partial(), Shard(grad_dim - 1)], | 249 | + [Partial(), Shard(grad_dim - 1)], |
| 220 | - [Shard(grad_dim - 1), Replicate(), Shard(tensor2_dim - 1), None] | 250 | + [Shard(grad_dim - 1), Replicate(), Shard(tensor2_dim - 1), None], |
| 221 | ) | 251 | ) |
| 222 | acceptable_shardings.append(strategy_1) | 252 | acceptable_shardings.append(strategy_1) |
| 223 | if grad.shape[-2] % grad.mesh.size(0) == 0: | 253 | if grad.shape[-2] % grad.mesh.size(0) == 0: |
| 224 | strategy_2 = ( | 254 | strategy_2 = ( |
| 225 | - [Shard(grad_dim - 2), Partial()], | 255 | + [Shard(grad_dim - 2), Partial()], |
| 226 | - [Shard(grad_dim - 2), Shard(tensor1_dim - 2), Replicate(), None] | 256 | + [Shard(grad_dim - 2), Shard(tensor1_dim - 2), Replicate(), None], |
| 227 | ) | 257 | ) |
| 228 | acceptable_shardings.append(strategy_2) | 258 | acceptable_shardings.append(strategy_2) |
| 229 | if tensor1.shape[-1] % tensor1.mesh.size(0) == 0: | 259 | if tensor1.shape[-1] % tensor1.mesh.size(0) == 0: |
| 230 | strategy_3 = ( | 260 | strategy_3 = ( |
| 231 | - [Shard(grad_dim - 1), Shard(grad_dim - 2)], | 261 | + [Shard(grad_dim - 1), Shard(grad_dim - 2)], |
| 232 | - [Replicate(), Shard(tensor1_dim - 1), Shard(tensor2_dim - 2), None] | 262 | + [Replicate(), Shard(tensor1_dim - 1), Shard(tensor2_dim - 2), None], |
| 233 | ) | 263 | ) |
| 234 | acceptable_shardings.append(strategy_3) | 264 | acceptable_shardings.append(strategy_3) |
| 235 | 265 | ||
| 236 | diff = abs(tensor1_dim - tensor2_dim) | 266 | diff = abs(tensor1_dim - tensor2_dim) |
| 237 | is_shape1_longer = tensor1_dim > tensor2_dim | 267 | is_shape1_longer = tensor1_dim > tensor2_dim |
| 238 | 268 | ||
| 239 | - for i in range(min(tensor1_dim, tensor2_dim) - 3, -1, -1): | 269 | + for i in range(min(tensor1_dim, tensor2_dim) - 3, -1, -1): |
| 240 | - shape1_shardable = tensor1.shape[i + diff] % tensor1.mesh.size(0) == 0 if is_shape1_longer \ | 270 | + shape1_shardable = ( |
| 271 | + tensor1.shape[i + diff] % tensor1.mesh.size(0) == 0 | ||
| 272 | + if is_shape1_longer | ||
| 241 | else tensor1.shape[i] % tensor1.mesh.size(0) == 0 | 273 | else tensor1.shape[i] % tensor1.mesh.size(0) == 0 |
| 242 | - shape2_shardable = tensor2.shape[i] % tensor2.mesh.size(0) == 0 if is_shape1_longer \ | 274 | + ) |
| 275 | + shape2_shardable = ( | ||
| 276 | + tensor2.shape[i] % tensor2.mesh.size(0) == 0 | ||
| 277 | + if is_shape1_longer | ||
| 243 | else tensor2.shape[i + diff] % tensor2.mesh.size(0) == 0 | 278 | else tensor2.shape[i + diff] % tensor2.mesh.size(0) == 0 |
| 279 | + ) | ||
| 244 | if shape1_shardable and shape2_shardable: | 280 | if shape1_shardable and shape2_shardable: |
| 245 | strategy_batch = ( | 281 | strategy_batch = ( |
| 246 | - [Shard(i + diff), Shard(i + diff)], | 282 | + ( |
| 247 | - [Shard(i + diff), Shard(i + diff), Shard(i), None] | 283 | + [Shard(i + diff), Shard(i + diff)], |
| 248 | - ) if is_shape1_longer else ( | 284 | + [Shard(i + diff), Shard(i + diff), Shard(i), None], |
| 249 | - [Shard(i + diff), Shard(i + diff)], | 285 | + ) |
| 250 | - [Shard(i + diff), Shard(i), Shard(i + diff), None] | 286 | + if is_shape1_longer |
| 287 | + else ( | ||
| 288 | + [Shard(i + diff), Shard(i + diff)], | ||
| 289 | + [Shard(i + diff), Shard(i), Shard(i + diff), None], | ||
| 290 | + ) | ||
| 251 | ) | 291 | ) |
| 252 | acceptable_shardings.append(strategy_batch) | 292 | acceptable_shardings.append(strategy_batch) |
| 253 | - | 293 | + |
| 254 | for i in range(diff - 1, -1, -1): | 294 | for i in range(diff - 1, -1, -1): |
| 255 | if is_shape1_longer and tensor1.shape[i] % tensor1.mesh.size(0) == 0: | 295 | if is_shape1_longer and tensor1.shape[i] % tensor1.mesh.size(0) == 0: |
| 256 | strategy_batch = ( | 296 | strategy_batch = ( |
| 257 | - [Shard(i), Partial()], | 297 | + [Shard(i), Partial()], |
| 258 | - [Shard(i), Shard(i), Replicate(), None] | 298 | + [Shard(i), Shard(i), Replicate(), None], |
| 259 | ) | 299 | ) |
| 260 | acceptable_shardings.append(strategy_batch) | 300 | acceptable_shardings.append(strategy_batch) |
| 261 | - elif not is_shape1_longer and tensor2.shape[i] % tensor2.mesh.size(0) == 0: | 301 | + elif ( |
| 302 | + not is_shape1_longer | ||
| 303 | + and tensor2.shape[i] % tensor2.mesh.size(0) == 0 | ||
| 304 | + ): | ||
| 262 | strategy_batch = ( | 305 | strategy_batch = ( |
| 263 | - [Partial(), Shard(i)], | 306 | + [Partial(), Shard(i)], |
| 264 | - [Shard(i), Replicate(), Shard(i), None] | 307 | + [Shard(i), Replicate(), Shard(i), None], |
| 265 | ) | 308 | ) |
| 266 | acceptable_shardings.append(strategy_batch) | 309 | acceptable_shardings.append(strategy_batch) |
| 267 | - | 310 | + |
| 268 | return acceptable_shardings | 311 | return acceptable_shardings |
| 269 | 312 | ||
| 270 | 313 | ||
| 271 | 314 | ||
| 272 | npu.npu_grouped_matmul.default, | 315 | npu.npu_grouped_matmul.default, |
| 273 | schema_info=RuntimeSchemaInfo( | 316 | schema_info=RuntimeSchemaInfo( |
| 274 | - static_kwargkey=["bias", "scale", "offset", "antiquant_scale", "antiquant_offset", "per_token_scale", | 317 | + static_kwargkey=[ |
| 275 | - "group_list", "activation_input", "activation_quant_scale", "activation_quant_offset"], | 318 | + "bias", |
| 276 | - needs_pytree=True | 319 | + "scale", |
| 277 | - ) | 320 | + "offset", |
| 321 | + "antiquant_scale", | ||
| 322 | + "antiquant_offset", | ||
| 323 | + "per_token_scale", | ||
| 324 | + "group_list", | ||
| 325 | + "activation_input", | ||
| 326 | + "activation_quant_scale", | ||
| 327 | + "activation_quant_offset", | ||
| 328 | + ], | ||
| 329 | + needs_pytree=True, | ||
| 330 | + ), | ||
| 278 | ) | 331 | ) |
| 279 | 332 | ||
| 280 | npu.npu_grouped_matmul.List, | 333 | npu.npu_grouped_matmul.List, |
| 281 | schema_info=RuntimeSchemaInfo( | 334 | schema_info=RuntimeSchemaInfo( |
| 282 | - static_kwargkey=["bias", "scale", "offset", "antiquant_scale", "antiquant_offset", "per_token_scale", | 335 | + static_kwargkey=[ |
| 283 | - "activation_input", "activation_quant_scale", "activation_quant_offset"], | 336 | + "bias", |
| 284 | - needs_pytree=True | 337 | + "scale", |
| 285 | - ) | 338 | + "offset", |
| 339 | + "antiquant_scale", | ||
| 340 | + "antiquant_offset", | ||
| 341 | + "per_token_scale", | ||
| 342 | + "activation_input", | ||
| 343 | + "activation_quant_scale", | ||
| 344 | + "activation_quant_offset", | ||
| 345 | + ], | ||
| 346 | + needs_pytree=True, | ||
| 347 | + ), | ||
| 286 | ) | 348 | ) |
| 287 | def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | 349 | def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: |
| 288 | # npu_grouped_matmul(Tensor[] x, Tensor[] weight, *, Tensor[]? bias=None, Tensor[]? scale=None, | 350 | # npu_grouped_matmul(Tensor[] x, Tensor[] weight, *, Tensor[]? bias=None, Tensor[]? scale=None, |
| @@ -296,28 +358,52 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 296 | x_num = len(x_src_strategy.children) | 358 | x_num = len(x_src_strategy.children) |
| 297 | weight_src_strategy: TupleStrategy = op_schema.args_schema[1] | 359 | weight_src_strategy: TupleStrategy = op_schema.args_schema[1] |
| 298 | weight_num = len(weight_src_strategy.children) | 360 | weight_num = len(weight_src_strategy.children) |
| 299 | - bias_src_strategy: Optional[Union[TupleStrategy, list]] = op_schema.kwargs_schema.get("bias", []) | 361 | + bias_src_strategy: TupleStrategy | list | None = op_schema.kwargs_schema.get( |
| 300 | - bias_num = len(bias_src_strategy.children) if isinstance(bias_src_strategy, TupleStrategy) else len(bias_src_strategy) | 362 | + "bias", [] |
| 301 | - group_list_num = 1 if ( | 363 | + ) |
| 302 | - op_schema.op == npu.npu_grouped_matmul.default and | 364 | + bias_num = ( |
| 303 | - op_schema.kwargs_schema.get("group_list", None) is not None | 365 | + len(bias_src_strategy.children) |
| 304 | - ) else 0 | 366 | + if isinstance(bias_src_strategy, TupleStrategy) |
| 367 | + else len(bias_src_strategy) | ||
| 368 | + ) | ||
| 369 | + group_list_num = ( | ||
| 370 | + 1 | ||
| 371 | + if ( | ||
| 372 | + op_schema.op == npu.npu_grouped_matmul.default | ||
| 373 | + and op_schema.kwargs_schema.get("group_list", None) is not None | ||
| 374 | + ) | ||
| 375 | + else 0 | ||
| 376 | + ) | ||
| 305 | split_item = op_schema.kwargs_schema.get("split_item", 0) | 377 | split_item = op_schema.kwargs_schema.get("split_item", 0) |
| 306 | - y_num = weight_num if split_item in (0, 1) else 1 # 0/1: multiple outputs, 2/3: single output | 378 | + y_num = ( |
| 379 | + weight_num if split_item in (0, 1) else 1 | ||
| 380 | + ) # 0/1: multiple outputs, 2/3: single output | ||
| 307 | 381 | ||
| 308 | strategies = [] | 382 | strategies = [] |
| 309 | 383 | ||
| 310 | all_replicate_strategy = [Replicate()] * y_num | 384 | all_replicate_strategy = [Replicate()] * y_num |
| 311 | - all_replicate_strategy.extend([Replicate()] * (len(op_schema.args_strategy) + len(op_schema.kwargs_strategy))) | 385 | + all_replicate_strategy.extend( |
| 386 | + [Replicate()] * (len(op_schema.args_strategy) + len(op_schema.kwargs_strategy)) | ||
| 387 | + ) | ||
| 312 | strategies.append(all_replicate_strategy) | 388 | strategies.append(all_replicate_strategy) |
| 313 | 389 | ||
| 314 | unsupported_arguments = [ | 390 | unsupported_arguments = [ |
| 315 | - "scale", "offset", "antiquant_scale", "antiquant_offset", "per_token_scale", # quant | 391 | + "scale", |
| 316 | - "activation_input", "activation_quant_scale", "activation_quant_offset", # reserved, unused now | 392 | + "offset", |
| 393 | + "antiquant_scale", | ||
| 394 | + "antiquant_offset", | ||
| 395 | + "per_token_scale", # quant | ||
| 396 | + "activation_input", | ||
| 397 | + "activation_quant_scale", | ||
| 398 | + "activation_quant_offset", # reserved, unused now | ||
| 317 | ] | 399 | ] |
| 318 | for key in unsupported_arguments: | 400 | for key in unsupported_arguments: |
| 319 | schema = op_schema.kwargs_schema.get(key, None) | 401 | schema = op_schema.kwargs_schema.get(key, None) |
| 320 | - if schema is not None and isinstance(schema, TupleStrategy) and len(schema.children) > 0: | 402 | + if ( |
| 403 | + schema is not None | ||
| 404 | + and isinstance(schema, TupleStrategy) | ||
| 405 | + and len(schema.children) > 0 | ||
| 406 | + ): | ||
| 321 | full_mesh_strategies = expand_to_full_mesh_op_strategy( | 407 | full_mesh_strategies = expand_to_full_mesh_op_strategy( |
| 322 | op_schema.get_mesh_from_args(), op_schema, strategies, input_index=y_num | 408 | op_schema.get_mesh_from_args(), op_schema, strategies, input_index=y_num |
| 323 | ) | 409 | ) |
| @@ -326,7 +412,9 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 326 | strategy.output_specs = [strategy.output_specs] | 412 | strategy.output_specs = [strategy.output_specs] |
| 327 | return full_mesh_strategies | 413 | return full_mesh_strategies |
| 328 | 414 | ||
| 329 | - if bias_num == 0: # if y is partial and bias exists, the bias will be added multiple times to the full tensor | 415 | + if ( |
| 416 | + bias_num == 0 | ||
| 417 | + ): # if y is partial and bias exists, the bias will be added multiple times to the full tensor | ||
| 330 | replicate_partial_strategy = [Partial()] * y_num | 418 | replicate_partial_strategy = [Partial()] * y_num |
| 331 | replicate_partial_strategy.extend([Replicate()] * x_num) | 419 | replicate_partial_strategy.extend([Replicate()] * x_num) |
| 332 | replicate_partial_strategy.extend([Partial()] * weight_num) | 420 | replicate_partial_strategy.extend([Partial()] * weight_num) |
| @@ -341,29 +429,35 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 341 | 429 | ||
| 342 | group_type = op_schema.kwargs_schema.get("group_type", None) | 430 | group_type = op_schema.kwargs_schema.get("group_type", None) |
| 343 | if group_type is not None and group_type > 0: | 431 | if group_type is not None and group_type > 0: |
| 344 | - raise NotImplementedError(f"npu_grouped_matmul does not support group_type={group_type} now.") | 432 | + raise NotImplementedError( |
| 433 | + f"npu_grouped_matmul does not support group_type={group_type} now." | ||
| 434 | + ) | ||
| 345 | 435 | ||
| 346 | - if x_num > 1 and weight_num > 1 and y_num > 1: # x_num, weight_num, y_num are equal | 436 | + if x_num > 1 and weight_num > 1 and y_num > 1: # x_num, weight_num, y_num are equal |
| 347 | pair_strategies = [] | 437 | pair_strategies = [] |
| 348 | # x: 2-6D, weight: 2D, weight: 1D (equals to weight.shape[1]) | 438 | # x: 2-6D, weight: 2D, weight: 1D (equals to weight.shape[1]) |
| 349 | # shard x | 439 | # shard x |
| 350 | x_ndim = x_src_strategy.children[0].ndim | 440 | x_ndim = x_src_strategy.children[0].ndim |
| 351 | for i in range(x_ndim - 1): | 441 | for i in range(x_ndim - 1): |
| 352 | - pair_strategies.append([Shard(i), Shard(i), Replicate(), Replicate()]) # y, x, weight, bias | 442 | + pair_strategies.append( |
| 443 | + [Shard(i), Shard(i), Replicate(), Replicate()] | ||
| 444 | + ) # y, x, weight, bias | ||
| 353 | # shard weight | 445 | # shard weight |
| 354 | pair_strategies.append([Shard(x_ndim - 1), Replicate(), Shard(1), Shard(0)]) | 446 | pair_strategies.append([Shard(x_ndim - 1), Replicate(), Shard(1), Shard(0)]) |
| 355 | # shard contracting dim | 447 | # shard contracting dim |
| 356 | if bias_num == 0: | 448 | if bias_num == 0: |
| 357 | pair_strategies.append([Partial(), Shard(x_ndim - 1), Shard(0), None]) | 449 | pair_strategies.append([Partial(), Shard(x_ndim - 1), Shard(0), None]) |
| 358 | # suppose that all pairs have the same shape and apply the same strategy | 450 | # suppose that all pairs have the same shape and apply the same strategy |
| 359 | - for (y_spec, x_spec, weight_spec, bias_spec) in pair_strategies: | 451 | + for y_spec, x_spec, weight_spec, bias_spec in pair_strategies: |
| 360 | strategy = [y_spec] * y_num | 452 | strategy = [y_spec] * y_num |
| 361 | strategy.extend([x_spec] * x_num) | 453 | strategy.extend([x_spec] * x_num) |
| 362 | strategy.extend([weight_spec] * weight_num) | 454 | strategy.extend([weight_spec] * weight_num) |
| 363 | strategy.extend([bias_spec] * bias_num) | 455 | strategy.extend([bias_spec] * bias_num) |
| 364 | strategy.extend([Replicate()] * group_list_num) | 456 | strategy.extend([Replicate()] * group_list_num) |
| 365 | strategies.append(strategy) | 457 | strategies.append(strategy) |
| 366 | - elif x_num == 1 and weight_num == 1 and y_num == 1: # npu_grouped_matmul.default only | 458 | + elif ( |
| 459 | + x_num == 1 and weight_num == 1 and y_num == 1 | ||
| 460 | + ): # npu_grouped_matmul.default only | ||
| 367 | # x: 2D, weight: 3D, bias: 2D, y: 2D, for each pair, define shape x: (m, k), weight: (k, n) | 461 | # x: 2D, weight: 3D, bias: 2D, y: 2D, for each pair, define shape x: (m, k), weight: (k, n) |
| 368 | if bias_num == 0: | 462 | if bias_num == 0: |
| 369 | k_shard_strategy = [Partial(), Shard(1), Shard(1)] | 463 | k_shard_strategy = [Partial(), Shard(1), Shard(1)] |
| @@ -373,7 +467,7 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 373 | n_shard_strategy.extend([Shard(1)] * bias_num) | 467 | n_shard_strategy.extend([Shard(1)] * bias_num) |
| 374 | n_shard_strategy.extend([Replicate()] * group_list_num) | 468 | n_shard_strategy.extend([Replicate()] * group_list_num) |
| 375 | strategies.append(n_shard_strategy) | 469 | strategies.append(n_shard_strategy) |
| 376 | - elif weight_num > 1: # x1wNy1, xNwNy1, x1wNyN | 470 | + elif weight_num > 1: # x1wNy1, xNwNy1, x1wNyN |
| 377 | # x: 2D, weight: 2D, bias: 1D, y: 2D | 471 | # x: 2D, weight: 2D, bias: 1D, y: 2D |
| 378 | if bias_num == 0: | 472 | if bias_num == 0: |
| 379 | k_shard_strategy = [Partial()] * y_num | 473 | k_shard_strategy = [Partial()] * y_num |
| @@ -388,8 +482,9 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 388 | n_shard_strategy.extend([Replicate()] * group_list_num) | 482 | n_shard_strategy.extend([Replicate()] * group_list_num) |
| 389 | strategies.append(n_shard_strategy) | 483 | strategies.append(n_shard_strategy) |
| 390 | 484 | ||
| 391 | - full_mesh_strategies = expand_to_full_mesh_op_strategy(op_schema.get_mesh_from_args(), op_schema, strategies, | 485 | + full_mesh_strategies = expand_to_full_mesh_op_strategy( |
| 392 | - input_index=y_num) | 486 | + op_schema.get_mesh_from_args(), op_schema, strategies, input_index=y_num |
| 487 | + ) | ||
| 393 | # output meta of npu_grouped_matmul is list, need convert output_spec here | 488 | # output meta of npu_grouped_matmul is list, need convert output_spec here |
| 394 | if y_num == 1: | 489 | if y_num == 1: |
| 395 | for strategy in full_mesh_strategies.strategies: | 490 | for strategy in full_mesh_strategies.strategies: |
| @@ -398,42 +493,51 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 398 | 493 | ||
| 399 | 494 | ||
| 400 | def _infer_npu_grouped_matmul_kwargs( | 495 | def _infer_npu_grouped_matmul_kwargs( |
| 401 | - op_schema: OpSchema, | 496 | + op_schema: OpSchema, output_sharding: OutputSharding |
| 402 | - output_sharding: OutputSharding | 497 | +) -> dict[str, DTensorSpec]: |
| 403 | -) -> Dict[str, DTensorSpec]: | ||
| 404 | output_spec = output_sharding.output_spec[0] | 498 | output_spec = output_sharding.output_spec[0] |
| 405 | kwargs_spec = {} | 499 | kwargs_spec = {} |
| 406 | for key, spec in op_schema.kwargs_schema.items(): | 500 | for key, spec in op_schema.kwargs_schema.items(): |
| 407 | - is_tensor_or_tenor_list_like = ( | 501 | + is_tensor_or_tenor_list_like = isinstance(spec, DTensorSpec) or ( |
| 408 | - isinstance(spec, DTensorSpec) or | 502 | + isinstance(spec, (list, tuple)) |
| 409 | - (isinstance(spec, (list, tuple)) and len(spec) > 0 and isinstance(spec[0], DTensorSpec)) | 503 | + and len(spec) > 0 |
| 504 | + and isinstance(spec[0], DTensorSpec) | ||
| 410 | ) | 505 | ) |
| 411 | if not is_tensor_or_tenor_list_like: | 506 | if not is_tensor_or_tenor_list_like: |
| 412 | kwargs_spec[key] = spec | 507 | kwargs_spec[key] = spec |
| 413 | continue | 508 | continue |
| 414 | 509 | ||
| 415 | - if key == 'group_list': # tensor | 510 | + if key == "group_list": # tensor |
| 416 | target_placement = [Replicate() for _ in output_spec.placements] | 511 | target_placement = [Replicate() for _ in output_spec.placements] |
| 417 | - kwargs_spec[key] = DTensorSpec(mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta) | 512 | + kwargs_spec[key] = DTensorSpec( |
| 513 | + mesh=spec.mesh, | ||
| 514 | + placements=target_placement, | ||
| 515 | + tensor_meta=spec.tensor_meta, | ||
| 516 | + ) | ||
| 418 | continue | 517 | continue |
| 419 | 518 | ||
| 420 | # tensor list | 519 | # tensor list |
| 421 | - if key == 'bias': | 520 | + if key == "bias": |
| 422 | target_placement = [ | 521 | target_placement = [ |
| 423 | Shard(0) if placement == Shard(output_spec.ndim - 1) else Replicate() | 522 | Shard(0) if placement == Shard(output_spec.ndim - 1) else Replicate() |
| 424 | for placement in output_spec.placements | 523 | for placement in output_spec.placements |
| 425 | ] | 524 | ] |
| 426 | - else: # unsupported sharding keys | 525 | + else: # unsupported sharding keys |
| 427 | target_placement = [Replicate() for _ in output_spec.placements] | 526 | target_placement = [Replicate() for _ in output_spec.placements] |
| 428 | - kwargs_spec[key] = [DTensorSpec(mesh=e.mesh, placements=target_placement, tensor_meta=e.tensor_meta) for e in spec] | 527 | + kwargs_spec[key] = [ |
| 528 | + DTensorSpec( | ||
| 529 | + mesh=e.mesh, placements=target_placement, tensor_meta=e.tensor_meta | ||
| 530 | + ) | ||
| 531 | + for e in spec | ||
| 532 | + ] | ||
| 429 | 533 | ||
| 430 | return kwargs_spec | 534 | return kwargs_spec |
| 431 | 535 | ||
| 432 | 536 | ||
| 433 | def _npu_grouped_matmul_handler( | 537 | def _npu_grouped_matmul_handler( |
| 434 | - op_call: torch._ops.OpOverload, | 538 | + op_call: torch._ops.OpOverload, |
| 435 | - args: Tuple[object, ...], | 539 | + args: tuple[object, ...], |
| 436 | - kwargs: Dict[str, object], | 540 | + kwargs: dict[str, object], |
| 437 | ) -> object: | 541 | ) -> object: |
| 438 | # extract local tensor and sharding infos to a OpInfo | 542 | # extract local tensor and sharding infos to a OpInfo |
| 439 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) | 543 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) |
| @@ -449,7 +553,9 @@ def _npu_grouped_matmul_handler( | |||
| 449 | if participating: | 553 | if participating: |
| 450 | # computation that happens in the current rank of the mesh, normal case | 554 | # computation that happens in the current rank of the mesh, normal case |
| 451 | local_args = get_redistributed_local_args(op_info, output_sharding) | 555 | local_args = get_redistributed_local_args(op_info, output_sharding) |
| 452 | - local_kwargs = get_redistributed_local_kwargs(_infer_npu_grouped_matmul_kwargs, op_info, output_sharding) | 556 | + local_kwargs = get_redistributed_local_kwargs( |
| 557 | + _infer_npu_grouped_matmul_kwargs, op_info, output_sharding | ||
| 558 | + ) | ||
| 453 | local_results = op_call(*local_args, **local_kwargs) | 559 | local_results = op_call(*local_args, **local_kwargs) |
| 454 | else: | 560 | else: |
| 455 | # For a non-participating device (happens on rank that does not belong to the device mesh), | 561 | # For a non-participating device (happens on rank that does not belong to the device mesh), |
| @@ -460,14 +566,28 @@ def _npu_grouped_matmul_handler( | |||
| 460 | 566 | ||
| 461 | 567 | ||
| 462 | 568 | ||
| 463 | -def npu_all_gather_base_mm_strategy(x1, x2, hcom, world_size, bias=None, x1_scale=None, x2_scale=None, gather_index=0, | 569 | +def npu_all_gather_base_mm_strategy( |
| 464 | - gather_output=True, comm_turn=0, output_dtype=None, comm_mode=None): | 570 | + x1, |
| 571 | + x2, | ||
| 572 | + hcom, | ||
| 573 | + world_size, | ||
| 574 | + bias=None, | ||
| 575 | + x1_scale=None, | ||
| 576 | + x2_scale=None, | ||
| 577 | + gather_index=0, | ||
| 578 | + gather_output=True, | ||
| 579 | + comm_turn=0, | ||
| 580 | + output_dtype=None, | ||
| 581 | + comm_mode=None, | ||
| 582 | +): | ||
| 465 | # npu_all_gather_base_mm(Tensor input, Tensor x2, str hcom, int world_size, *, Tensor? bias=None, | 583 | # npu_all_gather_base_mm(Tensor input, Tensor x2, str hcom, int world_size, *, Tensor? bias=None, |
| 466 | # Tensor? x1_scale=None, Tensor? x2_scale=None, int gather_index=0, bool gather_output=True, | 584 | # Tensor? x1_scale=None, Tensor? x2_scale=None, int gather_index=0, bool gather_output=True, |
| 467 | # int comm_turn=0, ScalarType? output_dtype=None, str? comm_mode=None) -> (Tensor, Tensor) | 585 | # int comm_turn=0, ScalarType? output_dtype=None, str? comm_mode=None) -> (Tensor, Tensor) |
| 468 | # op only support gather_index=0(i.e. allgather x1) now | 586 | # op only support gather_index=0(i.e. allgather x1) now |
| 469 | if gather_index != 0: | 587 | if gather_index != 0: |
| 470 | - raise NotImplementedError(f"npu_all_gather_base_mm only support gather_index=0 now, but got {gather_index}.") | 588 | + raise NotImplementedError( |
| 589 | + f"npu_all_gather_base_mm only support gather_index=0 now, but got {gather_index}." | ||
| 590 | + ) | ||
| 471 | 591 | ||
| 472 | # formula: output = allgather(x1)@x2 + bias | 592 | # formula: output = allgather(x1)@x2 + bias |
| 473 | # for all gather, x1: S(0) -> R | 593 | # for all gather, x1: S(0) -> R |
| @@ -478,37 +598,55 @@ def npu_all_gather_base_mm_strategy(x1, x2, hcom, world_size, bias=None, x1_scal | |||
| 478 | strategies = [] | 598 | strategies = [] |
| 479 | sharding_strategy_S0R = ( | 599 | sharding_strategy_S0R = ( |
| 480 | [ | 600 | [ |
| 481 | - Replicate(), # output | 601 | + Replicate(), # output |
| 482 | - Replicate() # gather_out | 602 | + Replicate(), # gather_out |
| 483 | ], | 603 | ], |
| 484 | [ | 604 | [ |
| 485 | - Shard(0), # x1 | 605 | + Shard(0), # x1 |
| 486 | - Replicate(), # x2 | 606 | + Replicate(), # x2 |
| 487 | - None, # hcom | 607 | + None, # hcom |
| 488 | - None, # world_size | 608 | + None, # world_size |
| 489 | - None if bias is None else Replicate(), # bias, global shape(n * world_size,) | 609 | + None |
| 490 | - None if x1_scale is None else Shard(0), # x1_scale follow x1, global shape(m * world_size, 1) | 610 | + if bias is None |
| 491 | - None if x2_scale is None else Replicate(), # x2_scale follow x2, global shape(1, n * world_size) | 611 | + else Replicate(), # bias, global shape(n * world_size,) |
| 492 | - None, None, None, None, None # gather_index, gather_output, comm_turn, output_dtype, comm_mode | 612 | + None |
| 493 | - ] | 613 | + if x1_scale is None |
| 614 | + else Shard(0), # x1_scale follow x1, global shape(m * world_size, 1) | ||
| 615 | + None | ||
| 616 | + if x2_scale is None | ||
| 617 | + else Replicate(), # x2_scale follow x2, global shape(1, n * world_size) | ||
| 618 | + None, | ||
| 619 | + None, | ||
| 620 | + None, | ||
| 621 | + None, | ||
| 622 | + None, # gather_index, gather_output, comm_turn, output_dtype, comm_mode | ||
| 623 | + ], | ||
| 494 | ) | 624 | ) |
| 495 | strategies.append(sharding_strategy_S0R) | 625 | strategies.append(sharding_strategy_S0R) |
| 496 | 626 | ||
| 497 | sharding_strategy_S0S1 = ( | 627 | sharding_strategy_S0S1 = ( |
| 498 | [ | 628 | [ |
| 499 | - Shard(1), # output | 629 | + Shard(1), # output |
| 500 | - Replicate() # gather_out | 630 | + Replicate(), # gather_out |
| 501 | ], | 631 | ], |
| 502 | [ | 632 | [ |
| 503 | - Shard(0), # x1 | 633 | + Shard(0), # x1 |
| 504 | - Shard(1), # x2 | 634 | + Shard(1), # x2 |
| 505 | - None, # hcom | 635 | + None, # hcom |
| 506 | - None, # world_size | 636 | + None, # world_size |
| 507 | - None if bias is None else Shard(0), # bias, global shape(n * world_size,) | 637 | + None if bias is None else Shard(0), # bias, global shape(n * world_size,) |
| 508 | - None if x1_scale is None else Shard(0), # x1_scale follow x1, global shape(m * world_size, 1) | 638 | + None |
| 509 | - None if x2_scale is None else Shard(1), # x2_scale follow x2, global shape(1, n * world_size) | 639 | + if x1_scale is None |
| 510 | - None, None, None, None, None # gather_index, gather_output, comm_turn, output_dtype, comm_mode | 640 | + else Shard(0), # x1_scale follow x1, global shape(m * world_size, 1) |
| 511 | - ] | 641 | + None |
| 642 | + if x2_scale is None | ||
| 643 | + else Shard(1), # x2_scale follow x2, global shape(1, n * world_size) | ||
| 644 | + None, | ||
| 645 | + None, | ||
| 646 | + None, | ||
| 647 | + None, | ||
| 648 | + None, # gather_index, gather_output, comm_turn, output_dtype, comm_mode | ||
| 649 | + ], | ||
| 512 | ) | 650 | ) |
| 513 | strategies.append(sharding_strategy_S0S1) | 651 | strategies.append(sharding_strategy_S0S1) |
| 514 | 652 | ||
| @@ -516,9 +654,8 @@ def npu_all_gather_base_mm_strategy(x1, x2, hcom, world_size, bias=None, x1_scal | |||
| 516 | 654 | ||
| 517 | 655 | ||
| 518 | def _infer_npu_all_gather_base_mm_kwargs( | 656 | def _infer_npu_all_gather_base_mm_kwargs( |
| 519 | - op_schema: OpSchema, | 657 | + op_schema: OpSchema, output_sharding: OutputSharding |
| 520 | - output_sharding: OutputSharding | 658 | +) -> dict[str, DTensorSpec]: |
| 521 | -) -> Dict[str, DTensorSpec]: | ||
| 522 | output_spec = output_sharding.output_spec[0] | 659 | output_spec = output_sharding.output_spec[0] |
| 523 | kwargs_spec = {} | 660 | kwargs_spec = {} |
| 524 | for key, spec in op_schema.kwargs_schema.items(): | 661 | for key, spec in op_schema.kwargs_schema.items(): |
| @@ -529,33 +666,48 @@ def _infer_npu_all_gather_base_mm_kwargs( | |||
| 529 | target_placement = [] | 666 | target_placement = [] |
| 530 | for placement in output_spec.placements: | 667 | for placement in output_spec.placements: |
| 531 | if placement == Replicate(): | 668 | if placement == Replicate(): |
| 532 | - if key == 'x1_scale': | 669 | + if key == "x1_scale": |
| 533 | target_placement.append(Shard(0)) | 670 | target_placement.append(Shard(0)) |
| 534 | - else: # bias, x2_scale | 671 | + else: # bias, x2_scale |
| 535 | target_placement.append(Replicate()) | 672 | target_placement.append(Replicate()) |
| 536 | elif placement == Shard(1): | 673 | elif placement == Shard(1): |
| 537 | - if key == 'x2_scale': | 674 | + if key == "x2_scale": |
| 538 | target_placement.append(Shard(1)) | 675 | target_placement.append(Shard(1)) |
| 539 | - else: # bias, x1_scale | 676 | + else: # bias, x1_scale |
| 540 | target_placement.append(Shard(0)) | 677 | target_placement.append(Shard(0)) |
| 541 | else: | 678 | else: |
| 542 | raise ValueError( | 679 | raise ValueError( |
| 543 | f"Unexpected output placement {placement} for npu_all_gather_base_mm." | 680 | f"Unexpected output placement {placement} for npu_all_gather_base_mm." |
| 544 | ) | 681 | ) |
| 545 | - kwargs_spec[key] = DTensorSpec(mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta) | 682 | + kwargs_spec[key] = DTensorSpec( |
| 683 | + mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta | ||
| 684 | + ) | ||
| 546 | 685 | ||
| 547 | return kwargs_spec | 686 | return kwargs_spec |
| 548 | 687 | ||
| 549 | 688 | ||
| 550 | 689 | ||
| 551 | -def npu_mm_reduce_scatter_base_strategy(x1, x2, hcom, world_size, reduce_op='sum', bias=None, x1_scale=None, | 690 | +def npu_mm_reduce_scatter_base_strategy( |
| 552 | - x2_scale=None, comm_turn=0, output_dtype=None, comm_mode=None): | 691 | + x1, |
| 692 | + x2, | ||
| 693 | + hcom, | ||
| 694 | + world_size, | ||
| 695 | + reduce_op="sum", | ||
| 696 | + bias=None, | ||
| 697 | + x1_scale=None, | ||
| 698 | + x2_scale=None, | ||
| 699 | + comm_turn=0, | ||
| 700 | + output_dtype=None, | ||
| 701 | + comm_mode=None, | ||
| 702 | +): | ||
| 553 | # npu_mm_reduce_scatter_base(Tensor self, Tensor x2, str hcom, int world_size, *, str reduce_op='sum', | 703 | # npu_mm_reduce_scatter_base(Tensor self, Tensor x2, str hcom, int world_size, *, str reduce_op='sum', |
| 554 | # Tensor? bias=None, Tensor? x1_scale=None, Tensor? x2_scale=None, int comm_turn=0, | 704 | # Tensor? bias=None, Tensor? x1_scale=None, Tensor? x2_scale=None, int comm_turn=0, |
| 555 | # ScalarType? output_dtype=None, str? comm_mode=None) -> Tensor | 705 | # ScalarType? output_dtype=None, str? comm_mode=None) -> Tensor |
| 556 | # op only support reduce_op='sum' now | 706 | # op only support reduce_op='sum' now |
| 557 | - if reduce_op != 'sum': | 707 | + if reduce_op != "sum": |
| 558 | - raise NotImplementedError(f"npu_mm_reduce_scatter_base only support reduce_op='sum' now, but got {reduce_op}.") | 708 | + raise NotImplementedError( |
| 709 | + f"npu_mm_reduce_scatter_base only support reduce_op='sum' now, but got {reduce_op}." | ||
| 710 | + ) | ||
| 559 | 711 | ||
| 560 | # formula: output = reducescatter(x1@x2 + bias) | 712 | # formula: output = reducescatter(x1@x2 + bias) |
| 561 | # for reduce_scatter, local_output: P -> S(0) | 713 | # for reduce_scatter, local_output: P -> S(0) |
| @@ -567,16 +719,22 @@ def npu_mm_reduce_scatter_base_strategy(x1, x2, hcom, world_size, reduce_op='sum | |||
| 567 | Shard(0) # output | 719 | Shard(0) # output |
| 568 | ], | 720 | ], |
| 569 | [ | 721 | [ |
| 570 | - Shard(1), # x1 | 722 | + Shard(1), # x1 |
| 571 | - Shard(0), # x2 | 723 | + Shard(0), # x2 |
| 572 | - None, # hcom | 724 | + None, # hcom |
| 573 | - None, # world_size | 725 | + None, # world_size |
| 574 | - None, # reduce_op | 726 | + None, # reduce_op |
| 575 | - None if bias is None else Shard(0), # bias, global shape(n * world_size,) | 727 | + None if bias is None else Shard(0), # bias, global shape(n * world_size,) |
| 576 | - None if x1_scale is None else Shard(1), # x1_scale follow x1, global shape(m, world_size) | 728 | + None |
| 577 | - None if x2_scale is None else Shard(0), # x2_scale follow x2, global shape(world_size, n) | 729 | + if x1_scale is None |
| 578 | - None, None, None # comm_turn, output_dtype, comm_mode | 730 | + else Shard(1), # x1_scale follow x1, global shape(m, world_size) |
| 579 | - ] | 731 | + None |
| 732 | + if x2_scale is None | ||
| 733 | + else Shard(0), # x2_scale follow x2, global shape(world_size, n) | ||
| 734 | + None, | ||
| 735 | + None, | ||
| 736 | + None, # comm_turn, output_dtype, comm_mode | ||
| 737 | + ], | ||
| 580 | ) | 738 | ) |
| 581 | strategies.append(sharding_strategy_S1S0) | 739 | strategies.append(sharding_strategy_S1S0) |
| 582 | 740 | ||
| @@ -584,9 +742,8 @@ def npu_mm_reduce_scatter_base_strategy(x1, x2, hcom, world_size, reduce_op='sum | |||
| 584 | 742 | ||
| 585 | 743 | ||
| 586 | def _infer_npu_mm_reduce_scatter_base_kwargs( | 744 | def _infer_npu_mm_reduce_scatter_base_kwargs( |
| 587 | - op_schema: OpSchema, | 745 | + op_schema: OpSchema, output_sharding: OutputSharding |
| 588 | - output_sharding: OutputSharding | 746 | +) -> dict[str, DTensorSpec]: |
| 589 | -) -> Dict[str, DTensorSpec]: | ||
| 590 | output_spec = output_sharding.output_spec | 747 | output_spec = output_sharding.output_spec |
| 591 | kwargs_spec = {} | 748 | kwargs_spec = {} |
| 592 | for key, spec in op_schema.kwargs_schema.items(): | 749 | for key, spec in op_schema.kwargs_schema.items(): |
| @@ -597,23 +754,25 @@ def _infer_npu_mm_reduce_scatter_base_kwargs( | |||
| 597 | target_placement = [] | 754 | target_placement = [] |
| 598 | for placement in output_spec.placements: | 755 | for placement in output_spec.placements: |
| 599 | if placement == Shard(0): | 756 | if placement == Shard(0): |
| 600 | - if key == 'x1_scale': | 757 | + if key == "x1_scale": |
| 601 | target_placement.append(Shard(1)) | 758 | target_placement.append(Shard(1)) |
| 602 | - else: # bias, x2_scale | 759 | + else: # bias, x2_scale |
| 603 | target_placement.append(Shard(0)) | 760 | target_placement.append(Shard(0)) |
| 604 | else: | 761 | else: |
| 605 | raise ValueError( | 762 | raise ValueError( |
| 606 | f"Unexpected output placement {placement} for npu_mm_reduce_scatter_base." | 763 | f"Unexpected output placement {placement} for npu_mm_reduce_scatter_base." |
| 607 | ) | 764 | ) |
| 608 | - kwargs_spec[key] = DTensorSpec(mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta) | 765 | + kwargs_spec[key] = DTensorSpec( |
| 766 | + mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta | ||
| 767 | + ) | ||
| 609 | 768 | ||
| 610 | return kwargs_spec | 769 | return kwargs_spec |
| 611 | 770 | ||
| 612 | 771 | ||
| 613 | def npu_comm_mm_fusion_handler( | 772 | def npu_comm_mm_fusion_handler( |
| 614 | - op_call: torch._ops.OpOverload, | 773 | + op_call: torch._ops.OpOverload, |
| 615 | - args: Tuple[object, ...], | 774 | + args: tuple[object, ...], |
| 616 | - kwargs: Dict[str, object], | 775 | + kwargs: dict[str, object], |
| 617 | ) -> object: | 776 | ) -> object: |
| 618 | # extract local tensor and sharding infos to a OpInfo | 777 | # extract local tensor and sharding infos to a OpInfo |
| 619 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) | 778 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) |
| @@ -634,11 +793,15 @@ def npu_comm_mm_fusion_handler( | |||
| 634 | new_shape[dim] = new_shape[dim] // world_size | 793 | new_shape[dim] = new_shape[dim] // world_size |
| 635 | elif op_call == npu.npu_mm_reduce_scatter_base.default: | 794 | elif op_call == npu.npu_mm_reduce_scatter_base.default: |
| 636 | new_shape[dim] = new_shape[dim] * world_size | 795 | new_shape[dim] = new_shape[dim] * world_size |
| 637 | - return TensorMeta(shape=torch.Size(new_shape), stride=tensor_meta.stride, dtype=tensor_meta.dtype) | 796 | + return TensorMeta( |
| 797 | + shape=torch.Size(new_shape), | ||
| 798 | + stride=tensor_meta.stride, | ||
| 799 | + dtype=tensor_meta.dtype, | ||
| 800 | + ) | ||
| 638 | 801 | ||
| 639 | if op_call == npu.npu_all_gather_base_mm.default: | 802 | if op_call == npu.npu_all_gather_base_mm.default: |
| 640 | world_size = args[3] | 803 | world_size = args[3] |
| 641 | - for spec in output_sharding.output_spec: # output, gather_out | 804 | + for spec in output_sharding.output_spec: # output, gather_out |
| 642 | spec.tensor_meta = get_output_meta(spec.tensor_meta, 0, world_size) | 805 | spec.tensor_meta = get_output_meta(spec.tensor_meta, 0, world_size) |
| 643 | elif op_call == npu.npu_mm_reduce_scatter_base.default: | 806 | elif op_call == npu.npu_mm_reduce_scatter_base.default: |
| 644 | world_size = args[3] | 807 | world_size = args[3] |
| @@ -669,9 +832,7 @@ def npu_comm_mm_fusion_handler( | |||
| 669 | return DTensor._op_dispatcher.wrap(local_results, output_sharding.output_spec) | 832 | return DTensor._op_dispatcher.wrap(local_results, output_sharding.output_spec) |
| 670 | 833 | ||
| 671 | 834 | ||
| 672 | -@register_op_strategy( | 835 | +@register_op_strategy([npu.npu_apply_adam_w.default, npu.npu_apply_adam_w.out]) |
| 673 | - [npu.npu_apply_adam_w.default, npu.npu_apply_adam_w.out] | ||
| 674 | -) | ||
| 675 | def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: | 836 | def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: |
| 676 | # npu_apply_adam_w( | 837 | # npu_apply_adam_w( |
| 677 | # Scalar beta1_power, Scalar beta2_power, Scalar lr, Scalar weight_decay, Scalar beta1, Scalar beta2, | 838 | # Scalar beta1_power, Scalar beta2_power, Scalar lr, Scalar weight_decay, Scalar beta1, Scalar beta2, |
| @@ -680,8 +841,10 @@ def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 680 | grad_arg_index = 7 | 841 | grad_arg_index = 7 |
| 681 | max_gard_norm_arg_index = 8 | 842 | max_gard_norm_arg_index = 8 |
| 682 | grad_strategy: OpStrategy = op_schema.args_schema[grad_arg_index] | 843 | grad_strategy: OpStrategy = op_schema.args_schema[grad_arg_index] |
| 683 | - if "out" in op_schema.kwargs_schema.keys(): | 844 | + if "out" in op_schema.kwargs_schema: |
| 684 | - grad_spec: DTensorSpec = op_schema.kwargs_schema["out"].children[0].strategies[0].output_spec | 845 | + grad_spec: DTensorSpec = ( |
| 846 | + op_schema.kwargs_schema["out"].children[0].strategies[0].output_spec | ||
| 847 | + ) | ||
| 685 | else: | 848 | else: |
| 686 | grad_spec: DTensorSpec = grad_strategy.strategies[0].output_spec | 849 | grad_spec: DTensorSpec = grad_strategy.strategies[0].output_spec |
| 687 | input_target_specs = [] | 850 | input_target_specs = [] |
| @@ -704,20 +867,20 @@ def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 704 | 867 | ||
| 705 | output_spec = [] | 868 | output_spec = [] |
| 706 | for k, values in op_schema.kwargs_schema.items(): | 869 | for k, values in op_schema.kwargs_schema.items(): |
| 707 | - if k == 'out': | 870 | + if k == "out": |
| 708 | for v in values.children: | 871 | for v in values.children: |
| 709 | output_spec.append(v.strategies[0].output_spec) | 872 | output_spec.append(v.strategies[0].output_spec) |
| 710 | - output_strategy = OpStrategy([ | 873 | + output_strategy = OpStrategy( |
| 711 | - OpSpec(output_specs=tuple(output_spec), input_specs=input_target_specs) | 874 | + [OpSpec(output_specs=tuple(output_spec), input_specs=input_target_specs)] |
| 712 | - ]) | 875 | + ) |
| 713 | 876 | ||
| 714 | return output_strategy | 877 | return output_strategy |
| 715 | 878 | ||
| 716 | 879 | ||
| 717 | def _npu_apply_adam_w_handler( | 880 | def _npu_apply_adam_w_handler( |
| 718 | - op_call: torch._ops.OpOverload, | 881 | + op_call: torch._ops.OpOverload, |
| 719 | - args: Tuple[object, ...], | 882 | + args: tuple[object, ...], |
| 720 | - kwargs: Dict[str, object], | 883 | + kwargs: dict[str, object], |
| 721 | ) -> object: | 884 | ) -> object: |
| 722 | # extract local tensor and sharding infos to a OpInfo | 885 | # extract local tensor and sharding infos to a OpInfo |
| 723 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) | 886 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) |
| @@ -739,12 +902,12 @@ def _npu_apply_adam_w_handler( | |||
| 739 | output_sharding.use_val_from_redistribute_schema, | 902 | output_sharding.use_val_from_redistribute_schema, |
| 740 | ) | 903 | ) |
| 741 | local_args = ( | 904 | local_args = ( |
| 742 | - pytree.tree_unflatten( | 905 | + pytree.tree_unflatten( |
| 743 | - cast(list[object], op_info.local_args), op_info.args_tree_spec | 906 | + cast(list[object], op_info.local_args), op_info.args_tree_spec |
| 744 | - ) | ||
| 745 | - if op_info.args_tree_spec | ||
| 746 | - else op_info.local_args | ||
| 747 | ) | 907 | ) |
| 908 | + if op_info.args_tree_spec | ||
| 909 | + else op_info.local_args | ||
| 910 | + ) | ||
| 748 | 911 | ||
| 749 | local_results = torch_npu.npu_apply_adam_w(*local_args, **op_info.local_kwargs) | 912 | local_results = torch_npu.npu_apply_adam_w(*local_args, **op_info.local_kwargs) |
| 750 | 913 | ||
| @@ -757,7 +920,7 @@ def _npu_apply_adam_w_handler( | |||
| 757 | out_dts = [] | 920 | out_dts = [] |
| 758 | spec_idx = 0 | 921 | spec_idx = 0 |
| 759 | for argument in op_call._schema.arguments: | 922 | for argument in op_call._schema.arguments: |
| 760 | - if argument.name == 'out': | 923 | + if argument.name == "out": |
| 761 | for value in kwargs[argument.name]: | 924 | for value in kwargs[argument.name]: |
| 762 | out_dt = cast(DTensor, value) | 925 | out_dt = cast(DTensor, value) |
| 763 | out_dt._spec = cast(DTensorSpec, output_specs[spec_idx]) | 926 | out_dt._spec = cast(DTensorSpec, output_specs[spec_idx]) |
| @@ -784,16 +947,17 @@ def custom_dropout_strategy(op_schema: OpSchema): | |||
| 784 | output_target_specs = [] | 947 | output_target_specs = [] |
| 785 | output_target_specs.append(input_spec) | 948 | output_target_specs.append(input_spec) |
| 786 | output_target_specs.append( | 949 | output_target_specs.append( |
| 787 | - DTensorSpec( | 950 | + DTensorSpec(mesh=input_spec.mesh, placements=[Shard(0)]) |
| 788 | - mesh=input_spec.mesh, | ||
| 789 | - placements=[Shard(0)] | ||
| 790 | - ) | ||
| 791 | ) | 951 | ) |
| 792 | input_target_specs = [] | 952 | input_target_specs = [] |
| 793 | input_target_specs.append(input_spec) | 953 | input_target_specs.append(input_spec) |
| 794 | - output_strategy = OpStrategy([ | 954 | + output_strategy = OpStrategy( |
| 795 | - OpSpec(output_specs=output_target_specs, input_specs=input_target_specs) | 955 | + [ |
| 796 | - ]) | 956 | + OpSpec( |
| 957 | + output_specs=output_target_specs, input_specs=input_target_specs | ||
| 958 | + ) | ||
| 959 | + ] | ||
| 960 | + ) | ||
| 797 | return output_strategy | 961 | return output_strategy |
| 798 | 962 | ||
| 799 | replicate_strategy = [Replicate(), Replicate(), Replicate()] | 963 | replicate_strategy = [Replicate(), Replicate(), Replicate()] |
| @@ -815,19 +979,30 @@ def custom_dropout_backward_strategy(op_schema: OpSchema): | |||
| 815 | if isinstance(spec, OpStrategy): | 979 | if isinstance(spec, OpStrategy): |
| 816 | input_target_specs.append(spec.strategies[0].output_spec) | 980 | input_target_specs.append(spec.strategies[0].output_spec) |
| 817 | 981 | ||
| 818 | - output_strategy = OpStrategy([ | 982 | + output_strategy = OpStrategy( |
| 819 | - OpSpec(output_specs=op_schema.args_schema[0].strategies[0].output_spec, input_specs=input_target_specs) | 983 | + [ |
| 820 | - ]) | 984 | + OpSpec( |
| 985 | + output_specs=op_schema.args_schema[0].strategies[0].output_spec, | ||
| 986 | + input_specs=input_target_specs, | ||
| 987 | + ) | ||
| 988 | + ] | ||
| 989 | + ) | ||
| 821 | 990 | ||
| 822 | return output_strategy | 991 | return output_strategy |
| 823 | 992 | ||
| 824 | 993 | ||
| 994 | + | ||
| 995 | +def custom_bmm_strategy(op_schema: OpSchema): | ||
| 996 | + mesh = op_schema.get_mesh_from_args() | ||
| 997 | + return _mm_like_strategy("bmk,bkn->bmn", mesh, op_schema) | ||
| 998 | + | ||
| 999 | + | ||
| 825 | customized_ops = { | 1000 | customized_ops = { |
| 826 | npu.npu_grouped_matmul.default: _npu_grouped_matmul_handler, | 1001 | npu.npu_grouped_matmul.default: _npu_grouped_matmul_handler, |
| 827 | npu.npu_grouped_matmul.List: _npu_grouped_matmul_handler, | 1002 | npu.npu_grouped_matmul.List: _npu_grouped_matmul_handler, |
| 828 | npu.npu_apply_adam_w.out: _npu_apply_adam_w_handler, | 1003 | npu.npu_apply_adam_w.out: _npu_apply_adam_w_handler, |
| 829 | npu.npu_all_gather_base_mm.default: npu_comm_mm_fusion_handler, | 1004 | npu.npu_all_gather_base_mm.default: npu_comm_mm_fusion_handler, |
| 830 | - npu.npu_mm_reduce_scatter_base.default: npu_comm_mm_fusion_handler | 1005 | + npu.npu_mm_reduce_scatter_base.default: npu_comm_mm_fusion_handler, |
| 831 | } | 1006 | } |
| 832 | 1007 | ||
| 833 | old_handlers = DTensor._op_dispatcher._custom_op_handlers | 1008 | old_handlers = DTensor._op_dispatcher._custom_op_handlers |
| @@ -1,14 +1,14 @@ | |||
| 1 | - | ||
| 2 | import torch | 1 | import torch |
| 3 | from torch.distributed.tensor._op_schema import OpSchema, RuntimeSchemaInfo | 2 | from torch.distributed.tensor._op_schema import OpSchema, RuntimeSchemaInfo |
| 4 | -from torch.distributed.tensor._ops.registration import register_op_strategy | ||
| 5 | from torch.distributed.tensor._ops._pointwise_ops import pointwise_strategy | 3 | from torch.distributed.tensor._ops._pointwise_ops import pointwise_strategy |
| 4 | +from torch.distributed.tensor._ops.registration import register_op_strategy | ||
| 6 | 5 | ||
| 7 | 6 | ||
| 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(op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"]))( | ||
| 26 | + custom_linear_pointwise_strategy | ||
| 27 | + ) | ||
| 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(op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"]))( |
| 26 | - op, schema_info=RuntimeSchemaInfo(static_kwargkey=["out"]) | 47 | + pointwise_strategy |
| 27 | - )(custom_pointwise_strategy) | 48 | + ) |
| @@ -0,0 +1,8 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.distributed.tensor._ops._view_ops import register_op_strategy_map | ||
| 3 | + | ||
| 4 | + | ||
| 5 | +npu = torch.ops.npu | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +register_op_strategy_map(npu.npu_transpose.default, torch.permute) | ||
| @@ -1,31 +1,51 @@ | |||
| 1 | -__all__ = ["npu_combine_tensors", "get_part_combined_tensor", "is_combined_tensor_valid", "FlopsCounter", | 1 | +__all__ = [ |
| 2 | - "set_thread_affinity", "reset_thread_affinity", "save_async", "get_cann_version"] | 2 | + "npu_combine_tensors", |
| 3 | - | 3 | + "get_part_combined_tensor", |
| 4 | -from torch_npu import _C | 4 | + "is_combined_tensor_valid", |
| 5 | -from torch_npu.npu.utils import get_cann_version | 5 | + "FlopsCounter", |
| 6 | -from ._module import _apply_module_patch | 6 | + "set_thread_affinity", |
| 7 | -from .tensor_methods import _add_tensor_methods | 7 | + "reset_thread_affinity", |
| 8 | -from .storage import _add_storage_methods | 8 | + "save_async", |
| 9 | -from .combine_tensors import npu_combine_tensors, get_part_combined_tensor, is_combined_tensor_valid | 9 | + "get_cann_version", |
| 10 | -from .serialization import _add_serialization_methods, save_async | 10 | +] |
| 11 | -from .npu_intercept import _cann_package_check, _add_intercept_methods | 11 | + |
| 12 | -from .dtensor import _register_ops_under_dtensor_rules | 12 | +from torch_npu import _C |
| 13 | -from .collect_env import _add_collect_env_methods | 13 | +from torch_npu.npu.utils import get_cann_version |
| 14 | -from ._dynamo import add_dynamo_methods | 14 | +from torch_npu.utils._error_code import ErrCode, pta_error |
| 15 | -from ._inductor import _inductor_register_device_op_overrides | 15 | + |
| 16 | -from ._optim import add_optim_method | 16 | +from ._dynamo import add_dynamo_methods |
| 17 | -from .asd_detector import set_asd_loss_scale, register_asd_hook | 17 | +from ._graph_tree import _apply_npugraph_tree_methods |
| 18 | -from .utils import _print_error_log, _print_warn_log, _print_info_log, _apply_npu_show_warning, _should_print_warning | 18 | +from ._inductor import _inductor_register_device_op_overrides |
| 19 | -from ._step import add_perf_dump_patch | 19 | +from ._module import _apply_module_patch |
| 20 | -from .flops_count import _FlopsCounter as FlopsCounter | 20 | +from ._npu_meta_registration import npu_patch_meta |
| 21 | -from .affinity import _set_thread_affinity as set_thread_affinity | 21 | +from ._optim import add_optim_method |
| 22 | -from .affinity import _reset_thread_affinity as reset_thread_affinity | 22 | +from ._step import add_perf_dump_patch |
| 23 | -from ._graph_tree import _apply_npugraph_tree_methods | 23 | +from .affinity import ( |
| 24 | -from .dlpack import _apply_dlpack_patch | 24 | + _reset_thread_affinity as reset_thread_affinity, |
| 25 | -from ._npu_meta_registration import npu_patch_meta | 25 | + _set_thread_affinity as set_thread_affinity, |
| 26 | -from torch_npu.utils._error_code import ErrCode, pta_error | 26 | +) |
| 27 | - | 27 | +from .asd_detector import register_asd_hook, set_asd_loss_scale |
| 28 | - | 28 | +from .collect_env import _add_collect_env_methods |
| 29 | -# init flopcount | 29 | +from .combine_tensors import ( |
| 30 | -if not _C._flops_count_init(): | 30 | + get_part_combined_tensor, |
| 31 | - raise RuntimeError("flopcount initialization failed" + pta_error(ErrCode.UNAVAIL)) | 31 | + is_combined_tensor_valid, |
| 32 | + npu_combine_tensors, | ||
| 33 | +) | ||
| 34 | +from .dlpack import _apply_dlpack_patch | ||
| 35 | +from .flops_count import _FlopsCounter as FlopsCounter | ||
| 36 | +from .npu_intercept import _add_intercept_methods, _cann_package_check | ||
| 37 | +from .serialization import _add_serialization_methods, save_async | ||
| 38 | +from .storage import _add_storage_methods | ||
| 39 | +from .tensor_methods import _add_tensor_methods | ||
| 40 | +from .utils import ( | ||
| 41 | + _apply_npu_show_warning, | ||
| 42 | + _print_error_log, | ||
| 43 | + _print_info_log, | ||
| 44 | + _print_warn_log, | ||
| 45 | + _should_print_warning, | ||
| 46 | +) | ||
| 47 | + | ||
| 48 | + | ||
| 49 | +# init flopcount | ||
| 50 | +if not _C._flops_count_init(): | ||
| 51 | + raise RuntimeError("flopcount initialization failed" + pta_error(ErrCode.UNAVAIL)) | ||
| @@ -1,51 +0,0 @@ | |||
| 1 | -import torch | ||
| 2 | -from torch.distributed.tensor._ops._common_rules import pointwise_rule | ||
| 3 | -from torch.distributed.tensor._ops.registration import register_prop_rule | ||
| 4 | -from torch.distributed.tensor._ops.utils import normalize_dims | ||
| 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) | ||