已关闭
[Bug]: npu_grouped_matmul meta fails on data-dependent SymInt under torch.compile #4520
wanghui0创建于  9月2日关闭于  9月3日
wanghui0
9月2日 创建

在提交新问题之前,请确保您已经在社区中搜索过相关问题,并使用了社区中提供的资源/工具后,仍未找到满意的解决方式。

⚠️ 安全信息提醒:请仔细检查提供的文本内容,确保其不包含敏感数据信息,包括但不限于:

  • API 令牌或密钥
  • 密码或身份验证凭证
  • 私有网址或接口地址
  • 个人或机密数据
  • ...

在分享配置信息或代码示例时,请将敏感信息脱敏处理,或使用 <TOKEN> 等占位符替代原有内容。

环境信息

- torch_npu: 2.14.0.dev20260808
- Python: 3.12
- torch: 2.14.0.dev20260719+cpu
- CANN: 海思日包 0824
- wheel 来源:
  https://cann-ai.obs.cn-north-4.myhuaweicloud.com/torchtitan-npu/simple/torch-npu/index.html
- 使用场景:
  - TorchTitan-NPU
  - MoE 低精度 grouped matmul
  - torch.compile 开启
  - SPMD backend: spmd_types
  - 量化 recipe: mix/block-fp8

🐛 问题描述

当 MoE routing 产生的 token 数量在编译期未知时,npu_grouped_matmul 的输入维度 m 是 data-dependent SymInt。

FakeTensor 执行 npu_grouped_matmul_meta 时,以下 Python 条件会尝试把SymBool 转换为普通 Python bool:

if m > 0 and n > 0:
    torch._check(k > 0, ...)

由于 m 是 unbacked/data-dependent SymInt,Dynamo 无法在编译期对 m > 0 生成静态 guard,最终抛出:

torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode:
Could not guard on data-dependent expression ... > 0

关闭 torch.compile 后可以运行,说明问题发生在 FakeTensor/meta shape 推导阶段,而不是 NPU grouped matmul 的实际执行阶段。

复现条件
使用动态 MoE routing,使 grouped matmul 的 M 维来自运行时 token 数量。
启用 torch.compile。
grouped matmul 输入使用 FP8、MXFP、A8W8 或 HiFloat8 等量化类型。
FakeTensor 执行 npu_grouped_matmul_meta。
也可以使用 ShapeEnv.create_unbacked_symint() 构造动态 M 维进行最小复现。

欢迎加入社区,感谢您对社区的贡献 🎉!

likedislike
TorchNPU-BotTorchNPU-Bot成员
9月2日 添加了label:triage-review
TorchNPU-Bot
TorchNPU-Bot成员
9月2日 评论:

issue待分派,添加triage-review标签

likedislike
ascend-robotascend-robot成员
9月2日 添加了label:bug
Wwanghui0
9月2日 关联了pull request:fix: support symbolic shapes in grouped matmul meta
Wwanghui0
9月2日 修改了issue 的描述
Wwanghui0
9月2日 关联了pull request:fix: 在 grouped matmul meta 中支持 symbolic shapes
Wwanghui0
9月3日 issue状态由 TODO 改变为 DONE
Wwanghui0
9月3日 关闭了 issue
ascend-robotascend-robot成员
9月3日 添加了label:resolved