Обучение

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

Прежде чем разбирать цикл, посмотрим, что он делает. По ходу обучения мы то и дело останавливались и просили Ростка продолжить четыре затравки. Вот что он писал.

Росток учится на глазах

Замедленная съёмка проходит через узнаваемые стадии. На нулевом шаге за затравкой идёт мешанина случайных токенов: её случайные предпочтения ещё не отражают закономерности текста. Через пару десятков шагов верх берут самые частые токены: точки, запятые, « the», « was», « and». Модель обнаружила, какие токены встречаются часто, — первое, что замечает любой, кто учит язык. Около 50-го шага появляются первые шаблоны — «Once upon a time, there was a girl named Lily», кавычки на своих местах, но грамматика ещё разваливается. Примерно к сотому шагу большинство предложений уже грамматичны, а диалоги записаны в правильном формате Имя: реплика, но история скачет с одного на другое. Через несколько сотен шагов истории держатся одной темы целый абзац. Дальше идёт медленная часть: меньше противоречий, всё длиннее связный текст, всё точнее выбор слов.

Следите и за числами над кривой. «Вариантов на токен» — наглядная подпись для перплексии $e^{\text{loss}}$ из главы 2. Это не число кандидатов, оставшихся в словаре: перплексия переводит среднюю ошибку в масштаб равновероятного угадывания. В начале она даже немного выше 8 192, потому что случайные предпочтения могут быть хуже равномерного прогноза. К концу она уменьшается до нескольких единиц. Это заметное улучшение предсказания текста, хотя одна эта метрика не измеряет ни правдивость, ни связность целой истории.

Один шаг обучения

Вот сердце train.py — часть, которая выполняется 10 070 раз. Всё остальное в файле (аргументы, журнал, контрольные точки) — леса вокруг этих тринадцати строчек:

T = context_at(step, steps, args.context, args.seq_warmup) x, y = train.get(args.batch_tokens // T, T) with ctx: _, loss = train_model(x, y) loss.backward() norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) f = lr_factor(step, steps, args.warmup, args.cooldown) for opt in opts: for g in opt.param_groups: g['lr'] = g['base_lr'] * f opt.step() opt.zero_grad(set_to_none=True) seen += x.numel()

По строчкам:

1. Длина окна. context_at сообщает длину окна на этом шаге: 128, 256 или 512 токенов. Зачем она меняется, расскажем ниже.
2. Батч. train.get вырезает из train.bin 32768 // T случайных окон: 256 окон по 128 токенов или 64 окна по 512. К каждому окну прилагается его копия, сдвинутая на один токен, — ответы y. В батче всегда 32 768 токенов, так что каждый шаг задаёт модели сразу 32 768 вопросов «что дальше?».
3. Прямой проход. Модель читает x и возвращает ошибку — среднее удивление по всем позициям всех окон. Строчка with ctx запускает её в bf16 — о нём чуть ниже.
4. Обратный проход. loss.backward() — это обратное распространение из главы 4, применённое к 17 миллионам параметров: после него у каждого параметра есть .grad, наклон ошибки по нему.
5. Обрезка. Если градиент подозрительно велик, его укорачивают (см. ниже).
6. Шаг. Скорость обучения на этот шаг берётся из расписания, каждый оптимизатор двигает свои параметры, и градиенты обнуляются для следующего круга.

Обрезка градиента

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

$$g \leftarrow g \cdot \min\left(1, \frac{1}{\lVert g \rVert}\right), \qquad \lVert g \rVert = \sqrt{\textstyle\sum_i g_i^2}.$$

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

bf16 и torch.compile: как ускорить вычисления

bf16 (bfloat16) — 16-битный формат с восемью битами порядка, как у float32, и семью явно хранимыми битами дробной части. Диапазон нормальных чисел близок к float32, но точность ниже: примерно две-три десятичные значащие цифры. В смешанной точности часть операций, прежде всего матричные умножения, выполняется в bf16, а другие остаются в float32. Правила torch.autocast зависят от операции и устройства. Это позволяет сократить объём передаваемых данных и использовать быстрое оборудование, сохраняя более точные вычисления там, где они нужны. Веса и состояние оптимизатора в нашем запуске хранятся во float32.

torch.compile захватывает вычисления модели в граф и старается выполнять их эффективнее: например, объединяет несколько операций в одно ядро для графического процессора. Меньше промежуточных записей в память — меньше лишней работы. Компиляция сама требует времени, а новые формы входа могут вызвать повторную компиляцию. В наших замерах на Mac выигрыш по сравнению с обычным режимом eager составил примерно два раза; на другом устройстве и другой модели результат нужно измерять заново.

Как сделать шаг

После обратного прохода мы знаем градиент $g$ — направление, в котором ошибка растёт быстрее всего. Самый простой шаг — обычный градиентный спуск из главы 3:

$$\theta \leftarrow \theta - \eta\, g.$$

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

Импульс превращает параметры в тяжёлый шар. Вместо текущего градиента мы шагаем по затухающей сумме прошлых: $m \leftarrow \beta m + g$, $\theta \leftarrow \theta - \eta m$ при $\beta = 0{,}9$. Зигзаги поперёк оврага гасят друг друга, а устойчивый толчок вдоль него накапливается.

Adam идёт дальше и даёт каждому параметру собственный размер шага. Кроме среднего градиента $m$ он следит за средним квадратом градиента $v$ и делит одно на другое:

$$m \leftarrow \beta_1 m + (1-\beta_1)\, g, \qquad v \leftarrow \beta_2 v + (1-\beta_2)\, g^2, \qquad \theta \leftarrow \theta - \eta\, \frac{\hat m}{\sqrt{\hat v} + \epsilon}.$$

Если градиент долго сохраняет примерно одну величину и знак, отношение $\hat m / \sqrt{\hat v}$ приближается к $+1$ или $-1$, если не учитывать $\epsilon$. Когда знак часто меняется, средний градиент может частично погаситься, и шаг уменьшается. Это полезная интуиция, но не точная формула размера каждого шага: два средних имеют разную память и зависят от истории. Шляпки обозначают поправку на их начальные нули. AdamW отдельно применяет затухание весов — небольшое уменьшение их величины, не смешанное с адаптивно масштабируемым градиентом.

Устроим гонку четырёх оптимизаторов на трёх ландшафтах. Ландшафты двумерные, чтобы их было видно, но трудности настоящие: узкий наклонный овраг, изогнутая долина-«банан» и плато, где градиент почти нулевой.

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

Muon: каждому направлению — одинаковый шаг

Adam смотрит на параметры как на длинный список независимых чисел. Но большая часть параметров Ростка сложена в матрицы, и у обновления матрицы есть своя структура. Здесь пригодится сингулярное разложение (SVD): любую матрицу можно записать как $G = U \Sigma V^\top$ — поворот, растяжение вдоль нескольких осей в $\sigma_1 \ge \sigma_2 \ge \dots$ раз и ещё один поворот. В градиенте весовой матрицы обычно несколько направлений имеют огромные $\sigma$, а остальные — крошечные. Обычный шаг тогда двигает веса в основном вдоль этих немногих главных направлений, а вдоль остальных — редких, но полезных — почти не двигает.

Muon (MomentUm Orthogonalized by Newton–Schulz) обрабатывает сглаженное импульсом обновление как матрицу. Его ориентир — убрать различие масштабов между сингулярными направлениями: заменить ненулевые сингулярные числа единицами. Для матрицы полного ранга это даёт ближайшую ортогональную матрицу, а в прямоугольном случае — матрицу с ортонормированными строками или столбцами:

$$G = U \Sigma V^\top \quad\longrightarrow\quad O = U V^\top,$$

У такой матрицы полного ранга все сингулярные числа равны единице. Реальная итерация ниже лишь приближает это выравнивание; строго нулевые направления она не создаёт заново.

Считать SVD на каждом шаге было бы долго. Muon пользуется трюком — полиномиальной итерацией, которой нужны только умножения матриц. Начинаем с $X_0 = G / \lVert G \rVert$ (тогда все сингулярные числа не больше 1) и повторяем

$$X \leftarrow aX + b\,(XX^\top)X + c\,(XX^\top)^2 X.$$

Вот оптимизатор, которым мы обучали Ростка, — весь muon.py, кроме шапки с импортами:

def orthogonalize(G, steps=5): """G: (..., rows, cols) -> the nearest semi-orthogonal matrices, batched.""" a, b, c = 3.4445, -4.7750, 2.0315 # tuned for fast convergence X = G.bfloat16() tall = X.size(-2) > X.size(-1) if tall: X = X.mT X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) for _ in range(steps): A = X @ X.mT X = a * X + (b * A + c * A @ A) @ X if tall: X = X.mT return X.to(G.dtype) class Muon(torch.optim.Optimizer): def __init__(self, params, lr=0.02, momentum=0.95, weight_decay=0.0, split=None): # split: {param: n} treats a fused weight (like q, k, v) as n separate matrices self.split = split or {} super().__init__(params, dict(lr=lr, momentum=momentum, weight_decay=weight_decay)) @torch.no_grad() def step(self): for group in self.param_groups: # matrices of the same shape are orthogonalised together, in one batch by_shape = defaultdict(list) for p in group['params']: if p.grad is None: continue state = self.state[p] if 'momentum' not in state: state['momentum'] = torch.zeros_like(p) buf = state['momentum'] buf.lerp_(p.grad, 1 - group['momentum']) g = p.grad.lerp(buf, group['momentum']) # Nesterov look-ahead n = self.split.get(p, 1) by_shape[(p.size(0) // n, p.size(1))].append((p, g.view(n, -1, p.size(1)))) for (rows, cols), items in by_shape.items(): updates = orthogonalize(torch.cat([g for _, g in items])) scale = max(1.0, rows / cols) ** 0.5 i = 0 for p, g in items: u = updates[i:i + g.size(0)].reshape_as(p) i += g.size(0) if group['weight_decay']: p.mul_(1 - group['lr'] * group['weight_decay']) p.add_(u, alpha=-group['lr'] * scale)

Несколько деталей, важных на практике. Импульс использует «заглядывание вперёд» по Нестерову: обновление смешивает текущий градиент с буфером импульса. Матрицы одинаковой формы ортогонализуются вместе, одним пакетом, поэтому Muon обходится так дёшево. Объединённая матрица, которая считает запросы, ключи и значения за один раз, сначала делится на три части — по смыслу это три разные матрицы. В Muon идут только двумерные веса внутри блоков; таблица эмбеддингов и коэффициенты RMSNorm остаются у AdamW.

Итерацию легко проверить в numpy: возьмём «градиент», в котором одно направление больше чем в сто раз сильнее другого, и ортогонализуем его.

import numpy as np def orthogonalize(G, steps=5): a, b, c = 3.4445, -4.7750, 2.0315 # the same coefficients as muon.py X = G / (np.linalg.norm(G) + 1e-7) for _ in range(steps): A = X @ X.T X = a * X + (b * A + c * A @ A) @ X return X rng = np.random.default_rng(0) # a "gradient" of a 4x6 weight matrix where a few directions dominate G = rng.normal(size=(4, 6)) @ np.diag([10, 3, 1, 0.3, 0.1, 0.03]) U, S, Vt = np.linalg.svd(G, full_matrices=False) print('singular values of G: ', S.round(3)) print('after orthogonalize(G): ', np.linalg.svd(orthogonalize(G), compute_uv=False).round(3)) print('exact U @ Vt: ', np.linalg.svd(U @ Vt, compute_uv=False).round(3))

Наш A/B

Оптимизатор — не дело вкуса: его выбирают по измерениям. Перед основным обучением мы трижды обучили одного и того же Ростка на одних и тех же 10 миллионах токенов: с AdamW и с Muon при двух скоростях обучения.

В этих коротких запусках Muon дал ошибку на проверке 2,393 против 2,709 у AdamW. Разница около 0,32 ната соответствует отношению перплексий $e^{0{,}32} \approx 1{,}37$. Muon был немного медленнее: 29,7 тысячи токенов в секунду против 31,8 тысячи. Но при одинаковом числе обучающих токенов качество оказалось лучше, поэтому для основного запуска мы выбрали Muon со скоростью 0,02. При 0,04 результат был близким. Этих трёх опытов достаточно для нашего рабочего выбора, но не для вывода, что один оптимизатор выигрывает всегда.

Что Muon делает с обновлением весовой матрицы, прежде чем применить его?

Делит каждый элемент на корень из среднего квадрата его градиента Сохраняет сингулярные направления и приближённо выравнивает ненулевые сингулярные числа Выбрасывает маленькие элементы градиента и оставляет большие

Это и есть ортогонализация: $U\Sigma V^\top \to UV^\top$. Поэлементное деление — это Adam; Muon смотрит на матрицу целиком и выравнивает её направления, чтобы редкие направления не тонули в главных.

Расписание скорости обучения

Скорость обучения $\eta$ не остаётся постоянной. В нашем расписании три части:

• Разогрев, 200 шагов: скорость линейно растёт от почти нуля до пика. В начале веса случайны, градиенты большие и хаотичные, а средние оптимизаторов ($m$, $v$ и буфер Muon) ещё пусты. Полноразмерный шаг в этот момент может забросить модель туда, откуда она будет выбираться тысячи шагов.
• Плато: большую часть обучения скорость держится на пике. Большие шаги быстро двигают модель, но шум батчей не даёт ей улечься на дно долины — она всё время дрожит вокруг него.
• Спуск, последние 30%: скорость линейно падает до нуля. Шаги мельчают, шум усредняется, и модель оседает на дно. Ошибка на этом участке обычно заметно снижается.

В следующем опыте можно собрать собственное расписание для игрушечной задачи: шумной долины в 24 измерениях. У неё специально сделаны направления разной крутизны, а вдали от дна стены круче в восемь раз. Такая конструкция помогает увидеть, зачем нужны осторожный старт и небольшие завершающие шаги. Это иллюстрация отдельных трудностей оптимизации, а не точная копия ландшафта языковой модели.

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

Вот две функции из train.py, которые задают расписание (без изменений), и то, что они дают для обучения Ростка:

def lr_factor(step, total, warmup, cooldown): """Warm up, hold, then decay linearly to zero over the last `cooldown` share.""" if step < warmup: return (step + 1) / warmup decay_start = total * (1 - cooldown) if step < decay_start: return 1.0 return max(0.0, (total - step) / (total - decay_start)) def context_at(step, total, context, warm): """Sequence-length warm-up: short windows first, full length later.""" if not warm: return context if step < 0.1 * total: return max(64, context // 4) if step < 0.3 * total: return context // 2 return context total = int(330e6 // 32768) # 10,070 steps for step in [0, 99, 199, 1006, 1007, 3021, 7048, 8500, 10069]: T = context_at(step, total, 512, True) print(f'step {step:5d} lr x {lr_factor(step, total, 200, 0.3):.3f} window {T} windows per batch {32768 // T}')

Вторая функция задаёт разогрев длины последовательности: первые 10% шагов используют окна по 128 токенов, затем до отметки 30% — по 256, дальше — по 512. Батч всё время содержит 32 768 токенов: коротких окон в нём просто больше. Для одного окна число пар внимания растёт как квадрат длины, а при фиксированном числе токенов в батче суммарное число пар растёт линейно с длиной окна. Поэтому короткие окна экономят часть вычислений. Сначала Росток учится на местных связях; после расширения окна он получает возможность обучаться и на более дальних.

Настоящее обучение

Теперь основной запуск: 10 070 шагов, 330 миллионов токенов, Muon и AdamW, расписание выше, всё на одном Mac с M4 Pro. Он занял 3 часа 41 минуту (13 289 секунд), в среднем 24,8 тысячи токенов в секунду вместе с проверками, а итоговая ошибка на проверке — 1,554 ната на токен. Каждые 20 шагов скрипт записывал ошибку на батче, скорость обучения, длину окна и норму градиента; каждые 250 шагов — измерял ошибку на проверочном тексте (каждый раз на одних и тех же 20 батчах окон по 512 токенов).

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

• Первая сотня шагов уводит ошибку примерно с 9 ниже 4: модель открывает частоты токенов и простейшие пары. Переключитесь на логарифмическую шкалу: на ней всё обучение видно лучше, а длинная середина превращается в пологий, почти прямой склон.
• Кривые обучения и проверки идут вместе. Большая часть корпуса остаётся непрочитанной, а повторов отдельных фрагментов сравнительно мало. Это помогает обучаться на разнообразных примерах. На графике нет явного расхождения кривых, но это не доказательство отсутствия запоминания.
• Смена окон оставляет след. Проверку всегда проводят на окнах по 512 токенов, а первые 10% обучения Росток не читал ничего длиннее 128. Экзамен был труднее уроков, и сразу после перехода на длинные окна ошибка на проверке падает заметно быстрее. Ошибка на обучении при каждой смене тоже скачком снижается: с более длинным контекстом следующий токен угадывать легче.
• Норма градиента (нижняя панель) начинается намного выше порога обрезки и опускается ниже него примерно за первые полторы сотни шагов.

В логарифмической шкале середина кривой — длинный пологий склон: каждое десятикратное увеличение числа шагов снимает похожий кусок ошибки, с каждым разом чуть меньший. Такие кривые описывают степенным законом: $\mathcal{L}(D) \approx E + A \cdot D^{-\alpha}$, где $E$ — часть ошибки, которую не убрать никаким количеством данных. Каждое десятикратное увеличение данных снимает одну и ту же долю того, что осталось сверх $E$ (на графике с логарифмами по обеим осям $\mathcal{L} - E$ — прямая). Это закон масштабирования, который Kaplan et al. (2020) нашли для моделей самых разных размеров и который уточнила Chinchilla в 2022 году. Благодаря ему ошибку огромной модели можно предсказать по серии маленьких, прежде чем тратить миллионы на её обучение. Мы вернёмся к нему в последней главе.

Что помогает Ростку сохранять близкие ошибки на обучающих и отложенных текстах в этом запуске?

Потому что проверочный текст взяли из обучающего Большой корпус даёт много новых примеров и сравнительно мало повторных предъявлений Потому что обрезка градиента не даёт переобучиться

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

Контрольные точки бок о бок

Во время обучения мы несколько раз сохраняли веса Ростка. Дадим всем версиям одну и ту же затравку и одно и то же случайное зерно и сравним. Каждая контрольная точка — полноценная модель на 17,5 МБ, поэтому при первом запуске они скачиваются одна за другой.

Лестница внизу показывает место Ростка среди всех моделей курса. Буквенная биграмма из главы 1, буквенные сети, сеть на токенах, трансформеры в один и четыре блока — каждая следующая модель немного снижала ошибку. Росток на шаге 200, всего после 6,5 миллиона токенов, уже обгоняет сеть на токенах. К шагу 1 000 он почти догоняет четырёхблочный трансформер из главы 8, а остаток обучения уводит его гораздо дальше: сравните ошибку на проверке на шаге 1 000 и в конце на кривых выше. Идеи всё время были те же; понадобились размер, данные и хорошо собранный цикл обучения.

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

Это результат предобучения — базовая модель. Она прочитала треть миллиарда токенов и продолжает любой текст в духе своего корпуса. Пишет детские истории с героями, у которых есть имена, держит диалог в формате Имя: реплика, знает, как история начинается и заканчивается. Но отвечать в роли ассистента её ещё не учили: спросите её о чём-нибудь — и она может продолжить писать, будто ваш вопрос — строчка из рассказа или пьесы: новые реплики, слова рассказчика, выдуманные собеседники, но не ответ вам. А слова она выбирает, бросая кости, которые мы для неё настроили. В следующей главе мы разберём эти кости — температуру, top-k, top-p — и увидим, почему жадный выбор застревает в петлях.