Локальні градієнти множаться
Прямий прохід
Прямий прохід ANDREA-120M проходить вхід через послідовність операцій:
x = embed(token_ids) # вбудовування токенів
for layer in 12_layers:
x = x + attn(LN(x)) # підшар attention
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 вагами прямий режим є нереальним.
Чому зворотний режим
Кожен Forward Op Отримує Зворотного Близнюка
Дисципліна Пара
microgpt_cuda.cu постачається з двома CUDA кернами для кожної операції: один, що обчислює forward вихід, один, що обчислює градієнти входів за градієнтами виходів. Пара є один-на-один:
| Ядро прямого проходу | Ядро зворотного проходу | Операція |
|---|---|---|
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. Двигун навчання зберігає їх у спеціальних буферах під час прямого проходу, а потім читає під час зворотного.
Вартість пам'яті: приблизно така ж форма, як вихід forward для кожного збереженого тензора. Для ANDREA-120M з batch=8, seq=1024, d_model=768, один збережений тензор — 8 × 1024 × 768 × 4 bytes = 25 MB. По 12 шарах та кількох збережених тензорах на шар, активації домінують VRAM під час тренування (~5-10 GB на карті 24 GB).
Трасування одного кроку backward
Де зберігаються градієнти в пам'яті
Один тензор градієнта на кожен тензор ваги
Кожен тензор ваги, що навчається, в 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. Більші ефективні батчі вимагають накопичення градієнтів: виконувати кілька прямих+зворотних проходів на малих батчах, сумуючи градієнти в той самий буфер, потім виконати один крок оптимізатора.
for microbatch in range(n_microbatches):
forward(microbatch)
backward() # ДОДАЄ до буферів градієнтів, не перезаписує
scale_grads(1.0 / n_microbatches) # усереднення по мікробатчах
optimizer_step()
zero_grads() # скидання для наступного кроку тренуванняКernelі backward використовують семантику +=, а не =. Кожен виклик додає вклад градієнтів до наявного буфера; буфер утримує поточну суму, доки 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 ГБ загалом. Комфортно нижче стелі 24 ГБ RTX 4090. 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 та експерименти з тензорними ядрами FP8 E4M3 в ANDREA вимагають точного контролю над тим, які тензори знаходяться в якій точності. Загальні фреймворки надають цей контроль через багатошарові API; власні CUDA-записи пишуть це безпосередньо.
Компроміс
Вартість власного CUDA: більше коду для написання, більше помилок для пошуку, відсутня екосистема спільноти. microgpt_cuda.cu від ANDREA — це ~6000 рядків написаного вручну CUDA, на налагодження якого пішли місяці. Кожна нова операція вимагає написання прямого ядра, зворотного ядра та тестів.
Що отримує ANDREA:
- Повна відтворюваність. Навчальний пайплайн — це один бінарний файл C плюс один Python-проксі. Немає зсувів версій між релізами PyTorch, немає невідповідностей версій CUDA з wheels фреймворку.
- Бітово-точні відновлення. SIGTERM запускає запис чекпоінту, що захоплює кожен тензор точно так, як його бачить GPU. Відновлення продовжує ту саму траєкторію втрат, на якій перебував запуск.
- Передбачувана пам'ять. ANDREA-120M навчалася 200K кроків без OOM. Пам'ять була врахована на старті двигуна.
- Прямий доступ до апаратного забезпечення. Розміри плиток tensor core, налаштування 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 GPU з ≥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]. Лічильник кроків, кількість ваг, потім ваги, за якими йдуть моменти Adam. Відновлює біт-екзактно. Проксі архимує .samples.json та .state.json окремо на polish; .loss.json ніколи не архивається (він накопичує повну історію тренування).