The f16 storage and compute plan

Дата: 2026-10-06. Заменяет вычислительную часть bf16-plan.md: тот план заблокирован, потому что bf16-матмула на этом бэкенде не существует (§2.1 AGENTS.md — в pliron→LLVM нет типа bf16). f16 — то же самое на hardware, который работает.

Решение и почему

	bf16	f16
экспонента	как fp32 (широкий диапазон)	как fp32/8 (узкий: 5.96e-8…65504)
tensor cores на нашем бэкенде	НЕТ	ЕСТЬ, 91–97% cuBLAS (§2.1, измерено)
поддержка в burn	DType::BF16 есть	DType::F16 есть (narrow_dtype_cpu.rs)

Ценность f16: диапазон уже, чем у fp32, но мы не теряем главного — ускорения на tensor cores. bf16 проигрывает по обоим пунктам: диапазон бесполезен (мы всё равно считаем в f32), а железа нет.

Два уровня, и оба нужны

L1 — вычисление. Граф остаётся fp32, но каждый matmul внутри идёт через f16 (каст x и веса в f16 → f16 GEMM → f32 accumulate → f32 результат). Даёт скорость, не даёт памяти.

L2 — хранение. Активации и градиенты живут в f16. Даёт и скорость, и 2× экономию памяти на веса/активации/градиенты — то, что разблокирует 1B на карте с 24 ГБ.

L2 зависит от одного факта: проходит ли f16 через burn-autodiff. Это гейт wt/f16spike; пока он не зелёный, L2 не начинается.

Абстракция: Prec вместо bf16: bool

Нынешний флаг смешивает хранение и вычисление. Разделим:

pub enum Prec { F32, F16, BF16 }

impl Prec {
    /// dtype хранения активаций; None = остаёмся f32.
    pub fn act_dtype(self) -> Option<FloatDType>;
    /// Нужен ли f32-каст перед Linear. Истина ТОЛЬКО для bf16:
    /// у него нет tensor-core пути, а f16 — есть, поэтому при F16 каст не нужен.
    pub fn needs_f32_matmul(self) -> bool;
}

Это убирает правило «каст в f32 перед каждым Linear» как обход отсутствия bf16-матмула. При F32 правила нет вовсе, при F16 его нет, при BF16 оно остаётся.

Где fp32 обязателен — и почему

место	причина
RMSNorm sum, softmax, logsumexp	редукция; накопление в f16 теряет точность
final norm + lm_head	логиты переполняют f16 (max 65504) → NaN в softmax/CE
CE / log_softmax	то же + loss-скаляр читается на хосте
NaN firewall mask	должен видеть inf/NaN, а не underflow-в-ноль
AdamW v	квадрат градиента; диапазон v ~1e-8 → в f16 ноль
мастер-веса	fp16-округление накапливается → дрейф (сначала оставляем f32)
KDA recurrent state	накопление по времени; при длинной последовательности уходит в underflow

Где f16 безопасен

LinearLike matmul · attention QK/PV · FFN активации · residual stream · controller · RoPE · градиенты (с loss scaling).

Конфигурация

--prec f32 | f16 | bf16          # storage dtype активаций (по умолчанию f32)
--loss-scale N                   # для f16; по умолчанию из свипа, НЕ 1.0

--quant fp16 (существующий) — это формат TSCT-факторов, ось ортогональна. Не путать.

Гейты (без них работа не считается сделанной)

1. f16-форвард == fp32 на трейнер-формах, полоса 1e-2 rel (narrow_dtype_cpu.rs:30).
2. f16-бэквард сходится с f32-эталоном — гейт wt/f16spike.
3. 0 NaN на 200 шагах трейнер-рецепта.
4. BPB A/B: 3 сида, чистая CE, окно 81920 B. Победа бьёт разброс контроля (§1.2).
5. Память и скорость измерены с индексом шага (шаг 0 — артефакт автотюнера, 18–23×).

Риски, названные заранее

• f16-градиенты underflow → без loss scaling обучение не взлетит. Свип значения.
• Смешанный dtype NaN (§2.2) — возможно, f16-активация × f32-вес тоже NaN’ит. Тогда f16 обязателен и на весах, и это отдельная работа.
• KDA-рекурренция — под длительными последовательностями состояние может уйти в ноль.
• Маленькие GEMM на f16 могут быть медленнее fp32 (vectorization). Мерить, а не верить.

Что НЕ входит

• Мульти-GPU. Наш код single-device; 24 ГБ — это весь бюджет на модель.
• Квантование состояния оптимизатора ниже f16 (stochastic rounding) — отдельный заход.