Tiny Universe, часть 1: Tiny-LLM — маленькая LLaMA2 с нуля

Полный перевод раздела TinyLLM из tiny-universe: обучаем свой токенизатор на TinyStories, готовим данные, тренируем LLaMA2-подобную модель на ~8,5 млн параметров (около 2 ГБ видеопамяти) и генерируем истории.

Tiny-LLM

Цель этого проекта — реализовать простую большую языковую модель: от обучения токенизатора к обучению модели и далее к генерации текста моделью. Простую LLM можно обучить, используя только NumPy и PyTorch; видеопамяти уходит около 2 ГБ. Ниже — демонстрация результата.

Рисунок в оригинале (images/model_show.png): снимок терминала, в котором обученная модель дописывает английскую детскую историю по началу-подсказке.

Ресурсов для обучения модели тоже нужно совсем немного: хватит одной видеокарты, видеопамяти уходит около 2 ГБ. Обучение недолгое — несколько часов.

Рисунок в оригинале (images/nvidia.png): вывод nvidia-smi во время обучения — процесс занимает около 2 ГБ видеопамяти.

Как пользоваться (Usage)

Пользоваться репозиторием просто: сначала установите PyTorch с поддержкой CUDA, затем — зависимости.

pip install -r requirements.txt

Потом обучайте по шагам:

Обучить токенизатор: python train_vocab.py --download True --vocab_size 4096
Предобработать данные: python preprocess.py
Обучить модель: python train.py
Сгенерировать текст моделью: python sample.py --prompt "One day, Lily met a Shoggoth"

Пример того, что генерирует модель (текст оригинала — модель обучена на английских сказках, поэтому пишет по-английски; в скобках — наш перевод первого абзаца):

One day, Lily met a Shoggoth and her mom said, “Some people borrowed, looking for seeds.” Lily wanted to help, so she decided to help. They walked to the store and bought seeds for the village and put it in a big bowl. As in the future, the children soon scattered the seeds into the ground and the bird wore yellow paper. Lily added some yummy tastes like celery and watched the plant grow until it grew and grew. After that, Lily and the man became good friends and always shared their food. They always got what they had and never brought away seeds to something from the lonely garden. Once upon a time, there was a toy, a big, powerful animal. It was so big that it could not fit in its tree. The little animal was very scared because it wanted to be into the powerful animal to see the angel. A little squirrel saw the pistol and wanted to help. The squirrel knew how to be safe, so he decided to sing like a tall tree. The little animal sang his song. But the giant was not nice, and he was not nice.

(«Однажды Лили встретила Шоггота, и мама сказала: „Кое-кто одолжил, ищут семена“. Лили захотела помочь…» Грамматика почти безупречна, смысл плывёт — типичное поведение крошечной модели.)

Разбор (Blog)

Шаг 1. Обучаем токенизатор

Сначала нужно обучить для обработки текста токенизатор (Tokenizer). Его задача — превращать текст в последовательность чисел, с которой модель может работать. Датасет у нас — TinyStories: небольшой набор коротких историй, сгенерированных GPT-3.5 и GPT-4, с ограниченным словарём. В этой задаче мы используем символьный токенизатор, который отображает каждый символ текста в соответствующее число. Скачать датасет и обучить токенизатор можно такой командой:

python train_vocab.py --download True --vocab_size 4096

Размер словаря LLaMA2 — 32 000, но датасет TinyStories маленький и словарь в нём ограниченный, поэтому мы задаём размер словаря 4096. После обучения токенизатор умеет превращать текст в последовательность чисел и восстанавливать текст из последовательности чисел.

def download_file(url: str, fname: str, chunk_size=1024):
    """Отправляет HTTP GET-запрос и скачивает файл потоком"""
    ···

def download():
    """Вызывает download_file, чтобы скачать датасет"""
    ···

def train_vocab(vocab_size: int=32000, num_shards: int=20):
    """
    vocab_size: int, размер словаря — определяет, сколько токенов знает токенизатор.
    num_shards: int, сколько фрагментов (шардов) данных обработать — ускоряет обучение словаря.
    """
    # Проверяем, что размер словаря положительный
    assert vocab_size > 0, "Vocab size must be positive"

    # Префикс пути модели SentencePiece — по нему сохраним токенизатор
    prefix = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}")

    # 1) Сливаем текст из нескольких шардов в один текстовый файл tiny.txt
    tiny_file = os.path.join(DATA_CACHE_DIR, "tiny.txt")
    data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
    shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))

    # Создаём tiny.txt и записываем в него тексты из заданного числа шардов
    print(f"Writing temporary file {tiny_file} with {num_shards} shards...")
    with open(tiny_file, "w", encoding="utf-8") as of:
        # Перебираем первые num_shards шардов
        for shard in tqdm(shard_filenames[:num_shards]):
            with open(shard, "r") as f:
                data = json.load(f)  # читаем JSON из шарда
            # Для каждого примера записываем текст истории в tiny.txt
            for example in data:
                text = example["story"]
                text = text.strip()  # убираем пробелы по краям
                of.write(text + "\n")  # каждая история — одна строка

    # Печатаем размер получившегося tiny.txt
    print(f"Size is: {os.path.getsize(tiny_file) / 1024 / 1024:.2f} MB")

    # 2) Обучаем токенизатор с помощью SentencePiece
    print("Will now train the vocab...")
    spm.SentencePieceTrainer.train(
        input=tiny_file,         # входной файл — созданный выше tiny.txt
        model_prefix=prefix,     # префикс пути модели
        model_type="bpe",        # обучаем токенизатор методом Byte-Pair Encoding (BPE)
        vocab_size=vocab_size,   # размер словаря
        self_test_sample_size=0, # размер выборки для самопроверки — 0
        input_format="text",     # формат входа — обычный текст
        character_coverage=1.0,  # покрыть все символы (включая редкие)
        num_threads=os.cpu_count(),  # число потоков = числу ядер CPU
        split_digits=True,       # разбивать числа на отдельные цифры
        allow_whitespace_only_pieces=True,  # разрешить токены только из пробелов
        byte_fallback=True,      # включить запасной байтовый уровень
        unk_surface=r" \342\201\207 ",  # как показывать неизвестный символ (UNK)
        normalization_rule_name="identity"  # правило нормализации «identity» (текст не меняем)
    )

    # 3) Необязательная уборка: спрашиваем, удалить ли временный tiny.txt
    dec = input(f"Delete the temporary file {tiny_file}? [y/N] ")
    if dec.lower() == "y":
        os.remove(tiny_file)  # удаляем временный файл
        print(f"Deleted {tiny_file}")

    # Печатаем, куда сохранена модель
    print(f"Trained tokenizer is in {prefix}.model")
    print("Done.")

В этой части мы обучаем свой токенизатор (Tokenizer) библиотекой SentencePiece. Сначала нужно извлечь текст из датасета TinyStories — это входные данные для обучения. SentencePiece — алгоритм токенизации на основе подсловных единиц (subword), он хорошо справляется с дроблением слов в разных языках.

При обучении SentencePiece сам создаёт два файла: tok4096.model и tok4096.vocab. tok4096.model — обученная модель токенизатора, она лежит в папке data. С её помощью текст превращается в последовательность токенов (Token), а последовательность токенов — обратно в текст.

Чтобы пользоваться токенизатором было удобнее, в файле tokenizer.py мы определили класс Tokenizer. Он оборачивает частые операции — кодирование и декодирование текста — и умеет загружать обученный файл модели. С этим классом легко превратить текст в последовательность чисел, которую понимает модель, или превратить предсказания модели в читаемый текст.

Подробности реализации — в файле tokenizer.py; дальше покажем, как этим классом обрабатывать истории из TinyStories.

class Tokenizer:
    def __init__(self, tokenizer_model=None):
        """
        Инициализация токенизатора: загружаем обученную модель SentencePiece и запоминаем id специальных токенов.

        Параметры:
        tokenizer_model: str, необязательный путь к модели токенизатора; если не указан, берём путь по умолчанию TOKENIZER_MODEL.
        """
        # Если путь к модели передан — используем его, иначе путь по умолчанию
        model_path = tokenizer_model if tokenizer_model else TOKENIZER_MODEL
        # Проверяем, что файл модели существует
        assert os.path.isfile(model_path), model_path

        # Загружаем модель SentencePiece
        self.sp_model = SentencePieceProcessor(model_file=model_path)
        self.model_path = model_path

        # Специальные токены и размер словаря
        self.n_words: int = self.sp_model.vocab_size()  # размер словаря
        self.bos_id: int = self.sp_model.bos_id()       # id токена начала последовательности (BOS)
        self.eos_id: int = self.sp_model.eos_id()       # id токена конца последовательности (EOS)
        self.pad_id: int = self.sp_model.pad_id()       # id токена-заполнителя (PAD)

        # Проверяем, что размер словаря согласован
        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()

    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
        """
        Кодирует строку в список id токенов. Можно добавить токен начала (BOS) и токен конца (EOS).

        Параметры:
        s: str, строка для кодирования.
        bos: bool, добавлять ли BOS в начало списка.
        eos: bool, добавлять ли EOS в конец списка.

        Возвращает:
        List[int]: список id токенов.
        """
        # Проверяем, что на входе строка
        assert type(s) is str
        # Кодируем строку в id токенов с помощью SentencePiece
        t = self.sp_model.encode(s)
        # Если нужен BOS — добавляем его в начало
        if bos:
            t = [self.bos_id] + t
        # Если нужен EOS — добавляем его в конец
        if eos:
            t = t + [self.eos_id]
        return t

    def decode(self, t: List[int]) -> str:
        """
        Декодирует список id токенов обратно в строку.

        Параметры:
        t: List[int], список id токенов.

        Возвращает:
        str: декодированная строка.
        """
        return self.sp_model.decode(t)

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

Затем определены два ключевых метода:

Метод encode превращает входной текст в последовательность id токенов. Загрузив обученную модель токенизатора, мы разбиваем текст на слова или подслова и отображаем их в числа. Такую последовательность чисел модель принимает при обучении и инференсе.
Метод decode делает обратное: превращает последовательность id токенов в читаемый текст. Числа переводятся обратно в токены, токены склеиваются в текст — так выход модели можно прочитать и показать.

С этими методами очень удобно переходить от текста к числам и обратно — это интерфейс для входа и выхода модели. Проверить работу Tokenizer можно таким кодом: убедитесь, что он правильно переводит текст в числа и числа в текст.

# Проверяем Tokenizer
enc = Tokenizer('./data/tok4096.model')  # загружаем токенизатор
text = 'Hello, world!'  # тестовый текст
print(enc.encode(text, bos=True, eos=True))  # кодируем текст
print(enc.decode(enc.encode(text, bos=True, eos=True)))  # декодируем обратно

# ВЫВОД:
# [1, 346, 2233, 4010, 1475, 4021, 2]
# Hello, world!

Шаг 2. Предобработка данных

Перед обучением данные нужно предобработать. Главная задача этого шага — превратить текст в последовательности чисел, понятные модели. Каждый символ, слово или подслово текста отображается в уникальный числовой id — только так модель может обрабатывать данные.

# Функция обработки одного шарда
def process_shard(args, vocab_size, tokenizer_model_path):
    """ Обрабатывает шард данных: токенизирует текст и сохраняет в двоичный файл """
    ···


# Функция предобработки — пакетно обрабатывает несколько шардов
def pretokenize(vocab_size):
    """ Предобрабатывает все шарды и сохраняет токенизированные данные в двоичные файлы """
    ···


class PretokDataset(torch.utils.data.IterableDataset):
    """Загружает с диска предобработанные (токенизированные) данные и отдаёт их как тензоры PyTorch."""

    def __init__(self, split, max_seq_len, vocab_size, vocab_source):
        """
        Инициализация датасета.

        Параметры:
        split: str, часть датасета ('train' или 'test').
        max_seq_len: int, максимальная длина последовательности для входа и выхода.
        vocab_size: int, размер словаря.
        vocab_source: str, источник словаря ('llama2' или 'custom').
        """
        super().__init__()
        self.split = split  # часть датасета (обучающая или тестовая)
        self.max_seq_len = max_seq_len  # максимальная длина последовательности
        self.vocab_size = vocab_size  # размер словаря
        self.vocab_source = vocab_source  # источник словаря

    def __iter__(self):
        """
        Возвращает итератор: по батчам загружает данные и формирует вход и выход модели.
        """
        # Информация о воркере DataLoader (для параллельной загрузки)
        worker_info = torch.utils.data.get_worker_info()
        worker_id = worker_info.id if worker_info else 0  # id воркера
        # Ранг процесса при распределённом обучении (для нескольких GPU)
        rank = dist.get_rank() if dist.is_initialized() else 0
        # Уникальное зерно случайности из worker_id и rank — чтобы данные у каждого воркера и ранга различались
        seed = 42 + worker_id + 1337 * rank
        rng = random.Random(seed)
        print(f"Created a PretokDataset with rng seed {seed}")

        # Путь к данным зависит от источника словаря
        if self.vocab_source == "llama2":
            # Со словарём Llama 2 файлы .bin лежат рядом с файлами .json
            bin_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
            shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))
        elif self.vocab_source == "custom":
            # Со своим словарём файлы .bin лежат в папке tok{N}
            bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{self.vocab_size}")
            shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))

        # Для разных частей датасета берём разные шарды:
        # обучающая часть — все шарды, кроме первого; тестовая — только первый
        shard_filenames = shard_filenames[1:] if self.split == "train" else shard_filenames[:1]
        assert len(shard_filenames) > 0, f"В {bin_dir} не найдено ни одного файла .bin"

        while True:
            # Перемешиваем порядок шардов
            rng.shuffle(shard_filenames)
            for shard in shard_filenames:
                # Читаем файл через memmap: данные остаются на диске, память почти не расходуется
                m = np.memmap(shard, dtype=np.uint16, mode="r")
                # Сколько кусков длины max_seq_len помещается в шард
                num_batches = len(m) // self.max_seq_len
                num_batches -= 1  # отбрасываем последний неполный кусок
                assert num_batches > 0, "Этот шард слишком маленький? Проверьте."
                # Перемешиваем индексы кусков
                ixs = list(range(num_batches))
                rng.shuffle(ixs)
                # Для каждого куска формируем вход x и цель y
                for ix in ixs:
                    start = ix * self.max_seq_len  # начало куска
                    end = start + self.max_seq_len + 1  # конец куска (на один токен длиннее)
                    # Переводим кусок в массив NumPy и копируем в оперативную память
                    chunk = torch.from_numpy((m[start:end]).astype(np.int64))
                    # Вход модели x — первые max_seq_len токенов
                    x = chunk[:-1]
                    # Цель y — те же токены, сдвинутые на один вперёд (следующий токен)
                    y = chunk[1:]
                    # Отдаём пару (x, y)
                    yield x, y


class Task:
    @staticmethod
    def iter_batches(batch_size, device, num_workers=0, **dataset_kwargs):
        ds = PretokDataset(**dataset_kwargs)
        dl = torch.utils.data.DataLoader(
            ds, batch_size=batch_size, pin_memory=True, num_workers=num_workers
        )
        for x, y in dl:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            yield x, y

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

Затем определена функция pretokenize для пакетной обработки нескольких шардов. Через неё все данные обрабатываются параллельно, что ещё сильнее ускоряет предобработку.

Далее мы спроектировали класс PretokDataset для загрузки предобработанного датасета. Он наследует torch.utils.data.IterableDataset — так с данными можно работать гибче и эффективнее. Главное в этом классе — метод __iter__, он порождает батчи для обучения.

Наконец, мы определили класс Task, который специально нужен для перебора датасета и выдачи входов и целевых выходов модели. Такая конструкция обеспечивает гладкую стыковку потока данных и стандартный вход для обучения модели. Проверить предобработанный датасет можно кодом ниже.

Самого «кода ниже» в оригинале нет — авторы, видимо, забыли его вставить. Минимальная проверка выглядит так (запускайте после python preprocess.py из папки code):

from preprocess import Task  # класс Task из файла preprocess.py проекта

batches = Task.iter_batches(
    batch_size=2, device="cpu",
    split="test", max_seq_len=16, vocab_size=4096, vocab_source="custom",
)
x, y = next(batches)          # берём первый батч
print(x.shape, y.shape)       # torch.Size([2, 16]) torch.Size([2, 16])
print(x[0, 1:6], y[0, :5])    # y — это x, сдвинутый на одну позицию

Обратите внимание на последнюю строку: «цель» — это тот же текст, сдвинутый на токен вперёд. Это и есть задача предсказания следующего токена из урока 4.1.1.

Шаг 3. Обучаем модель

Когда предобработка данных закончена, можно обучать модель. Наша модель — Decoder-only Transformer той же структуры, что LLaMA2, реализованный на PyTorch. Детали реализации — в файле model.py, здесь мы их не повторяем: в исходнике подробные комментарии на китайском, а саму архитектуру мы подробно разбирали в предыдущих статьях.

В части про модель советуем обратить особое внимание на то, как генеративная модель порождает токены. Посмотрите класс Transformer в файле model.py, особенно метод generate: он показывает, как модель по уже имеющемуся контексту порождает следующие токены.

@torch.inference_mode()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
    """
    Получает входную последовательность idx (тензор целых чисел формы (bz, seq_len)) и много раз
    порождает новые токены, дописывая последовательность.
    Работает в режиме model.eval(). Это неэффективная версия выборки — без KV-кэша.
    """
    for _ in range(max_new_tokens):
        # Если контекст слишком длинный — обрезаем его до максимальной длины
        idx_cond = idx if idx.size(1) <= self.args.max_seq_len else idx[:, -self.args.max_seq_len:]

        # Прямой проход: получаем логиты для последней позиции последовательности
        logits = self(idx_cond)
        logits = logits[:, -1, :]  # оставляем выход только последнего шага

        if temperature == 0.0:
            # Берём самый вероятный индекс
            _, idx_next = torch.topk(logits, k=1, dim=-1)
        else:
            # Масштабируем логиты температурой и применяем softmax
            logits = logits / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)

        # Добавляем выбранный индекс к последовательности и продолжаем
        idx = torch.cat((idx, idx_next), dim=1)

    return idx

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

python train.py

В train.py задано много гиперпараметров: размерность модели, число слоёв, скорость обучения и другие. Часть из них показана ниже; остальное смотрите в исходнике — там подробные комментарии на китайском, разобраться несложно.

# -----------------------------------------------------------------------------
# Настройки ввода-вывода: куда сохранять результаты, как часто писать логи и оценивать модель
out_dir = "output"  # папка для сохранения модели
eval_interval = 2000  # оценивать каждые 2000 шагов
log_interval = 1  # писать лог на каждом шаге
eval_iters = 100  # сколько итераций делать при каждой оценке
eval_only = False  # если True, скрипт завершится сразу после первой оценки
always_save_checkpoint = False  # если True, сохранять контрольную точку после каждой оценки
init_from = "scratch"  # обучать с нуля ('scratch') или продолжить с контрольной точки ('resume')

# Данные
batch_size = 8  # примеров в одном микробатче; с накоплением градиента реальный батч больше
max_seq_len = 256  # максимальная длина последовательности
vocab_size = 4096  # размер своего словаря

# Модель
dim = 288  # размерность скрытого слоя модели
n_layers = 8  # число слоёв Transformer
n_heads = 8  # число голов внимания
n_kv_heads = 4  # число групп K/V (группировка голов)
multiple_of = 32  # размерности некоторых слоёв должны быть кратны этому числу
dropout = 0.0  # вероятность Dropout

# Оптимизатор AdamW
gradient_accumulation_steps = 4  # шаги накопления градиента — имитируют больший батч
learning_rate = 5e-4  # максимальная скорость обучения
max_iters = 100000  # общее число итераций обучения
weight_decay = 1e-1  # коэффициент затухания весов
beta1 = 0.9  # параметр β1 оптимизатора AdamW
beta2 = 0.95  # параметр β2 оптимизатора AdamW
grad_clip = 1.0  # порог обрезки градиента, 0 — без обрезки

# Затухание скорости обучения
decay_lr = True  # включить ли затухание скорости обучения
warmup_iters = 1000  # число шагов разогрева

# Система
device = "cuda:0"  # устройство: 'cpu', 'cuda', 'cuda:0' и т. п.
dtype = "bfloat16"  # тип данных: 'float32', 'bfloat16', 'float16'

Шаг 4. Генерируем текст моделью

После обучения в папке output появится файл ckpt.pt — это и есть обученная модель. Сгенерировать текст можно такой командой:

python sample.py --prompt "One day, Lily met a Shoggoth"

Посмотрим на код в sample.py. В нём определён класс TextGenerator, который генерирует текст.

class TextGenerator:
    def __init__(self,
                 checkpoint='output/ckpt.pt',  # путь к контрольной точке модели
                 tokenizer_model_path='tok4096.model',  # путь к модели токенизатора
                 seed=1337,  # зерно случайности — для воспроизводимости
                 device=None,  # устройство: CUDA, если есть, иначе CPU
                 dtype="float32"):  # тип данных, по умолчанию float32; можно float16 или bfloat16
        """
        Инициализация TextGenerator: загружаем модель, выбираем устройство, загружаем токенизатор.
        """
        # Настройки загрузки модели
        self.checkpoint = checkpoint  # путь к сохранённой контрольной точке
        self.tokenizer_model_path = tokenizer_model_path  # путь к файлу модели токенизатора
        self.seed = seed  # зерно случайности — генерация воспроизводима
        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')  # выбираем устройство по наличию GPU
        self.dtype = dtype  # тип чисел с плавающей точкой для модели
        self.device_type = 'cuda' if 'cuda' in self.device else 'cpu'  # это CUDA или CPU?

        # Фиксируем зёрна случайности для воспроизводимости
        torch.manual_seed(seed)  # зерно для CPU
        torch.cuda.manual_seed(seed)  # зерно для CUDA
        torch.backends.cuda.matmul.allow_tf32 = True  # разрешаем CUDA умножать матрицы в TF32
        torch.backends.cudnn.allow_tf32 = True  # разрешаем cuDNN ускоряться за счёт TF32

        # По dtype выбираем контекст автоматической смешанной точности
        ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[self.dtype]
        self.ctx = nullcontext() if self.device_type == 'cpu' else torch.amp.autocast(device_type=self.device_type, dtype=ptdtype)

        # Загружаем файл контрольной точки
        checkpoint_dict = torch.load(self.checkpoint, map_location=self.device)  # загружаем параметры
        gptconf = ModelArgs(**checkpoint_dict['model_args'])  # восстанавливаем настройки модели
        self.model = Transformer(gptconf)  # создаём модель Transformer
        state_dict = checkpoint_dict['model']  # словарь состояния модели

        # Убираем из ключей словаря состояния лишний префикс
        unwanted_prefix = '_orig_mod.'  # этот префикс мог добавиться при сохранении (после torch.compile)
        for k, v in list(state_dict.items()):
            if k.startswith(unwanted_prefix):
                state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)  # убираем префикс

        # Загружаем параметры в модель
        self.model.load_state_dict(state_dict, strict=False)
        # Считаем число параметров модели
        num_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        print(f"Model has {num_params} parameters.")
        # Переводим модель в режим оценки, чтобы dropout и т. п. не влияли на результат
        self.model.eval()
        # Переносим модель на нужное устройство (GPU или CPU)
        self.model.to(self.device)
        # Загружаем токенизатор по указанному пути
        self.tokenizer = Tokenizer(tokenizer_model=self.tokenizer_model_path)

    def sample(self,
               start="Hello!",  # начальная подсказка, любая строка
               num_samples=3,  # сколько текстов сгенерировать, по умолчанию 3
               max_new_tokens=256,  # максимум новых токенов на один текст, по умолчанию 256
               temperature=1.0,  # управляет случайностью: 1.0 — стандарт, больше — случайнее
               top_k=300):  # оставлять только top_k самых вероятных токенов
        """
        Генерирует тексты по заданному началу.

        :param start: начальная подсказка
        :param num_samples: сколько текстов сгенерировать
        :param max_new_tokens: максимум новых токенов на текст
        :param temperature: случайность генерации — меньше значит увереннее, больше — случайнее
        :param top_k: из скольких самых вероятных токенов выбирать
        :return: список сгенерированных текстов
        """
        # Если start начинается с 'FILE:', начальный текст читаем из файла
        if start.startswith('FILE:'):
            with open(start[5:], 'r', encoding='utf-8') as f:
                start = f.read()  # содержимое файла — начальный текст

        # Кодируем начальный текст в последовательность id токенов
        start_ids = self.tokenizer.encode(start, bos=True, eos=False)  # bos=True — добавить начало, eos=False — без конца
        x = (torch.tensor(start_ids, dtype=torch.long, device=self.device)[None, ...])  # превращаем id в тензор PyTorch

        generated_texts = []  # сюда складываем сгенерированные тексты
        with torch.no_grad():  # отключаем подсчёт градиентов — так быстрее
            with self.ctx:  # контекст смешанной точности (на GPU с float16/bfloat16)
                for k in range(num_samples):  # генерируем нужное число текстов
                    y = self.model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)  # генерация
                    generated_texts.append(self.tokenizer.decode(y[0].tolist()))  # декодируем токены в текст

        return generated_texts  # возвращаем список текстов

Ссылки

llama2.c — проект Андрея Карпатого, на котором основан этот код.

От переводчика: как запустить на своей машине

Проекту хватает одной RTX 4080: укажите device = "cuda:0" (так и стоит по умолчанию). Остальные три GPU свободны — например, можно параллельно обучить вторую модель с другими гиперпараметрами на cuda:1 и сравнить.
Шаг 1 качает с ModelScope архив TinyStories_all_data.tar.gz и обучает токенизатор на CPU (num_threads = число ядер, так что 48 ядер EPYC здесь пригодятся).
Шаг 3 по оценке авторов занимает несколько часов. Чтобы быстро убедиться, что всё работает, поставьте max_iters = 2000 и eval_interval = 500: лосс должен быстро уйти вниз от стартовых ~8,3 (это $\ln 4096$ — столько даёт модель, которая угадывает все 4096 токенов равновероятно). Учтите, что косинусное затухание скорости обучения привязано к max_iters (lr_decay_iters = max_iters), так что короткий прогон — это полноценное, просто маленькое обучение.
Без GPU обучение тоже идёт (device = "cpu", dtype = "float32"), но в десятки раз медленнее — оставьте короткий прогон для проверки.

Что сделать руками:

Пройдите шаги 1–4 с коротким прогоном (max_iters = 2000) и сгенерируйте истории с temperature = 0, 0,5 и 1,0. Сравните разнообразие — это практическая версия урока 4.1.3.
Сравните train.py этого проекта с циклом обучения из урока 4.6.3: найдите разогрев скорости обучения, косинусное затухание, накопление градиента и обрезку градиента.
Добавьте в generate KV-кэш по уроку 4.6.4 и замерьте ускорение при max_new_tokens = 256.

Словарик

Русский	Пояснение
токенизатор (Tokenizer)	превращает текст в id токенов и обратно
словарь (vocab)	набор всех токенов; здесь 4096
шард (shard)	фрагмент большого датасета, отдельный файл
предобработка	токенизация и сохранение в двоичный формат
только-декодер (Decoder-only)	архитектура GPT/LLaMA
накопление градиента	несколько микробатчей на один шаг оптимизатора
разогрев скорости обучения (warmup)	плавный рост LR в начале
затухание весов (weight decay)	регуляризация в AdamW
контрольная точка (checkpoint)	сохранённые веса и настройки
сэмплирование, выборка (sampling)	выбор следующего токена по распределению
температура (temperature)	«острота» распределения перед выбором

Перевод выполнен по лицензии оригинала (CC-BY-NC-SA-4.0); автор оригинала — Datawhale (Сун Чжисюэ и соавторы проекта tiny-universe). Нашли неточность — сверяйтесь с оригиналом по ссылке выше. · обновлено 2026-10-05