Решения

Оптимизаторы для LLM: AdamW, Lion, AdaFactor

Задача страницы. Дать понятный выбор и рецепты: какой оптимизатор использовать в разных режимах обучения/дообучения, как снижать память оптимизатора и какие планировщики скорости обучения подключать.

Выбор оптимизатора: TL;DR

  • AdamW — стандарт де‑факто для трансформеров: устойчив, с корректным decoupled weight decay (лучше, чем «L2‑регуляризация внутри Adam»). Если доступна CUDA, включайте fused‑реализацию в PyTorch 2.x.
  • Lion — «знаковый» моментум‑оптимизатор: хранит только 1‑й момент (без второго), поэтому требует меньше памяти, чем Adam/AdamW; показал сравнимые результаты на ряде задач. Хороший кандидат для дообучения крупных моделей при дефиците VRAM.
  • AdaFactor — экономит память за счёт факторизации второго момента (память ≪ Adam/AdamW), исторически применялся в больших языковых моделях (например, T5), но требует аккуратной настройки и может быть менее стабильным.

Ускорители и «лайфхаки» для памяти:

  • 8‑битные оптимизаторы (bitsandbytes): квантуют состояния оптимизатора → до –75% памяти и значительный прирост скорости; часто работают без ретюнинга гиперпараметров. Есть и paged‑варианты (переносят часть состояний на CPU при нехватке GPU).
  • CPU‑offload (DeepSpeed ZeRO‑Offload + DeepSpeedCPUAdam) — перенос вычислений/состояний оптимизатора на CPU в больших тренировках.
  • Планировщики LR (linear/cosine + warmup) — базовая практика для LLM‑обучения.

Память и скорость: что закладывать в бюджет

  • AdamW: хранит m и v (1‑й и 2‑й моменты) → память под состояния ≈ 2× размер модели (в 32‑бит). Fused/foreach‑реализации снижают накладные расходы.
  • Lion: хранит только m → память состояний ≈ 1× размер модели.
  • AdaFactor: для матриц хранит факторизованные моменты (по строке/столбцу) → сублинейная память по сравнению с Adam.
  • 8‑бит оптимизаторы: квантуют состояния → минус ~75% памяти; paged ещё и «переливает» избыток на CPU.

Если VRAM на грани: начните с AdamW‑8bit или PagedAdamW‑8bit; при крайнем дефиците рассмотрите AdaFactor или CPUAdam (ZeRO‑Offload).

Практика: параметр‑группы и weight decay

Обычно не применяют decay к bias и LayerNorm‑параметрам:

				
					decay, no_decay = [], []
for n, p in model.named_parameters():
    (no_decay if any(k in n for k in ["bias", "norm", "ln"]) else decay).append(p)

param_groups = [
    {"params": decay,    "weight_decay": 0.01},
    {"params": no_decay, "weight_decay": 0.0},
]

Это сохраняет пользу weight_decay для весов, не штрафуя смещения/нормализации.

Планировщики скорости обучения (LR schedules)

Базовый набор (через Transformers): linear, cosine, cosine_restart, polynomial, constant(with_warmup), inverse_sqrt, и др. Тёплый старт (warmup) — почти всегда полезен.

				
					from transformers import get_cosine_schedule_with_warmup

num_training_steps = ...
num_warmup_steps = int(0.03 * num_training_steps)  # пример: 3% warmup
scheduler = get_cosine_schedule_with_warmup(
    optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps
)

Градиентный аккьюмулятор: как «нарастить» батч без OOM

Формула: effective_batch = micro_batch × grad_accum × world_size.

				
					grad_accum = 64
optimizer.zero_grad(set_to_none=True)

for step, batch in enumerate(loader):
    loss = model(**batch).loss / grad_accum
    loss.backward()
    if (step + 1) % grad_accum == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad(set_to_none=True)

Повышайте grad_accum, удерживая micro_batch на пороге VRAM.

Рецепты: AdamW (PyTorch), 8‑бит/CPU‑offload, AdaFactor, Lion

4.1 AdamW (PyTorch 2.x, fused)

				
					import torch
optimizer = torch.optim.AdamW(
    param_groups, lr=2e-4, betas=(0.9, 0.95), eps=1e-8, fused=True
)

fused=True включает слитую реализацию (если доступно для вашего девайса/типа).

4.2 AdamW‑8bit / PagedAdamW‑8bit (bitsandbytes) 

				
					import bitsandbytes as bnb

# 8‑битная версия
opt = bnb.optim.AdamW8bit(param_groups, lr=2e-4, betas=(0.9, 0.95))

# paged‑вариант: экономия VRAM + offload состояний на CPU при нехватке памяти
opt = bnb.optim.PagedAdamW8bit(param_groups, lr=2e-4, betas=(0.9, 0.95))

8‑битные оптимизаторы — drop‑in замена обычных; для NLP полезно использовать bnb.nn.StableEmbedding для устойчивости.
 PagedAdamW8bit дополнительно разгружает память (см. обзор torchtune).

4.3 CPU‑offload (DeepSpeed ZeRO‑Offload + CPUAdam)

ds_config.json (фрагмент):

				
					{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": { "device": "cpu", "pin_memory": true }
  },
  "optimizer": { "type": "CPUAdam", "params": { "lr": 0.0002, "betas": [0.9, 0.95], "eps": 1e-8, "weight_decay": 0.01 } }
}

CPUAdam — высокопроизводительная CPU‑реализация Adam(W), рекомендована для ZeRO‑Offload.

4.4 AdaFactor (Transformers)

				
					from transformers import Adafactor, AdafactorSchedule

optimizer = Adafactor(
    param_groups,
    scale_parameter=True, relative_step=True, warmup_init=True,
    lr=None, weight_decay=0.0
)
scheduler = AdafactorSchedule(optimizer)

AdaFactor экономит память за счёт факторизации; удобен через готовые классы Transformers.

4.5 Lion (для экономии памяти и быстрого дообучения)

				
					# Установите реализацию Lion для PyTorch (например, lion-pytorch или другой проверенный пакет)
from lion_pytorch import Lion
optimizer = Lion(param_groups, lr=2e-4, weight_decay=0.01)  # хранит только 1-й момент

Lion использует «sign‑momentum» и требует меньше памяти (нет второго момента). Тщательнее подбирайте lr и weight_decay.

Планировщики скорости обучения (LR schedules)

Базовый набор (через Transformers): linear, cosine, cosine_restart, polynomial, constant(with_warmup), inverse_sqrt, и др. Тёплый старт (warmup) — почти всегда полезен.

				
					from transformers import get_cosine_schedule_with_warmup

num_training_steps = ...
num_warmup_steps = int(0.03 * num_training_steps)  # пример: 3% warmup
scheduler = get_cosine_schedule_with_warmup(
    optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps
)
Градиентный аккьюмулятор: как «нарастить» батч без OOM

Формула: effective_batch = micro_batch × grad_accum × world_size.

				
					grad_accum = 64
optimizer.zero_grad(set_to_none=True)

for step, batch in enumerate(loader):
    loss = model(**batch).loss / grad_accum
    loss.backward()
    if (step + 1) % grad_accum == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad(set_to_none=True)

Повышайте grad_accum, удерживая micro_batch на пороге VRAM.

Что выбирать под разные задачи

  • Pretrain / крупные дообучения на нескольких GPU: AdamW (fused), при дефиците памяти — AdamW‑8bit или ZeRO‑Offload + CPUAdam.
  • PEFT/LoRA/QLoRA на одной/двух GPU: AdamW‑8bit или PagedAdamW‑8bit; для QLoRA paged‑оптимизаторы особенно полезны.
  • Экономия VRAM без 8‑бит: Lion (меньше памяти, чем AdamW). Для сверхжёстких ограничений — AdaFactor.

Диагностика и стабильность

  • Дивергенция лосса / NaN: уменьшите lr, увеличьте eps, включите grad clipping; для 8‑бит — проверьте StableEmbedding.
  • Слишком медленно / узкое место CPU: на PyTorch включите fused=True; в DeepSpeed — используйте FusedAdam (GPU) либо корректно настройте CPUAdam и ZeRO‑Offload.
  • OOM по состояниям оптимизатора: переход на AdamW‑8bit/PagedAdamW‑8bit или AdaFactor.

Куда дальше

Масштабирование/шардирование: /solutions/llm-training/fsdp-deepspeed/, /solutions/multi-gpu/

Точность и память: /solutions/llm-training/mixed-precision/, /solutions/llm-training/memory-opt/

Данные/I‑O/чекпоинты: /solutions/llm-training/datasets/, /solutions/llm-training/distributed-io/, /solutions/llm-training/checkpointing/

Стоимость и режимы аренды