Distributed Training: масштабирование на кластер

В 2020 году OpenAI опубликовала GPT-3: 175 млрд параметров, обучение заняло несколько месяцев на тысячах GPU. В 2023 году команда из 13 человек обучила LLaMA-65B на 2048 A100 за 21 день - с открытым кодом и подробным описанием distributed training стека. Разрыв между возможностями сократился с лет до месяцев.

• **Meta** обучает Llama 3 на кластерах из 16 000+ GPU используя комбинацию Data Parallel, Tensor Parallel и Pipeline Parallel - трёхмерный параллелизм, координация которого сама по себе инженерный подвиг.
• **Mistral AI** показала, что хорошо подобранный distributed training stack позволяет обучить 7B модель, конкурирующую с 13B моделями - эффективность инфраструктуры прямо влияет на качество результата.
• **DeepSpeed ZeRO** от Microsoft позволил командам из академии обучать 10B+ модели на кластерах из 8-16 GPU - сделав frontier-исследования доступными без 100-миллионного бюджета.

Parameter server, Horovod и ZeRO

Mu Li с соавторами представили архитектуру parameter server на OSDI 2014, масштабируя обучение через шардирование параметров по машинам. В 2017 году Alexander Sergeev и Mike Del Balso из Uber выпустили Horovod, использующий ring-allreduce для почти линейного масштабирования data-parallel обучения на множестве GPU с минимальными изменениями кода. Когда модели перестали помещаться в память одного устройства, ZeRO от Microsoft (часть DeepSpeed, 2020) распределил состояния оптимизатора, градиенты и параметры по устройствам, открыв обучение моделей с сотнями миллиардов параметров.

Предварительные знания

• Цикл обучения и вычисление градиентов
• Оптимизаторы и как градиенты обновляют параметры
• Базовые идеи GPU и обработки мини-батчей

• Backpropagation: как нейронные сети учатся
• Оптимизация: SGD -> AdamW

Data Parallelism: один граф, много GPU

GPT-3 с 175 миллиардами параметров требует минимум 700 GB только для хранения весов в float32. Один A100 GPU имеет 80 GB памяти. Это значит, что обучение таких моделей физически невозможно на одном устройстве - и распределенное обучение перестало быть оптимизацией, став необходимостью.

**Data Parallelism (DP):** каждый GPU хранит полную копию модели, получает разный мини-батч данных. После forward+backward pass градиенты агрегируются (AllReduce). Ограничение: модель должна помещаться в память одного GPU. **DDP (Distributed Data Parallel)** в PyTorch - синхронный, каждый GPU обменивается градиентами. **Масштабирование:** N GPU дают примерно N-кратное ускорение на большом батче (linear scaling, ограничен communication overhead).

Model Parallelism: разрезать модель по GPU

Когда модель не влезает в память одного GPU - её разрезают на части. Tensor Parallelism (TP) - разрезает отдельные матрицы: каждый GPU хранит и вычисляет часть одного слоя. Используется Megatron-LM, Tensor-cores A100 работают эффективнее с большими матрицами. Главный враг - communication overhead между GPU при каждом forward pass.

**Tensor Parallelism:** матрица attention Q размером (4096 x 4096) делится по столбцам между 8 GPU: каждый хранит (4096 x 512). Вычисление частичное на каждом GPU + AllReduce для агрегации. **Column Parallel:** матрица A делится по столбцам, результаты конкатенируются. **Row Parallel:** матрица B делится по строкам, результаты суммируются. **Требование:** GPU должны быть на одном сервере с NVLink/NVSwitch - PCIe слишком медленный для TP.

Pipeline Parallelism: микробатчи против простоя

Наивный Model Parallel страдает от pipeline bubble: пока GPU N обрабатывает данные, GPU 0...N-1 простаивают. Pipeline Parallelism решает это микробатчами: батч разбивается на M микробатчей, которые прокачиваются через pipeline как конвейер, минимизируя простой. GPT-3 обучался с комбинацией DP + TP + PP на тысячах GPU.

**Pipeline schedule:** в GPipe каждый GPU обрабатывает M микробатчей подряд forward, потом backward. Простой (bubble) = (p-1)/(p-1+m) где p - число стадий, m - число микробатчей. При m >> p bubble пренебрежимо мал. **1F1B schedule** (PipeDream): вместо all-forward-then-all-backward чередует forward и backward для снижения пиковой памяти. **3D parallelism:** DP + TP + PP комбинируются. TP внутри узла (NVLink), PP между узлами (InfiniBand), DP между репликами.

DeepSpeed ZeRO: уничтожение избыточности памяти

При обучении LLM на 16 GPU в DDP каждый GPU хранит полную копию весов, градиентов и состояния оптимизатора - это 3 идентичные копии на всех GPU. ZeRO (Zero Redundancy Optimizer) от Microsoft устраняет эту избыточность, разделяя хранение между GPU без увеличения communication overhead.

**ZeRO стадии:** ZeRO-1 - разделить состояние оптимизатора (Adam: m, v, fp32 веса) - ~4x экономия памяти. ZeRO-2 - добавить разделение градиентов - ~8x. ZeRO-3 - добавить разделение самих параметров - ~64x при N GPU. Каждый GPU хранит 1/N часть, при необходимости собирает полный тензор через AllGather. **ZeRO-Offload:** выгрузить оптимизатор на CPU/RAM - позволяет обучать 10B модели на одном GPU. **Практика:** DeepSpeed ZeRO-3 + bf16 + gradient checkpointing - стандартный стек для обучения 7-70B моделей.

Ключевые идеи

• **Data Parallelism** - каждый GPU держит полную модель, обрабатывает разные батчи, градиенты синхронизируются AllReduce. Ограничение: модель должна влезать в GPU.
• **Model Parallelism** (Tensor + Pipeline) - разрезает модель на части. TP - внутри узла по NVLink. PP - между узлами, микробатчи снижают pipeline bubble до <10%.
• **DeepSpeed ZeRO-3** устраняет избыточность хранения: каждый GPU держит 1/N параметров, градиентов и состояния оптимизатора - 64x экономия при N=8 против наивного DDP.

Связанные темы

Distributed Training применяется поверх техник из смежных уроков:

• Transfer Learning и Fine-Tuning — Fine-tuning больших моделей (7B+) требует ZeRO или LoRA для экономии памяти
• Обучение нейронных сетей — Оптимизаторы (Adam, ZeRO-1) и gradient checkpointing - основа эффективного distributed обучения

Вопросы для размышления

• При каком соотношении compute/communication имеет смысл переходить от DDP к ZeRO-3?
• Почему Pipeline Parallelism размещают между узлами, а Tensor Parallelism - внутри одного узла?
• Как gradient checkpointing влияет на компромисс между памятью и скоростью обучения?

Связанные уроки

• dl-11 — Fine-tuning больших моделей требует распределённой памяти
• dl-10 — Оптимизаторы и градиенты лежат в основе обучения
• dl-20 — Дизайн продакшен-систем строится на тренировочной инфраструктуре
• ml-54-distributed-training — Те же идеи параллелизма данных в классическом ML
• ds-05-replication — AllReduce похож на репликацию и согласованность в системах
• la-34-la-in-dl — Tensor parallelism делит матричные операции между устройствами