Расскажите про gradient accumulation

Gradient Accumulation в Deep Learning

Gradient Accumulation — техника, позволяющая тренировать модели с большим effective batch size на GPU с ограниченной памятью, путём накопления градиентов нескольких малых батчей перед обновлением весов.

Проблема: Ограничение памяти GPU

Оптимально:  batch_size = 256 (32 GB GPU память)
Реальность:  GPU память = 8 GB
             максимум: batch_size = 64
             Проблема: меньше batch_size → хуже обучение

# Проблема
model = BigTransformer()  # 1.2B параметров
batch_size = 256  # оптимально

try:
    output = model(batch)  # CUDA Out of Memory!
except:
    # GPU имеет только 8GB, нужно 32GB
    pass

Решение: Gradient Accumulation

Идея: вместо одного большого батча, используй несколько малых

Без Gradient Accumulation:
Батч 1 (размер 64) → forward → backward → обновить веса

С Gradient Accumulation (accumulation_steps=4):
Батч 1 (размер 64) → forward → backward → копить градиенты
Батч 2 (размер 64) → forward → backward → копить градиенты
Батч 3 (размер 64) → forward → backward → копить градиенты
Батч 4 (размер 64) → forward → backward → скопить + обновить веса

EFFECTIVE BATCH SIZE = 64 × 4 = 256

Реализация в PyTorch

import torch
from torch import nn, optim

model = BigTransformer()
optimizer = optim.Adam(model.parameters(), lr=1e-4)

accumulation_steps = 4
effective_batch_size = batch_size * accumulation_steps  # 256

for epoch in range(num_epochs):
    for step, (batch_x, batch_y) in enumerate(train_loader):
        # Forward pass
        output = model(batch_x)
        loss = criterion(output, batch_y)
        
        # Нормализуем loss по accumulation_steps
        loss_normalized = loss / accumulation_steps
        
        # Backward pass (накапливаем градиенты)
        loss_normalized.backward()
        
        # Обновляем веса каждые accumulation_steps
        if (step + 1) % accumulation_steps == 0:
            optimizer.step()      # Обновить веса
            optimizer.zero_grad() # Обнулить градиенты
            print(f"Updated weights at step {step}")

С использованием Hugging Face Transformers

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=32,  # физический batch size
    gradient_accumulation_steps=8,     # accumulate 8 times
    # Effective batch size = 32 × 8 = 256
    
    learning_rate=5e-5,
    warmup_steps=1000,
    weight_decay=0.01,
    fp16=True,  # mixed precision для экономии памяти
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

trainer.train()

Математически

# Без накопления градиентов
G = ∇L(batch)          # gradient for one batch
θ = θ - lr × G         # update weights immediately

# С накоплением (accumulation_steps=N)
G_accumulated = 0
for i in range(N):
    L_i = loss(model(batch_i), y_i)
    G_i = ∇(L_i / N)     # нормализуем loss!
    G_accumulated += G_i  # копим градиент

θ = θ - lr × G_accumulated  # обновляем один раз

Важно: теряем loss на N, чтобы получить корректный gradient!

loss_normalized = loss / accumulation_steps
loss_normalized.backward()  # правильный градиент

Пример из боевой практики

class TrainingLoop:
    def __init__(self, model, batch_size=32, accumulation_steps=8):
        self.model = model
        self.batch_size = batch_size
        self.accumulation_steps = accumulation_steps
        self.effective_batch_size = batch_size * accumulation_steps
        self.optimizer = optim.Adam(model.parameters())
    
    def train_epoch(self, dataloader):
        self.model.train()
        total_loss = 0
        
        for step, (inputs, labels) in enumerate(dataloader):
            # Forward pass
            outputs = self.model(inputs)
            loss = F.cross_entropy(outputs, labels)
            
            # Normalize for accumulation
            loss_scaled = loss / self.accumulation_steps
            
            # Backward pass
            loss_scaled.backward()
            
            total_loss += loss.item()
            
            # Update weights
            if (step + 1) % self.accumulation_steps == 0:
                # Clip gradients (optional)
                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(), max_norm=1.0
                )
                
                # Update
                self.optimizer.step()
                self.optimizer.zero_grad()
                
                effective_step = (step + 1) // self.accumulation_steps
                print(f"Effective step {effective_step}: "
                      f"loss={total_loss/effective_step:.4f}")
        
        return total_loss / len(dataloader)

Преимущества и недостатки

✓ Преимущества:

• Используешь большой effective batch size на маленькой GPU
• Лучше обучение (как на большом батче)
• Экономия памяти (обучаешь большие модели)
• Стабильнее训练 (благодаря большему batch size)

✗ Недостатки:

• Медленнее обучение (обновления реже)
• Сложнее отладка (нужно делить loss правильно)
• Может требовать регулировки learning rate

Сравнение

Без накопления (batch_size=256):
Шаг 1: обработка 256 образцов
Шаг 2: обработка 256 образцов
Шаг 3: обработка 256 образцов
Время на 1 обновление: 3 шага
Память: 32 GB

С накоплением (batch_size=64, accumulation_steps=4):
Шаг 1: обработка 64 образцов
Шаг 2: обработка 64 образцов
Шаг 3: обработка 64 образцов
Шаг 4: обработка 64 образцов → обновление
Время на 1 обновление: 4 шага (медленнее)
Память: 8 GB (эффективнее)
Эффективный batch size: 256 (как без накопления!)

Лучшие практики

# 1. Всегда нормализуй loss
loss = loss / accumulation_steps  # ✓ правильно

# 2. Обнуляй градиенты правильно
if (step + 1) % accumulation_steps == 0:
    optimizer.zero_grad()  # обнулить после обновления

# 3. Clip градиенты перед обновлением
if (step + 1) % accumulation_steps == 0:
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

# 4. Мониторь "эффективный" шаг обновления
effective_step = (step + 1) // accumulation_steps

# 5. С mixed precision (AMP)
with autocast():
    loss = model(batch)
loss_scaled = loss / accumulation_steps
scaler.scale(loss_scaled).backward()  # PyTorch autocast

Резюме

Gradient Accumulation:

• Накопление градиентов несколько батчей
• Эффективный batch size = batch_size × accumulation_steps
• Решает проблему ограничения памяти GPU
• Требует нормализации loss и правильного обнуления градиентов
• Медленнее, но более памяти-эффективно
• Essential для обучения больших моделей на слабых GPU