Масштабирование PyTorch: стратегии экспертов и практические примеры кода

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

Первый и наиболее доступный шаг — это использование нескольких GPU на одной машине (Data Parallelism). Модуль `torch.nn.DataParallel` позволяет легко обернуть вашу модель для параллельного обучения на нескольких GPU. Он автоматически разделяет входной батч на подбатчи, распределяет их по доступным GPU, собирает градиенты и выполняет обновление. Однако у этого подхода есть узкое место — главный GPU, который агрегирует результаты, что может привести к дисбалансу загрузки памяти.

Более современным и эффективным подходом является Distributed Data Parallel (DDP), доступный в `torch.nn.parallel.DistributedDataParallel`. DDP использует многопроцессный подход, где каждый процесс работает со своей моделью на отдельном GPU, и градиенты синхронизируются через коллективные операции связи (all-reduce). Это значительно уменьшает нагрузку на один узел и лучше масштабируется.

Рассмотрим базовый пример настройки DDP для обучения на одной машине с несколькими GPU.

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
 os.environ['MASTER_ADDR'] = 'localhost'
 os.environ['MASTER_PORT'] = '12355'
 dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
 dist.destroy_process_group()

def train(rank, world_size, model, dataset):
 setup(rank, world_size)
 torch.cuda.set_device(rank)
 model = model.to(rank)
 ddp_model = DDP(model, device_ids=[rank])
 # Создание DataLoader с DistributedSampler
 sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
 dataloader = DataLoader(dataset, sampler=sampler, batch_size=32)
 # Стандартный цикл обучения
 for epoch in range(epochs):
 sampler.set_epoch(epoch)
 for batch in dataloader:
 # ... forward, backward, optimization
 cleanup()

if __name__ == "__main__":
 world_size = torch.cuda.device_count()
 mp.spawn(train, args=(world_size, model, dataset), nprocs=world_size, join=True)

Этот код запускает отдельный процесс для каждого GPU. Ключевые моменты: инициализация процесса связи через `init_process_group`, использование `DistributedSampler` для не пересекающейся выборки данных каждым процессом и обертывание модели в DDP.

Следующий уровень — масштабирование на несколько машин (кластер). Принцип остается тем же, но требует корректной настройки сети. Здесь на помощь приходят инструменты оркестрации, такие как Kubernetes с оператором PyTorch или фреймворки высокого уровня, например PyTorch Lightning, который абстрагирует большую часть boilerplate-кода для распределенного обучения.

PyTorch Lightning делает реализацию DDP почти тривиальной.

import pytorch_lightning as pl
from pytorch_lightning import Trainer

class LitModel(pl.LightningModule):
 # Определение модели, training_step, configure_optimizers и т.д.
 ...

if __name__ == "__main__":
 model = LitModel()
 trainer = Trainer(accelerator="gpu", devices=4, strategy="ddp", max_epochs=10)
 trainer.fit(model, train_dataloader)

Всего несколько строк — и ваше обучение масштабируется на несколько GPU, причем Lightning сам заботится о деталях распределения.

Когда модель настолько велика, что не помещается в память даже одного GPU, на помощь приходит Pipeline Parallelism и Model Parallelism. Pipeline Parallelism делит модель на последовательные этапы, распределяя их по разным устройствам. Каждый микробатч проходит через эти этапы конвейером. PyTorch предоставляет экспериментальный API `torch.distributed.pipeline.sync.Pipe`. Более продвинутые решения предлагает FairScale от Facebook Research.

Model Parallelism предполагает ручное размещение разных частей модели на разных устройствах. Например, вы можете разместить первые слои на GPU 0, а последующие — на GPU 1.

class HybridModel(nn.Module):
 def __init__(self):
 super().__init__()
 self.part1 = nn.Linear(1000, 500).to('cuda:0')
 self.part2 = nn.Linear(500, 100).to('cuda:1')

 def forward(self, x):
 x = x.to('cuda:0')
 x = self.part1(x)
 x = x.to('cuda:1')
 x = self.part2(x)
 return x

Это требует аккуратного управления перемещением тензоров между устройствами.

Для сверхбольших моделей, таких как GPT-3 или T5, используется комбинация подходов: Data Parallelism, Pipeline Parallelism и Tensor Parallelism (разделение вычислений одного оператора, например матричного умножения, между устройствами). Фреймворки глубокого обучения, такие как DeepSpeed от Microsoft (интегрированный с PyTorch), предлагают оптимизации уровня ZeRO (Zero Redundancy Optimizer) для эффективного распределения памяти оптимизатора, градиентов и параметров по устройствам, что позволяет обучать модели с триллионами параметров.

Не забывайте и о масштабировании вывода (inference). Здесь полезны TorchScript для создания сериализуемых и оптимизируемых моделей, а также библиотеки типа TorchServe для развертывания моделей PyTorch в продакшн-среде с поддержкой многопоточности, батчинга и мониторинга.

В итоге, выбор стратегии масштабирования PyTorch зависит от ваших задач, инфраструктуры и размера модели. Начните с DataParallel или DDP для ускорения обучения на нескольких GPU. Для гигантских моделей изучите комбинированные подходы с использованием DeepSpeed или PyTorch Lightning. Инвестируйте время в освоение этих инструментов — они откроют путь к созданию и обучению state-of-the-art моделей, не ограничиваясь ресурсами одной видеокарты.