CS336, задание 2 (systems): профилирование, FlashAttention-2 и распределённое обучение — путеводитель

На чём основана страница

Источник — PDF-условие cs336_assignment2_systems.pdf (версия 26.1.3, весна 2026) из репозитория stanford-cs336/assignment2-systems (лицензия MIT). Условия всех задач (блоки «Problem») переведены по порядку, с баллами; длинные объяснения между ними (алгоритм FlashAttention-2, устройство DDP и FSDP) пересказаны кратко со ссылками на наши уроки. Разделы о вашей машине и подсказки написаны переводчиком; кода решений здесь нет.

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

Задание 2 — «сделай свою модель быстрой и распредели её по нескольким GPU». Вы измеряете, куда уходит время и память, пишете ядро FlashAttention-2 на Triton, затем строите собственные DDP, шардинг состояния оптимизатора и FSDP поверх torch.distributed и считаете, когда какой параллелизм упирается в связь.

• Когда: после модулей 5.4 и 5.7. Урок 5.7.3 даёт навыки Triton, 5.7.4 — опыт DDP/FSDP на ваших четырёх картах.
• Теория: лекции 5–6 (GPU, ядра, Triton) и лекции 7–8 (параллелизм).
• Код задания 1 не обязателен: в папке cs336-basics/ лежит эталонная модель из задания 1; можно подключить и свою, поменяв pyproject.toml.

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

Раздел	Что делаете	Баллы	Урок курса
2. Профилирование	Скрипт бенчмарка, Nsight Systems, смешанная точность, профиль памяти	4 + 5 + 1 + 2 + 4	5.7.1, 5.4.1
3. Память одной GPU	Остатки autograd, слияние операторов, оптимальный checkpointing	4	5.7.4 (перевычисление активаций)
4. Ядра GPU	Бенчмарк внимания, torch.compile, FlashAttention-2 forward (PyTorch + Triton), backward, бенчмарк	2 + 2 + 15 + 5 + 5	5.7.2, 5.7.3
5. DDP	Бенчмарк all-reduce, наивный DDP, плоские градиенты, перекрытие связи и вычислений	5 + 5 + 3 + 2 + 5 + 1	5.4.3, 5.7.4
6. Шардинг оптимизатора	Свой ZeRO-1 и его учёт	15 + 5	5.4.3
7. FSDP	Свой FSDP и его учёт	15 + 5	5.7.4
8. Анализ параллелизма	Кольцевой all-reduce, формулы для DP, FSDP, TP и 2D	1 + 3 + 3 + 4 + 6	лекция 7
9. Лидерборд	Самый быстрый шаг обучения модели 8B	10	—

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

Что делаем: (1) инструменты бенчмаркинга и профилирования; (2) перевычисление активаций (activation checkpointing); (3) ядро FlashAttention-2 на Triton; (4) распределённое обучение с параллелизмом по данным; (5) шардинг состояния оптимизатора; (6) полностью шардированный параллелизм по данным (FSDP).

Как устроен код: cs336-basics/ — эталонное решение задания 1; корень — пакет cs336_systems (пустой, пишете с нуля); tests/*.py — тесты, вызывающие адаптеры из tests/adapters.py; ваша реализация должна проходить исходный набор тестов.

Как сдают: writeup.pdf (ответы на письменные вопросы, набранные) и code.zip (собирается скриптом test_and_make_submission.sh) — в Gradescope. Авторы советуют генерировать таблицы для отчёта кодом (pandas.DataFrame.to_latex() или .to_typst()).

Размеры моделей (словарь 10 000, батч 4, контекст 512, если не сказано иное):

Размер	d_model	d_ff	слоёв	голов
small	768	3072	12	12
medium	1024	4096	24	16
large	1280	5120	36	20
xl	2560	10240	32	32
10B	4608	12288	50	36

Примечание переводчика: что поместится в 32 ГБ

Посчитаем параметры этих моделей с SwiGLU и словарём 10 000: small ≈ 0,13 млрд, medium ≈ 0,42 млрд, large ≈ 0,97 млрд, xl ≈ 3,4 млрд, 10B ≈ 12,8 млрд. Полный шаг обучения в FP32 с AdamW требует около 16 байт на параметр без активаций: large ≈ 14 ГБ (помещается в RTX 4080 32 ГБ), xl ≈ 51 ГБ (не помещается на одну карту), 10B — только с шардингом и выгрузкой. Поэтому на вашей машине задачи «для xl на одной GPU» делайте для large (и отметьте это в отчёте), а xl используйте там, где он помещается: прямой проход, BF16 с перевычислением активаций, шардинг на 4 картах.

---

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

2. Профилирование и бенчмаркинг

Три пути оценки производительности: (1) простой сквозной бенчмарк средствами стандартной библиотеки Python; (2) профилирование вычислений NVIDIA Nsight Systems — как время распределено по операциям CPU и GPU; (3) профилирование памяти. Главная ловушка — асинхронность CUDA: без torch.cuda.synchronize() вы меряете постановку ядра в очередь, а не его работу (подробно — в лекции 6).

Problem (benchmarking_script): скрипт бенчмарка (4 балла)

(a) Напишите скрипт сквозного бенчмарка прямого прохода, обратного прохода и шага оптимизатора. Скрипт должен: по гиперпараметрам (например, числу слоёв) создавать модель; генерировать случайный батч; делать $w$ шагов прогрева (до начала замера), затем замерять $n$ шагов (только прямой проход, прямой + обратный или прямой + обратный + шаг оптимизатора — в зависимости от аргумента). Для замера можно использовать модуль timeit (функцию timeit или timeit.default_timer() — самые точные системные часы, лучше, чем time.time()). После каждого шага вызывайте torch.cuda.synchronize(). Сдать: скрипт, который создаёт модель basics с заданными гиперпараметрами, случайный батч и замеряет только прямой проход, прямой + обратный и полный шаг с оптимизатором.

(b) Замерьте прямой проход, обратный и шаг оптимизатора для размеров моделей из таблицы. 5 шагов прогрева, среднее и стандартное отклонение по 10 замерам. Сколько занимает прямой проход? Обратный? Велик ли разброс? Сдать: 1–2 предложения с цифрами.

(c) Одна из ошибок бенчмаркинга — не делать прогрев. Повторите анализ без прогрева. Как это меняет результаты? Почему? Попробуйте также 1 или 2 шага прогрева. Почему результат всё ещё может отличаться? Сдать: 2–3 предложения.

Problem (nsys_profile): профилирование Nsight Systems (5 баллов)

Профилируйте nsys прямой проход, обратный и шаг оптимизатора для двух размеров модели из таблицы (на ваш выбор) и трёх длин контекста — степеней двойки больше 128, где самая большая — максимальная, что помещается в память. Выберите самые интересные сочетания. Для каждого профиля ответьте:

(a) Сколько всего времени занимает прямой проход? Совпадает ли это с замером стандартной библиотекой Python? Сдать: 1–2 предложения.

(b) Какое ядро CUDA занимает больше всего суммарного времени GPU в прямом проходе? Сколько раз оно вызывается за один прямой проход? То же ли ядро лидирует, когда идут и прямой, и обратный проходы? (Подсказка: смотрите «CUDA GPU Kernel Summary» в «Stats System View» и фильтруйте по диапазонам NVTX, чтобы понять, какие части модели вызывают какие ядра.) Сдать: 1–2 предложения.

(c) Хотя основная часть FLOPs приходится на матмулы, другие ядра тоже занимают заметную долю времени. Какие ядра, кроме матмулов, заметны в прямом проходе? Сдать: 1–2 предложения.

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

(e) Сравните время softmax и время матмулов внутри self-attention при прямом проходе. Как разница во времени соотносится с разницей во FLOPs? Сдать: 1–2 предложения.

2.1.5. Смешанная точность.

Problem (mixed_precision_accumulation): накопление в смешанной точности (1 балл)

Выполните код и прокомментируйте точность результатов.

s = torch.tensor(0, dtype=torch.float32)
for i in range(1000):
    s += torch.tensor(0.01, dtype=torch.float32)
print(s)
s = torch.tensor(0, dtype=torch.float16)
for i in range(1000):
    s += torch.tensor(0.01, dtype=torch.float16)
print(s)
s = torch.tensor(0, dtype=torch.float32)
for i in range(1000):
    s += torch.tensor(0.01, dtype=torch.float16)
print(s)
s = torch.tensor(0, dtype=torch.float32)
for i in range(1000):
    x = torch.tensor(0.01, dtype=torch.float16)
    s += x.type(torch.float32)
print(s)

Сдать: 2–3 предложения.

Problem (benchmarking_mixed_precision): бенчмарк смешанной точности (2 балла)

(a) Рассмотрим модель:

class ToyModel(nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super().__init__()
        self.fc1 = nn.Linear(in_features, 10, bias=False)
        self.ln = nn.LayerNorm(10)
        self.fc2 = nn.Linear(10, out_features, bias=False)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.ln(x)
        x = self.fc2(x)
        return x

Обучаем на GPU, параметры изначально в FP32, используем autocast со смешанной точностью FP16. Каковы типы данных: параметров модели внутри контекста autocast; выхода первого линейного слоя (fc1); выхода layer norm (ln); логитов; потерь; градиентов модели? Сдать: типы данных для каждого пункта.

(b) Вы должны были увидеть, что autocast FP16 обращается с layer norm иначе, чем с линейными слоями. Какие части layer norm чувствительны к смешанной точности? Если взять BF16 вместо FP16, нужно ли всё ещё особое обращение с layer norm? Почему? Сдать: 2–3 предложения.

(c) Добавьте в скрипт бенчмарка опцию смешанной точности BF16. Замерьте прямой и обратный проходы с ней и без неё для каждого размера модели. Сравните и прокомментируйте тренды с ростом модели. Может пригодиться пустой контекстный менеджер nullcontext. Сдать: 2–3 предложения с цифрами и комментарием.

2.1.6. Профиль памяти.

Problem (memory_profiling): профилирование памяти (4 балла)

Профилируйте полный шаг обучения (прямой, обратный, шаг оптимизатора) модели xl с длинами контекста 128 и 2048.

(a) Добавьте в скрипт опцию запуска через профайлер памяти (можно переиспользовать прежнюю инфраструктуру: смешанную точность, загрузку размеров моделей и т. п.). Получите профиль памяти xl для только инференса (прямого прохода) и для полного шага обучения. Как выглядят временные шкалы памяти? Можно ли по пикам понять, какой этап идёт? Сдать: два изображения «Active memory timeline» модели xl из инструмента memory_viz — для прямого прохода и для полного шага — и 2–3 предложения.

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

(c) Найдите пик памяти xl со смешанной точностью — для прямого прохода и полного шага. Сильно ли смешанная точность влияет на память? Сдать: 2–3 предложения.

(d) Для модели xl с эталонными гиперпараметрами: каков размер тензора активаций в остаточном потоке трансформера в одинарной точности? Дайте в МиБ (разделите байты на $1024^2$). Сдать: 1–2 предложения с выводом.

(e) Внимательно посмотрите на «Active Memory Timeline» из pytorch.org/memory_viz для снимка памяти прямого прохода xl. Если уменьшить уровень «Detail», инструмент скрывает мелкие выделения (например, при 10% показываются только 10% самых крупных). Каков размер самых крупных показанных выделений? Можно ли по стеку вызовов понять, откуда они? Сдать: 1–2 предложения.

(f) У Nsight Systems тоже есть флаги профилирования памяти; их можно сочетать с прежними флагами и видеть, какие выделения происходят на разных этапах. С помощью меток NVTX от PyTorch определите, сколько памяти сохраняется для обратного прохода (эти тензоры часто называют остатками, residuals) одним TransformerBlock. Назовите 5 крупнейших операций и их долю в общем объёме. В обратном проходе все эти тензоры освобождаются, но одновременно создаются тензоры градиентов. По профилю (сколько памяти выделилось в прямом проходе и как память меняется на каждом TransformerBlock в обратном) посчитайте, сколько занимают градиенты одного TransformerBlock. Совпадает ли с ожиданием? Сдать: снимки экрана Nsight Systems и ответ в 1–2 абзаца.

3. Память одной GPU

Раздел объясняет, какие тензоры autograd сохраняет для обратного прохода (residuals), как слияние операторов уменьшает их число и как перевычисление активаций (torch.utils.checkpoint) меняет память на вычисления. В условии приведена оценка: для модели xl с контекстом 2048 остатков набирается около 3,6 ГиБ на слой, то есть около 114 ГиБ на все слои — «просто сохранённых для обратного прохода».

Problem (gradient_checkpointing): оптимальное по памяти перевычисление (4 балла)

Возьмём трансформер из $N$ одинаковых последовательных блоков. Без перевычисления одновременно живут остатки всех $N$ блоков — пик памяти активаций $O(N)$. Любой фрагмент прямого прохода можно обернуть в checkpoint, в том числе вкладывать вызовы checkpoint друг в друга.

(a) Какая стратегия минимизирует пик памяти активаций, если не считать стоимость вычислений? Опишите, как расставить вызовы checkpoint (можно набросок кода), и дайте асимптотику пиковой памяти и вычислений как функции $N$. Считайте, что остатки одного блока много больше служебных расходов на каждый чекпоинт. Сдать: описание стратегии в 3–5 предложениях и её асимптотика памяти, плюс короткий набросок кода.

(b) Модель xl, батч 4, длина 2048. Если бюджет времени позволяет только один уровень перевычисления (без вложенных checkpoint), какая стратегия лучше всего снижает пик памяти? Проверьте гипотезу замером пика. Для уверенности сравните с соседними (меньшим и бо́льшим) размерами группы блоков в чекпоинте. Сдать: рассуждение в 3–5 предложениях и измеренный пик памяти для вашей стратегии.

4. Ядра GPU: FlashAttention-2

Problem (pytorch_attention): бенчмарк внимания PyTorch (2 балла)

(a) Замерьте свою реализацию внимания на разных масштабах. Скрипт должен: (i) зафиксировать батч 8 и не использовать многоголовость (убрать измерение голов); (ii) перебрать декартово произведение размерности $d_{model}$ из [16, 32, 64, 128] и длины последовательности из [256, 1024, 4096, 8192, 16384]; (iii) создать случайные $Q$, $K$, $V$ нужного размера; (iv) замерить 100 прямых проходов; (v) измерить, сколько памяти занято перед началом обратного прохода, и замерить 100 обратных проходов; (vi) не забыть прогрев и torch.cuda.synchronize() после каждого прохода. В зависимости от GPU часть конфигураций не поместится в память. Приведите времена (или ошибки нехватки памяти). На каком размере начинается нехватка памяти? Посчитайте память внимания для одной из самых маленьких не поместившихся конфигураций (можно использовать формулы памяти из задания 1). Как память, сохранённая для обратного прохода, зависит от длины последовательности? Что бы вы сделали, чтобы убрать эту стоимость? Сдать: таблица времён, расчёт памяти и ответ в 1–2 абзаца.

Problem (torch_compile): torch.compile (2 балла)

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

(b) Скомпилируйте всю модель в сквозном бенчмарке. Как меняется прямой проход? А прямой + обратный + шаг оптимизатора? Сдать: таблица сравнения обычной и скомпилированной модели.

4.2.1–4.2.2. Пример Triton и прямой проход FlashAttention-2. Условие сначала подробно разбирает ядро Triton для взвешенной суммы (с обратным проходом через torch.autograd.Function), затем — алгоритм 1 FlashAttention-2: запросы делятся на тайлы $B_q$, ключи и значения — на тайлы $B_k$; для каждого тайла запросов идём по тайлам ключей, поддерживая текущий максимум $m$, сумму $l$ и ненормированный выход $O$ (online softmax), и в конце сохраняем выход и logsumexp $L$ для обратного прохода. Подробный разбор с числами — в уроке 5.7.2.

Problem (flash_forward): прямой проход FlashAttention-2 (15 баллов)

(a) Напишите на чистом PyTorch (без Triton) autograd.Function, реализующую прямой проход FlashAttention-2. Она будет намного медленнее обычной реализации, но поможет отлаживать ядро Triton. На вход — $Q$, $K$, $V$ и флаг is_causal, на выход — $O$ и logsumexp $L$; флаг is_causal в этой задаче можно игнорировать. Метод forward должен сохранить $L, Q, K, V, O$ для обратного прохода и вернуть $O$ (первым параметром forward у autograd.Function всегда идёт контекст). Метод backward обязателен, но пока может просто бросать NotImplementedError. Для сравнения можно реализовать в PyTorch уравнения 4–6 и 12 условия. Интерфейс: def forward(ctx, Q, K, V, is_causal=False). Размеры тайлов выберите сами, но не меньше 16 × 16; тесты всегда используют размерности — степени двойки не меньше 16, так что о выходе за границы можно не беспокоиться. Сдать: подкласс torch.autograd.Function с прямым проходом FlashAttention-2. Реализуйте adapters.get_flashattention_autograd_function_pytorch и проверьте uv run pytest -k test_flash_forward_pass_pytorch.

(b) Напишите ядро Triton для прямого прохода FlashAttention-2 по алгоритму 1, затем ещё один подкласс torch.autograd.Function, который вызывает это (слитое) ядро вместо вычислений в PyTorch. Советы:

• для отладки сравнивайте результат каждой операции Triton с тайловой реализацией на PyTorch из части (a);
• сетка запуска — (T_q, batch_size): каждый экземпляр программы Triton загружает элементы одного индекса батча и читает/пишет только один тайл запросов в $Q$, $O$ и $L$;
• в ядре должен быть один цикл — по тайлам ключей $1 \le j \le T_k$;
• сдвигайте блочные указатели в конце цикла;
• используйте объявление функции из условия:

@triton.jit
def flash_fwd_kernel(
    Q_ptr, K_ptr, V_ptr,
    O_ptr, L_ptr,
    stride_qb, stride_qq, stride_qd,
    stride_kb, stride_kk, stride_kd,
    stride_vb, stride_vk, stride_vd,
    stride_ob, stride_oq, stride_od,
    stride_lb, stride_lq,
    N_QUERIES, N_KEYS,
    scale,
    D: tl.constexpr,
    Q_TILE_SIZE: tl.constexpr,
    K_TILE_SIZE: tl.constexpr,
):
    # Индексы программы
    query_tile_index = tl.program_id(0)
    batch_index = tl.program_id(1)

    # Сдвигаем каждый указатель на индекс батча, умноженный на шаг (stride) батча соответствующего тензора
    Q_block_ptr = tl.make_block_ptr(
        Q_ptr + batch_index * stride_qb,
        shape=(N_QUERIES, D),
        strides=(stride_qq, stride_qd),
        offsets=(query_tile_index * Q_TILE_SIZE, 0),
        block_shape=(Q_TILE_SIZE, D),
        order=(1, 0),
    )
    ...

где scale — $\frac{1}{\sqrt d}$, а Q_TILE_SIZE и K_TILE_SIZE — $B_q$ и $B_k$ (их можно подстроить позже). Чтобы избежать проблем с точностью: буферы на кристалле ($O_i$, $l$, $m$) — типа tl.float32; при накоплении в выходной буфер используйте аргумент acc (acc = tl.dot(..., acc=acc)); приводите $\tilde P_i^{(j)}$ к типу $V^{(j)}$ перед умножением, а $O_i$ — к нужному типу перед записью в глобальную память (tensor.to; тип тензора — tensor.dtype, тип указателя — *_block_ptr.type.element_ty). Сдать: подкласс torch.autograd.Function с прямым проходом на вашем ядре Triton. Реализуйте adapters.get_flash_autograd_function_triton и проверьте uv run pytest -k test_flash_forward_pass_triton.

Problem (flash_backward): обратный проход FlashAttention-2 (5 баллов)

Реализуйте обратный проход вашей autograd.Function на PyTorch (не Triton) с torch.compile. На вход — тензоры $Q, K, V, O, dO$ и $L$, на выход — $dQ$, $dK$ и $dV$. Не забудьте посчитать и использовать вектор $D$. Можно следовать уравнениям 13–19 условия. Сдать: проверка — uv run pytest -k test_flash_backward.

Problem (flash_benchmarking): бенчмарк FlashAttention-2 (5 баллов)

(a) Напишите бенчмарк на triton.testing.do_bench, сравнивающий вашу (частично) Triton-реализацию прямого и обратного проходов FlashAttention-2 с обычной реализацией на PyTorch (без FlashAttention). Таблица должна содержать задержки прямого, обратного и сквозного (прямой + обратный) прохода для обеих реализаций. Входы генерируйте случайно заранее; запуск — на одной B200; всегда батч 1 и причинная маска. Переберите декартово произведение длин последовательности — степеней двойки от 128 до 65 536, размерностей эмбеддинга — степеней двойки от 16 до 128, и точностей torch.bfloat16 и torch.float32. Скорее всего, размеры тайлов придётся подбирать под размер входа. Сдать: таблица сравнения с PyTorch по этим настройкам — задержки прямого, обратного и сквозного прохода.

4.2.3. (Необязательно) обратный проход на Triton — для тех, кто хочет больше практики и быстрый результат в лидерборде.

5. Распределённое обучение с параллелизмом по данным

Условие показывает основы torch.distributed (группа процессов, NCCL, all_reduce) и правила бенчмаркинга распределённых программ: прогрев, torch.cuda.synchronize(), несколько замеров, сбор результатов со всех рангов.

Problem (distributed_communication_single_node): распределённая связь на одном узле (5 баллов)

Напишите скрипт бенчмарка all-reduce в постановке «один узел, несколько процессов» (пример кода из условия — хорошая отправная точка). Меняйте: размер данных all-reduce — тензоры float32 размером 1 МБ, 10 МБ, 100 МБ, 1 ГБ; число GPU/процессов — 2, 4 или 6. Ресурсы: до 6 GPU, каждый запуск бенчмарка — меньше 5 минут. Сдать: графики и/или таблицы сравнения настроек и 2–3 предложения о результатах и о том, как факторы взаимодействуют.

5.2. Наивный DDP: после обратного прохода градиент каждого параметра сводится all-reduce по отдельности.

Problem (naive_ddp): наивный DDP (5 баллов)

Сдать: наивную форму DDP, которая после обратного прохода сводит all-reduce градиенты отдельных параметров. Тест: реализуйте adapters.get_ddp и (необязательно) adapters.ddp_on_after_backward, запустите uv run pytest tests/test_ddp.py.

Problem (naive_ddp_benchmarking): бенчмарк наивного DDP (3 балла)

В наивном DDP градиенты параметров сводятся all-reduce по отдельности после каждого обратного прохода. Чтобы понять накладные расходы, напишите скрипт бенчмарка модели из задания 1 при обучении с этим DDP. Измерьте общее время шага и долю времени на передачу градиентов. Постановка — один узел, 2 GPU, модель xl. Сдать: описание постановки, измеренное время итерации и время передачи градиентов.

5.3. Улучшаем минимальный DDP.

Problem (minimal_ddp_flat_benchmarking): минимальный DDP с плоскими градиентами (2 балла)

Измените минимальный DDP так, чтобы передавать один тензор со склеенными («сплющенными») градиентами всех параметров. Сравните с версией, где all-reduce вызывается для каждого параметра, в тех же условиях (1 узел × 2 GPU, xl). Сдать: время итерации и время передачи градиентов при одном пакетном вызове all-reduce и 1–2 предложения сравнения.

Problem (ddp_overlap_individual_parameters): DDP с перекрытием для отдельных параметров (5 баллов)

Реализуйте класс Python для распределённого обучения с параллелизмом по данным. Класс оборачивает произвольный nn.Module, рассылает веса перед обучением (чтобы у всех рангов были одинаковые начальные параметры) и запускает передачу для усреднения градиентов. Рекомендуемый интерфейс:

• def __init__(self, module: torch.nn.Module) — по созданному модулю построить контейнер DDP, синхронизирующий градиенты между рангами;
• def forward(self, *inputs, **kwargs) — вызвать forward() обёрнутого модуля с этими аргументами;
• def finish_gradient_synchronization(self) — дождаться завершения асинхронных вызовов связи на GPU.

Использование: обернуть модель, на каждом шаге вычислить потери и обратный проход, затем вызвать finish_gradient_synchronization() перед optimizer.step(), чтобы шаг оптимизатора (зависящий от градиентов) безопасно встал в очередь. Суть задачи — запускать асинхронный all-reduce градиента каждого параметра сразу, как только он посчитан в обратном проходе (через хуки), чтобы связь перекрывалась с вычислением. (В сложных случаях с несколькими потоками CUDA может понадобиться явная синхронизация потоков.) Сдать: класс с тестом из tests/test_ddp.py (см. адаптеры в условии).

Problem (ddp_overlap_individual_parameters_benchmarking): бенчмарк перекрытия (1 балл)

(a) Замерьте DDP с перекрытием обратного прохода и передачи градиентов отдельных параметров. Сравните с прежними вариантами (all-reduce на каждый параметр или один all-reduce склеенных градиентов) в той же постановке: 1 узел, 2 GPU, xl. Сдать: время итерации и 1–2 предложения сравнения.

(b) Профилируйте бенчмарк Nsight (та же постановка), сравнивая исходный DDP и вариант с перекрытием. Визуально сравните трассы и приведите снимок, показывающий, что одна реализация перекрывает вычисления и связь, а другая — нет. Сдать: 2 снимка экрана, наглядно показывающих, перекрывается связь с обратным проходом или нет.

6. Шардинг состояния оптимизатора

DDP требует, чтобы у каждого ранга была своя копия параметров и состояния оптимизатора; у AdamW это два числа на параметр — вдвое больше памяти, чем сами веса. По ZeRO [Rajbhandari et al.] состояние оптимизатора можно разделить между рангами.

Problem (optimizer_state_sharding): шардинг состояния оптимизатора (15 баллов)

Реализуйте класс Python, который оборачивает произвольный optim.Optimizer и синхронизирует обновлённые параметры после каждого шага оптимизатора. Рекомендуемый интерфейс:

• def __init__(self, params, optimizer_cls: Type[Optimizer], **kwargs: Any) — инициализировать шардированный оптимизатор. params — набор оптимизируемых параметров (или групп параметров с разными гиперпараметрами); эти параметры распределяются по всем рангам. optimizer_cls — тип оборачиваемого оптимизатора (например, optim.AdamW); остальные именованные аргументы передаются в его конструктор. Обязательно вызовите конструктор базового класса torch.optim.Optimizer.
• def step(self, closure, **kwargs) — вызвать step() обёрнутого оптимизатора с этими аргументами; после обновления параметров синхронизироваться с остальными рангами.
• def add_param_group(self, param_group: dict[str, Any]) — добавить группу параметров. Вызывается конструктором базового класса при создании и может вызываться во время обучения (например, при постепенной «разморозке» слоёв), поэтому должен распределять параметры модели между рангами.

Сдать: класс-контейнер для шардинга состояния оптимизатора. Реализуйте адаптер adapters.get_sharded_optimizer, запустите uv run pytest tests/test_sharded_optimizer.py. Советуем прогнать тесты несколько раз (например, 5), чтобы убедиться, что они проходят стабильно.

Problem (optimizer_state_sharding_accounting): учёт шардинга оптимизатора (5 баллов)

(a) Напишите скрипт профилирования пиковой памяти при обучении с шардингом состояния оптимизатора и без него. В стандартной постановке (1 узел, 2 GPU, xl) сообщите пик памяти после создания модели, непосредственно перед шагом оптимизатора и сразу после него. Совпадает ли с ожиданием? Разложите память в каждом случае (сколько на параметры, сколько на состояние оптимизатора и т. д.). Сдать: 2–3 предложения с пиками памяти и разбивкой по компонентам.

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

(c) Чем наш подход к шардингу отличается от ZeRO stage 1 (ZeRO-DP $P_{os}$ в статье Rajbhandari et al.)? Сдать: 2–3 предложения о различиях, особенно в памяти и объёме связи.

7. Fully-Sharded Data Parallel

Problem (fsdp): FSDP (15 баллов)

Реализуйте класс Python для полностью шардированного параллелизма по данным. Класс оборачивает произвольный nn.Module (всю модель) и подключается (хуками) к каждому слою Linear или Embedding внутри него — или оборачивает их. Рекомендуемый интерфейс:

• def __init__(self, module: torch.nn.Module, compute_dtype: torch.dtype | None = None) — построить модуль FSDP, который выполняет all-gather весов и reduce-scatter градиентов. Хуки или обёртки должны собирать веса (all-gather) вовремя для прямого прохода; чтобы ограничить память, начинайте сбор только после того, как слой «на два раньше» текущего завершил прямой проход. В обратном проходе хуки или обёртки снова собирают веса для вычислений; когда градиенты готовы, их нужно распределить reduce-scatter по рангам. Освобождайте собранные веса после использования. Если задан compute_dtype, приводите веса к этому типу перед передачей и вычислениями, но главные веса и обновления оптимизатора держите в FP32.
• def forward(self, *inputs, **kwargs) — вызвать forward() обёрнутого модуля.
• def finish_gradient_synchronization(self) — дождаться завершения асинхронных вызовов связи на GPU.

Сдать: класс-контейнер FSDP. Каждый шард должен работать со стандартным AdamW из задания 1. Реализуйте адаптер adapters.get_fsdp, запустите uv run pytest tests/test_fsdp.py. Советуем прогнать тесты несколько раз (например, 5), чтобы поймать состояния гонки.

Problem (fsdp_accounting): учёт FSDP (5 баллов)

(a) Исходя из анализа в разделе 6, сколько памяти от пика вы ожидаете сэкономить с FSDP? Размер заранее выделенных буферов для сбора весов можно не учитывать. Сдать: 2–3 предложения.

(b) Профилируйте модель xl на двух GPU, обращая внимание на сбор весов. Успевает ли связь завершиться к прямому проходу? Сдать: 2–3 предложения с цифрами и снимками Nsight в подтверждение.

8. Анализ стратегий параллелизма

Упрощённая постановка: $N$ устройств, каждая пара соединена, у каждого исходящая пропускная способность $W$ байт/с. Кольцевой all-gather тензора размера $S$ занимает $\frac{N-1}{N}\frac{S}{W}$ секунд, кольцевой reduce-scatter — столько же, all-reduce (reduce-scatter + all-gather) — $2\frac{N-1}{N}\frac{S}{W}$. Анализируется один слой FFN: $x_1 = xW_1$, $x_2 = xW_2$, $z = f(x_1) * x_2$, $y = zW_3$, где $x$ формы $(B, D)$, $W_1, W_2$ — $(D, D_{FF})$, $W_3$ — $(D_{FF}, D)$; в условии выписан и обратный проход. Все веса и активации — FP16 (2 байта), матмул $(A, B)(B, C)$ стоит $2ABC$ FLOPs, $C$ — скорость ускорителя (FLOP/с).

Problem (alternate_ring_all_reduce): альтернативный кольцевой all-reduce (1 балл)

Вместо reduce-scatter + all-gather используем алгоритм: на шагах $t = 1, \dots, N-1$ устройство $i$: при $t = 1$ ставит $y \leftarrow x^{(i)}$ (частичная сумма); отправляет $x^{((i - t + 1) \bmod N)}$ устройству $(i + 1) \bmod N$; получает $x^{((i - t) \bmod N)}$ от устройства $(i - 1) \bmod N$; обновляет $y \leftarrow y + x^{((i - t) \bmod N)}$. В тех же условиях (пропускная способность $W$, каждый $x^{(i)}$ размера $S$) — сколько времени занимает алгоритм? Сдать: ответ через $S$, $N$, $W$ и одно предложение обоснования.

Problem (data_parallel_calcs): расчёты для параллелизма по данным (3 балла)

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

(a) Сколько FLOPs нужно на обратный проход при $N_{DP}$-кратном параллелизме по данным (не-матмулы игнорируйте)? Сдать: ответ через $B$, $D$, $D_{FF}$, $N_{DP}$ и одно предложение обоснования.

(b) Сколько времени связи нужно в обратном проходе? Сдать: ответ через подмножество $B, D, D_{FF}, N_{DP}, W$ и одно предложение.

(c) При прочих равных, насколько большим может быть $N_{DP}$, прежде чем связь станет узким местом? Сдать: неравенство с $N_{DP}$ с одной стороны и выражением через подмножество $B, D, D_{FF}, C, W$ — с другой, и одно предложение.

Problem (fsdp_calcs): расчёты для FSDP (3 балла)

В той же постановке — когда FSDP упирается в связь?

(a) Сколько FLOPs на обратный проход при $N_{FSDP}$? На прямой? Сдать: два ответа через $B, D, D_{FF}, N_{FSDP}$ и два предложения обоснования.

(b) Сколько времени связи в обратном проходе? В прямом? Сдать: два ответа через подмножество $B, D, D_{FF}, N_{FSDP}, W$ и два предложения.

(c) Насколько большим может быть $N_{FSDP}$, прежде чем обратный проход упрётся в связь? А прямой? Сдать: два неравенства и два предложения обоснования.

Problem (tp_calcs): расчёты для тензорного параллелизма (4 балла)

В той же постановке — когда TP упирается в связь? (Стратегия TP из условия: $W_1^{(i)}$ и $W_2^{(i)}$ формы $(D, D_{FF}/N_{TP})$, $W_3^{(i)}$ — $(D_{FF}/N_{TP}, D)$.)

(a) По входу $dy$ формы $(B, D)$ выпишите обратный проход этой стратегии. Сдать: последовательность уравнений через $dy$, шардированные веса ($W_1^{(i)}, W_2^{(i)}, W_3^{(i)}$), сохранённые активации ($x, x_1^{(i)}, x_2^{(i)}, z^{(i)}, y^{(i)}$), коллективные операции и любые промежуточные переменные; уравнения должны давать градиенты каждого устройства $dW_1^{(i)}, dW_2^{(i)}, dW_3^{(i)}$ и выход обратного прохода $dx$. Можно взять нешардированный обратный проход из раздела 8.2 и изменить его.

(b) Сколько FLOPs на прямой проход при $N_{TP}$? На обратный? Сдать: два ответа через $B, D, D_{FF}, N_{TP}$ и два предложения.

(c) Сколько времени связи в прямом проходе? В обратном? Сдать: два ответа через подмножество $B, D, D_{FF}, N_{TP}, W$ и два предложения.

(d) Насколько большим может быть $N_{TP}$, прежде чем обратный проход упрётся в связь? А прямой? Сдать: два неравенства через подмножество $B, D, D_{FF}, C, W$ и два предложения.

Problem (fsdp_tp_calcs): расчёты для 2D-параллелизма (6 баллов)

В той же постановке — когда упирается в связь сочетание FSDP и TP?

(a) Сколько FLOPs на прямой проход при $N_{FSDP}$ FSDP + $N_{TP}$ TP? Сдать: ответ через $B, D, D_{FF}, N_{FSDP}, N_{TP}$ и одно предложение.

(b) Сколько времени связи в прямом проходе? Считайте, что связь по двум осям можно перекрыть (коллективы по оси FSDP перекрываются с коллективами по оси TP). Сдать: ответ через подмножество $B, D, D_{FF}, N_{FSDP}, N_{TP}, W$ и одно предложение. Подсказка: ответ — максимум из двух величин (стоимостей коллективов FSDP и TP).

(c) При оптимальных $N_{TP}$ и $N_{FSDP}$ насколько большим может быть $N = N_{TP} N_{FSDP}$, прежде чем прямой проход упрётся в связь? Сдать: неравенство с $N$ и выражением через подмножество $B, D, D_{FF}, C, W$, с несколькими предложениями и уравнениями обоснования.

(d) Теперь коллективы по осям FSDP и TP нельзя перекрыть (они делят одни сетевые ресурсы). При оптимальных $N_{TP}$ и $N_{FSDP}$ насколько большим может быть $N$? Округлять $N_{TP}$ и $N_{FSDP}$ до целых не нужно. Сдать: неравенство с обоснованием.

9. Лидерборд

Лидерборд проверяет скорость полного шага обучения модели на 8B: конфигурация ctx_len = 32768, vocab_size = 151936, d_model = 4096, d_ff = 11008, num_layers = 34, num_heads = 32, BF16, причинная маска, батч 2 — на двух B200. Модель специально трудно уместить в память. Нельзя менять поведение модели «вход–выход»; реализация должна быть своей. Идеи из условия: автоподбор тайлов Triton; настройка Triton и torch.compile; слитый AdamW; ядро, сливающее LM head и cross-entropy (не материализуя логиты [batch, seq_len, vocab]) вместе с обратным проходом; улучшение FlashAttention (обратный проход на Triton; два прохода — отдельно для $dQ$ и для $dK, dV$, без атомиков; ранняя остановка тайлов, целиком закрытых причинной маской; отдельная обработка диагональных тайлов; TMA на архитектурах новее Hopper); перевычисление активаций — только если без него не помещается.

Problem (leaderboard): самый быстрый шаг обучения (10 баллов)

Бенчмарк запускается с батчем 2 на двух GPU B200. Оценивается реальное время полного шага: прямой проход, потери, обратный проход и обновление AdamW. С пустого кэша PyTorch/Triton запуск должен уложиться в 10 минут, поэтому осторожнее со слишком агрессивными torch.compile и автоподбором Triton. Сдать: лучшее время полного шага прямой + обратный с AdamW. Ожидается, что результат будет лучше наивного базового уровня — 10 секунд. Отправка — в assignment2-systems-leaderboard.

---

Как выполнять на вашей машине (4 × RTX 4080, PCIe без NVLink)

1. Профилирование и Nsight — на одной карте. Nsight Systems (nsys) работает с GeForce. Задачи для xl, где нужен полный шаг обучения, делайте для large или с BF16 и перевычислением активаций (см. примечание о памяти выше). Это честно отметьте в отчёте — оценка для xl всё равно делается расчётом.

2. FlashAttention-2 на Triton — RTX 4080 подходит. Triton поддерживает архитектуру Ada (sm_89). Разделяемой памяти на SM меньше, чем у H100 (до 100 КБ против 228 КБ), поэтому для $d = 128$ берите тайлы поменьше (например, 64 × 64) и проверьте, что ядро запускается. Бенчмарк «на одной B200» делайте на одной RTX 4080 — длина 65 536 в FP32 для PyTorch-реализации, скорее всего, не поместится; это нормальный результат для таблицы.

3. Бенчмарк связи — 2 и 4 процесса. У вас нет 6 GPU: делайте 2 и 4. Ждите шинную скорость all-reduce в единицы — десяток ГБ/с (PCIe 4.0, P2P у GeForce обычно выключен драйвером, данные идут через RAM). Сравните с цифрами H100 из лекции 8 — эта разница и есть главный урок.

4. DDP, ZeRO-1, FSDP — 2 карты, модель меньше. Стандартная постановка условия — «1 узел, 2 GPU, xl». У вас xl в FP32 на две карты 32 ГБ без шардинга не влезет; берите large (или medium) для DDP-бенчмарков, а xl — для FSDP на 4 картах. Тесты test_ddp.py, test_sharded_optimizer.py, test_fsdp.py используют маленькие модели и подходят для ваших карт.

5. Перекрытие связи — особенно важно у вас. Медленная PCIe делает разницу между наивным DDP и DDP с перекрытием намного заметнее, чем на NVLink. Хорошая иллюстрация для отчёта: снимки Nsight, где all-reduce «спрятан» под обратным проходом.

6. Расчёты раздела 8 — бумага и карандаш. Подставьте в выведенные формулы свои $C$ (≈ 97,5 TFLOPS BF16 у RTX 4080, урок 5.7.1) и $W$ (ваш замер из шага 3) и оцените, сколько карт по PCIe имеет смысл для DP, FSDP и TP.

7. Лидерборд на модели 8B с контекстом 32 768 у вас не воспроизвести. Аналог для себя: самый быстрый шаг обучения large с контекстом 8192 на двух RTX 4080.

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

• Бенчмарк «врёт»? Почти всегда забыт torch.cuda.synchronize() или прогрев (первый запуск компилирует ядра и прогревает кэши аллокатора).
• FlashAttention расходится с эталоном на 1e-2? Проверьте, что $m$, $l$ и аккумулятор $O$ в FP32, что при обновлении максимума старые $l$ и $O$ умножаются на $e^{m_{old} - m_{new}}$ и что logsumexp сохраняется как $m + \log l$.
• Тесты DDP/FSDP иногда падают? Это гонки: ждите все асинхронные хэндлы (handle.wait()) перед шагом оптимизатора и не освобождайте собранные веса, пока на них ссылается граф autograd.
• Зависание на init_process_group? Свободен ли порт MASTER_PORT, все ли процессы дошли до коллективной операции в одном и том же порядке (урок 5.7.4, «Типичные ошибки»).

Что дальше

• Задание 3 — законы масштабирования.
• Лекции к заданию: 5–6 и 7–8.
• Обзор курса CS336.