CUDA内存不足，矩阵乘法

Question

我正在尝试将3个矩阵相乘，但是用完了CUDA内存。

# A: 3000 x 100 (~2MB)
# B: 100 x 100  (~0.05MB)
# C: 100 x 3MM  (~2GB)

A = np.random.randn(3000, 100)
B = np.random.randn(100, 100)
C = np.random.randn(100, 3e6)

A_gpu = torch.from_numpy(A).cuda()
B_gpu = torch.from_numpy(B).cuda()
C_gpu = torch.from_numpy(C).cuda()

R_gpu = (A_gpu @ B_gpu @ C_gpu)

Cuda正在为此操作请求大约90GB的内存。我不明白为什么。

Answer 1

乘以矩阵，您的输出大小将为3,000 x 3,000,000矩阵！因此，尽管A和B相对较小，但输出R却是巨大的：9 G 个元素。此外，我怀疑您的矩阵中的dtype是float64而不是float32（因为您使用numpy来初始化它们）。因此，R_gpu的每个9G元素都需要8个字节；仅使R_gpu的大小至少达到72 GB 个GPU内存。我怀疑中间结果和其他一些东西会占用您更多的GPU内存。