当前 DTensor 的 copy_ / zero_ / fill_ 在 PyTorch 和 MindSpore 两侧有各自独立的实现,存在以下问题:
copy_
zero_
fill_
PyTorch 侧:
copy_(fp16_src)
_local_tensor
MindSpore 侧:
_INPLACE_BYPASS_OPS
MindSpore DTensorBase.__new__:
DTensorBase.__new__
.from_local
.to("Ascend")
device != "meta"
platform/torch/dtensor.py
core/dtensor/dtensor.py
Tensor.copy_
src_is_scalar
dev.startswith("Ascend")
.to()
init_device
_should_bypass_dispatch
platform/mindspore/dtensor.py
core/shard/_op_dispatch.py
core/shard/ops/yaml/element_wise_ops.yaml
tests/ut/core/dtensor/test_dtensor.py
tests/ut/platform/mindspore/dtensor/test_dtensor.py
tests/torch/shard/ops/test_parallel_op_dtensor_inplace.py
tests/mindspore/st/shard/ops/test_parallel_op_dtensor_inplace.py
DTensor 重构:copy_/zero_/fill_ 下沉到 core 及 MindSpore new 优化
需求背景 & 价值
当前 DTensor 的
copy_/zero_/fill_在 PyTorch 和 MindSpore 两侧有各自独立的实现,存在以下问题:PyTorch 侧:
copy_(fp16_src)导致_local_tensor和 wrapper 的 dtype/storage 分裂)_local_tensor的方式绕过 leaf in-place 限制,破坏_local_tensor身份一致性MindSpore 侧:
_INPLACE_BYPASS_OPS兜底路径,缺少类型/placement 校验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_行为一致:src_is_scalar(numel==1)短路 placement 和广播双重校验,高频场景零开销2. MindSpore DTensorBase.new 优化
dev.startswith("Ascend")准确识别目标设备,同设备跳过.to()init_device设置逻辑3. GeLU 算子注册 & InplaceCopy 调度修正
_INPLACE_BYPASS_OPS_should_bypass_dispatch使 bypass 逻辑独立涉及文件
core/dtensor/dtensor.pyplatform/mindspore/dtensor.pyplatform/torch/dtensor.pycore/shard/_op_dispatch.pycore/shard/ops/yaml/element_wise_ops.yamltests/ut/core/dtensor/test_dtensor.pytests/ut/platform/mindspore/dtensor/test_dtensor.pytests/torch/shard/ops/test_parallel_op_dtensor_inplace.pytests/mindspore/st/shard/ops/test_parallel_op_dtensor_inplace.py