ЦОДы на вечной мерзлоте

Токенизацию в пайплайне обычно не трогают — болото Шрэка. Да и по моему опыту на собесах ее тоже так лишь интеллектуально ласкают. Но как говорится "идите туда, где страшно, :)".

Узкое место у хф и тиктокенов: regex (еще одно страшное место). Претокенизация через регулярки — около 50 МБ/с. Ручной автомат с таблицей переходов даёт сразу 10–20×. Дальше помогает байтовая арифметика без ветвлений и второй курсор по буферу: однопоток упирается не в CPU, а в цепочку зависимостей на токен.

Частые слова тоже решают многое. Уникальных претокенов мало, кэш закрывает почти все обращения, а 90% слов — один токен, так что ответ можно хранить прямо в кэше. Красивые схемы без ветвлений часто проигрывают: упираются в память. Многопоточность тоже не даёт честного ×16 — из-за сборки и memcpy по Амдалу выходит скорее ×3–4.


На триллионах токенов и оффлайн обработке даст еще пару часов (или суток) на эксперименты :^)

Завтра умный пост

Почему нельзя взять одну большую LLM и использовать её везде?

На практике каждая продуктовая задача предъявляет свои требования: например, к качеству, latency и стоимости инференса. Поэтому в RWB развивается не одна модель, а целое семейство BerryLM.

Узнали у Матвея Сапрыкина, AI Assistants Models Tech Lead, как оно устроено ⚙️

@rwb_delaet_ml

В пятницу пост

Ладно. Единственно важная математика сегодня — это расчет бенза на поездку до дачи и обратно

bf16 для DFlash: где matmul живёт, а softmax — нет

DFlash — block-diffusion draft: параллельно предсказывает блоки по 16 токенов, внимание через FlexAttention, loss — cross-entropy по огромному словарю.

Веса в bf16 — нормально. Проблема в другом: на каждом шаге DFlash делает два тяжёлых редукшена (softmax-подобных), и bf16 там систематически теряет точность.

> Экспонента с мантисой вышли покурить.
> Вернулись — а softmax уже другой.

1. FlexAttention + разреженная block-causal маска

512 draft-блоков. Каждый видит только:
— контекст до своего anchor (~половина последовательности)
— 16 своих draft-токенов

Остальное в маске. Видимых KV-позиций ~30–35% — экстремальная разреженность.

Softmax в Flash/FlexAttention считается онлайн: бежит max m, сумма l, перескалировка через exp(m_old − m_new). В bf16 (7 бит мантиссы ≈ 3 десятичных цифры) эти перескалировки округляются. Чем агрессивнее маска — тем меньше «коррекций» и тем сильнее дрейф.

> Маска выкинула 67% KV.
> Ядро сделало меньше перескалировок.
> bf16 округлил — и attention уже не тот.

PyTorch upstream ([#163588](https://github.com/pytorch/pytorch/issues/163588)): bf16 + маска → до нескольких % ошибки; fp32 → <0.001%.

2. Cross-entropy по словарю ~164k

CE = −z_y + log Σ exp(z_j) — log-sum-exp по 163 840 logits на каждый draft-токен.

Суммировать 164k экспонент в bf16 — прямой путь к ULP-ошибкам и редким NaN. PyTorch для log_softmax/CE давно аккумулирует в float — не из прихоти.

> 164 тысячи exp на один токен.
> Мантисса на семь бит.
> Кто-то из них обязан потеряться — вопрос только, когда это станет NaN.

Что делаем

Не «переходим на fp32-тренинг». Точечно:
→ FlexAttention: q, k, v в fp32, output cast обратно в bf16
→ CE: logits.float() перед cross_entropy

Matmul и веса — bf16. Редукции — fp32. Стандартный приём для LLM; RMSNorm в DFlash уже так работает.

> fp32 не на всё — fp32 только там, где exp, max и sum решают судьбу градиента.

Мораль

Когда тренируешь draft с block-sparse attention и большим словарём — смотри не на distributed, а на редукции. Экспонента тут не декорация: softmax и CE буквально живут на exp, max и sum.

DFlash собирает оба болевых места в одном forward. fp32 на этих двух участках — не перебор, а минимально необходимая страховка.


Кто найдет ошибки — тот молодец. Ему за это ничего не будет