[Bug] 分布式 MatMul 的 layout 推断不传播输入的 Partial 状态，导致链式 matmul（如 LoRA `(x@Aᵀ)@Bᵀ`）的反向梯度被破坏（数值减半 / 1-over-N）

环境

• hyper_parallel commit d16b218（MindSpore 后端）
• MindSpore 2.10.0；CANN 9.1.0-beta.1；Ascend 910B2
• 复现于 TP（tensor parallel）下 row-wise 切分的链式低秩 matmul

结论（先说定界）

这是 hyper_parallel 的 DTensor 分布式 MatMul bug，不是上层（MindFormers / LoRA 切分计划）的问题。 纯 DTensor 最小用例（不含任何 MindFormers / LoRA 代码、不含自定义 ParallelStyle）即可复现：对一个本应得到 Partial 输出的 matmul，框架按「自己的 contracting 维」推断 layout，不继承输入张量已有的 Partial 状态，于是把结果错误标成非 Partial，前向数值错误、反向梯度被破坏（对 Replicate 参数其梯度数值变为正确值的 1/N，N=TP 度）。

根因

链式 (x @ Aᵀ) @ Bᵀ，在 row-wise TP 下：

• x = Shard(-1)（in-dim 切分），A = Shard(1)，B = Replicate；
• 第一个 matmul 沿被切分的 in-dim 规约 → tmp 是 Partial（正确）；
• 第二个 matmul 沿未切分的 rank 维规约。MatMul layout 推断只看自己的 contracting 维（未切分→判定非 Partial），忽略了 tmp 已是 Partial，于是 lora_out 被错标为非 Partial；
• 前向把未规约的 Partial 当成已规约值使用 → 数值错误；反向时 B（Replicate）的梯度被错误规约（等价于做了平均而非求和）→ 梯度 = 正确值的 1/N。

注：Partial → Replicate 的 redistribute 本身是正确的（做 SUM）（见复现脚本测试1），问题仅在 matmul 不传播输入 Partial。

最小复现脚本（纯 hyper_parallel DTensor，无 MindFormers）

repro_dtensor.py，2 卡：msrun --worker_num=2 --local_worker_num=2 --master_port=59799 --join=True repro_dtensor.py

import os
os.environ.setdefault("HYPER_PARALLEL_PLATFORM", "mindspore")
import numpy as np, mindspore as ms
from mindspore import mint
from mindspore.mint.distributed import init_process_group
from mindspore.communication import get_rank, get_group_size
from hyper_parallel import init_device_mesh, DTensor
from hyper_parallel.core.dtensor.placement_types import Shard, Replicate, Partial

ms.set_context(mode=ms.PYNATIVE_MODE)
init_process_group()
RANK, WORLD = get_rank(), get_group_size()
mesh = init_device_mesh("npu", (WORLD,), mesh_dim_names=("tp",))
def log(*a):
    if RANK == 0: print("[REPRO]", *a, flush=True)

# 测试1：Partial->Replicate redistribute 应做 SUM
local = ms.Tensor(np.full((4,), float(RANK + 1), dtype=np.float32))   # rank r 持 (r+1)
rep = DTensor.from_local(local, mesh, (Partial(),)).redistribute(device_mesh=mesh, placements=(Replicate(),))
log("测试1 Partial->Replicate got=", float(rep.to_local().asnumpy()[0]), " (WORLD=2 时 SUM应=3, MEAN=1.5)")

# 测试2：链式 (x@Aᵀ)@Bᵀ，对 Replicate 的 B 求梯度，对照单卡 ground-truth
np.random.seed(0)
R, DIN, DOUT, BS = 8, 16, 32, 4
Xf = np.random.randn(BS, DIN).astype(np.float32)
Af = (np.random.randn(R, DIN) * 0.1).astype(np.float32)
Bf = (np.random.randn(DOUT, R) * 0.1).astype(np.float32)
tmp_full = Xf @ Af.T
out_truth = float((tmp_full @ Bf.T).sum())
gB_truth_norm = float(np.linalg.norm(np.tile(tmp_full.sum(0, keepdims=True), (DOUT, 1))))
sh = DIN // WORLD
x_loc = ms.Tensor(Xf[:, RANK*sh:(RANK+1)*sh]); a_loc = ms.Tensor(Af[:, RANK*sh:(RANK+1)*sh]); b_full = ms.Tensor(Bf)

def fwd(b_full_, a_loc_, x_loc_, reduce_tmp):
    x_dt = DTensor.from_local(x_loc_, mesh, (Shard(1),))
    a_dt = DTensor.from_local(a_loc_, mesh, (Shard(1),))
    b_dt = DTensor.from_local(b_full_, mesh, (Replicate(),))
    tmp = mint.matmul(x_dt, mint.transpose(a_dt, 1, 0))           # -> 应为 Partial
    if reduce_tmp and isinstance(tmp, DTensor) and any(p.is_partial() for p in tmp.placements):
        tmp = tmp.reduce_partial()                               # 规避：先 all-reduce 成 Replicate
    lora = mint.matmul(tmp, mint.transpose(b_dt, 1, 0))          # 沿未切分 rank 维规约
    return lora.redistribute(device_mesh=mesh, placements=(Replicate(),)).to_local().sum()

for reduce_tmp in (False, True):
    fv = float(fwd(b_full, a_loc, x_loc, reduce_tmp).asnumpy())
    _, gB = ms.value_and_grad(fwd, grad_position=0)(b_full, a_loc, x_loc, reduce_tmp)
    ratio = float(np.linalg.norm(gB.asnumpy()) / gB_truth_norm)
    log(f"测试2 reduce_partial={reduce_tmp}: forward={fv:.5f}(truth {out_truth:.5f}) grad_ratio={ratio:.4f}")

期望 vs 实测

即：只要在第二个 matmul 前把 Partial 显式规约成 Replicate，前向与梯度都恢复正确；不规约则前向错、梯度变 1/N。

MindFormers 端到端复现（动态图 LoRA + TP）

版本：mindformers feat/pynative-lora-v1 @ 047ae22d9，DeepSeek-V3 12 层 MoE，seq 4096，2 卡。

1. 触发点：mindformers/pynative/pet/lora_layer.py: LinearWithLoRA.construct，链式 tmp=x@lora_aᵀ; lora_out=tmp@lora_bᵀ。
2. 复现方式（等价性判据）：所有并行场景加载同一份 base 权重 W0（单卡 lr=0 存 ckpt），统一 global_batch_size=4 / lr=1e-4 / 20 步 / 仅加载权重，以单卡为 baseline 比较 loss 曲线逐步 diff。
3. yaml 关键项：parallelism.tensor_parallel: 2、sequence_parallel: true、model.compute_dtype: bfloat16、lora_config.target_modules 含 row-wise 的 self_attention.linear_proj、mlp.linear_fc2。
4. 现象（未规避前 / lora_layer.py 不调用 reduce_partial）：
   • 单卡 baseline lora_b.grad ‖·‖²=1.66e-11；TP=2 同一 W0/同 step1 = 4.18e-12 → 范数比 0.502 ≈ 1/tp（减半）；
   • loss 曲线单边漂移：step20 TP loss 比单卡高 +0.045，grad_norm 仅为单卡 ~57%；
   • 控制实验排除其它解释：fp32 下仍发散（非 bf16 噪声）、前向逐位对齐（非前向）、full-ft（无 LoRA）不漂移（非基座）。
5. 规避后（lora_layer.py 调用 tmp.reduce_partial()）：TP/TP×CP/DP×TP/DP×TP×CP loss diff 均回到零轴波动（mean≈+0.0001、max|Δ|≈0.003、正负步数 10/9），与单卡一致。

期望行为

分布式 MatMul 的 output layout 推断应继承输入张量的 Partial 状态（输入 Partial → 输出 Partial），使链式 matmul 的前向/反向自动正确，无需上层手动 reduce_partial 规避。

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 228
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/228

Start with the pure DTensor reproduction in repro_dtensor.py and trace mint.matmul layout inference for the first and second chained matmuls. Compare the reduce_partial=True and False cases, then inspect the linked MindFormers entry point at mindformers/pynative/pet/lora_layer.py: LinearWithLoRA.construct. Done means the chained matmul preserves the input Partial state and the reproduction matches the expected forward value and gradient ratio without a manual reduction.