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

经验分享 Tech Blogs 功能调试-Function Debugging
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())

根因分析

mindspore中的并行不是这么用的，mindspore支持多种并行的方式，根据描述应该是需要用多个卡进行算子级并行或者tensor并行推理来加速。

解决方案

多个卡进行算子级并行或者tensor并行推理来加速，可以参考以下相关文档：
算子级并行
高维张量并行

不过310P上目前并不支持，因为mindspore中任何并行方式都需要先调用mindspore.communication.init方法进行初始化，310P环境下init方法无法成功执行。