Multi gpu training overview

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

Глобально следует рассмотреть 3 сценария

Если вся модель, активации, данные оптимизатора влезают в память каждой GPU
Если вся модель, активации, данные оптимизатора влезают в память хотя бы одной GPU
Если хотя бы один слой (тензор), выходная активация с него вместе с данными оптимизатора не влезает в самую большую GPU

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

Сценарий 1.

DP
DDP
ZeRO

Сценарий 2.

Pipeline Parallel
Tensor Parallel

Сценарий 3.

Tensor parallel

Offtopic: реализация на торче

Данный блок добавлен для упрощения понимания кода ниже. Если вы пришли за теоретическим представлением - можете его пропускать

Каждый последующий пример будет использовать встроенные в торч инструменты для передачи данных между устройствами (gpu или разными серверами кластера). Так как, как я полагаю, большинство читателей с ними не знакома, думаю стоит оставить этот пример кода для лучшего ~~понимания~~ привыкания к инструментам торча

Отправка одного тензора между устройствами (нодами)

import torch  
import torch.distributed as dist  
  
dist.init_process_group(backend="gloo", # CPU only backend
						# https://pytorch.org/docs/stable/distributed.html
						init_method="tcp://node1_ip:8551", 
						rank=0, # номер ноды
						world_size=2) # число нод

t = torch.tensor([1.0]*10)
dist.send(tensor=t, 
          dst=1) # отправить на 1 ноду


# ------ second node --------
dist.init_process_group(backend="gloo",
						# https://pytorch.org/docs/stable/distributed.html
						init_method="tcp://node1_ip:8551", 
						rank=1,
						world_size=2)
tensor = torch.tensor([0.0]*10)  
dist.recv(tensor=tensor, src=0)  
print(f'rank 1 received tensor {tensor} from rank 0')

DP & DDP

Овервью в доке торча Пример для линейной модели DDP в 3 строчки для lightning

DataParallel и DistributedDataParallel довольно похожие стратегии многозадачности, распишем что происходит каждый степ

DDP:

Каждый графический процессор обрабатывает свой минибатч данных.
Как только градиенты по всем весам посчитаны, все ноды синхронизируют их между собой, затем усредняют
Каждая нода изменяет свои веса модели на одинаковую величину, благодаря чему веса моделей не разбегаются между разными устройствами

DP:

GPU 0 считывает батч и затем отправляет минибатч на каждый GPU.
Обновленная модель реплицируется с GPU 0 на каждый GPU.
Каждая нода независимо выполняют forward и отправляют y на GPU0 для вычисления лоса
Лосс распределяются от GPU 0 ко всем графическим процессорам, каждый процессор выполняет backward
Градиенты с каждого графического процессора передаются в GPU 0 и усредняются.

Плюсы и минусы
+

DDP требует в 5 раз меньше обменов данными
он нагружает все гпу равномерно, когда как DP перегружает GPU:0

-

Проще архитектурно, есть одна мастер нода, которая общается со всеми
Кажется, что DDP больше страдает от не детерминированных вычислений (?)
DP по честному считает градиенты по всему батчу на одной гпу, а DDP усредняет градиенты для каждого мини батча. В некоторых кейсах, это приводит к меньшей стабильности обучения (когда лосс по батчу != сумме лосов по элементам)
При реализации в коде, вам нужно написать на один символ меньше при импорте DP, чем DDP из torch.dist

Примеры на torch

1) Обернуть модель в `model = DistributedDataParallel(model, device_ids=[local_rank])`

2) Наладить communication group между устройствами (создать обьект dist)

3)

import torch  
import torch.distributed as dist  
# ...
for inp, label in dataloader:  
	optimizer.zero_grad()  
	out = model(inp)  
	loss = loss_fun(out, lable)  
	loss.backward()  
	
	# --start--  каждый трэин степ синхронизировать параметры
	for param in model.parameters():  
		dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM)  
		param.grad.data /= dist.get_world_size()  
	# --end--  
	
	optimizer.step()

Pipeline paralelism (model paralelism)

Дока торча

model paralelism:
Эта стратегия масштабируемости обучения предлагает разбить модель по слоям между GPU, а потом последовательно выполнять forward и backward pass-ы, передавая активации и частичные производные между гпу, затем, в конце каждого степа синхронизируя градиенты (веса модели между нодами)
pipeline paralelism:
Основная проблема модельного параллелизма в том, что каждая гпу простаивает, большую часть времени (например для 4х гпу, как на картинке - процент простоя ~80%).
Эту проблему решает pipeline параллелизм - мы разбиваем батч на минибатчи и рассчитываем их последовательно

Нативно торч поддерживает эти методы только если модель это nn.Sequential, для моделй с кастомными форвардами придется параллелить ее между gpu вручную
Также такой метод не очень подходит для сложных архитектур сетей, который нельзя разбить по слоям из-за skip connection и всякой другой экзотики в строении.

Пример на torch

# init rpc
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'
torch.distributed.rpc.init_rpc('worker', rank=0, world_size=1)

# Build pipe.
fc1 = nn.Linear(16, 8).cuda(0)
fc2 = nn.Linear(8, 4).cuda(1)
model = nn.Sequential(fc1, fc2)
model = Pipe(model, chunks=8)
input = torch.rand(16, 16).cuda(0)
output_rref = model(input)

Tensor parallelism

Дока Гитхаб

При обучении больших нейросетей некоторые тензоры не могут влезть на одну видеокарту целиком. Например Embedding слой трансформера с размером словаря в 128_000 токенов, а hidden_dim = 8192 будет содержать больше миллиарда параметров, что вместе с параметрами адама в fp16 будет весить десятки гигабайт.

Решение - разбить тензор на несколько гпу.

Разбиение основано на том, что $\forall$ поэлементной функции $\sigma : R \rightarrow R$

$Z = B \sigma(Ax)\cdot = B_1\sigma(A_1 x) + B_2\sigma(A_2x)$

(это же обобщается на любую комбинацию линейных слоев и активаций между ними.

Однако:

Это выходит медленно
Это очень сильно нагружает канал связи между гпу
Скриншот из документации торча

ZeRO

arxiv

Способ улучшить DDP (baseline) от майков
Каждая гпу, вместо хранения всей реплики модели хранит только ее часть, а по необходимости запрашивает недостающие параметры с других карточек.

Снижает занимаемый объем памяти, но более требовательно к скорости соединения

Комбинации подходов

Никто не мешает использовать сразу несколько подходов при обучении одной модели. Например дипмайнд при обучении моделей на нескольких датацентрах, вероятнее всего,
в рамках каждого датацентра они использовали PP, а затем синхронизировал градиенты между по сети между центрами (DDP).

Что еще почитать

1) Если вы занимаетесь ЛЛМ, то самое популярное решение на данный момент - DeepSpeed , паралелит обучение сильно лучше торча

2) Классная статья на медиуме

3) Овервью от huggingface в контексте ЛЛМ

4) Видео от торч лайтнинг

Кластерное обучение нейросетей