Локальные градиенты умножаются
Прямой проход
Прямой проход ANDREA-120M проходит входные данные через последовательность операций:
x = embed(token_ids) # вложения токенов
for layer in 12_layers:
x = x + attn(LN(x)) # подслой внимания
x = x + mlp(LN(x)) # подслой MLP
logits = LN(x) @ embed.T # связанная проекция выхода
loss = cross_entropy(logits, targets)Каждая операция читает входные тензоры и производит выходные тензоры. Прямой проход завершается одним скаляром: кросс-энтропийной потерей для этой партии.
Обратный проход
Обучение обновляет веса в направлении, уменьшающем потерю. Чтобы получить направления обновления, движок нуждается в:
dL/dW для каждого обучаемого W в моделиПравило цепочки даёт это. Для цепочки loss = f(g(h(x))):
dL/dx = (dL/df) * (df/dg) * (dg/dh) * (dh/dx)Каждый фактор — это локальный градиент: как выход одной операции изменяется, когда её вход изменяется на малую величину. Умножение локальных градиентов назад через граф распространяет сигнал потерь на каждый вес.
Дифференцирование в обратном режиме
Backprop вычисляет градиенты в обратном порядке: начиная с dL/dlogits = 1, затем проходя назад через кросс-энтропию, затем проекцию выхода, затем нормализацию слоя, затем двенадцать блоков трансформера, затем эмбеддинги. На каждом шаге умножайте входящий градиент на локальный якобиан.
Обратный режим эффективен, когда выход — это один скаляр (потери) и много входов (веса). Один обратный проход производит градиенты для каждого веса в модели. Прямой режим потребовал бы один проход на вес; для ANDREA-120M с ~120M весами прямой режим неосуществим.
Почему обратный режим
Каждой прямой операции соответствует обратная пара
Дисциплина парности
microgpt_cuda.cu поставляется с двумя CUDA-ядрами для каждой операции: одно вычисляет прямой выход, другое вычисляет градиенты входа при заданных градиентах выхода. Парность строго один-к-одному:
| Ядро прямого прохода | Ядро обратного прохода | Операция |
|---|---|---|
k_embed_fwd | k_embed_bwd | Поиск эмбеддингов токенов |
k_layernorm_fwd | k_layernorm_bwd | Нормализация слоя |
k_attn_qkv_fwd | k_attn_qkv_bwd | Проекции Q, K, V |
k_attn_fwd | k_attn_bwd | Масштабированное скалярное произведение внимания |
k_attn_out_fwd | k_attn_out_bwd | Проекция выхода W_O |
k_mlp_fwd | k_mlp_bwd | MLP (с GELU) |
k_residual_add | k_residual_add_bwd | Остаточное соединение |
k_loss_fwd | k_loss_bwd | Функция потерь кросс-энтропии |
Восемь пар операций охватывают весь трансформер. Плюс несколько вспомогательных ядер: k_grad_norm_partial, k_grad_norm_final, k_grad_scale для обрезки градиентов (см. активность 75).
Задача обратного ядра
Задан градиент, поступающий из последующих слоев (grad_output), обратное ядро вычисляет:
1. grad_input: градиент относительно входного тензора операции. Он передается дальше назад.
2. grad_weight: градиент относительно обучаемых параметров в операции. Он попадает в состояние оптимизатора.
Оба вычисляются в одном запуске ядра. Потоки CUDA сотрудничают над плитками тензора градиента параллельно.
Сохраненные тензоры
Обратное вычисление часто требует значений из прямого прохода. Например, k_layernorm_bwd нужны среднее и дисперсия, вычисленные во время прямого прохода; k_mlp_bwd нужна предактивация GELU. Движок обучения сохраняет их в специальных буферах во время прямого прохода, а затем читает во время обратного.
Затраты памяти: примерно такая же форма, как у выхода прямого прохода для каждого сохранённого тензора. Для ANDREA-120M с batch=8, seq=1024, d_model=768 один сохранённый тензор занимает 8 × 1024 × 768 × 4 bytes = 25 MB. По 12 слоям и нескольким сохранённым тензорам на слой активации доминируют в VRAM во время обучения (~5-10 GB на карте с 24 GB).
Отслеживание одного шага обратного прохода
Где хранятся градиенты в памяти
Один тензор градиента на каждый тензор весов
Каждый обучаемый тензор весов в ANDREA-120M имеет соответствующий тензор градиента одинаковой формы. Для каждого блока:
W_Q [768, 768] ↔ grad_W_Q [768, 768]
W_K [768, 768] ↔ grad_W_K [768, 768]
W_V [768, 768] ↔ grad_W_V [768, 768]
W_O [768, 768] ↔ grad_W_O [768, 768]
W_1 [768, 3072] ↔ grad_W_1 [768, 3072]
W_2 [3072, 768] ↔ grad_W_2 [3072, 768]
LN1.gamma [768] ↔ grad_LN1.gamma [768]
LN1.beta [768] ↔ grad_LN1.beta [768]
LN2.gamma [768] ↔ grad_LN2.gamma [768]
LN2.beta [768] ↔ grad_LN2.beta [768]Плюс вложения токенов, вложения позиций и финальная нормировка слоя. Общий объём памяти буфера градиентов соответствует объёму памяти весов: ~120M чисел с плавающей точкой, ~480 МБ при FP32, ~240 МБ при FP16.
Накопление по микробатчам
batch_size = 8 в ANDREA помещается в VRAM при FP16. Для больших эффективных батчей требуется накопление градиентов: выполняем несколько проходов forward+backward на малых батчах, суммируя градиенты в один буфер, затем делаем один шаг оптимизатора.
for microbatch in range(n_microbatches):
forward(microbatch)
backward() # ДОБАВЛЯЕТ к буферам градиентов, не перезаписывает
scale_grads(1.0 / n_microbatches) # усреднение по микробатчам
optimizer_step()
zero_grads() # сброс для следующего шага обученияBackward kernels используют семантику +=, а не =. Каждый вызов добавляет вклады градиентов к существующему буферу; буфер содержит текущую сумму до тех пор, пока zero_grads() не очистит его.
Состояние оптимизатора
AdamW (активность 73) содержит два дополнительных буфера на каждый вес: первый момент m и второй момент v. Общий объём памяти во время обучения:
веса: 1× количество весов
градиенты: 1× количество весов
Adam m: 1× количество весов
Adam v: 1× количество весов
сохранённые активации: ~2-4× в зависимости от слоёв и батча
──────────────────────────────────────────
всего: ~6-8× количество весовANDREA-120M на FP16: ~240 МБ × 4 буфера (веса, градиенты, m, v) + ~5-10 ГБ активации = ~10-12 ГБ всего. Комфортно ниже потолка RTX 4090 в 24 ГБ. ANDREA-12M обучалась в 1.4 ГБ; 10× масштабирование параметров приносит ~10× память.
Размер буферов градиентов
Полный контроль над памятью и точностью
Что стоит общим фреймворкам
PyTorch и JAX делают автодифференцирование удобным: пишете код на Python, получаете градиенты автоматически. Цена: общий слой диспетчеризации между вашим кодом и CUDA. Каждая операция проходит через накладные расходы интерпретатора Python, служебные записи фреймворка и динамический выбор ядер. При обучении небольшой языковой модели на одной GPU эти накладные расходы имеют значение.
Конкретные затраты, которых избегает ANDREA:
1. Задержка интерпретатора Python. Каждая операция PyTorch пересекает границу Python/C++. При ~100 запусках ядер на шаг обучения и ~9 шагах/мин это ~900 пересечений границы в минуту. Диспетчеризация на уровне C устраняет это.
2. Непредсказуемость аллокатора фреймворка. Кэширующий аллокатор PyTorch обеспечивает хорошую пропускную способность в среднем, но непредсказуемый пик памяти. Движок обучения ANDREA предварительно выделяет каждый буфер при запуске; нет перераспределения во время обучения, нет фрагментации, нет неожиданных OOM на шаге 100K.
3. Выбор универсальных ядер. PyTorch выбирает ядра во время выполнения с помощью эвристик. ANDREA выбирает ядра во время компиляции, настроенные под размеры тайлов тензорных ядер RTX 4090.
4. Прокладка смешанной точности. Путь FP16 cuBLAS в ANDREA-120M и эксперименты ANDREA с тензорными ядрами FP8 E4M3 требуют точного контроля над тем, какие тензоры находятся в какой точности. Общие фреймворки предоставляют этот контроль через многослойные API; кастомные CUDA пишут это напрямую.
Компромисс
Затраты на кастомный CUDA: больше кода для написания, больше ошибок для поиска, отсутствие экосистемы сообщества. Файл microgpt_cuda.cu в ANDREA — это ~6000 строк написанного вручную CUDA, на отладку которого ушли месяцы. Каждая новая операция требует написания прямого ядра, обратного ядра и тестов.
Что получает ANDREA:
- Полная воспроизводимость. Тренировочный пайплайн — это один бинарный файл C плюс один Python-прокси. Нет дрейфа версий между релизами PyTorch, нет несоответствий версий CUDA с колесами фреймворка.
- Бит-точные возобновления. SIGTERM запускает запись чекпоинта, которая захватывает каждый тензор точно так, как его видит GPU. Возобновление продолжает ту же траекторию лосса, на которой находилась тренировка.
- Предсказуемая память. ANDREA-120M обучалась 200K шагов без OOM. Память была учтена при запуске движка.
- Прямой доступ к аппаратному обеспечению. Размеры тайлов тензорных ядер, настройки FP8 E4M3, асинхронные копии памяти: всё напрямую доступно в CUDA, непрозрачно в общих фреймворках.
Воспроизводимость как миссия
Раздел 9 whitepaper ANDREA перечисляет полный стек воспроизводимости:
Движок обучения: microgpt/microgpt_cuda.cu
Прокси обучения: microgpt/training_proxy.py
Конфигурации экспериментов: experiments/ANDREA-*-TRAIN.json
Конвейер данных: scripts/pull-hermes3.py, scripts/prep-megachat.py
Панель управления: scripts/live-loss-dashboard.html
Спецификация Bandit: docs/FIREHOSE-BANDIT.md
Документация модели: docs/ANDREA.mdТребования к оборудованию: одна видеокарта NVIDIA с ≥8 ГБ VRAM (RTX 3060 или лучше). Любой может воспроизвести ANDREA-12M из этих артефактов. Собственный путь CUDA — часть причины: нет заморозки версий фреймворков, нет сюрпризов с зависимостями через пять лет.
Сигналы и контрольные точки
Цикл обучения CUDA реагирует на два сигнала POSIX:
- SIGTERM: записать немедленную контрольную точку, затем выйти. Используется при чистой остановке обучения.
- SIGUSR1: записать немедленный чекпоинт, продолжить обучение. Использовался во время polish pivot в v3 для захвата состояния без прерывания запуска.
Формат чекпоинта: [int32 step][int32 n_params][n_params × float32 weights][n_params × float32 m][n_params × float32 v]. Счетчик шагов, количество весов, затем веса, за которыми следуют моменты Адама. Возобновляет бит-в-бит точно. Прокси архивирует .samples.json & .state.json отдельно на polish; .loss.json никогда не архивируется (он накапливает полную историю обучения).