modify test case-pytorch-AtomGit

已合并

modify test case #29184

gcw_Dgfy2aKk创建于 1月7日从gcw_Dgfy2aKk/pytorch： correct_testcase_260107 合入到 Ascend/pytorch： master

modify test case #29184

已合并

gcw_Dgfy2aKk创建于 1月7日

从gcw_Dgfy2aKk/pytorch correct_testcase_260107合入到Ascend/pytorch master

共 3 个文件变更+11-16

Mtest/distributed/tensor/test_dtensor_ops.py+2-1

	@@ -3,10 +3,11 @@ from torch.distributed._tensor import distribute_tensor, Replicate
3	from torch_npu.testing.testcase import run_tests	3	from torch_npu.testing.testcase import run_tests
4	from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase	4	from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
5	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	5	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		6	+from torch_npu.testing.common_utils import SupportedDevices
6		7
7		8
8	class TestDTensorOps(NPUDTensorTestBase):	9	class TestDTensorOps(NPUDTensorTestBase):
9	- @skipIfUnsupportMultiNPU(4)	10	+ @skipIfUnsupportMultiNPU(2)
10	@with_comms	11	@with_comms
11	def test_torch_nn_functional_one_hot(self):	12	def test_torch_nn_functional_one_hot(self):
12	"""Test torch.nn.functional.one_hot with DTensor on NPU"""	13	"""Test torch.nn.functional.one_hot with DTensor on NPU"""

Mtest/distributed/tensor/test_sharded_optim.py+9-11

	@@ -27,7 +27,7 @@ class MyShardedModel(torch.nn.Module):
27	if spec is not None:	27	if spec is not None:
28	self.sharded_param = torch.nn.Parameter(	28	self.sharded_param = torch.nn.Parameter(
29	sharded_tensor.rand(	29	sharded_tensor.rand(
30	- spec, 20, 10, requires_grad=True, process_group=group	30	+ spec, 10, 10, requires_grad=True, process_group=group
31	)	31	)
32	)	32	)
33	else:	33	else:
	@@ -59,8 +59,6 @@ class MyShardedLinear(torch.nn.Module):
59	placements=[	59	placements=[
60	"rank:0/npu:0",	60	"rank:0/npu:0",
61	"rank:1/npu:1",	61	"rank:1/npu:1",
62	- "rank:2/npu:2",
63	- "rank:3/npu:3",
64	],	62	],
65	)	63	)
66		64
	@@ -69,8 +67,6 @@ class MyShardedLinear(torch.nn.Module):
69	placements=[	67	placements=[
70	"rank:0/npu:0",	68	"rank:0/npu:0",
71	"rank:1/npu:1",	69	"rank:1/npu:1",
72	- "rank:2/npu:2",
73	- "rank:3/npu:3",
74	],	70	],
75	)	71	)
76		72
	@@ -80,10 +76,16 @@ class MyShardedLinear(torch.nn.Module):
80	def forward(self, inp):	76	def forward(self, inp):
81	return self.linear2(self.gelu(self.linear1(inp)))	77	return self.linear2(self.gelu(self.linear1(inp)))
82		78
		79	+Test_GPU_NUM = 2
		80	+
83		81
84	class TestShardedOptimizer(ShardedTensorTestBase):	82	class TestShardedOptimizer(ShardedTensorTestBase):
		83	+ @property
		84	+ def world_size(self):
		85	+ return Test_GPU_NUM
		86	+
85	@with_comms(init_rpc=False)	87	@with_comms(init_rpc=False)
86	- @skip_if_lt_x_gpu(4)	88	+ @skip_if_lt_x_gpu(2)
87	@requires_nccl()	89	@requires_nccl()
88	def test_sharded_optim(self):	90	def test_sharded_optim(self):
89	rowwise_spec = ChunkShardingSpec(	91	rowwise_spec = ChunkShardingSpec(
	@@ -91,8 +93,6 @@ class TestShardedOptimizer(ShardedTensorTestBase):
91	placements=[	93	placements=[
92	"rank:0/npu:0",	94	"rank:0/npu:0",
93	"rank:1/npu:1",	95	"rank:1/npu:1",
94	- "rank:2/npu:2",
95	- "rank:3/npu:3",
96	],	96	],
97	)	97	)
98	local_model = MyShardedModel().npu()	98	local_model = MyShardedModel().npu()
	@@ -142,7 +142,7 @@ class TestShardedOptimizer(ShardedTensorTestBase):
142	self.assertEqual(new_val, local_model.param)	142	self.assertEqual(new_val, local_model.param)
143		143
144	@with_comms(init_rpc=False)	144	@with_comms(init_rpc=False)
145	- @skip_if_lt_x_gpu(4)	145	+ @skip_if_lt_x_gpu(2)
146	@requires_nccl()	146	@requires_nccl()
147	def test_named_params_with_sharded_tensor(self):	147	def test_named_params_with_sharded_tensor(self):
148	rowwise_spec = ChunkShardingSpec(	148	rowwise_spec = ChunkShardingSpec(
	@@ -150,8 +150,6 @@ class TestShardedOptimizer(ShardedTensorTestBase):
150	placements=[	150	placements=[
151	"rank:0/npu:0",	151	"rank:0/npu:0",
152	"rank:1/npu:1",	152	"rank:1/npu:1",
153	- "rank:2/npu:2",
154	- "rank:3/npu:3",
155	],	153	],
156	)	154	)
157	sharded_model = MyShardedModel(spec=rowwise_spec).npu()	155	sharded_model = MyShardedModel(spec=rowwise_spec).npu()

Mtorch_npu/distributed/tensor/_sharded_tensor_patch.py+0-4

	@@ -51,10 +51,6 @@ def _patched_sharded_tensor_npu(
51	"`torch.preserve_format` is supported!"	51	"`torch.preserve_format` is supported!"
52	)	52	)
53		53
54	- # Validate and process target device - only NPU device without index is allowed
55	- if device is not None:
56	- device = torch.device(device) if isinstance(device, str) else device
57	- raise ValueError("Only device without device id (e.g. 'npu') is expected for ShardedTensor!")
58		54
59	current_device = torch.device(torch.npu.current_device())	55	current_device = torch.device(torch.npu.current_device())
60	# returns a copy of ShardedTensor on NPU current device	56	# returns a copy of ShardedTensor on NPU current device