我需要乘以 3 个矩阵，A: 3000x100, B: 100x100, C: 100x3.6MM。我目前只是在 PyTorch 中使用普通矩阵乘法
A_gpu = torch.from_numpy(A)
B_gpu = torch.from_numpy(B)
C_gpu = torch.from_numpy(C)
D_gpu = (A_gpu @ B_gpu @ C_gpu.t()).t()
C 非常广泛，因此 gpu 上的数据重用是有限的，但是还有其他方法可以加快速度吗？我有一台配备 4 个 GPU 的机器。
【问题讨论】：
MM 是什么意思？是错字还是你的意思是 M 乘以 M，即 10^6 * 10^6 = 10^12 = T？
由于您有四个 GPU，您可以利用它们来执行高效的矩阵乘法。但请注意，乘法的结果大小为 3000x3600000，单精度浮点 (fp32) 占用 40GB。除非您有足够大的 RAM 供 CPU 使用，否则您无法将计算结果存储在 RAM 上。
对此可能的解决方案是将大矩阵C 分解为四个较小的块，在不同的 GPU 上执行每个块的矩阵乘法，并将结果保存在 GPU 上。假设每个 GPU 至少有 10GB 内存，那么您将有足够的内存用于此操作。
如果您也有足够的 CPU 内存，那么您可以将所有四个 GPU 的结果移动到 CPU 上并将它们连接起来（实际上，在这种情况下，您可以只使用单个 GPU 并将结果从 GPU 传输到CPU 每次）。否则，您可以将结果分块保存在 GPU 上，并且您需要记住并跟踪这四个块实际上是一个矩阵的一部分。
import numpy as np
import torch.nn as nn
import torch
number_of_gpus = 4
# create four matrics
A = np.random.normal(size=(3000,100))
B = np.random.normal(size=(100,100))
C = np.random.normal(size=(100,3600000))
# convert them to pytorch fp32 tensors
A = torch.from_numpy(A).float()
B = torch.from_numpy(B).float()
C = torch.from_numpy(C).float()
# calcualte `A@B`, which is easy
AB = A@B
# split the large matrix `C` into 4 smaller chunks along the second dimension.
# we assume here that the size of the second dimension of `C` is divisible by 4.
C_split = torch.split(C,C.shape[1]//number_of_gpus,dim=1)
# loop over the four GPUs, and perform the calculation on each using the corresponding chunk of `C`
D_split = []
for i in range(number_of_gpus):
device = 'cuda:{:d}'.format(i)
D_split.append( AB.to(device) @ C_split[i].to(device))
# DO THIS ONLY IF YOU HAVE ENOUGH CPU MEMORY!! :
D = torch.cat([d.cpu() for d in D_split],dim=1)
【讨论】：
根据您的矩阵C，稀疏矩阵可能会减少大小和计算时间，例如您只保存不为 0 的列，也许torch reference 可能会有所帮助。
【讨论】：
如果您有多个 GPU，则可以使用 PyTorch 的 DataParallel 在所有 GPU 上分配计算。它将在 GPU 之间拆分（并行化）矩阵 C_gpu 的列的乘法。
方法如下：
首先，导入模块并准备矩阵：
import torch
import torch.nn as nn
A_gpu = torch.from_numpy(A).float()
B_gpu = torch.from_numpy(B).float()
C_gpu = torch.from_numpy(C).float()
接下来，创建一个没有偏差的Linear“层”。这一层所做的正是矩阵乘法。输入大小将是C_gpu 每一列的大小，输出大小将是结果每一列的大小。
mat_mult = nn.Linear(in_features=C_gpu.shape[0],out_features=A_gpu.shape[0],bias=False)
将层的矩阵（=权重）设置为A_gpu @ B_gpu，这是一个无需并行化即可快速计算的小矩阵（尽管您也可以将其并行化）。
mat_mult.weight.data = A_gpu @ B_gpu
将图层转换为 DataParallel 实例。这意味着它将沿“批处理”维度自动并行计算。参数 device_ids 是您的 GPU 的索引列表（在您的情况下为 4 个）。
mat_mult_gpu = nn.DataParallel(mat_mult,device_ids=[0,1,2,3]).to('cuda:0')
现在您可以将矩阵C_gpu 输入层，计算将沿其大维度并行：
D_gpu = mat_mult_gpu(C_gpu.t())
重要提示：在撰写此答案时，我无法访问多个 GPU 来实际测试此提议的解决方案。如果有读者确认它确实有效（甚至更好 - 计时解决方案并与单个 GPU 进行比较），我将不胜感激
EDIT1：我现在在多个 GPU（四个 Nvidia Tesla P100）上尝试了这段代码，结果发现它给出了内存不足的错误。不过，我将在此处保留此解决方案作为参考，因为它确实适用于最大约 400K（而不是 3.6M）的大小。
此外，如果您将C 分成更小的块，将每个块送入mat_mult_gpu，然后在 CPU 上连接结果，则此解决方案仍适用于大小为 3.6M 的情况。请注意，您需要大量 CPU 内存才能工作，因为结果的大小为 3K-by-3.6M，在 fp32 中大约需要 40GB。 （或者，您可以将每个块保存到磁盘而不连接块）。
【讨论】：
C_gpu.shape[1]，因为这是转置矩阵时的列大小，删除了D_gpu = ... 中的转置并将.cuda() 添加到DataParallel 对象。我让 CUDA 内存不足。尝试分配 23GB。