numpy - 有什么方法可以使用多个 GPU 来提升矩阵乘法？

Question

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

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

这是我的问题:

1. 有没有python库可以自动实现我的想法？
2. 我想并行使用 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

Answer 1

任务 supports通过 CuPy 支持的数组在 GPU 上进行数组运算(包括矩阵乘法)。您可以将多节点、多 GPU 集群与 Dask-CUDA 一起使用.