feat: 接入 RotaryPositionEmbedding 算子分布式调度支持-hyper-parallel-AtomGit

已开启

feat: 接入 RotaryPositionEmbedding 算子分布式调度支持 #155

hedongdong创建于 5月19日

5月19日创建

背景

在 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