torch_npu在单进程中使用多张NPU卡
在PyTorch2.1.0及以上版本中,使用torch_npu可以在一个进程中使用多个device卡,根据指定的device id将任务下发到期望的卡上执行。通过单进程多卡的支撑,在torch_npu上可以执行更灵活的设备操作。
配置如下环境变量开启虚拟内存特性时,不能使用单进程多卡特性。
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
torch_npu在单进程中使用多张NPU卡具体示例如下:
import torch import torch_npu # 默认device可以设置也可以不设置 torch.npu.set_device(0) # 使用device 0 tensor_a = torch.randn(2,3, device="npu:0") # 使用devcie 1 tensor_b = torch.randn(2,3, device="npu:1") # 跨卡的拷贝 tensor_c = tensor_a.to("npu:1") # 在device 1上进行计算 tensor_d = tensor_b + tensor_c
父主题: 附录