已合并
fix error #29221
gcw_Dgfy2aKk创建于 1月7日
fix error #29221
已合并
共 4 个文件变更+274-1
| @@ -6,7 +6,7 @@ from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMult | |||
| 6 | 6 | ||
| 7 | 7 | ||
| 8 | class TestDTensorOps(NPUDTensorTestBase): | 8 | class TestDTensorOps(NPUDTensorTestBase): |
| 9 | - @skipIfUnsupportMultiNPU(4) | 9 | + @skipIfUnsupportMultiNPU(2) |
| 10 | 10 | ||
| 11 | def test_torch_nn_functional_one_hot(self): | 11 | def test_torch_nn_functional_one_hot(self): |
| 12 | """Test torch.nn.functional.one_hot with DTensor on NPU""" | 12 | """Test torch.nn.functional.one_hot with DTensor on NPU""" |
| @@ -0,0 +1,175 @@ | |||
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates | ||
| 2 | +# Owner(s): ["oncall: distributed"] | ||
| 3 | +from copy import deepcopy | ||
| 4 | + | ||
| 5 | +import torch | ||
| 6 | +import torch.optim as optim | ||
| 7 | +from torch.distributed._shard import shard_parameter, sharded_tensor | ||
| 8 | +from torch.distributed._shard.sharded_optim import ShardedOptimizer | ||
| 9 | +from torch.distributed._shard.sharding_spec import ChunkShardingSpec | ||
| 10 | +from torch.testing._internal.common_distributed import requires_nccl, skip_if_lt_x_gpu | ||
| 11 | +from torch.testing._internal.common_utils import run_tests | ||
| 12 | +from torch.testing._internal.distributed._shard.sharded_tensor import ( | ||
| 13 | + ShardedTensorTestBase, | ||
| 14 | + with_comms, | ||
| 15 | +) | ||
| 16 | + | ||
| 17 | +import torch_npu | ||
| 18 | +from torch_npu.contrib import transfer_to_npu | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +class MyShardedModel(torch.nn.Module): | ||
| 22 | + def __init__(self, spec=None, group=None): | ||
| 23 | + super().__init__() | ||
| 24 | + # Use same seed. | ||
| 25 | + torch.manual_seed(0) | ||
| 26 | + self.param = torch.nn.Parameter(torch.rand(5, 10)) | ||
| 27 | + if spec is not None: | ||
| 28 | + self.sharded_param = torch.nn.Parameter( | ||
| 29 | + sharded_tensor.rand( | ||
| 30 | + spec, 10, 10, requires_grad=True, process_group=group | ||
| 31 | + ) | ||
| 32 | + ) | ||
| 33 | + else: | ||
| 34 | + self.sharded_param = torch.nn.Parameter(torch.rand(5, 10)) | ||
| 35 | + | ||
| 36 | + def forward(self, x): | ||
| 37 | + if isinstance(self.sharded_param, sharded_tensor.ShardedTensor): | ||
| 38 | + return self.param + self.sharded_param.local_shards()[0].tensor + x | ||
| 39 | + else: | ||
| 40 | + return self.sharded_param + self.param + x | ||
| 41 | + | ||
| 42 | + | ||
| 43 | +class MyShardedLinear(torch.nn.Module): | ||
| 44 | + def __init__(self, rank=None): | ||
| 45 | + super().__init__() | ||
| 46 | + # Use same seed. | ||
| 47 | + torch.manual_seed(0) | ||
| 48 | + self.linear1 = torch.nn.Linear(17, 12) | ||
| 49 | + self.linear2 = torch.nn.Linear(12, 29) | ||
| 50 | + self.gelu = torch.nn.GELU() | ||
| 51 | + | ||
| 52 | + if rank: | ||
| 53 | + self.linear1.npu(rank) | ||
| 54 | + self.linear2.npu(rank) | ||
| 55 | + | ||
| 56 | + def shard_parameter(self): | ||
| 57 | + rowwise_sharding_spec = ChunkShardingSpec( | ||
| 58 | + dim=0, | ||
| 59 | + placements=[ | ||
| 60 | + "rank:0/npu:0", | ||
| 61 | + "rank:1/npu:1", | ||
| 62 | + ], | ||
| 63 | + ) | ||
| 64 | + | ||
| 65 | + colwise_sharding_spec = ChunkShardingSpec( | ||
| 66 | + dim=1, | ||
| 67 | + placements=[ | ||
| 68 | + "rank:0/npu:0", | ||
| 69 | + "rank:1/npu:1", | ||
| 70 | + ], | ||
| 71 | + ) | ||
| 72 | + | ||
| 73 | + shard_parameter(self.linear1, "weight", rowwise_sharding_spec) | ||
| 74 | + shard_parameter(self.linear2, "weight", colwise_sharding_spec) | ||
| 75 | + | ||
| 76 | + def forward(self, inp): | ||
| 77 | + return self.linear2(self.gelu(self.linear1(inp))) | ||
| 78 | + | ||
| 79 | +Test_GPU_NUM = 2 | ||
| 80 | + | ||
| 81 | + | ||
| 82 | +class TestShardedOptimizer(ShardedTensorTestBase): | ||
| 83 | + | ||
| 84 | + def world_size(self): | ||
| 85 | + return Test_GPU_NUM | ||
| 86 | + | ||
| 87 | + | ||
| 88 | + | ||
| 89 | + | ||
| 90 | + def test_sharded_optim(self): | ||
| 91 | + rowwise_spec = ChunkShardingSpec( | ||
| 92 | + dim=0, | ||
| 93 | + placements=[ | ||
| 94 | + "rank:0/npu:0", | ||
| 95 | + "rank:1/npu:1", | ||
| 96 | + ], | ||
| 97 | + ) | ||
| 98 | + local_model = MyShardedModel().npu() | ||
| 99 | + sharded_model = MyShardedModel(spec=rowwise_spec).npu() | ||
| 100 | + | ||
| 101 | + # copy the parameters from local model | ||
| 102 | + sharded_model.sharded_param.local_shards()[0].tensor = ( | ||
| 103 | + local_model.sharded_param.detach().clone().requires_grad_() | ||
| 104 | + ) | ||
| 105 | + | ||
| 106 | + local_optim = optim.SGD(local_model.parameters(), lr=0.1) | ||
| 107 | + sharded_model_params = dict(sharded_model.named_parameters()) | ||
| 108 | + sharded_optim = ShardedOptimizer(sharded_model_params, optim.SGD, lr=0.1) | ||
| 109 | + | ||
| 110 | + local_optim.zero_grad() | ||
| 111 | + sharded_optim.zero_grad() | ||
| 112 | + | ||
| 113 | + before_update = deepcopy(sharded_optim.named_params) | ||
| 114 | + | ||
| 115 | + inp = torch.rand([5, 10]).cuda(self.rank).requires_grad_() | ||
| 116 | + | ||
| 117 | + # run forward | ||
| 118 | + local_output = local_model(inp) | ||
| 119 | + sharded_output = sharded_model(inp) | ||
| 120 | + # backward | ||
| 121 | + local_output.sum().backward() | ||
| 122 | + sharded_output.sum().backward() | ||
| 123 | + | ||
| 124 | + # optimizer update | ||
| 125 | + local_optim.step() | ||
| 126 | + sharded_optim.step() | ||
| 127 | + | ||
| 128 | + # make sure the parameters (including sharded param) | ||
| 129 | + # get updated by the optimizer, and the updated | ||
| 130 | + # local params are the same as the sharded params | ||
| 131 | + for key, val in before_update.items(): | ||
| 132 | + new_val = sharded_optim.named_params[key] | ||
| 133 | + if isinstance(val, sharded_tensor.ShardedTensor): | ||
| 134 | + self.assertNotEqual( | ||
| 135 | + val.local_shards()[0].tensor, new_val.local_shards()[0].tensor | ||
| 136 | + ) | ||
| 137 | + self.assertEqual( | ||
| 138 | + new_val.local_shards()[0].tensor, local_model.sharded_param | ||
| 139 | + ) | ||
| 140 | + else: | ||
| 141 | + self.assertNotEqual(val, new_val) | ||
| 142 | + self.assertEqual(new_val, local_model.param) | ||
| 143 | + | ||
| 144 | + | ||
| 145 | + | ||
| 146 | + | ||
| 147 | + def test_named_params_with_sharded_tensor(self): | ||
| 148 | + rowwise_spec = ChunkShardingSpec( | ||
| 149 | + dim=0, | ||
| 150 | + placements=[ | ||
| 151 | + "rank:0/npu:0", | ||
| 152 | + "rank:1/npu:1", | ||
| 153 | + ], | ||
| 154 | + ) | ||
| 155 | + sharded_model = MyShardedModel(spec=rowwise_spec).npu() | ||
| 156 | + sharded_model_params = dict(sharded_model.named_parameters()) | ||
| 157 | + param_keys = list(sharded_model_params.keys()) | ||
| 158 | + self.assertEqual(len(param_keys), 2) | ||
| 159 | + self.assertTrue("param" in param_keys) | ||
| 160 | + self.assertTrue("sharded_param" in param_keys) | ||
| 161 | + | ||
| 162 | + sharded_linear = MyShardedLinear(rank=self.rank).npu() | ||
| 163 | + sharded_linear.shard_parameter() | ||
| 164 | + sharded_linear_params = dict(sharded_linear.named_parameters()) | ||
| 165 | + param_keys = list(sharded_linear_params.keys()) | ||
| 166 | + self.assertEqual(len(param_keys), 4) | ||
| 167 | + self.assertTrue("linear1.bias" in param_keys) | ||
| 168 | + self.assertTrue("linear2.bias" in param_keys) | ||
| 169 | + self.assertTrue("linear1.weight" in param_keys) | ||
| 170 | + self.assertTrue("linear2.weight" in param_keys) | ||
| 171 | + self.assertFalse("bias" in param_keys) | ||
| 172 | + | ||
| 173 | + | ||
| 174 | +if __name__ == "__main__": | ||
| 175 | + run_tests() | ||
| @@ -3,3 +3,4 @@ import torch_npu.distributed.tensor._matrix_ops | |||
| 3 | import torch_npu.distributed.tensor._attention | 3 | import torch_npu.distributed.tensor._attention |
| 4 | import torch_npu.distributed.tensor._math_ops | 4 | import torch_npu.distributed.tensor._math_ops |
| 5 | import torch_npu.distributed.tensor._moe_ops | 5 | import torch_npu.distributed.tensor._moe_ops |
| 6 | +import torch_npu.distributed.tensor._sharded_tensor_patch | ||
| @@ -0,0 +1,97 @@ | |||
| 1 | +import copy | ||
| 2 | + | ||
| 3 | +import torch | ||
| 4 | +from torch.distributed._shard.sharded_tensor import ShardedTensor | ||
| 5 | +from torch.distributed._shard.sharded_tensor.shard import Shard | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +def _patched_sharded_tensor_npu( | ||
| 9 | + self, | ||
| 10 | + device=None, | ||
| 11 | + non_blocking=False, | ||
| 12 | + memory_format=torch.preserve_format, | ||
| 13 | + process_group=None, | ||
| 14 | +) -> ShardedTensor: | ||
| 15 | + """ | ||
| 16 | + Returns a copy of this object in NPU memory, if the original ShardedTensor | ||
| 17 | + is on CPU, we will move the local shard to the current NPU device of each | ||
| 18 | + process in a SPMD fashion. | ||
| 19 | + If this ShardedTensor is already on NPU memory and local shards on each rank are | ||
| 20 | + already on current device, we still returns a new ShardedTensor object with new | ||
| 21 | + metadata, but no underlying data movements are performed. | ||
| 22 | + | ||
| 23 | + .. note:: When moving a ShardedTensor from CPU to NPU, the ShardedTensor might | ||
| 24 | + need to be managed by a different type of ProcessGroup that is compatible | ||
| 25 | + with NPU, it is the user's responsibility to explicitly pass in a new | ||
| 26 | + process_group that is compatible with NPU. | ||
| 27 | + | ||
| 28 | + Args: | ||
| 29 | + device (torch.device/str, optional): Target NPU device (only "npu" without index is supported). | ||
| 30 | + Defaults to None, which uses current NPU device. | ||
| 31 | + non_blocking (bool, optional): If True, the copy is asynchronous with respect to the host. | ||
| 32 | + Defaults to False. | ||
| 33 | + memory_format (torch.memory_format, optional): Desired memory format of the tensor after move. | ||
| 34 | + Only torch.preserve_format or torch.contiguous_format is supported. Defaults to torch.preserve_format. | ||
| 35 | + process_group (ProcessGroup, optional): The process group to use for the new ShardedTensor. | ||
| 36 | + Defaults to None, which uses the original process group of the ShardedTensor. | ||
| 37 | + | ||
| 38 | + Returns: | ||
| 39 | + ShardedTensor: A new ShardedTensor instance with all local shards moved to NPU device | ||
| 40 | + | ||
| 41 | + Raises: | ||
| 42 | + RuntimeError: If memory_format is not torch.preserve_format or torch.contiguous_format | ||
| 43 | + ValueError: If device is specified but not a valid NPU device without index | ||
| 44 | + """ | ||
| 45 | + if ( | ||
| 46 | + memory_format != torch.preserve_format | ||
| 47 | + and memory_format != torch.contiguous_format | ||
| 48 | + ): | ||
| 49 | + raise RuntimeError( | ||
| 50 | + "Only `torch.contiguous_format` or " | ||
| 51 | + "`torch.preserve_format` is supported!" | ||
| 52 | + ) | ||
| 53 | + | ||
| 54 | + current_device = torch.device(torch.npu.current_device()) | ||
| 55 | + # returns a copy of ShardedTensor on NPU current device | ||
| 56 | + list_shards: list[Shard] = [] | ||
| 57 | + # move all local shards to current device, and change metadata | ||
| 58 | + # if local shards already on the current device, there's no | ||
| 59 | + # real data movement, only the metadata are copied. | ||
| 60 | + for shard in self._local_shards: | ||
| 61 | + npu_tensor = shard.tensor.npu( | ||
| 62 | + device=current_device, | ||
| 63 | + non_blocking=non_blocking, | ||
| 64 | + memory_format=memory_format, | ||
| 65 | + ) # type: ignore[call-arg] | ||
| 66 | + metadata = copy.deepcopy(shard.metadata) | ||
| 67 | + metadata.placement._device = current_device # type: ignore[union-attr] | ||
| 68 | + | ||
| 69 | + list_shards.append(Shard(npu_tensor, metadata)) | ||
| 70 | + | ||
| 71 | + st_meta = copy.deepcopy(self.metadata()) | ||
| 72 | + for meta in st_meta.shards_metadata: | ||
| 73 | + if meta.placement.device().type != "npu": # type: ignore[union-attr] | ||
| 74 | + meta.placement._device = current_device # type: ignore[union-attr] | ||
| 75 | + | ||
| 76 | + pg = self._process_group if process_group is None else process_group | ||
| 77 | + # we need to use `init_from_local_shards` to communicate between ranks | ||
| 78 | + # and update the sharding spec/shards metadata. | ||
| 79 | + st_npu = ShardedTensor._init_from_local_shards_and_global_metadata( | ||
| 80 | + list_shards, | ||
| 81 | + sharded_tensor_metadata=st_meta, | ||
| 82 | + process_group=pg, | ||
| 83 | + init_rrefs=self._init_rrefs, | ||
| 84 | + ) | ||
| 85 | + return st_npu | ||
| 86 | + | ||
| 87 | + | ||
| 88 | +def _apply_sharded_tensor_npu_patch(): | ||
| 89 | + """ | ||
| 90 | + Applies the NPU patch to ShardedTensor by adding/overriding the npu() method. | ||
| 91 | + """ | ||
| 92 | + if not hasattr(ShardedTensor, "npu"): | ||
| 93 | + ShardedTensor.npu = _patched_sharded_tensor_npu | ||
| 94 | + | ||
| 95 | + | ||
| 96 | +# Execute the patch application when the module is imported | ||
| 97 | +_apply_sharded_tensor_npu_patch() | ||