mindspore-ai / mindspore-ai/hyper-parallel
[BUG] DSA 分布式算子要求 actual_seq_len 必须为 DTensor,无法兼容网络构造的普通 Tensor
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 53
- Forks
- 63
- Avg merge
- 23h 45m
- Merged PRs (30d)
- 63
Description
该问题是怎么引起的?
DSA 相关分布式算子(lightning_indexer、npu_dense_lightning_indexer_softmax_lse、npu_dense_lightning_indexer_grad_kl_loss、npu_sparse_lightning_indexer_grad_kl_loss、npu_sparse_flash_attention)在 preprocess 中对 actual_seq_len(q/kv 的累积序列长度)入参无条件调用 .to_local(),即要求它必须是 Replicate 的 DTensor。
但 actual_seq_len 是网络在前向过程中构造出来的普通 Tensor,用户调用分布式算子时无法保证把它包装成 DTensor。当传入普通 Tensor 时,.to_local() 抛 AttributeError。
现状(根因分析)
各算子 preprocess 对 actual_seq_len 的处理存在两类问题:
| 算子 | actual_seq_len 处理 | 问题 |
|---|---|---|
npu_dense_lightning_indexer_softmax_lse |
if x is not None: x.to_local() |
普通 Tensor → AttributeError |
lightning_indexer |
if x is not None: x.to_local() |
同上 |
npu_sparse_flash_attention |
宽松 _to_local 套在所有入参 |
actual_seq 可通过,但主张量传普通 Tensor 被静默放行,layout 无法推导 |
npu_dense_lightning_indexer_grad_kl_loss |
同上宽松 _to_local |
同上 |
npu_sparse_lightning_indexer_grad_kl_loss |
同上宽松 _to_local |
同上 |
两个矛盾点:
softmax_lse/lightning_indexer对actual_seq_len直接.to_local(),不兼容网络构造的普通 Tensor。- 另外三个算子把宽松
_to_local(isinstance(DTensor)判断)套在所有入参上,导致actual_seq_len之外的主张量(query/key/value/index/weights、rope 等)即使传普通 Tensor 也被静默透传,而这些入参必须是 DTensor 才能推导 layout。
修复方案
- 新增宽松 helper
_to_local_seq_len:仅当入参是 DTensor 时to_local,普通 Tensor /None原样透传。只用于actual_seq_len。 - 其余参与 layout 推导的入参收紧为严格
_to_local:None透传(可选参数未提供),否则直接to_local,普通 Tensor 触发AttributeError(明确拒绝,避免静默放行)。 - 同步更新 UT/ST:ST 把
actual_seq_len从 Replicate DTensor 改为普通 Tensor 传入;UT 增加「普通 Tensor 透传」「主张量拒绝普通 Tensor」的回归用例。
重现步骤
- 构造 query / key / weights 为 Replicate 或 Shard 的 DTensor。
actual_seq_qlen/actual_seq_klen用网络构造的普通 Tensor(非 DTensor)。- 调用
npu_dense_lightning_indexer_softmax_lse(q, k, w, actual_seq_qlen=..., actual_seq_klen=..., layout='TND')。 preprocess中actual_seq_qlen.to_local()触发报错。
报错信息
AttributeError: 'Tensor' object has no attribute 'to_local'
关联 PR
mindspore/hyper-parallel#871
schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 232
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/232
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 preprocess entry points for lightning_indexer, npu_dense_lightning_indexer_softmax_lse, npu_dense_lightning_indexer_grad_kl_loss, npu_sparse_lightning_indexer_grad_kl_loss, and npu_sparse_flash_attention, along with PR #871. Run the related UT/ST coverage; done means ordinary Tensor actual_seq_len values are accepted while non-DTensor layout-driving inputs are rejected explicitly.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100