Ascend for PyTorch 训练营学习笔记第五期

设备管理机制 Ascend for PyTorch 训练营(https://www.chaspark.com/#/s/AscendforPyTorch2026?multi=zh) 第五期

跨设备数据搬运（tensor.to('npu')）

1. Python 层：x.to("npu") 触发
2. Dispatch 路由：PyTorch dispatcher 识别目标设备，路由到 copy kernel
3. NPUGuard：c10_npu::NPUGuard 自动切换当前线程到目标 NPU 设备（构造时切、析构时还原）
4. 底层拷贝：调用 aclrtMemcpy（H2D）在 currentStream 上执行实际字节搬运

void copy_h2d_baseformat_opapi(...) {
    c10_npu::NPUGuard guard(dst.device());  // RAII 自动切设备
    copy_h2d_baseformat_dtype_contiguous_opapi(dst, src, non_blocking);
}

多流（Stream）并行

• 流 = 设备侧有序命令队列，host 提交后立即返回，device 按 FIFO 串行执行
• NPU 对应 aclrtStream（类比 CUDA 的 cudaStream_t）

with torch.npu.stream(s1):
    kernel1(tensors)          # 提交到流 s1
with torch.npu.stream(s2):
    kernel2(other_tensors)    # 提交到流 s2，与 s1 并发

目的：计算与通信/数据拷贝重叠，隐藏传输延迟，提升吞吐量。

PyTorch 2.x 引入 torch.accelerator，把"当前用哪个加速后端"抽象成统一入口，让同一份脚本跨 CUDA/NPU 通用，无需写死 torch.cuda 或 torch.npu。

import torch
# 不写死 torch.cuda / torch.npu
if torch.accelerator.is_available():
    dev = torch.accelerator.current_device()   # cuda:0 或 npu:0
    torch.accelerator.set_device(0)
    x = torch.randn(4, device=dev)
    torch.accelerator.synchronize()              # 统一同步入口

底层链路：

• torch.accelerator.* → at::getAccelerator() 选 DeviceType（PrivateUse1 > CUDA）
• → VirtualGuardImpl(type) 查 GuardImpl 注册表
• → 后端实现 CUDAGuardImpl / NPUGuardImpl

设备切换的两种触发方式

torch.accelerator 提供统一跨后端编程接口；设备切换既支持 Python 显式设置，也支持 C++ 层通过 guard 自动管理，最终都落到 ACL Runtime 的 aclrtSetDevice