Efficient way to implement matrix multiplication when one matrix is extremely wide?
I need to multiply 3 matrices, A: 3000x100, B: 100x100, C: 100x3.6MM.  I currently am just using normal matrix multiplication in 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 is very wide so the data reuse on gpu is limited but are there other ways to speed this up? I have a machine with 4x GPUs.

Since you have four GPUs, you can harness them to perform efficient matrix multiplication. Notice however that the results of the multiplication has size 3000x3600000, which takes up 40GB in single precision floating point (fp32). Unless you have a large enough RAM for the CPU, you cannot store the results of this computation on the RAM.

A possible solution for this is to break up the large matrix C into four smaller chunks, perform the matrix multiplication of each chunk on a different GPU, and keep the result on the GPU. Provided that each GPU has at least 10GB of memory, you will have enough memory for this.

If you do have also enough CPU memory, you can then move the results of all four GPUs onto the CPU and concatenate them (in fact, in this case you could have used only a single GPU and transfer the results from GPU to CPU each time). Otherwise, you can keep the results chunked on the GPUs, and you need to remember and keep track that the four chunks are actually part of one matrix.

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)

