Gradient Descent: как LLM учится на примерах

opensourceaillmgradient-descentoptimizationit
← Back to Blog

Введение: как модель учится?

Нейронная сеть не знает, как генерировать текст. Ей нужно научиться на примерах. Gradient Descent — основной механизм обучения.


Уровень 1: Случайные веса (до обучения)

Что делает модель без обучения?

# Инициализация весов
model = nn.LLM(hidden_size=4096, num_layers=32)

# Пробуем сгенерировать:
prompt = "Москва — это"
output = model.generate(prompt)
# "banana purple elephant quantum dancing unicorn..."

# Почему? Случайные веса = случайный output
# Каждый нейрон — случайная функция
# Сеть — случайное отображение input → output

Loss = полный хаос

Input: "Москва — это столица России"
Model output: "banana purple elephant..."

CrossEntropyLoss(output, target) = 9.21 (max possible for 100000 tokens)

Потому что модель предсказала "banana" с вероятностью 0.00001
когда правильный ответ — "столица"

Уровень 2: Что такое Gradient Descent?

Идея: итеративно улучшать веса

Аналогия: горный спуск

        /\
       /  \
      /    \  ← начальная позиция (случайные веса)
     /  ●   \
    / /| \   \
   / / | \   \
  / /  |  \   \
 / /   |   \   \
/______|____\___\ ← минимум (loss = 0)

● — модель "скатывается" к минимуму ошибки

Математика

Loss = f(weights)
Мы хотим: min(weights) Loss

Gradient = d(Loss) / d(weights)  ← направление наискорейшего роста

Step: weights = weights - learning_rate * gradient
       ↑                    ↑
       что обновляем       куда двигаться

Простой пример

# Предсказание: y = w * x + b
# Цель: y = 10*x + 5

w, b = 0.5, 0.1  # плохие начальные веса

for step in range(1000):
    # Forward pass
    pred = w * x + b
    loss = (pred - target) ** 2
    
    # Backward pass (градиенты)
    dw = 2 * (pred - target) * x
    db = 2 * (pred - target)
    
    # Update weights
    w = w - 0.01 * dw
    b = b - 0.01 * db

# Результат: w ≈ 10, b ≈ 5 ✓

Уровень 3: Forward Pass

Вычисляем loss

def forward_pass(model, batch):
    """Прогноз модели на батче данных"""
    for token_id in batch:
        # Каждый слой трансформера:
        # 1. Attention
        # 2. Add & Norm
        # 3. FFN
        # 4. Add & Norm
        
        output = model(token_id)
    
    # Cross-Entropy Loss:
    # Чем дальше prediction от target, тем выше loss
    loss = cross_entropy(output.logits, batch.targets)
    return loss

Cross-Entropy Loss

Target: "столица" (token_id = 42)

Model prediction:
  "это":      0.15
  "столица":  0.01  ← очень мало!
  "город":    0.10
  "Россия":   0.05
  "banana":   0.00001

Loss = -log(0.01) = 4.60  ← высокий loss!

Если бы prediction был правильным:
  "столица":  0.90
  Loss = -log(0.90) = 0.10  ← низкий loss ✓

Уровень 4: Backward Pass (Backpropagation)

Цепное правило (chain rule)

Loss → layer_32 → layer_31 → ... → layer_1 → input

d(Loss)/d(layer_1) = d(Loss)/d(layer_32) * ... * d(layer_2)/d(layer_1)

Каждый градиент = произведение градиентов всех последующих слоёв

Backpropagation шаг за шагом

# Forward pass
loss = forward(model, batch)

# Backward pass — вычисляем градиенты
loss.backward()

# Что происходит внутри:
# Для последнего слоя:
d_loss_d_weights32 = compute_gradient(loss, layer32_output)

# Для предпоследнего:
d_loss_d_weights31 = d_loss_d_weights32 * d_layer32_d_weights31

# ... и так до первого слоя
# Для каждого слоя:
#   d_loss_d_weight = d_loss_d_output * d_output_d_weight

Визуализация

Layer 32 (last):
  loss = 4.60
  d_loss/d_W32 = [0.02, -0.01, 0.03, ...]  ← градиенты для каждого веса

Layer 31:
  d_loss/d_W31 = [0.001, -0.0005, 0.002, ...]  ← градиенты меньше

Layer 1 (first):
  d_loss/d_W1 = [0.00001, -0.000005, 0.00002, ...]  ← ещё меньше

Vanishing gradients! (см. ниже)

Уровень 5: Update Weights

SGD (Stochastic Gradient Descent)

# Самый простой вариант
for param in model.parameters():
    param.data = param.data - learning_rate * param.grad

SGD с Momentum

# Как шарик, катящийся вниз:
# v = momentum * v_prev - learning_rate * gradient
# w = w + v

velocity = 0  # начальная скорость

for step in range(1000):
    gradient = compute_gradient(loss, model)
    
    # Обновляем скорость
    velocity = momentum * velocity - lr * gradient
    
    # Обновляем веса
    model.weights += velocity

# momentum = 0.9:
#   Шарик разгоняется в правильном направлении
#   Сглаживает шум градиентов
#   Ускоряет обучение

Adam Optimizer (стандарт для LLM)

# Adam = Momentum + RMSprop
# Для каждого веса храним:
#   1. momentum (среднее градиентов)
#   2. variance (среднее квадратов градиентов)

m_t = beta1 * m_{t-1} + (1 - beta1) * gradient
v_t = beta2 * v_{t-1} + (1 - beta2) * gradient^2

m_corrected = m_t / (1 - beta1^t)
v_corrected = v_t / (1 - beta2^t)

weights = weights - lr * m_corrected / (sqrt(v_corrected) + epsilon)

# beta1 = 0.9, beta2 = 0.999, epsilon = 1e-8

Почему Adam лучше SGD?

SGD:
  loss = [10.0, 8.5, 9.2, 7.8, 8.1, 6.9, ...]  ← шумит
  Скорость обучения: одна для всех

Adam:
  loss = [10.0, 7.5, 6.2, 5.3, 4.8, 4.2, ...]  ← плавно
  Скорость обучения: адаптивная для каждого веса

Adam адаптирует learning rate для каждого параметра:
  Частые градиенты → маленький step (шумный)
  Редкие градиенты → большой step (уверенный)

Уровень 6: Batch и Mini-batch

Почему не весь датасет сразу?

Полный датасет (Batch GD):
  1 трлн токенов = 1 трлн forward/backward pass
  Память: 1 трлн * 4 байта = 4 ТБ ← невозможно!
  Время: часы до одного шага ← слишком медленно

Mini-batch:
  32768 токенов на батч
  Память: 32768 * 4 байта = 128 КБ ✓
  Время: секунды до одного шага ✓

Размер батча

Batch size = 1:
  Очень шумный градиент
  Быстрые обновления
  Плохая параллелизация

Batch size = 32768:
  Сбалансированный градиент
  Хорошая параллелизация на GPU
  Стандарт для LLM

Batch size = 1048576:
  Очень стабильный градиент
  Требует огромной памяти
  Может ухудшить generalization

Gradient Accumulation

# Эмуляция большого батча на маленькой GPU
accumulation_steps = 32
actual_batch_size = 128 * 32 = 4096

for i, batch in enumerate(dataloader):
    loss = forward(model, batch) / accumulation_steps  # нормализуем
    loss.backward()  # градиенты накапливаются
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()  # обновляем веса
        optimizer.zero_grad()  # сбрасываем градиенты

Уровень 7: Learning Rate Schedule

Почему не постоянный LR?

Постоянный learning rate:
  В начале: слишком большой → loss explode
  В конце: слишком маленький → обучение останавливается

Нужно:
  В начале: большой LR → быстро учимся
  В конце: маленький LR → тонкая настройка

Warmup + Decay

# Стандартный schedule для LLM
total_steps = 100000
warmup_steps = 2000

if step < warmup_steps:
    # Linear warmup: от 0 до max_lr
    lr = max_lr * (step / warmup_steps)
else:
    # Cosine decay: от max_lr до 0
    lr = 0.5 * max_lr * (1 + cos(pi * (step - warmup_steps) / (total_steps - warmup_steps)))

# График:
# lr
#  |         ______
#  |        /      \
#  |       /        \____
#  |______/
#  0    2000          100000 step

Почему warmup нужен?

Без warmup:
  step 1: gradient = [100, -200, 150, ...]  ← случайный большой
  weights = weights - 3e-4 * [100, -200, ...]
  weights = [50, -100, 75, ...]  ← взрыв!

С warmup:
  step 1: lr = 3e-6 (в 100 раз меньше)
  weights = weights - 3e-6 * [100, -200, ...]
  weights = [0.5, -1.0, 0.75, ...]  ← аккуратно!

Уровень 8: Vanishing и Exploding Gradients

Vanishing Gradients

Проблема:
  gradient = d_loss/d_W1 = d_loss/d_W32 * d_W32/d_W31 * ... * d_W2/d_W1
  
  Если каждый множитель < 1:
    0.1^32 = 10^-32  ← градиент исчез!

Решения:
  1. Residual connections (skip connections)
     gradient может течь напрямую через слои
  
  2. Layer Normalization
     стабилизирует градиенты
  
  3. Good initialization (Xavier, Kaiming)
     начинает с правильным масштабом

Exploding Gradients

Проблема:
  gradient = 10 * 10 * 10 * ... = 10^32  ← взрыв!

Решение: Gradient Clipping

  if gradient_norm > max_norm:
      gradient = gradient * (max_norm / gradient_norm)

  # max_norm = 1.0 — стандарт для LLM

Gradient Norm

# Мониторинг градиентов
gradient_norm = torch.norm(gradient)

# Нормальный диапазон:
#   gradient_norm = 0.1 - 10.0

# Если > 10:
#   → gradient clipping сработает
#   → возможно, проблема в модели

# Если < 0.01:
#   → vanishing gradients
#   → проверь architecture

Уровень 9: Полный цикл обучения

One training step

for step, batch in enumerate(dataloader):
    # 1. Forward pass
    logits = model(batch.input_ids)
    loss = cross_entropy(logits, batch.labels)
    
    # 2. Backward pass
    loss.backward()
    
    # 3. Gradient clipping
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    
    # 4. Optimizer step
    optimizer.step()
    
    # 5. Learning rate schedule
    scheduler.step()
    
    # 6. Zero gradients
    optimizer.zero_grad()
    
    # 7. Logging
    if step % 100 == 0:
        print(f"Step {step}: loss = {loss.item():.4f}, lr = {scheduler.get_last_lr()[0]:.2e}")

Training loss curve

Loss
10 |●
 9 | ●
 8 |  ●
 7 |   ●
 6 |    ●
 5 |     ●●●
 4 |       ●●●●
 3 |          ●●●●●
 2 |              ●●●●●
 1 |                 ●●●●●
 0 +----------------------
    0   10k  20k  30k  40k  steps

Что происходит на каждом этапе

Epoch 1:
  loss = 5.0 → 3.0
  Модель учит: базовые паттерны языка
  
Epoch 5:
  loss = 2.0 → 1.5
  Модель учит: синтаксис, грамматику
  
Epoch 10:
  loss = 1.2 → 0.9
  Модель учит: факты, стили
  
Epoch 15:
  loss = 0.7 → 0.65
  Модель учит: тонкие нюансы, reasoning

Уровень 10: Практические аспекты

GPU memory и training

Память GPU (A100 80GB):
  Parameters:     40GB (Llama-7B)
  Gradients:      20GB
  Optimizer state: 20GB (Adam)
  Activations:    10GB (batch_size=32)
  Total:          90GB ← не влезает!

Решения:
  - Gradient accumulation
  - Gradient checkpointing
  - Mixed precision (FP16/BF16)
  - ZeRO sharding

Mixed Precision Training

# FP32 → FP16 / BF16
with torch.cuda.amp.autocast():
    logits = model(input_ids)
    loss = cross_entropy(logits, labels)

# Сокращение памяти:
#   FP32: 4 байта на число
#   FP16: 2 байта на число
#   Экономия: 2x!

# Но: потеря точности
# Решение: master weights в FP32 для update

Distributed Training

Data Parallel:
  GPU 0: batch часть 1 → loss 1
  GPU 1: batch часть 2 → loss 2
  GPU 2: batch часть 3 → loss 3
  GPU 3: batch часть 4 → loss 4
  
  Gradient sync → update all GPUs одинаково

Tensor Parallel:
  Layer 0: GPU 0 + GPU 1 (attention split)
  Layer 1: GPU 0 + GPU 1 (FFN split)
  ...

Итоги

  • Gradient Descent — итеративный метод оптимизации
  • Forward pass вычисляет loss, backward pass вычисляет градиенты
  • Adam — стандартный оптимизатор для LLM
  • Mini-batch — баланс между скоростью и стабильностью
  • Learning rate schedule (warmup + decay) — критичен для обучения
  • Gradient clipping — защита от exploding gradients
  • Mixed precision — 2x экономия памяти без потери качества