Open
【重构】 DTensor copy_/zero_/fill_ 下沉到 core 及 MindSpore __new__ 路径优化 #207
hedongdongAdd on  Jun 11
hedongdongCollaborators
Jun 11 Create

DTensor 重构:copy_/zero_/fill_ 下沉到 core 及 MindSpore new 优化

需求背景 & 价值

当前 DTensor 的 copy_ / zero_ / fill_ 在 PyTorch 和 MindSpore 两侧有各自独立的实现,存在以下问题:

PyTorch 侧:

  • leaf+grad 路径存在 dtype/storage 双重 desync bug(copy_(fp16_src) 导致 _local_tensor 和 wrapper 的 dtype/storage 分裂)
  • 依赖 rebind _local_tensor 的方式绕过 leaf in-place 限制,破坏 _local_tensor 身份一致性

MindSpore 侧:

  • 没有显式 override,依赖 _INPLACE_BYPASS_OPS 兜底路径,缺少类型/placement 校验
  • shape 不匹配在 NPU 上异步报错,traceback 指向无关代码行

MindSpore DTensorBase.__new__:

  • 热点路径(.from_local)每次都触发冗余 .to("Ascend"),即使 tensor 已在目标设备
  • device != "meta" 判断不准确,无法区分 CPU 和 Ascend

功能描述

1. DTensor copy_/zero_/fill_ 统一到 core

  • 三个方法从 platform/torch/dtensor.py 移到 core/dtensor/dtensor.py,两个平台共用
  • 采用 Tensor.copy_ 语义(方案B),与标准 Tensor.copy_ 行为一致:
    • version counter 正确 bump
    • requires_grad 叶 Tensor 抛异常(安全网有效)
    • 创建 grad_fn,梯度从 dst 流向 src
  • 强制 src 必须是 DTensor + 同 DeviceMesh
  • src_is_scalar(numel==1)短路 placement 和广播双重校验,高频场景零开销
  • 同步 shape 校验(解决 MS NPU 异步报错问题)
  • 删除 torch 侧三个 override

2. MindSpore DTensorBase.new 优化

  • 用 dev.startswith("Ascend") 准确识别目标设备,同设备跳过 .to()
  • 统一 has_init 的 init_device 设置逻辑
  • 合并两个分支末尾重复代码
  • 性能:24.56 us/call → 10.00 us/call(约 2.5×)

3. GeLU 算子注册 & InplaceCopy 调度修正

  • 在 MindSpore element_wise_ops.yaml 中注册 GeluExt 算子
  • 将 InplaceCopy 加入 _INPLACE_BYPASS_OPS
  • 从 whitelist 中移除 Inplace* 系列算子,修改 _should_bypass_dispatch 使 bypass 逻辑独立

涉及文件

文件 变更
core/dtensor/dtensor.py 新增 copy_/zero_/fill_/__is_broadcastable
platform/mindspore/dtensor.py 重构 DTensorBase.new
platform/torch/dtensor.py 删除 copy_/zero_/fill_ override
core/shard/_op_dispatch.py InplaceCopy bypass + _should_bypass_dispatch 修改
core/shard/ops/yaml/element_wise_ops.yaml 注册 GeluExt
tests/ut/core/dtensor/test_dtensor.py copy_/zero_/fill_ UT
tests/ut/platform/mindspore/dtensor/test_dtensor.py MS new UT
tests/torch/shard/ops/test_parallel_op_dtensor_inplace.py torch ST
tests/mindspore/st/shard/ops/test_parallel_op_dtensor_inplace.py MS ST
likedislike
HhedongdongCollaborators
Jun 13 Modify description of issue