已合并
refactor dtensor rules #34587
kisnwang创建于 4月28日
refactor dtensor rules #34587
已合并
共 11 个文件变更+777-409
| @@ -1,7 +1,12 @@ | |||
| 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._internal.common_dtensor import NPUDTensorTestBase | ||
| 8 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU, with_comms | ||
| 9 | + | ||
| 5 | import torch | 10 | import torch |
| 6 | from torch.distributed._tensor import DeviceMesh, distribute_tensor | 11 | from torch.distributed._tensor import DeviceMesh, distribute_tensor |
| 7 | from torch.distributed._tensor.api import DTensor | 12 | from torch.distributed._tensor.api import DTensor |
| @@ -13,9 +18,8 @@ from torch.distributed.tensor.placement_types import ( | |||
| 13 | ) | 18 | ) |
| 14 | from torch.testing._internal.common_utils import run_tests | 19 | from torch.testing._internal.common_utils import run_tests |
| 15 | 20 | ||
| 16 | -import torch_npu | 21 | + |
| 17 | -from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 22 | +npu = torch.ops.npu |
| 18 | -from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase | ||
| 19 | 23 | ||
| 20 | 24 | ||
| 21 | class DistMatrixOpsTest(NPUDTensorTestBase): | 25 | class DistMatrixOpsTest(NPUDTensorTestBase): |
| @@ -48,9 +52,13 @@ class DistMatrixOpsTest(NPUDTensorTestBase): | |||
| 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) |
| @@ -87,7 +95,7 @@ class DistMatrixOpsTest(NPUDTensorTestBase): | |||
| 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) |
| @@ -150,12 +158,12 @@ class DistMatrixOpsTest(NPUDTensorTestBase): | |||
| 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(NPUDTensorTestBase): | |||
| 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(NPUDTensorTestBase): | |||
| 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,27 +1,29 @@ | |||
| 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._internal.common_dtensor import NPUDTensorTestBase | ||
| 10 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU, with_comms | ||
| 8 | 11 | ||
| 12 | +import torch | ||
| 9 | import torch.utils._pytree as pytree | 13 | import torch.utils._pytree as pytree |
| 10 | from torch import Tensor | 14 | from torch import Tensor |
| 11 | - | ||
| 12 | from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor | 15 | from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor |
| 16 | +from torch.distributed.distributed_c10d import ReduceOp | ||
| 13 | from torch.distributed.tensor.placement_types import ( | 17 | from torch.distributed.tensor.placement_types import ( |
| 14 | Partial, | 18 | Partial, |
| 15 | Placement, | 19 | Placement, |
| 16 | Replicate, | 20 | Replicate, |
| 17 | Shard, | 21 | Shard, |
| 18 | ) | 22 | ) |
| 19 | -from torch.distributed.distributed_c10d import ReduceOp | ||
| 20 | from torch.testing._internal.common_utils import run_tests | 23 | from torch.testing._internal.common_utils import run_tests |
| 21 | 24 | ||
| 22 | -import torch_npu | 25 | + |
| 23 | -from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 26 | +npu = torch.ops.npu |
| 24 | -from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase | ||
| 25 | 27 | ||
| 26 | 28 | ||
| 27 | def no_op(): | 29 | def no_op(): |
| @@ -79,9 +81,9 @@ class DistElementwiseOpsTest(NPUDTensorTestBase): | |||
| 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(NPUDTensorTestBase): | |||
| 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, |
| @@ -161,6 +163,30 @@ class DistElementwiseOpsTest(NPUDTensorTestBase): | |||
| 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)], |
| @@ -257,16 +283,21 @@ class DistElementwiseOpsTest(NPUDTensorTestBase): | |||
| 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,11 +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._internal.common_dtensor import NPUDTensorTestBase | ||
| 7 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU, with_comms | ||
| 3 | 8 | ||
| 4 | import torch | 9 | import torch |
| 5 | import torch.distributed as dist | 10 | import torch.distributed as dist |
| 6 | from torch import rand, randn, Tensor | 11 | from torch import rand, randn, Tensor |
| 7 | -from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate, Shard | 12 | +from torch.distributed.tensor import DeviceMesh, distribute_tensor, Replicate, Shard |
| 8 | -from torch.distributed.tensor.debug import CommDebugMode | ||
| 9 | from torch.distributed.tensor._ops._view_ops import ( | 13 | from torch.distributed.tensor._ops._view_ops import ( |
| 10 | Broadcast, | 14 | Broadcast, |
| 11 | dim_maps, | 15 | dim_maps, |
| @@ -16,13 +20,13 @@ from torch.distributed.tensor._ops._view_ops import ( | |||
| 16 | Split, | 20 | Split, |
| 17 | view_groups, | 21 | view_groups, |
| 18 | ) | 22 | ) |
| 19 | -from torch.distributed._tensor.placement_types import Placement | 23 | +from torch.distributed.tensor.debug import CommDebugMode |
| 24 | +from torch.distributed.tensor.placement_types import Placement | ||
| 20 | from torch.testing._internal.common_utils import run_tests | 25 | from torch.testing._internal.common_utils import run_tests |
| 21 | from torch.utils import _pytree as pytree | 26 | from torch.utils import _pytree as pytree |
| 22 | 27 | ||
| 23 | -import torch_npu | 28 | + |
| 24 | -from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 29 | +npu = torch.ops.npu |
| 25 | -from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase | ||
| 26 | 30 | ||
| 27 | 31 | ||
| 28 | class TestViewOps(NPUDTensorTestBase): | 32 | class TestViewOps(NPUDTensorTestBase): |
| @@ -152,7 +156,7 @@ class TestViewOps(NPUDTensorTestBase): | |||
| 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(NPUDTensorTestBase): | |||
| 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 | # Apply monkey-patches. | 276 | # Apply monkey-patches. |
| @@ -220,8 +280,14 @@ _asd_patch() | |||
| 220 | _except_handler.patch_excepthook() | 280 | _except_handler.patch_excepthook() |
| 221 | 281 | ||
| 222 | _warn_msg = { | 282 | _warn_msg = { |
| 223 | - "DropoutWithByteMask" : "torch.nn.DropoutWithByteMask is deprecated and will be removed in future version. Use torch_npu.contrib.module.DropoutWithByteMask instead.", | 283 | + "DropoutWithByteMask": ( |
| 224 | - "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.", | 284 | + "torch.nn.DropoutWithByteMask is deprecated and will be removed in future version. " |
| 285 | + "Use torch_npu.contrib.module.DropoutWithByteMask instead." | ||
| 286 | + ), | ||
| 287 | + "dropout_with_byte_mask": ( | ||
| 288 | + "torch.nn.functional.dropout_with_byte_mask is deprecated and will be removed in future version. " | ||
| 289 | + "Use torch_npu.contrib.function.dropout_with_byte_mask instead." | ||
| 290 | + ), | ||
| 225 | } | 291 | } |
| 226 | 292 | ||
| 227 | 293 | ||
| @@ -230,9 +296,16 @@ def _wrap_torch_patch_warning_func(func): | |||
| 230 | def wrapper(*args, **kwargs): | 296 | def wrapper(*args, **kwargs): |
| 231 | warnings.warn(_warn_msg[func.__name__]) | 297 | warnings.warn(_warn_msg[func.__name__]) |
| 232 | return func(*args, **kwargs) | 298 | return func(*args, **kwargs) |
| 299 | + | ||
| 233 | return wrapper | 300 | return wrapper |
| 234 | -setattr(torch.nn, "DropoutWithByteMask", _wrap_torch_patch_warning_func(torch.nn.DropoutWithByteMask)) | 301 | + |
| 235 | -setattr(torch.nn.functional, "dropout_with_byte_mask", _wrap_torch_patch_warning_func(torch.nn.functional.dropout_with_byte_mask)) | 302 | + |
| 303 | +torch.nn.DropoutWithByteMask = _wrap_torch_patch_warning_func( | ||
| 304 | + torch.nn.DropoutWithByteMask | ||
| 305 | +) | ||
| 306 | +torch.nn.functional.dropout_with_byte_mask = _wrap_torch_patch_warning_func( | ||
| 307 | + torch.nn.functional.dropout_with_byte_mask | ||
| 308 | +) | ||
| 236 | # this must be placed at the end | 309 | # this must be placed at the end |
| 237 | torch_npu._C._initExtension() | 310 | torch_npu._C._initExtension() |
| 238 | 311 | ||
| @@ -241,32 +314,49 @@ def _new_process_group_hccl_helper(dist_backend_opts, pg_options): | |||
| 241 | store = dist_backend_opts.store | 314 | store = dist_backend_opts.store |
| 242 | group_rank = dist_backend_opts.group_rank | 315 | group_rank = dist_backend_opts.group_rank |
| 243 | group_size = dist_backend_opts.group_size | 316 | group_size = dist_backend_opts.group_size |
| 244 | - if pg_options is None or not isinstance(pg_options, torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options): | 317 | + if pg_options is None or not isinstance( |
| 318 | + pg_options, torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options | ||
| 319 | + ): | ||
| 245 | pg_options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options() | 320 | pg_options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options() |
| 246 | pg_options.is_high_priority_stream = False | 321 | pg_options.is_high_priority_stream = False |
| 247 | pg_options._timeout = dist_backend_opts.timeout | 322 | pg_options._timeout = dist_backend_opts.timeout |
| 248 | pg_options.global_ranks_in_group = dist_backend_opts.global_ranks_in_group | 323 | pg_options.global_ranks_in_group = dist_backend_opts.global_ranks_in_group |
| 249 | pg_options.group_id = dist_backend_opts.group_id | 324 | pg_options.group_id = dist_backend_opts.group_id |
| 250 | - return torch_npu._C._distributed_c10d.ProcessGroupHCCL(store, group_rank, group_size, pg_options) | 325 | + return torch_npu._C._distributed_c10d.ProcessGroupHCCL( |
| 326 | + store, group_rank, group_size, pg_options | ||
| 327 | + ) | ||
| 251 | 328 | ||
| 252 | 329 | ||
| 253 | def _new_process_group_lccl_helper(dist_backend_opts, pg_options): | 330 | def _new_process_group_lccl_helper(dist_backend_opts, pg_options): |
| 254 | store = dist_backend_opts.store | 331 | store = dist_backend_opts.store |
| 255 | group_rank = dist_backend_opts.group_rank | 332 | group_rank = dist_backend_opts.group_rank |
| 256 | group_size = dist_backend_opts.group_size | 333 | group_size = dist_backend_opts.group_size |
| 257 | - return torch_npu._C._distributed_c10d.ProcessGroupLCCL(store, group_rank, group_size) | 334 | + return torch_npu._C._distributed_c10d.ProcessGroupLCCL( |
| 335 | + store, group_rank, group_size | ||
| 336 | + ) | ||
| 258 | 337 | ||
| 259 | 338 | ||
| 260 | def _register_distributed_backend_for_npu(): | 339 | def _register_distributed_backend_for_npu(): |
| 261 | # init and register hccl backend | 340 | # init and register hccl backend |
| 262 | # Note: Since torch 2.8, the hccl backend must be registered at first to keep a right default_device_backend_map | 341 | # Note: Since torch 2.8, the hccl backend must be registered at first to keep a right default_device_backend_map |
| 263 | - torch.distributed.Backend.register_backend("hccl", lambda dist_backend_opts, pg_options: | 342 | + torch.distributed.Backend.register_backend( |
| 264 | - _new_process_group_hccl_helper(dist_backend_opts, pg_options), extended_api=True, devices=["npu"]) | 343 | + "hccl", |
| 344 | + lambda dist_backend_opts, pg_options: _new_process_group_hccl_helper( | ||
| 345 | + dist_backend_opts, pg_options | ||
| 346 | + ), | ||
| 347 | + extended_api=True, | ||
| 348 | + devices=["npu"], | ||
| 349 | + ) | ||
| 265 | 350 | ||
| 266 | # init and register lccl backend | 351 | # init and register lccl backend |
| 267 | - torch.distributed.Backend.register_backend("lccl", lambda dist_backend_opts, pg_options: | 352 | + torch.distributed.Backend.register_backend( |
| 268 | - _new_process_group_lccl_helper(dist_backend_opts, pg_options), extended_api=True, devices=["npu"]) | 353 | + "lccl", |
| 269 | - | 354 | + lambda dist_backend_opts, pg_options: _new_process_group_lccl_helper( |
| 355 | + dist_backend_opts, pg_options | ||
| 356 | + ), | ||
| 357 | + extended_api=True, | ||
| 358 | + devices=["npu"], | ||
| 359 | + ) | ||
| 270 | 360 | ||
| 271 | 361 | ||
| 272 | # init and register distributed backend | 362 | # init and register distributed backend |
| @@ -287,6 +377,7 @@ def _npu_shutdown(): | |||
| 287 | torch_npu.asd.asd.matmul_check._cleanup() | 377 | torch_npu.asd.asd.matmul_check._cleanup() |
| 288 | if torch_npu.npu.aclnn._use_static_aclnn_kernel: | 378 | if torch_npu.npu.aclnn._use_static_aclnn_kernel: |
| 289 | from torch_npu._inductor.npu_static_kernel import uninstall_static_kernel | 379 | from torch_npu._inductor.npu_static_kernel import uninstall_static_kernel |
| 380 | + | ||
| 290 | uninstall_static_kernel() | 381 | uninstall_static_kernel() |
| 291 | 382 | ||
| 292 | 383 | ||
| @@ -296,14 +387,11 @@ atexit.register(_npu_shutdown) | |||
| 296 | # init and register rpc npu backend | 387 | # init and register rpc npu backend |
| 297 | _rpc_backend_registry() | 388 | _rpc_backend_registry() |
| 298 | 389 | ||
| 299 | -# register rules for ops in dtensor | ||
| 300 | -_register_ops_under_dtensor_rules() | ||
| 301 | - | ||
| 302 | # register npu device interface for dynamo | 390 | # register npu device interface for dynamo |
| 303 | _dynamo_register_interface_for_device() | 391 | _dynamo_register_interface_for_device() |
| 304 | 392 | ||
| 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 |
| 7 | -import torch_npu.distributed.tensor._sharded_tensor_patch | 6 | +import torch_npu.distributed.tensor._pointwise_ops |
| 7 | +import torch_npu.distributed.tensor._sharded_tensor_patch | ||
| 8 | +import torch_npu.distributed.tensor._view_ops | ||
| @@ -1,22 +1,26 @@ | |||
| 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.utils import register_op_strategy, expand_to_full_mesh_op_strategy | ||
| 8 | 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 | ||
| 9 | from torch.distributed.tensor._op_schema import ( | 9 | from torch.distributed.tensor._op_schema import ( |
| 10 | OpInfo, | 10 | OpInfo, |
| 11 | OpSchema, | 11 | OpSchema, |
| 12 | - OpStrategy, | ||
| 13 | OpSpec, | 12 | OpSpec, |
| 13 | + OpStrategy, | ||
| 14 | OutputSharding, | 14 | OutputSharding, |
| 15 | RuntimeSchemaInfo, | 15 | RuntimeSchemaInfo, |
| 16 | - TupleStrategy | 16 | + TupleStrategy, |
| 17 | +) | ||
| 18 | +from torch.distributed.tensor._ops._matrix_ops import _mm_like_strategy | ||
| 19 | +from torch.distributed.tensor._ops.utils import ( | ||
| 20 | + expand_to_full_mesh_op_strategy, | ||
| 21 | + register_op_strategy, | ||
| 17 | ) | 22 | ) |
| 18 | 23 | ||
| 19 | -import torch_npu | ||
| 20 | 24 | ||
| 21 | try: | 25 | try: |
| 22 | from torch.utils import _cxx_pytree as pytree | 26 | from torch.utils import _cxx_pytree as pytree |
| @@ -24,11 +28,12 @@ except ImportError: | |||
| 24 | from torch.utils import _pytree as pytree | 28 | from torch.utils import _pytree as pytree |
| 25 | 29 | ||
| 26 | from ._common import ( | 30 | from ._common import ( |
| 31 | + get_empty_local_results, | ||
| 27 | get_redistributed_local_args, | 32 | get_redistributed_local_args, |
| 28 | get_redistributed_local_kwargs, | 33 | get_redistributed_local_kwargs, |
| 29 | - get_empty_local_results | ||
| 30 | ) | 34 | ) |
| 31 | 35 | ||
| 36 | + | ||
| 32 | aten = torch.ops.aten | 37 | aten = torch.ops.aten |
| 33 | npu = torch.ops.npu | 38 | npu = torch.ops.npu |
| 34 | 39 | ||
| @@ -36,7 +41,9 @@ npu = torch.ops.npu | |||
| 36 | def _get_max_shardable_dim(tensor): | 41 | def _get_max_shardable_dim(tensor): |
| 37 | shape = tensor.shape | 42 | shape = tensor.shape |
| 38 | world_size = torch.distributed.get_world_size() | 43 | world_size = torch.distributed.get_world_size() |
| 39 | - divisible_dims = [(idx, dim) for idx, dim in enumerate(shape) if dim % world_size == 0] | 44 | + divisible_dims = [ |
| 45 | + (idx, dim) for idx, dim in enumerate(shape) if dim % world_size == 0 | ||
| 46 | + ] | ||
| 40 | if divisible_dims: | 47 | if divisible_dims: |
| 41 | idx, _ = max(divisible_dims, key=lambda x: x[1]) | 48 | idx, _ = max(divisible_dims, key=lambda x: x[1]) |
| 42 | return idx | 49 | return idx |
| @@ -44,7 +51,7 @@ def _get_max_shardable_dim(tensor): | |||
| 44 | return -1 | 51 | return -1 |
| 45 | 52 | ||
| 46 | 53 | ||
| 47 | -def _handle_tensor_list_in_kwargs(kwargs: Dict[str, object], op_info: OpInfo) -> None: | 54 | +def _handle_tensor_list_in_kwargs(kwargs: dict[str, object], op_info: OpInfo) -> None: |
| 48 | for key, value in kwargs.items(): | 55 | for key, value in kwargs.items(): |
| 49 | if isinstance(value, list) and all(isinstance(e, DTensor) for e in value): | 56 | if isinstance(value, list) and all(isinstance(e, DTensor) for e in value): |
| 50 | new_schema = [] | 57 | new_schema = [] |
| @@ -52,12 +59,16 @@ def _handle_tensor_list_in_kwargs(kwargs: Dict[str, object], op_info: OpInfo) -> | |||
| 52 | for dtensor in value: | 59 | for dtensor in value: |
| 53 | new_schema.append(dtensor._spec) | 60 | new_schema.append(dtensor._spec) |
| 54 | new_local_tensors.append(dtensor._local_tensor) | 61 | new_local_tensors.append(dtensor._local_tensor) |
| 55 | - op_info.schema.kwargs_schema[key] = tuple(new_schema) # list is not hashable for cache | 62 | + op_info.schema.kwargs_schema[key] = tuple( |
| 63 | + new_schema | ||
| 64 | + ) # list is not hashable for cache | ||
| 56 | op_info.local_kwargs[key] = new_local_tensors | 65 | op_info.local_kwargs[key] = new_local_tensors |
| 57 | 66 | ||
| 58 | op_info.schema._recompute_comparison_key() | 67 | op_info.schema._recompute_comparison_key() |
| 59 | 68 | ||
| 69 | + | ||
| 60 | if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | 70 | if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": |
| 71 | + | ||
| 61 | 72 | ||
| 62 | def custom_matmul_strategy( | 73 | def custom_matmul_strategy( |
| 63 | tensor1: DTensorSpec, | 74 | tensor1: DTensorSpec, |
| @@ -110,13 +121,24 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 110 | diff = abs(len1 - len2) | 121 | diff = abs(len1 - len2) |
| 111 | is_shape1_longer = len1 > len2 | 122 | is_shape1_longer = len1 > len2 |
| 112 | 123 | ||
| 113 | - for i in range(min(len1, len2) - 3, -1, -1): | 124 | + for i in range(min(len1, len2) - 3, -1, -1): |
| 114 | - shape1_shardable = shape1[i + diff] % tensor1.mesh.size(0) == 0 if is_shape1_longer else shape1[i] % tensor1.mesh.size(0) == 0 | 125 | + shape1_shardable = ( |
| 115 | - shape2_shardable = shape2[i] % tensor2.mesh.size(0) == 0 if is_shape1_longer else shape2[i + diff] % tensor2.mesh.size(0) == 0 | 126 | + shape1[i + diff] % tensor1.mesh.size(0) == 0 |
| 127 | + if is_shape1_longer | ||
| 128 | + else shape1[i] % tensor1.mesh.size(0) == 0 | ||
| 129 | + ) | ||
| 130 | + shape2_shardable = ( | ||
| 131 | + shape2[i] % tensor2.mesh.size(0) == 0 | ||
| 132 | + if is_shape1_longer | ||
| 133 | + else shape2[i + diff] % tensor2.mesh.size(0) == 0 | ||
| 134 | + ) | ||
| 116 | 135 | ||
| 117 | if shape1_shardable and shape2_shardable: | 136 | if shape1_shardable and shape2_shardable: |
| 118 | - strategy_batch = ([Shard(i + diff)], [Shard(i + diff), Shard(i)]) if is_shape1_longer \ | 137 | + strategy_batch = ( |
| 138 | + ([Shard(i + diff)], [Shard(i + diff), Shard(i)]) | ||
| 139 | + if is_shape1_longer | ||
| 119 | else ([Shard(i + diff)], [Shard(i), Shard(i + diff)]) | 140 | else ([Shard(i + diff)], [Shard(i), Shard(i + diff)]) |
| 141 | + ) | ||
| 120 | acceptable_shardings.append(strategy_batch) | 142 | acceptable_shardings.append(strategy_batch) |
| 121 | 143 | ||
| 122 | for i in range(diff - 1, -1, -1): | 144 | for i in range(diff - 1, -1, -1): |
| @@ -128,26 +150,34 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 128 | acceptable_shardings.append(strategy_batch) | 150 | acceptable_shardings.append(strategy_batch) |
| 129 | # Shard tensor1 | 151 | # Shard tensor1 |
| 130 | if shape1[len(shape1) - 2] % tensor1.mesh.size(0) == 0: | 152 | if shape1[len(shape1) - 2] % tensor1.mesh.size(0) == 0: |
| 131 | - strategy_tensor1 = ([Shard(len(output_shape) - 2)], [Shard(len(shape1) - 2), Replicate()]) | 153 | + strategy_tensor1 = ( |
| 154 | + [Shard(len(output_shape) - 2)], | ||
| 155 | + [Shard(len(shape1) - 2), Replicate()], | ||
| 156 | + ) | ||
| 132 | acceptable_shardings.append(strategy_tensor1) | 157 | acceptable_shardings.append(strategy_tensor1) |
| 133 | # Shard tensor2 | 158 | # Shard tensor2 |
| 134 | if shape2[len(shape2) - 1] % tensor2.mesh.size(0) == 0: | 159 | if shape2[len(shape2) - 1] % tensor2.mesh.size(0) == 0: |
| 135 | - strategy_tensor2 = ([Shard(len(output_shape) - 1)], [Replicate(), Shard(len(shape2) - 1)]) | 160 | + strategy_tensor2 = ( |
| 161 | + [Shard(len(output_shape) - 1)], | ||
| 162 | + [Replicate(), Shard(len(shape2) - 1)], | ||
| 163 | + ) | ||
| 136 | acceptable_shardings.append(strategy_tensor2) | 164 | acceptable_shardings.append(strategy_tensor2) |
| 137 | # Shard tensor1 and tensor2 | 165 | # Shard tensor1 and tensor2 |
| 138 | if shape1[len(shape1) - 1] % tensor1.mesh.size(0) == 0: | 166 | if shape1[len(shape1) - 1] % tensor1.mesh.size(0) == 0: |
| 139 | - strategy_3 = ([Partial()], [Shard(len(shape1) - 1), Shard(len(shape2) - 2)]) | 167 | + strategy_3 = ( |
| 168 | + [Partial()], | ||
| 169 | + [Shard(len(shape1) - 1), Shard(len(shape2) - 2)], | ||
| 170 | + ) | ||
| 140 | acceptable_shardings.append(strategy_3) | 171 | acceptable_shardings.append(strategy_3) |
| 141 | 172 | ||
| 142 | return acceptable_shardings | 173 | return acceptable_shardings |
| 143 | 174 | ||
| 144 | - | ||
| 145 | 175 | ||
| 146 | def custom_matmul_backward_strategy( | 176 | def custom_matmul_backward_strategy( |
| 147 | grad: DTensorSpec, | 177 | grad: DTensorSpec, |
| 148 | tensor1: DTensorSpec, | 178 | tensor1: DTensorSpec, |
| 149 | tensor2: DTensorSpec, | 179 | tensor2: DTensorSpec, |
| 150 | - mask: List[bool], | 180 | + mask: list[bool], |
| 151 | ): | 181 | ): |
| 152 | grad_dim = len(grad.shape) | 182 | grad_dim = len(grad.shape) |
| 153 | tensor1_dim = len(tensor1.shape) | 183 | tensor1_dim = len(tensor1.shape) |
| @@ -156,8 +186,8 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 156 | acceptable_shardings = [] | 186 | acceptable_shardings = [] |
| 157 | 187 | ||
| 158 | replicate_strategy = ( | 188 | replicate_strategy = ( |
| 159 | - [Replicate(), Replicate()], | 189 | + [Replicate(), Replicate()], |
| 160 | - [Replicate(), Replicate(), Replicate(), None] | 190 | + [Replicate(), Replicate(), Replicate(), None], |
| 161 | ) | 191 | ) |
| 162 | acceptable_shardings.append(replicate_strategy) | 192 | acceptable_shardings.append(replicate_strategy) |
| 163 | 193 | ||
| @@ -166,21 +196,21 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 166 | elif tensor1_dim >= 2 and (tensor2_dim == 1 or tensor2_dim == 2): | 196 | elif tensor1_dim >= 2 and (tensor2_dim == 1 or tensor2_dim == 2): |
| 167 | if tensor2.shape[0] % tensor2.mesh.size(0) == 0: | 197 | if tensor2.shape[0] % tensor2.mesh.size(0) == 0: |
| 168 | strategy_1 = ( | 198 | strategy_1 = ( |
| 169 | - [Shard(tensor1_dim - 1), Shard(0)], | 199 | + [Shard(tensor1_dim - 1), Shard(0)], |
| 170 | - [Replicate(), Shard(tensor1_dim - 1), Shard(0), None] | 200 | + [Replicate(), Shard(tensor1_dim - 1), Shard(0), None], |
| 171 | ) | 201 | ) |
| 172 | acceptable_shardings.append(strategy_1) | 202 | acceptable_shardings.append(strategy_1) |
| 173 | for i in range(tensor1_dim - 1): | 203 | for i in range(tensor1_dim - 1): |
| 174 | if tensor1.shape[i] % tensor1.mesh.size(0) == 0: | 204 | if tensor1.shape[i] % tensor1.mesh.size(0) == 0: |
| 175 | strategy_2 = ( | 205 | strategy_2 = ( |
| 176 | - [Shard(i), Partial()], | 206 | + [Shard(i), Partial()], |
| 177 | - [Shard(i), Shard(i), Replicate(), None] | 207 | + [Shard(i), Shard(i), Replicate(), None], |
| 178 | ) | 208 | ) |
| 179 | acceptable_shardings.append(strategy_2) | 209 | acceptable_shardings.append(strategy_2) |
| 180 | if tensor2_dim == 2 and tensor2.shape[1] % tensor2.mesh.size(0) == 0: | 210 | if tensor2_dim == 2 and tensor2.shape[1] % tensor2.mesh.size(0) == 0: |
| 181 | strategy_3 = ( | 211 | strategy_3 = ( |
| 182 | - [Partial(), Shard(1)], | 212 | + [Partial(), Shard(1)], |
| 183 | - [Shard(grad_dim - 1), Replicate(), Shard(1), None] | 213 | + [Shard(grad_dim - 1), Replicate(), Shard(1), None], |
| 184 | ) | 214 | ) |
| 185 | acceptable_shardings.append(strategy_3) | 215 | acceptable_shardings.append(strategy_3) |
| 186 | return acceptable_shardings | 216 | return acceptable_shardings |
| @@ -188,100 +218,135 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1": | |||
| 188 | is_special = tensor2_dim == 2 and tensor1_dim == 1 | 218 | is_special = tensor2_dim == 2 and tensor1_dim == 1 |
| 189 | if tensor1.shape[-1] % tensor1.mesh.size(0) == 0: | 219 | if tensor1.shape[-1] % tensor1.mesh.size(0) == 0: |
| 190 | strategy_1 = ( | 220 | strategy_1 = ( |
| 191 | - [Shard(tensor1_dim if is_special else tensor1_dim - 1), Shard(tensor2_dim - 2)], | 221 | + [ |
| 192 | - [Replicate(), Shard(tensor1_dim - 1), Shard(tensor2_dim - 2), None] | 222 | + Shard(tensor1_dim if is_special else tensor1_dim - 1), |
| 223 | + Shard(tensor2_dim - 2), | ||
| 224 | + ], | ||
| 225 | + [Replicate(), Shard(tensor1_dim - 1), Shard(tensor2_dim - 2), None], | ||
| 193 | ) | 226 | ) |
| 194 | acceptable_shardings.append(strategy_1) | 227 | acceptable_shardings.append(strategy_1) |
| 195 | if tensor2.shape[-1] % tensor2.mesh.size(0) == 0: | 228 | if tensor2.shape[-1] % tensor2.mesh.size(0) == 0: |
| 196 | strategy_2 = ( | 229 | strategy_2 = ( |
| 197 | - [Partial(), Shard(tensor2_dim - 1)], | 230 | + [Partial(), Shard(tensor2_dim - 1)], |
| 198 | - [Shard(grad_dim - 1), Replicate(), Shard(tensor2_dim - 1), None] | 231 | + [Shard(grad_dim - 1), Replicate(), Shard(tensor2_dim - 1), None], |
| 199 | ) | 232 | ) |
| 200 | acceptable_shardings.append(strategy_2) | 233 | acceptable_shardings.append(strategy_2) |
| 201 | for i in range(tensor2_dim - 2): | 234 | for i in range(tensor2_dim - 2): |
| 202 | if tensor2.shape[i] % tensor2.mesh.size(0) == 0: | 235 | if tensor2.shape[i] % tensor2.mesh.size(0) == 0: |
| 203 | strategy_3 = ( | 236 | strategy_3 = ( |
| 204 | - [Partial(), Shard(i)], | 237 | + [Partial(), Shard(i)], |
| 205 | - [Shard(i), Replicate(), Shard(i)] | 238 | + [Shard(i), Replicate(), Shard(i)], |
| 206 | ) | 239 | ) |
| 207 | acceptable_shardings.append(strategy_3) | 240 | acceptable_shardings.append(strategy_3) |
| 208 | if tensor1_dim == 2 and tensor1.shape[0] % tensor1.mesh.size(0) == 0: | 241 | if tensor1_dim == 2 and tensor1.shape[0] % tensor1.mesh.size(0) == 0: |
| 209 | strategy_4 = ( | 242 | strategy_4 = ( |
| 210 | - [Shard(0), Partial()], | 243 | + [Shard(0), Partial()], |
| 211 | - [Shard(grad_dim - 2), Shard(0), Replicate(), None] | 244 | + [Shard(grad_dim - 2), Shard(0), Replicate(), None], |
| 212 | ) | 245 | ) |
| 213 | acceptable_shardings.append(strategy_4) | 246 | acceptable_shardings.append(strategy_4) |
| 214 | return acceptable_shardings | 247 | return acceptable_shardings |
| 215 | else: | 248 | else: |
| 216 | if grad.shape[-1] % grad.mesh.size(0) == 0: | 249 | if grad.shape[-1] % grad.mesh.size(0) == 0: |
| 217 | strategy_1 = ( | 250 | strategy_1 = ( |
| 218 | - [Partial(), Shard(grad_dim - 1)], | 251 | + [Partial(), Shard(grad_dim - 1)], |
| 219 | - [Shard(grad_dim - 1), Replicate(), Shard(tensor2_dim - 1), None] | 252 | + [Shard(grad_dim - 1), Replicate(), Shard(tensor2_dim - 1), None], |
| 220 | ) | 253 | ) |
| 221 | acceptable_shardings.append(strategy_1) | 254 | acceptable_shardings.append(strategy_1) |
| 222 | if grad.shape[-2] % grad.mesh.size(0) == 0: | 255 | if grad.shape[-2] % grad.mesh.size(0) == 0: |
| 223 | strategy_2 = ( | 256 | strategy_2 = ( |
| 224 | - [Shard(grad_dim - 2), Partial()], | 257 | + [Shard(grad_dim - 2), Partial()], |
| 225 | - [Shard(grad_dim - 2), Shard(tensor1_dim - 2), Replicate(), None] | 258 | + [Shard(grad_dim - 2), Shard(tensor1_dim - 2), Replicate(), None], |
| 226 | ) | 259 | ) |
| 227 | acceptable_shardings.append(strategy_2) | 260 | acceptable_shardings.append(strategy_2) |
| 228 | if tensor1.shape[-1] % tensor1.mesh.size(0) == 0: | 261 | if tensor1.shape[-1] % tensor1.mesh.size(0) == 0: |
| 229 | strategy_3 = ( | 262 | strategy_3 = ( |
| 230 | - [Shard(grad_dim - 1), Shard(grad_dim - 2)], | 263 | + [Shard(grad_dim - 1), Shard(grad_dim - 2)], |
| 231 | - [Replicate(), Shard(tensor1_dim - 1), Shard(tensor2_dim - 2), None] | 264 | + [Replicate(), Shard(tensor1_dim - 1), Shard(tensor2_dim - 2), None], |
| 232 | ) | 265 | ) |
| 233 | acceptable_shardings.append(strategy_3) | 266 | acceptable_shardings.append(strategy_3) |
| 234 | 267 | ||
| 235 | diff = abs(tensor1_dim - tensor2_dim) | 268 | diff = abs(tensor1_dim - tensor2_dim) |
| 236 | is_shape1_longer = tensor1_dim > tensor2_dim | 269 | is_shape1_longer = tensor1_dim > tensor2_dim |
| 237 | 270 | ||
| 238 | - for i in range(min(tensor1_dim, tensor2_dim) - 3, -1, -1): | 271 | + for i in range(min(tensor1_dim, tensor2_dim) - 3, -1, -1): |
| 239 | - shape1_shardable = tensor1.shape[i + diff] % tensor1.mesh.size(0) == 0 if is_shape1_longer \ | 272 | + shape1_shardable = ( |
| 273 | + tensor1.shape[i + diff] % tensor1.mesh.size(0) == 0 | ||
| 274 | + if is_shape1_longer | ||
| 240 | else tensor1.shape[i] % tensor1.mesh.size(0) == 0 | 275 | else tensor1.shape[i] % tensor1.mesh.size(0) == 0 |
| 241 | - shape2_shardable = tensor2.shape[i] % tensor2.mesh.size(0) == 0 if is_shape1_longer \ | 276 | + ) |
| 277 | + shape2_shardable = ( | ||
| 278 | + tensor2.shape[i] % tensor2.mesh.size(0) == 0 | ||
| 279 | + if is_shape1_longer | ||
| 242 | else tensor2.shape[i + diff] % tensor2.mesh.size(0) == 0 | 280 | else tensor2.shape[i + diff] % tensor2.mesh.size(0) == 0 |
| 281 | + ) | ||
| 243 | if shape1_shardable and shape2_shardable: | 282 | if shape1_shardable and shape2_shardable: |
| 244 | strategy_batch = ( | 283 | strategy_batch = ( |
| 245 | - [Shard(i + diff), Shard(i + diff)], | 284 | + ( |
| 246 | - [Shard(i + diff), Shard(i + diff), Shard(i), None] | 285 | + [Shard(i + diff), Shard(i + diff)], |
| 247 | - ) if is_shape1_longer else ( | 286 | + [Shard(i + diff), Shard(i + diff), Shard(i), None], |
| 248 | - [Shard(i + diff), Shard(i + diff)], | 287 | + ) |
| 249 | - [Shard(i + diff), Shard(i), Shard(i + diff), None] | 288 | + if is_shape1_longer |
| 289 | + else ( | ||
| 290 | + [Shard(i + diff), Shard(i + diff)], | ||
| 291 | + [Shard(i + diff), Shard(i), Shard(i + diff), None], | ||
| 292 | + ) | ||
| 250 | ) | 293 | ) |
| 251 | acceptable_shardings.append(strategy_batch) | 294 | acceptable_shardings.append(strategy_batch) |
| 252 | - | 295 | + |
| 253 | for i in range(diff - 1, -1, -1): | 296 | for i in range(diff - 1, -1, -1): |
| 254 | if is_shape1_longer and tensor1.shape[i] % tensor1.mesh.size(0) == 0: | 297 | if is_shape1_longer and tensor1.shape[i] % tensor1.mesh.size(0) == 0: |
| 255 | strategy_batch = ( | 298 | strategy_batch = ( |
| 256 | - [Shard(i), Partial()], | 299 | + [Shard(i), Partial()], |
| 257 | - [Shard(i), Shard(i), Replicate(), None] | 300 | + [Shard(i), Shard(i), Replicate(), None], |
| 258 | ) | 301 | ) |
| 259 | acceptable_shardings.append(strategy_batch) | 302 | acceptable_shardings.append(strategy_batch) |
| 260 | - elif not is_shape1_longer and tensor2.shape[i] % tensor2.mesh.size(0) == 0: | 303 | + elif ( |
| 304 | + not is_shape1_longer | ||
| 305 | + and tensor2.shape[i] % tensor2.mesh.size(0) == 0 | ||
| 306 | + ): | ||
| 261 | strategy_batch = ( | 307 | strategy_batch = ( |
| 262 | - [Partial(), Shard(i)], | 308 | + [Partial(), Shard(i)], |
| 263 | - [Shard(i), Replicate(), Shard(i), None] | 309 | + [Shard(i), Replicate(), Shard(i), None], |
| 264 | ) | 310 | ) |
| 265 | acceptable_shardings.append(strategy_batch) | 311 | acceptable_shardings.append(strategy_batch) |
| 266 | - | 312 | + |
| 267 | return acceptable_shardings | 313 | return acceptable_shardings |
| 268 | 314 | ||
| 269 | 315 | ||
| 270 | 316 | ||
| 271 | npu.npu_grouped_matmul.default, | 317 | npu.npu_grouped_matmul.default, |
| 272 | schema_info=RuntimeSchemaInfo( | 318 | schema_info=RuntimeSchemaInfo( |
| 273 | - static_kwargkey=["bias", "scale", "offset", "antiquant_scale", "antiquant_offset", "per_token_scale", | 319 | + static_kwargkey=[ |
| 274 | - "group_list", "activation_input", "activation_quant_scale", "activation_quant_offset"], | 320 | + "bias", |
| 275 | - needs_pytree=True | 321 | + "scale", |
| 276 | - ) | 322 | + "offset", |
| 323 | + "antiquant_scale", | ||
| 324 | + "antiquant_offset", | ||
| 325 | + "per_token_scale", | ||
| 326 | + "group_list", | ||
| 327 | + "activation_input", | ||
| 328 | + "activation_quant_scale", | ||
| 329 | + "activation_quant_offset", | ||
| 330 | + ], | ||
| 331 | + needs_pytree=True, | ||
| 332 | + ), | ||
| 277 | ) | 333 | ) |
| 278 | 334 | ||
| 279 | npu.npu_grouped_matmul.List, | 335 | npu.npu_grouped_matmul.List, |
| 280 | schema_info=RuntimeSchemaInfo( | 336 | schema_info=RuntimeSchemaInfo( |
| 281 | - static_kwargkey=["bias", "scale", "offset", "antiquant_scale", "antiquant_offset", "per_token_scale", | 337 | + static_kwargkey=[ |
| 282 | - "activation_input", "activation_quant_scale", "activation_quant_offset"], | 338 | + "bias", |
| 283 | - needs_pytree=True | 339 | + "scale", |
| 284 | - ) | 340 | + "offset", |
| 341 | + "antiquant_scale", | ||
| 342 | + "antiquant_offset", | ||
| 343 | + "per_token_scale", | ||
| 344 | + "activation_input", | ||
| 345 | + "activation_quant_scale", | ||
| 346 | + "activation_quant_offset", | ||
| 347 | + ], | ||
| 348 | + needs_pytree=True, | ||
| 349 | + ), | ||
| 285 | ) | 350 | ) |
| 286 | def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | 351 | def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: |
| 287 | # npu_grouped_matmul(Tensor[] x, Tensor[] weight, *, Tensor[]? bias=None, Tensor[]? scale=None, | 352 | # npu_grouped_matmul(Tensor[] x, Tensor[] weight, *, Tensor[]? bias=None, Tensor[]? scale=None, |
| @@ -295,28 +360,52 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 295 | x_num = len(x_src_strategy.children) | 360 | x_num = len(x_src_strategy.children) |
| 296 | weight_src_strategy: TupleStrategy = op_schema.args_schema[1] | 361 | weight_src_strategy: TupleStrategy = op_schema.args_schema[1] |
| 297 | weight_num = len(weight_src_strategy.children) | 362 | weight_num = len(weight_src_strategy.children) |
| 298 | - bias_src_strategy: Optional[Union[TupleStrategy, list]] = op_schema.kwargs_schema.get("bias", []) | 363 | + bias_src_strategy: TupleStrategy | list | None = op_schema.kwargs_schema.get( |
| 299 | - bias_num = len(bias_src_strategy.children) if isinstance(bias_src_strategy, TupleStrategy) else len(bias_src_strategy) | 364 | + "bias", [] |
| 300 | - group_list_num = 1 if ( | 365 | + ) |
| 301 | - op_schema.op == npu.npu_grouped_matmul.default and | 366 | + bias_num = ( |
| 302 | - op_schema.kwargs_schema.get("group_list", None) is not None | 367 | + len(bias_src_strategy.children) |
| 303 | - ) else 0 | 368 | + if isinstance(bias_src_strategy, TupleStrategy) |
| 369 | + else len(bias_src_strategy) | ||
| 370 | + ) | ||
| 371 | + group_list_num = ( | ||
| 372 | + 1 | ||
| 373 | + if ( | ||
| 374 | + op_schema.op == npu.npu_grouped_matmul.default | ||
| 375 | + and op_schema.kwargs_schema.get("group_list", None) is not None | ||
| 376 | + ) | ||
| 377 | + else 0 | ||
| 378 | + ) | ||
| 304 | split_item = op_schema.kwargs_schema.get("split_item", 0) | 379 | split_item = op_schema.kwargs_schema.get("split_item", 0) |
| 305 | - y_num = weight_num if split_item in (0, 1) else 1 # 0/1: multiple outputs, 2/3: single output | 380 | + y_num = ( |
| 381 | + weight_num if split_item in (0, 1) else 1 | ||
| 382 | + ) # 0/1: multiple outputs, 2/3: single output | ||
| 306 | 383 | ||
| 307 | strategies = [] | 384 | strategies = [] |
| 308 | 385 | ||
| 309 | all_replicate_strategy = [Replicate()] * y_num | 386 | all_replicate_strategy = [Replicate()] * y_num |
| 310 | - all_replicate_strategy.extend([Replicate()] * (len(op_schema.args_strategy) + len(op_schema.kwargs_strategy))) | 387 | + all_replicate_strategy.extend( |
| 388 | + [Replicate()] * (len(op_schema.args_strategy) + len(op_schema.kwargs_strategy)) | ||
| 389 | + ) | ||
| 311 | strategies.append(all_replicate_strategy) | 390 | strategies.append(all_replicate_strategy) |
| 312 | 391 | ||
| 313 | unsupported_arguments = [ | 392 | unsupported_arguments = [ |
| 314 | - "scale", "offset", "antiquant_scale", "antiquant_offset", "per_token_scale", # quant | 393 | + "scale", |
| 315 | - "activation_input", "activation_quant_scale", "activation_quant_offset", # reserved, unused now | 394 | + "offset", |
| 395 | + "antiquant_scale", | ||
| 396 | + "antiquant_offset", | ||
| 397 | + "per_token_scale", # quant | ||
| 398 | + "activation_input", | ||
| 399 | + "activation_quant_scale", | ||
| 400 | + "activation_quant_offset", # reserved, unused now | ||
| 316 | ] | 401 | ] |
| 317 | for key in unsupported_arguments: | 402 | for key in unsupported_arguments: |
| 318 | schema = op_schema.kwargs_schema.get(key, None) | 403 | schema = op_schema.kwargs_schema.get(key, None) |
| 319 | - if schema is not None and isinstance(schema, TupleStrategy) and len(schema.children) > 0: | 404 | + if ( |
| 405 | + schema is not None | ||
| 406 | + and isinstance(schema, TupleStrategy) | ||
| 407 | + and len(schema.children) > 0 | ||
| 408 | + ): | ||
| 320 | full_mesh_strategies = expand_to_full_mesh_op_strategy( | 409 | full_mesh_strategies = expand_to_full_mesh_op_strategy( |
| 321 | op_schema.get_mesh_from_args(), op_schema, strategies, input_index=y_num | 410 | op_schema.get_mesh_from_args(), op_schema, strategies, input_index=y_num |
| 322 | ) | 411 | ) |
| @@ -325,7 +414,9 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 325 | strategy.output_specs = [strategy.output_specs] | 414 | strategy.output_specs = [strategy.output_specs] |
| 326 | return full_mesh_strategies | 415 | return full_mesh_strategies |
| 327 | 416 | ||
| 328 | - if bias_num == 0: # if y is partial and bias exists, the bias will be added multiple times to the full tensor | 417 | + if ( |
| 418 | + bias_num == 0 | ||
| 419 | + ): # if y is partial and bias exists, the bias will be added multiple times to the full tensor | ||
| 329 | replicate_partial_strategy = [Partial()] * y_num | 420 | replicate_partial_strategy = [Partial()] * y_num |
| 330 | replicate_partial_strategy.extend([Replicate()] * x_num) | 421 | replicate_partial_strategy.extend([Replicate()] * x_num) |
| 331 | replicate_partial_strategy.extend([Partial()] * weight_num) | 422 | replicate_partial_strategy.extend([Partial()] * weight_num) |
| @@ -340,29 +431,35 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 340 | 431 | ||
| 341 | group_type = op_schema.kwargs_schema.get("group_type", None) | 432 | group_type = op_schema.kwargs_schema.get("group_type", None) |
| 342 | if group_type is not None and group_type > 0: | 433 | if group_type is not None and group_type > 0: |
| 343 | - raise NotImplementedError(f"npu_grouped_matmul does not support group_type={group_type} now.") | 434 | + raise NotImplementedError( |
| 435 | + f"npu_grouped_matmul does not support group_type={group_type} now." | ||
| 436 | + ) | ||
| 344 | 437 | ||
| 345 | - if x_num > 1 and weight_num > 1 and y_num > 1: # x_num, weight_num, y_num are equal | 438 | + if x_num > 1 and weight_num > 1 and y_num > 1: # x_num, weight_num, y_num are equal |
| 346 | pair_strategies = [] | 439 | pair_strategies = [] |
| 347 | # x: 2-6D, weight: 2D, weight: 1D (equals to weight.shape[1]) | 440 | # x: 2-6D, weight: 2D, weight: 1D (equals to weight.shape[1]) |
| 348 | # shard x | 441 | # shard x |
| 349 | x_ndim = x_src_strategy.children[0].ndim | 442 | x_ndim = x_src_strategy.children[0].ndim |
| 350 | for i in range(x_ndim - 1): | 443 | for i in range(x_ndim - 1): |
| 351 | - pair_strategies.append([Shard(i), Shard(i), Replicate(), Replicate()]) # y, x, weight, bias | 444 | + pair_strategies.append( |
| 445 | + [Shard(i), Shard(i), Replicate(), Replicate()] | ||
| 446 | + ) # y, x, weight, bias | ||
| 352 | # shard weight | 447 | # shard weight |
| 353 | pair_strategies.append([Shard(x_ndim - 1), Replicate(), Shard(1), Shard(0)]) | 448 | pair_strategies.append([Shard(x_ndim - 1), Replicate(), Shard(1), Shard(0)]) |
| 354 | # shard contracting dim | 449 | # shard contracting dim |
| 355 | if bias_num == 0: | 450 | if bias_num == 0: |
| 356 | pair_strategies.append([Partial(), Shard(x_ndim - 1), Shard(0), None]) | 451 | pair_strategies.append([Partial(), Shard(x_ndim - 1), Shard(0), None]) |
| 357 | # suppose that all pairs have the same shape and apply the same strategy | 452 | # suppose that all pairs have the same shape and apply the same strategy |
| 358 | - for (y_spec, x_spec, weight_spec, bias_spec) in pair_strategies: | 453 | + for y_spec, x_spec, weight_spec, bias_spec in pair_strategies: |
| 359 | strategy = [y_spec] * y_num | 454 | strategy = [y_spec] * y_num |
| 360 | strategy.extend([x_spec] * x_num) | 455 | strategy.extend([x_spec] * x_num) |
| 361 | strategy.extend([weight_spec] * weight_num) | 456 | strategy.extend([weight_spec] * weight_num) |
| 362 | strategy.extend([bias_spec] * bias_num) | 457 | strategy.extend([bias_spec] * bias_num) |
| 363 | strategy.extend([Replicate()] * group_list_num) | 458 | strategy.extend([Replicate()] * group_list_num) |
| 364 | strategies.append(strategy) | 459 | strategies.append(strategy) |
| 365 | - elif x_num == 1 and weight_num == 1 and y_num == 1: # npu_grouped_matmul.default only | 460 | + elif ( |
| 461 | + x_num == 1 and weight_num == 1 and y_num == 1 | ||
| 462 | + ): # npu_grouped_matmul.default only | ||
| 366 | # x: 2D, weight: 3D, bias: 2D, y: 2D, for each pair, define shape x: (m, k), weight: (k, n) | 463 | # x: 2D, weight: 3D, bias: 2D, y: 2D, for each pair, define shape x: (m, k), weight: (k, n) |
| 367 | if bias_num == 0: | 464 | if bias_num == 0: |
| 368 | k_shard_strategy = [Partial(), Shard(1), Shard(1)] | 465 | k_shard_strategy = [Partial(), Shard(1), Shard(1)] |
| @@ -372,7 +469,7 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 372 | n_shard_strategy.extend([Shard(1)] * bias_num) | 469 | n_shard_strategy.extend([Shard(1)] * bias_num) |
| 373 | n_shard_strategy.extend([Replicate()] * group_list_num) | 470 | n_shard_strategy.extend([Replicate()] * group_list_num) |
| 374 | strategies.append(n_shard_strategy) | 471 | strategies.append(n_shard_strategy) |
| 375 | - elif weight_num > 1: # x1wNy1, xNwNy1, x1wNyN | 472 | + elif weight_num > 1: # x1wNy1, xNwNy1, x1wNyN |
| 376 | # x: 2D, weight: 2D, bias: 1D, y: 2D | 473 | # x: 2D, weight: 2D, bias: 1D, y: 2D |
| 377 | if bias_num == 0: | 474 | if bias_num == 0: |
| 378 | k_shard_strategy = [Partial()] * y_num | 475 | k_shard_strategy = [Partial()] * y_num |
| @@ -387,8 +484,9 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 387 | n_shard_strategy.extend([Replicate()] * group_list_num) | 484 | n_shard_strategy.extend([Replicate()] * group_list_num) |
| 388 | strategies.append(n_shard_strategy) | 485 | strategies.append(n_shard_strategy) |
| 389 | 486 | ||
| 390 | - full_mesh_strategies = expand_to_full_mesh_op_strategy(op_schema.get_mesh_from_args(), op_schema, strategies, | 487 | + full_mesh_strategies = expand_to_full_mesh_op_strategy( |
| 391 | - input_index=y_num) | 488 | + op_schema.get_mesh_from_args(), op_schema, strategies, input_index=y_num |
| 489 | + ) | ||
| 392 | # output meta of npu_grouped_matmul is list, need convert output_spec here | 490 | # output meta of npu_grouped_matmul is list, need convert output_spec here |
| 393 | if y_num == 1: | 491 | if y_num == 1: |
| 394 | for strategy in full_mesh_strategies.strategies: | 492 | for strategy in full_mesh_strategies.strategies: |
| @@ -397,42 +495,51 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 397 | 495 | ||
| 398 | 496 | ||
| 399 | def _infer_npu_grouped_matmul_kwargs( | 497 | def _infer_npu_grouped_matmul_kwargs( |
| 400 | - op_schema: OpSchema, | 498 | + op_schema: OpSchema, output_sharding: OutputSharding |
| 401 | - output_sharding: OutputSharding | 499 | +) -> dict[str, DTensorSpec]: |
| 402 | -) -> Dict[str, DTensorSpec]: | ||
| 403 | output_spec = output_sharding.output_spec[0] | 500 | output_spec = output_sharding.output_spec[0] |
| 404 | kwargs_spec = {} | 501 | kwargs_spec = {} |
| 405 | for key, spec in op_schema.kwargs_schema.items(): | 502 | for key, spec in op_schema.kwargs_schema.items(): |
| 406 | - is_tensor_or_tenor_list_like = ( | 503 | + is_tensor_or_tenor_list_like = isinstance(spec, DTensorSpec) or ( |
| 407 | - isinstance(spec, DTensorSpec) or | 504 | + isinstance(spec, (list, tuple)) |
| 408 | - (isinstance(spec, (list, tuple)) and len(spec) > 0 and isinstance(spec[0], DTensorSpec)) | 505 | + and len(spec) > 0 |
| 506 | + and isinstance(spec[0], DTensorSpec) | ||
| 409 | ) | 507 | ) |
| 410 | if not is_tensor_or_tenor_list_like: | 508 | if not is_tensor_or_tenor_list_like: |
| 411 | kwargs_spec[key] = spec | 509 | kwargs_spec[key] = spec |
| 412 | continue | 510 | continue |
| 413 | 511 | ||
| 414 | - if key == 'group_list': # tensor | 512 | + if key == "group_list": # tensor |
| 415 | target_placement = [Replicate() for _ in output_spec.placements] | 513 | target_placement = [Replicate() for _ in output_spec.placements] |
| 416 | - kwargs_spec[key] = DTensorSpec(mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta) | 514 | + kwargs_spec[key] = DTensorSpec( |
| 515 | + mesh=spec.mesh, | ||
| 516 | + placements=target_placement, | ||
| 517 | + tensor_meta=spec.tensor_meta, | ||
| 518 | + ) | ||
| 417 | continue | 519 | continue |
| 418 | 520 | ||
| 419 | # tensor list | 521 | # tensor list |
| 420 | - if key == 'bias': | 522 | + if key == "bias": |
| 421 | target_placement = [ | 523 | target_placement = [ |
| 422 | Shard(0) if placement == Shard(output_spec.ndim - 1) else Replicate() | 524 | Shard(0) if placement == Shard(output_spec.ndim - 1) else Replicate() |
| 423 | for placement in output_spec.placements | 525 | for placement in output_spec.placements |
| 424 | ] | 526 | ] |
| 425 | - else: # unsupported sharding keys | 527 | + else: # unsupported sharding keys |
| 426 | target_placement = [Replicate() for _ in output_spec.placements] | 528 | target_placement = [Replicate() for _ in output_spec.placements] |
| 427 | - kwargs_spec[key] = [DTensorSpec(mesh=e.mesh, placements=target_placement, tensor_meta=e.tensor_meta) for e in spec] | 529 | + kwargs_spec[key] = [ |
| 530 | + DTensorSpec( | ||
| 531 | + mesh=e.mesh, placements=target_placement, tensor_meta=e.tensor_meta | ||
| 532 | + ) | ||
| 533 | + for e in spec | ||
| 534 | + ] | ||
| 428 | 535 | ||
| 429 | return kwargs_spec | 536 | return kwargs_spec |
| 430 | 537 | ||
| 431 | 538 | ||
| 432 | def _npu_grouped_matmul_handler( | 539 | def _npu_grouped_matmul_handler( |
| 433 | - op_call: torch._ops.OpOverload, | 540 | + op_call: torch._ops.OpOverload, |
| 434 | - args: Tuple[object, ...], | 541 | + args: tuple[object, ...], |
| 435 | - kwargs: Dict[str, object], | 542 | + kwargs: dict[str, object], |
| 436 | ) -> object: | 543 | ) -> object: |
| 437 | # extract local tensor and sharding infos to a OpInfo | 544 | # extract local tensor and sharding infos to a OpInfo |
| 438 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) | 545 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) |
| @@ -448,7 +555,9 @@ def _npu_grouped_matmul_handler( | |||
| 448 | if participating: | 555 | if participating: |
| 449 | # computation that happens in the current rank of the mesh, normal case | 556 | # computation that happens in the current rank of the mesh, normal case |
| 450 | local_args = get_redistributed_local_args(op_info, output_sharding) | 557 | local_args = get_redistributed_local_args(op_info, output_sharding) |
| 451 | - local_kwargs = get_redistributed_local_kwargs(_infer_npu_grouped_matmul_kwargs, op_info, output_sharding) | 558 | + local_kwargs = get_redistributed_local_kwargs( |
| 559 | + _infer_npu_grouped_matmul_kwargs, op_info, output_sharding | ||
| 560 | + ) | ||
| 452 | local_results = op_call(*local_args, **local_kwargs) | 561 | local_results = op_call(*local_args, **local_kwargs) |
| 453 | else: | 562 | else: |
| 454 | # For a non-participating device (happens on rank that does not belong to the device mesh), | 563 | # For a non-participating device (happens on rank that does not belong to the device mesh), |
| @@ -459,14 +568,28 @@ def _npu_grouped_matmul_handler( | |||
| 459 | 568 | ||
| 460 | 569 | ||
| 461 | 570 | ||
| 462 | -def npu_all_gather_base_mm_strategy(x1, x2, hcom, world_size, bias=None, x1_scale=None, x2_scale=None, gather_index=0, | 571 | +def npu_all_gather_base_mm_strategy( |
| 463 | - gather_output=True, comm_turn=0, output_dtype=None, comm_mode=None): | 572 | + x1, |
| 573 | + x2, | ||
| 574 | + hcom, | ||
| 575 | + world_size, | ||
| 576 | + bias=None, | ||
| 577 | + x1_scale=None, | ||
| 578 | + x2_scale=None, | ||
| 579 | + gather_index=0, | ||
| 580 | + gather_output=True, | ||
| 581 | + comm_turn=0, | ||
| 582 | + output_dtype=None, | ||
| 583 | + comm_mode=None, | ||
| 584 | +): | ||
| 464 | # npu_all_gather_base_mm(Tensor input, Tensor x2, str hcom, int world_size, *, Tensor? bias=None, | 585 | # npu_all_gather_base_mm(Tensor input, Tensor x2, str hcom, int world_size, *, Tensor? bias=None, |
| 465 | # Tensor? x1_scale=None, Tensor? x2_scale=None, int gather_index=0, bool gather_output=True, | 586 | # Tensor? x1_scale=None, Tensor? x2_scale=None, int gather_index=0, bool gather_output=True, |
| 466 | # int comm_turn=0, ScalarType? output_dtype=None, str? comm_mode=None) -> (Tensor, Tensor) | 587 | # int comm_turn=0, ScalarType? output_dtype=None, str? comm_mode=None) -> (Tensor, Tensor) |
| 467 | # op only support gather_index=0(i.e. allgather x1) now | 588 | # op only support gather_index=0(i.e. allgather x1) now |
| 468 | if gather_index != 0: | 589 | if gather_index != 0: |
| 469 | - raise NotImplementedError(f"npu_all_gather_base_mm only support gather_index=0 now, but got {gather_index}.") | 590 | + raise NotImplementedError( |
| 591 | + f"npu_all_gather_base_mm only support gather_index=0 now, but got {gather_index}." | ||
| 592 | + ) | ||
| 470 | 593 | ||
| 471 | # formula: output = allgather(x1)@x2 + bias | 594 | # formula: output = allgather(x1)@x2 + bias |
| 472 | # for all gather, x1: S(0) -> R | 595 | # for all gather, x1: S(0) -> R |
| @@ -477,37 +600,55 @@ def npu_all_gather_base_mm_strategy(x1, x2, hcom, world_size, bias=None, x1_scal | |||
| 477 | strategies = [] | 600 | strategies = [] |
| 478 | sharding_strategy_S0R = ( | 601 | sharding_strategy_S0R = ( |
| 479 | [ | 602 | [ |
| 480 | - Replicate(), # output | 603 | + Replicate(), # output |
| 481 | - Replicate() # gather_out | 604 | + Replicate(), # gather_out |
| 482 | ], | 605 | ], |
| 483 | [ | 606 | [ |
| 484 | - Shard(0), # x1 | 607 | + Shard(0), # x1 |
| 485 | - Replicate(), # x2 | 608 | + Replicate(), # x2 |
| 486 | - None, # hcom | 609 | + None, # hcom |
| 487 | - None, # world_size | 610 | + None, # world_size |
| 488 | - None if bias is None else Replicate(), # bias, global shape(n * world_size,) | 611 | + None |
| 489 | - None if x1_scale is None else Shard(0), # x1_scale follow x1, global shape(m * world_size, 1) | 612 | + if bias is None |
| 490 | - None if x2_scale is None else Replicate(), # x2_scale follow x2, global shape(1, n * world_size) | 613 | + else Replicate(), # bias, global shape(n * world_size,) |
| 491 | - None, None, None, None, None # gather_index, gather_output, comm_turn, output_dtype, comm_mode | 614 | + None |
| 492 | - ] | 615 | + if x1_scale is None |
| 616 | + else Shard(0), # x1_scale follow x1, global shape(m * world_size, 1) | ||
| 617 | + None | ||
| 618 | + if x2_scale is None | ||
| 619 | + else Replicate(), # x2_scale follow x2, global shape(1, n * world_size) | ||
| 620 | + None, | ||
| 621 | + None, | ||
| 622 | + None, | ||
| 623 | + None, | ||
| 624 | + None, # gather_index, gather_output, comm_turn, output_dtype, comm_mode | ||
| 625 | + ], | ||
| 493 | ) | 626 | ) |
| 494 | strategies.append(sharding_strategy_S0R) | 627 | strategies.append(sharding_strategy_S0R) |
| 495 | 628 | ||
| 496 | sharding_strategy_S0S1 = ( | 629 | sharding_strategy_S0S1 = ( |
| 497 | [ | 630 | [ |
| 498 | - Shard(1), # output | 631 | + Shard(1), # output |
| 499 | - Replicate() # gather_out | 632 | + Replicate(), # gather_out |
| 500 | ], | 633 | ], |
| 501 | [ | 634 | [ |
| 502 | - Shard(0), # x1 | 635 | + Shard(0), # x1 |
| 503 | - Shard(1), # x2 | 636 | + Shard(1), # x2 |
| 504 | - None, # hcom | 637 | + None, # hcom |
| 505 | - None, # world_size | 638 | + None, # world_size |
| 506 | - None if bias is None else Shard(0), # bias, global shape(n * world_size,) | 639 | + None if bias is None else Shard(0), # bias, global shape(n * world_size,) |
| 507 | - None if x1_scale is None else Shard(0), # x1_scale follow x1, global shape(m * world_size, 1) | 640 | + None |
| 508 | - None if x2_scale is None else Shard(1), # x2_scale follow x2, global shape(1, n * world_size) | 641 | + if x1_scale is None |
| 509 | - None, None, None, None, None # gather_index, gather_output, comm_turn, output_dtype, comm_mode | 642 | + else Shard(0), # x1_scale follow x1, global shape(m * world_size, 1) |
| 510 | - ] | 643 | + None |
| 644 | + if x2_scale is None | ||
| 645 | + else Shard(1), # x2_scale follow x2, global shape(1, n * world_size) | ||
| 646 | + None, | ||
| 647 | + None, | ||
| 648 | + None, | ||
| 649 | + None, | ||
| 650 | + None, # gather_index, gather_output, comm_turn, output_dtype, comm_mode | ||
| 651 | + ], | ||
| 511 | ) | 652 | ) |
| 512 | strategies.append(sharding_strategy_S0S1) | 653 | strategies.append(sharding_strategy_S0S1) |
| 513 | 654 | ||
| @@ -515,9 +656,8 @@ def npu_all_gather_base_mm_strategy(x1, x2, hcom, world_size, bias=None, x1_scal | |||
| 515 | 656 | ||
| 516 | 657 | ||
| 517 | def _infer_npu_all_gather_base_mm_kwargs( | 658 | def _infer_npu_all_gather_base_mm_kwargs( |
| 518 | - op_schema: OpSchema, | 659 | + op_schema: OpSchema, output_sharding: OutputSharding |
| 519 | - output_sharding: OutputSharding | 660 | +) -> dict[str, DTensorSpec]: |
| 520 | -) -> Dict[str, DTensorSpec]: | ||
| 521 | output_spec = output_sharding.output_spec[0] | 661 | output_spec = output_sharding.output_spec[0] |
| 522 | kwargs_spec = {} | 662 | kwargs_spec = {} |
| 523 | for key, spec in op_schema.kwargs_schema.items(): | 663 | for key, spec in op_schema.kwargs_schema.items(): |
| @@ -528,33 +668,48 @@ def _infer_npu_all_gather_base_mm_kwargs( | |||
| 528 | target_placement = [] | 668 | target_placement = [] |
| 529 | for placement in output_spec.placements: | 669 | for placement in output_spec.placements: |
| 530 | if placement == Replicate(): | 670 | if placement == Replicate(): |
| 531 | - if key == 'x1_scale': | 671 | + if key == "x1_scale": |
| 532 | target_placement.append(Shard(0)) | 672 | target_placement.append(Shard(0)) |
| 533 | - else: # bias, x2_scale | 673 | + else: # bias, x2_scale |
| 534 | target_placement.append(Replicate()) | 674 | target_placement.append(Replicate()) |
| 535 | elif placement == Shard(1): | 675 | elif placement == Shard(1): |
| 536 | - if key == 'x2_scale': | 676 | + if key == "x2_scale": |
| 537 | target_placement.append(Shard(1)) | 677 | target_placement.append(Shard(1)) |
| 538 | - else: # bias, x1_scale | 678 | + else: # bias, x1_scale |
| 539 | target_placement.append(Shard(0)) | 679 | target_placement.append(Shard(0)) |
| 540 | else: | 680 | else: |
| 541 | raise ValueError( | 681 | raise ValueError( |
| 542 | f"Unexpected output placement {placement} for npu_all_gather_base_mm." | 682 | f"Unexpected output placement {placement} for npu_all_gather_base_mm." |
| 543 | ) | 683 | ) |
| 544 | - kwargs_spec[key] = DTensorSpec(mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta) | 684 | + kwargs_spec[key] = DTensorSpec( |
| 685 | + mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta | ||
| 686 | + ) | ||
| 545 | 687 | ||
| 546 | return kwargs_spec | 688 | return kwargs_spec |
| 547 | 689 | ||
| 548 | 690 | ||
| 549 | 691 | ||
| 550 | -def npu_mm_reduce_scatter_base_strategy(x1, x2, hcom, world_size, reduce_op='sum', bias=None, x1_scale=None, | 692 | +def npu_mm_reduce_scatter_base_strategy( |
| 551 | - x2_scale=None, comm_turn=0, output_dtype=None, comm_mode=None): | 693 | + x1, |
| 694 | + x2, | ||
| 695 | + hcom, | ||
| 696 | + world_size, | ||
| 697 | + reduce_op="sum", | ||
| 698 | + bias=None, | ||
| 699 | + x1_scale=None, | ||
| 700 | + x2_scale=None, | ||
| 701 | + comm_turn=0, | ||
| 702 | + output_dtype=None, | ||
| 703 | + comm_mode=None, | ||
| 704 | +): | ||
| 552 | # npu_mm_reduce_scatter_base(Tensor self, Tensor x2, str hcom, int world_size, *, str reduce_op='sum', | 705 | # npu_mm_reduce_scatter_base(Tensor self, Tensor x2, str hcom, int world_size, *, str reduce_op='sum', |
| 553 | # Tensor? bias=None, Tensor? x1_scale=None, Tensor? x2_scale=None, int comm_turn=0, | 706 | # Tensor? bias=None, Tensor? x1_scale=None, Tensor? x2_scale=None, int comm_turn=0, |
| 554 | # ScalarType? output_dtype=None, str? comm_mode=None) -> Tensor | 707 | # ScalarType? output_dtype=None, str? comm_mode=None) -> Tensor |
| 555 | # op only support reduce_op='sum' now | 708 | # op only support reduce_op='sum' now |
| 556 | - if reduce_op != 'sum': | 709 | + if reduce_op != "sum": |
| 557 | - raise NotImplementedError(f"npu_mm_reduce_scatter_base only support reduce_op='sum' now, but got {reduce_op}.") | 710 | + raise NotImplementedError( |
| 711 | + f"npu_mm_reduce_scatter_base only support reduce_op='sum' now, but got {reduce_op}." | ||
| 712 | + ) | ||
| 558 | 713 | ||
| 559 | # formula: output = reducescatter(x1@x2 + bias) | 714 | # formula: output = reducescatter(x1@x2 + bias) |
| 560 | # for reduce_scatter, local_output: P -> S(0) | 715 | # for reduce_scatter, local_output: P -> S(0) |
| @@ -566,16 +721,22 @@ def npu_mm_reduce_scatter_base_strategy(x1, x2, hcom, world_size, reduce_op='sum | |||
| 566 | Shard(0) # output | 721 | Shard(0) # output |
| 567 | ], | 722 | ], |
| 568 | [ | 723 | [ |
| 569 | - Shard(1), # x1 | 724 | + Shard(1), # x1 |
| 570 | - Shard(0), # x2 | 725 | + Shard(0), # x2 |
| 571 | - None, # hcom | 726 | + None, # hcom |
| 572 | - None, # world_size | 727 | + None, # world_size |
| 573 | - None, # reduce_op | 728 | + None, # reduce_op |
| 574 | - None if bias is None else Shard(0), # bias, global shape(n * world_size,) | 729 | + None if bias is None else Shard(0), # bias, global shape(n * world_size,) |
| 575 | - None if x1_scale is None else Shard(1), # x1_scale follow x1, global shape(m, world_size) | 730 | + None |
| 576 | - None if x2_scale is None else Shard(0), # x2_scale follow x2, global shape(world_size, n) | 731 | + if x1_scale is None |
| 577 | - None, None, None # comm_turn, output_dtype, comm_mode | 732 | + else Shard(1), # x1_scale follow x1, global shape(m, world_size) |
| 578 | - ] | 733 | + None |
| 734 | + if x2_scale is None | ||
| 735 | + else Shard(0), # x2_scale follow x2, global shape(world_size, n) | ||
| 736 | + None, | ||
| 737 | + None, | ||
| 738 | + None, # comm_turn, output_dtype, comm_mode | ||
| 739 | + ], | ||
| 579 | ) | 740 | ) |
| 580 | strategies.append(sharding_strategy_S1S0) | 741 | strategies.append(sharding_strategy_S1S0) |
| 581 | 742 | ||
| @@ -583,9 +744,8 @@ def npu_mm_reduce_scatter_base_strategy(x1, x2, hcom, world_size, reduce_op='sum | |||
| 583 | 744 | ||
| 584 | 745 | ||
| 585 | def _infer_npu_mm_reduce_scatter_base_kwargs( | 746 | def _infer_npu_mm_reduce_scatter_base_kwargs( |
| 586 | - op_schema: OpSchema, | 747 | + op_schema: OpSchema, output_sharding: OutputSharding |
| 587 | - output_sharding: OutputSharding | 748 | +) -> dict[str, DTensorSpec]: |
| 588 | -) -> Dict[str, DTensorSpec]: | ||
| 589 | output_spec = output_sharding.output_spec | 749 | output_spec = output_sharding.output_spec |
| 590 | kwargs_spec = {} | 750 | kwargs_spec = {} |
| 591 | for key, spec in op_schema.kwargs_schema.items(): | 751 | for key, spec in op_schema.kwargs_schema.items(): |
| @@ -596,23 +756,25 @@ def _infer_npu_mm_reduce_scatter_base_kwargs( | |||
| 596 | target_placement = [] | 756 | target_placement = [] |
| 597 | for placement in output_spec.placements: | 757 | for placement in output_spec.placements: |
| 598 | if placement == Shard(0): | 758 | if placement == Shard(0): |
| 599 | - if key == 'x1_scale': | 759 | + if key == "x1_scale": |
| 600 | target_placement.append(Shard(1)) | 760 | target_placement.append(Shard(1)) |
| 601 | - else: # bias, x2_scale | 761 | + else: # bias, x2_scale |
| 602 | target_placement.append(Shard(0)) | 762 | target_placement.append(Shard(0)) |
| 603 | else: | 763 | else: |
| 604 | raise ValueError( | 764 | raise ValueError( |
| 605 | f"Unexpected output placement {placement} for npu_mm_reduce_scatter_base." | 765 | f"Unexpected output placement {placement} for npu_mm_reduce_scatter_base." |
| 606 | ) | 766 | ) |
| 607 | - kwargs_spec[key] = DTensorSpec(mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta) | 767 | + kwargs_spec[key] = DTensorSpec( |
| 768 | + mesh=spec.mesh, placements=target_placement, tensor_meta=spec.tensor_meta | ||
| 769 | + ) | ||
| 608 | 770 | ||
| 609 | return kwargs_spec | 771 | return kwargs_spec |
| 610 | 772 | ||
| 611 | 773 | ||
| 612 | def npu_comm_mm_fusion_handler( | 774 | def npu_comm_mm_fusion_handler( |
| 613 | - op_call: torch._ops.OpOverload, | 775 | + op_call: torch._ops.OpOverload, |
| 614 | - args: Tuple[object, ...], | 776 | + args: tuple[object, ...], |
| 615 | - kwargs: Dict[str, object], | 777 | + kwargs: dict[str, object], |
| 616 | ) -> object: | 778 | ) -> object: |
| 617 | # extract local tensor and sharding infos to a OpInfo | 779 | # extract local tensor and sharding infos to a OpInfo |
| 618 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) | 780 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) |
| @@ -633,11 +795,15 @@ def npu_comm_mm_fusion_handler( | |||
| 633 | new_shape[dim] = new_shape[dim] // world_size | 795 | new_shape[dim] = new_shape[dim] // world_size |
| 634 | elif op_call == npu.npu_mm_reduce_scatter_base.default: | 796 | elif op_call == npu.npu_mm_reduce_scatter_base.default: |
| 635 | new_shape[dim] = new_shape[dim] * world_size | 797 | new_shape[dim] = new_shape[dim] * world_size |
| 636 | - return TensorMeta(shape=torch.Size(new_shape), stride=tensor_meta.stride, dtype=tensor_meta.dtype) | 798 | + return TensorMeta( |
| 799 | + shape=torch.Size(new_shape), | ||
| 800 | + stride=tensor_meta.stride, | ||
| 801 | + dtype=tensor_meta.dtype, | ||
| 802 | + ) | ||
| 637 | 803 | ||
| 638 | if op_call == npu.npu_all_gather_base_mm.default: | 804 | if op_call == npu.npu_all_gather_base_mm.default: |
| 639 | world_size = args[3] | 805 | world_size = args[3] |
| 640 | - for spec in output_sharding.output_spec: # output, gather_out | 806 | + for spec in output_sharding.output_spec: # output, gather_out |
| 641 | spec.tensor_meta = get_output_meta(spec.tensor_meta, 0, world_size) | 807 | spec.tensor_meta = get_output_meta(spec.tensor_meta, 0, world_size) |
| 642 | elif op_call == npu.npu_mm_reduce_scatter_base.default: | 808 | elif op_call == npu.npu_mm_reduce_scatter_base.default: |
| 643 | world_size = args[3] | 809 | world_size = args[3] |
| @@ -668,9 +834,7 @@ def npu_comm_mm_fusion_handler( | |||
| 668 | return DTensor._op_dispatcher.wrap(local_results, output_sharding.output_spec) | 834 | return DTensor._op_dispatcher.wrap(local_results, output_sharding.output_spec) |
| 669 | 835 | ||
| 670 | 836 | ||
| 671 | -@register_op_strategy( | 837 | +@register_op_strategy([npu.npu_apply_adam_w.default, npu.npu_apply_adam_w.out]) |
| 672 | - [npu.npu_apply_adam_w.default, npu.npu_apply_adam_w.out] | ||
| 673 | -) | ||
| 674 | def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: | 838 | def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: |
| 675 | # npu_apply_adam_w( | 839 | # npu_apply_adam_w( |
| 676 | # Scalar beta1_power, Scalar beta2_power, Scalar lr, Scalar weight_decay, Scalar beta1, Scalar beta2, | 840 | # Scalar beta1_power, Scalar beta2_power, Scalar lr, Scalar weight_decay, Scalar beta1, Scalar beta2, |
| @@ -679,8 +843,10 @@ def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 679 | grad_arg_index = 7 | 843 | grad_arg_index = 7 |
| 680 | max_gard_norm_arg_index = 8 | 844 | max_gard_norm_arg_index = 8 |
| 681 | grad_strategy: OpStrategy = op_schema.args_schema[grad_arg_index] | 845 | grad_strategy: OpStrategy = op_schema.args_schema[grad_arg_index] |
| 682 | - if "out" in op_schema.kwargs_schema.keys(): | 846 | + if "out" in op_schema.kwargs_schema: |
| 683 | - grad_spec: DTensorSpec = op_schema.kwargs_schema["out"].children[0].strategies[0].output_spec | 847 | + grad_spec: DTensorSpec = ( |
| 848 | + op_schema.kwargs_schema["out"].children[0].strategies[0].output_spec | ||
| 849 | + ) | ||
| 684 | else: | 850 | else: |
| 685 | grad_spec: DTensorSpec = grad_strategy.strategies[0].output_spec | 851 | grad_spec: DTensorSpec = grad_strategy.strategies[0].output_spec |
| 686 | input_target_specs = [] | 852 | input_target_specs = [] |
| @@ -703,20 +869,20 @@ def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 703 | 869 | ||
| 704 | output_spec = [] | 870 | output_spec = [] |
| 705 | for k, values in op_schema.kwargs_schema.items(): | 871 | for k, values in op_schema.kwargs_schema.items(): |
| 706 | - if k == 'out': | 872 | + if k == "out": |
| 707 | for v in values.children: | 873 | for v in values.children: |
| 708 | output_spec.append(v.strategies[0].output_spec) | 874 | output_spec.append(v.strategies[0].output_spec) |
| 709 | - output_strategy = OpStrategy([ | 875 | + output_strategy = OpStrategy( |
| 710 | - OpSpec(output_specs=tuple(output_spec), input_specs=input_target_specs) | 876 | + [OpSpec(output_specs=tuple(output_spec), input_specs=input_target_specs)] |
| 711 | - ]) | 877 | + ) |
| 712 | 878 | ||
| 713 | return output_strategy | 879 | return output_strategy |
| 714 | 880 | ||
| 715 | 881 | ||
| 716 | def _npu_apply_adam_w_handler( | 882 | def _npu_apply_adam_w_handler( |
| 717 | - op_call: torch._ops.OpOverload, | 883 | + op_call: torch._ops.OpOverload, |
| 718 | - args: Tuple[object, ...], | 884 | + args: tuple[object, ...], |
| 719 | - kwargs: Dict[str, object], | 885 | + kwargs: dict[str, object], |
| 720 | ) -> object: | 886 | ) -> object: |
| 721 | # extract local tensor and sharding infos to a OpInfo | 887 | # extract local tensor and sharding infos to a OpInfo |
| 722 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) | 888 | op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs) |
| @@ -738,12 +904,12 @@ def _npu_apply_adam_w_handler( | |||
| 738 | output_sharding.use_val_from_redistribute_schema, | 904 | output_sharding.use_val_from_redistribute_schema, |
| 739 | ) | 905 | ) |
| 740 | local_args = ( | 906 | local_args = ( |
| 741 | - pytree.tree_unflatten( | 907 | + pytree.tree_unflatten( |
| 742 | - cast(list[object], op_info.local_args), op_info.args_tree_spec | 908 | + cast(list[object], op_info.local_args), op_info.args_tree_spec |
| 743 | - ) | ||
| 744 | - if op_info.args_tree_spec | ||
| 745 | - else op_info.local_args | ||
| 746 | ) | 909 | ) |
| 910 | + if op_info.args_tree_spec | ||
| 911 | + else op_info.local_args | ||
| 912 | + ) | ||
| 747 | 913 | ||
| 748 | local_results = torch_npu.npu_apply_adam_w(*local_args, **op_info.local_kwargs) | 914 | local_results = torch_npu.npu_apply_adam_w(*local_args, **op_info.local_kwargs) |
| 749 | 915 | ||
| @@ -756,7 +922,7 @@ def _npu_apply_adam_w_handler( | |||
| 756 | out_dts = [] | 922 | out_dts = [] |
| 757 | spec_idx = 0 | 923 | spec_idx = 0 |
| 758 | for argument in op_call._schema.arguments: | 924 | for argument in op_call._schema.arguments: |
| 759 | - if argument.name == 'out': | 925 | + if argument.name == "out": |
| 760 | for value in kwargs[argument.name]: | 926 | for value in kwargs[argument.name]: |
| 761 | out_dt = cast(DTensor, value) | 927 | out_dt = cast(DTensor, value) |
| 762 | out_dt._spec = cast(DTensorSpec, output_specs[spec_idx]) | 928 | out_dt._spec = cast(DTensorSpec, output_specs[spec_idx]) |
| @@ -783,16 +949,17 @@ def custom_dropout_strategy(op_schema: OpSchema): | |||
| 783 | output_target_specs = [] | 949 | output_target_specs = [] |
| 784 | output_target_specs.append(input_spec) | 950 | output_target_specs.append(input_spec) |
| 785 | output_target_specs.append( | 951 | output_target_specs.append( |
| 786 | - DTensorSpec( | 952 | + DTensorSpec(mesh=input_spec.mesh, placements=[Shard(0)]) |
| 787 | - mesh=input_spec.mesh, | ||
| 788 | - placements=[Shard(0)] | ||
| 789 | - ) | ||
| 790 | ) | 953 | ) |
| 791 | input_target_specs = [] | 954 | input_target_specs = [] |
| 792 | input_target_specs.append(input_spec) | 955 | input_target_specs.append(input_spec) |
| 793 | - output_strategy = OpStrategy([ | 956 | + output_strategy = OpStrategy( |
| 794 | - OpSpec(output_specs=output_target_specs, input_specs=input_target_specs) | 957 | + [ |
| 795 | - ]) | 958 | + OpSpec( |
| 959 | + output_specs=output_target_specs, input_specs=input_target_specs | ||
| 960 | + ) | ||
| 961 | + ] | ||
| 962 | + ) | ||
| 796 | return output_strategy | 963 | return output_strategy |
| 797 | 964 | ||
| 798 | replicate_strategy = [Replicate(), Replicate(), Replicate()] | 965 | replicate_strategy = [Replicate(), Replicate(), Replicate()] |
| @@ -814,19 +981,30 @@ def custom_dropout_backward_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 814 | if isinstance(spec, OpStrategy): | 981 | if isinstance(spec, OpStrategy): |
| 815 | input_target_specs.append(spec.strategies[0].output_spec) | 982 | input_target_specs.append(spec.strategies[0].output_spec) |
| 816 | 983 | ||
| 817 | - output_strategy = OpStrategy([ | 984 | + output_strategy = OpStrategy( |
| 818 | - OpSpec(output_specs=op_schema.args_schema[0].strategies[0].output_spec, input_specs=input_target_specs) | 985 | + [ |
| 819 | - ]) | 986 | + OpSpec( |
| 987 | + output_specs=op_schema.args_schema[0].strategies[0].output_spec, | ||
| 988 | + input_specs=input_target_specs, | ||
| 989 | + ) | ||
| 990 | + ] | ||
| 991 | + ) | ||
| 820 | 992 | ||
| 821 | return output_strategy | 993 | return output_strategy |
| 822 | 994 | ||
| 823 | 995 | ||
| 996 | + | ||
| 997 | +def custom_bmm_strategy(op_schema: OpSchema): | ||
| 998 | + mesh = op_schema.get_mesh_from_args() | ||
| 999 | + return _mm_like_strategy("bmk,bkn->bmn", mesh, op_schema) | ||
| 1000 | + | ||
| 1001 | + | ||
| 824 | customized_ops = { | 1002 | customized_ops = { |
| 825 | npu.npu_grouped_matmul.default: _npu_grouped_matmul_handler, | 1003 | npu.npu_grouped_matmul.default: _npu_grouped_matmul_handler, |
| 826 | npu.npu_grouped_matmul.List: _npu_grouped_matmul_handler, | 1004 | npu.npu_grouped_matmul.List: _npu_grouped_matmul_handler, |
| 827 | npu.npu_apply_adam_w.out: _npu_apply_adam_w_handler, | 1005 | npu.npu_apply_adam_w.out: _npu_apply_adam_w_handler, |
| 828 | npu.npu_all_gather_base_mm.default: npu_comm_mm_fusion_handler, | 1006 | npu.npu_all_gather_base_mm.default: npu_comm_mm_fusion_handler, |
| 829 | - npu.npu_mm_reduce_scatter_base.default: npu_comm_mm_fusion_handler | 1007 | + npu.npu_mm_reduce_scatter_base.default: npu_comm_mm_fusion_handler, |
| 830 | } | 1008 | } |
| 831 | 1009 | ||
| 832 | old_handlers = DTensor._op_dispatcher._custom_op_handlers | 1010 | 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.utils 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.utils 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)) | ||
Dtorch_npu/utils/dtensor.py+0-50
| @@ -1,50 +0,0 @@ | |||
| 1 | -import torch | ||
| 2 | -from torch.distributed.tensor._ops._common_rules import pointwise_rule | ||
| 3 | -from torch.distributed.tensor._ops.utils import register_prop_rule, normalize_dims | ||
| 4 | -from torch.distributed.tensor._ops._matrix_ops import bmm_strategy | ||
| 5 | -from torch.distributed.tensor._ops._view_ops import ( | ||
| 6 | - register_op_strategy_map, | ||
| 7 | - dim_maps, | ||
| 8 | - InputDim | ||
| 9 | -) | ||
| 10 | -import torch_npu | ||
| 11 | - | ||
| 12 | -__all__ = [] | ||
| 13 | - | ||
| 14 | - | ||
| 15 | -def _register_ops_under_dtensor_rules(): | ||
| 16 | - npu = torch.ops.npu | ||
| 17 | - aten = torch.ops.aten | ||
| 18 | - | ||
| 19 | - pointwise_ops = [ | ||
| 20 | - # please keep the entries below alphabetically sorted | ||
| 21 | - # native ops | ||
| 22 | - aten.isclose.default, | ||
| 23 | - aten.isfinite.default, | ||
| 24 | - # custom ops | ||
| 25 | - npu.fast_gelu.default, | ||
| 26 | - npu.npu_fast_gelu.default, | ||
| 27 | - npu.npu_layer_norm_eval.default, | ||
| 28 | - # backward point-wise ops | ||
| 29 | - # please keep the entries below alphabetically sorted | ||
| 30 | - npu.npu_fast_gelu_backward.default | ||
| 31 | - ] | ||
| 32 | - | ||
| 33 | - matrix_ops = [ | ||
| 34 | - npu.npu_bmmV2.default | ||
| 35 | - ] | ||
| 36 | - # pointwise rule | ||
| 37 | - for op in pointwise_ops: | ||
| 38 | - register_prop_rule(op)(pointwise_rule) | ||
| 39 | - | ||
| 40 | - # bmm rules | ||
| 41 | - for op in matrix_ops: | ||
| 42 | - register_prop_rule(op)(bmm_strategy) | ||
| 43 | - | ||
| 44 | - # reshape_prop under view_ops | ||
| 45 | - dim_maps.update({ | ||
| 46 | - torch_npu.npu_transpose: lambda input, dims: tuple( | ||
| 47 | - InputDim(i) for i in normalize_dims(dims, input.ndim) | ||
| 48 | - ) | ||
| 49 | - }) | ||
| 50 | - register_op_strategy_map(npu.npu_transpose.default, torch_npu.npu_transpose) | ||