已合并
feat: add distributed op RotaryPositionEmbedding with UT and MindSpore ST #661
hedongdong创建于 5月19日
feat: add distributed op RotaryPositionEmbedding with UT and MindSpore ST #661
已合并
Pull Request已成功合入, 合并人@MindSpore-Bot
（感谢 hedongdong 的贡献）5月19日 创建了 pull request，commit aaafe6b1
5月19日 关联了issue：feat: 接入 RotaryPositionEmbedding 算子分布式调度支持
5月19日 添加了label：mindspore-cla/yes
MindSpore-Bot
5月19日 评论：
5月19日 评论：
5月19日 添加了label：pr-check-pass
5月19日 添加了label：no-pass-all-review
MindSpore-Bot
5月19日 评论：
5月19日 评论：
hedongdong
5月19日 评论：
5月19日 评论：
/retest
5月19日 添加了label：ci-pipeline-running
5月19日 通过审查
MindSpore-Bot
5月19日 评论：
5月19日 评论：
hyper_parallel/core/shard/ops/parallel_rotary_position_embedding.py
@@ -0,0 +96,4 @@
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):
🚨 [关键问题] 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";
}
司小南（机器人）
5月19日 评论：
5月19日 评论：
📖 Python注释检查：
根据您提供的代码差异，我仔细审查了其中的英文文档修改，发现以下不符合Python注释格式规范的地方：
1.💡问题描述：Args部分参数格式缺失类型说明，且参数mode有默认值但未按规范标注optional和Default。
- 📄 文件:
hyper_parallel/core/shard/ops/parallel_rotary_position_embedding.py - ✍ 修改建议：
- Args: - x: Input tensor. - cos: Cosine position encoding tensor. - sin: Sine position encoding tensor. - mode: Rotation mode. 0=rotate_half, 1=rotate_interleaved, 2=quarter, - 3=interleave-half. Defaults to 0. + Args: + x (Tensor): Input tensor. + cos (Tensor): Cosine position encoding tensor. + sin (Tensor): Sine position encoding tensor. + mode (int, optional): Rotation mode. 0=rotate_half, 1=rotate_interleaved, 2=quarter, + 3=interleave-half. Default: 0.
2.💡问题描述：Args部分参数格式缺失类型说明。
- 📄 文件:
hyper_parallel/core/shard/ops/parallel_rotary_position_embedding.py - ✍ 修改建议：
- Args: - x_layout: Layout of the x tensor. - cos_layout: Layout of the cos tensor. - sin_layout: Layout of the sin tensor. + Args: + x_layout (Layout): Layout of the x tensor. + cos_layout (Layout): Layout of the cos tensor. + sin_layout (Layout): Layout of the sin tensor.
3.💡问题描述：Args部分参数格式缺失类型说明。
- 📄 文件:
hyper_parallel/core/shard/ops/parallel_rotary_position_embedding.py - ✍ 修改建议：
- Args: - args: Positional arguments, may include DTensors. - kwargs: Keyword arguments. + Args: + args (tuple): Positional arguments, may include DTensors. + kwargs (dict): Keyword arguments.
4.💡问题描述：Args部分参数格式缺失类型说明。
- 📄 文件:
hyper_parallel/core/shard/ops/parallel_rotary_position_embedding.py - ✍ 修改建议：
- Args: - cache_values: [x_layout, cos_layout, sin_layout] + Args: + cache_values (list): [x_layout, cos_layout, sin_layout]
5.💡问题描述：Args部分参数格式缺失类型说明。
- 📄 文件:
tests/mindspore/st/shard/ops/rotary_position_embedding_shard_in_python.py - ✍ 修改建议：
- Args: - mode: Rotation mode (0=rotate_half, 1=rotate_interleaved, - 2=quarter, 3=interleave-half). + Args: + mode (int): Rotation mode (0=rotate_half, 1=rotate_interleaved, + 2=quarter, 3=interleave-half).
6.💡问题描述：Args部分参数格式缺失类型说明，且参数mode有默认值但未按规范标注optional和Default。
- 📄 文件:
tests/mindspore/st/shard/ops/rotary_position_embedding_shard_in_python.py - ✍ 修改建议：
- Args: - x_np: float32 ndarray, shape (B, N, S, D). - cos_np: float32 ndarray, shape (B, N, S, D) or (1, 1, S, D). - sin_np: float32 ndarray, same shape as cos_np. - mode: RPE rotation mode. + Args: + x_np (np.ndarray): float32 ndarray, shape (B, N, S, D). + cos_np (np.ndarray): float32 ndarray, shape (B, N, S, D) or (1, 1, S, D). + sin_np (np.ndarray): float32 ndarray, same shape as cos_np. + mode (int, optional): RPE rotation mode. Default: 0.
7.💡问题描述：Args部分参数格式缺失类型说明，且参数mode有默认值但未按规范标注optional和Default。
- 📄 文件:
tests/mindspore/st/shard/ops/rotary_position_embedding_shard_in_python.py - ✍ 修改建议：
- Args: - x_np: float32 ndarray, shape (B, N, S, D). - cos_np: float32 ndarray, shape (B, N, S, D) or (1, 1, S, D). - sin_np: float32 ndarray, same shape as cos_np. - mode: RPE rotation mode (must be 0 or 1). + Args: + x_np (np.ndarray): float32 ndarray, shape (B, N, S, D). + cos_np (np.ndarray): float32 ndarray, shape (B, N, S, D) or (1, 1, S, D). + sin_np (np.ndarray): float32 ndarray, same shape as cos_np. + mode (int, optional): RPE rotation mode (must be 0 or 1). Default: 0.
8.💡问题描述：Args部分参数格式缺失类型说明。
- 📄 文件:
tests/mindspore/st/shard/ops/rotary_position_embedding_shard_in_python.py - ✍ 修改建议：
- Args: - d_inputs: Tuple of DTensor inputs (dx, dcos, dsin). + Args: + d_inputs (tuple): Tuple of DTensor inputs (dx, dcos, dsin).
9.💡问题描述：Args部分参数格式缺失类型说明。
- 📄 文件:
tests/mindspore/st/shard/ops/rotary_position_embedding_shard_in_python.py - ✍ 修改建议：
- Args: - raw_grads: Tuple of local gradient tensors from ms.grad. - d_inputs: Tuple of DTensor inputs (dx, dcos, dsin). - ref_grads: Tuple of float32 reference gradient ndarrays. - tag: Human-readable label for assertion messages. + Args: + raw_grads (tuple): Tuple of local gradient tensors from ms.grad. + d_inputs (tuple): Tuple of DTensor inputs (dx, dcos, dsin). + ref_grads (tuple): Tuple of float32 reference gradient ndarrays. + tag (str): Human-readable label for assertion messages.
🤗 上述内容仅供参考。
5月19日 添加了label：ai-reviewed
司小南（机器人）
5月19日 评论：
5月19日 评论：
| Project Name | Build_Stage | Build Result | Details |
|---|---|---|---|
| Hyper-parallel_Atomgit_Gate | - | ✅ SUCCESS | 2570 |
5月19日 删除了label：ci-pipeline-running
5月19日 添加了label：ci-pipeline-passed
5月19日 通过审查
5月19日 删除了label：no-pass-all-review
5月19日 合入了pull request，合并节点 SHA：f417a67a046efda75ca836e63070fc14355a71b0
What type of PR is this?
/kind feature
What does this PR do / why do we need it:
为 HyperParallel 接入
RotaryPositionEmbedding算子的分布式调度支持：RotaryPositionEmbedding是 Transformer 模型中广泛使用的位置编码算子（公式：y = x * cos + x_rotate * sin），在分布式并行训练中需要正确推断 DTensor 的输出分布状态。HyperParallel 的工作包括：
preprocess()的to_local()提取本地 tensor 后传给 CANN 内核infer_layout()基于输入分片方式推断输出 layout分片规则
测试覆盖
test_parallel_rotary_position_embedding.py(22 cases)_normalize_rpe_args/ YAML 注册 /infer_layout正例（全 Replicate、B-DP、N-TP、2D/3D mesh、cos/sin 广播）/ 错误处理（D 维切分、cos/sin 切分不一致、Partial 输入）rotary_position_embedding_shard_in_python.py(8 cases)Which issue(s) this PR fixes:
Fixes #155
Test Plan and Test result：What scenarios were tested, and what were the verification results（Function, performance, reliability, etc.）：
UT（单元测试）
tests/ut/core/shard/ops/test_parallel_rotary_position_embedding.py— 22 cases，覆盖参数归一化、YAML 注册、infer_layout（全 Replicate / B-DP / N-TP / 2D-3D mesh / cos/sin 广播）、错误处理（D 维切分、cos/sin 切分不一致、Partial 输入）ST（系统测试 - distributed, MindSpore）
tests/mindspore/st/shard/ops/test_rotary_position_embedding_shard_in_python.py— 4 组：验证结果
Self-checklist:（请自检，在[ ]内打上x，我们将检视你的完成情况，否则会导致pr无法合入）