mindspore-ai / mindspore-ai/hyper-parallel

[BUG] DSA 分布式算子要求 actual_seq_len 必须为 DTensor,无法兼容网络构造的普通 Tensor

Open
#683 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
53
Forks
63
Avg merge
23h 45m
Merged PRs (30d)
63

Description

该问题是怎么引起的?

DSA 相关分布式算子(lightning_indexernpu_dense_lightning_indexer_softmax_lsenpu_dense_lightning_indexer_grad_kl_lossnpu_sparse_lightning_indexer_grad_kl_lossnpu_sparse_flash_attention)在 preprocess 中对 actual_seq_len(q/kv 的累积序列长度)入参无条件调用 .to_local(),即要求它必须是 Replicate 的 DTensor。

actual_seq_len 是网络在前向过程中构造出来的普通 Tensor,用户调用分布式算子时无法保证把它包装成 DTensor。当传入普通 Tensor 时,.to_local()AttributeError

现状(根因分析)

各算子 preprocessactual_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 同上

两个矛盾点:

  1. softmax_lse / lightning_indexeractual_seq_len 直接 .to_local(),不兼容网络构造的普通 Tensor。
  2. 另外三个算子把宽松 _to_localisinstance(DTensor) 判断)套在所有入参上,导致 actual_seq_len 之外的主张量(query/key/value/index/weights、rope 等)即使传普通 Tensor 也被静默透传,而这些入参必须是 DTensor 才能推导 layout。
修复方案
  1. 新增宽松 helper _to_local_seq_len:仅当入参是 DTensor 时 to_local,普通 Tensor / None 原样透传。用于 actual_seq_len
  2. 其余参与 layout 推导的入参收紧为严格 _to_localNone 透传(可选参数未提供),否则直接 to_local,普通 Tensor 触发 AttributeError(明确拒绝,避免静默放行)。
  3. 同步更新 UT/ST:ST 把 actual_seq_len 从 Replicate DTensor 改为普通 Tensor 传入;UT 增加「普通 Tensor 透传」「主张量拒绝普通 Tensor」的回归用例。
重现步骤
  1. 构造 query / key / weights 为 Replicate 或 Shard 的 DTensor。
  2. actual_seq_qlen / actual_seq_klen 用网络构造的普通 Tensor(非 DTensor)。
  3. 调用 npu_dense_lightning_indexer_softmax_lse(q, k, w, actual_seq_qlen=..., actual_seq_klen=..., layout='TND')
  4. preprocessactual_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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.