Своё дообучение

У нас есть чат-модель: она отвечает простыми фразами, рассказывает истории и иногда напоминает о своих ограничениях. Теперь попробуем изменить её манеру речи. Например, добавить пиратские словечки или научить отвечать в рифму. Хотелось бы сохранить всё полезное, чему Росток уже научился, и не хранить ради каждого стиля ещё одну полную копию модели. Как сделать такую настройку компактной?

Можно снова разрешить изменяться всем весам, как в прошлой главе. Но при полном дообучении нужны память для градиентов и состояния оптимизатора, а результат обычно сохраняют как новый набор весов. Для маленького Ростка это посильно; с миллиардами параметров расходы становятся заметными. Есть способ обучать и хранить гораздо меньше дополнительных чисел.

Этот способ называется LoRA, low-rank adaptation — адаптация с помощью матриц низкого ранга. Исходные веса на время обучения замораживают, а к выбранным линейным слоям добавляют небольшие обучаемые поправки. Насколько такой поправки хватит для задачи, предстоит проверить. Сначала разберём на картинках, что означает ранг матрицы и почему две небольшие таблицы иногда могут заменить одну большую.

Дообучить всё целиком

Вспомним, что было в главе 12. Мы взяли базового Ростка и продолжили его учить, теперь на диалогах: та же функция потерь, тот же спуск, только данные новые. Меняться мог любой из 17 309 056 параметров. Это называется полным дообучением.

Для оценки памяти представим полное дообучение с Adam, где веса, градиенты и оба скользящих средних хранятся во float32: четыре числа по четыре байта, то есть 16 байт на параметр. Для Ростка это около 277 МБ, для модели на 8 миллиардов параметров — около 128 ГБ. Это ещё без активаций и временных буферов. Конкретный расход зависит от оптимизатора и форматов: у Muon другое состояние, а смешанная точность, квантование и распределение по устройствам меняют расчёт.

Есть и вторая цена. На выходе получается целая новая модель. Нужны пират и поэт? Ещё две копии. Сто клиентов, и каждому свой тон? Сто копий.

А ведь интуиция подсказывает, что новый стиль — изменение маленькое. Росток-пират по-прежнему знает грамматику, по-прежнему помнит, что Луна вращается вокруг Земли, по-прежнему ставит слова в правильном порядке. Меняется только манера: «arr», «matey», «be» вместо «is». Неужели ради маленькой перемены нужно 17 миллионов новых чисел? Чтобы ответить, сначала научимся измерять, насколько сложна перемена в матрице.

Сколько чисел нужно перемене

Большинство параметров Ростка собрано в матрицы — прямоугольные таблицы чисел; у нормализации есть и одномерные коэффициенты. Например, объединённая проекция запросов, ключей и значений преобразует 384 входных числа в 1 152 выходных и содержит 442 368 весов. Чтобы почувствовать структуру матрицы, представим её картинкой: число в каждой клетке задаёт яркость одного пикселя.

Возьмём столбец чисел $u$ (по одному на строку) и строку чисел $v$ (по одному на столбец) и заполним таблицу по правилу «клетка = число её строки × число её столбца»:

$$W_{ij} = u_i \, v_j, \qquad W = u\,v^\top .$$

Это внешнее произведение, и картинка из него получается совсем простая: каждая строка — один и тот же узор $v$, только умноженный на $u_i$, то есть темнее или светлее. Похоже на шотландку, где нити перекрещиваются под прямым углом. Если оба вектора ненулевые, у такой матрицы ранг один, и чтобы её хранить, достаточно $m + n$ чисел вместо $m \cdot n$.

Если сложить несколько таких шотландок, можно нарисовать почти что угодно. Наименьшее число слоёв ранга один, которые в сумме дают матрицу в точности, называется её рангом $r$. Матрицу ранга $r$ можно хранить в $r(m + n)$ числах: $r$ столбцов и $r$ строк.

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

То же самое на numpy: SVD — одна строка, лучшая копия ранга $r$ — ещё одна. Запустите и сверьте последнюю строку вывода со строкой для ранга 8 — они должны совпасть, как и обещает теорема.

import numpy as np m, n = 64, 96 i, j = np.mgrid[0:m, 0:n] # a "picture": a round sun, a striped sky, a dark strip of ground W = 1.0 * ((i - 24) ** 2 + (j - 60) ** 2 < 14 ** 2) W += 0.3 * np.sin(i / 5) * np.cos(j / 9) W[54:, :] = 0.8 U, S, Vt = np.linalg.svd(W, full_matrices=False) for r in [1, 2, 4, 8, 9, 10]: W_r = (U[:, :r] * S[:r]) @ Vt[:r] # the best rank-r copy of W err = np.linalg.norm(W - W_r) / np.linalg.norm(W) print(f"rank {r:2d}: {r * (m + n):5d} numbers instead of {m * n}, error {err:5.1%}") # Eckart-Young: the error is exactly what the dropped singular values carry r = 8 print(f"check: {np.sqrt((S[r:] ** 2).sum()) / np.linalg.norm(W):.1%}")

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

Дообучение меняет матрицу весов на $\Delta W = W_\text{после} - W_\text{до}$. Может ли полезная поправка иметь небольшой ранг? Исследования внутренней размерности дообучения показали, что для некоторых задач достаточно менять параметры в пространстве значительно меньшей размерности. Но малая внутренняя размерность и низкий ранг отдельной матрицы — разные понятия. LoRA предлагает практическое ограничение: сразу искать поправки низкого ранга и проверять, хватает ли их для новой задачи.

LoRA: обходной путь рядом с весами

Вот вся LoRA целиком. Слой Ростка вычисляет $y = Wx$. Мы замораживаем $W$: ни одно её число больше не изменится. А рядом ставим обходной путь из двух тонких матриц:

$$y = W x + \frac{\alpha}{r}\, B A\, x, \qquad A \in \mathbb{R}^{r \times d_\text{in}},\quad B \in \mathbb{R}^{d_\text{out} \times r}.$$

$A$ сжимает вход до $r$ чисел, $B$ разворачивает их обратно до размера выхода. Их произведение $BA$ — матрица $d_\text{out} \times d_\text{in}$, той же формы, что и $W$, но ранг у неё не больше $r$. Учатся только $A$ и $B$. Для матрицы внимания 384 × 1 152 обходной путь ранга 8 хранит $8 \cdot (384 + 1152) = 12\,288$ чисел вместо 442 368 — это 2,8%.

Множитель $\alpha/r$ задаёт масштаб вклада адаптера. В нашем lora.py $\alpha=16$ и $r=8$, поэтому поправка умножается на 2. Деление на ранг помогает контролировать масштаб, но не гарантирует одинакового поведения при любом $r$: при смене ранга всё равно стоит проверить скорость обучения и качество. Ниже мы упомянем rsLoRA, где используют другое правило масштабирования.

Есть одна тонкость: с чего начинать. LoRA заполняет $A$ маленькими случайными числами, а $B$ — ровно нулями. Тогда на первом шаге $BA = 0$, и модель в точности совпадает с чат-моделью: адаптер начинает с того, что ничего не делает, и учится уже отсюда. Но почему бы не обнулить обе? Посчитаем градиенты.

Всё это можно проверить и заодно посмотреть на процесс целиком — на игрушке. Роль замороженных весов $W$ здесь играет картинка Ростка. «Дообучение» — это костюм: цель — Росток, наряженный пиратом или поэтом. Перемена $\Delta W$ — сам костюм, и LoRA ранга $r$ должна найти её градиентным спуском, начав с $B = 0$.

На картинках хорошо видно, что разным изменениям нужен разный ранг. Берет и шарф поэта состоят преимущественно из горизонтальных и вертикальных полос. Диагональный ремешок пиратской повязки сложнее представить несколькими внешними произведениями: тёмная точка сдвигается от строки к строке. При ранге 8 ошибка у пирата заметно больше. Пунктир показывает лучшую достижимую ошибку данного ранга по теореме Эккарта — Янга. Сравните с ним результат оптимизатора: насколько близко он подойдёт за отведённое число шагов, зависит от настроек.

Что случилось бы, если бы и A, и B начинались с нулей?

Адаптер учился бы вдвое быстрее Ничего: оба градиента были бы нулями, и адаптер не сдвинулся бы никогда Модель забыла бы всё на первом же шаге

В градиент по A входит B, а в градиент по B входит A. Если обе нули, каждый шаг спуска нулевой. Одна из двух должна начинаться со случайных чисел; LoRA выбирает A, чтобы произведение BA всё равно стартовало с нуля и модель в начале оставалась прежней.

Сколько это экономит

Теперь посчитаем на настоящих моделях. Наш lora.py ставит обходной путь рядом с каждой матрицей внутри блоков: общей матрицей Q K V во внимании, выходной проекцией внимания и всеми тремя матрицами сети прямого распространения (SwiGLU). Таблицу эмбеддингов он не трогает. Выберите модель, ранг и место для адаптеров:

Для Ростка LoRA прежде всего удобна: адаптеры маленькие, их легко хранить и менять. У большой модели экономия существеннее. Например, 8 миллиардов замороженных весов в 16-битном формате занимают примерно 16 ГБ; обучаемые адаптеры, их градиенты и состояния оптимизатора добавляют сравнительно немного. Но остаются активации и рабочие буферы, которые могут требовать много памяти. QLoRA сокращает память под исходные веса с помощью 4-битного хранения. Поместится ли весь запуск на ноутбуке, зависит также от длины контекста, батча, ранга и реализации.

В коде всё ещё короче, чем на словах. Весь слой LoRA из lora.py:

class LoRALinear(nn.Module): def __init__(self, base, rank=8, alpha=16): super().__init__() self.base = base self.A = nn.Parameter(torch.randn(rank, base.in_features) / math.sqrt(base.in_features)) self.B = nn.Parameter(torch.zeros(base.out_features, rank)) self.scale = alpha / rank def forward(self, x): return self.base(x) + (x @ self.A.T @ self.B.T) * self.scale def merge(self): self.base.weight.data += (self.B @ self.A).to(self.base.weight.dtype) * self.scale return self.base

Три детали стоит сверить с математикой выше. A стартует случайной с масштабом $1/\sqrt{d_\text{in}}$, B — нулевой. forward сначала умножает на A.T и только потом на B.T: полная матрица $BA$ размером $d_\text{out} \times d_\text{in}$ так и не строится, и обходной путь стоит всего $r(d_\text{in} + d_\text{out})$ умножений на токен. А merge — одна строка, к ней мы сейчас вернёмся.

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

TARGETS = [('attn', 'qkv'), ('attn', 'proj'), ('ffn', 'w1'), ('ffn', 'w2'), ('ffn', 'w3')] def add_lora(model, rank, alpha): """Freeze everything, then wrap every projection in the blocks with a LoRA detour.""" for p in model.parameters(): p.requires_grad_(False) for block in model.blocks: for part, name in TARGETS: parent = getattr(block, part) setattr(parent, name, LoRALinear(getattr(parent, name), rank, alpha)) return [p for p in model.parameters() if p.requires_grad] # ... trainable = add_lora(model, args.rank, args.alpha) # ... opt = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=0.0)

Всё остальное в lora.py — цикл обучения из главы 12: те же render() и pack(), та же функция потерь, которая учитывает только слова самого Ростка.

Влить обратно

После обучения у нас есть замороженная $W$ и обходной путь $BA$. Можно было бы на каждом токене считать оба пути. Но обходной путь линейный, а значит, два пути склеиваются в одну матрицу:

$$W x + s\,BAx = (W + s\,BA)\,x = W' x .$$

Считаем $W' = W + s\,BA$ один раз, записываем туда, где была $W$, а обходной путь выбрасываем. Остаётся обычный Росток: те же формы, те же 17 309 056 чисел, та же скорость. Поэтому Росток-пират в этой главе — обычный файл модели на 17,5 МБ, и браузерный движок загружает его, ничего не зная про LoRA.

def merge_lora(model): for block in model.blocks: for part, name in TARGETS: parent = getattr(block, part) layer = getattr(parent, name) if isinstance(layer, LoRALinear): setattr(parent, name, layer.merge()) # ... adapter = {n: p.detach().cpu() for n, p in model.named_parameters() if p.requires_grad} torch.save({'rank': args.rank, 'alpha': args.alpha, 'adapter': adapter}, f'{args.out}/adapter.pt') merge_lora(model) torch.save({'config': cfg.to_dict(), 'model': model.state_dict(), 'step': steps}, f'{args.out}/ckpt_final.pt')

lora.py сохраняет и то и другое: отдельно адаптер (adapter.pt, те самые 417 792 числа) и уже слитую модель. Сливать не обязательно. Если держать адаптеры отдельно, одна базовая модель в памяти может обслуживать сразу много стилей, подключая к каждому запросу свой маленький файл: так S-LoRA (Шэн и др., 2023) обслуживает тысячи адаптеров с одной видеокарты. Мы сливаем, потому что Росток маленький, а браузерный движок хочется держать простым.

Соберём теперь всю идею примерно в двадцать пять строк numpy, которые можно запустить: замороженный слой, спрятанная перемена ранга 2, которую адаптеру предстоит найти, обходной путь ранга 4, обученный ровно теми градиентами, что мы вывели, и слияние в конце.

import numpy as np rng = np.random.default_rng(1) d_in, d_out, r = 64, 48, 4 W = rng.standard_normal((d_out, d_in)) / 8 # frozen, "pre-trained" weight # the new habit: a rank-2 change the adapter has to discover dW = rng.standard_normal((d_out, 2)) @ rng.standard_normal((2, d_in)) / 8 X = rng.standard_normal((256, d_in)) # 256 inputs Y = X @ (W + dW).T # what the tuned layer should say A = rng.standard_normal((r, d_in)) / np.sqrt(d_in) # random, as in lora.py B = np.zeros((d_out, r)) # zero: step 0 changes nothing s, lr = 2.0, 0.01 # s = alpha / r; lora.py has 16 / 8 for step in range(151): h = X @ A.T # the detour: x -> A x Y_hat = X @ W.T + s * h @ B.T # frozen path + s * B A x err = Y_hat - Y loss = (err ** 2).sum() / len(X) g = 2 * err / len(X) # dLoss / dY_hat dB = s * g.T @ h # s * g (Ax)^T dA = s * (g @ B).T @ X # s * B^T g x^T A -= lr * dA B -= lr * dB if step % 25 == 0: print(f"step {step:3d} loss {loss:.5f}") W_merged = W + s * B @ A # merge: an ordinary matrix again print("merged gives the same outputs:", np.allclose(X @ W_merged.T, X @ W.T + s * (X @ A.T) @ B.T)) print("trained", A.size + B.size, "numbers instead of", W.size)

Функция потерь падает с 89 до нуля, а после слияния слой выдаёт то же самое, что замороженный путь плюс обходной. Теперь поменяйте в строке d_in, d_out, r = 64, 48, 4 последнее число на 1 и запустите снова: функция потерь застрянет около 41. Искомая перемена имеет ранг 2, и обходной путь ранга 1 не может её выразить, сколько его ни учи.

Три голоса Ростка

Пора за настоящее дело. Мы обучили два адаптера поверх чат-модели sprout-chat. Данные для каждого — 500 обычных диалогов Ростка, в которых каждый ответ переписан: для одного адаптера пиратской речью, для другого — в рифму. Вопросы остались прежними, переоделись только ответы. Настройки — те, что в lora.py по умолчанию: ранг 8, α = 16, AdamW с шагом обучения 0,003, девять проходов по данным, и функция потерь считается только на словах Ростка, как в главе 12. Число проходов мы подобрали по отложенным диалогам: их ошибка ниже всего примерно на седьмом-восьмом проходе, а потом снова растёт — 500 диалогов легко вызубрить. Это всего 34 шага для пирата и 41 для поэта: меньше минуты на Mac.

Задайте всем троим один вопрос. У них общее зерно случайности, так что ответы различаются только из-за весов.

Сколько пришлось выучить каждому адаптеру? Вот их настоящие кривые обучения. Толстая линия — ошибка на диалогах, на которых адаптер не учился (lora.py откладывает для этого каждый двадцатый); тонкая — ошибка на обучающих батчах.

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

Что модель забывает

Чему бы нейросеть ни училась, учится она, меняя числа, которые уже чем-то были заняты. Когда дообучение тянет веса к новой задаче, старые навыки могут портиться. В крайнем случае это называют катастрофическим забыванием: сеть, обученная задаче Б, начисто забывает задачу А. Маклоски и Коэн описали его ещё в 1989 году на простейших сетях, и никуда оно с тех пор не делось.

Один из способов заметить ухудшение старого навыка — снова измерить ошибку на текстах прежнего типа. Дадим четырём версиям Ростка одинаковые примеры и сравним среднее $-\ln p$ на токен. Рост ошибки означает, что эта версия хуже предсказывает выбранный текст. Несколько коротких примеров дают полезную подсказку, но для оценки забывания навыка целиком нужна более широкая проверка.

На что смотреть? Если адаптеры выучили свои стили, в пиратской строке самым спокойным должен быть Росток-пират, а в стихотворной — Росток-поэт. Цена — в остальных строках: насколько сильнее каждый адаптер удивляется обычному ответу или обычной истории, чем чат-модель, из которой он вырос? Эта разница показывает ухудшение предсказания конкретного примера, а не всех старых навыков сразу. На наших моделях оба адаптера свои стили выучили (в пиратской строке 2,93 против 4,14 у чат-модели, в стихотворной — 3,72 против 5,14), а забывание маленькое, но неравное: пират удивляется обычному ответу всего на 0,13 ната сильнее чат-модели, а обычной истории — не сильнее вовсе; поэт платит за обычный ответ 0,56 ната. Стихи уводят слова от обычной речи дальше, чем пиратские словечки, — сравните, где начинаются две кривые выше.

Ограничить его можно несколькими способами:

• Маленький ранг и мало шагов. Это ограничивает сложность поправки и время обучения, но не её величину: даже матрица ранга один может содержать очень большие числа.
• Немного старых данных в новой смеси. Несколько процентов обычных диалогов напоминают модели её привычную манеру; это называют повторением (replay).
• Вовремя остановиться. Следить за ошибкой на отложенных данных обоих видов и остановиться, когда старая начнёт расти.
• Беречь важные веса. Метод EWC (Киркпатрик и др., 2017) штрафует за изменение тех весов, которые больше всего значили для старой задачи.

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

Своё дообучение

Всё, что было выше, можно повторить на своих данных. Вот рецепт целиком.

1. Данные. Файл, где на каждой строке один диалог, в том же формате, что и в главе 12:

{"messages": [{"role": "user", "content": "how do birds fly?"}, {"role": "assistant", "content": "Arr, they flap their wings and push the air down, matey, and the air pushes them up, just like wind in a sail!"}]}

В нашем опыте у каждого адаптера было 500 диалогов. Это отправная точка для похожего стиля, а не минимальное число для любой задачи. Важны разнообразие вопросов и согласованность ответов: одни «arr» на каждую реплику ещё не учат отвечать по существу. Начните с коротких фраз и знакомой Ростку лексики, оставьте часть тем для проверки. Если смешиваете новый стиль с обычным, оценивайте, когда и как модель выбирает каждый из них.

2. Обучение. Начинаем с чат-модели: стиль ложится поверх умения разговаривать.

# обучить адаптер и влить его в веса
python lora.py --base runs/chat/ckpt_final.pt --chats my-style.jsonl --out runs/my-style
# поговорить с результатом в терминале
python chat.py runs/my-style/ckpt_final.pt
# упаковать для браузера: int8, как в главе 11
python export.py runs/my-style/ckpt_final.pt sprout-my-style.bin

lora.py выбирает доступное устройство: CUDA для поддерживаемой NVIDIA, MPS для Apple silicon или обычный процессор. Нужны скрипты, data/tokenizer.json, контрольная точка чат-модели и файл ваших диалогов. Росток достаточно мал для экспериментов без большого сервера, но время зависит от устройства и настроек. Продолжительность наших двух запусков указана на графиках выше.

3. Проверка. Следите за проверочной ошибкой в log.jsonl: если она начинает расти, а ошибка на обучении всё падает, адаптер зубрит ваши примеры вместо того, чтобы учить стиль, — уменьшите --epochs. Поговорите с результатом на темы, которых нет в ваших данных: там и проявится забывание. Крутить стоит --rank (сравните 2, 8 и 32 по проверочной ошибке), --lr и --epochs.

После слияния Росток-пират пишет текст у вас в браузере. Насколько он быстрее или медленнее обычной чат-модели?

Медленнее: каждый токен должен пройти ещё и обходной путь B·A·x Число операций на токен прежнее: после слияния меняются веса, а не форма модели Быстрее: LoRA убрала большую часть параметров

У W′ = W + s·BA та же форма, что у W. Движок выполняет те же умножения матриц, что и раньше, — различаются только числа внутри.

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

Росток научился менять манеру речи, не переучиваясь разговаривать заново. Каждый стиль — поправка ранга 8 из 417 792 чисел, обученная на 500 диалогах и влитая обратно в веса. Переключайте голоса прямо посреди разговора и смотрите, кто что помнит. В последней главе мы соберём все версии Ростка из этого курса в одном саду и посмотрим, куда двигаться дальше.