mindspore-ai / mindspore-ai/hyper-parallel
feat: 接入 RotaryPositionEmbedding 算子分布式调度支持
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 53
- Forks
- 63
- Avg merge
- 23h 45m
- Merged PRs (30d)
- 63
Description
背景
在 Tensor Parallelism（TP）/ 数据并行等分布式训练场景中，RotaryPositionEmbedding（RPE）是 Transformer 模型中广泛使用的位置编码算子（公式：y = x * cos + x_rotate * sin）。HyperParallel 当前缺少 RPE 的分布式算子调度实现：使用 DTensor 对象调用时无法推断输出 DTensor 的分布状态，下游算子无法正确组合分布式计算图。
目标
为 RotaryPositionEmbedding 提供分布式调度支持：
- 通过
preprocess()的to_local()提取本地 tensor 传给 CANN 内核 - 通过
infer_layout()基于输入分片方式推断输出 layout：- D（最后维度）不可切分（内核在 D 维度内旋转）
- B / N / S 维可自由切分
- 输出 layout 为 x layout 的深拷贝（输出 shape 等于 x shape）
- cos/sin 若在非-D 维度切分，切分方式须与 x 一致或为 Replicate（广播）
支持的分片场景
| 场景 | mesh | 用例 |
|---|---|---|
| 全 Replicate | 1-D dp | test_rpe_replicated |
| B 维 DP | 1-D dp, x Shard(0) | test_rpe_dp_b（fwd+bwd） |
| N 维 TP | 1-D tp, x/cos/sin Shard(1) | test_rpe_tp_n（fwd+bwd） |
| dp×tp 2-D | dp=2, tp=2; x/cos/sin (Shard(0), Shard(1)) | test_rpe_dp_tp（fwd+bwd） |
| dp×sp 2-D | dp=2, sp=2; x (Shard(0), Shard(2)); cos/sin (Replicate, Shard(2)) 广播 | test_rpe_dp_sp（fwd） |
| tp×sp 2-D | tp=2, sp=2; x/cos/sin (Shard(1), Shard(2)) | test_rpe_tp_sp（fwd） |
| dp×tp cos 广播 | dp=2, tp=2; x (Shard(0), Shard(1)); cos/sin (1,1,S,D) Replicate | test_rpe_dp_tp_cos_full（fwd） |
| dp×tp×sp 3-D | dp=2, tp=2, sp=2; x (Shard(0), Shard(1), Shard(2)); cos/sin (Replicate, Replicate, Shard(2)) | test_rpe_dp_tp_sp（fwd） |
约束
| 约束 | 说明 |
|---|---|
| D 维必须 Replicate | 所有输入（x/cos/sin）的最后维度不可切分 |
| cos/sin 切分须与 x 一致 | 若 cos/sin 在非-D 维切分，切分 mesh axis 须与 x 相同 |
| 不允许 Partial 输入 | x/cos/sin 均不得有 Partial 状态 |
| mode 不影响 layout | 4 种旋转模式（0/1/2/3）不影响输出 layout 推断 |
| CANN backward 限制 | mode=2/3 无 backward CANN 内核，仅支持 forward |
UT：test_parallel_rotary_position_embedding.py（22 cases）
ST：rotary_position_embedding_shard_in_python.py（8 cases，4 组，覆盖 2/4/8 卡）
相关 PR
#661
schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 155
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/155
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reviewing the existing rotary-position-embedding implementation and related PR #661, then read test_parallel_rotary_position_embedding.py and rotary_position_embedding_shard_in_python.py. Implement the preprocess() and infer_layout() behavior described for the listed sharding cases, and run the 22 UT cases plus the 8 ST cases to verify forward, backward, broadcasting, and layout constraints.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 65/100