Оптимизация в распределённом обучении

LLaMA-3 (70B параметров) обучался на 1.4 триллиона токенов, 2048 GPU. GPT-4 - на ещё большем масштабе. Это невозможно без систематического подхода к распределённому обучению: как разбить модель между GPU, как эффективно синхронизировать градиенты, как масштабировать оптимизатор? Это не просто «запустить на N машинах» - это отдельная дисциплина.

• **Megatron-LM (NVIDIA)**: реализация 3D parallelism (TP+PP+DP), используется для обучения GPT-4, LLaMA - без этой системы современные LLM просто невозможны
• **BERT 76 минут**: LAMB + batch 65536 на 1024 TPU - сокращение времени обучения с 3 дней сделало BERT-like модели практичными
• **Gradient compression в федеративном обучении**: TopK и 1-bit compression позволяют обучать модели на мобильных устройствах с ограниченной пропускной способностью сети

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

• Optimization for LLM Training

Data Parallelism: DDP и Ring-AllReduce

**Data Parallelism** - простейший и самый распространённый способ распределённого обучения: каждый GPU хранит **полную копию модели**, но обрабатывает разные части данных. После backward pass градиенты суммируются между всеми GPU.

**DistributedDataParallel (DDP)** в PyTorch оптимизирует AllReduce: **Bucketing**: объединяет малые тензоры в bucket (25MB по умолчанию) для одного AllReduce вызова - снижает latency overhead. **Overlap**: начинает AllReduce градиентов слоя сразу как он вычислен, не ожидая конца всего backward pass. Пока CPU/GPU считают градиенты следующего слоя - сеть передаёт градиенты предыдущего. **Hook на backward**: `register_reduce_hook()` автоматически запускает AllReduce для каждого bucket.

Model Parallelism: Tensor и Pipeline

Когда модель не умещается на одном GPU даже без данных, нужен **Model Parallelism**: разбиение самой модели между GPU. Два основных вида: **Tensor Parallelism** (разбиение матриц внутри слоя) и **Pipeline Parallelism** (разбиение слоёв по GPU).

Реальные LLM используют все три вида параллелизма одновременно: **Tensor Parallelism (tp)**: 2-8 GPU внутри одного сервера (NVLink - быстрая связь ~600GB/s) **Pipeline Parallelism (pp)**: несколько серверов вдоль «трубы» модели **Data Parallelism (dp)**: несколько таких «трубок» обрабатывают разные batch-и Общее число GPU = tp × pp × dp Principle: tp внутри сервера (быстрая inter-GPU связь), pp между серверами (медленная inter-node связь, но мало коммуникаций), dp на верхнем уровне (AllReduce только раз в итерацию).

Сжатие градиентов: bandwidth bottleneck

AllReduce градиентов - основной коммуникационный bottleneck в Data Parallelism. Для модели с Ψ параметрами AllReduce передаёт 2Ψ байт данных. При Ψ=7B (LLaMA-7B) в fp16 это 14GB за итерацию. **Gradient Compression** уменьшает объём передаваемых данных.

**Без Error Feedback**: если отбросить малые градиенты, они потеряются навсегда → смещение оценки градиента → расходимость. **С Error Feedback**: несохранённые градиенты накапливаются в `residual`. В следующей итерации они добавляются к новым градиентам и участвуют в следующей передаче. Математически доказано: TopK с error feedback сходится к тем же решениям, что и без сжатия, только немного медленнее. Практически: при TopK с K=0.1% замедление ≈ 5-10%, экономия bandwidth ≈ 1000×.

Масштабирование оптимизатора: batch size и LARS/LAMB

При увеличении числа GPU в Data Parallelism обычно пропорционально увеличивают batch size (больше GPU → больше данных за итерацию). Это поднимает вопрос: как масштабировать learning rate, чтобы модель обучилась так же хорошо?

Linear Scaling Rule работает только при **стабильном состоянии** обучения, но не в начале. В начале обучения: - Веса случайны, градиенты нестабильны - Большой lr с momentum сразу на большом batch → расходимость **Правило**: при увеличении batch size в k раз использовать **gradual warmup**: - Начать с lr = base_lr (маленький) - За T_warmup шагов линейно увеличить до k × base_lr - T_warmup = 5 эпох в экспериментах Facebook Это стабилизирует обучение независимо от масштаба batch size.

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

• **Data Parallelism**: каждый GPU - полная модель, разные батчи, Ring-AllReduce с O(1) per-GPU bandwidth независимо от N GPU
• **Tensor Parallelism**: разбивает матрицы слоёв (intra-node NVLink), Pipeline Parallelism - слои между GPU (inter-node, micro-batches для уменьшения bubble)
• **Gradient Compression**: TopK/1-bit/PowerSGD - до 1000× сжатие передаваемых данных; Error Feedback гарантирует корректность
• **Linear Scaling + LARS**: lr ∝ batch_size для умеренных масштабов; LARS/LAMB адаптирует lr по-слойно для batch 32K+ с gradual warmup

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

Распределённое обучение использует все аспекты оптимизации:

• Оптимизация при обучении LLM — ZeRO sharding - часть распределённого обучения; mixed precision и warmup применяются вместе с DDP и tensor parallelism
• Адаптивные методы оптимизации — LAMB - адаптивный оптимизатор для распределённого обучения; Adam scaling с большими батчами требует LARS-like adaptation
• Стохастическая оптимизация — Distributed SGD - расширение стохастической оптимизации; noise в мини-батче влияет на сходимость при разных уровнях параллелизма

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

• У вас 4 сервера по 8 GPU (A100 80GB) и модель 70B параметров. Как вы настроите TP, PP, DP и почему именно так?
• Почему gradient compression (TopK) практически не используется для transformer pretraining, хотя даёт 1000× сжатие?
• Linear Scaling Rule работает хорошо до batch ≈32K, но ломается при больших батчах. Что происходит и почему LARS решает эту проблему?

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

• calc-01-sequences