Оптимизация при обучении LLM

GPT-3 стоил 5M на обучение. Небольшая ошибка в learning rate schedule или precision strategy - и всё нужно начинать заново. Как обучить модель с 175 миллиардами параметров, когда она не умещается ни в один GPU? Как сделать так, чтобы fp16 не «обнулил» градиенты? Оптимизация обучения LLM - это отдельная дисциплина на стыке системного программирования и численных методов.

• **GPT-4, LLaMA-3**: используют ZeRO-3 + tensor parallelism + pipeline parallelism - без этого обучение на тысячах GPU невозможно
• **bf16 vs fp16**: переход на bf16 в A100 GPU позволил избавиться от GradScaler - bf16 имеет тот же диапазон что fp32, нет overflow
• **Cosine LR с warmup**: стандарт во всех современных LLM (GPT-3, LLaMA, Mistral, Falcon) - отступление от него требует тщательного обоснования

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

• Adaptive Optimization Methods
• Neural Network Loss Landscapes

Mixed Precision Training: fp16 + fp32

**Mixed Precision Training** (Micikevicius et al., 2018) позволяет обучать большие модели в 2× быстрее и с вдвое меньшей памятью, используя 16-битные числа для большинства операций, но сохраняя 32-битную точность там, где это критично.

Проблема: градиенты в начале обучения очень маленькие - в fp16 они становятся нулём (underflow). **Loss Scaling**: умножить loss на большое число S (например, 2¹⁶=65536) перед backward pass. Все градиенты масштабируются на S. Затем перед обновлением параметров разделить на S. **Dynamic Loss Scaling**: автоматически подбирает S: - Если overflow (inf/nan в градиентах) → уменьшить S вдвое, пропустить шаг - Если overflow не было 2000 шагов → увеличить S вдвое PyTorch: `torch.cuda.amp.GradScaler()` делает всё автоматически.

Gradient Checkpointing: компромисс памяти и вычислений

**Gradient Checkpointing** (также - activation recomputation) - техника уменьшения потребления памяти при обучении глубоких сетей за счёт **пересчёта активаций** во время backward pass вместо их хранения.

Не все слои одинаково выгодно делать checkpoint: **Дорогие по памяти, дешёвые по вычислению** → хорошо делать checkpoint: - Attention layers (O(n²) активаций по длине последовательности) - MLP blocks (большие матричные умножения) **Дешёвые по памяти** → не стоит делать checkpoint: - Embedding layers - LayerNorm - Residual connections В трансформерах обычно делают checkpoint каждый transformer block - это даёт 10-20× экономию памяти с ~30% замедлением.

ZeRO: шардинг состояния оптимизатора

**ZeRO** (Zero Redundancy Optimizer, Rajbhandari et al., 2020) - система DeepSpeed для обучения моделей, не умещающихся в память одного GPU. Она устраняет **избыточность** хранения: стандартный DDP хранит полную копию состояния оптимизатора на каждом GPU.

**ZeRO-Infinity** (2021) расширяет шардинг на CPU RAM и NVMe SSD: - **ZeRO-Offload**: optimizer state на CPU RAM (медленнее, но дёшево) - **ZeRO-Infinity**: параметры частично на NVMe через NVMe-Direct-IO Это позволяет обучать модели в триллион параметров на кластере обычных серверов, не требуя специального оборудования типа TPU Pod. Связь с pipeline parallelism: ZeRO-3 + pipeline + tensor parallelism = «3D parallelism» используемый для обучения GPT-4, LLaMA-3 и т.д.

Learning Rate Schedule для LLM

**Learning Rate Schedule** - как меняется lr в процессе обучения. Для LLM правильный schedule критичен: слишком большой lr в начале → расходимость, слишком маленький в конце → недообучение. Warmup + cosine decay стал стандартом.

**Интуиция**: в начале обучения Adam's second moment v₀=0. Bias correction частично компенсирует это, но на первых шагах v̂ₜ всё ещё нестабильна - она оценивает дисперсию только по t точкам. **Mathematically**: при t=1, v̂₁ = g₁²/(1-β₂) может быть очень большим для параметров с первоначально большими градиентами → адаптивный lr нестабилен. **Empirically**: Lion (2023) и другие работы показали, что warmup помогает любому оптимизатору, не только Adam - это явление более фундаментальное, связанное с ранней стадией «выравнивания» весов модели.

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

• **Mixed Precision**: fp16 для forward/backward, fp32 мастер-веса + loss scaling → 2× скорость, 2× память; bf16 предпочтительнее на A100
• **Gradient Checkpointing**: сохранять √n активаций, пересчитывать промежуточные → O(√n) памяти за +33% вычислений
• **ZeRO-1/2/3**: шардинг optimizer state / gradients / params между GPU - до 64× экономии памяти при N=64 GPU
• **Warmup + Cosine**: linear warmup (1-3% шагов) + cosine decay до 10% lr_max - стандарт для стабильного обучения LLM

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

Техники обучения LLM строятся на основе стандартных методов оптимизации:

• Адаптивные методы оптимизации — AdamW - базовый оптимизатор для всех LLM; понимание bias correction объясняет необходимость warmup
• Ландшафты потерь нейронных сетей — Warmup и gradient clipping помогают навигировать по острым частям ландшафта в начале обучения
• Оптимизация в распределённом обучении — ZeRO - часть большой экосистемы distributed training: tensor parallelism, pipeline parallelism, data parallelism

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

• Если у вас 8 GPU по 80GB и модель требует 320GB - какие техники вы комбинируете и в каком порядке включаете?
• Почему bf16 стал предпочтительнее fp16 для обучения LLM, несмотря на меньшую точность мантиссы?
• Как правильно масштабировать learning rate при увеличении batch size в 4× (linear scaling rule)?

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

• ml-09-gradient-descent