Практика: трансформер с нуля

Блок в одной строке

Всё, что было в этом блоке, собирается в два выражения, и оба уже разобраны по частям.

$$x \leftarrow x + \htmlData{k=attn}{\text{Attn}}\big(\htmlData{k=norm}{\text{LN}}(x)\big), \qquad x \leftarrow x + \htmlData{k=mlp}{\text{MLP}}\big(\htmlData{k=norm}{\text{LN}}(x)\big)$$

Единственное разделение труда, которое здесь есть: внимание работает между позициями, MLP — внутри позиции. Больше в блоке ничего нет.

Что реализовать

1. Десять операций через einsum

Начните с нотации, потому что дальше она сэкономит все ошибки с формами. Выпишите каждую операцию блока как строку индексов и посчитайте стоимость до написания кода.

операция	спецификация
проекция $Q$	bnD,Dhd->bhnd
оценки	bhqd,bhkd->bhqk
выход внимания	bhqk,bhkd->bhqd
склейка голов	bhnd,hdD->bnD
MLP вверх	bnD,DF->bnF
MLP вниз	bnF,FD->bnD
логиты	bnD,VD->bnV

Проверьте себя на первой строке: она содержит переформовку (h d) -> h d внутри контракции, и именно поэтому проекция и разбиение на головы — это одна операция, а не две.

проекция и головыоценкиMLP вверхлогиты

bnD, Dhd → bhnd

индекс	размер	роль
b	1	остаётся
n	256	остаётся
D	384	суммируется
h	6	остаётся
d	64	остаётся

форма результата

[1, 6, 256, 64]

элементов

98,304

умножений-сложений

37,748,736

операций на элемент

384

2. Блок целиком

Масштаб nanoGPT: $D = 384$, $6$ слоёв, $6$ голов, контекст $256$. На таком размере всё считается на одном CPU за минуты, а все эффекты блока видны.

Обязательные детали, каждая из которых разбиралась:

• pre-norm, а не post-norm (урок 090);
• инициализация He для MLP, выходные проекции ветвей поделены на $\sqrt{2L}$ (урок 100);
• причинная маска большим отрицательным числом, а не $-\infty$ (урок 080);
• деление оценок на $\sqrt{d}$ (урок 070);
• softmax со сдвигом на максимум (урок 060).

3. Пять измерений вместо доверия

Здесь главная часть работы. Каждое утверждение блока проверяется одной строкой на вашей же реализации.

Измерение 1: $\sqrt{d}$ действительно нужен. Отключите деление и напечатайте стандартное отклонение оценок и энтропию строк весов. Ожидание: разброс порядка $\sqrt{d} = 8$, энтропия падает почти до нуля ещё до обучения.

Измерение 2: инициализация даёт нужный масштаб. Напечатайте стандартное отклонение активаций после каждого слоя на случайном входе. Ожидание: с $1/\sqrt{2L}$ на выходных проекциях итоговая дисперсия residual-потока около $2$ при любом числе слоёв — проверьте на $L = 6$ и $L = 24$.

Измерение 3: RoPE относителен точно. Посчитайте оценку для пар $(0, 5)$, $(10, 15)$, $(100, 105)$. Ожидание: совпадение до последнего знака. Сделайте то же с прибавленной синусоидальной кодировкой — там расхождение будет порядка самого размаха оценок.

Измерение 4: кэш ничего не меняет в результате. Сгенерируйте двадцать токенов дважды — с кэшем и наивным пересчётом всей последовательности — и сравните полученные последовательности. Ожидание: побитовое совпадение. Если расходятся, ошибка в кэше, и это самый надёжный способ её найти.

Измерение 5: LoRA стартует ровно с исходной модели. Добавьте $BA$ к проекциям внимания при $B = 0$ и сравните логиты до и после добавления. Ожидание: побитовое равенство. Затем сделайте один шаг оптимизации и убедитесь, что логиты изменились.

4. Генерация с кэшем

Реализуйте два режима и сравните не только результат, но и время:

режим	что считает на шаге
наивный	весь прямой проход для всех $n$ токенов
с кэшем	один токен, ключи и значения читаются

Ожидание из урока 140: разница растёт линейно по длине. Заодно посчитайте размер кэша своей модели: $2 n L d_{\text{kv}}$ элементов — на этом масштабе он мал, и полезно увидеть, во что он превращается при $D = 4096$ и $L = 32$.

Отдельная деталь, которую легко пропустить: на шаге генерации логиты нужны только для последней позиции. Проверьте по таблице einsum выше, сколько работы экономит этот срез — на маленькой модели со словарём $50257$ это самая дорогая операция.

5. Обучение на маленьком корпусе

Возьмите любой текст в несколько мегабайт и обучите до осмысленных фрагментов. Что стоит наблюдать:

• лосс на первом шаге обязан быть около $\ln V$ — это энтропия равномерного распределения по словарю (блок 4, урок 010). Отклонение означает ошибку в инициализации выходного слоя;
• перплексия $e^{\text{loss}}$ — «во сколько раз модель неуверенна» (блок 4, урок 020);
• warmup и cosine-расписание (блок 5, урок 090) — сравните с постоянным шагом;
• градиентный клиппинг: посмотрите, как часто он срабатывает в начале и в конце.

Куда смотреть, если не сходится

симптом	что проверять первым
лосс на первом шаге $\ne \ln V$	инициализацию выходного слоя
лосс nan через несколько шагов	сдвиг в softmax, маску $-\infty$, слишком большой шаг
лосс стоит на $\ln V$	маску (не видит ли модель будущее наоборот), знак в residual
качество внезапно ниже на коротких строках	padding-маску
генерация с кэшем расходится с наивной	порядок применения RoPE: позиции в кэше и у нового токена
активации растут по слоям	множитель $1/\sqrt{2L}$ на выходных проекциях

Пятая строка — самая частая ошибка в реализациях с кэшем: при повороте нового ключа нужно взять его абсолютную позицию, а не нулевую. Проверяется измерением 4 и никак иначе.

Чем это стоит закончить

Прогоните готовую модель через утверждения блока и убедитесь, что все пять измерений дают ожидаемое. Именно этот набор — а не работающая генерация — показывает, что блок понят: генерация получается и при трёх ошибках, которые взаимно компенсируются.