Привет, Хабр! Меня зовут Владимир, и это третья часть цикла статей по написании и обучению небольшой decoder-only LLM с нуля. В первой части мы обучили токенизатор и собрали pretrain-датасет, во второй части реализовали класс Трансформер-блока. В этой части будем собирать и предобучать нашу LLM.

Содержание цикла

1. Подготовка и токенизация данных

2. Трансформер

3. Сборка и обучение LLM (вы находитесь здесь)

4. SFT этап (дообучение на Вопрос-Ответ)

5. Интеграция с Hugging Face

Содержание может меняться и дополняться ссылками по мере написания

Но для начала немного подробнее разберём decoder-only архитектуру.

Decoder-only архитектура

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

Схематично путь данных внутри модели такой:

1. Токены превращаются в векторы (эмбеддинги).

2. К эмбеддингам токенов добавляется информация о позиции (чтобы модель учитывала позиции токенов относительно друг друга).

3. Последовательность проходит через несколько Transformer-блоков.

4. Финальный слой превращает внутреннее представление в последовательность логитов.

Кратко про Embedding-слой

Ещё раз уточню, что модель не “понимает слова” напрямую. Для неё токен - это просто число. Embedding-слой превращает это число в вектор, с которым уже можно работать математически. В процессе обучения модель подбирает такие векторы токенов, чтобы по ним было проще предсказывать продолжение текста: токены, часто встречающиеся в похожих контекстах, получают похожие представления, а важные для смысла различия сохраняются в разных направлениях этого векторного пространства.

Позиционные эмбеддинги

В отличие от RNN, где токены подаются в сеть один за другим, трансформер получает все токены одновременно, и сам по себе механизм внимания не знает порядка токенов. Для него последовательность без дополнительных признаков похожа на набор элементов: он может сравнивать токены друг с другом, но не понимает, какой был первым, какой вторым, а какой десятым. Например, фразы “кот укусил пса” и “пёс укусил кота” состоят из одного и того же набора слов, но смысл у них противоположный. Если модель не знает порядок токенов, ей трудно понять, кто совершил действие, а кто оказался объектом действия.

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

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

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

Синусоидальные (фиксированные) позиционные кодировки

Данный вид кодировок был показан в оригинальной статье “Attention Is All You Need” (2017). Положение кодируется аналитически через sin/cos. Это дёшево по памяти и вычислениям (с точки зрения затрат на обучение / инференс) и позволяет экстраполировать на длинные последовательности, но жёсткая функциональная форма ограничивает “выразительность” - модель не может адаптировать позиционное представление под данные в процессе обучения.

Обучаемые абсолютные позиционные эмбеддинги

Простейший и широко используемый подход: таблица nn.Embedding(max_len, d_model), обучаемая вместе с моделью. Обладает высокой выразительностью. Существенное ограничение данного способа - нельзя напрямую обрабатывать последовательности длиннее max_len.

Относительные позиционные кодировки

В этом подходе модели сообщают не “этот токен стоит на позиции 10”, а “этот токен находится на 3 позиции левее/правее другого токена”. То есть важен не номер места в строке, а расстояние между токенами. Это полезно, потому что многие языковые связи зависят именно от взаимного расположения слов: например, соседние слова обычно сильнее связаны друг с другом, чем слова через несколько абзацев.

Rotary Position Embedding (RoPE)

Данный метод используется во многих современных моделях: LLaMA, Qwen, Mistral. В отличие от обычных позиционных эмбеддингов, здесь позиция не прибавляется к вектору токена отдельной таблицей. Вместо этого модель “поворачивает” query- и key-векторы на угол, зависящий от позиции токена. Для этого query- и key-векторы разбивают на пары координат (вида (x0, x1), (x2, x3) и т.д) и каждую пару рассматривают как точку на плоскости. Для токена на позиции 0 поворот один, для позиции 1 - чуть другой, для позиции 10 - ещё сильнее. В результате один и тот же токен на разных позициях получает разные query/key-представления.

Главный плюс RoPE проявляется в момент, когда токены попадают в механизм внимания. Attention считает скалярное произведение между query одного токена и key другого: чем оно больше, тем сильнее один токен “смотрит” на другой. А скалярное произведение напрямую зависит от угла между ними. RoPE заранее поворачивает query/key-векторы в зависимости от их позиций, поэтому угол между ними автоматически несёт информацию о расстоянии между токенами.

ALiBi (Attention with Linear Biases)

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

Ну и перед тем, как перейти к реализации модели, рассмотрим ещё один вопрос.

Зачем много Transformer-блоков подряд

Блок Transformer делает две главные вещи: сначала attention-часть позволяет каждому токену посмотреть на предыдущие токены и собрать из них полезный контекст, затем feed-forward сеть независимо обрабатывает представление каждого токена. Если attention отвечает за обмен информацией между позициями, то feed-forward часть отвечает за нелинейную переработку уже собранной информации.

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

Ну всё, с теорией закончили, переходим к torch-реализации.

Реализация LLM

Хоть может показаться, что самый сложный этап ванильной GPT-like LLM - это написание MHA класса, со всеми его преобразованиями и транспонированиями размерностей, это не так. Самый сложный этап вот сейчас - как назвать свой класс LLM. Ведь под этим именем ей (модели) существовать в дальнейшем… Да я детям быстрее с именем определился

Итак, встречайте: Lingua Laboratorium Mechanicus или сокращённо LLM:

import torch.nn as nn
from .transformer import TransformerBlock

class LinguaLaboratoriumMechanicus(nn.Module):
    def __init__(
        self, vocab_size: int, emb_dim: int, n_layers: int, n_heads: int, max_context_length: int = 1024, dropout: float = 0.1, qkv_bias: bool = False
    ):
        super().__init__()

        cfg = {
            'emb_dim': emb_dim, 'n_heads': n_heads, 'context_length': max_context_length, 'dropout': dropout, 'qkv_bias': qkv_bias
        }

        self.vocab_size = vocab_size
        self.emb_dim = emb_dim
        self.max_context_length = max_context_length

        self.token_emb = nn.Embedding(vocab_size, emb_dim)
        self.pos_emb = nn.Embedding(max_context_length, emb_dim)
        self.drop_emb = nn.Dropout(dropout)

        self.blocks = nn.Sequential(*[TransformerBlock(**cfg) for _ in range(n_layers)])

        self.final_norm = nn.LayerNorm(emb_dim)
        self.out_head = nn.Linear(emb_dim, vocab_size, bias=False)

        self.apply(self._init_weights)

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

• token_emb - переводит id токенов из словаря размера vocab_size во внутреннее векторное пространство размерности emb_dim;

• pos_emb - хранит обучаемые позиционные эмбеддинги для позиций от 0 до max_context_length - 1;

• drop_emb - применяет dropout к сумме token embedding и positional embedding;

• blocks - последовательность из n_layers Transformer-блоков;

• final_norm - нормализует итоговые скрытые представления;

• out_head - переводит внутренние векторы размерности emb_dim в логиты по всему словарю.

Перед обучением модель нужно инициализировать. Метод _init_weights проходит по линейным слоям и embedding-таблицам и заполняет их веса маленькими случайными числами из нормального распределения. Это даёт модели нейтральную, но не полностью одинаковую стартовую точку: разные веса смогут обучаться по-разному, а небольшое начальное значение весов помогает избежать нестабильных активаций в начале обучения. Bias у линейных слоёв обнуляется - для него случайная инициализация обычно не нужна:

class LinguaLaboratoriumMechanicus(nn.Module):
    # Предыдущий код
    @staticmethod
    def _init_weights(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                torch.nn.init.zeros_(module.bias)

После создания слоёв остаётся описать прямой проход модели - то есть что происходит с батчем токенов, когда мы передаём его в LLM. На вход метод forward получает input_ids - тензор с токенами формы (batch_size, n_tokens). Для каждого токена берутся token embedding и positional embedding, складываются и проходят через dropout. После этого получившаяся последовательность векторов прогоняется через стек Transformer-блоков, нормализуется и отправляется в выходной слой.

class LinguaLaboratoriumMechanicus(nn.Module):
    # Предыдущий код
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        _, n_tokens = input_ids.size()
        if n_tokens > self.max_context_length:
            raise ValueError(f'Длина входной последовательности ({n_tokens}) превышает ' f'максимальную заданную ({self.max_context_length}).')
        x = self.drop_emb(
            self.token_emb(input_ids) +
            self.pos_emb(torch.arange(n_tokens, device=input_ids.device).unsqueeze(0))
        )
        x = self.blocks(x)
        return self.out_head(self.final_norm(x))

На выходе получается тензор логитов формы (batch_size, n_tokens, vocab_size). Для каждой позиции в последовательности модель выдаёт по одному числу на каждый токен словаря.

Перед тем, как переходить к функции обучения, необходимо реализовать ещё одну - функцию генерации текста.

Наша свеженаписанная LLM (да и любая другая) за один проход генерирует на выходе N токенов (вообще N распределений логитов, но для упрощения пойдёт) при N входных токенах, и на этом считает свою работу выполненной. Для продолжения генерации текста (мы же не Эллочку-людоедочку хотим) необходимо добавлять последний сгенерированный токен (так как формально только он новый и нужный) ко входной последовательности и снова подавать на вход модели для генерации ещё одного токена. Повторять данный процесс надо или до генерации спец токена EOS или до достижения определённого порога генерации. Собственно, этим процессом и будет заниматься метод генерации:

import torch
import torch.nn.functional as F
from transformers import PreTrainedTokenizerFast

@torch.no_grad()
def generate(
    model: torch.nn.Module, tokenizer: PreTrainedTokenizerFast, prompt: str, max_new_tokens: int = 100, temperature: float = 1.0,
    top_k: int | None = None, eos_token_id: int | None = None, device: str = 'cpu',
) -> str:
    model.eval()
    input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
    prompt_length = input_ids.size(1)
    generated = input_ids.clone()
    if eos_token_id is None:
        eos_token_id = tokenizer.eos_token_id

    for _ in range(max_new_tokens):
        if generated.size(1) > model.max_context_length:
            generated = generated[:, -model.max_context_length:]

        logits = model(generated)
        logits_last = logits[0, -1, :]
        if temperature >= 1e-5:
            logits_last = logits_last / temperature
            logits_last = top_k_filtering(logits_last.unsqueeze(0), top_k).squeeze(0)
            probs = F.softmax(logits_last, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1).unsqueeze(0)
        else:
            next_id = torch.argmax(logits_last, dim=-1, keepdim=True).unsqueeze(0)
        generated = torch.cat([generated, next_id], dim=-1)
        input_ids = torch.cat([input_ids, next_id], dim=-1)

        if next_id.item() == eos_token_id:
            break

    model.train()
    return tokenizer.decode(input_ids[0, prompt_length:], skip_special_tokens=True)

Метод начинаем с декоратора, отключающего вычисление градиентов, и перевода модели в режим инференса (для отключения слоёв dropout). Токенизируем входной промпт (запрос пользователя, по-русски) и полного копирования токенизированного запроса (надо для исключения входного промпта из ответа).

В цикле генерации проверяем, что нагенерированная последовательность не превышает наше контекстное окно и, при необходимости, обрезаем начало последовательности. Далее вызываем метод forward модели и забираем с выхода логиты последнего токена (помним, что выходная размерность (batch_size, n_tokens, vocab_size) а промпт - одна строка, поэтому [0, -1, :]).

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

• temperature управляет “резкостью” распределения вероятностей. При низкой температуре модель чаще выбирает самые вероятные токены. При высокой температуре распределение сглаживается, и у менее вероятных вариантов появляется больше шансов попасть в ответ. Для получения такого результата логиты просто делят на температурный коэффициент.

• top_k ограничивает выбор только k самыми вероятными токенами.

• top_p работает немного иначе: он оставляет не фиксированное количество токенов, а минимальный набор самых вероятных токенов, суммарная вероятность которых достигает порога p. Например, при top_p = 0.9 модель выбирает только из тех токенов, которые вместе покрывают 90% вероятностной массы.

Вернёмся к коду. В своей реализации я оставил только temperature и top_k:

# Повтор кода
if temperature >= 1e-5:
    logits_last = logits_last / temperature
    logits_last = top_k_filtering(logits_last.unsqueeze(0), top_k).squeeze(0)
    probs = F.softmax(logits_last, dim=-1)
    next_id = torch.multinomial(probs, num_samples=1).unsqueeze(0)
else:
    next_id = torch.argmax(logits_last, dim=-1, keepdim=True).unsqueeze(0)

При положительной температуре - поделили на неё, выбрали top_k логитов, перевели их в вероятности с помощью F.softmax и выбираем токен с помощью взвешенной случайной выборки (torch.multinomial). Если температура околонулевая, то просто берём токен с максимальным логитом. В обоих случаях unsqueeze(0) добавляет в начало тензора измерение для дальнейшей конкатенации.

Реализация top_k_filtering приведена ниже:

def top_k_filtering(logits: torch.Tensor, top_k: int | None) -> torch.Tensor:
    if top_k is None or top_k <= 0:
        return logits
    threshold = torch.topk(logits, min(top_k, logits.size(-1)), dim=-1).values[:, -1].unsqueeze(-1)
    return logits.masked_fill(logits < threshold, float('-inf'))

После проверки на невалидное значение параметра мы вычисляем порог сравнения для логитов (минимально допустимое значение). Для этого методом torch.topk мы выбираем k наибольших значений в тензоре logits в последнем измерении - после logits[0, -1, :] в методе генерации имеем тензор размерности (vocab_size,), а вызов logits_last.unsqueeze(0) дополнит размерность до (1, vocab_size). Метод .values[:, -1] выбирает нам наименьший из k логитов (получаем размерность (batch_size, )), а .unsqueeze(-1) добавит в конец тензора дополнительное измерение, необходимое для дальнейшего сравнения. min(top_k, logits.size(-1)) нужен на случай, если заданный top_k больше длины тензора.

С помощью masked_fill мы заменяем все логиты, которые меньше нашего порога, на минус -inf (это обеспечит нулевую вероятность токена после операции softmax).

Напомню, что там в генерации после выбора токена:

@torch.no_grad()
def generate(...) -> str:
    # Повтор хвоста метода
        generated = torch.cat([generated, next_id], dim=-1)
        input_ids = torch.cat([input_ids, next_id], dim=-1)

        if next_id.item() == eos_token_id:
            break

    model.train()
    return tokenizer.decode(input_ids[0, prompt_length:], skip_special_tokens=True)

Токен добавляем в массив generated (рабочая копия, которая в LLM подаётся) и input_ids (массив с входными и выходными токенами). Если встретили EOS, то прерываем генерацию, иначе продолжаем. Возвращаем из метода детокенизированный текст с отрезанной входной последовательностью.

Теперь можно перейти к методу тренировки.

Тренировочный цикл

Для хоть какой-то организации, все параметры модели и тренировки я решил собрать в @dataclass. Получилось 22 скучных параметра, поэтому код я спрятал

@dataclass
class Config:
    tokenizer_path: str = 'tokenizer/tokenizer_config'
    json_data_dir: str = 'dataset/json_data'
    save_dir: str = 'checkpoints'
    vocab_size: int = 50257
    emb_dim: int = 768
    n_layers: int = 12
    n_heads: int = 12
    max_context_length: int = 1024
    dropout: float = 0.1
    qkv_bias: bool = False
    batch_size: int = 4
    lr: float = 3e-4
    max_epochs: int = 10
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
    eval_prompts: tuple[str] = ('Кто такой Император?', )
    eval_max_new_tokens: int = 100
    eval_temperature: float = 0.8
    eval_top_k: int = 40
    force_reprocess: bool = False
    warmup_steps: int = 500
    grad_clip: float = 1.0
    hub_repo_id: str | None = None

Едем дальше. Тренируются обычно эпохами, поэтому для упорядочивания кода напишем метод обучающего прогона одной эпохи:

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.utils.data import DataLoader

def train_one_epoch(
    model: torch.nn.Module, dataloader: DataLoader, optimizer: AdamW,
    scheduler, vocab_size: int, device: str, grad_clip: float,
) -> list[float]:
    model.train()
    losses: list[float] = []

    for inputs, targets in tqdm(dataloader, desc='Train', leave=False):
        inputs = inputs.to(device)
        targets = targets.to(device)

        logits = model(inputs)
        loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))

        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)

        optimizer.step()
        optimizer.zero_grad()

        if scheduler is not None:
            scheduler.step()

        losses.append(loss.item())

    return losses

Начало цикла обучения почти torch-каноничное - каждый батч переводим на тот же device, на котором находится модель, делаем прямой проход, считаем loss и выполняем обратный проход (расчёт градиентов). Далее небольшая оптимизация - ограничение нормы градиентов, чтобы предотвратить “взрыв”. Далее идет пересчёт весов модели, обнуление градиентов и пересчёт LR с помощью шедулера. Лоссы копим для истории.

Теперь основной train loop:

def train(cfg: Config) -> None:
    save_dir = Path(cfg.save_dir)
    save_dir.mkdir(exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_path)

    dataset = W40kDataset(
        tokenizer=tokenizer, json_path=cfg.json_data_dir, max_length=cfg.max_context_length, force_reprocess=cfg.force_reprocess)
    dataloader = DataLoader(dataset, batch_size=cfg.batch_size, shuffle=True)

    model = LinguaLaboratoriumMechanicus(
        vocab_size=cfg.vocab_size, emb_dim=cfg.emb_dim, n_layers=cfg.n_layers, n_heads=cfg.n_heads,
        max_context_length=cfg.max_context_length, dropout=cfg.dropout, qkv_bias=cfg.qkv_bias,
    ).to(cfg.device)

    optimizer = AdamW(model.parameters(), lr=cfg.lr)
    scheduler = get_cosine_schedule_with_warmup(
        optimizer=optimizer, num_warmup_steps=cfg.warmup_steps, num_training_steps=len(dataloader) * cfg.max_epochs)

    for epoch in range(cfg.max_epochs):
        start = time.perf_counter()
        losses = train_one_epoch(
            model=model, dataloader=dataloader, optimizer=optimizer, scheduler=scheduler, vocab_size=cfg.vocab_size, 
            device=cfg.device, grad_clip=cfg.grad_clip)
        elapsed = time.perf_counter() - start

        cur_lr = scheduler.get_last_lr()[0]
        avg_loss = sum(losses) / len(losses)

        torch.save(losses, save_dir / f'loss_epoch{epoch + 1:02d}.pt')
        print(f'\n=== Эпоха {epoch + 1}/{cfg.max_epochs} [{elapsed:.1f}s] ===')
        print(f'AvgLoss: {avg_loss:.4f}  LR: {cur_lr:.2e}')

        run_eval(model, tokenizer, cfg)

        torch.save({
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'epoch': epoch, 'loss': avg_loss,
        }, save_dir / f'checkpoint_epoch{epoch + 1:02d}.pt')

    print('Обучение завершено.')

Наверное, кроме шедулера, ничего сверхъестественного нет. Загружаем наш кастомный токенизатор, создаём датасет, преобразуем его в torch.utils.data.DataLoader. Далее создаём наш mini-GPT с 12 головами, с внутренней размерностью 768, 12 Transformer-блоками и длиной контекста в 1024 токена. Для оптимизации выбран AdamW - стандартный оптимизатор для обучения и дообучения трансформеров. Да и читал где-то, что в любой непонятной ситуации используй AdamW.

Помимо оптимизатора используется шедулер - планировщик learning rate. Я использую get_cosine_schedule_with_warmup: сначала короткий “разогрев” (линейный рост от нуля до заданного порога cfg.lr в течение cfg.warmup_steps шагов), потом плавное уменьшение по косинусной кривой до нуля к последнему шагу обучения. Это нужно, чтобы в начале обучения, когда веса ещё почти случайные, не делать слишком большие шаги и не разогнать loss. Общее число шагов считается как len(dataloader) * cfg.max_epochs, то есть сколько батчей пройдёт модель за все эпохи (scheduler.step() вызывается на каждом батче).

Собственно, всё. Запускаем цикл по количеству эпох, прогоняем train_one_epoch, логгируем loss и LR в конце эпохи, сохраняем лоссы и чекпоинт.

Перед запуском цикла обучения я решил немного посчитать теоретическое использование памяти.

1. Блок MHA содержит слой QKV матриц (emb_dim x 3·emb_dim = 1769K параметров) и выходной слой (emb_dim x emb_dim = 590K параметров). Это примерно 2359 тысяч параметров.

2. Блок FF содержит два слоя размерами emb_dim x 4·emb_dim. Это примерно 4718 тысяч параметров.

3. Размер блока трансформера (пренебрегая копейками нормализации) получается 7077 тысяч параметров.

4. Сама LLM состоит из одинаковых по размеру входа-выхода (vocab_size x emb_dim = 38.6М) и 12-ти блоков трансформера. Итого около 163М параметров.

Каждый параметр - это 4 байта, соответственно сама модель занимает в памяти всего 0.6 ГБ. Фигня вроде, но градиенты на каждом шаге занимают ещё столько же, а оптимизатор - в два раза больше. Получаем уже 2.5-3ГБ. Уже серьёзнее. Но это только постоянная часть, которая всегда “висит” в памяти. А есть ещё временные коэффициенты - активации, которые создаёт PyTorch на каждом шагу для вычисления градиентов.

В статье Reducing Activation Recomputation in Large Transformer Models память активаций одного Transformer-блока оценивается формулой:

S × B × H × (34 + 5 × A × S / H)

где S — длина последовательности, B — размер батча, H — внутренняя размерность, а A — количество голов внимания.

Для нашей модели (S=1024, B=1, H=768, A=12) получается около 85 МБ на один блок, или примерно 1 ГБ на все 12 блоков при вычислениях в float16. Модель обучается в float32, поэтому, грубо говоря, оценка удваивается - около 2 ГБ активаций на одну последовательность. Активации растут почти линейно вместе с размером батча: для батча из восьми последовательностей это уже около 16 ГБ.

А теперь картинка из практики:

Памяти занято около 25 ГБ.

Сам процесс обучения при этом около 7 часов (около 40 минут на эпоху). Пруфов в виде скриншотов, к сожалению, не сохранил.

Если попробовать оценить время одной эпохи чисто теоретически, то можно воспользоваться формулой 6 × N × T, где N — количество параметров модели, а T — количество обработанных токенов, а 6 - коэффициент, который примерно учитывает прямой и обратный проходы. Для модели на 163М параметров и датасета на 71М токенов получим, что нужно около 69 PFLOPs вычислительных ресурсов.

Моя RTX 5090 имеет на борту около 105 TFLOPS, так что теоретически обучение на ней должно было занять 69 PFLOPs / 105 TFLOPs = 657 секунд (около 11 минут). Но это прям сферическая цифра в вакууме. На практике же нужно время на обращения к памяти, создание тензоров, вычисления softmax, запуск CUDA-операций и прочие накладные расходы PyTorch.

Посмотрим теперь, как менялся loss за эти десять эпох:

На графике серым показан loss каждого отдельного батча, а оранжевым - скользящее среднее по 200 шагам. Пунктирная горизонтальная линия соответствует среднему loss последней эпохи - около 3. В начале loss быстро падает: модель осваивает базовую структуру текста и распределение токенов. Затем снижение постепенно замедляется, а к последним эпохам график выходит на плато. Это означает, что при выбранных данных, архитектуре и параметрах обучения дальнейшие эпохи уже вряд ли дали бы существенное улучшение.

В результате обучения получилась базовая модель, которая генерирует более-менее связанный текст. Недостаток данных для обучения (по законам масштабирования для модели 163М необходимо около 3.2 миллиарда токенов против моих 72М) сказывается. Как я писал под спойлером в первой статье

Законы масштабирования не обманешь, получилось не очень - текст генерит, иногда даже связанный. Но в целом - это не модель, которая живёт в мире Империума и Варпа.

На этом на сегодня всё. В следующей статье рассмотрим SFT-этап - дообучение модели на вопрос-ответ с целью обучения её вести диалог.

Код обучения тут

Базовую модель можно попробовать использовать отсюда