文档
注册
评分
提单
论坛
小AI

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
搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容,请尝试其他搜索词