有没有办法使用多个 GPU 来提升矩阵乘法？答案

【问题标题】：Is there any way to boost matrix multiplication using multiple GPUs?有没有办法使用多个 GPU 来提升矩阵乘法？
【发布时间】：2022-01-27 15:58:59
【问题描述】：

我想将两个巨大的矩阵相乘，大小超过 100,000 行和列。我在有多个 GPU 的服务器上运行任务，比如说 8 个 RTX 3090 GPU，它们的内存大小是 24GB，显然，矩阵无法容纳它，所以我不能直接使用 cupy.array。这是我的想法：

在主内存中存储两个矩阵，使用 numpy.array
将它们切成块，可能是 4 块或 9 块
将块发送到 GPU，进行计算
将生成的块检索到主内存，重新组装它们

这是我的问题：

python中有没有可以自动实现我的想法的库？
我想并行使用GPU，我认为瓶颈是主内存和GPU内存之间的数据传输，即numpy.array -> cupy.array。我可以使用多处理库并行移动数据吗？ PCIe 总线呢？

注意：

假设矩阵不是稀疏的。

[[a1,b1],   *   [[a2,b2],   =   [[a1a2+b1c2, a1b2+b1d2],
 [c1,d1]]        [c2,d2]]        [c1a2+d1c2, c1b2+d1d2]]

import cupy as cp
import numpy as np

N = 27000
P = 27000

# init two matrices
source1 = np.random.random((N * 2, P * 2))
source2 = np.random.random((N * 2, P * 2))

# cut them in blocks
a1 = source1[:N, :P]
b1 = source1[:N, P:]
c1 = source1[N:, :P]
d1 = source1[N:, P:]

a2 = source2[:N, :P]
b2 = source2[:N, P:]
c2 = source2[N:, :P]
d2 = source2[N:, P:]

# move a1 and a2 to one gpu
m1 = cp.array(a1)
m2 = cp.array(a2)
r1 = m1 * m2
# free memory so that m3 and m4 can fit in gpu's ram
del m1
del m2

# move b1 and c2 to one gpu
m3 = cp.array(b1)
m4 = cp.array(c2)
r2 = m3 * m4
del m3
del m4
r1 += r2

【问题讨论】：

您是否要求在 GPU 上完成？
是的，因为矩阵非常大，在 CPU 上将它们相乘可能需要几个小时。根据我的实验，使用一个 GPU 只需几分钟。
考虑 pytorch（或者可能是 tensorflow）。它得到很好的支持并与 numpy 紧密集成。我对 pyopencl 和 numba 的结果好坏参半。

标签： python numpy gpu matrix-multiplication cupy

【解决方案1】：

Dask supports 通过 CuPy 支持的数组在 GPU 上进行数组运算（包括矩阵乘法）。您可以通过Dask-CUDA 使用多节点、多 GPU 集群。

【讨论】：

【解决方案2】：

查看“cuBLAS 多 GPU 扩展”：https://developer.nvidia.com/cublas

您必须申请抢先体验计划。现有的 python 库可能不会利用此扩展，但您可以在更新 CUDA 库后启用它。获得访问权限后，您必须阅读文档。

【讨论】：

【解决方案3】：

Python有一个特殊的库：https://documen.tician.de/pycuda/

简单例子：

import pycuda.autoinit
import pycuda.driver as drv
import numpy

from pycuda.compiler import SourceModule
mod = SourceModule("""
__global__ void multiply_them(float *dest, float *a, float *b)
{
  const int i = threadIdx.x;
  dest[i] = a[i] * b[i];
}
""")

multiply_them = mod.get_function("multiply_them")

a = numpy.random.randn(400).astype(numpy.float32)
b = numpy.random.randn(400).astype(numpy.float32)

dest = numpy.zeros_like(a)
multiply_them(
        drv.Out(dest), drv.In(a), drv.In(b),
        block=(400,1,1), grid=(1,1))

print dest-a*b

【讨论】：