← Росток Своя LLM с нуля Словарик Код EN

Глава 8 из 17 40 мин

Трансформер

Проследим путь от эмбеддинга до следующего токена. Соберём блок Ростка, разберём нормировку и позиции, затем посчитаем все параметры модели.

В этой главе

  • собрать блок трансформера из внимания, сети SwiGLU, остаточных связей и RMSNorm
  • понять, как RoPE кодирует позиции поворотом и учитывает расстояние между токенами
  • посчитать параметры Ростка и любой похожей модели
  • заглянуть внутрь модели логит-линзой и увидеть, как предсказание созревает от слоя к слою

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

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

Росток на одной схеме

Вот вся модель. Текст входит сверху, предсказание выходит снизу. Толстая линия слева — остаточный поток, пунктирная рамка — блок, который повторяется 8 раз, жёлтый пунктир справа — одна и та же матрица, работающая и на входе, и на выходе. Нажмите на любую деталь, и рядом со схемой появится её код — дословно из файла, по которому обучали Ростка, с номерами строк.

Код в панели берётся прямо из scripts/micro-llm/model.py. Переключатель сверху меняет только числа в конфигурации: у Ростка, у gpt-4 этой главы и у gpt-1 из прошлой код один и тот же.

Обратите внимание на переключатель: gpt-1 из прошлой главы, gpt-4 из этой и полный Росток — это один и тот же код. Меняются лишь несколько чисел: ширина потока, число блоков и голов, размер сети. Вся разница между маленькой и большой моделью — в этих числах, данных и времени обучения. Дальше пройдём по схеме сверху вниз.

Остаточный поток

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

$$x \leftarrow x + \operatorname{Attn}\big(\operatorname{Norm}(x)\big),$$ $$x \leftarrow x + \operatorname{FFN}\big(\operatorname{Norm}(x)\big).$$

Представления, которые проходят через последовательные блоки с такими прибавлениями, называют остаточным потоком. Его можно представить общей рабочей записью: внимание и сеть читают текущее состояние и добавляют свои поправки. Но это не архив отдельных страниц. Все поправки складываются в одном векторе фиксированной длины, поэтому могут усиливать, ослаблять и даже отменять прежние компоненты. Гарантии сохранения всей исходной информации нет.

Почему не проще — $x \leftarrow f(x)$, как в сетях из главы 4? Причин две.

Градиенту нужна дорога. При обратном проходе (глава 4) производные слоёв перемножаются по цепочке. Если каждая из них хоть немного меньше единицы, через восемь блоков сигнал заметно ослабнет, а через сотню почти исчезнет. С прибавлением всё иначе:

$$\frac{\partial}{\partial x}\big(x + f(x)\big) = 1 + f'(x).$$

Для вектора вместо единицы в формуле стоит единичная матрица. Эта прямая ветвь даёт градиенту дополнительный путь, по которому ему не нужно проходить через производную нелинейной части блока. Остаточные связи облегчают обучение глубоких сетей, хотя сами по себе не гарантируют, что градиенты нигде не затухнут и не вырастут. Они были ключевой частью ResNet (He и др., 2015) и затем вошли в трансформеры.

Небольшие поправки в начале. В инициализации выходные матрицы внимания и сети (proj и w2) получают веса в $\sqrt{2 \cdot 8}=4$ раза меньшего масштаба. Это уменьшает начальный вклад каждого блока в остаточный поток. Модель всё ещё учитывает контекст через внимание, но исходные эмбеддинги не сразу перекрываются суммой больших случайных поправок. По мере обучения блоки учатся менять представления полезным для предсказания образом.

RMSNorm: держать числа в узде

У постоянного дописывания есть обратная сторона: размер вектора в потоке растёт. Мы измерили на gpt-4 из этой главы: типичное число в потоке (среднеквадратичное, в среднем по отложенному тексту) сразу после эмбеддинга — около 0,18, после первого блока — уже около 10, а после четвёртого — около 35, в двести раз больше. А слою удобнее всего, когда вход всегда примерно одного масштаба, где бы в модели этот слой ни стоял.

Поэтому перед каждым слоем вектор нормируют. Росток использует RMSNorm — среднеквадратичную нормировку:

$$\operatorname{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2},$$ $$\hat x_i = \frac{x_i}{\sqrt{\operatorname{RMS}(x)^2 + \varepsilon}}\, g_i .$$

Каждую координату делим на корень из среднего квадрата координат с добавкой $\varepsilon$. До обучаемого масштабирования среднеквадратичная величина результата близка к единице, если исходный вектор не слишком мал. Затем каждую координату умножаем на свой коэффициент $g_i$, который учится вместе с моделью. Так слой получает управляемый масштаб входа, сохраняя возможность усиливать разные координаты по-разному. Всего таких коэффициентов $d$.

Нажмите «× 10»: вход вырос вдесятеро, а нормированный выход не изменился. Потом переключитесь на LayerNorm и нажмите «+ 3 ко всем»: вот в чём разница между двумя нормировками.

Обратите внимание, где стоит нормировка: на входе в слой, а не в самом потоке. Нормируется копия, которую слой читает, а поток остаётся нетронутым — иначе мы бы снова «переписывали тетрадь». Такую схему называют pre-norm. В самом первом трансформере 2017 года нормировка стояла после сложения (post-norm), и глубокие модели с ней обучались капризно; Сюн и др. (2020) показали, почему pre-norm стабильнее. Ещё одна нормировка стоит в самом конце, перед выходным слоем.

Каждый токен думает сам: SwiGLU

Внимание — единственное место, где токены обмениваются сведениями. Но собрать сведения мало, их надо переработать. Для этого в каждом блоке после внимания стоит маленькая нейросеть из двух слоёв, почти такая же, как в главе 4, — сеть прямого распространения (feed-forward network). Она обрабатывает каждый токен отдельно, одними и теми же весами на всех позициях. Здесь живёт большая часть параметров: $3 \times 384 \times 1\,024 = 1\,179\,648$ на блок против $589\,824$ у внимания — две трети блока.

В первом трансформере сеть была такой: $\operatorname{ReLU}(x W_1)\,W_2$ — расширить вектор в 4 раза, обнулить отрицательное, сжать обратно. Росток, как и Llama, использует вариант SwiGLU с тремя матрицами:

$$\operatorname{FFN}(x) = \big(\operatorname{silu}(x W_1) \odot x W_3\big)\,W_2,$$ $$\operatorname{silu}(a) = \frac{a}{1 + e^{-a}}.$$

Здесь $\odot$ — поэлементное умножение. Вектор расширяется до 1 024 чисел дважды, двумя разными матрицами. Первая копия проходит через SiLU — гладкую родственницу ReLU, — вторая остаётся как есть, и они перемножаются: одна копия служит воротами для другой. Нейрон может сказать «пропускаю, только если выполнены оба условия» или даже поменять знак сигнала — одиночной ReLU такое не под силу.

Один из 1 024 нейронов сети. Справа — его выход при всех сочетаниях $a$ и $g$: синее — положительный вклад, рыжее — отрицательный. Двигайте точку по плоскости или ползунками.

А откуда странное число 1 024? В классической сети скрытый слой в 4 раза шире потока, и матриц две: $2 \times 4d^2 = 8d^2$ параметров. У SwiGLU матриц три, поэтому ширину берут равной $\tfrac{2}{3}$ от обычной, чтобы параметров вышло столько же: $\tfrac23 \cdot 4 \cdot 384$ — ровно $1\,024$. Вот сеть Ростка целиком:

class FeedForward(nn.Module): """SwiGLU: silu(x·W1) ⊙ (x·W3), projected back with W2.""" def __init__(self, cfg): super().__init__() self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x))

Где я? Позиции поворотом

В конце прошлой главы выяснилось, что внимание не различает порядок слов: для него «the dog bit the boy» и «the boy bit the dog» — один и тот же набор векторов. Порядок надо как-то сообщить.

Самое прямолинейное решение — завести ещё одну таблицу эмбеддингов, на этот раз для позиций, и прибавлять её строку к вектору токена. Так сделано в GPT-2 (2019). Но тогда модель знает только абсолютные номера и отдельно учит, что «слово двумя позициями раньше» на десятом месте и на трёхсотом — одно и то же отношение. А для языка важнее как раз относительное: насколько далеко стоит нужное слово.

RoPE (rotary position embedding; Су и др., 2021) ничего не прибавляет — он поворачивает. Запрос одной головы — это 64 числа; разобьём их на 32 пары, и каждая пара станет точкой на плоскости. Токен на позиции $m$ поворачивает свою $i$-ю пару на угол $m\,\theta_i$, где

$$\theta_i = 10000^{-2i/64}, \qquad i = 0, 1, \dots, 31.$$

Ключи поворачиваются так же, каждый по своей позиции $n$. Значения не поворачивают.

Для фиксированных запроса и ключа поворот сохраняет длины, а их скалярное произведение зависит от разности углов поворота. Поэтому явная зависимость оценки от позиций входит через относительное смещение $m-n$. Это не значит, что внимание зависит только от расстояния: сами векторы несут содержание текста, а в последующих слоях — и результаты предыдущих вычислений. Сравним одну и ту же пару векторов на разных позициях.

У разных пар координат разный темп поворота. Первая пара делает оборот примерно за 6,3 позиции; последняя — примерно за 47 тысяч. Один поворот периодичен, а сочетание разных частот позволяет различать больше взаимных расположений. Это похоже на несколько часовых стрелок с разной скоростью. Но аналогия не обещает точного отсчёта на произвольной длине: способность модели использовать расстояния зависит также от обучения и длины доступного контекста.

Восемь из 32 пар одной головы с настоящими частотами Ростка. Нажмите «Сдвинуть оба на +10»: все стрелки повернутся, каждая со своей скоростью, но угол между синей и рыжей в каждой паре — и итоговая оценка — останутся прежними.

В коде это две функции. Первая один раз готовит таблицы косинусов и синусов для всех позиций и частот. Вторая поворачивает пары $(x_i, x_{i+32})$: первая половина вектора — это «иксы» всех пар, вторая — «игреки».

def rope_tables(head_dim, context, base): """cos/sin for every position and every frequency, shape (context, head_dim/2).""" inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) angles = torch.outer(torch.arange(context).float(), inv_freq) return angles.cos(), angles.sin() def apply_rope(x, cos, sin): """Rotate pairs (x[i], x[i + d/2]) by a position-dependent angle. x: (batch, heads, time, head_dim); cos/sin: (time, head_dim/2). """ d = x.size(-1) // 2 x1, x2 = x[..., :d], x[..., d:] return torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)

Проверим главное свойство на numpy. Одни и те же запрос и ключ на позициях 5 и 2 и на позициях 105 и 102 — расстояние одно, значит, и оценка должна совпасть:

import numpy as np d = 64 inv_freq = 1.0 / (10000 ** (np.arange(0, d, 2) / d)) # the 32 thetas def rope(x, pos): """Rotate pairs (x[i], x[i + 32]) by pos * theta_i.""" a = pos * inv_freq x1, x2 = x[:d // 2], x[d // 2:] return np.concatenate([x1 * np.cos(a) - x2 * np.sin(a), x1 * np.sin(a) + x2 * np.cos(a)]) rng = np.random.default_rng(1) q, k = rng.standard_normal(d), rng.standard_normal(d) for m, n in [(5, 2), (105, 102), (400, 397), (5, 4)]: s = rope(q, m) @ rope(k, n) print(f"m={m:3d} n={n:3d} distance {m - n}: q.k = {s:+.4f}") print("length kept:", np.isclose(np.linalg.norm(rope(q, 77)), np.linalg.norm(q)))

Запрос стоит на позиции 100, ключ — на позиции 90. Потом ту же пару векторов переносим на позиции 20 и 10. Что станет с оценкой внимания?

Повороты на $m\theta$ и $n\theta$ входят в скалярное произведение только через разность $(m-n)\theta$. Это и проверила ячейка выше: пары позиций (5, 2), (105, 102) и (400, 397) дали одно и то же число.

В нашем варианте RoPE частоты заданы заранее и не обучаются. Формулу поворота можно вычислить для любого номера позиции, но реализация Ростка заранее строит таблицы только для настроенного контекста. Она не получает неограниченную память автоматически. Даже если расширить таблицы, качество за пределами обученной длины нужно проверять отдельно. Именно различие между вычислимым поворотом и освоенным диапазоном важно помнить, обсуждая длинный контекст.

Блок целиком

Теперь у нас есть все части блока. Вот он. Весь его смысл умещается в двух строчках метода forward, а комментарии к ним пересказывают эту главу лучше любой схемы:

class Block(nn.Module): def __init__(self, cfg): super().__init__() self.norm1 = RMSNorm(cfg.d_model) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.d_model) self.ffn = FeedForward(cfg) def forward(self, x, cos, sin): x = x + self.attn(self.norm1(x), cos, sin) # tokens exchange information x = x + self.ffn(self.norm2(x)) # each token thinks on its own return x

Блок трансформера — это два шага: сначала токены обмениваются сведениями (внимание), потом каждый обдумывает их сам (сеть SwiGLU). Оба шага обновляют вектор прибавлением поправки.

Восемь этажей и выход

Блоки ставятся друг на друга: у Ростка их восемь, одинаковых по устройству, но каждый со своими весами. Зачем больше одного? Потому что каждый следующий блок читает то, что записали предыдущие, и может на этом строить. В прошлой главе мы видели индукционную голову: ей нужно, чтобы кто-то раньше записал в каждый токен, какое слово стояло перед ним. Одному слою такое не под силу, а двум — уже да. Глубина позволяет строить цепочки таких шагов.

Проверим на деле. Ниже три модели читают один и тот же текст: однослойная gpt-1 из прошлой главы, четырёхблочная gpt-4 из этой и полный Росток.

Для каждого токена — вероятность, которую ему дала модель, прочитав всё, что было до него. Строка под каждой моделью сравнивает одни и те же повторяющиеся кусочки текста в первый и во второй раз.

В примере с повторяющимся «Zorbin Plax» четырёхблочная модель заметно повышает вероятность уже встречавшихся продолжений: около 98% для «or» и 93% для «bin». У однослойной модели вероятность «or» остаётся ниже процента. Это показывает разницу в поведении на данном примере; сами числа не доказывают, что результат обеспечила именно пара голов, рассмотренная ранее. У полного Ростка вероятность «or» при округлении доходит до 100%, а «bin» — до 98%. Округлённые 100% здесь не означают математически точную единицу.

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

self.head.weight = self.embed.weight   # weight tying

Это называется связанными весами (weight tying; Пресс и Вольф, 2017). Смысл простой: на входе строка таблицы говорит «вот что значит этот токен», на выходе та же строка спрашивает «похоже ли то, что модель хочет сказать, на этот токен?». Одна таблица учится с двух сторон, а модель экономит $8\,192 \times 384 = 3\,145\,728$ параметров — без этого трюка Ростку понадобилось бы 20,45 млн чисел вместо 17,31 млн.

Дальше всё знакомо: softmax превращает логиты в вероятности, а при обучении кросс-энтропия с настоящим следующим токеном даёт ошибку (глава 2). Маска loss_mask понадобится, когда мы будем учить Ростка разговаривать (глава 12).

Считаем параметры

Теперь можно разобрать Ростка по косточкам и пересчитать всё до последнего числа. Обозначим ширину потока $d = 384$, число блоков $L = 8$, ширину сети $d_{ff} = 1\,024$ и размер словаря $V = 8\,192$:

  • эмбеддинг (он же выходной слой): $V d = 3\,145\,728$;
  • внимание в одном блоке: $W_Q, W_K, W_V, W_O$ — это $4d^2 = 589\,824$, плюс $2 \times 64$ усилений QK-norm;
  • сеть SwiGLU в одном блоке: $3 d\, d_{ff} = 1\,179\,648$;
  • две нормировки в блоке: $2d = 768$, и одна финальная: $384$.
$$3\,145\,728 + 8 \times (589\,824 + 128 + 1\,179\,648 + 768) + 384 = 17\,309\,056.$$

Это и есть 17,31 миллиона параметров Ростка. Посчитайте свою модель:

Размер файла — такой, какой скачивает браузер: веса квантованы в int8, по одному числу масштаба на строку (глава 11). Скорость — грубая оценка по темпу, измеренному для Ростка; настоящая зависит от устройства.

Пощёлкайте заготовки. У крошечной gpt-1 почти всё — эмбеддинг: 80% параметров уходит на словарь, и на «мышление» остаётся совсем немного. У Ростка на словарь приходится уже 18%, а у модели размером с GPT-2 small с её огромным словарём доля снова растёт. Видно и то, что цена блока растёт как квадрат ширины потока.

Ширину потока $d$ удвоили, $d_{ff}$ тоже удвоили (чтобы сохранить пропорцию), число блоков оставили. Во сколько раз вырастет число параметров внутри блоков?

И внимание ($4d^2$), и сеть ($3d \cdot d_{ff}$, где $d_{ff}$ растёт вместе с $d$) — квадратичны по ширине. Удвоение $d$ даёт $2^2 = 4$; только крошечные усиления нормировок растут вдвое. Эмбеддинг $Vd$ за пределами блоков тоже вырастет лишь вдвое.

Логит-линза: промежуточные прогнозы

Обычно выходной слой читает остаточный поток только после последнего блока. Для исследования можно применить ту же финальную нормировку и выходную матрицу к промежуточным состояниям. Получим распределения по словарю, с которыми удобно сравнивать окончательный прогноз. Этот приём называют логит-линзой (logit lens). Он не раскрывает внутренний «план» модели: мы намеренно читаем промежуточный вектор способом, для которого он отдельно не обучался.

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

На что посмотреть:

  • Строка «эмбеддинг» часто повторяет входной токен. Здесь одна таблица служит и для чтения токенов, и для получения выходных оценок: собственная строка токена нередко даёт большое скалярное произведение с его вектором. Это не математическая гарантия — другая строка может дать ещё большее значение. На этом этапе блоки трансформера ещё не обработали контекст.
  • Где впервые появляется правильный ответ? Рамка отмечает совпадение с настоящим следующим словом. Для лёгких продолжений ответ часто созревает в средних слоях, для трудных — только в последних.
  • Как меняется распределение? Цвет показывает вероятность самого вероятного токена при чтении через линзу. На трёх заготовленных текстах её среднее значение было 77% после седьмого блока и 59% после восьмого. Это наблюдение о выбранных примерах, а не доказательство, что последний блок специально «убирает самоуверенность». Промежуточные состояния не обучались напрямую выдавать вероятности через эту линзу. Нажмите на клетку и сравните пять лучших вариантов.
Весь model.py

Вот файл полностью, от первой строки до последней. Теперь вам понятно в нём всё, кроме функции generate: выбору слов посвящена глава 11. Одно предупреждение: числа в Config — лишь значения по умолчанию, а настоящие размеры Ростка (поток из 384 чисел, 6 голов, сеть на 1 024) передаёт ему train.py.

"""Sprout: a tiny GPT, all of it in one file. A decoder-only transformer in the style of today's open models: - pre-norm blocks with RMSNorm, - rotary position embeddings (RoPE), - multi-head causal self-attention with QK-norm, - a SwiGLU feed-forward layer, - input and output embeddings shared (weight tying), - no biases anywhere. """ import math from dataclasses import dataclass, asdict import torch import torch.nn as nn import torch.nn.functional as F @dataclass class Config: vocab_size: int = 8192 # how many different tokens the model knows context: int = 512 # how many tokens it can look back at n_layer: int = 8 # transformer blocks stacked on top of each other n_head: int = 8 # attention heads per block d_model: int = 512 # width of the residual stream d_ff: int = 1408 # hidden width of the feed-forward layer rope_base: float = 10000.0 def to_dict(self): return asdict(self) class RMSNorm(nn.Module): """Rescale a vector to unit root-mean-square, then apply a learned gain.""" def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): return F.rms_norm(x.float(), (x.size(-1),), self.weight.float(), self.eps).type_as(x) def rope_tables(head_dim, context, base): """cos/sin for every position and every frequency, shape (context, head_dim/2).""" inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) angles = torch.outer(torch.arange(context).float(), inv_freq) return angles.cos(), angles.sin() def apply_rope(x, cos, sin): """Rotate pairs (x[i], x[i + d/2]) by a position-dependent angle. x: (batch, heads, time, head_dim); cos/sin: (time, head_dim/2). """ d = x.size(-1) // 2 x1, x2 = x[..., :d], x[..., d:] return torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1) class Attention(nn.Module): def __init__(self, cfg): super().__init__() self.n_head = cfg.n_head self.head_dim = cfg.d_model // cfg.n_head self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) def forward(self, x, cos, sin): B, T, C = x.shape q, k, v = self.qkv(x).split(C, dim=-1) # (B, T, C) -> (B, heads, T, head_dim) q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) q, k = self.q_norm(q), self.k_norm(k) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) # softmax(q·k / sqrt(d)) · v, every token sees only itself and the past y = F.scaled_dot_product_attention(q, k, v, is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.proj(y) class FeedForward(nn.Module): """SwiGLU: silu(x·W1) ⊙ (x·W3), projected back with W2.""" def __init__(self, cfg): super().__init__() self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) class Block(nn.Module): def __init__(self, cfg): super().__init__() self.norm1 = RMSNorm(cfg.d_model) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.d_model) self.ffn = FeedForward(cfg) def forward(self, x, cos, sin): x = x + self.attn(self.norm1(x), cos, sin) # tokens exchange information x = x + self.ffn(self.norm2(x)) # each token thinks on its own return x class GPT(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model) self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer)) self.norm = RMSNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) self.head.weight = self.embed.weight # weight tying cos, sin = rope_tables(cfg.d_model // cfg.n_head, cfg.context, cfg.rope_base) self.register_buffer('cos', cos, persistent=False) self.register_buffer('sin', sin, persistent=False) self.apply(self._init) # residual projections start small so every block begins close to "do nothing" for name, p in self.named_parameters(): if name.endswith('proj.weight') or name.endswith('w2.weight'): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer)) @staticmethod def _init(m): if isinstance(m, (nn.Linear, nn.Embedding)): nn.init.normal_(m.weight, mean=0.0, std=0.02) def num_params(self): return sum(p.numel() for p in self.parameters()) def forward(self, idx, targets=None, loss_mask=None): T = idx.size(1) x = self.embed(idx) cos, sin = self.cos[:T], self.sin[:T] for block in self.blocks: x = block(x, cos, sin) logits = self.head(self.norm(x)) if targets is None: return logits, None loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), targets.reshape(-1), reduction='none') if loss_mask is None: return logits, loss.mean() mask = loss_mask.reshape(-1).float() return logits, (loss * mask).sum() / mask.sum().clamp(min=1) @torch.no_grad() def generate(self, idx, max_new, temperature=0.8, top_k=None, top_p=0.95, stop=None): for _ in range(max_new): logits, _ = self(idx[:, -self.cfg.context:]) logits = logits[:, -1, :].float() / max(temperature, 1e-5) if top_k: kth = torch.topk(logits, top_k).values[:, -1, None] logits[logits < kth] = -float('inf') probs = F.softmax(logits, dim=-1) if top_p and top_p < 1.0: sorted_p, order = probs.sort(descending=True) drop = sorted_p.cumsum(-1) - sorted_p > top_p sorted_p[drop] = 0 probs = torch.zeros_like(probs).scatter(-1, order, sorted_p) probs /= probs.sum(-1, keepdim=True) nxt = torch.multinomial(probs, 1) idx = torch.cat((idx, nxt), dim=1) if stop is not None and nxt.item() in stop: break return idx

Росток сейчас

Это уже Росток в миниатюре: тот же model.py, только поток шириной 192 числа вместо 384 и четыре блока вместо восьми. На отложенном тексте он ошибается в среднем на 2,26 ната на токен — против 2,87 у однослойной gpt-1 из прошлой главы. Разница кажется небольшой, но вспомните перплексию из главы 2: $e^{2{,}87} \approx 18$, а $e^{2{,}26} \approx 9{,}6$. Раньше модель в среднем колебалась между восемнадцатью вариантами следующего токена, теперь — между девятью с половиной.

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

Главы

  1. 0 Знакомство
  2. 1 Считаем буквы
  3. 2 Мера удивления
  4. 3 Градиентный спуск
  5. 4 Обратное распространение
  6. 5 Эмбеддинги
  7. 6 Токены
  8. 7 Внимание
  9. 8 Трансформер
    1. Росток на одной схеме
    2. Остаточный поток
    3. RMSNorm: держать числа в узде
    4. Каждый токен думает сам: SwiGLU
    5. Где я? Позиции поворотом
    6. Блок целиком
    7. Восемь этажей и выход
    8. Считаем параметры
    9. Логит-линза: промежуточные прогнозы
    10. Росток сейчас
  10. 9 Корпус
  11. 10 Обучение
  12. 11 Выборка
  13. 12 Разговор
  14. 13 LoRA
  15. 14 Больше данных
  16. 15 Модель побольше
  17. 16 Думать и Python
  18. 17 Что дальше