microgpt

Это краткий гид по моему новому арт-проекту microgpt — одному файлу в 200 строк чистого Python без зависимостей, который обучает GPT и запускает её инференс. В этом файле содержится всё алгоритмически необходимое: датасет документов, токенизатор, движок autograd, нейросетевая архитектура в духе GPT-2, оптимизатор Adam, цикл обучения и цикл инференса. Всё остальное — это лишь вопросы эффективности. Упростить дальше я уже не могу. Этот скрипт — кульминация нескольких проектов (micrograd, makemore, nanogpt и др.) и десятилетней одержимости свести LLM к их сути, и я считаю, что это красиво 🥹. Он даже идеально разбивается на 3 колонки:

Где найти:

В этом GitHub gist лежит весь исходный код: microgpt.py. Также доступен на этой веб-странице: https://karpathy.ai/microgpt.html. И в виде Google Colab notebook. NEW: купить microgpt как триптих можно в моём арт-магазине karpathy.art :)

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

Датасет

Топливо больших языковых моделей — поток текстовых данных, опционально разделённый на набор документов. В промышленных приложениях каждый документ — это веб-страница из интернета, но для microgpt мы используем более простой пример: 32 000 имён по одному на строку:

# Пусть есть входной датасет `docs`: list[str] документов (например, датасет имён) if not os.path.exists('input.txt'): import urllib.request names_url = 'https://raw.githubusercontent.com/karpathy/makemore/refs/heads/master/names.txt' urllib.request.urlretrieve(names_url, 'input.txt') docs = [l.strip() for l in open('input.txt').read().strip().split('\n') if l.strip()] # list[str] документов random.shuffle(docs) print(f"num docs: {len(docs)}")

Датасет выглядит так. Каждое имя — это документ:

emma olivia ava isabella sophia charlotte mia amelia harper ... (далее ~32 000 имён)

Цель модели — выучить закономерности в данных и затем генерировать похожие новые документы, которые разделяют те же статистические паттерны. В качестве превью: к концу скрипта наша модель будет генерировать («галлюцинировать»!) новые правдоподобно звучащие имена. Забегая вперёд, мы получим:

sample 1: kamon sample 2: ann sample 3: karai sample 4: jaire sample 5: vialan sample 6: karia sample 7: yeran sample 8: anna sample 9: areli sample 10: kaina sample 11: konna sample 12: keylen sample 13: liole sample 14: alerin sample 15: earan sample 16: lenne sample 17: kana sample 18: lara sample 19: alela sample 20: anton

Выглядит не слишком впечатляюще, но с точки зрения модели вроде ChatGPT ваш диалог с ней — это просто странновато выглядящий «документ». Когда вы инициализируете документ своим промптом, ответ модели с её точки зрения — это просто статистическое продолжение документа.

Токенизатор

Под капотом нейросети работают с числами, а не с символами, поэтому нам нужен способ конвертировать текст в последовательность целочисленных id токенов и обратно. Промышленные токенизаторы вроде tiktoken (используется GPT-4) для эффективности оперируют группами символов, но простейший возможный токенизатор просто присваивает по одному целому числу каждому уникальному символу в датасете:

# Пусть есть Tokenizer для перевода строк в дискретные символы и обратно uchars = sorted(set(''.join(docs))) # уникальные символы в датасете становятся id токенов 0..n-1 BOS = len(uchars) # id специального токена Beginning of Sequence (BOS) vocab_size = len(uchars) + 1 # общее число уникальных токенов, +1 — это BOS print(f"vocab size: {vocab_size}")

В коде выше мы собираем все уникальные символы по датасету (это просто все строчные буквы a-z), сортируем их, и каждая буква получает id по своему индексу. Заметьте, что сами целочисленные значения не имеют никакого смысла; каждый токен — это просто отдельный дискретный символ. Вместо 0, 1, 2 это вполне могли бы быть разные эмодзи. Кроме того, мы создаём ещё один специальный токен под названием BOS (Beginning of Sequence), который служит разделителем: он говорит модели «здесь начинается/заканчивается новый документ». Позже во время обучения каждый документ оборачивается BOS с обеих сторон: [BOS, e, m, m, a, BOS]. Модель учится тому, что BOS начинает новое имя, а следующий BOS его заканчивает. Итого мы получаем финальный словарь из 27 элементов (26 возможных строчных символов a-z плюс токен BOS).

Autograd

Для обучения нейросети нужны градиенты: для каждого параметра модели нам надо знать «если я чуть-чуть увеличу это число, loss вырастет или упадёт, и насколько?». У вычислительного графа много входов (параметры модели и входные токены), но он сходится к одному скалярному выходу: loss (мы определим, что такое loss, ниже). Backpropagation начинается с этого единственного выхода и движется в обратную сторону по графу, вычисляя градиент loss по каждому входу. Он опирается на цепное правило из матанализа. В промышленных решениях это автоматически делают библиотеки вроде PyTorch. Здесь мы реализуем его с нуля в одном классе под названием Value:

class Value: __slots__ = ('data', 'grad', '_children', '_local_grads') def __init__(self, data, children=(), local_grads=()): self.data = data # скалярное значение этого узла, посчитанное при forward pass self.grad = 0 # производная loss по этому узлу, считается при backward pass self._children = children # дети этого узла в вычислительном графе self._local_grads = local_grads # локальная производная этого узла по его детям def __add__(self, other): other = other if isinstance(other, Value) else Value(other) return Value(self.data + other.data, (self, other), (1, 1)) def __mul__(self, other): other = other if isinstance(other, Value) else Value(other) return Value(self.data * other.data, (self, other), (other.data, self.data)) def __pow__(self, other): return Value(self.data**other, (self,), (other * self.data**(other-1),)) def log(self): return Value(math.log(self.data), (self,), (1/self.data,)) def exp(self): return Value(math.exp(self.data), (self,), (math.exp(self.data),)) def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),)) def __neg__(self): return self * -1 def __radd__(self, other): return self + other def __sub__(self, other): return self + (-other) def __rsub__(self, other): return other + (-self) def __rmul__(self, other): return self * other def __truediv__(self, other): return self * other**-1 def __rtruediv__(self, other): return other * self**-1 def backward(self): topo = [] visited = set() def build_topo(v): if v not in visited: visited.add(v) for child in v._children: build_topo(child) topo.append(v) build_topo(self) self.grad = 1 for v in reversed(topo): for child, local_grad in zip(v._children, v._local_grads): child.grad += local_grad * v.grad

Я понимаю, что это самая математически и алгоритмически насыщенная часть, и у меня есть 2,5-часовое видео по ней: видео про micrograd. Если коротко, Value оборачивает одно скалярное число (.data) и отслеживает, как оно было получено. Представьте каждую операцию как маленький лего-блок: он принимает на вход какие-то входы, выдаёт выход (forward pass) и знает, как его выход изменится по отношению к каждому из входов (локальный градиент). Это вся информация, которая нужна autograd от каждого блока. Всё остальное — это просто цепное правило, нанизывающее блоки друг на друга.

Каждый раз, когда вы делаете математику над объектами Value (сложение, умножение и т. д.), результатом становится новый Value, помнящий свои входы (_children) и локальную производную этой операции (_local_grads). Например, __mul__ запоминает, что \(\frac{\partial(a \cdot b)}{\partial a} = b\) и \(\frac{\partial(a \cdot b)}{\partial b} = a\). Полный набор лего-блоков:

Метод backward() обходит этот граф в обратном топологическом порядке (начиная от loss и заканчивая параметрами), применяя цепное правило на каждом шаге. Если loss обозначить как \(L\), а у узла \(v\) есть ребёнок \(c\) с локальным градиентом \(\frac{\partial v}{\partial c}\), то:

Это выглядит слегка пугающе, если вы не очень в ладах с матанализом, но на самом деле тут просто перемножаются два числа интуитивным образом. Один из способов это увидеть: «Если машина едет в два раза быстрее велосипеда, а велосипед — в четыре раза быстрее идущего человека, то машина едет в 2 × 4 = 8 раз быстрее человека». Цепное правило — это та же идея: вы перемножаете скорости изменения вдоль пути.

Запускаем всё, выставив self.grad = 1 на узле loss, потому что \(\frac{\partial L}{\partial L} = 1\): скорость изменения loss по самому себе тривиально равна 1. Дальше цепное правило просто перемножает локальные градиенты вдоль каждого пути обратно к параметрам.

Обратите внимание на += (накопление, а не присваивание). Когда значение используется в нескольких местах графа (т. е. граф ветвится), градиенты текут обратно по каждой ветке независимо и должны суммироваться. Это следствие многомерного цепного правила: если \(c\) вносит вклад в \(L\) по нескольким путям, общая производная — это сумма вкладов от каждого пути.

После завершения backward() у каждого Value в графе есть .grad, содержащий \(\frac{\partial L}{\partial v}\), который говорит нам, как изменился бы итоговый loss, если бы мы чуть подёргали это значение.

Вот конкретный пример. Заметьте, что a используется дважды (граф ветвится), поэтому её градиент — сумма по обоим путям:

a = Value(2.0) b = Value(3.0) c = a * b # c = 6.0 L = c + a # L = 8.0 L.backward() print(a.grad) # 4.0 (dL/da = b + 1 = 3 + 1, по обоим путям) print(b.grad) # 2.0 (dL/db = a = 2)

Ровно это же даёт .backward() в PyTorch:

import torch a = torch.tensor(2.0, requires_grad=True) b = torch.tensor(3.0, requires_grad=True) c = a * b L = c + a L.backward() print(a.grad) # tensor(4.) print(b.grad) # tensor(2.)

Это тот же алгоритм, который выполняет loss.backward() в PyTorch, только на скалярах, а не на тензорах (массивах скаляров) — алгоритмически идентично, существенно меньше и проще, но, разумеется, заметно менее эффективно.

Давайте раскроем то, что нам выше дал .backward(). Autograd посчитал, что если L = a*b + a, и a=2, и b=3, то a.grad = 4.0 говорит нам о локальном влиянии a на L. Если вы пошевелите вход a, в какую сторону изменится L? Здесь производная L по a равна 4.0, то есть если мы увеличим a на крошечную величину (скажем, 0.001), L увеличится примерно в 4 раза больше (0.004). Аналогично b.grad = 2.0 означает, что такое же подталкивание b увеличило бы L примерно вдвое больше этой величины (0.002). Иначе говоря, эти градиенты сообщают нам направление (положительное или отрицательное в зависимости от знака) и крутизну (величину) влияния каждого отдельного входа на итоговый выход (loss). Это и позволяет нам итеративно подталкивать параметры нашей нейросети, чтобы снижать loss и тем самым улучшать её предсказания.

Параметры

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

n_embd = 16 # размерность эмбеддинга n_head = 4 # число attention-голов n_layer = 1 # число слоёв block_size = 16 # максимальная длина последовательности head_dim = n_embd // n_head # размерность каждой головы matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)] state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd), 'lm_head': matrix(vocab_size, n_embd)} for i in range(n_layer): state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd) state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd) state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd) state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd) state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd) state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd) params = [p for mat in state_dict.values() for row in mat for p in row] print(f"num params: {len(params)}")

Каждый параметр инициализируется небольшим случайным числом из гауссовского распределения. state_dict организует их в именованные матрицы (заимствуя терминологию PyTorch): таблицы эмбеддингов, веса attention, веса MLP и финальная выходная проекция. Также мы складываем все параметры в один плоский список params, чтобы оптимизатору потом было удобно по ним пройтись. В нашей крошечной модели получается 4192 параметра. У GPT-2 было 1,6 миллиарда, а у современных LLM — сотни миллиардов.

Архитектура

Архитектура модели — это функция без состояния: она принимает токен, позицию, параметры и закэшированные keys/values с предыдущих позиций, и возвращает logits (оценки) того, какой токен, по мнению модели, должен идти следующим в последовательности. Мы следуем GPT-2 с небольшими упрощениями: RMSNorm вместо LayerNorm, без bias-ов и ReLU вместо GeLU. Сначала три маленькие вспомогательные функции:

def linear(x, w): return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]

linear — это умножение матрицы на вектор. Она принимает вектор x и матрицу весов w и вычисляет по одному скалярному произведению на каждую строку w. Это фундаментальный строительный блок нейросетей: обучаемое линейное преобразование.

def softmax(logits): max_val = max(val.data for val in logits) exps = [(val - max_val).exp() for val in logits] total = sum(exps) return [e / total for e in exps]

softmax переводит вектор сырых оценок (logits), которые могут лежать в диапазоне от \(-\infty\) до \(+\infty\), в распределение вероятностей: все значения оказываются в \([0, 1]\) и в сумме дают 1. Сначала мы вычитаем максимум для численной устойчивости (математически это не меняет результат, но предотвращает переполнение в exp).

def rmsnorm(x): ms = sum(xi * xi for xi in x) / len(x) scale = (ms + 1e-5) ** -0.5 return [xi * scale for xi in x]

rmsnorm (Root Mean Square Normalization) перенормирует вектор так, чтобы его значения имели единичное среднеквадратичное значение. Это удерживает активации от взрывного роста или затухания по мере прохождения через сеть, что стабилизирует обучение. Это более простой вариант LayerNorm, использовавшегося в оригинальной GPT-2.

Теперь сама модель:

def gpt(token_id, pos_id, keys, values): tok_emb = state_dict['wte'][token_id] # эмбеддинг токена pos_emb = state_dict['wpe'][pos_id] # эмбеддинг позиции x = [t + p for t, p in zip(tok_emb, pos_emb)] # совместный эмбеддинг токена и позиции x = rmsnorm(x) for li in range(n_layer): # 1) Multi-head attention block x_residual = x x = rmsnorm(x) q = linear(x, state_dict[f'layer{li}.attn_wq']) k = linear(x, state_dict[f'layer{li}.attn_wk']) v = linear(x, state_dict[f'layer{li}.attn_wv']) keys[li].append(k) values[li].append(v) x_attn = [] for h in range(n_head): hs = h * head_dim q_h = q[hs:hs+head_dim] k_h = [ki[hs:hs+head_dim] for ki in keys[li]] v_h = [vi[hs:hs+head_dim] for vi in values[li]] attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))] attn_weights = softmax(attn_logits) head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)] x_attn.extend(head_out) x = linear(x_attn, state_dict[f'layer{li}.attn_wo']) x = [a + b for a, b in zip(x, x_residual)] # 2) MLP block x_residual = x x = rmsnorm(x) x = linear(x, state_dict[f'layer{li}.mlp_fc1']) x = [xi.relu() for xi in x] x = linear(x, state_dict[f'layer{li}.mlp_fc2']) x = [a + b for a, b in zip(x, x_residual)] logits = linear(x, state_dict['lm_head']) return logits

Функция обрабатывает один токен (с id token_id) в конкретной позиции во времени (pos_id) и некий контекст из предыдущих итераций, суммированный активациями в keys и values, известными как KV Cache. Вот что происходит шаг за шагом:

Эмбеддинги. Нейросеть не может напрямую обрабатывать сырое id токена вроде 5. Она умеет работать только с векторами (списками чисел). Поэтому мы связываем с каждым возможным токеном обучаемый вектор и подаём его в качестве нейронной подписи токена. id токена и id позиции выбирают по одной строке из соответствующих таблиц эмбеддингов (wte и wpe). Эти два вектора складываются вместе, давая модели представление, кодирующее одновременно и что это за токен, и где он находится в последовательности. Современные LLM обычно отказываются от позиционных эмбеддингов и вводят другие схемы относительного позиционирования, например RoPE.

Attention block. Текущий токен проецируется в три вектора: query (Q), key (K) и value (V). Интуитивно query говорит «что я ищу?», key говорит «что у меня внутри?», а value говорит «что я предлагаю, если меня выберут?». Например, в имени «emma», когда модель находится на второй «m» и пытается предсказать, что идёт дальше, она может выучить query вроде «какие гласные встречались недавно?». У более ранней «e» был бы key, хорошо совпадающий с этим query, поэтому она получит высокий attention-вес, и её value (информация о том, что это гласная) перетечёт в текущую позицию. Key и value добавляются в KV cache, чтобы предыдущие позиции оставались доступны. Каждая attention-голова считает скалярные произведения между своим query и всеми закэшированными keys (масштабированные на \(\sqrt{d_{head}}\)), применяет softmax для получения attention-весов и берёт взвешенную сумму закэшированных values. Выходы всех голов конкатенируются и проецируются через attn_wo. Стоит подчеркнуть, что Attention block — это единственное место, где токен в позиции t вообще «смотрит» на токены в прошлом 0..t-1. Attention — это механизм коммуникации между токенами.

MLP block. MLP — сокращение от «multilayer perceptron» (многослойный перцептрон), это двухслойная feed-forward сеть: проекция вверх до 4× размерности эмбеддинга, ReLU, проекция обратно. Это место, где модель делает большую часть «мышления» в каждой позиции. В отличие от attention, эта вычисление полностью локально во времени t. Transformer чередует коммуникацию (Attention) с вычислением (MLP).

Residual connections. Оба блока — attention и MLP — добавляют свой выход обратно к своему входу (x = [a + b for ...]). Это позволяет градиентам напрямую течь через сеть и делает более глубокие модели обучаемыми.

Output. Финальное скрытое состояние проецируется к размеру словаря через lm_head, давая по одному logit на каждый токен в словаре. В нашем случае это всего 27 чисел. Чем выше logit, тем более вероятным считает модель появление соответствующего токена следующим.

Вы можете заметить, что мы используем KV cache во время обучения, что необычно. Люди обычно ассоциируют KV cache только с инференсом. Но KV cache концептуально всегда там, даже во время обучения. В промышленных реализациях он просто спрятан внутри сильно векторизованного attention-вычисления, которое обрабатывает все позиции последовательности одновременно. Поскольку microgpt обрабатывает по одному токену за раз (никакой batch-размерности, никаких параллельных временных шагов), мы строим KV cache явно. И в отличие от типичного сценария инференса, где KV cache хранит detached тензоры, здесь закэшированные keys и values — это живые узлы Value в вычислительном графе, так что мы реально через них делаем backprop.

Цикл обучения

Теперь связываем всё вместе. Цикл обучения многократно: (1) выбирает документ, (2) прогоняет модель вперёд по его токенам, (3) вычисляет loss, (4) делает backprop, чтобы получить градиенты, и (5) обновляет параметры.

# Пусть будет Adam, благословенный оптимизатор, и его буферы learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8 m = [0.0] * len(params) # буфер первого момента v = [0.0] * len(params) # буфер второго момента # Повторяем последовательно num_steps = 1000 # число шагов обучения for step in range(num_steps): # Берём один документ, токенизируем его, оборачиваем спецтокеном BOS с обеих сторон doc = docs[step % len(docs)] tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS] n = min(block_size, len(tokens) - 1) # Прогоняем последовательность токенов через модель, строя вычислительный граф вплоть до loss. keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)] losses = [] for pos_id in range(n): token_id, target_id = tokens[pos_id], tokens[pos_id + 1] logits = gpt(token_id, pos_id, keys, values) probs = softmax(logits) loss_t = -probs[target_id].log() losses.append(loss_t) loss = (1 / n) * sum(losses) # итоговый средний loss по последовательности документа. Пусть будет низким. # Делаем backward по loss, считая градиенты по всем параметрам модели. loss.backward() # Обновление оптимизатора Adam: обновляем параметры модели по соответствующим градиентам. lr_t = learning_rate * (1 - step / num_steps) # линейный спад learning rate for i, p in enumerate(params): m[i] = beta1 * m[i] + (1 - beta1) * p.grad v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2 m_hat = m[i] / (1 - beta1 ** (step + 1)) v_hat = v[i] / (1 - beta2 ** (step + 1)) p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam) p.grad = 0 print(f"step {step+1:4d} / {num_steps:4d} | loss {loss.data:.4f}")

Пройдёмся по каждому куску:

Токенизация. Каждый шаг обучения выбирает один документ и оборачивает его BOS с обеих сторон: имя «emma» становится [BOS, e, m, m, a, BOS]. Задача модели — предсказывать каждый следующий токен по предыдущим.

Forward pass и loss. Мы подаём токены в модель по одному, попутно строя KV cache. На каждой позиции модель выдаёт 27 logits, которые мы через softmax превращаем в вероятности. Loss в каждой позиции — это отрицательный логарифм вероятности правильного следующего токена: \(-\log p(\text{target})\). Это называется cross-entropy loss. Интуитивно loss измеряет степень ошибки предсказания: насколько модель удивлена тем, что в реальности идёт следующим. Если модель присваивает правильному токену вероятность 1.0, она совсем не удивлена, и loss равен 0. Если она присваивает вероятность близкую к 0, модель очень удивлена, и loss уходит в \(+\infty\). Мы усредняем попозиционные значения loss по всему документу, чтобы получить один скаляр.

Backward pass. Один вызов loss.backward() прогоняет backpropagation через весь вычислительный граф — от loss через softmax, через модель, до каждого параметра. После этого .grad каждого параметра говорит, как его изменить, чтобы уменьшить loss.

Оптимизатор Adam. Мы могли бы просто делать p.data -= lr * p.grad (gradient descent), но Adam умнее. Он поддерживает два бегущих среднего на каждый параметр: m отслеживает среднее недавних градиентов (импульс, как катящийся шар), а v отслеживает среднее недавних квадратов градиентов (адаптируя learning rate для каждого параметра). m_hat и v_hat — это bias-коррекции, учитывающие, что m и v инициализированы нулями и нуждаются в разогреве. Learning rate спадает линейно в течение обучения. После обновления мы сбрасываем .grad = 0 для следующего шага.

За 1000 шагов loss падает примерно с 3.3 (случайное угадывание среди 27 токенов: \(-\log(1/27) \approx 3.3\)) до примерно 2.37. Чем ниже, тем лучше, и нижний предел — 0 (идеальные предсказания), так что место для улучшений ещё есть, но модель явно учит статистические паттерны имён.

Inference

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

temperature = 0.5 # в (0, 1], контролирует «творчество» сгенерированного текста, от низкого к высокому print("\n--- inference (новые, галлюцинированные имена) ---") for sample_idx in range(20): keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)] token_id = BOS sample = [] for pos_id in range(block_size): logits = gpt(token_id, pos_id, keys, values) probs = softmax([l / temperature for l in logits]) token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0] if token_id == BOS: break sample.append(uchars[token_id]) print(f"sample {sample_idx+1:2d}: {''.join(sample)}")

Мы начинаем каждый сэмпл с токена BOS, говорящего модели «начни новое имя». Модель выдаёт 27 logits, мы переводим их в вероятности и случайно сэмплируем один токен в соответствии с этими вероятностями. Этот токен подаётся обратно как следующий вход, и мы повторяем, пока модель снова не выдаст BOS (что означает «я закончила») или пока мы не упрёмся в максимальную длину последовательности.

Параметр temperature управляет случайностью. Перед softmax мы делим logits на temperature. Temperature, равная 1.0, сэмплирует напрямую из выученного распределения модели. Меньшие значения temperature (как 0.5 здесь) заостряют распределение, делая модель более консервативной и склонной выбирать свои топ-варианты. Temperature, стремящаяся к 0, всегда выбирала бы один-единственный самый вероятный токен (greedy decoding). Более высокие temperature сглаживают распределение и дают более разнообразный, но потенциально менее связный вывод.

Запуск

Всё, что вам нужно, — это Python (никаких pip install, никаких зависимостей):

python train.py

На моём MacBook скрипт работает около 1 минуты. Вы будете видеть, как loss печатается на каждом шаге:

train.py num docs: 32033 vocab size: 27 num params: 4192 step 1 / 1000 | loss 3.3660 step 2 / 1000 | loss 3.4243 step 3 / 1000 | loss 3.1778 step 4 / 1000 | loss 3.0664 step 5 / 1000 | loss 3.2209 step 6 / 1000 | loss 2.9452 step 7 / 1000 | loss 3.2894 step 8 / 1000 | loss 3.3245 step 9 / 1000 | loss 2.8990 step 10 / 1000 | loss 3.2229 step 11 / 1000 | loss 2.7964 step 12 / 1000 | loss 2.9345 step 13 / 1000 | loss 3.0544 ...

Понаблюдайте, как он падает с ~3.3 (случайность) к ~2.37. Чем ниже это число, тем лучше предсказания сети о том, какой токен идёт следующим в последовательности. К концу обучения знания о статистических паттернах обучающих последовательностей токенов сконденсированы в параметрах модели. Зафиксировав эти параметры, мы теперь можем генерировать новые, галлюцинированные имена. Вы увидите (снова):

sample 1: kamon sample 2: ann sample 3: karai sample 4: jaire sample 5: vialan sample 6: karia sample 7: yeran sample 8: anna sample 9: areli sample 10: kaina sample 11: konna sample 12: keylen sample 13: liole sample 14: alerin sample 15: earan sample 16: lenne sample 17: kana sample 18: lara sample 19: alela sample 20: anton

Как альтернатива запуску скрипта на своём компьютере, вы можете попробовать запустить его прямо в этом Google Colab notebook и позадавать про него вопросы Gemini. Поиграйтесь со скриптом! Можете попробовать другой датасет. Или потренировать дольше (увеличить num_steps), или увеличить размер модели, чтобы получать всё более хорошие результаты.

Прогрессия

Чтобы увидеть, как код наращивается кусочек за кусочком, слоями луковицы, рекомендованная прогрессия выглядит примерно так:

Я создал Gist под названием build_microgpt.py, где в Revisions можно увидеть все эти версии и diff между каждым шагом. Думаю, это может быть одним из полезных способов пройтись по кодовой базе, добавляя по одному компоненту за раз.

Серьёзные вещи

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

Данные. Вместо 32K коротких имён промышленные модели обучаются на триллионах токенов интернет-текста: веб-страниц, книг, кода и т. д. Данные дедуплицируются, фильтруются по качеству и аккуратно смешиваются по доменам.

Токенизатор. Вместо одиночных символов промышленные модели используют subword-токенизаторы вроде BPE (Byte Pair Encoding), которые учатся объединять часто встречающиеся вместе последовательности символов в одиночные токены. Распространённые слова вроде «the» становятся одним токеном, редкие слова разбиваются на куски. Это даёт словарь из ~100K токенов и куда эффективнее, потому что модель видит больше содержания на каждой позиции.

Autograd. microgpt оперирует скалярными объектами Value на чистом Python. Промышленные системы используют тензоры (большие многомерные массивы чисел) и работают на GPU/TPU, выполняющих миллиарды операций с плавающей точкой в секунду. Библиотеки вроде PyTorch обрабатывают autograd над тензорами, а CUDA-ядра вроде FlashAttention сливают несколько операций ради скорости. Математика идентична, она просто соответствует множеству скаляров, обрабатываемых параллельно.

Архитектура. У microgpt 4192 параметра. У моделей класса GPT-4 — сотни миллиардов. В целом это всё та же похожая Transformer-нейросеть, только гораздо шире (размерности эмбеддинга 10000+) и гораздо глубже (100+ слоёв). Современные LLM также включают несколько дополнительных типов лего-блоков и меняют их порядок: например, RoPE (Rotary Position Embeddings) вместо выученных позиционных эмбеддингов, GQA (Grouped Query Attention) для сокращения размера KV cache, gated linear activations вместо ReLU, слои Mixture of Experts (MoE) и т. д. Но базовая структура из Attention (коммуникация) и MLP (вычисление), чередующихся на residual stream, прекрасно сохраняется.

Обучение. Вместо одного документа за шаг промышленное обучение использует большие батчи (миллионы токенов за шаг), gradient accumulation, mixed precision (float16/bfloat16) и тщательную настройку гиперпараметров. Обучение фронтирной модели занимает тысячи GPU, работающих месяцами.

Оптимизация. microgpt использует Adam с простым линейным спадом learning rate, и это, в общем-то, всё. В масштабе оптимизация становится отдельной дисциплиной. Модели обучаются в пониженной точности (bfloat16 или даже fp8) и через большие GPU-кластеры ради эффективности, что вносит свои численные сложности. Настройки оптимизатора (learning rate, weight decay, beta-параметры, warmup-расписание, decay-расписание) должны быть подобраны точно, и правильные значения зависят от размера модели, размера батча и состава датасета. Scaling laws (например, Chinchilla) подсказывают, как распределить фиксированный compute-бюджет между размером модели и числом обучающих токенов. Ошибка в любой из этих деталей в масштабе может стоить миллионов долларов compute, поэтому команды проводят обширные эксперименты в меньшем масштабе, чтобы предсказать правильные настройки, прежде чем коммититься в полноценный обучающий запуск.

Post-training. Базовая модель, выходящая из обучения (называемая «pretrained»-моделью), — это дописыватель документов, а не чат-бот. Превращение её в ChatGPT происходит в два этапа. Первый — SFT (Supervised Fine-Tuning): вы просто меняете документы на курируемые диалоги и продолжаете обучать. Алгоритмически ничего не меняется. Второй — RL (Reinforcement Learning): модель генерирует ответы, их оценивают (люди, другая модель-«судья» или алгоритм), и модель учится на этой обратной связи. По сути модель всё ещё обучается на документах, но эти документы теперь составлены из токенов, выходящих из самой же модели.

Inference. Обслуживание модели для миллионов пользователей требует своего инженерного стека: батчинга запросов вместе, управления и paging-а KV cache (vLLM и т. п.), speculative decoding для скорости, квантования (запуска в int8/int4 вместо float16) для сокращения памяти и распределения модели на несколько GPU. По сути мы всё так же предсказываем следующий токен в последовательности, но с большим объёмом инженерии, потраченной на то, чтобы делать это быстрее.

Всё это — важные инженерные и исследовательские вклады, но если вы понимаете microgpt, вы понимаете алгоритмическую суть.

FAQ

«Понимает» ли модель что-нибудь? Это философский вопрос, но механически — никакой магии не происходит. Модель — это большая математическая функция, отображающая входные токены в распределение вероятностей по следующему токену. Во время обучения параметры подстраиваются так, чтобы правильный следующий токен становился более вероятным. Считать ли это «пониманием» — на ваше усмотрение, но механизм полностью умещается в 200 строках выше.

Почему это работает? У модели тысячи настраиваемых параметров, и оптимизатор каждый шаг чуть-чуть подталкивает их, чтобы loss падал. За много шагов параметры оседают в значениях, схватывающих статистические закономерности данных. Для имён это означает такие вещи как: имена часто начинаются с согласных, «qu» имеет тенденцию появляться вместе, в именах редко идут три согласные подряд и т. д. Модель не учит явных правил, она учит распределение вероятностей, которое случайно их отражает.

Как это связано с ChatGPT? ChatGPT — это тот же базовый цикл (предсказать следующий токен, сэмплировать, повторить), огромно отмасштабированный, с post-training для превращения его в разговорный. Когда вы общаетесь с ним, system prompt, ваше сообщение и его ответ — это всё просто токены в последовательности. Модель достраивает документ по одному токену за раз, точно так же, как microgpt достраивает имя.

Что за история с «галлюцинациями»? Модель генерирует токены, сэмплируя из распределения вероятностей. У неё нет понятия истины, она знает только, какие последовательности статистически правдоподобны при имеющихся обучающих данных. microgpt, «галлюцинирующий» имя вроде «karia», — это то же явление, что и ChatGPT, уверенно заявляющий ложный факт. Оба — это правдоподобно звучащие продолжения, которые случайно не являются настоящими.

Почему это так медленно? microgpt обрабатывает по одному скаляру за раз на чистом Python. Один шаг обучения занимает секунды. Та же математика на GPU обрабатывает миллионы скаляров параллельно и работает на порядки быстрее.

Могу ли я заставить его генерировать имена получше? Да. Тренируйте дольше (увеличьте num_steps), сделайте модель больше (n_embd, n_layer, n_head) или используйте больший датасет. Это те же самые ручки, которые важны и в масштабе.

А что если я поменяю датасет? Модель выучит те паттерны, которые есть в данных. Подсуньте файл с названиями городов, именами покемонов, английскими словами или короткими стихами — и модель научится генерировать их. Остальной код менять не нужно.