04 · Блок трансформера: один проход от token_id до логитов

🌐 English · Русский · Eesti

Это урок, на котором машина перестаёт быть схемой и становится конкретным, конечным объектом. Внутрь заходит одно число (id токена), происходит тридцать с чем-то операций — каждая из них видна ниже — и наружу выходят 27 чисел. GPT-4 и microGPT различаются здесь только количеством: больше блоков в стопке, шире векторы, больше голов. Форма пути, который ты сейчас проследишь, — та же самая, которую индустрия на триллион долларов прокручивает миллиарды раз в секунду.

Прежде чем нырнуть. Этот урок — сборочный конвейер: вектор токена (его эмбеддинг, список из 16 чисел — 00 · Основы §3) заходит, проходит фиксированную последовательность станций и выходит как логиты — по одной сырой оценке на каждый возможный следующий символ (§5). Каждая станция — это что-то уже знакомое: умножения матриц (скалярные произведения пачкой), рецепт внимания из урока 03, сложения и две маленькие новинки. RMSNorm перемасштабирует вектор, чтобы его числа держались здорового общего размера — нормализация, и ничего больше. ReLU — простейшая кривая в глубоком обучении: отрицательное становится 0, положительное проходит как есть. Не заучивай порядок станций — карта тебе и есть песочница.

Теория

Уроки 01–03 смотрели издалека (цикл прямой проход → потери → сэмпл) и вблизи (автоград, внимание). Этот урок собирает деталь посередине: единственную функцию gpt(), которая превращает одну пару (token_id, pos_id) в вектор логитов по следующему символу. Для microGPT это ровно один блок трансформера (n_layer = 1), n_embd = 16, n_head = 4, head_dim = 4.

Путь данных, в том порядке, в котором его выполняет эталонный код:

Эмбеддинг — взять строку токена из wte и строку позиции из wpe, сложить их поэлементно → вектор x длины 16.
RMSNorm (начальный) — Карпатый применяет rmsnorm один раз прямо здесь, до блока. Его легко не заметить, и он выглядит лишним рядом с нормализацией внутри блока, но комментарий автора однозначен: «not redundant due to backward pass via the residual connection» (не избыточен из-за обратного прохода через остаточную связь). Он меняет то, что несёт остаточная ветка.
Подблок внимания (pre-norm + остаточная связь):
- сохранить x_residual = x (ветка ①),
- применить rmsnorm к копии,
- многоголовое внимание — те же q·kᵀ/√head_dim → софтмакс → ·v из урока 03, по каждой голове с конкатенацией,
- спроецировать через attn_wo,
- прибавить обратно сохранённую ветку ①.
Подблок MLP (pre-norm + остаточная связь):
- сохранить x_residual = x (ветка ②),
- применить rmsnorm к копии,
- mlp_fc1 расширяет 16 → 64,
- ReLU (не GeLU — в эталоне max(0, x)),
- mlp_fc2 сжимает 64 → 16,
- прибавить обратно сохранённую ветку ②.
LM Head — linear(x, lm_head) проецирует финальный 16-вектор в один логит на каждый токен словаря. Никакой финальной нормализации перед lm_head в эталоне нет.

Это весь блок. Чего в microGPT сознательно нет — и чего поэтому нет и в песочнице: LayerNorm (используется RMSNorm), GeLU (используется ReLU), dropout и смещений (bias) в любых линейных слоях.

Замечание о двух остаточных связях: это pre-norm трансформер. Каждый подблок нормализует копию x, прогоняет свой подслой и прибавляет результат обратно к ненормализованному x, который сохранил. Этот сохранённый-и-прибавленный обход нарисован двумя дугами справа в сцене.

Инициализация параметров

Прежде чем случится хоть один прямой проход, веса должны откуда-то взяться. microGPT строит их один раз — как обычные гауссовские случайные скаляры, обёрнутые в Value (строки py 99–114):


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]

Каждая матрица — nout × nin значений random.gauss(0, 0.08): обычное нормальное распределение со стандартным отклонением 0.08. Это вся инициализация: ни Xavier, ни Kaiming, ни особого масштабирования. При n_embd = 16, block_size = 16, n_head = 4 и vocab_size = len(uchars) + 1 в state_dict лежат:

матрица	форма (nout × nin)	роль
`wte`	vocab_size × 16	таблица эмбеддингов токенов
`wpe`	16 × 16	таблица эмбеддингов позиций (block_size × n_embd)
`attn_wq` / `wk` / `wv` / `wo`	16 × 16 каждая	проекции Q/K/V и выходная проекция, на слой
`mlp_fc1`	64 × 16	расширяющая проекция MLP (4·n_embd × n_embd)
`mlp_fc2`	16 × 64	сжимающая проекция MLP (n_embd × 4·n_embd)
`lm_head`	vocab_size × 16	финальная проекция в логиты

linear(x, w) читает каждую матрицу весов как [nout][nin], так что выход j — это скалярное произведение w[j] со входом. Наконец, params разворачивает каждый скаляр каждой матрицы в один плоский список — ровно по нему ходит цикл Adam из 05 · Обучение и генерация: по одному буферу m/v и одному обновлению на скаляр, каждый шаг.

Аннотированный код

Блок живёт в src/microgpt_annotated.py, подраздел attention-multihead (помощники linear / softmax / rmsnorm — в overview-pipeline-helpers):


def gpt(token_id, pos_id, keys, values):
    tok_emb = state_dict['wte'][token_id]          # token embedding
    pos_emb = state_dict['wpe'][pos_id]            # position embedding
    x = [t + p for t, p in zip(tok_emb, pos_emb)]  # joint embedding
    x = rmsnorm(x)  # note: not redundant due to backward pass via the residual connection
 
    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

Порт на TypeScript в src/inference/model.ts считает тот же путь. Различие чисто механическое: Python вызывает gpt() по разу на позицию с растущим KV-кэшем, а порт берёт всю последовательность сразу и применяет явную причинную маску j ≤ i — та же математика, другой поток управления (та же мысль, которую урок 03 проговаривает про внимание).

Песочница

Каждый модуль на пути — блок, по которому можно кликнуть и увидеть его форму входа → выхода и точную строку Python, которую он выполняет. Нажми play (или мотай вручную), чтобы отправить импульс данных по пути; две зелёные дуги — остаточные обходы (сохраняются в ①/② и прибавляются обратно на соответствующих станциях Add). Станция внимания сжато показывает то же вычисление, что разобрано в уроке 03, — этот урок про то, где именно внимание сидит внутри целого блока. Это карта структуры блока и порядка выполнения, а не постадийный инспектор реальных значений тензоров (показанные формы — это статические размерности слоёв; за настоящими числами внимания — в урок 03).

Попробуй сам.

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

Кликни станцию mlp_fc1 и проверь её форму: 16 на входе, 64 на выходе. Затем mlp_fc2: 64 на входе, 16 на выходе. Блок дышит — расшириться, подумать, сжаться.

Найди станцию, чью строку Python ты теперь смог бы написать по памяти. (Кандидат: сложение эмбеддингов — от него один zip до сложения векторов из урока 00.)

По таблице параметров выше прикинь, какая доля всех ~4 000 ручек живёт в таблицах эмбеддингов и lm_head, а какая — внутри самого блока. Удивился?