CS336, задание 1 (basics): трансформер с нуля — путеводитель и перевод условий

Что нужно сделать в первом задании Stanford CS336 (весна 2026): byte-level BPE-токенизатор, трансформер с RMSNorm, SwiGLU и RoPE, cross-entropy, AdamW, расписание, цикл обучения, генерация и эксперименты на TinyStories и OpenWebText. Перевод всех условий задач, правила, связь с уроками курса и план выполнения на 4 × RTX 4080.

Источник — PDF-условие cs336_assignment1_basics.pdf (версия 26.0.3, весна 2026) из репозитория stanford-cs336/assignment1-basics, который распространяется под лицензией MIT. Все условия задач (блоки «Problem») переведены полностью, по порядку, с баллами; объяснительный текст между ними пересказан кратко — за подробностями и формулами идите в оригинал или в наши уроки по ссылкам. Код решений здесь намеренно не приводится. Разделы «Как выполнять на вашей машине» и «Подсказки» написаны переводчиком.

Что это и когда делать

Задание 1 — «собери GPT своими руками, не используя готовых слоёв». За ~45 страниц условия вы пишете byte-level BPE, все слои трансформера, функцию потерь, AdamW, расписание learning rate, загрузчик данных, чекпоинты, цикл обучения и генерацию, а потом обучаете модель на TinyStories и OpenWebText и проводите абляции. Код проверяют автотесты из папки tests/.

Когда: после модуля 4.6 «GPT с нуля». Наши домашки s4-hw01-bpe-tokenizer, s4-hw02-attention и s4-hw03-gpt-pytorch — уменьшенные версии частей 2–5 этого задания; сделайте их сначала.
Теория: лекции 1–2 (токенизация, учёт FLOPs и памяти) и лекция 3 (почему pre-norm, RMSNorm, SwiGLU, RoPE).
Объём: в Стэнфорде на задание отводят около двух недель (по расписанию весны 2026 года); на странице «Учимся у создателей» задания CS336 оценены в 20–60 часов каждое. Самые долгие части — быстрый BPE и отладка обучения.

Карта задания

Раздел условия	Что делаете	Баллы	Урок нашего курса
2. BPE-токенизатор	Юникод, обучение BPE, кодирование/декодирование, эксперименты	1 + 3 + 15 + 2 + 2 + 15 + 4	4.2.2, 4.2.3
3. Архитектура	Linear, Embedding, RMSNorm, SwiGLU, RoPE, softmax, внимание, блок, LM, учёт FLOPs	1 + 1 + 1 + 2 + 2 + 1 + 5 + 5 + 3 + 3 + 5	4.4, 4.5.2, 4.5.3
4. Обучение	Cross-entropy, SGD и learning rate, AdamW и его учёт, косинусное расписание, обрезка градиента	1 + 1 + 2 + 2 + 1 + 1	1.4.2, 4.6.3
5. Цикл обучения	Загрузчик, чекпоинты, скрипт обучения	2 + 1 + 4	4.6.3
6. Генерация	Температура, top-p	3	4.1.3
7. Эксперименты	Лог экспериментов, learning rate, батч, генерация, 4 абляции, OpenWebText, лидерборд	3 + 3 + 1 + 1 + 4 × 1 + 2 + 6	4.6.3, 3.4

Правила и как сдают

Что делаем. BPE-токенизатор (раздел 2), трансформерную LM (раздел 3), cross-entropy и AdamW (раздел 4), цикл обучения с сохранением и загрузкой состояния (раздел 5).

Что запускаем. Обучение BPE на TinyStories; кодирование датасета в последовательность id; обучение трансформера на TinyStories; генерацию и замер перплексии; обучение на OpenWebText и отправку результата в таблицу лидеров.

Что можно использовать. Каждый компонент пишется с нуля. Из torch.nn, torch.nn.functional и torch.optim нельзя брать ничего, кроме: torch.nn.Parameter; классов-контейнеров (Module, ModuleList, Sequential и т. п.); базового класса torch.optim.Optimizer. Остальной PyTorch разрешён. В сомнительных случаях спрашивайте себя, не нарушает ли это дух «с нуля».

ИИ-инструменты. ИИ разрешено использовать для высокоуровневых вопросов о понятиях и как справочник по API (сигнатуры функций, параметры библиотек). Нельзя использовать ИИ, чтобы реализовывать любую часть задания: ни агентов (Cursor Agents, Codex, Claude Code), ни ИИ-автодополнение (Cursor Tab, GitHub Copilot). Если всё же работаете с агентом, он должен использовать файл AGENTS.md из репозитория. Автодополнение имён функций без ИИ — можно. Авторы настоятельно советуют отключить ИИ-автодополнение: прежние студенты отмечали, что без него глубже разбираешься в материале.

Как устроен код.

cs336_basics/* — сюда пишете свой код; там пусто, структура на ваше усмотрение.
adapters.py — набор функций, которые должен уметь ваш код; для каждой (например, run_scaled_dot_product_attention) вы просто вызываете свою реализацию. В адаптерах не должно быть содержательной логики — только «клей».
test_*.py — тесты, которые нужно пройти; они вызывают адаптеры. Тесты не редактируйте.

Как сдают. make_submission.sh собирает архив. В Gradescope отправляют writeup.pdf (ответы на все письменные вопросы, набранные, а не от руки) и code.zip (весь код). Результат для лидерборда — pull request в assignment1-basics-leaderboard.

Данные. TinyStories [Eldan, Li, 2023] и выборка OpenWebText [Gokaslan et al., 2019] — два больших текстовых файла. Команды скачивания — в README.md репозитория (с Hugging Face).

По всему условию разбросаны синие блоки Low-Resource Tip — как сделать задание с малым числом GPU или вовсе без них (уменьшить датасет или модель, запустить на CPU или Apple Silicon). Решение авторов обучает на чипе Apple M4 Max с 36 ГБ памяти модель, пишущую довольно связные детские рассказы, менее чем за 5 минут на Metal GPU (MPS) и примерно за 30 минут на CPU. Ваша машина намного мощнее — эти советы пригодятся только для быстрой отладки.

Перевод условий задач

2. Byte-pair encoding (BPE) токенизатор

Сначала вы обучаете и реализуете byte-level BPE [Sennrich et al., 2016; Wang et al., 2019]: любая юникод-строка представляется последовательностью байтов, и BPE обучается на этой последовательности.

2.1. Стандарт Юникод. Юникод сопоставляет символам целые числа — кодовые точки (в Unicode 17.0, сентябрь 2025, — 159 801 символ в 172 письменностях). ord('ж') → 1078, chr(1078) → 'ж' (в оригинале пример с китайским иероглифом).

Problem (unicode1): понимание Юникода (1 балл)

(a) Какой символ Юникода возвращает chr(0)? Сдать: ответ в одно предложение.

(b) Чем строковое представление этого символа (__repr__()) отличается от его печатного представления? Сдать: одно предложение.

>>> chr(0)
>>> print(chr(0))
>>> "this is a test" + chr(0) + "string"
>>> print("this is a test" + chr(0) + "string")

Сдать: одно предложение.

2.2. Кодировки Юникода. Обучать токенизатор прямо на кодовых точках непрактично (словарь ~150 тыс., разреженный), поэтому берут кодировку — UTF-8, UTF-16 или UTF-32; UTF-8 доминирует (больше 98% веб-страниц). Байтовый словарь из 256 значений удобен, и при byte-level токенизации не бывает «неизвестных» токенов.

Problem (unicode2): кодировки Юникода (3 балла)

(a) Почему токенизатор лучше обучать на байтах UTF-8, а не UTF-16 или UTF-32? Полезно сравнить выход этих кодировок для разных строк. Сдать: 1–2 предложения.

(b) Вот (неправильная) функция, которая должна декодировать байтовую строку UTF-8 в юникод-строку. Почему она неверна? Приведите пример входа, на котором она даёт неправильный результат.

def decode_utf8_bytes_to_str_wrong(bytestring: bytes):
    return "".join([bytes([b]).decode("utf-8") for b in bytestring])

>>> decode_utf8_bytes_to_str_wrong("hello".encode("utf-8"))
'hello'

Сдать: пример байтовой строки, на которой функция ошибается, и одно предложение, почему она неверна.

(c) Приведите последовательность из двух байтов, которая не декодируется ни в какой символ Юникода. Сдать: пример и одно предложение объяснения.

2.3–2.4. Подсловная токенизация и обучение BPE. BPE итеративно заменяет («сливает») самую частую пару байтов новым токеном. Обучение состоит из трёх шагов:

Инициализация словаря: 256 байтов (плюс специальные токены).
Предтокенизация: грубое разбиение корпуса, чтобы считать пары внутри «слов» и не сливать, например, dog! и dog.. Используется регулярное выражение в стиле GPT-2: PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" (нужен пакет regex); в коде используйте re.finditer, чтобы не хранить все предтокены.
Вычисление слияний: пары через границы предтокенов не рассматриваются; при равенстве частот выбирается лексикографически бо́льшая пара (max([("A","B"), ("A","C"), ("B","ZZ"), ("BA","A")]) → ('BA','A')).

Специальные токены (например, <|endoftext|>) никогда не делятся и добавляются в словарь. В условии разобран пример из статьи Sennrich: корпус low low low low low / lower lower widest widest widest / newest ×6 даёт слияния ['s t', 'e st', 'o w', 'l ow', 'w est', 'n e', 'ne west', 'w i', 'wi d', 'wid est', 'low e', 'lowe r'].

2.5. Эксперименты с обучением BPE. Главное узкое место — предтокенизация; её стоит распараллелить через multiprocessing, разрезая корпус по границам специального токена (можно взять готовый pretokenization_example.py). Перед предтокенизацией вырежьте специальные токены и разбейте текст по ним (re.split с "|".join(...) и re.escape), чтобы слияния не шли через границы документов; сами специальные токены в подсчёт пар не входят. Шаг слияний ускоряется индексированием счётчиков пар и их инкрементальным обновлением (после слияния меняются только пары, пересекающиеся со слитой); в Python эта часть не параллелится. Советы: профилируйте (cProfile, py-spy) и сначала отлаживайтесь на малом «отладочном» датасете (например, на валидации TinyStories — 22 тыс. документов вместо 2,12 млн).

Problem (train_bpe): обучение BPE-токенизатора (15 баллов)

Сдать: функцию, которая по пути к текстовому файлу обучает (byte-level) BPE. Она должна принимать как минимум:

input_path: str — путь к файлу с обучающими данными;
vocab_size: int — положительное целое, максимальный итоговый размер словаря (включая начальные байты, токены от слияний и специальные токены);
special_tokens: list[str] — строки, добавляемые в словарь; при обучении они — жёсткие границы, через которые нельзя сливать, но в статистику слияний они не входят.

Функция возвращает:

vocab: dict[int, bytes] — словарь: id токена → байты токена;
merges: list[tuple[bytes, bytes]] — список слияний в порядке создания; каждый элемент (<token1>, <token2>) значит, что <token1> слит с <token2>.

Для проверки реализуйте адаптер adapters.run_train_bpe и запустите uv run pytest tests/test_train_bpe.py — должны проходить все тесты. По желанию (это может занять много времени) ключевые части можно написать на системном языке — C++ (cppyy или nanobind) или Rust (PyO3); тогда следите, какие операции копируют память, а какие читают прямо из памяти Python, и оставьте инструкции сборки (или пусть всё собирается из pyproject.toml). Учтите, что регулярное выражение GPT-2 плохо поддерживается большинством движков, а где поддерживается — работает слишком медленно; Oniguruma достаточно быстр и умеет отрицательный просмотр вперёд, а пакет regex для Python, пожалуй, ещё быстрее.

Problem (train_bpe_tinystories): BPE на TinyStories (2 балла)

(a) Обучите byte-level BPE на TinyStories с максимальным словарём 10 000. Добавьте в словарь специальный токен <|endoftext|>. Сохраните словарь и слияния на диск для анализа. Сколько времени и памяти заняло обучение? Какой токен в словаре самый длинный? Логично ли это? Ресурсы: не больше 30 минут (без GPU), не больше 30 ГБ RAM. Подсказка: можно уложиться меньше чем в 2 минуты, распараллелив предтокенизацию и учтя два факта: (1) <|endoftext|> разделяет документы в файлах; (2) <|endoftext|> обрабатывается как особый случай до применения слияний. Сдать: 1–2 предложения.

(b) Профилируйте код. Какая часть обучения токенизатора занимает больше всего времени? Сдать: 1–2 предложения.

Problem (train_bpe_expts_owt): BPE на OpenWebText (2 балла)

(a) Обучите byte-level BPE на OpenWebText с максимальным словарём 32 000. Сохраните словарь и слияния. Какой токен самый длинный? Логично ли? Ресурсы: не больше 12 часов (без GPU), не больше 100 ГБ RAM. Сдать: 1–2 предложения.

(b) Сравните токенизаторы, обученные на TinyStories и OpenWebText. Сдать: 1–2 предложения.

2.6. Кодирование и декодирование. Кодирование повторяет обучение: (1) предтокенизация, каждый предтокен — байты UTF-8; (2) слияния применяются к каждому предтокену отдельно в порядке их создания. Пример из условия: строка 'the cat ate', словарь {0: b' ', 1: b'a', 2: b'c', 3: b'e', 4: b'h', 5: b't', 6: b'th', 7: b' c', 8: b' a', 9: b'the', 10: b' at'}, слияния [(b't', b'h'), (b' ', b'c'), (b' ', b'a'), (b'th', b'e'), (b' a', b't')] дают [9, 7, 1, 5, 10, 3]. Специальные токены пользователя должны сохраняться целиком. Для больших файлов кодируйте поток по кускам, следя, чтобы токен не пересекал границу куска. При декодировании байты склеиваются и декодируются; неправильные последовательности заменяются символом U+FFFD (errors='replace').

Problem (tokenizer): реализация токенизатора (15 баллов)

Сдать: класс Tokenizer, который по словарю и списку слияний кодирует текст в id и декодирует id в текст. Он должен поддерживать пользовательские специальные токены (добавляя их в словарь, если их там нет). Рекомендуемый интерфейс:

def __init__(self, vocab, merges, special_tokens=None) — создать токенизатор по словарю (dict[int, bytes]), слияниям (list[tuple[bytes, bytes]]) и, необязательно, списку специальных токенов (list[str] | None);
def from_files(cls, vocab_filepath, merges_filepath, special_tokens=None) — метод класса: создать токенизатор из сохранённых словаря и слияний (в том формате, который выдаёт ваш код обучения);
def encode(self, text: str) -> list[int] — закодировать текст в id;
def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int] — по итерируемому набору строк (например, файловому объекту) лениво выдавать id; нужно для экономного по памяти кодирования больших файлов;
def decode(self, ids: list[int]) -> str — декодировать id в текст.

Реализуйте адаптер adapters.get_tokenizer и запустите uv run pytest tests/test_tokenizer.py — должны проходить все тесты.

Problem (tokenizer_experiments): эксперименты с токенизаторами (4 балла)

(a) Возьмите по 10 документов из TinyStories и OpenWebText. Закодируйте их своими токенизаторами (10K для TinyStories и 32K для OpenWebText). Какова степень сжатия каждого (байт на токен)? Сдать: 1–2 предложения.

(b) Что будет, если закодировать выборку OpenWebText токенизатором TinyStories? Сравните степень сжатия и/или опишите словами. Сдать: 1–2 предложения.

(c) Оцените пропускную способность токенизатора (например, байт в секунду). Сколько займёт кодирование датасета Pile (825 ГБ текста)? Сдать: 1–2 предложения.

(d) Закодируйте обучающие и валидационные части обоих датасетов в последовательности id — они понадобятся для обучения. Советуем сохранять id как массив NumPy типа uint16. Почему uint16 подходит? Сдать: 1–2 предложения.

3. Архитектура трансформерной LM

Модель принимает батч id формы (batch_size, sequence_length) и выдаёт распределение по словарю формы (batch_size, sequence_length, vocab_size). Схема: эмбеддинг токенов → num_layers pre-norm блоков → финальная нормализация → линейный слой «LM head» → логиты. Блок: $y = x + \mathrm{MultiHeadSelfAttention}(\mathrm{RMSNorm}(x))$ , затем $z = y + \mathrm{FFN}(\mathrm{RMSNorm}(y))$ . В разделе 3.2 условие советует работать с «батчевыми» измерениями через einsum/einops и помнить, что PyTorch хранит данные построчно (row-major).

Инициализация (3.3.1): веса Linear — $\mathcal N(0, \sigma^2 = \frac{2}{d_{in} + d_{out}})$ , обрезанное на $[-3\sigma, 3\sigma]$ ; эмбеддинги — $\mathcal N(0, 1)$ , обрезанное на $[-3, 3]$ ; RMSNorm — единицы. Используйте torch.nn.init.trunc_normal_. Свободных членов (bias) нет, как в большинстве современных LLM.

Problem (linear): линейный слой (1 балл)

Сдать: класс Linear, наследник torch.nn.Module, выполняющий линейное преобразование $y = Wx$ . Интерфейс — как у nn.Linear, но без bias. Рекомендуется:

def __init__(self, in_features, out_features, device=None, dtype=None) — in_features: int (последняя размерность входа), out_features: int (последняя размерность выхода), device, dtype — где и в каком типе хранить параметры;
def forward(self, x: torch.Tensor) -> torch.Tensor — применить преобразование.

Обязательно: наследуйте nn.Module; вызовите конструктор базового класса; храните параметр как $W$ (а не $W^\top$ ) в nn.Parameter; конечно, не используйте nn.Linear и nn.functional.linear. Инициализация — как выше, через trunc_normal_. Тест: адаптер adapters.run_linear (загрузите в модуль данные веса, например через Module.load_state_dict), затем uv run pytest -k test_linear.

Problem (embedding): слой эмбеддингов (1 балл)

Сдать: класс Embedding, наследник torch.nn.Module, выполняющий поиск эмбеддинга по индексу; интерфейс — как у nn.Embedding:

def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None) — размер словаря, размерность векторов ( $d_{model}$ ), устройство, тип;
def forward(self, token_ids: torch.Tensor) -> torch.Tensor — вернуть эмбеддинги для данных id.

Обязательно: наследуйте nn.Module; вызовите конструктор базового класса; матрица — nn.Parameter формы (vocab_size, d_model) (d_model — последнее измерение); не используйте nn.Embedding и nn.functional.embedding. Тест: adapters.run_embedding, uv run pytest -k test_embedding.

3.4.1. RMSNorm: $\mathrm{RMSNorm}(a_i) = \frac{a_i}{\mathrm{RMS}(a)} g_i$ , где $\mathrm{RMS}(a) = \sqrt{\frac{1}{d_{model}}\sum_i a_i^2 + \varepsilon}$ , $g$ — обучаемый «коэффициент усиления», $\varepsilon = 10^{-5}$ . Перед возведением в квадрат приводите вход к float32, а результат — обратно к исходному типу.

Problem (rmsnorm): RMSNorm (1 балл)

Сдать: RMSNorm как torch.nn.Module. Рекомендуется: def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None) и def forward(self, x) — вход формы (batch_size, sequence_length, d_model), выход той же формы. Не забудьте привести вход к torch.float32 перед нормализацией (и вернуть исходный тип). Тест: adapters.run_rmsnorm, uv run pytest -k test_rmsnorm.

3.4.2. Позиционная FFN: $\mathrm{FFN}(x) = W_2(\mathrm{SiLU}(W_1 x) \odot W_3 x)$ , $\mathrm{SiLU}(x) = x \cdot \sigma(x)$ , $d_{ff} \approx \frac{8}{3} d_{model}$ (можно округлить до кратного 64). Условие цитирует знаменитую фразу Шазира о «божественной благосклонности» — подробнее в лекции 3.

Problem (positionwise_feedforward): позиционная FFN (2 балла)

Сдать: FFN на SwiGLU — активация SiLU плюс вентиль GLU. Здесь для численной устойчивости можно использовать torch.sigmoid. $d_{ff}$ берите примерно $\frac{8}{3} d_{model}$ , но кратным 64 — для эффективной работы железа. Тест: adapters.run_swiglu, uv run pytest -k test_swiglu.

3.4.3. RoPE. Для запроса в позиции $i$ пары координат $(2k-1, 2k)$ поворачиваются на угол $\theta_{i,k} = \frac{i}{\Theta^{(2k-2)/d}}$ ; ключи — так же. Строить всю матрицу $d \times d$ не нужно; синусы и косинусы можно один раз посчитать и хранить в буфере (self.register_buffer(..., persistent=False)), общем для всех слоёв. Обучаемых параметров у слоя нет. Подробности и числовой пример — урок 4.5.2.

Problem (rope): RoPE (2 балла)

Сдать: класс RotaryPositionalEmbedding, применяющий RoPE ко входному тензору. Рекомендуется:

def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None) — $\Theta$ , размерность запросов и ключей, максимальная длина последовательности, устройство для буфера;
def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor — вход формы (..., seq_len, d_k), выход той же формы; число батчевых измерений произвольно; token_positions формы (..., seq_len) задаёт позиции токенов. По позициям выбирайте нужные строки (возможно, заранее посчитанных) тензоров cos и sin.

Тест: adapters.run_rope, uv run pytest -k test_rope.

3.4.4. Softmax и внимание. Softmax численно устойчив, если вычесть максимум. Внимание: $\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$ . Маска $M$ из True/False: True в $(i, j)$ значит «запрос $i$ смотрит на ключ $j$ »; для False к логиту до softmax прибавляют $-\infty$ .

Problem (softmax): softmax (1 балл)

Сдать: функцию softmax по измерению $i$ тензора: на входе тензор и номер измерения, на выходе тензор той же формы, где по $i$ -му измерению — нормированное распределение. Вычитайте максимум по $i$ -му измерению. Тест: adapters.run_softmax, uv run pytest -k test_softmax_matches_pytorch.

Problem (scaled_dot_product_attention): внимание со скалярным произведением (5 баллов)

Сдать: функцию scaled dot-product attention. Ключи и запросы формы (batch_size, ..., seq_len, d_k), значения — (batch_size, ..., seq_len, d_v), где ... — любое число других «батчевых» измерений; выход формы (batch_size, ..., seq_len, d_v). Поддержите необязательную булеву маску формы (seq_len, seq_len): вероятности внимания на позициях с True в сумме дают 1, на позициях с False — ноль. Тест: adapters.run_scaled_dot_product_attention; uv run pytest -k test_scaled_dot_product_attention (тензоры третьего порядка) и uv run pytest -k test_4d_scaled_dot_product_attention (четвёртого порядка).

3.4.5. Причинное многоголовое внимание: $\mathrm{MultiHeadSelfAttention}(x) = W_O \cdot \mathrm{MultiHead}(W_Q x, W_K x, W_V x)$ ; проекции K, Q, V — всего за три матричных умножения (в идеале — за одно). Причинная маска запрещает смотреть на будущие токены (torch.triu или сравнение индексов). RoPE применяется к запросам и ключам, но не к значениям, одинаково для всех голов (голова — батчевое измерение).

Problem (multihead_self_attention): причинное многоголовое self-attention (5 баллов)

Сдать: причинное многоголовое self-attention как torch.nn.Module, принимающее как минимум d_model: int (размерность входа блока) и num_heads: int (число голов). Как у Vaswani et al., $d_k = d_v = d_{model}/h$ . Тест: adapters.run_multihead_self_attention, uv run pytest -k test_multihead_self_attention.

Problem (transformer_block): блок трансформера (3 балла)

Реализуйте pre-norm блок трансформера (раздел 3.4, рисунок 2). Блок принимает как минимум d_model: int, num_heads: int и d_ff: int (размерность внутреннего слоя FFN). Тест: adapters.run_transformer_block, uv run pytest -k test_transformer_block. Сдать: код блока, проходящий тесты.

Problem (transformer_lm): трансформерная LM (3 балла)

Пора собрать всё вместе! Реализуйте языковую модель (раздел 3.1, рисунок 1). Кроме параметров блока она принимает как минимум: vocab_size: int (размер словаря — для матрицы эмбеддингов), context_length: int (максимальная длина контекста — для буферов sin и cos RoPE), num_layers: int (число блоков). Тест: adapters.run_transformer_lm, uv run pytest -k test_transformer_lm. Сдать: модуль LM, проходящий тесты.

Учёт ресурсов. Почти все FLOPs трансформера — матмулы. Правило: для $A \in \mathbb R^{m \times n}$ и $B \in \mathbb R^{n \times p}$ произведение $AB$ стоит $2mnp$ FLOPs.

Problem (transformer_accounting): учёт ресурсов трансформерной LM (5 баллов)

(a) Возьмите модель размера GPT-2 XL в архитектуре задания: vocab_size 50 257, context_length 1024, num_layers 48, d_model 1600, num_heads 25, d_ff 4288 (ближайшее кратное 64 к $\frac{8}{3} \times 1600$ ). Сколько у неё обучаемых параметров? Сколько памяти нужно, чтобы просто загрузить модель, если каждый параметр хранится в одинарной точности? Сдать: 1–2 предложения.

(b) Перечислите матричные умножения прямого прохода этой модели. Сколько всего FLOPs они требуют? Входная последовательность — context_length токенов. Сдать: список матмулов (с описаниями) и общее число FLOPs.

(d) Повторите анализ для GPT-2 small (12 слоёв, d_model 768, 12 голов), medium (24 слоя, 1024, 16 голов) и large (36 слоёв, 1280, 20 голов). Какие части трансформера занимают пропорционально больше или меньше FLOPs с ростом модели? Сдать: для каждой модели — разбивку по компонентам с долей в общих FLOPs прямого прохода и 1–2 предложения о том, как размер модели меняет эти доли.

(e) Увеличьте у GPT-2 XL длину контекста до 16 384. Как меняются общие FLOPs одного прямого прохода? Как меняется вклад компонентов? Сдать: 1–2 предложения.

4. Обучение трансформерной LM

4.1. Cross-entropy: $\ell(\theta; D) = \frac{1}{|D| m}\sum_{x \in D}\sum_{i=1}^{m} -\log p_\theta(x_{i+1} \mid x_{1:i})$ ; перплексия последовательности — $\exp\left(\frac{1}{m}\sum_i \ell_i\right)$ .

Problem (cross_entropy): cross-entropy (1 балл)

Сдать: функцию, которая по логитам $o_i$ и целям $x_{i+1}$ считает $\ell_i = -\log \mathrm{softmax}(o_i)[x_{i+1}]$ . Функция должна: вычитать максимум для численной устойчивости; сокращать log и exp, где возможно; поддерживать дополнительные батчевые измерения (они всегда идут перед измерением словаря) и возвращать среднее по батчу. Тест: adapters.run_cross_entropy, uv run pytest -k test_cross_entropy.

4.2. SGD в PyTorch. Оптимизатор — наследник torch.optim.Optimizer с методами __init__(self, params, ...) и step(self); в условии разобран пример SGD с убывающим шагом $\theta_{t+1} = \theta_t - \frac{\alpha}{\sqrt{t+1}}\nabla L$ и минимальный цикл обучения на (weights**2).mean().

Problem (learning_rate_tuning): подбор learning rate (1 балл)

Learning rate — один из самых влияющих гиперпараметров. Запустите пример с SGD с тремя другими значениями: 1e1, 1e2 и 1e3, всего на 10 итераций. Что происходит с потерями при каждом значении: убывают быстрее, медленнее или расходятся (растут в ходе обучения)? Сдать: 1–2 предложения о наблюдаемом поведении.

4.3. AdamW (алгоритм 2 из Loshchilov, Hutter): для каждого параметра хранятся оценки первого и второго моментов; гиперпараметры $\alpha$ , $(\beta_1, \beta_2)$ (обычно $(0{,}9;\ 0{,}999)$ , у LLM часто $(0{,}9;\ 0{,}95)$ ), $\varepsilon$ (например, $10^{-8}$ ) и $\lambda$ (weight decay). Шаг $t$ (с 1): $\alpha_t = \alpha\frac{\sqrt{1 - \beta_2^t}}{1 - \beta_1^t}$ ; $\theta \leftarrow \theta - \alpha\lambda\theta$ ; $m \leftarrow \beta_1 m + (1 - \beta_1) g$ ; $v \leftarrow \beta_2 v + (1 - \beta_2) g^2$ ; $\theta \leftarrow \theta - \alpha_t \frac{m}{\sqrt v + \varepsilon}$ .

Problem (adamw): AdamW (2 балла)

Сдать: оптимизатор AdamW — наследник torch.optim.Optimizer. Класс принимает в __init__ learning rate $\alpha$ и гиперпараметры $\beta$ , $\varepsilon$ , $\lambda$ . Для хранения состояния используйте словарь self.state базового класса (параметр → словарь с нужными величинами, у AdamW это оценки моментов). Тест: adapters.get_adamw_cls, uv run pytest -k test_adamw.

Problem (adamw_accounting): учёт ресурсов обучения с AdamW (2 балла)

Посчитаем память и вычисления для AdamW. Все тензоры в float32.

(a) Сколько пиковой памяти нужно для AdamW? Разложите на параметры, активации, градиенты и состояние оптимизатора. Выразите через batch_size и гиперпараметры модели (vocab_size, context_length, num_layers, d_model, num_heads). $d_{ff} = \frac{8}{3} d_{model}$ . Для активаций учитывайте только: в блоке трансформера — RMSNorm; подслой внимания (проекции QKV, умножение $QK^\top$ , softmax, взвешенная сумма значений, выходная проекция); FFN на SwiGLU ( $W_1$ , $W_2$ , SiLU на ветке вентиля, поэлементное произведение, $W_3$ ); финальный RMSNorm; выходной эмбеддинг; cross-entropy по логитам. Сдать: алгебраические выражения для параметров, активаций, градиентов, состояния оптимизатора и итога.

(b) Подставьте модель формы GPT-2 XL и получите выражение, зависящее только от batch_size. Какой максимальный батч поместится в 80 ГБ? Сдать: выражение вида $a \cdot \text{batch\_size} + b$ с числами $a$ , $b$ и максимальный батч.

(d) Model FLOPs utilization (MFU) — отношение наблюдаемой пропускной способности (токенов в секунду) к теоретическому пику FLOPs железа. Пик NVIDIA H100 — 495 TFLOP/s для операций «float32» (на деле TensorFloat-32, то есть фактически «bfloat19»). Если удаётся получить 50% MFU, сколько займёт обучение GPT-2 XL на 400 тыс. шагов с батчем 1024 на одной H100? Как у Kaplan et al. и Hoffmann et al., считайте, что обратный проход вдвое дороже прямого. Сдать: число часов с кратким обоснованием.

4.4. Расписание learning rate (косинусное, как у LLaMA). Прогрев: при $t < T_w$ — $\alpha_t = \frac{t}{T_w}\alpha_{max}$ . Косинусный спад: при $T_w \le t \le T_c$ — $\alpha_t = \alpha_{min} + \frac{1}{2}\left(1 + \cos\left(\frac{t - T_w}{T_c - T_w}\pi\right)\right)(\alpha_{max} - \alpha_{min})$ . После: при $t > T_c$ — $\alpha_t = \alpha_{min}$ .

Problem (learning_rate_schedule): косинусное расписание с прогревом (1 балл)

Напишите функцию, которая по $t$ , $\alpha_{max}$ , $\alpha_{min}$ , $T_w$ и $T_c$ возвращает $\alpha_t$ по расписанию выше. Затем реализуйте адаптер adapters.get_lr_cosine_schedule и убедитесь, что проходит uv run pytest -k test_get_lr_cosine_schedule.

4.5. Обрезка градиента: если $\|g\|_2$ больше $M$ , градиент умножают на $\frac{M}{\|g\|_2 + \varepsilon}$ ( $\varepsilon = 10^{-6}$ ).

Problem (gradient_clipping): обрезка градиента (1 балл)

Напишите функцию обрезки градиента: принимает список параметров и максимальную $\ell_2$ -норму и меняет градиент каждого параметра на месте. $\varepsilon = 10^{-6}$ (значение PyTorch по умолчанию). Адаптер adapters.run_gradient_clipping, тест uv run pytest -k test_gradient_clipping.

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

5.1. Загрузчик. Все документы склеены в одну последовательность токенов (с разделителем <|endoftext|>); батч — $B$ случайных отрезков длины $m$ и сдвинутые на один токен цели. Любое $1 \le i \le n - m$ даёт пример; паддинг не нужен.

Problem (data_loading): загрузка данных (2 балла)

Сдать: функцию, которая по массиву NumPy x (id токенов), batch_size, context_length и строке устройства PyTorch (например, 'cpu' или 'cuda:0') возвращает пару тензоров: входные последовательности и цели (следующие токены). Оба тензора формы (batch_size, context_length) и лежат на нужном устройстве. Тест: adapters.run_get_batch, uv run pytest -k test_get_batch.

Если датасет не помещается в память, используйте np.memmap (или np.load(..., mmap_mode='r')) с правильным dtype и проверьте, что значения не выходят за размер словаря.

5.2. Чекпоинты. Чекпоинт должен содержать всё для продолжения обучения: веса модели, состояние оптимизатора и номер итерации (state_dict(), load_state_dict(), torch.save, torch.load).

Problem (checkpointing): чекпоинты (1 балл)

Реализуйте две функции:

def save_checkpoint(model, optimizer, iteration, out) — записать всё состояние модели, оптимизатора и номер итерации в out (путь или файловый объект); обычно это словарь, но формат — на ваше усмотрение, лишь бы потом загрузился. Параметры: model: torch.nn.Module, optimizer: torch.optim.Optimizer, iteration: int, out: str | os.PathLike | typing.BinaryIO | typing.IO[bytes].
def load_checkpoint(src, model, optimizer) — загрузить чекпоинт из src (путь или файловый объект), восстановить состояния модели и оптимизатора и вернуть сохранённый номер итерации.

Адаптеры adapters.run_save_checkpoint и adapters.run_load_checkpoint, тест uv run pytest -k test_checkpointing.

Problem (training_together): собираем всё вместе (4 балла)

Сдать: скрипт с циклом обучения модели на данных пользователя. Рекомендуется, чтобы скрипт позволял как минимум: настраивать гиперпараметры модели и оптимизатора; экономно загружать большие обучающие и валидационные датасеты через np.memmap; сохранять чекпоинты по пути пользователя; периодически логировать качество на обучении и валидации (в консоль и/или во внешний сервис вроде Weights & Biases).

6. Генерация текста

Шаг декодирования: по $x_{1 \dots t}$ модель даёт логиты $v$ последней позиции, следующий токен сэмплируется из $\mathrm{softmax}(v)$ ; повторяем до <|endoftext|> или лимита. Два приёма: температура $\mathrm{softmax}(v, \tau)_i = \frac{\exp(v_i/\tau)}{\sum_j \exp(v_j/\tau)}$ и top-p (nucleus) — оставляем наименьшее множество самых вероятных токенов с суммой вероятностей не меньше $p$ и перенормируем.

Problem (decoding): декодирование (3 балла)

Сдать: функцию генерации из вашей LM. Рекомендуется поддержать: генерацию продолжения для промпта пользователя (по $x_{1 \dots t}$ сэмплировать до токена <|endoftext|>); ограничение максимального числа генерируемых токенов; масштабирование softmax температурой; top-p сэмплирование [Holtzman et al., 2020] с порогом пользователя.

7. Эксперименты

Модель — около 17 млн параметров без эмбеддингов; данные — TinyStories. Эксперименты нужно вести систематически и записывать: периодически считать потери на валидации и записывать номер шага и реальное время.

Problem (experiment_log): журнал экспериментов (3 балла)

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

Базовые гиперпараметры TinyStories (7.2.1): словарь 10 000; контекст 256; d_model 512; d_ff 1344; $\Theta$ RoPE 10 000; 4 слоя, 16 голов; всего обработать 327 680 000 токенов (батч × число шагов × длина контекста). Learning rate, прогрев, $\beta_1$ , $\beta_2$ , $\varepsilon$ и weight decay подберите сами. При правильной и эффективной реализации обучение занимает примерно 20–30 минут на одной B200. Советы по отладке (7.2.3): пользуйтесь отладчиком IDE; сначала переобучите модель на одном мини-батче (потери должны быстро упасть почти до нуля); проверяйте формы промежуточных тензоров; следите за нормами активаций, весов и градиентов.

Problem (learning_rate): подбор learning rate (2 часа B200) (3 балла)

Learning rate — один из важнейших гиперпараметров. На базовой модели ответьте:

(a) Переберите несколько learning rate и сообщите итоговые потери (или отметьте расходимость). Сдать: кривые обучения для нескольких learning rate и объяснение стратегии поиска. Сдать: модель с потерями на валидации TinyStories (на токен) не больше 1,45.

Совет для малых ресурсов: на CPU или MPS уменьшите число токенов до 40 000 000 — этого хватит для довольно связного текста, а целевые потери можно поднять с 1,45 до 2,00. Решение авторов на M4 Max (36 ГБ): 32 × 5000 × 256 = 40 960 000 токенов, 1 ч 22 мин на CPU и 36 мин на MPS, потери на валидации 1,80 на шаге 5000. Косинусный спад при $N$ шагах заканчивайте ровно на шаге $N$ ; на MPS не включайте TF32 (torch.set_float32_matmul_precision('high')) — авторы видели молча сломанные ядра; torch.compile ускоряет (на MPS — backend="aot_eager").

(b) Народная мудрость: лучший learning rate — «на краю устойчивости». Исследуйте, как точка, где learning rate начинает расходиться, связана с лучшим learning rate. Сдать: кривые обучения с растущим learning rate, включая хотя бы один расходящийся запуск, и анализ связи со скоростью сходимости.

Problem (batch_size_experiment): размер батча (1 час B200) (1 балл)

Меняйте размер батча от 1 до предела памяти GPU. Попробуйте несколько промежуточных значений, включая типичные 64 и 128. Сдать: кривые обучения для разных батчей (при необходимости заново подберите learning rate). Сдать: несколько предложений о том, как размер батча влияет на обучение.

Problem (generate): генерация текста (1 балл)

С помощью своего декодера и обученного чекпоинта приведите сгенерированный текст. Возможно, придётся подобрать параметры (температуру, top-p и т. д.), чтобы получить связный текст. Сдать: не меньше 256 токенов текста (или до первого <|endoftext|>) и краткий комментарий о связности и хотя бы двух факторах, влияющих на качество.

7.3. Абляции.

Problem (layer_norm_ablation): убрать RMSNorm (0,5 часа B200) (1 балл)

Уберите из трансформера все RMSNorm и обучите. Что происходит при прежнем оптимальном learning rate? Можно ли добиться устойчивости меньшим learning rate? Сдать: кривую обучения без RMSNorm и кривую для лучшего learning rate. Сдать: несколько предложений о роли RMSNorm.

Problem (pre_norm_ablation): post-norm (0,5 часа B200) (1 балл)

Переделайте pre-norm трансформер в post-norm ( $z = \mathrm{RMSNorm}(x + \mathrm{MHA}(x))$ , $y = \mathrm{RMSNorm}(z + \mathrm{FFN}(z))$ ) и обучите. Сдать: кривую обучения post-norm трансформера в сравнении с pre-norm.

Problem (no_pos_emb): NoPE (0,5 часа B200) (1 балл)

Уберите из реализации с RoPE позиционную информацию полностью (NoPE) и посмотрите, что будет. (Декодер с причинной маской теоретически может сам вывести позицию [Tsai et al., 2019; Kazemnejad et al., 2023].) Сдать: кривую обучения, сравнивающую RoPE и NoPE.

Problem (swiglu_ablation): SwiGLU против SiLU (0,5 часа B200) (1 балл)

Сравните FFN на SwiGLU с FFN на SiLU без вентиля: $\mathrm{FFN}_{SiLU}(x) = W_2 \mathrm{SiLU}(W_1 x)$ , при $d_{ff} = 4 d_{model}$ , чтобы число параметров примерно совпало с SwiGLU (у которого три матрицы вместо двух). Сдать: кривую обучения, сравнивающую SwiGLU и SiLU при примерно равном числе параметров. Сдать: несколько предложений о выводах.

Совет для малых ресурсов: обучение на OpenWebText до связного текста долгое, поэтому студентам с ограниченным доступом к GPU советуют и дальше проверять изменения на TinyStories (по потерям на валидации).

7.4. OpenWebText. Более реалистичный, сложный и разнообразный веб-текст; возможно, придётся заново подобрать learning rate или батч.

Problem (main_experiment): эксперимент на OWT (2 часа B200) (2 балла)

Обучите модель на OpenWebText с той же архитектурой и тем же числом итераций, что на TinyStories. Насколько хорошо она справляется? Сдать: кривую обучения на OpenWebText и объяснение разницы в потерях с TinyStories — как их понимать? Сдать: сгенерированный текст в том же формате. Насколько он связный? Почему качество хуже при той же модели и том же бюджете вычислений?

7.5. Своя модификация + лидерборд. Правила: не больше 45 минут на одной B200; только выданный обучающий датасет OpenWebText; остальное — на ваше усмотрение. Идеи — в открытых моделях (Llama 3, Qwen 2.5) и в репозитории modded-nanogpt (например, связывание весов входного и выходного эмбеддингов — тогда, возможно, придётся уменьшить стандартное отклонение их инициализации). Проверяйте идеи на малом подмножестве OWT или на TinyStories. Не всё, что помогает здесь, переносится на большой масштаб — об этом в разделе о законах масштабирования.

Problem (leaderboard): лидерборд (10 часов B200) (6 баллов)

Обучите модель по правилам лидерборда, минимизируя потери на валидации за 0,75 часа B200. Сдать: итоговые потери на валидации, кривую обучения с осью реального времени, явно меньше 45 минут, и описание того, что вы сделали. Ожидается, что результат будет лучше наивного базового уровня — потерь 5,0. Отправка — в assignment1-basics-leaderboard.

Как выполнять на вашей машине (4 × RTX 4080 32 ГБ)

Окружение. Склонируйте репозиторий, установите uv (урок 0.1.2) и запустите uv run pytest — все тесты должны падать с NotImplementedError. Данные скачиваются командами из README.md (обучающий файл TinyStories — порядка 2 ГБ, выборка OWT после распаковки — порядка 12 ГБ; проверьте место на NVMe).
BPE — это CPU-задача. У вас 48 ядер EPYC и 128 ГБ RAM: распараллельте предтокенизацию по 32–48 процессам. Тогда TinyStories обучается за минуты, а лимит OWT (12 часов, 100 ГБ RAM) реально уложить в 1–2 часа, если слияния обновляются инкрементально. Следите за памятью при подсчёте предтокенов OWT.
Модель и тесты — на одной карте или CPU. Все тесты частей 3–6 крошечные и проходят на CPU за секунды.
Обучение TinyStories — одна RTX 4080. Модель ~23 млн параметров с эмбеддингами; $6ND \approx 6 \cdot 2{,}3 \cdot 10^7 \cdot 3{,}3 \cdot 10^8 \approx 4{,}5 \cdot 10^{16}$ FLOPs. При 25–40 TFLOPS реальной производительности в BF16 это примерно 20–40 минут (оценка переводчика; без смешанной точности — в 2–3 раза дольше). Включите torch.autocast(dtype=torch.bfloat16) и torch.compile — это разрешено, ведь вы используете их поверх своих слоёв.
Перебор гиперпараметров — четыре запуска параллельно. Для learning_rate, batch_size_experiment и абляций запускайте по одному эксперименту на карту: CUDA_VISIBLE_DEVICES=0 uv run python train.py --lr 1e-3 &, затем 1, 2, 3. DDP для такой маленькой модели не нужен и по PCIe только замедлит. «Часы B200» из условия на RTX 4080 примерно соответствуют нескольким часам — планируйте ночные прогоны.
OpenWebText и «лидерборд». Официальный лидерборд для вас недоступен по правилам (45 минут на B200), но можно поставить себе аналог: «лучшие потери на валидации OWT за 45 минут на одной RTX 4080» или «за 45 минут на 4 картах с DDP» — тогда пригодится урок 5.7.4.

Подсказки (без решений)

BPE медленный? 90% времени — предтокенизация (распараллельте) и пересчёт пар после каждого слияния (храните счётчики пар и индекс «пара → в каких предтокенах встречается»).
Тест BPE не совпадает на одной позиции? Почти всегда виноваты правило выбора при равенстве (лексикографически бо́льшая пара сравнивается как bytes, а не как строка) или специальные токены, попавшие в подсчёт.
Потери застряли около $\ln(\text{vocab}) \approx 9{,}2$ ? Модель ничего не учит: проверьте причинную маску, сдвиг целей на один токен и то, что оптимизатор получает параметры модели.
NaN после нескольких сотен шагов? Проверьте приведение к float32 в RMSNorm и softmax, прогрев learning rate и обрезку градиента.
Медленное обучение? Проверьте, что загрузчик не копирует весь массив на каждом шаге и что валидация считается не слишком часто.

Что дальше

Задание 2 — профилирование этой же модели, FlashAttention-2 на Triton и распределённое обучение: путеводитель по заданию 2.
Перевод лекций, нужных для задания: 1–2 и 3–4.
Обзор всего курса CS336.

Конспект и пересказ — собственный текст курса; оригинал: CS336 Staff, Stanford University, «CS336 Assignment 1 (basics): Building a Transformer LM — Version 26.0.3 (Spring 2026)». Нашли неточность — сверяйтесь с оригиналом по ссылке выше. · обновлено 2026-10-04