Как построить LLM с нуля: пошаговое руководство от Bigram до Transformer

С чего начинается LLM? С биграммы, конечно.

Каждая великая нейросеть когда-то была просто таблицей с парами слов. И это нормально. Когда слышишь "LLM с нуля", кажется, что нужно сразу писать Transformer на 8 GPU. Но на самом деле весь путь — от убогой Bigram до GPT — помещается в один ноутбук. И я это докажу.

Зачем вам это? Чтобы понять, как работают современные монстры, нужно прощупать руками каждый кирпичик. Без этого вы будете повторять чужие архитектуры вслепую. А ещё это чертовски весело — смотреть, как модель учится писать осмысленный бред.

1 Bigram Language Model — минимально живая модель

Биграмма — это модель, которая предсказывает следующий токен на основе ровно одного предыдущего. Звучит убого? Так и есть. Но она работает, и это отличная отправная точка.

Нам нужен датасет. Возьмём тексты Шекспира — классика. Разобьём на символы, построим словарь, закодируем. Потом будем подавать в модель пары (x, y), где x — текущий символ, y — следующий. Всё просто.

Вот как это выглядит на PyTorch:

import torch
import torch.nn as nn
import torch.nn.functional as F

# Гипотетический словарь из 100 символов
vocab_size = 100
embedding_dim = 10

class BigramModel(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        super().__init__()
        # Простая таблица соответствий: токен -> вектор -> логиты
        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
        self.lm_head = nn.Linear(embed_dim, vocab_size)

    def forward(self, idx):
        # idx: (batch, seq_len) — но для биграммы seq_len=1
        x = self.token_embedding(idx)   # (batch, seq_len, embed_dim)
        logits = self.lm_head(x)        # (batch, seq_len, vocab_size)
        return logits

model = BigramModel(vocab_size, embedding_dim)

Потери считаем кросс-энтропией. Обучаем стохастическим градиентным спуском. Уже после пары сотен итераций модель начнёт выдавать осмысленные последовательности — не слова, а хотя бы распределение символов, похожее на текст.

def generate(model, idx, max_new_tokens=100):
    for _ in range(max_new_tokens):
        logits = model(idx)                # предсказание для последнего токена
        probs = F.softmax(logits[:, -1, :], dim=-1)  # берём последний шаг
        idx_next = torch.multinomial(probs, num_samples=1)  # семплируем
        idx = torch.cat((idx, idx_next), dim=1)
    return idx

Результат? Текст типа "eoUMyq" — уже осмысленно по распределению, но ещё не слова. Но мы только начали.

2 N-gram? Нет, RNN — контекст растёт

Один токен — это смешно. Нам нужно учитывать историю произвольной длины. Тут на сцену выходят рекуррентные нейронные сети (RNN). Они проносят скрытое состояние через всю последовательность.

Проблема: RNN страдают от затухающих градиентов. LSTM и GRU отчасти спасают, но не идеально. Тем не менее, для небольших задач — самое то.

Вот минимальная RNN-языковая модель:

class RNNModel(nn.Module):
    def __init__(self, vocab_size, hidden_size, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.rnn = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True)
        self.lm_head = nn.Linear(hidden_size, vocab_size)

    def forward(self, idx, hidden=None):
        emb = self.embedding(idx)
        out, hidden = self.rnn(emb, hidden)
        logits = self.lm_head(out)
        return logits, hidden

Обучается дольше, но результат — уже почти читаемые фразы. На этом этапе можно остановиться и получить неплохую модель для генерации текста в стиле "автодополнение".

Но мы хотим Transformer — идём дальше.

3 Attention с нуля — как заглянуть в любой токен

В RNN каждый шаг зависит только от предыдущего скрытого состояния. Это последовательное узкое горлышко. Attention позволяет модели на каждом шаге смотреть на любые токены из прошлого (и будущего, если нет маски).

Идея: для каждого токена мы вычисляем Query (запрос), Key (ключ) и Value (значение). Скалярное произведение Query и Key показывает, насколько токену нужно "обратить внимание" на другой токен. Полученные веса нормируем softmax и взвешиваем Values.

Реализация простой attention head:

class AttentionHead(nn.Module):
    def __init__(self, d_model, head_dim):
        super().__init__()
        self.q = nn.Linear(d_model, head_dim, bias=False)
        self.k = nn.Linear(d_model, head_dim, bias=False)
        self.v = nn.Linear(d_model, head_dim, bias=False)

    def forward(self, x, mask=None):
        # x: (batch, seq, d_model)
        Q = self.q(x)  # (batch, seq, head_dim)
        K = self.k(x)
        V = self.v(x)

        attn_weights = (Q @ K.transpose(-2, -1)) / (head_dim ** 0.5)  # scale
        if mask is not None:
            attn_weights = attn_weights.masked_fill(mask == 0, float('-inf'))
        attn_weights = F.softmax(attn_weights, dim=-1)
        out = attn_weights @ V
        return out

Ключевой момент — маска внимания. Для авторегрессивных моделей (как GPT) нужна каузальная маска: чтобы токен не видел будущие токены. Иначе модель будет "жульничать".

Соединяем несколько голов — получаем Multi-Head Attention. Каждая голова учится разным паттернам зависимости (синтаксис, семантика, позиция).

4 Transformer: собираем всё в кучу

Теперь у нас есть все ингредиенты: эмбеддинги, self-attention, Feed-Forward сети, layer norm и residual connections. Собираем блок Transformer-decoder (как в GPT).

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Self-attention с residual
        attn_out, _ = self.attention(x, x, x, attn_mask=mask)
        x = self.ln1(x + attn_out)
        # FFN с residual
        ffn_out = self.ffn(x)
        x = self.ln2(x + ffn_out)
        return x

class GPT(nn.Module):
    def __init__(self, vocab_size, d_model=256, n_heads=8, n_layers=6, d_ff=1024, max_seq=512):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.pos_embedding = nn.Embedding(max_seq, d_model)
        self.blocks = nn.Sequential(*[
            TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)
        ])
        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size)

    def forward(self, idx):
        B, T = idx.shape
        positions = torch.arange(T, device=idx.device).unsqueeze(0)
        x = self.token_embedding(idx) + self.pos_embedding(positions)
        # Создаём каузальную маску
        mask = torch.triu(torch.ones(T, T, device=idx.device) * float('-inf'), diagonal=1)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        return logits

Обратите внимание на маску: torch.triu(..., diagonal=1) — верхняя треугольная матрица с -inf, чтобы токен не мог подсматривать вперёд.

Поздравляю, у вас есть свой мини-GPT. Можно обучать его на тексте Шекспира или на GitHub-коде. Потери падают, текст становится почти связным.

5 Обучение и генерация: запускаем зверя

Тренировочный цикл стандартный: батчи, forward, loss, backward, optimizer. Важно правильно масштабировать: начинайте с маленького датасета (100KB), потом переходите на большие.

def train(model, data_loader, optimizer, epochs=10):
    model.train()
    for epoch in range(epochs):
        for x, y in data_loader:
            logits = model(x)
            loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch}: loss = {loss.item():.4f}")

# Генерация: как и раньше, но с передачей всей последовательности
def generate(model, start_tokens, max_new=200, temperature=1.0):
    model.eval()
    idx = start_tokens
    for _ in range(max_new):
        logits = model(idx)[:, -1, :] / temperature
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, idx_next], dim=1)
    return idx

Параметр temperature контролирует случайность: <1 — уверенные предсказания, >1 — хаос. Играйтесь, находите баланс.

Главные ошибки, которые превращают код в тыкву

Напишу список того, что ломало мне ночи:

• Неправильная маска внимания. Если маска не каузальная, модель подглядит правильные ответы и loss быстро упадёт до смешного, но на генерации будет полный бред. Проверяйте: в маске должно быть -inf для будущих позиций.
• Слишком большой learning rate. Трансформеры чувствительны к LR. Используйте планировщик (cosine annealing, warm-up).
• Забыть про dropout. Особенно на малом датасете — переобучение гарантировано.
• Не нормировать логиты при семплировании. Если не разделить на temperature, распределение может быть слишком острым или плоским.
• Думать, что большая модель = хорошо. Начать лучше с d_model=128, n_layers=4, а потом масштабировать. Иначе GPU просто ляжет.

Что дальше? Не торопитесь писать свой GPT-4

Вот вы собрали рабочую модель. Она умеет генерировать текст с логикой на уровне "cat sat on mat". Что теперь?

Во-первых, не нужно сразу бежать закупать 8 H100. Начните с расширения датасета: соберите корпус текстов размером 10-100 МБ. Используйте токенизацию BPE (через HuggingFace Tokenizers). Затем поставьте модель побольше (d_model=512, n_layers=12) и обучите на одной видеокарте — это займёт день-два, но результат будет уже похож на настоящий LLM.

Во-вторых, попробуйте разные архитектурные модификации. Например, заменить обычный FFN на смесь экспертов или добавить механизм управления контекстом как в RLM: как заставить LLM управлять своим контекстом, пока он не сгнил.

В-третьих, не зацикливайтесь только на тексте. Можно дообучить модель для работы с изображениями — почитайте "Архитектура Vision-Language моделей: как дообучить текстовую LLM для работы с изображениями" — это следующий логичный шаг.

И наконец, не верьте, что для быстрого обучения нужны кастомные CUDA ядра. Для прототипа хватит обычного PyTorch. Когда модель вырастет до размеров, где каждая микросекунда на счету — тогда да, стоит задуматься. Но это уже другая история.

Если вы дошли до этого абзаца — вы уже понимаете, как работает современная генеративная магия. Остальное — инженерная работа и набивка шишек. Удачи.

Подписаться на канал