CS336, лекции 1–2: зачем строить LLM с нуля, токенизация, PyTorch и учёт ресурсов

Полный перевод двух первых исполняемых лекций Stanford CS336 (весна 2025): карта курса, история языковых моделей, токенизаторы от символов до BPE, тензоры, типы чисел, FLOPs, MFU, оптимизаторы и цикл обучения.

Полный переводс языка: английский

Оригинал: Stanford CS336 (Spring 2025) — Lecture 1: Overview, tokenization; Lecture 2: PyTorch, resource accounting · исполняемые лекции на Python (lecture_01.py, lecture_02.py) и видео на YouTube · весна 2025
Автор: Перси Лян (Percy Liang), Тацунори Хашимото (Tatsunori Hashimoto), Stanford University
Лицензия оригинала: MIT (© 2025 Stanford University)

Текст оригинала переведён целиком, его лицензия это разрешает.

К урокам курса

CS336 читают необычно: часть лекций — это исполняемые программы на Python. Каждый вызов text("…") выводит строку на «слайд», image(…) — картинку, link(…) — ссылку, а комментарий # @inspect x показывает значение переменной прямо в просмотрщике лекции. Ниже переведён весь текст лекций 1 и 2 по порядку, код приведён с русскими комментариями, а значения, которые просмотрщик показывает рядом с кодом, даны в комментариях # → … (взяты из трасс выполнения в репозитории лекций). Картинки не копируем: где они важны, даём ссылку и описание. Оригиналы: lecture_01.py, lecture_02.py; видео — плейлист 2025 года, лекции 1 и 2.

Как читать вместе с курсом

Лекция 1 — обзор всего CS336 и тема токенизации. Её полезно прочитать сразу после модуля 4.2 «Токенизация»: BPE здесь тот же, что в уроке 4.2.2, но с другим акцентом — на эффективность.
Лекция 2 — «бухгалтерия ресурсов»: память в байтах и вычисления в FLOPs. Требует уверенного PyTorch (урок 3.3.1); формула «6 × токены × параметры» отсюда пригодится в уроках 5.1.1 и 5.7.1.
Новичку можно пропустить разделы о логистике Стэнфорда — они переведены для полноты.

Лекция 1. Обзор курса и токенизация

CS336: Language Models From Scratch (весна 2025)

[Фото преподавателей и ассистентов курса — images/course-staff.png.]

Это второй запуск CS336. Стэнфордский поток вырос на 50%. Лекции будут выложены на YouTube и станут доступны всему миру.

Зачем мы сделали этот курс?

Спросим GPT-4 (GPT-4 Technical Report): «Зачем преподавать курс о построении языковых моделей с нуля? Ответь одним предложением».

response = query_gpt4o(prompt="Why teach a course on building language models from scratch? Answer in one sentence.")
# → «Курс о построении языковых моделей с нуля даёт фундаментальное понимание методов обработки
#    естественного языка, поощряет инновации и позволяет эффективно адаптировать модели к конкретным задачам
#    и предметным областям».

Проблема: исследователи теряют связь с технологией, на которой стоят.

8 лет назад исследователи сами реализовывали и обучали свои модели.
6 лет назад исследователи скачивали модель (например, BERT) и дообучали её.
Сегодня исследователи просто пишут промпты к проприетарной модели (GPT-4, Claude, Gemini).

Подъём по уровням абстракции повышает продуктивность, но:

эти абстракции «протекают» (в отличие от языков программирования или операционных систем);
остаются фундаментальные исследования, для которых нужно разбирать весь стек.

Полное понимание этой технологии необходимо для фундаментальных исследований.

Этот курс: понимание через построение. Но есть одна маленькая проблема…

Индустриализация языковых моделей

[Иллюстрация: старинная гравюра с заводскими трубами — эпоха индустриализации.]

У GPT-4, по слухам, 1,8 трлн параметров (статья).
Обучение GPT-4, по слухам, стоило 100 млн долл. (статья).
xAI строит кластер из 200 000 H100 для обучения Grok (статья).
Stargate (OpenAI, NVIDIA, Oracle) вкладывает 500 млрд долл. за 4 года (анонс).

Кроме того, нет публичных подробностей о том, как строятся фронтирные модели. Из технического отчёта GPT-4: [снимок абзаца, где OpenAI объясняет, что из-за конкуренции и безопасности не раскрывает архитектуру, размер модели, железо, вычисления, данные и метод обучения — images/gpt4-no-details.png].

Больше — значит по-другому

Фронтирные модели нам недоступны. Но маленькие языковые модели (в этом курсе — меньше 1 млрд параметров) могут плохо представлять большие.

Пример 1: доля FLOPs на внимание и на MLP меняется с масштабом (пост Стивена Роллера в X). [График: у маленьких моделей заметная доля вычислений уходит на внимание, у больших почти всё съедает MLP.]
Пример 2: появление новых способностей с ростом масштаба (Emergent Abilities of Large Language Models). [Графики: точность на ряде задач держится около случайной и резко растёт после определённого бюджета вычислений.]

Что из этого курса переносится на фронтирные модели?

Есть три вида знаний:

Механика: как всё работает (что такое трансформер, как параллелизм моделей использует GPU).
Образ мышления: выжимать из железа максимум, относиться к масштабу всерьёз (законы масштабирования).
Интуиции: какие решения о данных и модели дают хорошую точность.

Механике и образу мышления мы научить можем (они переносятся). Интуициям — лишь частично (они не обязательно переносятся между масштабами).

Интуиции? 🤷

Некоторые проектные решения просто (пока) нельзя обосновать — они пришли из экспериментов. Пример: статья Ноама Шазира, в которой появился SwiGLU (GLU Variants Improve Transformer). [Снимок финала статьи: автор пишет, что не предлагает объяснения, почему эти архитектуры работают, и приписывает их успех, как и всё остальное, «божественной благосклонности» — images/divine-benevolence.png.]

Горький урок

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

точность = эффективность × ресурсы

На большом масштабе эффективность даже гораздо важнее (там нельзя позволить себе расточительность). Measuring the Algorithmic Efficiency of Neural Networks показала, что алгоритмическая эффективность на ImageNet с 2012 по 2019 год выросла в 44 раза.

Постановка задачи: какую лучшую модель можно построить при заданном бюджете вычислений и данных? Иными словами, максимизируйте эффективность!

Текущий ландшафт

До нейросетей (до 2010-х)

Языковая модель для измерения энтропии английского языка (Шеннон, 1950).
Много работ об n-граммных языковых моделях — для машинного перевода и распознавания речи (Brants et al., 2007).

Нейросетевые ингредиенты (2010-е)

Первая нейросетевая языковая модель (Бенджио и др., 2003).
Моделирование «последовательность в последовательность» для перевода (Sutskever et al., 2014).
Оптимизатор Adam (Kingma, Ba, 2014).
Механизм внимания для перевода (Bahdanau et al., 2014).
Архитектура трансформер для перевода (Attention Is All You Need, 2017).
Смесь экспертов (Shazeer et al., 2017).
Параллелизм моделей (GPipe, ZeRO, Megatron-LM).

Ранние фундаментальные модели (конец 2010-х)

ELMo: предобучение на LSTM, дообучение помогает в задачах (статья).
BERT: предобучение трансформера, дообучение помогает в задачах (статья).
T5 от Google (11 млрд): всё сводится к задаче «текст в текст» (статья).

Ставка на масштаб, больше закрытости

GPT-2 от OpenAI (1,5 млрд): связный текст, первые признаки zero-shot, поэтапный выпуск (статья).
Законы масштабирования: надежда и предсказуемость при масштабировании (Kaplan et al., 2020).
GPT-3 от OpenAI (175 млрд): обучение в контексте, закрытая модель (статья).
PaLM от Google (540 млрд): огромный масштаб, недообучена (статья).
Chinchilla от DeepMind (70 млрд): вычислительно-оптимальные законы масштабирования (статья).

Открытые модели

Открытые датасеты (The Pile) и модели (GPT-J) EleutherAI (The Pile, GPT-J).
OPT от Meta (175 млрд): воспроизведение GPT-3, масса проблем с железом (статья).
BLOOM от Hugging Face и BigScience: упор на происхождение данных (статья).
Модели Llama от Meta (LLaMA, Llama 2, Llama 3).
Модели Qwen от Alibaba (Qwen2.5).
Модели DeepSeek (DeepSeek LLM, DeepSeek-V2, DeepSeek-V3).
OLMo 2 от AI2 (OLMo, OLMo 2).

Уровни открытости

Закрытые модели (например, GPT-4o): только доступ через API.
Модели с открытыми весами (например, DeepSeek): веса доступны, статья с деталями архитектуры, часть деталей обучения, никаких деталей о данных.
Модели с открытым исходным кодом (например, OLMo): доступны веса и данные, статья с большинством деталей (но не обязательно с обоснованиями и неудачными экспериментами).

Фронтирные модели сегодня (весна 2025): o3 от OpenAI, Claude Sonnet 3.7 от Anthropic, Grok 3 от xAI, Gemini 2.5 от Google, Llama 3.3 от Meta, r1 от DeepSeek, Qwen 2.5 Max от Alibaba, Hunyuan-T1 от Tencent.

Что это за программа

Это исполняемая лекция — программа, выполнение которой и есть содержание лекции. Исполняемые лекции позволяют:

смотреть и запускать код (ведь всё — это код!);

total = 0             # → total = 0
for x in [1, 2, 3]:   # → x = 1, затем 2, затем 3
    total += x        # → total = 1, 3, 6

видеть иерархическую структуру лекции;
переходить к определениям и понятиям (в оригинале — ссылка на функцию supervised_finetuning).

Логистика курса

Вся информация — на сайте курса (архив весны 2025).

Это курс на 5 кредитов. Комментарий из отзывов студентов весны 2024: «Всё задание по объёму было примерно как все 5 заданий CS 224n плюс финальный проект. И это только первое домашнее задание».

Почему вам стоит взять этот курс

У вас навязчивая потребность понимать, как всё устроено.
Вы хотите накачать мускулы инженера-исследователя.

Почему вам не стоит брать этот курс

Вы на самом деле хотите в этом семестре сделать исследование (поговорите с научным руководителем).
Вам интересны самые горячие новые техники ИИ — мультимодальность, RAG и т. п. (для этого возьмите семинар).
Вы хотите хороших результатов в своей прикладной области (просто пишите промпты или дообучите существующую модель).

Как заниматься дома

Все материалы лекций и задания выкладываются в открытый доступ — занимайтесь самостоятельно.
Лекции записываются через CGOE (бывший SCPD) и выкладываются на YouTube (с задержкой).
Мы планируем провести курс снова в следующем году.

Задания

5 заданий: основы, системы, законы масштабирования, данные, выравнивание.
Никакого готового каркаса кода, но мы даём модульные тесты и интерфейсы-адаптеры, чтобы проверить правильность.
Реализуете локально, проверяете правильность, затем запускаете на кластере для замеров (точность и скорость).
В некоторых заданиях есть лидерборд (минимизировать перплексию при заданном бюджете обучения).
ИИ-инструменты (CoPilot, Cursor и т. п.) могут отнять у вас обучение — пользуйтесь на свой страх и риск.

Кластер

Спасибо Together AI за вычислительный кластер. 🙏
Прочитайте руководство по работе с кластером.
Начинайте задания заранее: ближе к дедлайну кластер забивается!

Всё дело в эффективности

Ресурсы: данные + железо (вычисления, память, пропускная способность связи). Как обучить лучшую модель при фиксированном наборе ресурсов? Пример: есть дамп Common Crawl и 32 H100 на 2 недели — что делать?

Проектные решения: [схема images/design-decisions.png — дерево решений от данных (обработка, фильтрация) и токенизации через архитектуру и обучение к выравниванию].

Обзор курса

Основы

Цель: заставить работать базовую версию всего конвейера. Компоненты: токенизация, архитектура модели, обучение.

Токенизация. Токенизаторы переводят строки в последовательности целых чисел (токенов) и обратно. [Пример: фраза, раскрашенная по токенам.] Интуиция: разбить строку на популярные фрагменты.

В этом курсе — токенизатор Byte-Pair Encoding (BPE) (Sennrich et al., 2016). Подходы без токенизатора: ByT5, MEGABYTE, BLT, T-FREE — работают прямо с байтами; многообещающе, но пока не масштабированы до фронтира.

Архитектура. Отправная точка — исходный трансформер (Attention Is All You Need). Варианты:

функции активации: ReLU, SwiGLU (Shazeer, 2020);
позиционное кодирование: синусоидальное, RoPE (RoFormer);
нормализация: LayerNorm, RMSNorm (LayerNorm, RMSNorm);
место нормализации: pre-norm против post-norm (Xiong et al., 2020);
MLP: плотный, смесь экспертов (Shazeer et al., 2017);
внимание: полное, со скользящим окном, линейное (Mistral 7B, линейное внимание);
внимание пониженной размерности: GQA, MLA (GQA, DeepSeek-V2);
модели пространства состояний: Hyena (статья).

Обучение.

Оптимизатор (AdamW, Muon, SOAP).
Расписание learning rate (косинусное, WSD).
Размер батча (критический размер батча, An Empirical Model of Large-Batch Training).
Регуляризация (dropout, weight decay).
Гиперпараметры (число голов, скрытая размерность): поиск по сетке.

Задание 1 (GitHub):

реализовать токенизатор BPE;
реализовать трансформер, кросс-энтропию, оптимизатор AdamW, цикл обучения;
обучить на TinyStories и OpenWebText;
лидерборд: минимизировать перплексию на OpenWebText за 90 минут на одной H100.

Системы

Цель: выжать из железа максимум. Компоненты: ядра, параллелизм, инференс.

Ядра. Как выглядит GPU (A100): [фото платы и схема чипа]. Аналогия: склад относится к DRAM, как фабрика к SRAM [рисунок Хораса Хе: фабрика (вычисления) и склад (память), соединённые дорогой — пропускной способностью]. Приём: организовать вычисления так, чтобы максимально загрузить GPU, минимизируя перемещение данных. Ядра пишут на CUDA, Triton, CUTLASS, ThunderKittens.

Параллелизм. А если GPU несколько (8 A100)? [Схема топологии типичного сервера 8×A100.] Перемещение данных между GPU ещё медленнее, но принцип «минимизируй перемещение данных» тот же. Используем коллективные операции (gather, reduce, all-reduce). Шардируем (параметры, активации, градиенты, состояния оптимизатора) по GPU. Как делить вычисления: параллелизм по данным, тензорный, конвейерный, по последовательности.

Инференс. Цель: сгенерировать токены по промпту (это нужно, чтобы моделями вообще пользоваться!). Инференс нужен и для обучения с подкреплением, вычислений во время инференса (test-time compute), оценки. В масштабах мира вычисления на инференс (каждое использование) превышают вычисления на обучение (разовая стоимость).

Две фазы: prefill и decode [схема images/prefill-decode.png].

Prefill (похож на обучение): токены даны, их можно обработать все сразу (упирается в вычисления).
Decode: токены генерируются по одному (упирается в память).

Как ускорить декодирование:

взять модель подешевле (прунинг, квантование, дистилляция);
спекулятивное декодирование: дешёвая «черновая» модель генерирует несколько токенов, а полная модель проверяет их параллельно (декодирование точное!);
системные оптимизации: KV-кэш, батчинг.

Задание 2 (в лекции дана ссылка на версию 2024 года):

реализовать слитое ядро RMSNorm на Triton;
реализовать распределённое обучение с параллелизмом по данным;
реализовать шардирование состояний оптимизатора;
замерить и профилировать реализации.

Законы масштабирования

Цель: ставить эксперименты в малом масштабе и предсказывать гиперпараметры и loss в большом.

Вопрос: при бюджете FLOPs $C$ взять модель побольше ( $N$ ) или обучать на большем числе токенов ( $D$ )? Вычислительно-оптимальные законы масштабирования: Kaplan et al., 2020, Chinchilla. [График IsoFLOP-кривых из статьи Chinchilla.]

Коротко: $D^* = 20 N^*$ (например, модель на 1,4 млрд параметров нужно обучать на 28 млрд токенов). Но здесь не учтена стоимость инференса!

Задание 3 (в лекции — версия 2024 года):

мы задаём «API обучения» (гиперпараметры → loss), построенный на прошлых запусках;
вы отправляете «задачи обучения» (в рамках бюджета FLOPs) и собираете точки данных;
подгоняете закон масштабирования к точкам;
отправляете предсказания для гиперпараметров в большом масштабе;
лидерборд: минимизировать loss при бюджете FLOPs.

Данные

Вопрос: какие способности мы хотим от модели? Многоязычность? Код? Математика? [Круговая диаграмма состава The Pile.]

Оценка

Перплексия: хрестоматийная оценка языковых моделей.
Стандартизованные тесты (MMLU, HellaSwag, GSM8K).
Следование инструкциям (AlpacaEval, IFEval, WildBench).
Масштабирование вычислений во время инференса: цепочки рассуждений, ансамбли.
LM-as-a-judge: оценка генеративных задач.
Система целиком: RAG, агенты.

Курирование данных

Данные не падают с неба. [В лекции код скачивает три файла Common Crawl (дамп CC-MAIN-2024-18), берёт 10 случайных документов и превращает их в Markdown — файл var/sample-documents.txt.] Там пустошь! Данные нужно серьёзно обрабатывать.
Источники: веб-страницы из интернета, книги, статьи arXiv, код с GitHub и т. д.
Ссылаться на добросовестное использование (fair use), чтобы обучаться на данных под авторским правом? (Foundation Models and Fair Use)
Может понадобиться лицензировать данные (например, Google и данные Reddit, статья Reuters).
Форматы: HTML, PDF, каталоги файлов (не текст!).

Обработка данных

Преобразование: HTML/PDF в текст (сохранить содержание, часть структуры, переписывание).
Фильтрация: оставить качественные данные, убрать вредный контент (с помощью классификаторов).
Дедупликация: экономит вычисления, мешает запоминанию; фильтры Блума или MinHash.

Задание 4 (в лекции — версия 2024 года):

преобразовать HTML из Common Crawl в текст;
обучить классификаторы для фильтрации по качеству и вредному контенту;
дедупликация с MinHash;
лидерборд: минимизировать перплексию при бюджете токенов.

Выравнивание

Пока что базовая модель — это сырой потенциал: она очень хорошо дописывает следующий токен. Выравнивание делает модель по-настоящему полезной.

Цели выравнивания:

научить языковую модель следовать инструкциям;
настроить стиль (формат, длину, тон и т. п.);
встроить безопасность (например, отказы отвечать на вредные вопросы).

Две фазы.

Обучение с учителем на инструкциях (SFT). Данные инструкций — пары (промпт, ответ):

@dataclass(frozen=True)
class Turn:
    role: str      # роль: system, user или assistant
    content: str   # текст реплики

@dataclass(frozen=True)
class ChatExample:
    turns: list[Turn]

sft_data: list[ChatExample] = [
    ChatExample(
        turns=[
            Turn(role="system", content="You are a helpful assistant."),  # «Ты — полезный ассистент»
            Turn(role="user", content="What is 1 + 1?"),                 # «Сколько будет 1 + 1?»
            Turn(role="assistant", content="The answer is 2."),          # «Ответ: 2»
        ],
    ),
]

Такие данные часто размечают люди. Интуиция: у базовой модели навыки уже есть, нужно лишь немного примеров, чтобы их проявить (LIMA: Less Is More for Alignment). Обучение с учителем: дообучаем модель, максимизируя $p(\text{ответ} \mid \text{промпт})$ .

Обучение на обратной связи. Теперь у нас есть предварительная модель, следующая инструкциям. Улучшим её без дорогой разметки.

Данные предпочтений. Модель генерирует несколько ответов (например, A и B) на один промпт, пользователь указывает предпочтение (A хуже B или A лучше B):

@dataclass(frozen=True)
class PreferenceExample:
    history: list[Turn]   # диалог до ответа
    response_a: str       # первый вариант ответа
    response_b: str       # второй вариант ответа
    chosen: str           # какой выбран: "a" или "b"

preference_data: list[PreferenceExample] = [
    PreferenceExample(
        history=[
            Turn(role="system", content="You are a helpful assistant."),
            Turn(role="user", content="What is the best way to train a language model?"),  # «Как лучше всего обучать языковую модель?»
        ],
        response_a="You should use a large dataset and train for a long time.",   # «Большой датасет и долгое обучение»
        response_b="You should use a small dataset and train for a short time.",  # «Маленький датасет и короткое обучение»
        chosen="a",
    )
]

Верификаторы. Формальные (для кода, математики) и обученные — обучение против LM-as-a-judge.

Алгоритмы.

Proximal Policy Optimization (PPO) из обучения с подкреплением (PPO, InstructGPT).
Direct Preference Optimization (DPO): для данных предпочтений, проще (DPO).
Group Relative Policy Optimization (GRPO): убирает функцию ценности (DeepSeekMath).

Задание 5 (в лекции — версия 2024 года): реализовать SFT, DPO и GRPO.

Эффективность определяет проектные решения

Сегодня мы ограничены вычислениями, поэтому проектные решения отражают стремление выжать максимум из данного железа.

Обработка данных: не тратить драгоценные вычисления на обновления по плохим или нерелевантным данным.
Токенизация: работать с сырыми байтами элегантно, но при нынешних архитектурах вычислительно неэффективно.
Архитектура модели: многие изменения мотивированы экономией памяти или FLOPs (общий KV-кэш, внимание со скользящим окном).
Обучение: можно обойтись одной эпохой!
Законы масштабирования: подбирать гиперпараметры на маленьких моделях, тратя меньше вычислений.
Выравнивание: если сильнее подстроить модель под нужные сценарии, хватит базовой модели поменьше.

Завтра мы станем ограничены данными…

Токенизация

Эта часть навеяна видео Андрея Карпатого о токенизации — посмотрите его! (видео; у нас есть урок по его мотивам — 4.2.2).

Введение

Сырой текст обычно представлен строками Unicode: string = "Hello, 🌍! Привет!". (Примечание переводчика: в оригинале в этой строке китайское приветствие; мы заменили его русским и пересчитали все числа в примерах ниже.)

Языковая модель задаёт распределение вероятностей над последовательностями токенов (обычно это целочисленные индексы): indices = [15496, 11, 995, 0].

Значит, нужна процедура, которая кодирует строки в токены, и процедура, которая декодирует токены обратно в строки. Токенизатор — это класс, реализующий методы encode и decode:

class Tokenizer(ABC):
    """Абстрактный интерфейс токенизатора."""
    def encode(self, string: str) -> list[int]:
        raise NotImplementedError

    def decode(self, indices: list[int]) -> str:
        raise NotImplementedError

Размер словаря — число возможных токенов (целых чисел).

Примеры токенизации

Чтобы почувствовать, как работают токенизаторы, поиграйте с интерактивным сайтом.

Наблюдения

Слово и пробел перед ним входят в один токен (например, " world").
Слово в начале и в середине текста кодируется по-разному (например, "hello hello").
Числа разбиваются на токены по нескольку цифр.

Вот токенизатор GPT-2 от OpenAI (tiktoken) в действии. Проверим, что encode() и decode() дают исходную строку:

import tiktoken

def get_gpt2_tokenizer():
    # Код: https://github.com/openai/tiktoken
    # Для токенизатора gpt-3.5-turbo или gpt-4 можно взять cl100k_base
    return tiktoken.get_encoding("gpt2")

def get_compression_ratio(string: str, indices: list[int]) -> float:
    """Степень сжатия: сколько байт UTF-8 приходится на один токен."""
    num_bytes = len(bytes(string, encoding="utf-8"))
    num_tokens = len(indices)
    return num_bytes / num_tokens

tokenizer = get_gpt2_tokenizer()
string = "Hello, 🌍! Привет!"
indices = tokenizer.encode(string)
# → [15496, 11, 12520, 234, 235, 0, 12466, 253, 21169, 18849, 38857, 16843, 20375, 0]
reconstructed_string = tokenizer.decode(indices)   # → "Hello, 🌍! Привет!"
assert string == reconstructed_string
compression_ratio = get_compression_ratio(string, indices)   # → 26 байт / 14 токенов ≈ 1.86

Посимвольная токенизация

Строка Unicode — это последовательность символов Unicode. Каждый символ превращается в кодовую точку (целое число) функцией ord и обратно — функцией chr:

assert ord("a") == 97
assert ord("🌍") == 127757
assert chr(97) == "a"
assert chr(127757) == "🌍"

Построим Tokenizer и убедимся, что он возвращает исходную строку:

class CharacterTokenizer(Tokenizer):
    """Представляет строку как последовательность кодовых точек Unicode."""
    def encode(self, string: str) -> list[int]:
        return list(map(ord, string))

    def decode(self, indices: list[int]) -> str:
        return "".join(map(chr, indices))

tokenizer = CharacterTokenizer()
indices = tokenizer.encode("Hello, 🌍! Привет!")
# → [72, 101, 108, 108, 111, 44, 32, 127757, 33, 32, 1055, 1088, 1080, 1074, 1077, 1090, 33]
vocabulary_size = max(indices) + 1      # это лишь нижняя граница → 127758
compression_ratio = get_compression_ratio("Hello, 🌍! Привет!", indices)   # → 26 / 17 ≈ 1.53

Символов Unicode примерно 150 тысяч (Википедия).

Проблема 1: это очень большой словарь.
Проблема 2: многие символы довольно редки (например, 🌍) — словарь используется неэффективно.

Побайтовая токенизация

Строку Unicode можно представить последовательностью байтов, а каждый байт — целым числом от 0 до 255. Самая распространённая кодировка Unicode — UTF-8. Некоторые символы занимают один байт, другие — несколько:

assert bytes("a", encoding="utf-8") == b"a"
assert bytes("🌍", encoding="utf-8") == b"\xf0\x9f\x8c\x8d"   # четыре байта


class ByteTokenizer(Tokenizer):
    """Представляет строку как последовательность байтов."""
    def encode(self, string: str) -> list[int]:
        string_bytes = string.encode("utf-8")
        indices = list(map(int, string_bytes))
        return indices

    def decode(self, indices: list[int]) -> str:
        string_bytes = bytes(indices)
        string = string_bytes.decode("utf-8")
        return string

tokenizer = ByteTokenizer()
indices = tokenizer.encode("Hello, 🌍! Привет!")
# → [72, 101, 108, 108, 111, 44, 32, 240, 159, 140, 141, 33, 32, 208, 159, 209, 128, 208, 184, 208, 178, 208, 181, 209, 130, 33]
vocabulary_size = 256
compression_ratio = get_compression_ratio("Hello, 🌍! Привет!", indices)   # → 1.0
assert compression_ratio == 1

Словарь приятно маленький: байт принимает 256 значений. А степень сжатия? Она ужасна — последовательности получаются слишком длинными. А контекст трансформера ограничен (внимание квадратично по длине), так что выглядит это невесело…

Пословная токенизация

Другой подход (ближе к классической обработке языка) — делить строки на слова:

import regex

string = "I'll say supercalifragilisticexpialidocious!"
segments = regex.findall(r"\w+|.", string)
# → ["I", "'", "ll", " ", "say", " ", "supercalifragilisticexpialidocious", "!"]

Это регулярное выражение держит вместе все буквенно-цифровые символы (слова). Вот вариант поизощрённее — регулярное выражение токенизатора GPT-2:

# https://github.com/openai/tiktoken/blob/main/tiktoken_ext/openai_public.py#L23
GPT2_TOKENIZER_REGEX = \
    r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""

segments = regex.findall(GPT2_TOKENIZER_REGEX, string)
# → ["I", "'ll", " say", " supercalifragilisticexpialidocious", "!"]
compression_ratio = get_compression_ratio(string, segments)   # → 44 байта / 5 сегментов = 8.8

Чтобы превратить это в Tokenizer, нужно сопоставить каждому сегменту целое число. Но есть проблемы:

Слов огромное количество (как и символов Unicode).
Многие слова редки, и модель мало о них узнает.
Такой подход не даёт фиксированного размера словаря.

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

Byte Pair Encoding (BPE)

(Википедия.) Алгоритм BPE предложил Филип Гейдж в 1994 году для сжатия данных. Его адаптировали для нейросетевого машинного перевода (Sennrich et al., 2016); до этого в статьях использовали пословную токенизацию. Затем BPE взяли в GPT-2.

Основная идея: обучить токенизатор на сыром тексте, чтобы словарь определился автоматически. Интуиция: частые последовательности символов представляются одним токеном, редкие — многими.

В статье GPT-2 текст сначала разбивали на начальные сегменты пословной токенизацией, а затем запускали исходный алгоритм BPE на каждом сегменте. Набросок: начинаем с того, что каждый байт — токен, и раз за разом сливаем самую частую пару соседних токенов.

Обучение токенизатора

@dataclass(frozen=True)
class BPETokenizerParams:
    """Всё, что нужно, чтобы задать BPETokenizer."""
    vocab: dict[int, bytes]             # индекс -> байты
    merges: dict[tuple[int, int], int]  # (индекс1, индекс2) -> новый индекс


def merge(indices: list[int], pair: tuple[int, int], new_index: int) -> list[int]:
    """Вернуть `indices`, где все вхождения пары `pair` заменены на `new_index`."""
    new_indices = []
    i = 0
    while i < len(indices):
        if i + 1 < len(indices) and indices[i] == pair[0] and indices[i + 1] == pair[1]:
            new_indices.append(new_index)
            i += 2
        else:
            new_indices.append(indices[i])
            i += 1
    return new_indices


def train_bpe(string: str, num_merges: int) -> BPETokenizerParams:
    # Начинаем со списка байтов строки
    indices = list(map(int, string.encode("utf-8")))
    merges: dict[tuple[int, int], int] = {}                       # (индекс1, индекс2) => индекс слияния
    vocab: dict[int, bytes] = {x: bytes([x]) for x in range(256)}  # индекс -> байты

    for i in range(num_merges):
        # Считаем, сколько раз встречается каждая пара соседних токенов
        counts = defaultdict(int)
        for index1, index2 in zip(indices, indices[1:]):  # для каждой соседней пары
            counts[(index1, index2)] += 1

        # Находим самую частую пару
        pair = max(counts, key=counts.get)
        index1, index2 = pair

        # Сливаем эту пару
        new_index = 256 + i
        merges[pair] = new_index
        vocab[new_index] = vocab[index1] + vocab[index2]
        indices = merge(indices, pair, new_index)

    return BPETokenizerParams(vocab=vocab, merges=merges)


params = train_bpe("the cat in the hat", num_merges=3)

Что происходит на трёх шагах (значения из трассы лекции):

Шаг	Самая частая пара	Новый токен	Последовательность после слияния
старт	—	—	18 байтов `[116, 104, 101, 32, 99, 97, 116, …]`
1	`(116, 104)` = `t`+`h` (встречается 2 раза)	256 = `th`	`[256, 101, 32, 99, 97, 116, 32, 105, 110, 32, 256, 101, 32, 104, 97, 116]`
2	`(256, 101)` = `th`+`e`	257 = `the`	`[257, 32, 99, 97, 116, 32, 105, 110, 32, 257, 32, 104, 97, 116]`
3	`(257, 32)` = `the`+пробел	258 = `the`	`[258, 99, 97, 116, 32, 105, 110, 32, 258, 104, 97, 116]`

(Пара (97, 116) = at тоже встречается дважды, но max возвращает первую пару с максимальным счётом.)

Использование токенизатора. Теперь новый текст можно закодировать:

class BPETokenizer(Tokenizer):
    """Токенизатор BPE по заданным слияниям и словарю."""
    def __init__(self, params: BPETokenizerParams):
        self.params = params

    def encode(self, string: str) -> list[int]:
        indices = list(map(int, string.encode("utf-8")))
        # Замечание: это очень медленная реализация
        for pair, new_index in self.params.merges.items():
            indices = merge(indices, pair, new_index)
        return indices

    def decode(self, indices: list[int]) -> str:
        bytes_list = list(map(self.params.vocab.get, indices))
        string = b"".join(bytes_list).decode("utf-8")
        return string

tokenizer = BPETokenizer(params)
indices = tokenizer.encode("the quick brown fox")
# → [258, 113, 117, 105, 99, 107, 32, 98, 114, 111, 119, 110, 32, 102, 111, 120]
#    258 = b"the ", дальше по байту на символ
assert tokenizer.decode(indices) == "the quick brown fox"

В задании 1 вы пойдёте дальше:

encode() сейчас перебирает все слияния. Перебирайте только те, что имеют значение.
Находите и сохраняйте специальные токены (например, <|endoftext|>).
Используйте предразбиение (pre-tokenization), например регулярное выражение GPT-2.
Постарайтесь сделать реализацию как можно быстрее.

Итоги

Токенизатор: строки ⟷ токены (индексы).
Посимвольная, побайтовая и пословная токенизации сильно неоптимальны.
BPE — эффективная эвристика, опирающаяся на статистику корпуса.
Токенизация — необходимое зло; может быть, однажды мы будем работать прямо с байтами…

В следующий раз: строительные блоки PyTorch, учёт ресурсов.

Лекция 2. PyTorch и учёт ресурсов

Прошлая лекция: обзор, токенизация.

Обзор этой лекции:

Обсудим все примитивы, нужные для обучения модели.
Пойдём снизу вверх: от тензоров к моделям, оптимизаторам и циклу обучения.
Будем пристально следить за эффективностью (использованием ресурсов).

В частности, учитываем два вида ресурсов:

память (ГБ);
вычисления (FLOPs).

Мотивирующие вопросы

Сделаем прикидку «на салфетке».

Вопрос: сколько времени займёт обучение модели на 70 млрд параметров на 15 трлн токенов на 1024 H100?

total_flops = 6 * 70e9 * 15e12          # → 6.3e24 FLOPs
assert h100_flop_per_sec == 1979e12 / 2  # пиковая скорость H100 в BF16 без разреженности
mfu = 0.5                                # считаем, что используем половину пика
flops_per_day = h100_flop_per_sec * mfu * 1024 * 60 * 60 * 24   # → 4.38e22 FLOPs в сутки
days = total_flops / flops_per_day       # → ≈ 144 дня

Вопрос: какую самую большую модель можно обучить на 8 H100 с AdamW (наивно)?

h100_bytes = 80e9                         # 80 ГБ памяти у одной H100
bytes_per_parameter = 4 + 4 + (4 + 4)     # параметры, градиенты, состояние оптимизатора → 16
num_parameters = (h100_bytes * 8) / bytes_per_parameter   # → 4e10, то есть 40 млрд

Оговорка 1: мы наивно храним параметры и градиенты во float32. Можно хранить их в bf16 (2 + 2) и держать дополнительную копию параметров во float32 (4). Память это не экономит, зато быстрее (ZeRO).
Оговорка 2: активации не учтены (зависят от размера батча и длины последовательности).

Это грубая прикидка «на обороте конверта».

Трансформер мы разбирать не будем. Есть отличные объяснения: условие задания 1, математическое описание, Illustrated Transformer, Illustrated GPT-2. Вместо этого будем работать с моделями попроще.

Какие знания унести:

Механика: прямолинейная (просто PyTorch).
Образ мышления: учёт ресурсов (не забывайте его делать).
Интуиции: широкими мазками (без больших моделей).

Учёт памяти

Основы тензоров

Тензоры — базовый строительный блок для хранения всего: параметров, градиентов, состояния оптимизатора, данных, активаций (документация PyTorch). Создавать их можно по-разному:

x = torch.tensor([[1., 2, 3], [4, 5, 6]])  # из списка
x = torch.zeros(4, 8)   # матрица 4×8 из нулей
x = torch.ones(4, 8)    # матрица 4×8 из единиц
x = torch.randn(4, 8)   # матрица 4×8 из независимых выборок Normal(0, 1)

Выделить память, но не инициализировать значения:

x = torch.empty(4, 8)   # матрица 4×8 с неинициализированными значениями (там мусор вроде 4e-34 и -2.8e12)

…потому что вы хотите задать значения позже своей логикой:

nn.init.trunc_normal_(x, mean=0, std=1, a=-2, b=2)   # усечённое нормальное распределение на [-2, 2]

Память тензоров

Почти всё (параметры, градиенты, активации, состояния оптимизатора) хранится в виде чисел с плавающей точкой.

float32 (Википедия). [Схема битов: 1 бит знака, 8 бит экспоненты, 23 бита мантиссы.] Тип float32 (он же fp32, одинарная точность) — по умолчанию. Традиционно в научных вычислениях float32 — базовый уровень; в некоторых случаях берут двойную точность (float64). В глубоком обучении можно быть гораздо небрежнее.

Посмотрим на расход памяти. Память определяется (1) числом значений и (2) типом данных каждого значения:

x = torch.zeros(4, 8)
assert x.dtype == torch.float32          # тип по умолчанию
assert x.numel() == 4 * 8
assert x.element_size() == 4             # float занимает 4 байта
assert get_memory_usage(x) == 4 * 8 * 4  # 128 байт

# Одна матрица слоя прямого распространения GPT-3:
assert get_memory_usage(torch.empty(12288 * 4, 12288)) == 2304 * 1024 * 1024  # 2.3 ГБ

…это много!

float16 (Википедия). [Схема: 1 бит знака, 5 бит экспоненты, 10 бит мантиссы.] Тип float16 (fp16, половинная точность) сокращает память:

x = torch.zeros(4, 8, dtype=torch.float16)
assert x.element_size() == 2

Однако динамический диапазон (особенно для маленьких чисел) неважный:

x = torch.tensor([1e-8], dtype=torch.float16)   # → [0.0]
assert x == 0  # потеря значимости (underflow)!

Если такое случится при обучении, возможна нестабильность.

bfloat16 (Википедия). [Схема: 1 бит знака, 8 бит экспоненты, 7 бит мантиссы.] Google Brain разработала bfloat (brain floating point) в 2018 году, чтобы решить эту проблему. bfloat16 занимает столько же памяти, сколько float16, но имеет тот же динамический диапазон, что float32! Подвох только в том, что разрешение хуже, но для глубокого обучения это менее важно.

x = torch.tensor([1e-8], dtype=torch.bfloat16)   # → [1.0012e-08]
assert x != 0  # потери значимости нет!

Сравним динамические диапазоны и расход памяти типов:

Тип	Байт	Наибольшее число	Наименьшее нормальное	eps (шаг около 1)
float32	4	$3.4 \cdot 10^{38}$	$1.18 \cdot 10^{-38}$	$1.19 \cdot 10^{-7}$
float16	2	65 504	$6.1 \cdot 10^{-5}$	$9.8 \cdot 10^{-4}$
bfloat16	2	$3.39 \cdot 10^{38}$	$1.18 \cdot 10^{-38}$	$7.8 \cdot 10^{-3}$

(Значения — вывод torch.finfo для каждого типа в трассе лекции.)

fp8. В 2022 году стандартизовали FP8 — под нагрузки машинного обучения (руководство NVIDIA). H100 поддерживает два варианта FP8: E4M3 (диапазон [-448, 448]) и E5M2 ([-57344, 57344]). Ссылка: FP8 Formats for Deep Learning.

Что это значит для обучения:

Обучение во float32 работает, но требует много памяти.
Обучение в fp8, float16 и даже bfloat16 рискованно: возможна нестабильность.
Решение (позже): смешанная точность, см. раздел о ней в конце лекции.

Учёт вычислений

Тензоры на GPU

По умолчанию тензоры хранятся в памяти CPU:

x = torch.zeros(32, 32)
assert x.device == torch.device("cpu")

Но чтобы воспользоваться массовым параллелизмом GPU, их надо перенести в память GPU. [Схема: CPU и его RAM, GPU и его память, связаны шиной.] Сначала проверим, есть ли у нас GPU:

num_gpus = torch.cuda.device_count()
for i in range(num_gpus):
    properties = torch.cuda.get_device_properties(i)   # имя, объём памяти и т. д.

memory_allocated = torch.cuda.memory_allocated()

# Переносим тензор в память GPU (устройство 0)
y = x.to("cuda:0")
assert y.device == torch.device("cuda", 0)

# Или создаём тензор прямо на GPU
z = torch.zeros(32, 32, device="cuda:0")

new_memory_allocated = torch.cuda.memory_allocated()
memory_used = new_memory_allocated - memory_allocated
assert memory_used == 2 * (32 * 32 * 4)  # две матрицы 32×32 из 4-байтовых float

Операции с тензорами

Большинство тензоров получается операциями над другими тензорами. У каждой операции есть последствия для памяти и вычислений.

Хранение тензоров. Что такое тензор в PyTorch? Это указатель в выделенную память… с метаданными, описывающими, как добраться до любого элемента (документация о stride). [Рисунок: двумерный тензор, разложенный в одномерный массив, и шаги по строкам и столбцам.]

x = torch.tensor([
    [0., 1, 2, 3],
    [4, 5, 6, 7],
    [8, 9, 10, 11],
    [12, 13, 14, 15],
])
# Чтобы перейти к следующей строке (dim 0), пропустить 4 элемента хранилища
assert x.stride(0) == 4
# Чтобы перейти к следующему столбцу (dim 1), пропустить 1 элемент
assert x.stride(1) == 1
# Чтобы найти элемент:
r, c = 1, 2
index = r * x.stride(0) + c * x.stride(1)   # → 6
assert index == 6

Срезы. Многие операции просто дают другое представление (view) тензора. Копия не создаётся, поэтому изменение одного тензора меняет и другой.

x = torch.tensor([[1., 2, 3], [4, 5, 6]])

y = x[0]               # строка 0 → [1., 2., 3.]
assert same_storage(x, y)
y = x[:, 1]            # столбец 1 → [2., 5.]
assert same_storage(x, y)
y = x.view(3, 2)       # матрица 2×3 как 3×2 → [[1, 2], [3, 4], [5, 6]]
assert same_storage(x, y)
y = x.transpose(1, 0)  # транспонирование → [[1, 4], [2, 5], [3, 6]]
assert same_storage(x, y)

# Проверим, что изменение x меняет и y
x[0][0] = 100          # → x = [[100, 2, 3], [4, 5, 6]], y = [[100, 4], [2, 5], [3, 6]]
assert y[0][0] == 100

Заметьте: некоторые представления — несмежные (non-contiguous) записи, и дальнейшие view для них невозможны:

x = torch.tensor([[1., 2, 3], [4, 5, 6]])
y = x.transpose(1, 0)
assert not y.is_contiguous()
try:
    y.view(2, 3)
    assert False
except RuntimeError as e:
    assert "view size is not compatible with input tensor's size and stride" in str(e)

# Можно сначала сделать тензор смежным:
y = x.transpose(1, 0).contiguous().view(2, 3)   # → [[1, 4, 2], [5, 3, 6]]
assert not same_storage(x, y)

Представления бесплатны, копирование стоит и (дополнительной) памяти, и вычислений.

Поэлементные операции применяют операцию к каждому элементу и возвращают (новый) тензор той же формы:

x = torch.tensor([1, 4, 9])
assert torch.equal(x.pow(2), torch.tensor([1, 16, 81]))
assert torch.equal(x.sqrt(), torch.tensor([1, 2, 3]))
assert torch.equal(x.rsqrt(), torch.tensor([1, 1 / 2, 1 / 3]))  # i -> 1/sqrt(x_i)

assert torch.equal(x + x, torch.tensor([2, 8, 18]))
assert torch.equal(x * 2, torch.tensor([2, 8, 18]))
assert torch.equal(x / 0.5, torch.tensor([2, 8, 18]))

# triu берёт верхнетреугольную часть матрицы
x = torch.ones(3, 3).triu()   # → [[1, 1, 1], [0, 1, 1], [0, 0, 1]]

Это полезно для вычисления причинной маски внимания, где M[i, j] — вклад i в j.

Умножение матриц — хлеб с маслом глубокого обучения:

x = torch.ones(16, 32)
w = torch.ones(32, 2)
y = x @ w
assert y.size() == torch.Size([16, 2])

Обычно операции выполняются для каждого примера в батче и каждого токена в последовательности [рисунок: тензор «батч × последовательность × признаки»]:

x = torch.ones(4, 8, 16, 32)
w = torch.ones(32, 2)
y = x @ w
assert y.size() == torch.Size([4, 8, 16, 2])

Здесь мы перебираем значения первых двух измерений x и умножаем на w.

einops

Традиционный код PyTorch:

x = torch.ones(2, 2, 3)  # батч, последовательность, скрытое измерение
y = torch.ones(2, 2, 3)  # батч, последовательность, скрытое измерение
z = x @ y.transpose(-2, -1)  # батч, последовательность, последовательность → все элементы равны 3

Легко запутаться в измерениях (что такое -2 и -1?)…

Einops — библиотека для работы с тензорами с именованными измерениями. Она вдохновлена эйнштейновской нотацией суммирования (Эйнштейн, 1916) (учебник einops).

jaxtyping. Как следить за измерениями тензоров?

# Старый способ:
x = torch.ones(2, 2, 1, 3)  # batch seq heads hidden

# Новый способ (jaxtyping):
x: Float[torch.Tensor, "batch seq heads hidden"] = torch.ones(2, 2, 1, 3)

Замечание: это только документация (ничего не проверяется).

einsum — обобщённое умножение матриц с хорошим учётом измерений:

x: Float[torch.Tensor, "batch seq1 hidden"] = torch.ones(2, 3, 4)
y: Float[torch.Tensor, "batch seq2 hidden"] = torch.ones(2, 3, 4)

# Старый способ:
z = x @ y.transpose(-2, -1)  # батч, последовательность, последовательность

# Новый способ (einops):
z = einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")   # → все элементы равны 4

Измерения, не названные в выходе, суммируются. Можно использовать ... для трансляции (broadcasting) по любому числу измерений:

z = einsum(x, y, "... seq1 hidden, ... seq2 hidden -> ... seq1 seq2")

reduce. Свернуть один тензор какой-то операцией (sum, mean, max, min):

x: Float[torch.Tensor, "batch seq hidden"] = torch.ones(2, 3, 4)
# Старый способ:
y = x.sum(dim=-1)
# Новый способ (einops):
y = reduce(x, "... hidden -> ...", "sum")   # → [[4, 4, 4], [4, 4, 4]]

rearrange. Иногда одно измерение на самом деле представляет два… и вы хотите работать с одним из них:

x: Float[torch.Tensor, "batch seq total_hidden"] = torch.ones(2, 3, 8)
# ...где total_hidden — «сплющенное» представление heads * hidden1
w: Float[torch.Tensor, "hidden1 hidden2"] = torch.ones(4, 4)

# Разбиваем total_hidden на два измерения (heads и hidden1):
x = rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2)
# Применяем преобразование w:
x = einsum(x, w, "... hidden1, hidden1 hidden2 -> ... hidden2")
# Склеиваем heads и hidden2 обратно:
x = rearrange(x, "... heads hidden2 -> ... (heads hidden2)")

FLOPs операций

Разобрав все операции, посмотрим на их вычислительную стоимость.

Операция с плавающей точкой (FLOP) — базовая операция вроде сложения ( $x + y$ ) или умножения ( $x y$ ).

Два ужасно путающих сокращения (произносятся одинаково!):

FLOPs — число операций с плавающей точкой (мера объёма вычислений);
FLOP/s — операций с плавающей точкой в секунду (пишут и FLOPS), мера скорости железа.

Интуиции

Обучение GPT-3 (2020) заняло $3.14 \cdot 10^{23}$ FLOPs (статья).
Обучение GPT-4 (2023), по предположениям, — $2 \cdot 10^{25}$ FLOPs (статья).
Указ президента США: о любой фундаментальной модели, обученной с $\ge 10^{26}$ FLOPs, нужно сообщать правительству (отменён в 2025 году).
Пиковая производительность A100 — 312 TFLOP/s (спецификация).
Пиковая производительность H100 — 1979 TFLOP/s с разреженностью, без неё — 50% от этого (спецификация).

8 H100 за 2 недели:

total_flops = 8 * (60 * 60 * 24 * 7) * h100_flop_per_sec   # → 4.79e21

(Примечание переводчика: в коде лекции умножено на одну неделю, $60 \cdot 60 \cdot 24 \cdot 7$ секунд; за две недели получится вдвое больше, $\approx 9.6 \cdot 10^{21}$ .)

Линейная модель. Для мотивации представим линейную модель:

у нас $n$ точек;
каждая точка — $d$ -мерная;
линейная модель переводит каждый $d$ -мерный вектор в $k$ выходов.

if torch.cuda.is_available():
    B = 16384  # число точек
    D = 32768  # размерность
    K = 8192   # число выходов
else:
    B = 1024
    D = 256
    K = 64

device = get_device()
x = torch.ones(B, D, device=device)
w = torch.randn(D, K, device=device)
y = x @ w

На каждую тройку $(i, j, k)$ — одно умножение (x[i][j] * w[j][k]) и одно сложение:

actual_num_flops = 2 * B * D * K

FLOPs других операций

Поэлементная операция над матрицей $m \times n$ требует $O(mn)$ FLOPs.
Сложение двух матриц $m \times n$ требует $mn$ FLOPs.

В целом ни одна другая операция, которую вы встретите в глубоком обучении, не стоит столько, сколько умножение достаточно больших матриц.

Интерпретация:

$B$ — число точек данных;
$DK$ — число параметров;
FLOPs прямого прохода = 2 × (число токенов) × (число параметров).

Оказывается, это обобщается на трансформеры (в первом приближении).

Как наши подсчёты FLOPs переводятся во время по часам (секунды)? Замерим!

def time_matmul(a: torch.Tensor, b: torch.Tensor) -> float:
    """Вернуть число секунд, за которое выполняется a @ b."""
    # Дождаться завершения предыдущих потоков CUDA
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    def run():
        # Выполнить операцию
        a @ b
        # Дождаться завершения потоков CUDA
        if torch.cuda.is_available():
            torch.cuda.synchronize()

    # Повторить операцию num_trials раз
    num_trials = 5
    total_time = timeit.timeit(run, number=num_trials)
    return total_time / num_trials

actual_time = time_matmul(x, w)
actual_flop_per_sec = actual_num_flops / actual_time

У каждого GPU есть паспорт с пиковой производительностью (A100, H100). Заметьте: FLOP/s сильно зависят от типа данных!

def get_promised_flop_per_sec(device: str, dtype: torch.dtype) -> float:
    """Вернуть пиковые FLOP/s устройства `device` для типа `dtype`."""
    properties = torch.cuda.get_device_properties(device)
    if "A100" in properties.name:
        if dtype == torch.float32:
            return 19.5e12
        if dtype in (torch.bfloat16, torch.float16):
            return 312e12
    if "H100" in properties.name:
        if dtype == torch.float32:
            return 67.5e12
        if dtype in (torch.bfloat16, torch.float16):
            return 1979e12 / 2  # 1979 — для разреженных, для плотных вдвое меньше
    raise ValueError(f"Unknown device: {device}")

Утилизация FLOPs модели (MFU). Определение: (фактические FLOP/s) / (обещанные FLOP/s), без учёта коммуникаций и накладных расходов.

mfu = actual_flop_per_sec / promised_flop_per_sec

Обычно MFU ≥ 0.5 — это весьма хорошо (и будет выше, если доминируют матричные умножения).

Сделаем то же в bfloat16:

x = x.to(torch.bfloat16)
w = w.to(torch.bfloat16)
bf16_actual_time = time_matmul(x, w)
bf16_actual_flop_per_sec = actual_num_flops / bf16_actual_time
bf16_promised_flop_per_sec = get_promised_flop_per_sec(device, x.dtype)
bf16_mfu = bf16_actual_flop_per_sec / bf16_promised_flop_per_sec

Заметьте: в bfloat16 фактические FLOP/s выше, чем во float32. MFU здесь довольно низкий — вероятно, обещанные FLOPs несколько оптимистичны.

(Примечание переводчика: трасса в репозитории снята на CPU, поэтому чисел MFU для GPU в ней нет. На RTX 4080 повторите замер сами: пик BF16 у неё 97.5 TFLOP/s, FP32 — 48.7, см. урок 5.7.1.)

Итоги

Доминируют умножения матриц: $2mnp$ FLOPs.
FLOP/s зависят от железа (H100 ≫ A100) и типа данных (bfloat16 ≫ float32).
MFU = (фактические FLOP/s) / (обещанные FLOP/s).

Градиенты

До сих пор мы строили тензоры (параметры или данные) и прогоняли их через операции (прямой проход). Теперь вычислим градиент (обратный проход).

Простой пример — линейная модель $y = 0.5\,(x \cdot w - 5)^2$ :

# Прямой проход: считаем loss
x = torch.tensor([1., 2, 3])
w = torch.tensor([1., 1, 1], requires_grad=True)  # хотим градиент
pred_y = x @ w
loss = 0.5 * (pred_y - 5).pow(2)

# Обратный проход: считаем градиенты
loss.backward()
assert loss.grad is None
assert pred_y.grad is None
assert x.grad is None
assert torch.equal(w.grad, torch.tensor([1, 2, 3]))

(Проверка: $x \cdot w = 6$ , производная loss по $w$ равна $(6 - 5) \cdot x = x = [1, 2, 3]$ .)

FLOPs градиентов. Вернёмся к линейной модели, теперь с двумя слоями:

x = torch.ones(B, D, device=device)
w1 = torch.randn(D, D, device=device, requires_grad=True)
w2 = torch.randn(D, K, device=device, requires_grad=True)

# Модель: x --w1--> h1 --w2--> h2 -> loss
h1 = x @ w1
h2 = h1 @ w2
loss = h2.pow(2).mean()

Число FLOPs прямого прохода:

умножить x[i][j] * w1[j][k], прибавить к h1[i][k];
умножить h1[i][j] * w2[j][k], прибавить к h2[i][k].

num_forward_flops = (2 * B * D * D) + (2 * B * D * K)   # на CPU-размерах: 167 772 160

Сколько FLOPs у обратного прохода?

h1.retain_grad()  # для отладки
h2.retain_grad()  # для отладки
loss.backward()

Напомним: h1.grad = d loss / d h1, h2.grad = d loss / d h2, w1.grad = d loss / d w1, w2.grad = d loss / d w2. Сосредоточимся на параметре w2 и применим цепное правило.

num_backward_flops = 0

# w2.grad[j,k] = sum_i h1[i,j] * h2.grad[i,k]
assert w2.grad.size() == torch.Size([D, K])
assert h1.size() == torch.Size([B, D])
assert h2.grad.size() == torch.Size([B, K])
# Для каждой тройки (i, j, k) — умножение и сложение
num_backward_flops += 2 * B * D * K

# h1.grad[i,j] = sum_k w2[j,k] * h2.grad[i,k]
assert h1.grad.size() == torch.Size([B, D])
assert w2.size() == torch.Size([D, K])
assert h2.grad.size() == torch.Size([B, K])
# Для каждой тройки (i, j, k) — умножение и сложение
num_backward_flops += 2 * B * D * K

# Это было только для w2 (D*K параметров).
# То же можно сделать для w1 (D*D параметров), хотя x.grad не нужен.
num_backward_flops += (2 + 2) * B * D * D

Хорошая визуализация — статья Дмитрия Богданова о «FLOPs-исчислении» [анимация: на каждый параметр приходится 2 FLOPs вперёд и 4 FLOPs назад].

Собираем вместе:

прямой проход: 2 × (число точек данных) × (число параметров) FLOPs;
обратный проход: 4 × (число точек данных) × (число параметров) FLOPs;
итого: 6 × (число точек данных) × (число параметров) FLOPs.

Модели

Параметры модуля

Параметры модели хранятся в PyTorch как объекты nn.Parameter:

input_dim = 16384
output_dim = 32

w = nn.Parameter(torch.randn(input_dim, output_dim))
assert isinstance(w, torch.Tensor)   # ведёт себя как тензор
assert type(w.data) == torch.Tensor  # доступ к лежащему внутри тензору

Инициализация параметров. Посмотрим, что получится:

x = nn.Parameter(torch.randn(input_dim))
output = x @ w     # → значения порядка сотен: 18.9, 151.1, -167.9, 263.8, ...
assert output.size() == torch.Size([output_dim])

Заметьте: каждый элемент output растёт как sqrt(input_dim). Большие значения могут привести к взрыву градиентов и нестабильному обучению.

Нам нужна инициализация, не зависящая от input_dim. Для этого просто масштабируем на 1/sqrt(input_dim):

w = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))
output = x @ w     # → значения порядка единицы: -1.53, -0.87, -1.75, 0.19, ...

Теперь каждый элемент output порядка константы. С точностью до множителя это инициализация Ксавье (статья, обсуждение на StackExchange). Для пущей надёжности усекаем нормальное распределение до $[-3, 3]$ , чтобы исключить выбросы:

w = nn.Parameter(nn.init.trunc_normal_(torch.empty(input_dim, output_dim),
                                       std=1 / np.sqrt(input_dim), a=-3, b=3))

Своя модель

Построим простую глубокую линейную модель на nn.Parameter:

class Linear(nn.Module):
    """Простой линейный слой."""
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x @ self.weight


class Cruncher(nn.Module):
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([
            Linear(dim, dim)
            for i in range(num_layers)
        ])
        self.final = Linear(dim, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Применяем линейные слои
        B, D = x.size()
        for layer in self.layers:
            x = layer(x)

        # Применяем финальную голову
        x = self.final(x)
        assert x.size() == torch.Size([B, 1])

        # Убираем последнее измерение
        x = x.squeeze(-1)
        assert x.size() == torch.Size([B])

        return x


D = 64  # размерность
num_layers = 2
model = Cruncher(dim=D, num_layers=num_layers)

param_sizes = [(name, param.numel()) for name, param in model.state_dict().items()]
assert param_sizes == [
    ("layers.0.weight", D * D),
    ("layers.1.weight", D * D),
    ("final.weight", D),
]
num_parameters = get_num_parameters(model)
assert num_parameters == (D * D) + (D * D) + D

# Не забудьте перенести модель на GPU
device = get_device()
model = model.to(device)

# Прогоняем модель на данных
B = 8  # размер батча
x = torch.randn(B, D, device=device)
y = model(x)
assert y.size() == torch.Size([B])

Цикл обучения и лучшие практики

О случайности

Случайность возникает во многих местах: инициализация параметров, dropout, порядок данных и т. д. Для воспроизводимости рекомендуем всегда передавать отдельный random seed для каждого использования случайности. Детерминированность особенно полезна при отладке — так можно выследить ошибку.

Seed задают в трёх местах — делайте это во всех сразу для надёжности:

# Torch
seed = 0
torch.manual_seed(seed)

# NumPy
import numpy as np
np.random.seed(seed)

# Python
import random
random.seed(seed)

Загрузка данных

В языковом моделировании данные — последовательность целых чисел (выход токенизатора). Их удобно сериализовать как массивы numpy (это делает токенизатор):

orig_data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=np.int32)
orig_data.tofile("data.npy")

Загрузить их обратно можно как массивы numpy. Не хочется загружать все данные в память сразу (данные LLaMA — 2,8 ТБ). Используйте memmap, чтобы лениво подгружать в память только нужные части:

data = np.memmap("data.npy", dtype=np.int32)
assert np.array_equal(data, orig_data)

Загрузчик данных генерирует батч последовательностей для обучения:

def get_batch(data: np.array, batch_size: int, sequence_length: int, device: str) -> torch.Tensor:
    # Выбираем batch_size случайных позиций в data
    start_indices = torch.randint(len(data) - sequence_length, (batch_size,))
    assert start_indices.size() == torch.Size([batch_size])

    # Индексируем данные
    x = torch.tensor([data[start:start + sequence_length] for start in start_indices])
    assert x.size() == torch.Size([batch_size, sequence_length])

    # Закреплённая (pinned) память: по умолчанию тензоры CPU лежат в страничной памяти,
    # её можно явно закрепить
    if torch.cuda.is_available():
        x = x.pin_memory()

    # Это позволяет копировать x из CPU в GPU асинхронно
    x = x.to(device, non_blocking=True)
    return x

B = 2  # размер батча
L = 4  # длина последовательности
x = get_batch(data, batch_size=B, sequence_length=L, device=get_device())
assert x.size() == torch.Size([B, L])

Асинхронное копирование позволяет делать две вещи параллельно (здесь не сделано): подгружать следующий батч в CPU и обрабатывать x на GPU. Подробнее: блог NVIDIA, обсуждение на GitHub.

Оптимизатор

Вспомним нашу глубокую линейную модель (B = 2, D = 4, два слоя). Определим оптимизатор AdaGrad. Родословная:

momentum = SGD + экспоненциальное усреднение градиента;
AdaGrad = SGD + усреднение по квадрату градиента (статья);
RMSProp = AdaGrad + экспоненциальное усреднение квадрата градиента;
Adam = RMSProp + momentum.

class SGD(torch.optim.Optimizer):
    def __init__(self, params: Iterable[nn.Parameter], lr: float = 0.01):
        super(SGD, self).__init__(params, dict(lr=lr))

    def step(self):
        for group in self.param_groups:
            lr = group["lr"]
            for p in group["params"]:
                grad = p.grad.data
                p.data -= lr * grad


class AdaGrad(torch.optim.Optimizer):
    def __init__(self, params: Iterable[nn.Parameter], lr: float = 0.01):
        super(AdaGrad, self).__init__(params, dict(lr=lr))

    def step(self):
        for group in self.param_groups:
            lr = group["lr"]
            for p in group["params"]:
                # Состояние оптимизатора
                state = self.state[p]
                grad = p.grad.data

                # Сумма квадратов градиентов g2 = sum_{i<t} g_i^2
                g2 = state.get("g2", torch.zeros_like(grad))

                # Обновляем состояние оптимизатора
                g2 += torch.square(grad)
                state["g2"] = g2

                # Обновляем параметры
                p.data -= lr * grad / torch.sqrt(g2 + 1e-5)


model = Cruncher(dim=4, num_layers=2).to(get_device())
optimizer = AdaGrad(model.parameters(), lr=0.01)

# Считаем градиенты
x = torch.randn(2, 4, device=get_device())
y = torch.tensor([4., 5.], device=get_device())
pred_y = model(x)
loss = F.mse_loss(input=pred_y, target=y)
loss.backward()

# Делаем шаг: каждый вес сдвигается примерно на ±0.01
optimizer.step()

# Освобождаем память (необязательно)
optimizer.zero_grad(set_to_none=True)

(Почему «примерно на ±0.01»: на первом шаге g2 = grad², и сдвиг равен lr · grad / |grad| = ±lr. В трассе так и видно: например, вес −0.1253 стал −0.1353.)

Память

num_parameters = (D * D * num_layers) + D      # параметры → 36
num_activations = B * D * num_layers           # активации → 16
num_gradients = num_parameters                 # градиенты → 36
num_optimizer_states = num_parameters          # состояния оптимизатора → 36
# Всё вместе, во float32:
total_memory = 4 * (num_parameters + num_activations + num_gradients + num_optimizer_states)   # → 496 байт

Вычисления (на один шаг)

flops = 6 * B * num_parameters   # → 432

Трансформеры. Учёт для трансформера сложнее, но идея та же. Задание 1 попросит вас его сделать. Блог о расходе памяти при обучении трансформера: erees.dev. Блог о FLOPs трансформера: Адам Кассон.

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

Сгенерируем данные линейной функцией с весами $(0, 1, 2, \ldots, D-1)$ :

D = 16
true_w = torch.arange(D, dtype=torch.float32, device=get_device())
def get_batch(B: int) -> tuple[torch.Tensor, torch.Tensor]:
    x = torch.randn(B, D).to(get_device())
    true_y = x @ true_w
    return (x, true_y)


def train(name: str, get_batch,
          D: int, num_layers: int,
          B: int, num_train_steps: int, lr: float):
    model = Cruncher(dim=D, num_layers=0).to(get_device())
    optimizer = SGD(model.parameters(), lr=0.01)

    for t in range(num_train_steps):
        # Получаем данные
        x, y = get_batch(B=B)

        # Прямой проход (считаем loss)
        pred_y = model(x)
        loss = F.mse_loss(pred_y, y)

        # Обратный проход (считаем градиенты)
        loss.backward()

        # Обновляем параметры
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

# Базовый запуск
train("simple", get_batch, D=D, num_layers=0, B=4, num_train_steps=10, lr=0.01)
# Немного подбора гиперпараметров
train("simple", get_batch, D=D, num_layers=0, B=4, num_train_steps=10, lr=0.1)

(Примечание переводчика: в коде лекции train игнорирует аргументы num_layers и lr — внутри зашиты num_layers=0 и lr=0.01, поэтому два запуска отличаются только случайными данными. Исправьте сами: Cruncher(dim=D, num_layers=num_layers) и SGD(model.parameters(), lr=lr) — и сравните loss.)

Чекпоинты

Обучение языковых моделей длится долго и наверняка где-нибудь упадёт. Терять весь прогресс не хочется. Поэтому во время обучения полезно периодически сохранять на диск состояние модели и оптимизатора:

model = Cruncher(dim=64, num_layers=3).to(get_device())
optimizer = AdaGrad(model.parameters(), lr=0.01)

# Сохраняем чекпоинт
checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}
torch.save(checkpoint, "model_checkpoint.pt")

# Загружаем чекпоинт
loaded_checkpoint = torch.load("model_checkpoint.pt")

Обучение со смешанной точностью

У выбора типа данных (float32, bfloat16, fp8) есть компромиссы:

выше точность: точнее и стабильнее, больше памяти, больше вычислений;
ниже точность: менее точно и стабильно, меньше памяти, меньше вычислений.

Как взять лучшее из обоих миров? Решение: по умолчанию float32, но bfloat16 или fp8 — где возможно.

Конкретный план:

bfloat16 или fp8 — для прямого прохода (активации);
float32 — для остального (параметры, градиенты).

Ссылки: Mixed Precision Training; в PyTorch есть библиотека автоматической смешанной точности (AMP, руководство NVIDIA). Transformer Engine от NVIDIA поддерживает FP8 для линейных слоёв. FP8 повсеместно на всём обучении: FP8-LM. У нас эта тема — урок 5.4.1.

Вспомогательные функции лекции

def get_memory_usage(x: torch.Tensor):
    return x.numel() * x.element_size()   # число элементов × байт на элемент

def same_storage(x: torch.Tensor, y: torch.Tensor):
    # Два тензора делят одно хранилище, если у них совпадает адрес начала данных
    return x.untyped_storage().data_ptr() == y.untyped_storage().data_ptr()

def get_num_parameters(model: nn.Module) -> int:
    return sum(param.numel() for param in model.parameters())

def get_device(index: int = 0) -> torch.device:
    """По возможности GPU, иначе CPU."""
    if torch.cuda.is_available():
        return torch.device(f"cuda:{index}")
    else:
        return torch.device("cpu")

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

Перепишите train_bpe и BPETokenizer из лекции 1 и обучите на строке с русским текстом и эмодзи, например "мама мыла раму, 👋👋". Сколько слияний нужно, чтобы мама стала одним токеном? (Подсказка: одна кириллическая буква — 2 байта UTF-8, эмодзи — 4. Ответ: 👋 появляется на 7-м слиянии, мама — на 8-м. По дороге возникают токены из «половинок» символов, например b"\xd0\xbc\xd1" или b"\xf0\x9f" — их нельзя декодировать по отдельности, поэтому decode склеивает байты всех токенов и только потом переводит в строку.)
Посчитайте степень сжатия токенизатора GPT-2 (tiktoken, pip install tiktoken) на русском абзаце и на английском. Почему по-русски она ниже?
Повторите «прикидку на салфетке» лекции 2 для своей машины: сколько дней займёт обучение модели на 125 млн параметров на 2,5 млрд токенов на 4 × RTX 4080 при MFU 0.4 и пике BF16 97.5 TFLOP/s на карту? (Ответ: $6 \cdot 1.25 \cdot 10^8 \cdot 2.5 \cdot 10^9 / (4 \cdot 0.4 \cdot 97.5 \cdot 10^{12}) \approx 12\,000$ с, то есть около 3,3 часа.)
Запустите time_matmul на своей видеокарте для float32 и bfloat16 с B = 16384, D = 32768, K = 8192 (матрицы займут около 4 ГБ) и посчитайте MFU. Затем исправьте train (см. примечание) и сравните lr = 0.01 и lr = 0.1.

Словарик

EN	RU
executable lecture	исполняемая лекция
tokenizer, vocabulary size	токенизатор, размер словаря
compression ratio	степень сжатия (байт на токен)
code point	кодовая точка (номер символа Unicode)
pre-tokenization	предразбиение (претокенизация)
merge	слияние (правило BPE)
resource accounting	учёт ресурсов
FLOPs / FLOP/s	число операций / операций в секунду
MFU (model FLOPs utilization)	утилизация FLOPs модели
underflow, dynamic range	потеря значимости, динамический диапазон
view, stride, contiguous	представление, шаг, смежный (тензор)
mixed precision	смешанная точность
checkpoint	чекпоинт (сохранённое состояние)
pinned memory	закреплённая память

Перевод выполнен по лицензии оригинала (MIT (© 2025 Stanford University)); автор оригинала — Перси Лян (Percy Liang), Тацунори Хашимото (Tatsunori Hashimoto), Stanford University. Нашли неточность — сверяйтесь с оригиналом по ссылке выше. · обновлено 2026-10-04