Ascend for PyTorch test_state_dict_utils 用例修复
**作者**：昇腾实战派
**知识地图**：[https://blog.csdn.net/Lumos\_Lovegood/article/details/161601003](https://blog.csdn.net/Lumos_Lovegood/article/details/161601003)
## 1. 背景概述
在 PyTorch 分布式训练框架中，`test_state_dict_utils` 测试用例用于验证状态字典（state dict）的创建与转换功能。其中 `test_create_cpu_state_dict` 用例旨在测试在 CPU 上创建共享内存且支持 pinned memory 的状态字典。然而，在 NPU 环境下运行该用例时，出现了 `AttributeError` 异常，导致测试失败。本文记录了该问题的排查过程、根因分析以及修复方案，为类似环境适配问题提供参考。
## 2. 问题现象
### 2.1 测试命令
```bash
cd /path/to/pytorch-2.7.1/test
CI= python distributed/checkpoint/test_state_dict_utils.py -k test_create_cpu_state_dict
```
### 2.2 失败用例
`TestStateDictUtils.test_create_cpu_state_dict`
### 2.3 原始报错
关键错误信息：
```
AttributeError: 'NoneType' object has no attribute 'cudaHostRegister'
```
退出时可能出现：
```
AttributeError: 'NoneType' object has no attribute 'cudaHostUnregister'
```
### 2.4 调用栈分析
- 测试入口：`test/distributed/checkpoint/test_state_dict_utils.py::test_create_cpu_state_dict`
- 进入函数：`torch.distributed._state_dict_utils._create_cpu_state_dict`
- 触发条件：当 `share_memory=True` 且 `pin_memory=True` 时，执行 `torch.cuda.cudart().cudaHostRegister/Unregister(...)`
## 3. 根因分析
| 维度 | 说明 |
|------|------|
| **触发条件** | `_create_cpu_state_dict(..., share_memory=True, pin_memory=True)` |
| **PyTorch 行为** | 为让共享内存也支持 pinned，PyTorch 直接调用 CUDA cudart 的 `cudaHostRegister/Unregister` 对 CPU 内存进行注册/反注册，并使用 `assert succ == 0` 进行强校验 |
| **环境现状** | 当前环境无 CUDA，`torch.cuda.cudart()` 返回 `None` |
| **直接原因** | 对 `None` 对象调用 `cudaHostRegister(...)` / `cudaHostUnregister(...)` 导致 `AttributeError` |
| **关键约束** | 不允许修改 PyTorch 源码，需优先在 `torch_npu` 侧修复 |
**结论**：该问题是 PyTorch 的 CUDA 专用 pinned-shared 路径在 NPU/无 CUDA 环境下缺少等价实现所导致的失败。
## 4. 修复方案
### 4.1 方案取舍
**曾考虑方案**：返回一个假的 cudart 对象，让 `cudaHostRegister` 直接返回 0 以规避崩溃。
**否决原因**：
- 逻辑上虽然通过了断言，但未真实完成 host memory 注册，行为无法与 GPU 对齐
- 不符合 API 一致性要求
**最终方案**：在 NPU 环境下，使用 AscendCL（ACL）的真实接口完成 host 内存注册/反注册，并对外提供 CUDA cudart 兼容入口。
### 4.2 实现思路
在 `torch_npu/contrib/transfer_to_npu.py` 中：
1. **新增 `_AclCudart` 类**：使用 `ctypes` 动态加载 AscendCL 动态库，封装以下接口：
- `aclrt_host_register` → `aclrtHostRegister`
- `aclrt_host_unregister` → `aclrtHostUnregister`
2. **保留 CUDA 兼容入口**：`cudaHostRegister` / `cudaHostUnregister` 仅做转发，供 PyTorch 上游继续调用
3. **通过 `_patch_cuda_cudart()` 对 `torch.cuda.cudart` 进行包装**：
- 若原本 `torch.cuda.cudart()` 返回非空（真实 CUDA 环境），不干预
- 若返回 `None` 且 `torch.npu.is_available()`，返回 `_AclCudart()` 实例
### 4.3 核心代码
修改文件：`/path/to/torch_npu/contrib/transfer_to_npu.py`
```python
class _AclCudart:
def __init__(self):
import ctypes
from ctypes.util import find_library
lib_names = []
for cand in ("ascendcl", "acl"):
name = find_library(cand)
if name:
lib_names.append(name)
lib_names += ["libascendcl.so", "libacl.so"]
last_err = None
lib = None
for name in lib_names:
try:
lib = ctypes.CDLL(name)
break
except OSError as e:
last_err = e
if lib is None:
raise RuntimeError(
f"Failed to load AscendCL library for host register: {last_err}"
)
self._ctypes = ctypes
self._lib = lib
self._lib.aclrtHostRegister.argtypes = [
ctypes.c_void_p,
ctypes.c_uint64,
ctypes.c_int, # aclrtHostRegisterType
ctypes.POINTER(ctypes.c_void_p),
]
self._lib.aclrtHostRegister.restype = ctypes.c_int
self._lib.aclrtHostUnregister.argtypes = [ctypes.c_void_p]
self._lib.aclrtHostUnregister.restype = ctypes.c_int
def aclrt_host_register(self, ptr, size, register_type: int = 0):
dev_ptr = self._ctypes.c_void_p()
return int(
self._lib.aclrtHostRegister(
self._ctypes.c_void_p(int(ptr)),
int(size),
int(register_type),
self._ctypes.byref(dev_ptr),
)
)
def aclrt_host_unregister(self, ptr):
return int(self._lib.aclrtHostUnregister(self._ctypes.c_void_p(int(ptr))))
# Compatibility entrypoints for upstream CUDA-style cudart APIs.
def cudaHostRegister(self, ptr, size, flags=0):
# Upstream passes flags=1 (cudaHostRegisterPortable). AscendCL uses a register
# type (aclrtHostRegisterType) instead of CUDA flags; we map to MAPPED (0).
return self.aclrt_host_register(ptr, size, register_type=0)
def cudaHostUnregister(self, ptr):
return self.aclrt_host_unregister(ptr)
def _patch_cuda_cudart():
cudart_fn = getattr(torch.cuda, "cudart", None)
if cudart_fn is None or not callable(cudart_fn):
return
@wraps(cudart_fn)
def _wrapped_cudart(*args, **kwargs):
ret = cudart_fn(*args, **kwargs)
if ret is not None:
return ret
if torch.npu.is_available():
return _AclCudart()
return ret
torch.cuda.cudart = _wrapped_cudart
```
## 5. 验证步骤
### 5.1 修改工作副本
按上述代码修改 `/path/to/torch_npu/contrib/transfer_to_npu.py`。
### 5.2 优先加载工作副本验证
使用 `PYTHONPATH` 确保导入的是工作副本而非 site-packages：
```bash
cd /path/to/pytorch-2.7.1/test
PYTHONPATH=/path/to/torch_npu:$PYTHONPATH CI= python distributed/checkpoint/test_state_dict_utils.py -k test_create_cpu_state_dict
```
### 5.3 同步到安装目录
```bash
yes | cp /path/to/torch_npu/contrib/transfer_to_npu.py \
/usr/local/python311/lib/python3.11/site-packages/torch_npu/contrib/transfer_to_npu.py
```
### 5.4 按原命令复测
```bash
cd /path/to/pytorch-2.7.1/test
CI= python distributed/checkpoint/test_state_dict_utils.py -k test_create_cpu_state_dict
```
### 5.5 预期结果
- 退出码为 0
- 输出 `Ran 1 test ... OK`
- 不再出现 `cudaHostRegister`/`cudaHostUnregister` 的 `NoneType` 报错
## 6. 经验小结
| 要点 | 说明 |
|------|------|
| **看清触发条件** | 仅在 `share_memory=True && pin_memory=True` 时，PyTorch 才会走 cudart 的 host register 路径 |
| **拒绝“假成功”** | 直接返回 0 虽能让断言通过，但不等价 GPU pinned 行为；更正确的做法是用 NPU 侧真实接口实现注册/反注册 |
| **最小侵入** | 不修改 PyTorch 源码，只在 `torch_npu` 侧通过 `torch.cuda.cudart` 的包装补齐能力 |
| **参数不完全对等** | CUDA `flags=1`（portable）在 ACL `aclrtHostRegister` 中没有同构 flags；当前映射为 `ACL_HOST_REGISTER_MAPPED (0)`，满足“注册/反注册”的核心语义与断言要求 |
| **同步很关键** | 不同步到 site-packages，实际运行仍会加载旧逻辑，表现为报错仍在 `torch.distributed._state_dict_utils` 的 `cudaHostRegister` 行 |
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。
更多推荐
0
0- 0
已为社区贡献1条内容
所有评论(0)