register strategy for npu custom ops-pytorch-AtomGit

Mtest/distributed/tensor/test_math_ops.py+160-1

	@@ -1,12 +1,37 @@
		1	+import itertools
		2	+
1	import torch	3	import torch
2	from torch.distributed._tensor import distribute_tensor, Replicate, Shard	4	from torch.distributed._tensor import distribute_tensor, Replicate, Shard
		5	+from torch.testing._internal.common_utils import (
		6	+ instantiate_parametrized_tests,
		7	+ parametrize,
		8	+ run_tests,
		9	+)
3		10
4	import torch_npu	11	import torch_npu
5	-from torch_npu.testing.testcase import run_tests	12	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
6	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU	13	from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
7	from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase	14	from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
8		15
9		16
		17	+def get_shape_from_layout(batch: int, num_head: int, seq_length: int, dimension: int, layout: str):
		18	+ layout_map = {
		19	+ "B": batch,
		20	+ "N": num_head,
		21	+ "S": seq_length,
		22	+ "D": dimension,
		23	+ "1": 1,
		24	+ }
		25	+ shape = []
		26	+ for dim in layout:
		27	+ if dim in layout_map:
		28	+ shape.append(layout_map[dim])
		29	+ else:
		30	+ raise ValueError(f"Invalid layout character: {dim}")
		31	+
		32	+ return tuple(shape)
		33	+
		34	+
10	class TestMathOps(NPUDTensorTestBase):	35	class TestMathOps(NPUDTensorTestBase):
11	@skipIfUnsupportMultiNPU(4)	36	@skipIfUnsupportMultiNPU(4)
12	@with_comms	37	@with_comms
	@@ -61,6 +86,140 @@ class TestMathOps(NPUDTensorTestBase):
61	self.assertEqual(dist_dx.full_tensor(), dx)	86	self.assertEqual(dist_dx.full_tensor(), dx)
62	self.assertEqual(dist_dw.full_tensor(), dw)	87	self.assertEqual(dist_dw.full_tensor(), dw)
63		88
		89	+ @skipIfUnsupportMultiNPU(4)
		90	+ @with_comms
		91	+ def test_npu_add_rms_norm_forward(self):
		92	+ device_mesh = self.build_device_mesh()
		93	+
		94	+ x1 = torch.randn((1, 128, 64), dtype=torch.float32).npu()
		95	+ x2 = torch.randn((1, 128, 64), dtype=torch.float32).npu()
		96	+ gamma = torch.randn(64, dtype=torch.float32).npu()
		97	+
		98	+ y, rstd, x = torch_npu.npu_add_rms_norm(x1, x2, gamma)
		99	+
		100	+ def test_placement_comb(placements1, placements2):
		101	+ dist_x1 = distribute_tensor(x1, device_mesh, placements1)
		102	+ dist_x2 = distribute_tensor(x2, device_mesh, placements2)
		103	+ dist_gamma = distribute_tensor(gamma, device_mesh, [Replicate()])
		104	+ dist_y, dist_rstd, dist_x = torch_npu.npu_add_rms_norm(dist_x1, dist_x2, dist_gamma)
		105	+ self.assertEqual(dist_y.full_tensor(), y)
		106	+ self.assertEqual(dist_rstd.full_tensor(), rstd)
		107	+ self.assertEqual(dist_x.full_tensor(), x)
		108	+
		109	+ placement = [Shard(0), Shard(1), Shard(2), Replicate()]
		110	+ placement_combs = itertools.product(placement, placement)
		111	+ for comb in placement_combs:
		112	+ test_placement_comb([comb[0]], [comb[1]])
		113	+
		114	+ @skipIfUnsupportMultiNPU(4)
		115	+ @with_comms
		116	+ @parametrize(
		117	+ "rotary_mode,input_layout,sin_cos_layout",
		118	+ [
		119	+ ("half", "BNSD", "11SD"),
		120	+ ("half", "BNSD", "B1SD"),
		121	+ ("half", "BNSD", "BNSD"),
		122	+ ("half", "BSND", "1S1D"),
		123	+ ("half", "BSND", "BS1D"),
		124	+ ("half", "BSND", "BSND"),
		125	+ ("half", "SBND", "S11D"),
		126	+ ("half", "SBND", "SB1D"),
		127	+ ("half", "SBND", "SBND"),
		128	+ ("interleave", "BNSD", "11SD"),
		129	+ ("interleave", "BSND", "1S1D"),
		130	+ ("interleave", "SBND", "S11D"),
		131	+ ]
		132	+ )
		133	+ def test_npu_rotary_mul_forward(self, rotary_mode, input_layout, sin_cos_layout):
		134	+ device_mesh = self.build_device_mesh()
		135	+
		136	+ B = 8
		137	+ N = 8
		138	+ S = 64
		139	+ D = 32
		140	+ x_shape = get_shape_from_layout(B, N, S, D, input_layout)
		141	+ x = torch.randn(x_shape, dtype=torch.float32, device="npu")
		142	+ sin_cos_shape = get_shape_from_layout(B, N, S, D, sin_cos_layout)
		143	+ sin = torch.randn(sin_cos_shape, dtype=torch.float32, device="npu") * 2 - 1
		144	+ cos = torch.randn(sin_cos_shape, dtype=torch.float32, device="npu") * 2 - 1
		145	+
		146	+ y = torch_npu.npu_rotary_mul(x, cos, sin, rotary_mode=rotary_mode)
		147	+
		148	+ def test_placement_comb(x_placements, sin_placements, cos_placements):
		149	+ dist_x = distribute_tensor(x, device_mesh, x_placements)
		150	+ dist_sin = distribute_tensor(sin, device_mesh, sin_placements)
		151	+ dist_cos = distribute_tensor(cos, device_mesh, cos_placements)
		152	+ dist_y = torch_npu.npu_rotary_mul(dist_x, dist_cos, dist_sin, rotary_mode=rotary_mode)
		153	+ self.assertEqual(dist_y.full_tensor(), y)
		154	+
		155	+ placements = [Shard(0), Shard(1), Shard(2), Replicate()]
		156	+ for placement in placements:
		157	+ if isinstance(placement, Shard) and sin_cos_shape[placement.dim] == 1:
		158	+ test_placement_comb([placement], [Replicate()], [Replicate()])
		159	+ else:
		160	+ test_placement_comb([placement], [placement], [placement])
		161	+
		162	+ @skipIfUnsupportMultiNPU(4)
		163	+ @with_comms
		164	+ @parametrize(
		165	+ "rotary_mode,input_layout,sin_cos_layout",
		166	+ [
		167	+ ("half", "BNSD", "11SD"),
		168	+ ("half", "BNSD", "B1SD"),
		169	+ ("half", "BNSD", "BNSD"),
		170	+ ("half", "BSND", "1S1D"),
		171	+ ("half", "BSND", "BS1D"),
		172	+ ("half", "BSND", "BSND"),
		173	+ ("half", "SBND", "S11D"),
		174	+ ("half", "SBND", "SB1D"),
		175	+ ("half", "SBND", "SBND"),
		176	+ ("interleave", "BNSD", "11SD"),
		177	+ ("interleave", "BSND", "1S1D"),
		178	+ ("interleave", "SBND", "S11D"),
		179	+ ]
		180	+ )
		181	+ def test_npu_rotary_mul_backward(self, rotary_mode, input_layout, sin_cos_layout):
		182	+ device_mesh = self.build_device_mesh()
		183	+
		184	+ B = 8
		185	+ N = 8
		186	+ S = 64
		187	+ D = 32
		188	+ x_shape = get_shape_from_layout(B, N, S, D, input_layout)
		189	+ x = torch.randn(x_shape, dtype=torch.float32, device="npu", requires_grad=True)
		190	+ sin_cos_shape = get_shape_from_layout(B, N, S, D, sin_cos_layout)
		191	+ sin = torch.randn(sin_cos_shape, dtype=torch.float32, device="npu") * 2 - 1
		192	+ cos = torch.randn(sin_cos_shape, dtype=torch.float32, device="npu") * 2 - 1
		193	+ sin.requires_grad = True
		194	+ cos.requires_grad = True
		195	+
		196	+
		197	+ y = torch_npu.npu_rotary_mul(x, cos, sin, rotary_mode=rotary_mode)
		198	+ grad_y = torch.ones_like(y, dtype=torch.float32, device="npu")
		199	+ y.backward(grad_y)
		200	+
		201	+ def test_placement_comb(x_placements, sin_placements, cos_placements):
		202	+ dist_x = distribute_tensor(x, device_mesh, x_placements)
		203	+ dist_sin = distribute_tensor(sin, device_mesh, sin_placements)
		204	+ dist_cos = distribute_tensor(cos, device_mesh, cos_placements)
		205	+ dist_y = torch_npu.npu_rotary_mul(dist_x, dist_cos, dist_sin, rotary_mode=rotary_mode)
		206	+ dist_grad_y = distribute_tensor(grad_y, device_mesh, dist_y.placements)
		207	+ dist_y.backward(dist_grad_y)
		208	+ self.assertEqual(dist_y.full_tensor(), y)
		209	+ self.assertEqual(dist_x.grad.full_tensor(), x.grad)
		210	+ self.assertEqual(dist_sin.grad.full_tensor(), sin.grad)
		211	+ self.assertEqual(dist_cos.grad.full_tensor(), cos.grad)
		212	+
		213	+ placements = [Shard(0), Shard(1), Shard(2), Replicate()]
		214	+ for placement in placements:
		215	+ if isinstance(placement, Shard) and sin_cos_shape[placement.dim] == 1:
		216	+ test_placement_comb([placement], [Replicate()], [Replicate()])
		217	+ else:
		218	+ test_placement_comb([placement], [placement], [placement])
		219	+
		220	+
		221	+instantiate_parametrized_tests(TestMathOps)
		222	+
64		223
65	if __name__ == "__main__":	224	if __name__ == "__main__":
66	run_tests()	225	run_tests()

Atest/distributed/tensor/test_moe_ops.py+185-0

@@ -0,0 +1,185 @@
	1	+import itertools
	2	+
	3	+import torch
	4	+from torch.distributed._tensor import distribute_tensor, Replicate, Shard
	5	+
	6	+import torch_npu
	7	+from torch_npu.testing.testcase import run_tests
	8	+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
	9	+from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
	10	+
	11	+
	12	+class TestMoeOps(NPUDTensorTestBase):
	13	+ @skipIfUnsupportMultiNPU(4)
	14	+ @with_comms
	15	+ def test_npu_moe_token_permute_forward(self):
	16	+ device_mesh = self.build_device_mesh()
	17	+
	18	+ num_tokens = 16
	19	+ hidden_size = 8
	20	+ topk = 4
	21	+ tokens = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="npu")
	22	+ indices = torch.randint(0, 4, (num_tokens, topk), dtype=torch.int32, device="npu")
	23	+
	24	+ permuted_tokens, sorted_indices = torch_npu.npu_moe_token_permute(tokens, indices)
	25	+
	26	+ def test_placement_comb(placements1, placements2):
	27	+ dist_tokens = distribute_tensor(tokens, device_mesh, placements1)
	28	+ dist_indices = distribute_tensor(indices, device_mesh, placements2)
	29	+ dist_permuted_tokens, dist_sorted_indices = torch_npu.npu_moe_token_permute(dist_tokens, dist_indices)
	30	+ self.assertEqual(dist_permuted_tokens.full_tensor(), permuted_tokens)
	31	+ self.assertEqual(dist_sorted_indices.full_tensor(), sorted_indices)
	32	+
	33	+ placement = [Shard(0), Shard(1), Replicate()]
	34	+ placement_combs = itertools.product(placement, placement)
	35	+ for comb in placement_combs:
	36	+ test_placement_comb([comb[0]], [comb[1]])
	37	+
	38	+ @skipIfUnsupportMultiNPU(4)
	39	+ @with_comms
	40	+ def test_npu_moe_token_permute_backward(self):
	41	+ device_mesh = self.build_device_mesh()
	42	+
	43	+ num_tokens = 16
	44	+ hidden_size = 8
	45	+ topk = 4
	46	+ tokens = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="npu", requires_grad=True)
	47	+ indices = torch.randint(0, 4, (num_tokens, topk), dtype=torch.int32, device="npu")
	48	+
	49	+ permuted_tokens, sorted_indices = torch_npu.npu_moe_token_permute(tokens, indices)
	50	+ grad_permuted_tokens = torch.ones_like(permuted_tokens, dtype=torch.bfloat16, device="npu")
	51	+ permuted_tokens.backward(grad_permuted_tokens)
	52	+
	53	+ def test_placement_comb(placements1, placements2):
	54	+ dist_tokens = distribute_tensor(tokens, device_mesh, placements1)
	55	+ dist_indices = distribute_tensor(indices, device_mesh, placements2)
	56	+ dist_permuted_tokens, dist_sorted_indices = torch_npu.npu_moe_token_permute(dist_tokens, dist_indices)
	57	+ dist_grad_permuted_tokens = distribute_tensor(
	58	+ grad_permuted_tokens, device_mesh, dist_permuted_tokens.placements
	59	+ )
	60	+ dist_permuted_tokens.backward(dist_grad_permuted_tokens)
	61	+ self.assertEqual(dist_permuted_tokens.full_tensor(), permuted_tokens)
	62	+ self.assertEqual(dist_sorted_indices.full_tensor(), sorted_indices)
	63	+ self.assertEqual(dist_tokens.grad.full_tensor(), tokens.grad)
	64	+
	65	+ placement = [Shard(0), Shard(1), Replicate()]
	66	+ placement_combs = itertools.product(placement, placement)
	67	+ for comb in placement_combs:
	68	+ test_placement_comb([comb[0]], [comb[1]])
	69	+
	70	+ @skipIfUnsupportMultiNPU(4)
	71	+ @with_comms
	72	+ def test_npu_moe_token_permute_clip(self):
	73	+ device_mesh = self.build_device_mesh()
	74	+
	75	+ num_tokens = 16
	76	+ hidden_size = 8
	77	+ topk = 4
	78	+ num_out_tokens = 10
	79	+ tokens = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="npu", requires_grad=True)
	80	+ indices = torch.randint(0, 4, (num_tokens, topk), dtype=torch.int32, device="npu")
	81	+
	82	+ permuted_tokens, sorted_indices = torch_npu.npu_moe_token_permute(
	83	+ tokens, indices, num_out_tokens=num_out_tokens
	84	+ )
	85	+ permuted_tokens.sum().backward()
	86	+
	87	+ dist_tokens = distribute_tensor(tokens, device_mesh, [Shard(1)])
	88	+ dist_indices = distribute_tensor(indices, device_mesh, [Replicate()])
	89	+ dist_permuted_tokens, dist_sorted_indices = torch_npu.npu_moe_token_permute(
	90	+ dist_tokens, dist_indices, num_out_tokens=num_out_tokens
	91	+ )
	92	+ dist_permuted_tokens.sum().backward()
	93	+
	94	+ self.assertEqual(dist_permuted_tokens.full_tensor(), permuted_tokens)
	95	+ self.assertEqual(dist_sorted_indices.full_tensor(), sorted_indices)
	96	+ self.assertEqual(dist_tokens.grad.full_tensor(), tokens.grad)
	97	+
	98	+ @skipIfUnsupportMultiNPU(4)
	99	+ @with_comms
	100	+ def test_npu_moe_token_unpermute_forward(self):
	101	+ device_mesh = self.build_device_mesh()
	102	+
	103	+ num_tokens = 16
	104	+ hidden_size = 8
	105	+ topk = 4
	106	+ permuted_tokens = torch.randn(num_tokens * topk, hidden_size, dtype=torch.bfloat16, device="npu")
	107	+ sorted_indices = torch.randperm(num_tokens * topk, dtype=torch.int32, device="npu")
	108	+
	109	+ tokens = torch_npu.npu_moe_token_unpermute(permuted_tokens, sorted_indices)
	110	+
	111	+ def test_placement_comb(placements1, placements2):
	112	+ dist_permuted_tokens = distribute_tensor(permuted_tokens, device_mesh, placements1)
	113	+ dist_sorted_indices = distribute_tensor(sorted_indices, device_mesh, placements2)
	114	+ dist_tokens = torch_npu.npu_moe_token_unpermute(dist_permuted_tokens, dist_sorted_indices)
	115	+ self.assertEqual(dist_tokens.full_tensor(), tokens)
	116	+
	117	+ permuted_tokens_placement = [Shard(0), Shard(1), Replicate()]
	118	+ sorted_indices_placement = [Shard(0), Replicate()]
	119	+ placement_combs = itertools.product(permuted_tokens_placement, sorted_indices_placement)
	120	+ for comb in placement_combs:
	121	+ test_placement_comb([comb[0]], [comb[1]])
	122	+
	123	+ @skipIfUnsupportMultiNPU(4)
	124	+ @with_comms
	125	+ def test_npu_moe_token_unpermute_backward(self):
	126	+ device_mesh = self.build_device_mesh()
	127	+
	128	+ num_tokens = 8
	129	+ hidden_size = 4
	130	+ topk = 2
	131	+ permuted_tokens = torch.randn(
	132	+ num_tokens * topk, hidden_size, dtype=torch.bfloat16, device="npu", requires_grad=True
	133	+ )
	134	+ sorted_indices = torch.randperm(num_tokens * topk, dtype=torch.int32, device="npu")
	135	+ probs = torch.randn(num_tokens, topk, dtype=torch.bfloat16, device="npu", requires_grad=True)
	136	+
	137	+ tokens = torch_npu.npu_moe_token_unpermute(permuted_tokens, sorted_indices, probs)
	138	+ grad_tokens = torch.ones_like(tokens, dtype=torch.bfloat16, device="npu")
	139	+ tokens.backward(grad_tokens)
	140	+
	141	+ def test_placement_comb(placements1, placements2):
	142	+ dist_permuted_tokens = distribute_tensor(permuted_tokens, device_mesh, placements1)
	143	+ dist_sorted_indices = distribute_tensor(sorted_indices, device_mesh, placements2)
	144	+ dist_probs = distribute_tensor(probs, device_mesh, [Shard(0)])
	145	+ dist_tokens = torch_npu.npu_moe_token_unpermute(dist_permuted_tokens, dist_sorted_indices, dist_probs)
	146	+ dist_grad_tokens = distribute_tensor(
	147	+ grad_tokens, device_mesh, dist_tokens.placements
	148	+ )
	149	+ dist_tokens.backward(dist_grad_tokens)
	150	+ self.assertEqual(dist_tokens.full_tensor(), tokens)
	151	+ self.assertEqual(dist_permuted_tokens.grad.full_tensor(), permuted_tokens.grad)
	152	+ self.assertEqual(dist_probs.grad.full_tensor(), probs.grad)
	153	+
	154	+ permuted_tokens_placement = [Shard(0), Shard(1), Replicate()]
	155	+ sorted_indices_placement = [Shard(0), Replicate()]
	156	+ placement_combs = itertools.product(permuted_tokens_placement, sorted_indices_placement)
	157	+ for comb in placement_combs:
	158	+ test_placement_comb([comb[0]], [comb[1]])
	159	+
	160	+ @skipIfUnsupportMultiNPU(4)
	161	+ @with_comms
	162	+ def test_npu_moe_token_permute_unpermute(self):
	163	+ device_mesh = self.build_device_mesh()
	164	+
	165	+ num_tokens = 16
	166	+ hidden_size = 8
	167	+ topk = 4
	168	+ tokens = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="npu")
	169	+ indices = torch.randint(0, 4, (num_tokens, topk), dtype=torch.int32, device="npu")
	170	+
	171	+ permuted_tokens, sorted_indices = torch_npu.npu_moe_token_permute(tokens, indices)
	172	+ reconstruct_tokens = torch_npu.npu_moe_token_unpermute(permuted_tokens, sorted_indices)
	173	+
	174	+ dist_tokens = distribute_tensor(tokens, device_mesh, [Shard(1)])
	175	+ dist_indices = distribute_tensor(indices, device_mesh, [Replicate()])
	176	+ dist_permuted_tokens, dist_sorted_indices = torch_npu.npu_moe_token_permute(dist_tokens, dist_indices)
	177	+ dist_reconstruct_tokens = torch_npu.npu_moe_token_unpermute(dist_permuted_tokens, dist_sorted_indices)
	178	+
	179	+ self.assertEqual(dist_permuted_tokens.full_tensor(), permuted_tokens)
	180	+ self.assertEqual(dist_sorted_indices.full_tensor(), sorted_indices)
	181	+ self.assertEqual(dist_reconstruct_tokens.full_tensor(), reconstruct_tokens)
	182	+
	183	+
	184	+if __name__ == "__main__":
	185	+ run_tests()

Mtest/distributed/test_zero_redundancy_optimizer.py+1-1

	@@ -1153,7 +1153,7 @@ class TestZeroRedundancyOptimizerDistributed(TestZeroRedundancyOptimizer):
1153	"parameters_as_bucket_view",	1153	"parameters_as_bucket_view",
1154	[False, True],	1154	[False, True],
1155	)	1155	)
1156	- @skipIfUnsupportMultiNPU(2)	1156	+ @skipIfUnsupportMultiNPU(4)
1157	def test_zero_model_parallel(	1157	def test_zero_model_parallel(
1158	self,	1158	self,
1159	parameters_as_bucket_view: bool,	1159	parameters_as_bucket_view: bool,

Mtorch_npu/distributed/tensor/__init__.py+1-0

	@@ -1,3 +1,4 @@
1	import torch_npu.distributed.tensor._matrix_ops	1	import torch_npu.distributed.tensor._matrix_ops
2	import torch_npu.distributed.tensor._attention	2	import torch_npu.distributed.tensor._attention
3	import torch_npu.distributed.tensor._math_ops	3	import torch_npu.distributed.tensor._math_ops
		4	+import torch_npu.distributed.tensor._moe_ops

Mtorch_npu/distributed/tensor/_math_ops.py+146-0

	@@ -18,6 +18,8 @@ from torch.distributed.tensor._ops._math_ops import (
18	_infer_reduce_dims_map,	18	_infer_reduce_dims_map,
19	map_placements_after_reduction)	19	map_placements_after_reduction)
20	from torch.distributed.tensor._utils import normalize_to_torch_size	20	from torch.distributed.tensor._utils import normalize_to_torch_size
		21	+from torch.distributed.tensor import Partial, Replicate, Shard
		22	+from torch.distributed.tensor.experimental import register_sharding
21		23
22	npu = torch.ops.npu	24	npu = torch.ops.npu
23		25
	@@ -158,3 +160,147 @@ def npu_rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
158	)	160	)
159		161
160	return out_tuple_strategy	162	return out_tuple_strategy
		163	+
		164	+
		165	+@register_op_strategy(npu.npu_add_rms_norm.default)
		166	+def npu_add_rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
		167	+ # func: npu_add_rms_norm(Tensor x1, Tensor x2, Tensor gamma, float epsilon=1e-06) -> (Tensor, Tensor, Tensor)
		168	+ mesh = op_schema.get_mesh_from_args(validate=False)
		169	+ expected_args_len = 3
		170	+ (
		171	+ x1_strategy,
		172	+ x2_strategy,
		173	+ gamma_strategy,
		174	+ ) = op_schema.args_schema[:expected_args_len]
		175	+
		176	+ normalized_shape = gamma_strategy.shape
		177	+ normalized_size = normalize_to_torch_size(normalized_shape)
		178	+
		179	+ # x1, x2 should have the same shape
		180	+ input_ndim = x1_strategy.ndim
		181	+ axis = input_ndim - len(normalized_size)
		182	+
		183	+ output_strategy = OpStrategy([])
		184	+ for idx, x1_placement_strategy in enumerate(x1_strategy.strategies):
		185	+ op_args_target_specs = []
		186	+ redistribute_costs = []
		187	+
		188	+ # dims to calculate rms norm should be Replicate
		189	+ x1_src_spec = x1_placement_strategy.output_spec
		190	+ x1_target_placements = _replicate_dims_start_at(x1_src_spec.placements, axis)
		191	+ x1_target_spec = DTensorSpec(
		192	+ mesh=mesh,
		193	+ placements=x1_target_placements,
		194	+ tensor_meta=x1_src_spec.tensor_meta,
		195	+ )
		196	+ op_args_target_specs.append(x1_target_spec)
		197	+ redistribute_costs.append(
		198	+ generate_redistribute_costs(x1_strategy, x1_target_spec)
		199	+ )
		200	+
		201	+ # x2 follows x1
		202	+ if x2_strategy is not None:
		203	+ x2_src_spec = x2_strategy.strategies[idx].output_spec
		204	+ x2_target_spec = DTensorSpec(
		205	+ mesh,
		206	+ placements=x1_target_placements,
		207	+ tensor_meta=x2_src_spec.tensor_meta,
		208	+ )
		209	+ op_args_target_specs.append(x2_target_spec)
		210	+ redistribute_costs.append(
		211	+ generate_redistribute_costs(x2_strategy, x2_target_spec)
		212	+ )
		213	+
		214	+ if gamma_strategy is not None:
		215	+ gamma_src_spec = gamma_strategy.strategies[idx].output_spec
		216	+ gamma_target_spec = DTensorSpec(
		217	+ mesh=mesh,
		218	+ placements=_replicate_dims_start_at(gamma_src_spec.placements),
		219	+ tensor_meta=gamma_src_spec.tensor_meta,
		220	+ )
		221	+ op_args_target_specs.append(gamma_target_spec)
		222	+ redistribute_costs.append(
		223	+ generate_redistribute_costs(gamma_strategy, gamma_target_spec)
		224	+ )
		225	+
		226	+ y_target_spec = x1_target_spec
		227	+ rstd_target_spec = DTensorSpec(
		228	+ mesh=mesh,
		229	+ placements=x1_target_placements[:axis],
		230	+ tensor_meta=x1_src_spec.tensor_meta,
		231	+ )
		232	+ x_target_spec = x1_target_spec
		233	+ output_target_spec = (y_target_spec, rstd_target_spec, x_target_spec)
		234	+
		235	+ output_strategy.strategies.append(
		236	+ OpSpec(
		237	+ output_specs=output_target_spec,
		238	+ input_specs=op_args_target_specs,
		239	+ redistribute_cost=redistribute_costs,
		240	+ )
		241	+ )
		242	+
		243	+ return output_strategy
		244	+
		245	+
		246	+@register_sharding(npu.npu_rotary_mul.default)
		247	+def npu_rotary_mul_strategy(x, r1, r2, rotary_mode="half"):
		248	+ # func: npu_rotary_mul(Tensor self, Tensor r1, Tensor r2, str rotary_mode='half') -> Tensor
		249	+ acceptable_shardings = []
		250	+
		251	+ # all replicate strategy
		252	+ replicate_strategy = (
		253	+ [Replicate()], # output
		254	+ [Replicate(), Replicate(), Replicate(), None] # x, r1, r2, rotary_mode
		255	+ )
		256	+ acceptable_shardings.append(replicate_strategy)
		257	+
		258	+ # sharding strategy
		259	+ # for any layout of x, the last dim always be D, which is not shardable
		260	+ for i in range(x.ndim - 1):
		261	+ # x, r1, r2 have the same layout, while B/N in r1/r2 can be 1, r1 and r2 have the same shape
		262	+ if r1.shape[i] == 1:
		263	+ sharding_strategy = (
		264	+ [Shard(i)],
		265	+ [Shard(i), Replicate(), Replicate(), None]
		266	+ )
		267	+ else:
		268	+ sharding_strategy = (
		269	+ [Shard(i)],
		270	+ [Shard(i), Shard(i), Shard(i), None]
		271	+ )
		272	+ acceptable_shardings.append(sharding_strategy)
		273	+
		274	+ return acceptable_shardings
		275	+
		276	+
		277	+@register_sharding(npu.npu_rotary_mul_backward.default)
		278	+def npu_rotary_mul_backward_strategy(grad, x, r1, r2, rotary_mode='half'):
		279	+ # func: npu_rotary_mul_backward(Tensor grad, Tensor self, Tensor r1, Tensor r2, str rotary_mode='half')
		280	+ # -> (Tensor, Tensor, Tensor)
		281	+ acceptable_shardings = []
		282	+
		283	+ # all replicate strategy
		284	+ replicate_strategy = (
		285	+ [Replicate(), Replicate(), Replicate()], # dx, d_r1, d_r2
		286	+ [Replicate(), Replicate(), Replicate(), Replicate(), None] # grad, x, r1, r2, rotary_mode
		287	+ )
		288	+ acceptable_shardings.append(replicate_strategy)
		289	+
		290	+ # sharding strategy
		291	+ # for any layout of x, the last dim always be D, which is not shardable
		292	+ for i in range(x.ndim - 1):
		293	+ # x, r1, r2 have the same layout, while B/N in r1, r2 can be 1, r1 and r2 have the same shape
		294	+ if r1.shape[i] == 1:
		295	+ sharding_strategy = (
		296	+ [Shard(i), Partial(), Partial()],
		297	+ [Shard(i), Shard(i), Replicate(), Replicate(), None]
		298	+ )
		299	+ else:
		300	+ sharding_strategy = (
		301	+ [Shard(i), Shard(i), Shard(i)],
		302	+ [Shard(i), Shard(i), Shard(i), Shard(i), None]
		303	+ )
		304	+ acceptable_shardings.append(sharding_strategy)
		305	+
		306	+ return acceptable_shardings

Atorch_npu/distributed/tensor/_moe_ops.py+100-0

@@ -0,0 +1,100 @@
	1	+import torch
	2	+from torch.distributed._tensor import Partial, Replicate, Shard
	3	+from torch.distributed._tensor.experimental import register_sharding
	4	+
	5	+npu = torch.ops.npu
	6	+
	7	+
	8	+@register_sharding(npu.npu_moe_token_permute.default)
	9	+def npu_moe_token_permute_strategy(tokens, indices, num_out_tokens=None, padded_mode=False):
	10	+ # func: npu_moe_token_permute(Tensor tokens, Tensor indices, int? num_out_tokens=None, bool padded_mode=False)
	11	+ # -> (Tensor, Tensor)
	12	+ strategies = []
	13	+
	14	+ # all replicate strategy
	15	+ replicate_strategy = (
	16	+ [Replicate(), Replicate()], # output
	17	+ [Replicate(), Replicate(), None, None] # input
	18	+ )
	19	+ strategies.append(replicate_strategy)
	20	+
	21	+ # hidden_size dim sharding strategy
	22	+ hidden_size_sharding_strategy = (
	23	+ [Shard(1), Replicate()],
	24	+ [Shard(1), Replicate(), None, None]
	25	+ )
	26	+ strategies.append(hidden_size_sharding_strategy)
	27	+
	28	+ return strategies
	29	+
	30	+
	31	+@register_sharding(npu.npu_moe_token_permute_grad.default)
	32	+def npu_moe_token_permute_grad_strategy(tokens, grad_permuted_tokens, indices, sorted_indices, padded_mode=False):
	33	+ # func: npu_moe_token_permute_grad(Tensor tokens, Tensor grad_permuted_tokens, Tensor indices,
	34	+ # Tensor sorted_indices, bool padded_mode=False) -> Tensor
	35	+ strategies = []
	36	+
	37	+ # all replicate strategy
	38	+ replicate_strategy = (
	39	+ [Replicate()], # output
	40	+ [Replicate(), Replicate(), Replicate(), Replicate(), None] # input
	41	+ )
	42	+ strategies.append(replicate_strategy)
	43	+
	44	+ # hidden_size dim sharding strategy
	45	+ hidden_size_sharding_strategy = (
	46	+ [Shard(1)],
	47	+ [Shard(1), Shard(1), Replicate(), Replicate(), None]
	48	+ )
	49	+ strategies.append(hidden_size_sharding_strategy)
	50	+
	51	+ return strategies
	52	+
	53	+
	54	+@register_sharding(npu.npu_moe_token_unpermute.default)
	55	+def npu_moe_token_unpermute_strategy(permuted_tokens, sorted_indices, probs=None, padded_mode=False,
	56	+ restore_shape=None):
	57	+ # func: npu_moe_token_unpermute(Tensor permuted_tokens, Tensor sorted_indices, Tensor? probs=None,
	58	+ # bool padded_mode=False, int[]? restore_shape=None) -> Tensor
	59	+ strategies = []
	60	+
	61	+ # all replicate strategy
	62	+ replicate_strategy = (
	63	+ [Replicate()], # output
	64	+ [Replicate(), Replicate(), None if probs is None else Replicate(), None, None] # input
	65	+ )
	66	+ strategies.append(replicate_strategy)
	67	+
	68	+ # hidden_size dim sharding strategy
	69	+ hidden_size_sharding_strategy = (
	70	+ [Shard(1)],
	71	+ [Shard(1), Replicate(), None if probs is None else Replicate(), None, None]
	72	+ )
	73	+ strategies.append(hidden_size_sharding_strategy)
	74	+
	75	+ return strategies
	76	+
	77	+
	78	+@register_sharding(npu.npu_moe_token_unpermute_grad.default)
	79	+def npu_moe_token_unpermute_grad_strategy(permuted_tokens, grad_unpermuted_tokens, sorted_indices, probs=None,
	80	+ padded_mode=False, restore_shape=None):
	81	+ # func: npu_moe_token_unpermute_grad(Tensor permuted_tokens, Tensor grad_unpermuted_tokens, Tensor sorted_indices,
	82	+ # Tensor? probs=None, bool padded_mode=False, int[]? restore_shape=None)
	83	+ # -> (Tensor, Tensor)
	84	+ strategies = []
	85	+
	86	+ # all replicate strategy
	87	+ replicate_strategy = (
	88	+ [Replicate(), None if probs is None else Replicate()], # permuted_tokens grad, probs grad
	89	+ [Replicate(), Replicate(), Replicate(), None if probs is None else Replicate(), None, None] # input
	90	+ )
	91	+ strategies.append(replicate_strategy)
	92	+
	93	+ # hidden_size dim sharding strategy
	94	+ hidden_size_sharding_strategy = (
	95	+ [Shard(1), Partial()],
	96	+ [Shard(1), Shard(1), Replicate(), None if probs is None else Replicate(), None, None]
	97	+ )
	98	+ strategies.append(hidden_size_sharding_strategy)
	99	+
	100	+ return strategies

Mtorch_npu/testing/common_distributed.py+3-2

	@@ -31,10 +31,11 @@ TEST_SKIPS = {
31		31
32	def skipIfUnsupportMultiNPU(npu_number_needed):	32	def skipIfUnsupportMultiNPU(npu_number_needed):
33	def skip_dec(func):	33	def skip_dec(func):
34	- def wrapper(self):	34	+ @wraps(func)
		35	+ def wrapper(self, args, *kwargs):
35	if not torch.npu.is_available() or torch.npu.device_count() < npu_number_needed:	36	if not torch.npu.is_available() or torch.npu.device_count() < npu_number_needed:
36	raise unittest.SkipTest(f"Multi-NPU {npu_number_needed} condition not satisfied")	37	raise unittest.SkipTest(f"Multi-NPU {npu_number_needed} condition not satisfied")
37	- return func(self)	38	+ return func(self, args, *kwargs)
38	return wrapper	39	return wrapper
39	return skip_dec	40	return skip_dec
40		41