[fix] with_comms can't set self.device in 2.8.0+ bugfix-pytorch-AtomGit

Mtest/distributed/_tensor/test_api.py+7-8

	@@ -9,11 +9,10 @@ from torch.distributed._tensor import (
9	Shard,	9	Shard,
10	)	10	)
11	from torch.testing._internal.common_utils import run_tests	11	from torch.testing._internal.common_utils import run_tests
12	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
13	-
14		12
15	import torch_npu	13	import torch_npu
16	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	14	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		15	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
17		16
18		17
19	class MyModel(nn.Module):	18	class MyModel(nn.Module):
	@@ -31,15 +30,15 @@ class MyModel(nn.Module):
31	m.reset_parameters()	30	m.reset_parameters()
32		31
33		32
34	-class DTensorAPITest(DTensorTestBase):	33	+class DTensorAPITest(NPUDTensorTestBase):
35	@property	34	@property
36	def world_size(self) -> int:	35	def world_size(self) -> int:
37	# hard code world size to 4 as we need to test	36	# hard code world size to 4 as we need to test
38	# at least with 2d mesh	37	# at least with 2d mesh
39	return 4	38	return 4
40		39
41	- @with_comms
42	@skipIfUnsupportMultiNPU(4)	40	@skipIfUnsupportMultiNPU(4)
		41	+ @with_comms
43	def test_distribute_tensor(self):	42	def test_distribute_tensor(self):
44	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	43	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
45	shard_spec = [Shard(0)]	44	shard_spec = [Shard(0)]
	@@ -56,8 +55,8 @@ class DTensorAPITest(DTensorTestBase):
56	self.assertTrue(dist_tensor.requires_grad)	55	self.assertTrue(dist_tensor.requires_grad)
57	self.assertTrue(dist_tensor.is_leaf)	56	self.assertTrue(dist_tensor.is_leaf)
58		57
59	- @with_comms
60	@skipIfUnsupportMultiNPU(4)	58	@skipIfUnsupportMultiNPU(4)
		59	+ @with_comms
61	def test_distribute_tensor_uneven_sharding(self):	60	def test_distribute_tensor_uneven_sharding(self):
62	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	61	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
63	input_sizes_and_shard_dims = [	62	input_sizes_and_shard_dims = [
	@@ -79,8 +78,8 @@ class DTensorAPITest(DTensorTestBase):
79	local_tensor = dist_tensor.to_local()	78	local_tensor = dist_tensor.to_local()
80	self.assertEqual(local_tensor, splitted_tensor_list[self.rank])	79	self.assertEqual(local_tensor, splitted_tensor_list[self.rank])
81		80
82	- @with_comms
83	@skipIfUnsupportMultiNPU(4)	81	@skipIfUnsupportMultiNPU(4)
		82	+ @with_comms
84	def test_distribute_module(self):	83	def test_distribute_module(self):
85	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	84	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
86	# fully shard all linear modules on dim 0	85	# fully shard all linear modules on dim 0
	@@ -143,8 +142,8 @@ class DTensorAPITest(DTensorTestBase):
143	else:	142	else:
144	self.assertEqual(param.placements, replica_spec)	143	self.assertEqual(param.placements, replica_spec)
145		144
146	- @with_comms
147	@skipIfUnsupportMultiNPU(4)	145	@skipIfUnsupportMultiNPU(4)
		146	+ @with_comms
148	def test_distribute_module_input_fn_output_fn(self):	147	def test_distribute_module_input_fn_output_fn(self):
149	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	148	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
150		149
	@@ -188,8 +187,8 @@ class DTensorAPITest(DTensorTestBase):
188	self.assertTrue(isinstance(param_grad, DTensor))	187	self.assertTrue(isinstance(param_grad, DTensor))
189	self.assertTrue(isinstance(param_grad.placements[0], Replicate))	188	self.assertTrue(isinstance(param_grad.placements[0], Replicate))
190		189
191	- @with_comms
192	@skipIfUnsupportMultiNPU(4)	190	@skipIfUnsupportMultiNPU(4)
		191	+ @with_comms
193	def test_distribute_module_meta(self):	192	def test_distribute_module_meta(self):
194	# If the model is too big, the user may first the create entire model on the meta device and then initialize	193	# If the model is too big, the user may first the create entire model on the meta device and then initialize
195	# it on the device in the partition function.	194	# it on the device in the partition function.

Mtest/distributed/_tensor/test_basic_strategy.py+3-3

	@@ -1,15 +1,15 @@
1	import torch	1	import torch
2	from torch.distributed._tensor import DeviceMesh	2	from torch.distributed._tensor import DeviceMesh
3	-from torch.distributed._tensor.ops._einsum_strategy import (	3	+from torch.distributed.tensor._ops._einsum_strategy import (
4	EinsumDims,	4	EinsumDims,
5	gen_einsum_strategies,	5	gen_einsum_strategies,
6	)	6	)
7		7
8	from torch.testing._internal.common_utils import run_tests, TestCase	8	from torch.testing._internal.common_utils import run_tests, TestCase
9	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
10		9
11	import torch_npu	10	import torch_npu
12	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	11	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		12	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
13		13
14		14
15	class TestEinsumDims(TestCase):	15	class TestEinsumDims(TestCase):
	@@ -76,7 +76,7 @@ class TestEinsumDims(TestCase):
76	self.assertEqual(edims.rhs_out_only_dims, ["f"])	76	self.assertEqual(edims.rhs_out_only_dims, ["f"])
77		77
78		78
79	-class TestEinsumStrategies(DTensorTestBase):	79	+class TestEinsumStrategies(NPUDTensorTestBase):
80	@property	80	@property
81	def world_size(self) -> int:	81	def world_size(self) -> int:
82	return 4	82	return 4

Mtest/distributed/_tensor/test_common_rules.py+2-5

	@@ -4,17 +4,14 @@ from torch.distributed._tensor.placement_types import DTensorSpec, TensorMeta
4	from torch.distributed.tensor._op_schema import OpSchema	4	from torch.distributed.tensor._op_schema import OpSchema
5	from torch.distributed.tensor._ops._common_rules import einop_rule, pointwise_rule	5	from torch.distributed.tensor._ops._common_rules import einop_rule, pointwise_rule
6	from torch.testing._internal.common_utils import run_tests	6	from torch.testing._internal.common_utils import run_tests
7	-from torch.testing._internal.distributed._tensor.common_dtensor import (
8	- DTensorTestBase,
9	- with_comms,
10	-)
11	import torch_npu	7	import torch_npu
12	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	8	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		9	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
13		10
14	aten = torch.ops.aten	11	aten = torch.ops.aten
15		12
16		13
17	-class CommonRulesTest(DTensorTestBase):	14	+class CommonRulesTest(NPUDTensorTestBase):
18	@property	15	@property
19	def world_size(self) -> int:	16	def world_size(self) -> int:
20	# hard code world size to 4 as we need to test	17	# hard code world size to 4 as we need to test

Mtest/distributed/_tensor/test_dtensor.py+8-10

	@@ -13,7 +13,7 @@ from torch.distributed._tensor import (
13	DTensor,	13	DTensor,
14	init_device_mesh,	14	init_device_mesh,
15	)	15	)
16	-from torch.distributed._tensor.placement_types import _Partial, Replicate, Shard	16	+from torch.distributed.tensor.placement_types import Partial, Replicate, Shard
17	from torch.distributed.tensor.parallel import (	17	from torch.distributed.tensor.parallel import (
18	ColwiseParallel,	18	ColwiseParallel,
19	parallelize_module,	19	parallelize_module,
	@@ -21,12 +21,10 @@ from torch.distributed.tensor.parallel import (
21	)	21	)
22		22
23	from torch.testing._internal.common_utils import run_tests	23	from torch.testing._internal.common_utils import run_tests
24	-from torch.testing._internal.distributed._tensor.common_dtensor import (
25	- DTensorTestBase,
26	-)
27		24
28	import torch_npu	25	import torch_npu
29	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	26	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		27	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
30		28
31		29
32	class DummyMLP(torch.nn.Module):	30	class DummyMLP(torch.nn.Module):
	@@ -47,7 +45,7 @@ class DummyMLP(torch.nn.Module):
47	self.net2.bias.fill_(1.2)	45	self.net2.bias.fill_(1.2)
48		46
49		47
50	-class DTensorTest(DTensorTestBase):	48	+class DTensorTest(NPUDTensorTestBase):
51	@skipIfUnsupportMultiNPU(4)	49	@skipIfUnsupportMultiNPU(4)
52	@with_comms	50	@with_comms
53	def test_dtensor_constructor(self):	51	def test_dtensor_constructor(self):
	@@ -177,7 +175,7 @@ class DTensorTest(DTensorTestBase):
177	ddp_tensor = DTensor.from_local(local_tensor, device_mesh, replica_spec)	175	ddp_tensor = DTensor.from_local(local_tensor, device_mesh, replica_spec)
178	self.assertEqual(ddp_tensor.size(), local_tensor.size())	176	self.assertEqual(ddp_tensor.size(), local_tensor.size())
179		177
180	- partial_spec = [_Partial()]	178	+ partial_spec = [Partial()]
181	partial_tensor = DTensor.from_local(local_tensor, device_mesh, partial_spec)	179	partial_tensor = DTensor.from_local(local_tensor, device_mesh, partial_spec)
182	self.assertEqual(partial_tensor.size(), local_tensor.size())	180	self.assertEqual(partial_tensor.size(), local_tensor.size())
183		181
	@@ -336,7 +334,7 @@ class DTensorTest(DTensorTestBase):
336		334
337	sharded_dtensor = distribute_tensor(global_tensor, device_mesh, placements)	335	sharded_dtensor = distribute_tensor(global_tensor, device_mesh, placements)
338	local_out = sharded_dtensor.redistribute(placements=[Replicate()]).to_local(	336	local_out = sharded_dtensor.redistribute(placements=[Replicate()]).to_local(
339	- grad_placements=[_Partial()]	337	+ grad_placements=[Partial()]
340	)	338	)
341	local_out.sum().backward()	339	local_out.sum().backward()
342		340
	@@ -363,7 +361,7 @@ class DTensorTest(DTensorTestBase):
363	global_tensor = torch.ones(8, 3, requires_grad=True)	361	global_tensor = torch.ones(8, 3, requires_grad=True)
364		362
365	sharded_dtensor = distribute_tensor(global_tensor, device_mesh, placements)	363	sharded_dtensor = distribute_tensor(global_tensor, device_mesh, placements)
366	- local_out = sharded_dtensor.full_tensor(grad_placements=[_Partial()])	364	+ local_out = sharded_dtensor.full_tensor(grad_placements=[Partial()])
367	local_out.sum().backward()	365	local_out.sum().backward()
368		366
369	replica_grad = sharded_dtensor.grad.full_tensor()	367	replica_grad = sharded_dtensor.grad.full_tensor()
	@@ -522,7 +520,7 @@ class DTensorTest(DTensorTestBase):
522	self.assertEqual(sharded_tensor, reloaded_st)	520	self.assertEqual(sharded_tensor, reloaded_st)
523		521
524		522
525	-class DTensorMeshTest(DTensorTestBase):	523	+class DTensorMeshTest(NPUDTensorTestBase):
526	@property	524	@property
527	def world_size(self):	525	def world_size(self):
528	return 8	526	return 8
	@@ -762,7 +760,7 @@ class DTensorMeshTest(DTensorTestBase):
762	)	760	)
763		761
764		762
765	-class TestDTensorPlacementTypes(DTensorTestBase):	763	+class TestDTensorPlacementTypes(NPUDTensorTestBase):
766	@property	764	@property
767	def world_size(self):	765	def world_size(self):
768	return 8	766	return 8

Mtest/distributed/_tensor/test_dtensor_compile.py+3-3

	@@ -30,7 +30,6 @@ from torch.testing._internal.common_utils import (
30	run_tests,	30	run_tests,
31	)	31	)
32	from torch.testing._internal.distributed._tensor.common_dtensor import (	32	from torch.testing._internal.distributed._tensor.common_dtensor import (
33	- DTensorTestBase,
34	MLPModule,	33	MLPModule,
35	)	34	)
36	from torch.testing._internal.distributed.fake_pg import FakeStore	35	from torch.testing._internal.distributed.fake_pg import FakeStore
	@@ -40,6 +39,7 @@ from torch._dynamo.backends.common import aot_autograd
40		39
41	import torch_npu	40	import torch_npu
42	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	41	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		42	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
43		43
44		44
45	class SimpleModel(nn.Module):	45	class SimpleModel(nn.Module):
	@@ -70,7 +70,7 @@ aot_eager_graph = aot_autograd(
70	)	70	)
71		71
72		72
73	-class TestDTensorCompile(DTensorTestBase):	73	+class TestDTensorCompile(NPUDTensorTestBase):
74	@property	74	@property
75	def world_size(self) -> int:	75	def world_size(self) -> int:
76	return 2	76	return 2
	@@ -170,7 +170,7 @@ class TestDTensorCompile(DTensorTestBase):
170	self.assertEqual(res, ref)	170	self.assertEqual(res, ref)
171		171
172		172
173	-class TestDTensorCompileE2E(DTensorTestBase):	173	+class TestDTensorCompileE2E(NPUDTensorTestBase):
174	@property	174	@property
175	def world_size(self):	175	def world_size(self):
176	return 4	176	return 4

Mtest/distributed/_tensor/test_dtensor_custom_ops.py+2-2

	@@ -6,13 +6,13 @@ from torch.distributed._tensor import DeviceMesh, distribute_tensor
6	from torch.distributed._tensor.api import DTensor	6	from torch.distributed._tensor.api import DTensor
7	from torch.distributed._tensor.placement_types import Replicate, Shard	7	from torch.distributed._tensor.placement_types import Replicate, Shard
8	from torch.testing._internal.common_utils import run_tests	8	from torch.testing._internal.common_utils import run_tests
9	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
10		9
11	import torch_npu	10	import torch_npu
12	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	11	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		12	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
13		13
14		14
15	-class TestDTensorCustomOps(DTensorTestBase):	15	+class TestDTensorCustomOps(NPUDTensorTestBase):
16	@property	16	@property
17	def world_size(self):	17	def world_size(self):
18	# hard code world size to 4 as we need to test	18	# hard code world size to 4 as we need to test

Mtest/distributed/_tensor/test_embedding_ops.py+2-2

	@@ -4,16 +4,16 @@ import torch
4	from torch.distributed._tensor import distribute_tensor, DTensor, DeviceMesh	4	from torch.distributed._tensor import distribute_tensor, DTensor, DeviceMesh
5	from torch.distributed._tensor.placement_types import Replicate, Shard	5	from torch.distributed._tensor.placement_types import Replicate, Shard
6	from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN	6	from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
7	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
8		7
9	import torch_npu	8	import torch_npu
10	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	9	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		10	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
11		11
12	if TEST_WITH_DEV_DBG_ASAN:	12	if TEST_WITH_DEV_DBG_ASAN:
13	raise RuntimeError("Skip dev-asan as torch + multiprocessing spawn have known issues")	13	raise RuntimeError("Skip dev-asan as torch + multiprocessing spawn have known issues")
14		14
15		15
16	-class TestEmbeddingOp(DTensorTestBase):	16	+class TestEmbeddingOp(NPUDTensorTestBase):
17	def _run_embedding_op_test(	17	def _run_embedding_op_test(
18	self,	18	self,
19	shard_dim,	19	shard_dim,

Mtest/distributed/_tensor/test_init.py+3-3

	@@ -1,13 +1,13 @@
1	import torch	1	import torch
2	from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard, zeros	2	from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard, zeros
3	from torch.testing._internal.common_utils import run_tests	3	from torch.testing._internal.common_utils import run_tests
4	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
5		4
6	import torch_npu	5	import torch_npu
7	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	6	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		7	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
8		8
9		9
10	-class DTensorInitOpsTest(DTensorTestBase):	10	+class DTensorInitOpsTest(NPUDTensorTestBase):
11	def _run_init_op(self, init_op, args, *kwargs):	11	def _run_init_op(self, init_op, args, *kwargs):
12	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	12	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
13	shard_spec = [Shard(0)]	13	shard_spec = [Shard(0)]
	@@ -28,7 +28,7 @@ class DTensorInitOpsTest(DTensorTestBase):
28	self._run_init_op(torch.nn.init.constant_, 2.4)	28	self._run_init_op(torch.nn.init.constant_, 2.4)
29		29
30		30
31	-class DTensorConstructorTest(DTensorTestBase):	31	+class DTensorConstructorTest(NPUDTensorTestBase):
32	@property	32	@property
33	def world_size(self):	33	def world_size(self):
34	return 4	34	return 4

Mtest/distributed/_tensor/test_math_ops.py+2-2

	@@ -5,13 +5,13 @@ import torch
5	from torch.distributed._tensor import distribute_tensor, DeviceMesh	5	from torch.distributed._tensor import distribute_tensor, DeviceMesh
6	from torch.distributed._tensor.placement_types import Replicate, Shard	6	from torch.distributed._tensor.placement_types import Replicate, Shard
7	from torch.testing._internal.common_utils import run_tests	7	from torch.testing._internal.common_utils import run_tests
8	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
9		8
10	import torch_npu	9	import torch_npu
11	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	10	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		11	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
12		12
13		13
14	-class DistMathOpsTest(DTensorTestBase):	14	+class DistMathOpsTest(NPUDTensorTestBase):
15	@skipIfUnsupportMultiNPU(4)	15	@skipIfUnsupportMultiNPU(4)
16	@with_comms	16	@with_comms
17	def test_sum(self):	17	def test_sum(self):

Mtest/distributed/_tensor/test_matrix_ops.py+8-8

	@@ -5,20 +5,20 @@ from unittest import skip
5	import torch	5	import torch
6	from torch.distributed._tensor import DeviceMesh, distribute_tensor	6	from torch.distributed._tensor import DeviceMesh, distribute_tensor
7	from torch.distributed._tensor.api import DTensor	7	from torch.distributed._tensor.api import DTensor
8	-from torch.distributed._tensor.placement_types import (	8	+from torch.distributed.tensor.placement_types import (
9	- _Partial,	9	+ Partial,
10	Placement,	10	Placement,
11	Replicate,	11	Replicate,
12	Shard,	12	Shard,
13	)	13	)
14	from torch.testing._internal.common_utils import run_tests	14	from torch.testing._internal.common_utils import run_tests
15	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
16		15
17	import torch_npu	16	import torch_npu
18	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	17	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		18	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
19		19
20		20
21	-class DistMatrixOpsTest(DTensorTestBase):	21	+class DistMatrixOpsTest(NPUDTensorTestBase):
22	@skipIfUnsupportMultiNPU(4)	22	@skipIfUnsupportMultiNPU(4)
23	@with_comms	23	@with_comms
24	def test_addmm(self):	24	def test_addmm(self):
	@@ -60,7 +60,7 @@ class DistMatrixOpsTest(DTensorTestBase):
60		60
61	# test if addmm output is a partial	61	# test if addmm output is a partial
62	self.assertIsInstance(dist_res, DTensor)	62	self.assertIsInstance(dist_res, DTensor)
63	- self.assertIsInstance(dist_res.placements[0], _Partial)	63	+ self.assertIsInstance(dist_res.placements[0], Partial)
64		64
65	# test if result is the same as tensor	65	# test if result is the same as tensor
66	replica_res = dist_res.redistribute(device_mesh, replica_spec)	66	replica_res = dist_res.redistribute(device_mesh, replica_spec)
	@@ -130,10 +130,10 @@ class DistMatrixOpsTest(DTensorTestBase):
130	da = distribute_tensor(a, device_mesh, [Shard(1)])	130	da = distribute_tensor(a, device_mesh, [Shard(1)])
131	db = distribute_tensor(b, device_mesh, [Shard(0)])	131	db = distribute_tensor(b, device_mesh, [Shard(0)])
132		132
133	- # mm(da, db) should return a _Partial tensor.	133	+ # mm(da, db) should return a Partial tensor.
134	- # transposing it should keep it _Partial	134	+ # transposing it should keep it Partial
135	dc = torch.mm(da, db).t()	135	dc = torch.mm(da, db).t()
136	- self.assertTrue(isinstance(dc.placements[0], _Partial))	136	+ self.assertTrue(isinstance(dc.placements[0], Partial))
137	# check that the local and distributed op results match	137	# check that the local and distributed op results match
138	self.assertEqual(	138	self.assertEqual(
139	c,	139	c,

Mtest/distributed/_tensor/test_pointwise_ops.py+7-7

	@@ -7,18 +7,18 @@ import torch.utils._pytree as pytree
7	from torch import Tensor	7	from torch import Tensor
8		8
9	from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor	9	from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor
10	-from torch.distributed._tensor.placement_types import (	10	+from torch.distributed.tensor.placement_types import (
11	- _Partial,	11	+ Partial,
12	Placement,	12	Placement,
13	Replicate,	13	Replicate,
14	Shard,	14	Shard,
15	)	15	)
16	from torch.distributed.distributed_c10d import ReduceOp	16	from torch.distributed.distributed_c10d import ReduceOp
17	from torch.testing._internal.common_utils import run_tests	17	from torch.testing._internal.common_utils import run_tests
18	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
19		18
20	import torch_npu	19	import torch_npu
21	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	20	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		21	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
22		22
23		23
24	def no_op():	24	def no_op():
	@@ -69,7 +69,7 @@ def deepcopy_convert_from_dtensor(val: Any) -> Any:
69	return pytree.tree_map(f, [val])[0]	69	return pytree.tree_map(f, [val])[0]
70		70
71		71
72	-class DistElementwiseOpsTest(DTensorTestBase):	72	+class DistElementwiseOpsTest(NPUDTensorTestBase):
73	def _compare_pairwise_ops(	73	def _compare_pairwise_ops(
74	self,	74	self,
75	*,	75	*,
	@@ -137,8 +137,8 @@ class DistElementwiseOpsTest(DTensorTestBase):
137	@with_comms	137	@with_comms
138	def test_partial_add(self):	138	def test_partial_add(self):
139	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	139	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
140	- d_1 = DTensor.from_local(torch.rand(2, 2), device_mesh, [_Partial()])	140	+ d_1 = DTensor.from_local(torch.rand(2, 2), device_mesh, [Partial()])
141	- d_2 = DTensor.from_local(torch.rand(2, 2), device_mesh, [_Partial()])	141	+ d_2 = DTensor.from_local(torch.rand(2, 2), device_mesh, [Partial()])
142	d_3 = d_1 + d_2	142	d_3 = d_1 + d_2
143	self.assertEqual(d_3._spec.placements[0].is_partial(), True)	143	self.assertEqual(d_3._spec.placements[0].is_partial(), True)
144		144
	@@ -218,7 +218,7 @@ class DistElementwiseOpsTest(DTensorTestBase):
218	with self.assertRaisesRegex(RuntimeError, "supported"):	218	with self.assertRaisesRegex(RuntimeError, "supported"):
219	self._run_sharded_elementwise_ops(	219	self._run_sharded_elementwise_ops(
220	device_mesh=device_mesh,	220	device_mesh=device_mesh,
221	- placements=[_Partial(ReduceOp.SUM)],	221	+ placements=[Partial(ReduceOp.SUM)],
222	input_size=(8, 5),	222	input_size=(8, 5),
223	op=torch.nn.functional.dropout,	223	op=torch.nn.functional.dropout,
224	)	224	)

Mtest/distributed/_tensor/test_random_ops.py+3-4

	@@ -13,14 +13,13 @@ from torch.distributed.tensor._random import is_rng_supported_mesh, manual_seed
13		13
14	from torch.distributed.distributed_c10d import broadcast_object_list	14	from torch.distributed.distributed_c10d import broadcast_object_list
15		15
16	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
17	-
18	import torch_npu	16	import torch_npu
19	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	17	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		18	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
20	from torch_npu.testing.testcase import run_tests	19	from torch_npu.testing.testcase import run_tests
21		20
22		21
23	-class DistTensorRandomInitTest(DTensorTestBase):	22	+class DistTensorRandomInitTest(NPUDTensorTestBase):
24	def _run_init_op(self, init_op, args, *kwargs):	23	def _run_init_op(self, init_op, args, *kwargs):
25	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	24	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
26	shard_spec = [Shard(0)]	25	shard_spec = [Shard(0)]
	@@ -75,7 +74,7 @@ class DistTensorRandomInitTest(DTensorTestBase):
75	self._run_init_op(torch.nn.init.uniform_, a=0, b=1.2)	74	self._run_init_op(torch.nn.init.uniform_, a=0, b=1.2)
76		75
77		76
78	-class DistTensorRandomOpTest(DTensorTestBase):	77	+class DistTensorRandomOpTest(NPUDTensorTestBase):
79	@skipIfUnsupportMultiNPU(4)	78	@skipIfUnsupportMultiNPU(4)
80	@with_comms	79	@with_comms
81	def test_rng_tracker_init(self):	80	def test_rng_tracker_init(self):

Mtest/distributed/_tensor/test_redistribute.py+7-7

	@@ -2,15 +2,15 @@ import itertools
2		2
3	import torch	3	import torch
4	from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor	4	from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor
5	-from torch.distributed._tensor.placement_types import _Partial, Replicate, Shard	5	+from torch.distributed.tensor.placement_types import Partial, Replicate, Shard
6	from torch.testing._internal.common_utils import run_tests	6	from torch.testing._internal.common_utils import run_tests
7	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
8		7
9	import torch_npu	8	import torch_npu
10	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	9	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		10	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
11		11
12		12
13	-class RedistributeTest(DTensorTestBase):	13	+class RedistributeTest(NPUDTensorTestBase):
14	@skipIfUnsupportMultiNPU(4)	14	@skipIfUnsupportMultiNPU(4)
15	@with_comms	15	@with_comms
16	def test_shard_to_replicate_forward_backward(self):	16	def test_shard_to_replicate_forward_backward(self):
	@@ -116,7 +116,7 @@ class RedistributeTest(DTensorTestBase):
116	# backward should work as expected	116	# backward should work as expected
117	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	117	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
118	partial_local = torch.ones(12, 3, device=self.device_type, requires_grad=True)	118	partial_local = torch.ones(12, 3, device=self.device_type, requires_grad=True)
119	- partial_spec = [_Partial()]	119	+ partial_spec = [Partial()]
120	replica_spec = [Replicate()]	120	replica_spec = [Replicate()]
121	# test partial -> replicate, which trigger all_reduce	121	# test partial -> replicate, which trigger all_reduce
122	partial_tensor = DTensor.from_local(partial_local, device_mesh, partial_spec)	122	partial_tensor = DTensor.from_local(partial_local, device_mesh, partial_spec)
	@@ -139,11 +139,11 @@ class RedistributeTest(DTensorTestBase):
139	def test_replicate_to_partial(self):	139	def test_replicate_to_partial(self):
140	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	140	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
141	local_tensor = torch.randn(12, 3, device=self.device_type, requires_grad=True)	141	local_tensor = torch.randn(12, 3, device=self.device_type, requires_grad=True)
142	- partial_spec = _Partial()	142	+ partial_spec = Partial()
143	replica_spec = Replicate()	143	replica_spec = Replicate()
144	# 1) test replicate -> partial forward	144	# 1) test replicate -> partial forward
145	replica_tensor = distribute_tensor(local_tensor, device_mesh, [replica_spec])	145	replica_tensor = distribute_tensor(local_tensor, device_mesh, [replica_spec])
146	- with self.assertRaisesRegex(RuntimeError, "Can not redistribute to _Partial"):	146	+ with self.assertRaisesRegex(RuntimeError, "Can not redistribute to Partial"):
147	partial_tensor = replica_tensor.redistribute(device_mesh, [partial_spec])	147	partial_tensor = replica_tensor.redistribute(device_mesh, [partial_spec])
148		148
149	from torch.distributed._tensor.redistribute import Redistribute	149	from torch.distributed._tensor.redistribute import Redistribute
	@@ -179,7 +179,7 @@ class RedistributeTest(DTensorTestBase):
179	@with_comms	179	@with_comms
180	def test_partial_to_shard(self):	180	def test_partial_to_shard(self):
181	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	181	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
182	- partial_spec = [_Partial()]	182	+ partial_spec = [Partial()]
183	my_rank = device_mesh.get_rank()	183	my_rank = device_mesh.get_rank()
184		184
185	input_sizes_and_shard_dim = [	185	input_sizes_and_shard_dim = [

Mtest/distributed/_tensor/test_tensor_ops.py+7-7

	@@ -1,17 +1,17 @@
1	import torch	1	import torch
2	from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor	2	from torch.distributed._tensor import DeviceMesh, distribute_tensor, DTensor
3	-from torch.distributed._tensor.placement_types import _Partial, Replicate, Shard	3	+from torch.distributed.tensor.placement_types import Partial, Replicate, Shard
4	from torch.testing._internal.common_utils import run_tests	4	from torch.testing._internal.common_utils import run_tests
5	from torch.testing._internal.distributed._tensor.common_dtensor import (	5	from torch.testing._internal.distributed._tensor.common_dtensor import (
6	DTensorConverter,	6	DTensorConverter,
7	- DTensorTestBase
8	)	7	)
9		8
10	import torch_npu	9	import torch_npu
11	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	10	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		11	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
12		12
13		13
14	-class DistTensorOpsTest(DTensorTestBase):	14	+class DistTensorOpsTest(NPUDTensorTestBase):
15	@skipIfUnsupportMultiNPU(4)	15	@skipIfUnsupportMultiNPU(4)
16	@with_comms	16	@with_comms
17	def test_aten_contiguous(self):	17	def test_aten_contiguous(self):
	@@ -67,7 +67,7 @@ class DistTensorOpsTest(DTensorTestBase):
67	# test inplace op self and other dtensor with other specs	67	# test inplace op self and other dtensor with other specs
68	# and make sure out spec not change	68	# and make sure out spec not change
69	shard_spec = [Shard(0)]	69	shard_spec = [Shard(0)]
70	- partial_spec = [_Partial()]	70	+ partial_spec = [Partial()]
71	dt_to_inplace_add = distribute_tensor(input_tensor, mesh, shard_spec)	71	dt_to_inplace_add = distribute_tensor(input_tensor, mesh, shard_spec)
72	partial_grad = DTensor.from_local(torch.randn(12, 3), mesh, partial_spec)	72	partial_grad = DTensor.from_local(torch.randn(12, 3), mesh, partial_spec)
73	res = dt_to_inplace_add.add_(partial_grad)	73	res = dt_to_inplace_add.add_(partial_grad)
	@@ -149,7 +149,7 @@ class DistTensorOpsTest(DTensorTestBase):
149	@with_comms	149	@with_comms
150	def test_ones_like_partial_sum(self):	150	def test_ones_like_partial_sum(self):
151	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	151	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
152	- shard_spec = [_Partial()]	152	+ shard_spec = [Partial()]
153		153
154	input_tensor = torch.randn(4, 8, requires_grad=True)	154	input_tensor = torch.randn(4, 8, requires_grad=True)
155	dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)	155	dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)
	@@ -166,7 +166,7 @@ class DistTensorOpsTest(DTensorTestBase):
166	@with_comms	166	@with_comms
167	def test_fill_inplace_partial_sum(self):	167	def test_fill_inplace_partial_sum(self):
168	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	168	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
169	- shard_spec = [_Partial()]	169	+ shard_spec = [Partial()]
170		170
171	input_tensor = torch.randn(4, 8, requires_grad=True)	171	input_tensor = torch.randn(4, 8, requires_grad=True)
172	dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)	172	dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)
	@@ -182,7 +182,7 @@ class DistTensorOpsTest(DTensorTestBase):
182	@with_comms	182	@with_comms
183	def test_zeros_like_partial_sum(self):	183	def test_zeros_like_partial_sum(self):
184	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))	184	device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
185	- shard_spec = [_Partial()]	185	+ shard_spec = [Partial()]
186		186
187	input_tensor = torch.randn(4, 8, requires_grad=True)	187	input_tensor = torch.randn(4, 8, requires_grad=True)
188	dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)	188	dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)

Mtest/distributed/_tensor/test_utils.py+8-9

	@@ -1,26 +1,25 @@
1	import itertools	1	import itertools
2	import torch	2	import torch
3	from torch.distributed._tensor import distribute_tensor	3	from torch.distributed._tensor import distribute_tensor
4	-from torch.distributed._tensor._utils import (	4	+from torch.distributed.tensor._utils import (
5	- compute_local_shape,
6	compute_local_shape_and_global_offset,	5	compute_local_shape_and_global_offset,
7	)	6	)
8	from torch.distributed._tensor.device_mesh import DeviceMesh	7	from torch.distributed._tensor.device_mesh import DeviceMesh
9	from torch.distributed._tensor.placement_types import Replicate, Shard	8	from torch.distributed._tensor.placement_types import Replicate, Shard
10		9
11	from torch.testing._internal.common_utils import run_tests	10	from torch.testing._internal.common_utils import run_tests
12	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
13		11
14	import torch_npu	12	import torch_npu
15	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	13	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		14	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
16		15
17		16
18	-class UtilTest(DTensorTestBase):	17	+class UtilTest(NPUDTensorTestBase):
19	@property	18	@property
20	def world_size(self):	19	def world_size(self):
21	return 8	20	return 8
22		21
23	- @skipIfUnsupportMultiNPU(4)	22	+ @skipIfUnsupportMultiNPU(8)
24	@with_comms	23	@with_comms
25	def test_compute_local_shape_2d_uneven(self):	24	def test_compute_local_shape_2d_uneven(self):
26	# mesh: 4 * 2	25	# mesh: 4 * 2
	@@ -31,7 +30,7 @@ class UtilTest(DTensorTestBase):
31		30
32	# replicate, shard	31	# replicate, shard
33	placements2 = [Replicate(), Shard(0)]	32	placements2 = [Replicate(), Shard(0)]
34	- local_size2 = compute_local_shape(size, mesh, placements2)	33	+ local_size2, _ = compute_local_shape_and_global_offset(size, mesh, placements2)
35	if rank_coordinates[1] < 1:	34	if rank_coordinates[1] < 1:
36	self.assertEqual(local_size2, torch.Size([4, 7]))	35	self.assertEqual(local_size2, torch.Size([4, 7]))
37	else:	36	else:
	@@ -39,7 +38,7 @@ class UtilTest(DTensorTestBase):
39		38
40	# shard, shard	39	# shard, shard
41	placements3 = [Shard(0), Shard(1)]	40	placements3 = [Shard(0), Shard(1)]
42	- local_size3 = compute_local_shape(size, mesh, placements3)	41	+ local_size3, _ = compute_local_shape_and_global_offset(size, mesh, placements3)
43	# first dim	42	# first dim
44	if rank_coordinates[0] < 3:	43	if rank_coordinates[0] < 3:
45	self.assertEqual(local_size3[0], 2)	44	self.assertEqual(local_size3[0], 2)
	@@ -51,7 +50,7 @@ class UtilTest(DTensorTestBase):
51	else:	50	else:
52	self.assertEqual(local_size3[1], 3)	51	self.assertEqual(local_size3[1], 3)
53		52
54	- @skipIfUnsupportMultiNPU(4)	53	+ @skipIfUnsupportMultiNPU(8)
55	@with_comms	54	@with_comms
56	def test_compute_local_shape_and_global_offset_1D(self):	55	def test_compute_local_shape_and_global_offset_1D(self):
57	one_d_placements = [[Shard(0)], [Replicate()]]	56	one_d_placements = [[Shard(0)], [Replicate()]]
	@@ -77,7 +76,7 @@ class UtilTest(DTensorTestBase):
77	global_tensor[dim0_start:dim0_end],	76	global_tensor[dim0_start:dim0_end],
78	)	77	)
79		78
80	- @skipIfUnsupportMultiNPU(4)	79	+ @skipIfUnsupportMultiNPU(8)
81	@with_comms	80	@with_comms
82	def test_compute_local_shape_and_global_offset_2D(self):	81	def test_compute_local_shape_and_global_offset_2D(self):
83	two_d_placements_options = [Shard(0), Shard(1), Replicate()]	82	two_d_placements_options = [Shard(0), Shard(1), Replicate()]

Mtest/distributed/_tensor/test_view_ops.py+4-4

	@@ -5,8 +5,8 @@ import torch
5	import torch.distributed as dist	5	import torch.distributed as dist
6	from torch import rand, randn, Tensor	6	from torch import rand, randn, Tensor
7	from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate, Shard	7	from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate, Shard
8	-from torch.distributed._tensor.debug import CommDebugMode	8	+from torch.distributed.tensor.debug import CommDebugMode
9	-from torch.distributed._tensor.ops._view_ops import (	9	+from torch.distributed.tensor._ops._view_ops import (
10	Broadcast,	10	Broadcast,
11	dim_maps,	11	dim_maps,
12	Flatten,	12	Flatten,
	@@ -18,14 +18,14 @@ from torch.distributed._tensor.ops._view_ops import (
18	)	18	)
19	from torch.distributed._tensor.placement_types import Placement	19	from torch.distributed._tensor.placement_types import Placement
20	from torch.testing._internal.common_utils import run_tests	20	from torch.testing._internal.common_utils import run_tests
21	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
22	from torch.utils import _pytree as pytree	21	from torch.utils import _pytree as pytree
23		22
24	import torch_npu	23	import torch_npu
25	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	24	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		25	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
26		26
27		27
28	-class TestViewOps(DTensorTestBase):	28	+class TestViewOps(NPUDTensorTestBase):
29	@skipIfUnsupportMultiNPU(4)	29	@skipIfUnsupportMultiNPU(4)
30	def test_view_groups(self):	30	def test_view_groups(self):
31	self.assertEqual(	31	self.assertEqual(

Mtest/distributed/checkpoint/test_fsdp_model_state.py+4-7

	@@ -9,18 +9,15 @@ from torch.distributed.checkpoint.default_planner import (
9	DefaultSavePlanner,	9	DefaultSavePlanner,
10	DefaultLoadPlanner,	10	DefaultLoadPlanner,
11	)	11	)
12	-
13	-from torch.testing._internal.distributed._tensor.common_dtensor import (
14	- DTensorTestBase,
15	-)
16	from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir	12	from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
17		13
18	import torch_npu	14	import torch_npu
19	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	15	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		16	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
20	from torch_npu.testing.testcase import run_tests	17	from torch_npu.testing.testcase import run_tests
21		18
22		19
23	-class FsdpModelStateCheckpoint(DTensorTestBase):	20	+class FsdpModelStateCheckpoint(NPUDTensorTestBase):
24	def _test_fsdp_model_state(self, process_group) -> None:	21	def _test_fsdp_model_state(self, process_group) -> None:
25	CHECKPOINT_DIR = self.temp_dir	22	CHECKPOINT_DIR = self.temp_dir
26		23
	@@ -65,8 +62,8 @@ class FsdpModelStateCheckpoint(DTensorTestBase):
65	self.assertEqual(model.weight, model_2.weight)	62	self.assertEqual(model.weight, model_2.weight)
66	self.assertEqual(model.bias, model_2.bias)	63	self.assertEqual(model.bias, model_2.bias)
67		64
68	- @with_comms
69	@skipIfUnsupportMultiNPU(4)	65	@skipIfUnsupportMultiNPU(4)
		66	+ @with_comms
70	@with_temp_dir	67	@with_temp_dir
71	def test_fsdp_model_state_no_resharding(self):	68	def test_fsdp_model_state_no_resharding(self):
72	self._test_fsdp_model_state(process_group=None)	69	self._test_fsdp_model_state(process_group=None)
	@@ -86,8 +83,8 @@ class FsdpModelStateCheckpoint(DTensorTestBase):
86		83
87	return my_fsdp	84	return my_fsdp
88		85
89	- @with_comms
90	@skipIfUnsupportMultiNPU(4)	86	@skipIfUnsupportMultiNPU(4)
		87	+ @with_comms
91	@with_temp_dir	88	@with_temp_dir
92	def test_fsdp_model_state_with_resharding(self):	89	def test_fsdp_model_state_with_resharding(self):
93	self._test_fsdp_model_state(process_group=self._create_new_dist_group())	90	self._test_fsdp_model_state(process_group=self._create_new_dist_group())

Mtest/distributed/checkpoint/test_fsdp_optim_state.py+3-6

	@@ -9,20 +9,17 @@ from torch.distributed.checkpoint.default_planner import (
9	DefaultSavePlanner,	9	DefaultSavePlanner,
10	DefaultLoadPlanner,	10	DefaultLoadPlanner,
11	)	11	)
12	-
13	-from torch.testing._internal.distributed._tensor.common_dtensor import (
14	- DTensorTestBase,
15	-)
16	from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir	12	from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
17		13
18	import torch_npu	14	import torch_npu
19	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	15	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		16	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
20	from torch_npu.testing.testcase import run_tests	17	from torch_npu.testing.testcase import run_tests
21		18
22		19
23	-class FsdpOptimStateCheckpoint(DTensorTestBase):	20	+class FsdpOptimStateCheckpoint(NPUDTensorTestBase):
24	- @with_comms
25	@skipIfUnsupportMultiNPU(4)	21	@skipIfUnsupportMultiNPU(4)
		22	+ @with_comms
26	@with_temp_dir	23	@with_temp_dir
27	def test_distributed_tensor_planner(self) -> None:	24	def test_distributed_tensor_planner(self) -> None:
28	CHECKPOINT_DIR = self.temp_dir	25	CHECKPOINT_DIR = self.temp_dir

Mtest/distributed/tensor/parallel/test_parallelize_api.py+2-2

	@@ -14,13 +14,13 @@ from torch.distributed.tensor.parallel.style import (
14	)	14	)
15	from torch.testing._internal.common_utils import run_tests	15	from torch.testing._internal.common_utils import run_tests
16	from torch.testing._internal.distributed._tensor.common_dtensor import (	16	from torch.testing._internal.distributed._tensor.common_dtensor import (
17	- DTensorTestBase,
18	MLPModule,	17	MLPModule,
19	MLPStacked,	18	MLPStacked,
20	)	19	)
21		20
22	import torch_npu	21	import torch_npu
23	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	22	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		23	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
24		24
25		25
26	class DummyModule(torch.nn.Module):	26	class DummyModule(torch.nn.Module):
	@@ -31,7 +31,7 @@ class DummyModule(torch.nn.Module):
31	return x	31	return x
32		32
33		33
34	-class TensorParallelAPITests(DTensorTestBase):	34	+class TensorParallelAPITests(NPUDTensorTestBase):
35	@property	35	@property
36	def world_size(self):	36	def world_size(self):
37	return 2	37	return 2

Mtest/distributed/tensor/parallel/test_tp_examples.py+2-2

	@@ -33,7 +33,6 @@ from torch.testing._internal.common_utils import (
33	run_tests,	33	run_tests,
34	)	34	)
35	from torch.testing._internal.distributed._tensor.common_dtensor import (	35	from torch.testing._internal.distributed._tensor.common_dtensor import (
36	- DTensorTestBase,
37	MLPModule,	36	MLPModule,
38	ModelArgs,	37	ModelArgs,
39	skip_unless_torch_gpu,	38	skip_unless_torch_gpu,
	@@ -42,6 +41,7 @@ from torch.testing._internal.distributed._tensor.common_dtensor import (
42		41
43	import torch_npu	42	import torch_npu
44	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	43	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		44	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
45		45
46		46
47	c10d_functional = torch.ops.c10d_functional	47	c10d_functional = torch.ops.c10d_functional
	@@ -58,7 +58,7 @@ class ExpCommCounts(NamedTuple):
58	optim: Optional[dict] = None	58	optim: Optional[dict] = None
59		59
60		60
61	-class DistTensorParallelExampleTest(DTensorTestBase):	61	+class DistTensorParallelExampleTest(NPUDTensorTestBase):
62	@property	62	@property
63	def world_size(self):	63	def world_size(self):
64	return 2	64	return 2

Mtest/distributed/tensor/parallel/test_tp_style.py+2-2

	@@ -25,19 +25,19 @@ from torch.distributed.tensor.parallel.style import (
25	from torch.distributed.tensor.placement_types import _Partial	25	from torch.distributed.tensor.placement_types import _Partial
26	from torch.testing._internal.common_utils import run_tests	26	from torch.testing._internal.common_utils import run_tests
27	from torch.testing._internal.distributed._tensor.common_dtensor import (	27	from torch.testing._internal.distributed._tensor.common_dtensor import (
28	- DTensorTestBase,
29	RMSNormPython,	28	RMSNormPython,
30	)	29	)
31		30
32	import torch_npu	31	import torch_npu
33	from torch_npu.testing.common_utils import SupportedDevices	32	from torch_npu.testing.common_utils import SupportedDevices
34	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	33	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		34	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
35		35
36		36
37	c10d_functional = torch.ops.c10d_functional	37	c10d_functional = torch.ops.c10d_functional
38		38
39		39
40	-class TensorParallelStyleTest(DTensorTestBase):	40	+class TensorParallelStyleTest(NPUDTensorTestBase):
41	@property	41	@property
42	def world_size(self):	42	def world_size(self):
43	return 2	43	return 2

Mtest/distributed/tensor/test_math_ops.py+2-2

	@@ -1,13 +1,13 @@
1	import torch	1	import torch
2	from torch.distributed._tensor import distribute_tensor, Replicate, Shard	2	from torch.distributed._tensor import distribute_tensor, Replicate, Shard
3	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
4		3
5	import torch_npu	4	import torch_npu
6	from torch_npu.testing.testcase import run_tests	5	from torch_npu.testing.testcase import run_tests
7	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	6	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		7	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
8		8
9		9
10	-class TestMathOps(DTensorTestBase):	10	+class TestMathOps(NPUDTensorTestBase):
11	@skipIfUnsupportMultiNPU(4)	11	@skipIfUnsupportMultiNPU(4)
12	@with_comms	12	@with_comms
13	def test_npu_rms_norm_forward(self):	13	def test_npu_rms_norm_forward(self):

Mtest/distributed/test_device_mesh.py+2-38

	@@ -25,12 +25,12 @@ from torch.distributed.tensor._collective_utils import (
25	unpad_tensor,	25	unpad_tensor,
26	)	26	)
27	from torch.distributed.tensor.placement_types import _Partial, Shard	27	from torch.distributed.tensor.placement_types import _Partial, Shard
28	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
29	from torch.testing._internal.distributed.fake_pg import FakeStore	28	from torch.testing._internal.distributed.fake_pg import FakeStore
30	from torch.utils._typing_utils import not_none	29	from torch.utils._typing_utils import not_none
31		30
32	import torch_npu	31	import torch_npu
33	-from torch_npu.testing.common_distributed import init_pg, skipIfUnsupportMultiNPU, TEST_SKIPS	32	+from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
		33	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
34	from torch_npu.testing.testcase import run_tests	34	from torch_npu.testing.testcase import run_tests
35		35
36		36
	@@ -53,42 +53,6 @@ def _set_env_var(addr="localhost", port="29500", world_size=1, rank=0):
53	os.environ["RANK"] = f"{rank}"	53	os.environ["RANK"] = f"{rank}"
54		54
55		55
56	-def with_comms(func):
57	- if func is None:
58	- raise RuntimeError("Test function is None.")
59	-
60	- def get_device_type(self):
61	- if torch.npu.is_available() and torch.npu.device_count() >= self.world_size:
62	- return "npu"
63	- return "cpu"
64	-
65	- @wraps(func) # pyre-ignore[6]
66	- def wrapper(
67	- self, args: Tuple[object], *kwargs: Dict[str, Any] # type: ignore[misc]
68	- ) -> None:
69	-
70	- pg_backend = (
71	- "hccl" if get_device_type(self) == "npu" else "gloo"
72	- )
73	- if pg_backend == "hccl" and torch.npu.device_count() < self.world_size:
74	- raise RuntimeError(TEST_SKIPS[f"multi-npu-{self.world_size}"].message)
75	-
76	- init_pg(backend=pg_backend, world_size=self.world_size, rank=self.rank, file_name=self.file_name)
77	-
78	- torch.npu.manual_seed(0)
79	- torch.npu.initial_seed()
80	- func(self, args, *kwargs) # type: ignore[misc]
81	- self.destroy_pg()
82	-
83	- return wrapper
84	-
85	-
86	-class NPUDTensorTestBase(DTensorTestBase):
87	- @property
88	- def device_type(self):
89	- return "npu"
90	-
91	-
92	class DeviceMeshTest(NPUDTensorTestBase):	56	class DeviceMeshTest(NPUDTensorTestBase):
93	@property	57	@property
94	def world_size(self):	58	def world_size(self):

Mtest/distributed/test_register_sharding.py+2-3

	@@ -1,14 +1,13 @@
1	import torch	1	import torch
2	from torch.distributed._tensor import distribute_tensor, Replicate	2	from torch.distributed._tensor import distribute_tensor, Replicate
3	from torch.testing._internal.common_utils import run_tests	3	from torch.testing._internal.common_utils import run_tests
4	-from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
5		4
6	import torch_npu	5	import torch_npu
7	from torch_npu.testing.common_distributed import with_comms	6	from torch_npu.testing.common_distributed import with_comms
		7	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
8		8
9		9
10	-class TestRegisterSharding(DTensorTestBase):	10	+class TestRegisterSharding(NPUDTensorTestBase):
11	-
12	def _run_matmul(self, shape1, shape2, device_mesh):	11	def _run_matmul(self, shape1, shape2, device_mesh):
13	x = torch.rand(shape1, device=self.device_type)	12	x = torch.rand(shape1, device=self.device_type)
14	dist_x = distribute_tensor(x, device_mesh, [Replicate()])	13	dist_x = distribute_tensor(x, device_mesh, [Replicate()])

Atorch_npu/testing/_internal/common_dtensor.py+7-0

@@ -0,0 +1,7 @@
	1	+from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
	2	+
	3	+
	4	+class NPUDTensorTestBase(DTensorTestBase):
	5	+ @property
	6	+ def device_type(self):
	7	+ return "npu"

Mtorch_npu/testing/common_distributed.py+8-9

	@@ -13,7 +13,6 @@ import torch
13	import torch.distributed as dist	13	import torch.distributed as dist
14	import torch_npu	14	import torch_npu
15		15
16	-
17	TestSkip = namedtuple('TestSkip', 'exit_code, message')	16	TestSkip = namedtuple('TestSkip', 'exit_code, message')
18	TEST_SKIPS = {	17	TEST_SKIPS = {
19	"multi-npu": TestSkip(75, "Need at least 2 ASCEND devices"),	18	"multi-npu": TestSkip(75, "Need at least 2 ASCEND devices"),
	@@ -34,7 +33,7 @@ def skipIfUnsupportMultiNPU(npu_number_needed):
34	def skip_dec(func):	33	def skip_dec(func):
35	def wrapper(self):	34	def wrapper(self):
36	if not torch.npu.is_available() or torch.npu.device_count() < npu_number_needed:	35	if not torch.npu.is_available() or torch.npu.device_count() < npu_number_needed:
37	- return unittest.SkipTest("Multi-NPU condition not satisfied")	36	+ raise unittest.SkipTest("Multi-NPU condition not satisfied")
38	return func(self)	37	return func(self)
39	return wrapper	38	return wrapper
40	return skip_dec	39	return skip_dec
	@@ -44,18 +43,18 @@ def with_comms(func):
44	if func is None:	43	if func is None:
45	raise RuntimeError("Test function is None.")	44	raise RuntimeError("Test function is None.")
46		45
		46	+ def get_device_type(self):
		47	+ if torch.npu.is_available() and torch.npu.device_count() >= self.world_size:
		48	+ return "npu"
		49	+ return "cpu"
		50	+
47	@wraps(func) # pyre-ignore[6]	51	@wraps(func) # pyre-ignore[6]
48	def wrapper(	52	def wrapper(
49	- self, args: Tuple[object], *kwargs: Dict[str, Any] # type: ignore[misc]	53	+ self, args: Tuple[object], *kwargs: Dict[str, Any] # type: ignore[misc]
50	) -> None:	54	) -> None:
51	- # if backend not specified, and npu available, then use hccl, else gloo
52	- if torch.npu.is_available() and torch.npu.device_count() >= self.world_size:
53	- self.device_type = "npu"
54	- else:
55	- self.device_type = "cpu"
56		55
57	pg_backend = (	56	pg_backend = (
58	- "hccl" if self.device_type == "npu" else "gloo"	57	+ "hccl" if get_device_type(self) == "npu" else "gloo"
59	)	58	)
60	if pg_backend == "hccl" and torch.npu.device_count() < self.world_size:	59	if pg_backend == "hccl" and torch.npu.device_count() < self.world_size:
61	raise RuntimeError(TEST_SKIPS[f"multi-npu-{self.world_size}"].message)	60	raise RuntimeError(TEST_SKIPS[f"multi-npu-{self.world_size}"].message)