numpy 有没有办法使用多个GPU来提高矩阵乘法的性能?

r1zhe5dt  于 2022-12-18  发布在  其他
关注(0)|答案(3)|浏览(147)

我想把两个巨大的矩阵相乘,大小超过100,000行和列。我运行的服务器上有几个GPU的任务,让我们说8 RTX 3090 GPU,他们的内存大小是24 GB,显然,矩阵无法容纳它,所以我不能直接使用cupy.array。以下是我的想法:
1.使用numpy.array在主存中存储两个矩阵
1.把它们切成块,可能是4块或9块
1.将数据块发送到GPU进行计算
1.检索结果块到主存储器,重新组合它们
以下是我的问题:

  1. python中有没有什么库可以自动实现我的想法?
    1.我想并行使用GPU,我认为瓶颈是主内存和GPU内存之间的数据传输,即numpy.array -〉cupy. array。我可以使用多处理库并行移动数据吗?PCIe总线如何?
    注:
    1.假设矩阵不是稀疏的。
[[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
xmjla07d

xmjla07d1#

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

5cg8jx4n

5cg8jx4n2#

了解“cuBLAS多GPU扩展”:https://developer.nvidia.com/cublas
你必须申请早期访问计划。现有的python库可能不会利用这个扩展,但你可以在更新CUDA库后启用它。你必须在获得访问权限后阅读文档。

kpbwa7wx

kpbwa7wx3#

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

相关问题