Matmul on multiple GPUs

Uncategorized
2

There isn’t an automatic way to do this. If A is a list of Tensors, each on a separate GPU, I presume A is a large matrix, with rows 0 to i on GPU0, i to j on GPU1, etc.

Let’s assume B is only on GPU 0, because you didn’t mention anything about B.

Here’s a sample snippet showing how to parallelize this operation over multiple GPUs and collect the result on GPU0.

import torch

ngpu = torch.cuda.device_count()

N, M, K = 128, 256, 512 # example matmul sizes

A = []
B = torch.randn(M, K, device='cuda:0')

# randomly initialize A
for i in range(ngpu):
    # each GPU has a slice of A
    A.append(torch.randn(N // ngpu, M, device='cuda:' + str(i)))

# now let's matmul

# Step 1: make a copy of B on each GPU
B_ = [B]
for i in range(ngpu):
    if i != 0:
        B_.append(B.to('cuda:' + str(i)))

# Step 2: issue the matmul on each GPU
C_ = []
for i in range(ngpu):
    C_.append(torch.matmul(A[i], B_[i]))

C = torch.empty(N, K)
for i in range(ngpu):
    start_index = i * (N//ngpu)
    C[start_index:start_index+(N//ngpu), :].copy_(C_[i])

# C is the final result gathered on GPU-0