MindSpore如何在多卡环境下对计算进行加速

经验分享 Tech Blogs 其他干货-Others
1

系统环境

硬件环境(Ascend/GPU/CPU): Ascend300
MindSpore版本: mindspore=2.4.0
执行模式（PyNative/ Graph）: 不限
Python版本: Python=3.9
操作系统平台: linux

报错信息

问题描述

在Ascend300计算卡上执行一个简单的矩阵乘程序，希望能够让计算卡的两块Ascend 310P并行计算矩阵乘的一个部分。
例如，a矩阵的大小是(1024, 256)，b矩阵的大小是(256, 256)，需要计算a和b执行矩阵乘的结果。希望第一块310P负责a[0:512]和b矩阵的相乘，第二块310P负责a[512:1024]和b矩阵相乘。
但是，如果多线程调用mindspore.set_device方法的话会报错。如果多进程调用mindspore.set_device方法的话，则通过npu-smi发现两张卡无法并行执行计算。

报错信息

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

脚本信息

多线程调用

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())

多进程调用

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())

根因分析

此处由用户填写

解决方案

此处由用户填写
包含文字方案和最终脚本代码
请将正确的脚本打包并上传附件