Математику SGD и Adam мы разбирали в треке «Математика для ML». Здесь практика: как настроить обучение нейросети, чтобы оно не разваливалось, влезало в память и шло быстро.
Выбор оптимизатора
- AdamW по умолчанию, особенно для трансформеров. Weight decay отделён от адаптивного шага.
- SGD с momentum часто для CNN с BatchNorm и хорошим расписанием.
Weight decay обычно не применяют к смещениям, γ нормализаций и некоторым эмбеддингам: их выносят в группу параметров с weight_decay = 0.
Learning rate
- С нуля: для AdamW около 1e-3 на маленьких моделях и 1e-4..3e-4 на трансформерах.
- Fine-tuning энкодера: около 1e-5..5e-5. Для LoRA обычно выше.
- Warmup первые сотни или тысячи шагов, затем косинусное или линейное затухание.
Расписание считают в шагах оптимизатора. Если scheduler пошаговый (OneCycle, cosine с warmup), scheduler.step() вызывают после каждого батча, а не раз в эпоху.
Gradient clipping
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
Норма считается по всем параметрам сразу, направление сохраняется.
Накопление градиентов
Батч не влезает в память: делаем несколько forward и backward с loss, делённым на число шагов накопления, затем один step и zero_grad. Эффективный батч = микро-батч × шаги накопления × число GPU.
Mixed precision
- fp16: 5 бит экспоненты, маленькие градиенты обнуляются. Нужен GradScaler: loss умножают на большой множитель, перед step градиенты делят обратно, шаги с inf пропускают.
- bf16: диапазон как у fp32, GradScaler обычно не нужен, точность мантиссы ниже.
scaler = torch.amp.GradScaler("cuda")
optimizer.zero_grad()
with torch.autocast("cuda", dtype=torch.float16):
loss = criterion(model(x), y)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
scheduler.step()
Перед клиппингом обязательно unscale_, иначе порог применится к масштабированным градиентам.
Модель не учится: чек-лист
- Переобучиться на одном маленьком батче. Не выходит: баг в коде.
- Проверить пары вход-метка глазами после всех преобразований.
- Логиты в CrossEntropyLoss, правильные размерности.
- zero_grad, backward, step на месте, параметры переданы оптимизатору, model.train().
- Начальный loss для K классов около ln K.
- Нормы градиентов по слоям, поиск lr по логарифмической сетке.
Большие модели
- Память: веса, градиенты и состояние Adam это около 16 байт на параметр в fp32, плюс активации.
- Gradient checkpointing экономит активации ценой лишнего прямого прохода.
- DDP реплицирует модель на каждую GPU, FSDP или ZeRO шардирует веса, градиенты и состояние оптимизатора.
- Всплески loss: логировать нормы градиентов, держать частые чекпоинты, откатываться и пропускать проблемные батчи, снижать lr.
Частые ошибки
- OneCycleLR со step раз в эпоху.
- Клиппинг без unscale_ при fp16.
- Накопление без деления loss на число шагов.
- Большой lr при fine-tuning.
Что спросят на собесе
- Какой оптимизатор и lr взять, зачем warmup и затухание.
- Как эмулировать большой батч, где ставить клиппинг.
- fp16 против bf16, зачем GradScaler.
- Ваши действия, если loss не падает или стал NaN.