feat: add distributed op RotaryPositionEmbedding with UT and MindSpore ST-hyper-parallel-AtomGit

已合并

feat: add distributed op RotaryPositionEmbedding with UT and MindSpore ST #661

hedongdong创建于 5月19日从david-he91/hyper-parallel： master_dist_op 合入到 mindspore/hyper-parallel： master

feat: add distributed op RotaryPositionEmbedding with UT and MindSpore ST #661

已合并

hedongdong创建于 5月19日

从david-he91/hyper-parallel master_dist_op合入到mindspore/hyper-parallel master

共 5 个文件变更+1177-0

Ahyper_parallel/core/shard/ops/parallel_rotary_position_embedding.py+157-0

@@ -0,0 +1,157 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Distributed implementation for RotaryPositionEmbedding operator."""
	16	+import copy
	17	+from typing import Optional, Tuple
	18	+
	19	+from .parallel_ops import DistributedOp
	20	+
	21	+
	22	+def _normalize_rpe_args(x, cos, sin, mode=0):
	23	+ """Normalize positional and keyword arguments into a canonical positional tuple.
	24	+
	25	+ Args:
	26	+ x: Input tensor.
	27	+ cos: Cosine position encoding tensor.
	28	+ sin: Sine position encoding tensor.
	29	+ mode: Rotation mode. 0=rotate_half, 1=rotate_interleaved, 2=quarter,
	30	+ 3=interleave-half. Defaults to 0.
	31	+
	32	+ Returns:
	33	+ tuple: (positional_args_tuple, empty_kwargs_dict)
	34	+ """
	35	+ return (x, cos, sin, mode), {}
	36	+
	37	+
	38	+class RotaryPositionEmbeddingDistributedOp(DistributedOp):
	39	+ """Distributed operator for RotaryPositionEmbedding.
	40	+
	41	+ Computes rotary position embedding element-wise:
	42	+ y = x * cos + x_rotate * sin
	43	+
	44	+ where x_rotate is obtained by rotating within the last (D) dimension.
	45	+ Output shape equals x shape exactly.
	46	+
	47	+ Sharding constraints:
	48	+ - D (last dim) must be replicated for x, cos, and sin: the kernel rotates
	49	+ within D and the operation is indivisible along that axis.
	50	+ - B, N, S dims are fully independent across positions and can be freely
	51	+ sharded.
	52	+ - cos/sin may have any subset of non-D dims replicated (broadcast case),
	53	+ but if cos/sin is sharded on a dimension, it must match x's sharding
	54	+ on that dimension.
	55	+
	56	+ MODE does not affect layout inference: all modes produce output shape == x
	57	+ shape and leave B/N/S independence unchanged. MODE is passed through as a
	58	+ kernel parameter.
	59	+
	60	+ Output:
	61	+ Single tensor with the same shape and layout as x.
	62	+ """
	63	+
	64	+ @staticmethod
	65	+ def _validate_input_layouts(x_layout, cos_layout, sin_layout) -> None:
	66	+ """Validate sharding constraints for all input tensors.
	67	+
	68	+ Rules (applied to both 4-D BNSD/BSND/SBND and 3-D TND layouts):
	69	+ - x's last dim (D) must be replicated.
	70	+ - cos and sin's last dim (D) must be replicated.
	71	+ - For any non-D dimension d: if cos/sin is sharded there, the mesh
	72	+ axis must equal x's mesh axis on the same dimension.
	73	+
	74	+ Args:
	75	+ x_layout: Layout of the x tensor.
	76	+ cos_layout: Layout of the cos tensor.
	77	+ sin_layout: Layout of the sin tensor.
	78	+
	79	+ Raises:
	80	+ ValueError: If D is sharded for any input, or if cos/sin sharding
	81	+ is inconsistent with x on any non-D dimension.
	82	+ """
	83	+ op = "rotary_position_embedding"
	84	+ x_tm = x_layout.tensor_map
	85	+
	86	+ if x_tm[-1] != -1:
	87	+ raise ValueError(
	88	+ f"For {op}, D (last dim) of x must be replicated, "
	89	+ f"but got tensor_map={x_tm}"
	90	+ )
	91	+
	92	+ for name, layout in [('cos', cos_layout), ('sin', sin_layout)]:
	93	+ tm = layout.tensor_map
	94	+ if tm[-1] != -1:
	95	+ raise ValueError(
	96	+ f"For {op}, D (last dim) of {name} must be replicated, "
	97	+ f"but got tensor_map={tm}"
	98	+ )
	99	+ for d in range(len(tm) - 1):
	司小南（机器人）5月19日 🚨 [关键问题] cos/sin 分片验证逻辑不完整，可能漏检非法分片问题陈述： cos/sin 分片参数的校验逻辑存在漏洞，仅验证了分片参数的基础有效性（如非负或非空），却遗漏了对分片完整性（Coverage）和与输入张量维度一致性（Consistency）的校验，导致非法分片配置可能绕过检查进入计算流程。证据支撑：推断校验逻辑处（如 `rotary_embedding_kernel.cpp` 或相关校验函数中）仅包含形如 `assert(split_size > 0)` 的基础检查，未对分片总和进行约束缺失关键校验逻辑：`assert(accumulate(split_sizes) == input_dim)` 以及 `assert(cos.shape == expected_split_shape)` 问题影响：若分片总和小于输入维度，部分数据将被静默忽略，导致计算结果错误；若分片总和大于输入维度，将引发内存越界访问；若 cos/sin 形状与分片不匹配，将导致计算结果非预期修改建议：完善校验逻辑，增加分片总和与输入维度的对齐检查，以及 cos/sin 张量形状与分片配置的匹配检查： `// 修复建议示例 // 1. 校验分片总和覆盖整个维度 int total_split = std::accumulate(split_sizes.begin(), split_sizes.end(), 0); CHECK_EQ(total_split, hidden_dim) << "Split sizes must sum to hidden dimension (" << hidden_dim << ")"; // 2. 校验 cos/sin 张量形状与分片配置匹配 for (size_t i = 0; i < split_sizes.size(); ++i) { CHECK_EQ(cos_tensors[i].size(0), split_sizes[i]) << "Cos tensor " << i << " size mismatch"; }`
	100	+ x_d = x_tm[d] if d < len(x_tm) - 1 else -1
	101	+ if tm[d] != -1 and tm[d] != x_d:
	102	+ raise ValueError(
	103	+ f"For {op}, {name} sharding on dim {d} must match x or be replicated, "
	104	+ f"but got x={x_d}, {name}={tm[d]}"
	105	+ )
	106	+
	107	+ def preprocess(self, args: tuple, kwargs: dict) -> Optional[tuple]:
	108	+ """Extract local tensors and build the layout cache.
	109	+
	110	+ All inputs (x, cos, sin) are expected to be DTensors.
	111	+
	112	+ Args:
	113	+ args: Positional arguments, may include DTensors.
	114	+ kwargs: Keyword arguments.
	115	+
	116	+ Returns:
	117	+ tuple: (local_args, local_kwargs, cache_values) where
	118	+ local_args = (x_local, cos_local, sin_local, mode),
	119	+ local_kwargs = {},
	120	+ cache_values = [x_layout, cos_layout, sin_layout].
	121	+ """
	122	+ norm_args, _ = _normalize_rpe_args(args, *kwargs)
	123	+ x = norm_args[0]
	124	+ cos = norm_args[1]
	125	+ sin = norm_args[2]
	126	+ mode = norm_args[3]
	127	+
	128	+ local_args = (x.to_local(), cos.to_local(), sin.to_local(), mode)
	129	+ cache_values = [x.layout, cos.layout, sin.layout]
	130	+ return local_args, {}, cache_values
	131	+
	132	+ def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
	133	+ """Infer output layout for the single output tensor.
	134	+
	135	+ Rules:
	136	+ 1. Partial inputs are rejected.
	137	+ 2. D (last dim) must be replicated for x, cos, and sin.
	138	+ 3. cos/sin sharding on non-D dims must match x or be replicated.
	139	+ 4. Output layout = deep copy of x_layout (output shape == x shape).
	140	+
	141	+ Args:
	142	+ cache_values: [x_layout, cos_layout, sin_layout]
	143	+
	144	+ Returns:
	145	+ tuple: ((output_layout,), None)
	146	+
	147	+ Raises:
	148	+ ValueError: If any input has Partial status, D is sharded,
	149	+ or cos/sin sharding is inconsistent with x.
	150	+ """
	151	+ x_layout = cache_values[0]
	152	+ cos_layout = cache_values[1]
	153	+ sin_layout = cache_values[2]
	154	+
	155	+ self._check_partial_inputs([x_layout, cos_layout, sin_layout])
	156	+ self._validate_input_layouts(x_layout, cos_layout, sin_layout)
	157	+ return (copy.deepcopy(x_layout),), None

Ahyper_parallel/core/shard/ops/yaml/rotary_position_embedding_ops.yaml+4-0

@@ -0,0 +1,4 @@
	1	+RotaryPositionEmbedding:
	2	+ dist_op_name: _rotary_position_embedding_dist_op
	3	+ distributed_op_class: RotaryPositionEmbeddingDistributedOp
	4	+ distributed_op_file: parallel_rotary_position_embedding

Atests/mindspore/st/shard/ops/rotary_position_embedding_shard_in_python.py+418-0

@@ -0,0 +1,418 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""MindSpore distributed ST for RotaryPositionEmbedding operator.
	16	+
	17	+Formula: y = x * cos + x_rotate * sin, where x_rotate rotates within the last
	18	+(D) dimension. All B/N/S positions are element-wise independent, so they can
	19	+be freely sharded. D must remain replicated.
	20	+
	21	+Shapes used throughout (BNSD layout):
	22	+ B=4, N=4, S=16, D=64 → BN=16 ≤ S8=128 (CANN float16 constraint). ✓
	23	+
	24	+cos/sin broadcast variants:
	25	+ full-shape: (B, N, S, D) — shards on the same dims as x.
	26	+ broadcast: (1, 1, S, D) — replicated on B/N; may shard on S when x shards S.
	27	+
	28	+Precision tolerance:
	29	+ atol=rtol=1e-3. RPE is pure element-wise with no cross-rank reduction in the
	30	+ forward pass, so distributed and standalone invoke the same NPU kernel on the
	31	+ same data — expected max_diff ≈ 0. For backward Partial("sum") gradients the
	32	+ all-reduce adds at most one float16 rounding step (epsilon_f16 ≈ 1e-3), so
	33	+ 1e-3 is both correct and tight.
	34	+
	35	+mode values exercised:
	36	+ 0=rotate_half, 1=rotate_interleaved, 2=quarter, 3=interleave-half.
	37	+ CANN backward is only supported for mode=0 and mode=1. mode=2/3 are
	38	+ validated in forward-only tests (_run_standalone_fwd, no ms.grad call).
	39	+"""
	40	+import numpy as np
	41	+import mindspore as ms
	42	+import mindspore.communication.management as D
	43	+from mindspore import Tensor
	44	+from mindspore.ops import rotary_position_embedding
	45	+
	46	+from hyper_parallel import init_device_mesh
	47	+from hyper_parallel.core.dtensor.dtensor import DTensor, distribute_tensor
	48	+from hyper_parallel.core.dtensor.placement_types import Shard, Replicate, Partial
	49	+
	50	+
	51	+def setup_module():
	52	+ """Initialize MindSpore context and communication."""
	53	+ ms.context.set_context(mode=ms.context.PYNATIVE_MODE, device_target="Ascend")
	54	+ D.init()
	55	+
	56	+
	57	+np.random.seed(42)
	58	+ms.set_seed(42)
	59	+ms.set_deterministic(True)
	60	+
	61	+B = 4
	62	+N = 4
	63	+S = 16
	64	+D_DIM = 64
	65	+
	66	+_INPUT_NAMES = ("x", "cos", "sin")
	67	+_TOL = {"atol": 1e-3, "rtol": 1e-3}
	68	+
	69	+
	70	+def _np_to_f16(arr: np.ndarray) -> Tensor:
	71	+ """Convert float32 numpy array to MindSpore float16 Tensor."""
	72	+ return Tensor(arr, ms.float16)
	73	+
	74	+
	75	+def _to_f32np(t) -> np.ndarray:
	76	+ """Convert any MindSpore tensor to float32 numpy array."""
	77	+ return t.astype(ms.float32).asnumpy()
	78	+
	79	+
	80	+def _make_rpe(mode: int):
	81	+ """Return a 3-argument RPE callable with the given rotation mode captured.
	82	+
	83	+ Args:
	84	+ mode: Rotation mode (0=rotate_half, 1=rotate_interleaved,
	85	+ 2=quarter, 3=interleave-half).
	86	+
	87	+ Returns:
	88	+ Callable (x, cos, sin) -> y that ms.grad can differentiate for mode=0/1.
	89	+ """
	90	+ def _fn(x, cos, sin):
	91	+ return rotary_position_embedding(x, cos, sin, mode=mode)
	92	+ return _fn
	93	+
	94	+
	95	+def _run_standalone_fwd(x_np, cos_np, sin_np, mode=0):
	96	+ """Compute standalone RPE forward pass only.
	97	+
	98	+ Used for mode=2 and mode=3 whose CANN backward kernels are not available.
	99	+
	100	+ Args:
	101	+ x_np: float32 ndarray, shape (B, N, S, D).
	102	+ cos_np: float32 ndarray, shape (B, N, S, D) or (1, 1, S, D).
	103	+ sin_np: float32 ndarray, same shape as cos_np.
	104	+ mode: RPE rotation mode.
	105	+
	106	+ Returns:
	107	+ fwd_np: float32 ndarray, forward output.
	108	+ """
	109	+ return _to_f32np(_make_rpe(mode)(_np_to_f16(x_np), _np_to_f16(cos_np), _np_to_f16(sin_np)))
	110	+
	111	+
	112	+def _run_standalone(x_np, cos_np, sin_np, mode=0):
	113	+ """Compute standalone RPE forward and backward.
	114	+
	115	+ Use only for mode=0 and mode=1, which have CANN backward support.
	116	+
	117	+ Args:
	118	+ x_np: float32 ndarray, shape (B, N, S, D).
	119	+ cos_np: float32 ndarray, shape (B, N, S, D) or (1, 1, S, D).
	120	+ sin_np: float32 ndarray, same shape as cos_np.
	121	+ mode: RPE rotation mode (must be 0 or 1).
	122	+
	123	+ Returns:
	124	+ tuple: (fwd_np, (grad_x_np, grad_cos_np, grad_sin_np)) — all float32.
	125	+ """
	126	+ x = _np_to_f16(x_np)
	127	+ cos = _np_to_f16(cos_np)
	128	+ sin = _np_to_f16(sin_np)
	129	+ fn = _make_rpe(mode)
	130	+ fwd_np = _to_f32np(fn(x, cos, sin))
	131	+ grads = ms.grad(fn, (0, 1, 2))(x, cos, sin)
	132	+ return fwd_np, tuple(_to_f32np(g) for g in grads)
	133	+
	134	+
	135	+def _get_grad_placements(d_inputs: tuple) -> list:
	136	+ """Derive the correct DTensor placement for each input's gradient.
	137	+
	138	+ Rule: if any input is Shard on mesh dimension d, then a Replicate input on
	139	+ that dimension receives gradient contributions from all ranks → Partial("sum").
	140	+ Shard inputs keep their Shard placement.
	141	+
	142	+ Args:
	143	+ d_inputs: Tuple of DTensor inputs (dx, dcos, dsin).
	144	+
	145	+ Returns:
	146	+ list of placement tuples, one per input.
	147	+ """
	148	+ n_dims = len(d_inputs[0].layout.placements)
	149	+ dim_has_shard = [
	150	+ any(isinstance(d_inp.layout.placements[d], Shard) for d_inp in d_inputs)
	151	+ for d in range(n_dims)
	152	+ ]
	153	+ result = []
	154	+ for d_inp in d_inputs:
	155	+ grad_placements = tuple(
	156	+ Partial("sum") if isinstance(p, Replicate) and dim_has_shard[dim_idx] else p
	157	+ for dim_idx, p in enumerate(d_inp.layout.placements)
	158	+ )
	159	+ result.append(grad_placements)
	160	+ return result
	161	+
	162	+
	163	+def _assert_fwd(dist_out, ref_np, tag):
	164	+ """Assert gathered distributed forward output matches standalone reference."""
	165	+ dist_np = _to_f32np(dist_out.full_tensor())
	166	+ assert np.allclose(dist_np, ref_np, **_TOL), (
	167	+ f"{tag} fwd mismatch: "
	168	+ f"max_diff={np.abs(dist_np - ref_np).max()}"
	169	+ )
	170	+
	171	+
	172	+def _assert_bwd(raw_grads, d_inputs, ref_grads, tag):
	173	+ """Assert all input gradients match standalone reference.
	174	+
	175	+ Reconstructs each gradient as a DTensor using placements from
	176	+ _get_grad_placements (handles both Shard and Partial("sum") cases), gathers
	177	+ via full_tensor(), and compares to ref_grads.
	178	+
	179	+ Args:
	180	+ raw_grads: Tuple of local gradient tensors from ms.grad.
	181	+ d_inputs: Tuple of DTensor inputs (dx, dcos, dsin).
	182	+ ref_grads: Tuple of float32 reference gradient ndarrays.
	183	+ tag: Human-readable label for assertion messages.
	184	+ """
	185	+ grad_placements = _get_grad_placements(d_inputs)
	186	+ for name, raw_g, d_inp, ref, g_pl in zip(
	187	+ _INPUT_NAMES, raw_grads, d_inputs, ref_grads, grad_placements):
	188	+ dist_g = (
	189	+ DTensor.from_local(raw_g, d_inp.layout.mesh, g_pl)
	190	+ .full_tensor()
	191	+ .astype(ms.float32)
	192	+ .asnumpy()
	193	+ )
	194	+ assert np.allclose(dist_g, ref, **_TOL), (
	195	+ f"{tag} backward grad_{name} mismatch: "
	196	+ f"max_diff={np.abs(dist_g - ref).max()}"
	197	+ )
	198	+
	199	+
	200	+# ---------------------------------------------------------------------------
	201	+# 2-card tests
	202	+# ---------------------------------------------------------------------------
	203	+
	204	+def test_rpe_replicated():
	205	+ """
	206	+ Feature: RotaryPositionEmbedding (MindSpore) all replicated, 2-card dp mesh.
	207	+ Description:
	208	+ - 1-D dp mesh (2 cards); x/cos/sin all Replicate.
	209	+ - cos/sin shape (1, 1, S, D) — broadcast over B and N.
	210	+ - float16, mode=0 (rotate_half).
	211	+ Expectation: Distributed forward matches standalone within tolerance.
	212	+ """
	213	+ mode = 0
	214	+ x_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	215	+ cos_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	216	+ sin_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	217	+ ref_fwd = _run_standalone_fwd(x_np, cos_np, sin_np, mode=mode)
	218	+
	219	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("dp",))
	220	+ dx = distribute_tensor(_np_to_f16(x_np), mesh, (Replicate(),))
	221	+ dcos = distribute_tensor(_np_to_f16(cos_np), mesh, (Replicate(),))
	222	+ dsin = distribute_tensor(_np_to_f16(sin_np), mesh, (Replicate(),))
	223	+
	224	+ _assert_fwd(_make_rpe(mode)(dx, dcos, dsin), ref_fwd, "replicated mode=0")
	225	+
	226	+
	227	+def test_rpe_dp_b():
	228	+ """
	229	+ Feature: RotaryPositionEmbedding (MindSpore) B-dim data parallel, 2-card dp mesh.
	230	+ Description:
	231	+ - 1-D dp mesh (2 cards); x Shard(0); cos/sin Replicate, broadcast (1,1,S,D).
	232	+ - float16, mode=1 (rotate_interleaved).
	233	+ - Forward verified; all three gradients verified:
	234	+ grad_x → Shard(0); grad_cos/sin → Partial("sum") via _get_grad_placements.
	235	+ Expectation: Distributed forward and all gradients match standalone within tolerance.
	236	+ """
	237	+ mode = 1
	238	+ x_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	239	+ cos_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	240	+ sin_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	241	+ ref_fwd, ref_grads = _run_standalone(x_np, cos_np, sin_np, mode=mode)
	242	+
	243	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("dp",))
	244	+ dx = distribute_tensor(_np_to_f16(x_np), mesh, (Shard(0),))
	245	+ dcos = distribute_tensor(_np_to_f16(cos_np), mesh, (Replicate(),))
	246	+ dsin = distribute_tensor(_np_to_f16(sin_np), mesh, (Replicate(),))
	247	+ d_inputs = (dx, dcos, dsin)
	248	+
	249	+ fn = _make_rpe(mode)
	250	+ _assert_fwd(fn(*d_inputs), ref_fwd, "dp_b mode=1")
	251	+
	252	+ raw_grads = ms.grad(fn, (0, 1, 2))(*d_inputs)
	253	+ _assert_bwd(raw_grads, d_inputs, ref_grads, "dp_b mode=1")
	254	+
	255	+
	256	+def test_rpe_tp_n():
	257	+ """
	258	+ Feature: RotaryPositionEmbedding (MindSpore) N-dim tensor parallel, 2-card tp mesh.
	259	+ Description:
	260	+ - 1-D tp mesh (2 cards); x/cos/sin all Shard(1).
	261	+ - cos/sin full shape (B, N, S, D) — same sharding as x.
	262	+ - float16, mode=0 (rotate_half).
	263	+ - Forward and all three gradients verified; all grad placements = Shard(1).
	264	+ Expectation: Distributed forward and gradients match standalone within tolerance.
	265	+ """
	266	+ mode = 0
	267	+ x_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	268	+ cos_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	269	+ sin_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	270	+ ref_fwd, ref_grads = _run_standalone(x_np, cos_np, sin_np, mode=mode)
	271	+
	272	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2,), mesh_dim_names=("tp",))
	273	+ dx = distribute_tensor(_np_to_f16(x_np), mesh, (Shard(1),))
	274	+ dcos = distribute_tensor(_np_to_f16(cos_np), mesh, (Shard(1),))
	275	+ dsin = distribute_tensor(_np_to_f16(sin_np), mesh, (Shard(1),))
	276	+ d_inputs = (dx, dcos, dsin)
	277	+
	278	+ fn = _make_rpe(mode)
	279	+ _assert_fwd(fn(*d_inputs), ref_fwd, "tp_n mode=0")
	280	+
	281	+ raw_grads = ms.grad(fn, (0, 1, 2))(*d_inputs)
	282	+ _assert_bwd(raw_grads, d_inputs, ref_grads, "tp_n mode=0")
	283	+
	284	+
	285	+# ---------------------------------------------------------------------------
	286	+# 4-card tests
	287	+# ---------------------------------------------------------------------------
	288	+
	289	+def test_rpe_dp_tp():
	290	+ """
	291	+ Feature: RotaryPositionEmbedding (MindSpore) 2-D dp×tp mesh, B by dp, N by tp.
	292	+ Description:
	293	+ - 4-card 2-D mesh (dp=2, tp=2); x/cos/sin all (Shard(0), Shard(1)).
	294	+ - cos/sin full shape (B, N, S, D) — same placements as x.
	295	+ - float16, mode=0 (rotate_half).
	296	+ - Forward and all three gradients verified; all grad placements = (Shard(0), Shard(1)).
	297	+ Expectation: Distributed fwd and all gradients match standalone within tolerance.
	298	+ """
	299	+ mode = 0
	300	+ x_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	301	+ cos_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	302	+ sin_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	303	+ ref_fwd, ref_grads = _run_standalone(x_np, cos_np, sin_np, mode=mode)
	304	+
	305	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
	306	+ placements = (Shard(0), Shard(1))
	307	+ dx = distribute_tensor(_np_to_f16(x_np), mesh, placements)
	308	+ dcos = distribute_tensor(_np_to_f16(cos_np), mesh, placements)
	309	+ dsin = distribute_tensor(_np_to_f16(sin_np), mesh, placements)
	310	+ d_inputs = (dx, dcos, dsin)
	311	+
	312	+ fn = _make_rpe(mode)
	313	+ _assert_fwd(fn(*d_inputs), ref_fwd, "dp_tp mode=0")
	314	+
	315	+ raw_grads = ms.grad(fn, (0, 1, 2))(*d_inputs)
	316	+ _assert_bwd(raw_grads, d_inputs, ref_grads, "dp_tp mode=0")
	317	+
	318	+
	319	+def test_rpe_dp_sp():
	320	+ """
	321	+ Feature: RotaryPositionEmbedding (MindSpore) 2-D dp×sp mesh, B by dp, S by sp.
	322	+ Description:
	323	+ - 4-card 2-D mesh (dp=2, sp=2); x (Shard(0), Shard(2)).
	324	+ - cos/sin broadcast (1, 1, S, D): dp Replicate, sp Shard(2).
	325	+ Local S slices of cos/sin align with the corresponding S slices of x.
	326	+ - float16, mode=2 (quarter); forward only (CANN backward not available for mode=2).
	327	+ Expectation: Distributed forward matches standalone within tolerance.
	328	+ """
	329	+ mode = 2
	330	+ x_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	331	+ cos_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	332	+ sin_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	333	+ ref_fwd = _run_standalone_fwd(x_np, cos_np, sin_np, mode=mode)
	334	+
	335	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "sp"))
	336	+ dx = distribute_tensor(_np_to_f16(x_np), mesh, (Shard(0), Shard(2)))
	337	+ dcos = distribute_tensor(_np_to_f16(cos_np), mesh, (Replicate(), Shard(2)))
	338	+ dsin = distribute_tensor(_np_to_f16(sin_np), mesh, (Replicate(), Shard(2)))
	339	+
	340	+ _assert_fwd(_make_rpe(mode)(dx, dcos, dsin), ref_fwd, "dp_sp mode=2")
	341	+
	342	+
	343	+def test_rpe_tp_sp():
	344	+ """
	345	+ Feature: RotaryPositionEmbedding (MindSpore) 2-D tp×sp mesh, N by tp, S by sp.
	346	+ Description:
	347	+ - 4-card 2-D mesh (tp=2, sp=2); x/cos/sin full shape (B, N, S, D).
	348	+ Placements: (Shard(1), Shard(2)) — tp shards N, sp shards S.
	349	+ - float16, mode=3 (interleave-half); forward only (CANN backward not available for mode=3).
	350	+ Expectation: Distributed forward matches standalone within tolerance.
	351	+ """
	352	+ mode = 3
	353	+ x_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	354	+ cos_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	355	+ sin_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	356	+ ref_fwd = _run_standalone_fwd(x_np, cos_np, sin_np, mode=mode)
	357	+
	358	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("tp", "sp"))
	359	+ placements = (Shard(1), Shard(2))
	360	+ dx = distribute_tensor(_np_to_f16(x_np), mesh, placements)
	361	+ dcos = distribute_tensor(_np_to_f16(cos_np), mesh, placements)
	362	+ dsin = distribute_tensor(_np_to_f16(sin_np), mesh, placements)
	363	+
	364	+ _assert_fwd(_make_rpe(mode)(dx, dcos, dsin), ref_fwd, "tp_sp mode=3")
	365	+
	366	+
	367	+def test_rpe_dp_tp_cos_full():
	368	+ """
	369	+ Feature: RotaryPositionEmbedding (MindSpore) 2-D dp×tp mesh, broadcast cos/sin.
	370	+ Description:
	371	+ - 4-card 2-D mesh (dp=2, tp=2); x (Shard(0), Shard(1)).
	372	+ - cos/sin broadcast (1, 1, S, D): both Replicate on dp and tp.
	373	+ Validates that the kernel broadcasts B=1/N=1 cos/sin against the
	374	+ sharded x without requiring cos/sin to be co-sharded.
	375	+ - float16, mode=2 (quarter); forward only (CANN backward not available for mode=2).
	376	+ Expectation: Distributed forward matches standalone within tolerance.
	377	+ """
	378	+ mode = 2
	379	+ x_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	380	+ cos_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	381	+ sin_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	382	+ ref_fwd = _run_standalone_fwd(x_np, cos_np, sin_np, mode=mode)
	383	+
	384	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
	385	+ dx = distribute_tensor(_np_to_f16(x_np), mesh, (Shard(0), Shard(1)))
	386	+ dcos = distribute_tensor(_np_to_f16(cos_np), mesh, (Replicate(), Replicate()))
	387	+ dsin = distribute_tensor(_np_to_f16(sin_np), mesh, (Replicate(), Replicate()))
	388	+
	389	+ _assert_fwd(_make_rpe(mode)(dx, dcos, dsin), ref_fwd, "dp_tp_cos_full mode=2")
	390	+
	391	+
	392	+# ---------------------------------------------------------------------------
	393	+# 8-card tests
	394	+# ---------------------------------------------------------------------------
	395	+
	396	+def test_rpe_dp_tp_sp():
	397	+ """
	398	+ Feature: RotaryPositionEmbedding (MindSpore) 3-D dp×tp×sp mesh, B/N/S all sharded.
	399	+ Description:
	400	+ - 8-card 3-D mesh (dp=2, tp=2, sp=2).
	401	+ - x (Shard(0), Shard(1), Shard(2)): B by dp, N by tp, S by sp.
	402	+ - cos/sin broadcast (1, 1, S, D): dp/tp Replicate, sp Shard(2).
	403	+ Each rank's cos/sin S-slice aligns with its x S-slice.
	404	+ - float16, mode=3 (interleave-half); forward only (CANN backward not available for mode=3).
	405	+ Expectation: Distributed forward matches standalone within tolerance.
	406	+ """
	407	+ mode = 3
	408	+ x_np = np.random.randn(B, N, S, D_DIM).astype(np.float32)
	409	+ cos_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	410	+ sin_np = np.random.randn(1, 1, S, D_DIM).astype(np.float32)
	411	+ ref_fwd = _run_standalone_fwd(x_np, cos_np, sin_np, mode=mode)
	412	+
	413	+ mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2, 2), mesh_dim_names=("dp", "tp", "sp"))
	414	+ dx = distribute_tensor(_np_to_f16(x_np), mesh, (Shard(0), Shard(1), Shard(2)))
	415	+ dcos = distribute_tensor(_np_to_f16(cos_np), mesh, (Replicate(), Replicate(), Shard(2)))
	416	+ dsin = distribute_tensor(_np_to_f16(sin_np), mesh, (Replicate(), Replicate(), Shard(2)))
	417	+
	418	+ _assert_fwd(_make_rpe(mode)(dx, dcos, dsin), ref_fwd, "dp_tp_sp mode=3")

Atests/mindspore/st/shard/ops/test_rotary_position_embedding_shard_in_python.py+91-0

@@ -0,0 +1,91 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Test runner for RotaryPositionEmbedding distributed ST (MindSpore)."""
	16	+from pathlib import Path
	17	+
	18	+from tests.common.mark_utils import arg_mark
	19	+from tests.common.parallel_case import parallel_run, MindSporeCase
	20	+
	21	+IMPL_FILE = str(
	22	+ Path(__file__).resolve().parent / "rotary_position_embedding_shard_in_python.py"
	23	+)
	24	+
	25	+
	26	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	27	+def test_parallel_rpe_ms_group1():
	28	+ """
	29	+ Feature: parallel run case in rotary_position_embedding_shard_in_python (2-card, 1-D mesh)
	30	+ Description:
	31	+ 1. test_rpe_replicated — all replicated on dp mesh, fwd vs standalone
	32	+ 2. test_rpe_dp_b — B-dim DP (x Shard(0), cos/sin broadcast Replicate),
	33	+ fwd + grad_x vs standalone
	34	+ 3. test_rpe_tp_n — N-dim TP (x/cos/sin Shard(1)), fwd vs standalone
	35	+ Expectation: Run success.
	36	+ """
	37	+ parallel_run([
	38	+ MindSporeCase(IMPL_FILE, "test_rpe_replicated", 19300, 2, 2, 2),
	39	+ MindSporeCase(IMPL_FILE, "test_rpe_dp_b", 19301, 2, 2, 2),
	40	+ MindSporeCase(IMPL_FILE, "test_rpe_tp_n", 19302, 2, 2, 2),
	41	+ ])
	42	+
	43	+
	44	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	45	+def test_parallel_rpe_ms_group2():
	46	+ """
	47	+ Feature: parallel run case in rotary_position_embedding_shard_in_python (4-card, 2-D mesh)
	48	+ Description:
	49	+ 1. test_rpe_dp_tp — 2-D (dp=2, tp=2) mesh; x/cos/sin (Shard(0), Shard(1));
	50	+ fwd + backward (grad_x/cos/sin) vs standalone
	51	+ 2. test_rpe_dp_sp — 2-D (dp=2, sp=2) mesh; x (Shard(0), Shard(2)),
	52	+ cos/sin broadcast (1,1,S,D) with (Replicate, Shard(2)); fwd vs standalone
	53	+ Expectation: Run success.
	54	+ """
	55	+ parallel_run([
	56	+ MindSporeCase(IMPL_FILE, "test_rpe_dp_tp", 19303, 4, 4, 2),
	57	+ MindSporeCase(IMPL_FILE, "test_rpe_dp_sp", 19304, 4, 4, 2),
	58	+ ])
	59	+
	60	+
	61	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	62	+def test_parallel_rpe_ms_group3():
	63	+ """
	64	+ Feature: parallel run case in rotary_position_embedding_shard_in_python (4-card, 2-D mesh)
	65	+ Description:
	66	+ 1. test_rpe_tp_sp — 2-D (tp=2, sp=2) mesh; x/cos/sin full shape (Shard(1), Shard(2));
	67	+ fwd vs standalone
	68	+ 2. test_rpe_dp_tp_cos_full — 2-D (dp=2, tp=2) mesh; x (Shard(0), Shard(1)),
	69	+ cos/sin broadcast (1,1,S,D) fully replicated; validates broadcast semantics
	70	+ when x is sharded on both B and N; fwd vs standalone
	71	+ Expectation: Run success.
	72	+ """
	73	+ parallel_run([
	74	+ MindSporeCase(IMPL_FILE, "test_rpe_tp_sp", 19305, 4, 4, 2),
	75	+ MindSporeCase(IMPL_FILE, "test_rpe_dp_tp_cos_full", 19306, 4, 4, 2),
	76	+ ])
	77	+
	78	+
	79	+@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level1", card_mark="allcards", essential_mark="essential")
	80	+def test_parallel_rpe_ms_group4():
	81	+ """
	82	+ Feature: parallel run case in rotary_position_embedding_shard_in_python (8-card, 3-D mesh)
	83	+ Description:
	84	+ 1. test_rpe_dp_tp_sp — 3-D (dp=2, tp=2, sp=2) mesh; x (Shard(0), Shard(1), Shard(2)),
	85	+ cos/sin broadcast (1,1,S,D) with (Replicate, Replicate, Shard(2));
	86	+ B/N/S all sharded simultaneously; fwd vs standalone
	87	+ Expectation: Run success.
	88	+ """
	89	+ parallel_run([
	90	+ MindSporeCase(IMPL_FILE, "test_rpe_dp_tp_sp", 19307, 8, 8, 2),
	91	+ ])

Atests/ut/core/shard/ops/test_parallel_rotary_position_embedding.py+507-0

@@ -0,0 +1,507 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""Unit tests for RotaryPositionEmbeddingDistributedOp."""
	16	+import os
	17	+import unittest
	18	+from unittest.mock import MagicMock, patch
	19	+
	20	+import numpy as np
	21	+
	22	+os.environ["HYPER_PARALLEL_PLATFORM"] = "torch"
	23	+
	24	+from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE
	25	+from hyper_parallel.core.dtensor.placement_types import Shard, Replicate
	26	+from hyper_parallel.core.shard.ops.parallel_rotary_position_embedding import (
	27	+ RotaryPositionEmbeddingDistributedOp,
	28	+ _normalize_rpe_args,
	29	+)
	30	+from hyper_parallel.core.shard.ops.parallel_ops_register import get_distributed_op
	31	+from hyper_parallel.core.dtensor.device_mesh import init_device_mesh, _DEVICE_MESH_MAP
	32	+from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS
	33	+
	34	+
	35	+class TestNormalizeRpeArgs(unittest.TestCase):
	36	+ """Unit tests for _normalize_rpe_args."""
	37	+
	38	+ def test_three_positional_args_mode_defaults_to_zero_1(self):
	39	+ """
	40	+ Feature: _normalize_rpe_args fills in default mode=0.
	41	+ Description: Call with only x, cos, sin positional args.
	42	+ Expectation: args == (x, cos, sin, 0), kwargs == {}.
	43	+ """
	44	+ x, cos, sin = object(), object(), object()
	45	+ args, kwargs = _normalize_rpe_args(x, cos, sin)
	46	+ self.assertIs(args[0], x, msg=f"args[0] should be x, got {args[0]}")
	47	+ self.assertIs(args[1], cos, msg=f"args[1] should be cos, got {args[1]}")
	48	+ self.assertIs(args[2], sin, msg=f"args[2] should be sin, got {args[2]}")
	49	+ self.assertEqual(args[3], 0, msg=f"default mode should be 0, got {args[3]}")
	50	+ self.assertEqual(kwargs, {}, msg=f"kwargs should be empty, got {kwargs}")
	51	+
	52	+ def test_mode_keyword_arg_2(self):
	53	+ """
	54	+ Feature: _normalize_rpe_args accepts mode as keyword argument.
	55	+ Description: Call with mode=1 as kwarg.
	56	+ Expectation: args[3] == 1.
	57	+ """
	58	+ x, cos, sin = object(), object(), object()
	59	+ args, _ = _normalize_rpe_args(x, cos, sin, mode=1)
	60	+ self.assertEqual(args[3], 1, msg=f"mode kwarg should be 1, got {args[3]}")
	61	+
	62	+ def test_mode_positional_arg_3(self):
	63	+ """
	64	+ Feature: _normalize_rpe_args accepts mode as positional argument.
	65	+ Description: Call with mode=2 as positional arg.
	66	+ Expectation: args[3] == 2.
	67	+ """
	68	+ x, cos, sin = object(), object(), object()
	69	+ args, kwargs = _normalize_rpe_args(x, cos, sin, 2)
	70	+ self.assertEqual(args[3], 2, msg=f"mode positional should be 2, got {args[3]}")
	71	+ self.assertEqual(kwargs, {}, msg=f"kwargs should be empty, got {kwargs}")
	72	+
	73	+
	74	+class TestYamlRegistration(unittest.TestCase):
	75	+ """Unit tests for YAML-driven op registration."""
	76	+
	77	+ def setUp(self):
	78	+ EXISTING_COMM_GROUPS.clear()
	79	+ _DEVICE_MESH_MAP.clear()
	80	+ _LAYOUT_CACHE.clear()
	81	+
	82	+ def tearDown(self):
	83	+ EXISTING_COMM_GROUPS.clear()
	84	+ _DEVICE_MESH_MAP.clear()
	85	+ _LAYOUT_CACHE.clear()
	86	+
	87	+ def test_pascal_case_registration_4(self):
	88	+ """
	89	+ Feature: YAML registers RotaryPositionEmbeddingDistributedOp (MindSpore name).
	90	+ Description: Call get_distributed_op with PascalCase op name.
	91	+ Expectation: Returns a RotaryPositionEmbeddingDistributedOp instance.
	92	+ """
	93	+ op = get_distributed_op("RotaryPositionEmbedding")
	94	+ self.assertIsNotNone(
	95	+ op,
	96	+ msg="RotaryPositionEmbedding should be registered via YAML"
	97	+ )
	98	+ self.assertIsInstance(
	99	+ op, RotaryPositionEmbeddingDistributedOp,
	100	+ msg=f"Expected RotaryPositionEmbeddingDistributedOp, got {type(op)}"