Как реализовать параллелизм модели в PyTorch?

Question

Параллелизм модели в PyTorch достигается путем размещения различных частей модели на разных GPU. Для этого необходимо вручную распределить слои модели между устройствами и управлять перемещением данных.

YeaHub · Accepted Answer

Основная идея: Параллелизм модели позволяет разделить модель на несколько частей, каждая из которых выполняется на своем устройстве (например, на разных GPU). Это полезно, когда модель слишком велика, чтобы поместиться в память одного устройства.Пример реализации:В этом примере создается модель с двумя слоями, каждый из которых размещен на отдельном GPU.import torch
import torch.nn as nn

class ModelParallel(nn.Module):
    	def __init__(self, device1, device2):
        	super(ModelParallel, self).__init__()
        	self.layer1 = nn.Linear(10, 20).to(device1)
        	self.layer2 = nn.Linear(20, 30).to(device2)

def forward(self, x):
        	x = self.layer1(x.to(self.layer1.device))
        	return self.layer2(x.to(self.layer2.device))

device1 = torch.device('cuda:0')
device2 = torch.device('cuda:1')
model = ModelParallel(device1, device2) Основные моменты:Данные и слои модели должны быть вручную перенесены на соответствующие устройства.Во время прямого прохода данные перемещаются на нужное устройство перед тем, как попасть в слой. Когда использовать: Параллелизм модели полезен при обучении очень больших моделей, которые не помещаются в память одного устройства.