已关闭

[Bug]: NPU TypedStorage._deepcopy dtype 丢失问题 #4510

xuyun15创建于 24 天前关闭于 6 天前

xuyun15成员

24 天前创建

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

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

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

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

环境信息

- 操作系统：不限制
- 昇腾硬件信息：不限制
- CANN软件版本：不限制
- 安装的对应软件版本：不限制

🐛 问题描述

NPU TypedStorage._deepcopy dtype 丢失问题分析与修复

问题现象

以下代码在 CPU 上正常运行，在 NPU 上报错：

import torch
import torch_npu

weights = torch.randint(0, 255, (1024,), dtype=torch.uint8, device="npu")
qscale_view = weights[512:576].view(torch.float16)
copy.deepcopy(qscale_view)
# RuntimeError: Expected a Storage of type c10::Half or an UntypedStorage,
#              but got type Byte for argument 1 'storage'

核心场景：在 uint8 存储上创建 float16 视图，对该视图执行 deepcopy 时 NPU 报错。

一、完整调用栈

copy.deepcopy(qscale_view)            # 用户代码
└─ Tensor.__deepcopy__                # pytorch/torch/_tensor.py:134
   ├─ new_storage = self._typed_storage()._deepcopy(memo)   # 第174行
   │  └─ _typed_storage() 返回 TypedStorage(dtype=Half, wrap_storage=uint8_untyped)
   │     # pytorch/torch/_tensor.py:309-313
   │  └─ torch_npu _deepcopy(self, memo)
   │     # ascend_pytorch/pytorch/torch_npu/utils/storage.py:54
   │     ├─ self.device.type == 'npu' → 进入非 cpu 分支 (第55行)
   │     ├─ src_tensor = torch_npu._C._tensor_construct_from_storage(self)  # 第57行
   │     │  └─ C++: THNPModule_tensor_construct_from_storage
   │     │     # ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp:2021
   │     │     ├─ parser.parse → _r.storage(0, storage_scalar_type, is_typed_storage)
   │     │     │  → 提取到 storage_scalar_type = Half, 但未传递给下游函数
   │     │     └─ at_npu::native::set_tensor_with_storage_format(storage)  # 第2042行
   │     │        └─ ascend_pytorch/pytorch/torch_npu/csrc/aten/common/SetNpu.cpp:120
   │     │           ├─ desc = GetNpuStorageImpl(src)->npu_desc_
   │     │           ├─ desc.data_type_ = Byte (来自原始 uint8 weights 创建时记录)
   │     │           ├─ dist_tensor = NPUNativeFunctions::empty(
   │     │           │     {0}, desc.data_type_.toScalarType()=Byte, ...)
   │     │           │  → 创建了 dtype=Byte 的 tensor!
   │     │           └─ set_storage_nd_npu(dist_tensor, src, ...)
   │     │              → 返回 dtype=Byte 的 src_tensor
   │     ├─ dst_tensor = src_tensor.clone()                  # 第58行, dtype 仍为 Byte
   │     ├─ dst_tensor = torch_npu.npu_format_cast(...)      # 第59行
   │     └─ new_storage = dst_tensor._typed_storage()        # 第60行 → dtype=Byte 的 TypedStorage!
   ├─ new_tensor = torch.empty_like(self)                    # 第193行, dtype=Half
   └─ new_tensor.set_(new_storage, self.storage_offset(), self.size(), self.stride())  # 第234行
      │  ↑ new_storage.dtype=Byte, new_tensor.dtype=Half → 不匹配!
      └─ set_.source_Storage_storage_offset 绑定
         # pytorch/tools/autograd/templates/python_variable_methods.cpp:1183
         └─ TORCH_CHECK(storage_scalar_type == self.dtype() || !is_typed_storage)  # 第1197行 → 失败!
            → "Expected a Storage of type c10::Half or an UntypedStorage, but got type Byte"

二、根因详解（逐层代码证据）

2.1 deepcopy 入口：Tensor.deepcopy

pytorch/torch/_tensor.py 第174行和第234行：

new_storage = self._typed_storage()._deepcopy(memo)    # 第174行
...
new_tensor.set_(new_storage, self.storage_offset(), self.size(), self.stride())  # 第234行

此时 self.dtype() = Half（float16），因为 qscale_view 是 weights[512:576].view(torch.float16)。

2.2 _typed_storage() 正确构造了 dtype=Half 的 TypedStorage

pytorch/torch/_tensor.py 第309-313行：

def _typed_storage(self):
    untyped_storage = self.untyped_storage()
    return torch.TypedStorage(
        wrap_storage=untyped_storage, dtype=self.dtype, _internal=True
    )

self.dtype = Half，所以返回的 TypedStorage 的 dtype = Half，但底层 untyped_storage 是 uint8 字节。

2.3 torch_npu 的 _deepcopy 覆盖了原生实现——问题开始

ascend_pytorch/pytorch/torch_npu/utils/storage.py 第54-66行：

def _deepcopy(self, memo):
    if self.device.type != 'cpu':
        memo = memo.setdefault('torch', {})
        if self._cdata in memo:
            return memo[self._cdata]
        src_tensor = torch_npu._C._tensor_construct_from_storage(self)   # ← 关键调用
        dst_tensor = src_tensor.clone()
        dst_tensor = torch_npu.npu_format_cast(dst_tensor, torch_npu.get_npu_format(src_tensor))
        new_storage = dst_tensor._typed_storage()   # ← 从 clone 后的 tensor 提取 storage
        memo[self._cdata] = new_storage
        return new_storage
    else:
        return self._new_wrapped_storage(copy.deepcopy(self._untyped_storage, memo))

第93-95行注册了 monkey-patch：

def _add_storage_methods():
    torch.storage.UntypedStorage.cpu = _cpu
    torch.storage.TypedStorage._deepcopy = _deepcopy

这覆盖了原生 torch 的 _deepcopy（pytorch/torch/storage.py 第1155-1156行）。

2.4 C++ 层：_tensor_construct_from_storage 提取了 dtype 但丢弃

ascend_pytorch/pytorch/torch_npu/csrc/npu/Module.cpp 第2021-2042行：

PyObject* THNPModule_tensor_construct_from_storage(PyObject* self, PyObject* args) {
    HANDLE_TH_ERRORS
    static torch::PythonArgParser parser(
        {"set_storage_with_format_(Storage source)"},
        /* traceable= */ false);

    torch::ParsedArgs<1> parsed_args;
    auto _r = parser.parse(args, nullptr, parsed_args);

    at::ScalarType storage_scalar_type;        // ← 会提取到 Half
    bool is_typed_storage = true;
    c10::Storage storage = _r.storage(0, storage_scalar_type, is_typed_storage);
    // ↑ storage_scalar_type = Half 被提取出来了
    return THPVariable_Wrap(
        at_npu::native::set_tensor_with_storage_format(storage));
    // ↑ 但只传了 storage，storage_scalar_type 被丢弃！
    END_HANDLE_TH_ERRORS
}

_r.storage() 通过 createStorageGetType（pytorch/torch/csrc/DynamicTypes.cpp 第79-112行）从 TypedStorage 的 dtype 属性提取了 storage_scalar_type = Half，但这个值没有传递给 set_tensor_with_storage_format。

2.5 set_tensor_with_storage_format：用 StorageDesc.data_type_ 而非 TypedStorage.dtype

ascend_pytorch/pytorch/torch_npu/csrc/aten/common/SetNpu.cpp 第120-161行（修复前）：

at::Tensor set_tensor_with_storage_format(c10::Storage src) {
    if (StorageDescHelper::CheckDescInit(src)) {
        auto desc =
            torch_npu::NPUBridge::GetNpuStorageImpl(src.unsafeGetStorageImpl())
                ->npu_desc_;
        // ↑ desc.data_type_ = Byte，因为原始 weights 用 torch.randint(dtype=uint8) 创建时记录的
        auto dist_tensor = NPUNativeFunctions::empty(
            {0}, desc.data_type_.toScalarType(), c10::nullopt, src.device(),
            false, c10::MemoryFormat::Contiguous);
        // ↑ 用 desc.data_type_ = Byte 创建 tensor，而非传入的 TypedStorage.dtype = Half
        set_storage_nd_npu(dist_tensor, src, 0,
            desc.base_sizes_.size(), desc.base_sizes_, desc.base_strides_);
        return dist_tensor;   // ← 返回的 tensor dtype = Byte
    }
    ...
}

desc.data_type_ 是在原始 weights = torch.randint(0, 255, (nbytes,), dtype=torch.uint8, device=device) 创建时，通过 StorageDescHelper::SetDesc（ascend_pytorch/pytorch/torch_npu/csrc/framework/StorageDescHelper.cpp 第275-294行）写入的，值为 Byte。

后续 view(torch.float16) 只改变了 tensor 的 dtype，不会更新底层 storage 的 npu_desc.data_type_。

2.6 回到 deepcopy：set_ 校验失败

dst_tensor._typed_storage() 返回 dtype=Byte 的 TypedStorage。

回到 pytorch/torch/_tensor.py 第234行：

new_tensor.set_(new_storage, self.storage_offset(), self.size(), self.stride())

此时 new_tensor.dtype() = Half（由 torch.empty_like(self) 创建），new_storage.dtype = Byte。

set_ 的 Python 绑定（pytorch/tools/autograd/templates/python_variable_methods.cpp 第1193-1200行）执行校验：

TORCH_CHECK(storage_scalar_type == self.dtype() || !is_typed_storage,
    "Expected a Storage of type ", self.dtype(),
    " or an UntypedStorage, but got type ", storage_scalar_type,
    " for argument 1 'storage'");

Byte != Half 且 is_typed_storage = true → 校验失败，报错。

三、原生 torch CUDA 是否有问题？——无问题

CUDA 走原生默认 _deepcopy（torch_npu 不安装时）。

原生 _deepcopy（pytorch/torch/storage.py 第1155-1156行）：

def _deepcopy(self, memo):
    return self._new_wrapped_storage(copy.deepcopy(self._untyped_storage, memo))

_new_wrapped_storage（pytorch/torch/storage.py 第905-918行）用 self.dtype（= Half）创建新 TypedStorage：

def _new_wrapped_storage(self, untyped_storage):
    new_ts = TypedStorage(
        wrap_storage=untyped_storage,
        dtype=self.dtype,          # ← 正确使用 Half
        _internal=True)
    ...
    return new_ts

因此 deepcopy 后的 new_storage.dtype = Half，与 new_tensor.dtype() = Half 匹配，set_ 校验通过。

结论：原生 torch CUDA 无此问题，因为原生 _deepcopy 通过 _new_wrapped_storage 正确保留了 TypedStorage 的 dtype。

四、问题本质总结

对比项	原生 torch	torch_npu
`_deepcopy` 实现	`_new_wrapped_storage(deepcopy(untyped_storage))`	`_tensor_construct_from_storage` → C++ 构造
dtype 来源	`self.dtype`（TypedStorage 的 dtype = Half）	`desc.data_type_`（StorageDesc 的 data_type_ = Byte）
dtype 是否正确	正确（Half）	错误（Byte）
`set_` 校验	通过（Half == Half）	失败（Byte != Half）

根本缺陷：set_tensor_with_storage_format（SetNpu.cpp:120）函数签名只接收 c10::Storage，不接收 dtype 参数；而 THNPModule_tensor_construct_from_storage（Module.cpp:2021）虽然从 Python 侧提取了 storage_scalar_type = Half，但没有传递给该函数。函数内部转而依赖 NPU StorageDesc 的 data_type_，而该值是 storage 创建时（uint8）记录的，view(torch.float16) 不会更新它。

五、修复方案

共修改 3 个文件，核心思路：将 THNPModule_tensor_construct_from_storage 中已被提取但被丢弃的 storage_scalar_type 透传到 set_tensor_with_storage_format，并在该函数中优先使用 TypedStorage 的 dtype 而非 StorageDesc 的 data_type_。

5.1 文件 1: torch_npu/csrc/aten/common/SetNpu.h（声明）

函数签名增加带默认值的 storage_scalar_type 参数，保持向后兼容：

at::Tensor set_tensor_with_storage_format(
    c10::Storage src,
    at::ScalarType storage_scalar_type = at::ScalarType::Undefined);

5.2 文件 2: torch_npu/csrc/aten/common/SetNpu.cpp（实现）

在 CheckDescInit 为 true 的分支中，优先使用传入的 storage_scalar_type，仅在未指定（Undefined）时回退到 desc.data_type_：

at::Tensor set_tensor_with_storage_format(
    c10::Storage src,
    at::ScalarType storage_scalar_type) {
  if (StorageDescHelper::CheckDescInit(src)) {
    auto desc =
        torch_npu::NPUBridge::GetNpuStorageImpl(src.unsafeGetStorageImpl())
            ->npu_desc_;
    // Prefer the TypedStorage's dtype (storage_scalar_type) over the
    // StorageDesc's data_type_, which may differ when a dtype view is
    // created over storage of a different underlying type (e.g. fp16 view
    // over uint8 storage).
    at::ScalarType dtype = (storage_scalar_type != at::ScalarType::Undefined)
        ? storage_scalar_type
        : desc.data_type_.toScalarType();
    auto dist_tensor = NPUNativeFunctions::empty(
        {0}, dtype, c10::nullopt, src.device(),
        false, c10::MemoryFormat::Contiguous);
    set_storage_nd_npu(
        dist_tensor, src, 0,
        desc.base_sizes_.size(), desc.base_sizes_, desc.base_strides_);
    return dist_tensor;
  } else {
    // (else 分支保持不变)
    ...
  }
}

5.3 文件 3: torch_npu/csrc/npu/Module.cpp（调用方）

在 Python 绑定层，当来源是 TypedStorage 时把提取到的 storage_scalar_type 传出，UntypedStorage 时传 Undefined（让被调方走回退逻辑）：

at::ScalarType storage_scalar_type;
bool is_typed_storage = true;
c10::Storage storage = _r.storage(0, storage_scalar_type, is_typed_storage);
// Only pass the TypedStorage's dtype when it is a typed storage;
// for UntypedStorage, leave it as Undefined so that the callee falls
// back to the StorageDesc's data_type_.
at::ScalarType dtype =
    is_typed_storage ? storage_scalar_type : at::ScalarType::Undefined;
return THPVariable_Wrap(
    at_npu::native::set_tensor_with_storage_format(storage, dtype));

六、修复前后对比

场景	修复前	修复后
TypedStorage(dtype=Half, 底层 uint8)	用 `desc.data_type_`=Byte 创建 tensor → 报错	用 `storage_scalar_type`=Half 创建 tensor → 正常
UntypedStorage	走 `desc.data_type_`（原逻辑不变）	`storage_scalar_type`=Undefined → 回退到 `desc.data_type_`（原逻辑不变）
普通 TypedStorage(dtype 与 storage 一致)	`desc.data_type_` 与 `storage_scalar_type` 相同	优先用 `storage_scalar_type`，结果一致