Эффективный способ реализации умножения матриц, когда одна матрица чрезвычайно широка?

Question

Мне нужно умножить 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 ограничено, но есть ли другие способы ускорить это?У меня есть машина с 4x графическими процессорами.

Answer 1

Если у вас есть несколько графических процессоров, вы можете распределить вычисления по всем из них, используя PyTorch DataParallel.Он разделит (распараллелит) умножение столбцов матрицы 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)

Установите для матрицы (= вес) слоя значениеbe A_gpu @ B_gpu, которая представляет собой небольшую матрицу, которая может быть быстро вычислена без распараллеливания (хотя вы также можете распараллелить ее, если хотите).

mat_mult.weight.data = A_gpu @ B_gpu

Преобразовать слой в экземпляр DataParallel.Это означает, что он будет автоматически распараллеливать вычисления по «пакетному» измерению.Аргумент device_ids представляет собой список индексов ваших графических процессоров (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)

---

EDIT1: Я сейчас попробовал этокод на нескольких графических процессорах (четыре Nvidia Tesla P100), и оказывается, что выдает ошибку нехватки памяти.Я оставлю это решение здесь для справки, поскольку оно подходит для размеров до 400 КБ (вместо 3,6 МБ).

Кроме того, это решение будет работать и для размеров 3,6M, если вы разделите C на более мелкие куски, подадите каждый кусок в mat_mult_gpu, а затем объедините результаты в ЦП.Обратите внимание, что вам нужно много памяти процессора для этого, так как результат имеет размер 3K-by-3.6M, который в fp32 занимает около 40 ГБ.(в качестве альтернативы вы можете сохранить каждый чанк на диск без объединения чанков).

Answer 2

Поскольку у вас есть четыре графических процессора, вы можете использовать их для эффективного умножения матриц.Однако обратите внимание, что результаты умножения имеют размер 3000x3600000, который занимает 40 ГБ с плавающей запятой одинарной точности (fp32).Если у вас недостаточно большой ОЗУ для ЦП, вы не сможете сохранить результаты этих вычислений в ОЗУ.

Возможное решение для этого - разбить большую матрицу C на четыре меньших блока, выполнитьумножение матрицы каждого куска на другом графическом процессоре и сохранение результата на графическом процессоре.При условии, что у каждого GPU есть по крайней мере 10 ГБ памяти, у вас будет достаточно памяти для этого.

Если у вас также достаточно памяти CPU, вы можете затем переместить результаты всех четырех GPU в CPU и объединить их(на самом деле, в этом случае вы могли бы использовать только один графический процессор и каждый раз передавать результаты из графического процессора в центральный процессор).В противном случае вы можете сохранить результаты в виде фрагментов на графических процессорах, и вам нужно помнить и отслеживать, что четыре фрагмента фактически являются частью одной матрицы.

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)

Answer 3

В зависимости от вашей матрицы C, разреженная матрица может уменьшить размер и время вычислений, например, вы сохраняете только столбцы, которые не равны 0, может быть, ссылка на горелку может помочь.