使用MindSpore如何在多卡环境下对计算进行加速
收藏回复举报
使用MindSpore如何在多卡环境下对计算进行加速
t('forum.solved') 已解决
发表于2025-06-04 17:04:30
0 查看

我需要在Ascend 300I DUO计算卡上执行一个简单的矩阵乘程序,我希望能够让计算卡的两块Ascend 310P并行计算矩阵乘的一个部分。例如,a矩阵的大小是(1024, 256),b矩阵的大小是(256, 256),我需要计算a和b执行矩阵乘的结果。我希望第一块310P负责a[0:512]和b矩阵的相乘,第二块310P负责a[512:1024]和b矩阵相乘。但是,如果多线程调用mindspore.set_device方法的话,即

import mindspore as ms
from concurrent.futures import ThreadPoolExecutor
def calc(a, b, device_id):
    ms.set_device('Ascend', device_id=device_id)
    a = a.move_to('Ascend')
    b = b.move_to('Ascend')
    return ms.mint.mm(a, b).move_to('CPU')

with ThreadPoolExecutor(max_workers=2) as executor:
    futures = []
    futures.append(executor.submit(calc, a[:512], b, 0))
    futures.append(executor.submit(calc, a[512:], b, 1))
print(futures[0].result())
print(futures[1].result())

会报错提示

RuntimeError: The 'mindspore.set_device' can not be modified.

如果多进程调用mindspore.set_device方法的话,即

import mindspore as ms
from concurrent.futures import ProcessPoolExecutor
def calc(a, b, device_id):
    ms.set_device('Ascend', device_id=device_id)
    a = a.move_to('Ascend')
    b = b.move_to('Ascend')
    return ms.mint.mm(a, b).move_to('CPU')

with ProcessPoolExecutor(max_workers=2) as executor:
    futures = []
    futures.append(executor.submit(calc, a[:512], b, 0))
    futures.append(executor.submit(calc, a[512:], b, 1))
print(futures[0].result())
print(futures[1].result())

则通过npu-smi发现两张卡无法并行执行计算,请问正确的多卡并行加速应该怎么做?

我要发帖子