Ultra-Scale Playbook, часть 1: память и вычисления на одной GPU

Это перевод вступления, обзорной главы, главы «Первые шаги: обучение на одной GPU» и приложения A0 («Ускоренный курс параллельного программирования») из книги The Ultra-Scale Playbook. Книга опубликована командой Hugging Face в 2025 году в виде интерактивной страницы (Space на Hugging Face) под лицензией Apache-2.0. Остальные части перевода:

• часть 2 — параллелизм по данным, ZeRO, тензорный параллелизм и параллелизм по последовательности;
• часть 3 — контекстный, конвейерный и экспертный параллелизм, 5D-параллелизм и поиск лучшей конфигурации;
• часть 4 — внутри GPU: ядра, слияние операций, FlashAttention, смешанная точность.

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

Читайте после уроков 5.4.1 «Смешанная точность» и 5.4.3 «Распределённое обучение»: там те же идеи даны короче и с упражнениями. Эта часть особенно полезна, чтобы научиться считать память — без этого не понять, зачем нужны все виды параллелизма из следующих частей. В оригинале много интерактивных виджетов и графиков; мы их не копируем, а описываем словами. Ссылки на рисунки и виджеты ведут к оригиналу.

Ваши 4× RTX 4080 соединены через PCIe без NVLink, а книга ведёт эксперименты на кластере из узлов по 8× H100 с NVLink внутри узла и InfiniBand между узлами. Когда в тексте говорится о «быстрой связи внутри узла», для вашей машины это не так: обмен между картами у вас примерно на порядок медленнее, чем у H100 в одном узле.

The Ultra-Scale Playbook: обучение LLM на кластерах GPU

Тысячи GPU, гудящих в идеальной гармонии. Вот что нужно, чтобы обучить самые мощные современные модели ИИ, — симфония вычислительной мощности, которая ещё недавно была доступна только элитным исследовательским лабораториям. Открытый код изменил эту картину, но не до конца. Да, можно скачать свежие модели Llama или DeepSeek. Да, можно прочитать их технические и экспериментальные отчёты. Но самая трудная часть — код обучения, знания и приёмы, нужные, чтобы согласовать работу GPU при обучении таких огромных систем, — остаётся скрытой за сложностью и разбросанной по разрозненным статьям и часто закрытым кодовым базам.

Время чтения: 2–4 дня. Для удобства не читайте с мобильного телефона.

Эта открытая книга призвана это изменить. Начиная с основ, мы пройдём путь, нужный, чтобы масштабировать обучение больших языковых моделей (LLM) с одной GPU на десятки, сотни и даже тысячи GPU, — с теорией, практическими примерами кода и воспроизводимыми бенчмарками.

По мере роста кластеров, на которых обучают такие модели, появились разные приёмы — параллелизм по данным, тензорный параллелизм, конвейерный параллелизм, контекстный параллелизм, а также ZeRO и слияние ядер (kernel fusion), — чтобы GPU всё время были загружены работой. Это заметно сокращает время обучения и позволяет использовать дорогое железо максимально эффективно. Эти приёмы распределённого обучения важны не только для создания исходных моделей: они стали необходимы и для дообучения больших моделей на специализированных данных, которое часто даёт лучшие результаты. В книге мы постепенно разберём все эти приёмы — от самых простых до самых изощрённых, — сохраняя единую сюжетную линию, чтобы было понятно, откуда взялся каждый метод.

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

Мы предполагаем, что вы в общих чертах знаете архитектуры современных LLM и примерно понимаете, как обучаются модели глубокого обучения, но распределённое обучение может быть для вас новым. Основы обучения моделей при необходимости можно найти в отличных курсах DeepLearning.AI или в учебниках PyTorch. Эту книгу можно считать второй частью трилогии, после нашего поста об обработке данных для предобучения (так называемый «пост о FineWeb»). Прочитав оба, вы будете знать почти всё основное о том, как сегодня строят сильные LLM; не хватит только «секретного соуса» — смешивания данных и выбора архитектуры (следите за третьей частью…).

Примечание переводчика

Третья часть трилогии вышла осенью 2025 года — это The Smol Training Playbook, наш пересказ есть в этом же разделе.

Мы очень благодарны всей команде distill.pub за шаблон, на котором построен этот пост.

Книга стоит на трёх общих основаниях.

1. Короткие введения в теорию и понятия. Прежде чем погрузиться в код и эксперименты, мы хотим, чтобы вы поняли, как каждый метод работает на высоком уровне и каковы его преимущества и ограничения. Например, вы узнаете, какие части языковой модели «съедают» память и в какой момент обучения это происходит. Вы узнаете, как обойти ограничения памяти, распараллелив модель, и как увеличить пропускную способность, добавляя GPU. В результате вы поймёте, как работает виджет, считающий разбивку памяти модели-трансформера.

Виджет в оригинале

В оригинале здесь интерактивный калькулятор памяти: число голов внимания, микробатч, размер скрытого слоя, длина последовательности, словарь, число слоёв, перевычисление (нет / выборочное / полное), уровень ZeRO (0–3), тензорный и параллелизм по данным и пресеты Llama 3 (Tiny, 8B, 70B, 405B). Конвейерного параллелизма в виджете нет — авторы оставили его «как упражнение читателю». «Не волнуйтесь, если не понимаете, что происходит в этом виджете. Для того мы здесь и собрались!» Есть и отдельный инструмент predict_memory, предсказывающий потребление памяти во время обучения.

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

• репозиторий Picotron создан для обучения, поэтому обычно реализует каждое понятие в одном коротком самодостаточном файле;
• чтобы посмотреть на код промышленного уровня, мы ссылаемся на реализации в Nanotron — кодовой базе, на которой в Hugging Face обучают модели.

Если вместо чтения блога или кода picotron хочется посмотреть видео о распределённом обучении, загляните на YouTube-канал Фердинанда.

3. Реальные бенчмарки эффективности обучения. Как на самом деле масштабировать обучение LLM, зависит от вашей инфраструктуры: какие чипы, какая связь между ними и т. д., поэтому единого рецепта мы дать не можем. Зато мы дадим способ сравнить несколько конфигураций — так мы сделали на своём кластере. Мы провели более 4100 распределённых экспериментов (больше 16 тысяч с учётом тестовых запусков) на числе GPU до 512, перебирая множество возможных схем распределённого обучения и размеров моделей.

Как видите, впереди много работы. Прежде чем спуститься в окопы распределённого обучения, бегло посмотрим на задачи, которые разбираются в книге.

Обзор с высоты

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

1. Потребление памяти. Это жёсткое ограничение: если шаг обучения не помещается в память, обучение не пойдёт.
2. Эффективность вычислений. Мы хотим, чтобы железо большую часть времени считало, поэтому нужно сократить время на передачу данных и на ожидание, пока другие GPU сделают свою работу.
3. Накладные расходы на коммуникацию. Их нужно минимизировать, потому что во время обмена GPU простаивают. Для этого мы будем стараться как можно лучше использовать пропускную способность внутри узла (быструю) и между узлами (медленнее), а также как можно сильнее перекрывать коммуникацию вычислениями.

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

Книга охватывает много тем, поэтому мы сделали шпаргалку, которая поможет ориентироваться и уловить главные выводы. Держите её под рукой в этих бурных водах!

Первые шаги: обучение на одной GPU

Заметка

В оригинале здесь можно послушать подкаст: ведущие NotebookLM обсуждают первые разделы книги.

Прежде чем масштабироваться на много GPU, быстро вспомним самые основы обучения модели. Когда модель обучается на одной GPU, обучение обычно состоит из трёх шагов:

1. прямой проход (forward pass) — входы проходят через модель, и получаются выходы;
2. обратный проход (backward pass) — вычисляются градиенты;
3. шаг оптимизации — по градиентам обновляются параметры.

Как увидим дальше, эти шаги могут повторяться или переплетаться, но пока начнём с простого.

flowchart LR
subgraph F["Прямой проход"]
  direction LR
  A1["Слой 1"] --> A2["Слой 2"] --> A3["Слой 3"] --> LOSS["Функция потерь"]
end
subgraph B["Обратный проход"]
  direction RL
  G3["Градиент слоя 3"] --> G2["Градиент слоя 2"] --> G1["Градиент слоя 1"]
end
LOSS --> G3
G1 --> OPT["Шаг оптимизатора: обновить веса"]

Схема переводчика вместо интерактивного рисунка оригинала: прямой проход, обратный проход, шаг оптимизации

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

Размер батча (batch size, $bs$) — один из важных гиперпараметров обучения: он влияет и на сходимость модели, и на пропускную способность.

Маленький батч бывает полезен в начале обучения, чтобы быстро пройти по ландшафту функции потерь к хорошей точке. Но дальше маленький батч оставляет градиенты шумными, и модель может не сойтись к наилучшему итоговому качеству. Другая крайность — большой батч: он даёт очень точные оценки градиента, но, как правило, хуже использует каждый обучающий токен, из-за чего сходимость медленнее и вычисления могут тратиться впустую. Хорошее раннее обсуждение темы — в статье OpenAI об обучении с большими батчами или в разделе 4.2 техотчёта MiniMax-01.

Например, при обучении DeepSeek-V3/R1 размер батча постепенно увеличивали с 3072 входных последовательностей до 15 360 на первых 469 млрд токенов, а затем держали 15 360 до конца обучения.

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

В сообществе, занимающемся предобучением LLM, размер батча обычно указывают в токенах, а не в примерах ($bst$ — batch size tokens, размер батча в токенах). Так числа обучения почти не зависят от точной длины входных последовательностей.

В простейшем случае, при обучении на одной машине, $bs$ (в примерах) и $bst$ связаны через длину входной последовательности модели ($seq$):

$$bst = bs \cdot seq$$

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

Удачный диапазон для современных LLM — порядка 4–60 млн токенов в батче. Размер батча и размер обучающего корпуса год от года растут: Llama 1 обучали с батчем около 4 млн токенов на 1,4 трлн токенов, а DeepSeek — с батчем около 60 млн токенов на 14 трлн токенов.

Первая проблема возникает, когда мы пытаемся масштабировать обучение до таких больших батчей: нехватка памяти (out-of-memory, OOM). Что делать, если в GPU не помещается целый батч нужного размера?

Сначала разберёмся, что вообще приводит к OOM. Это даст полезную интуицию о том, сколько памяти нужно для обучения.

Потребление памяти в трансформерах

При обучении нейросети в памяти хранится несколько вещей:

• веса модели;
• градиенты модели;
• состояния оптимизатора;
• активации, нужные для вычисления градиентов.

Заметка

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

• ядра CUDA обычно занимают 1–2 ГБ памяти GPU; это легко проверить, выполнив import torch; torch.ones((1, 1)).to("cuda") и посмотрев память GPU через nvidia-smi;
• часть памяти уходит на буферы и промежуточные результаты, а часть нельзя использовать из-за фрагментации.

Эти два слагаемых мы будем игнорировать: обычно они малы и постоянны.

Всё это хранится как тензоры разных форм и разной точности. Формы определяются гиперпараметрами: размером батча, длиной последовательности, размерностями скрытых слоёв модели, числом голов внимания, размером словаря и, как увидим позже, возможным шардированием модели. Точность — это форматы вроде FP32, BF16 или FP8, которым нужно соответственно 4, 2 или 1 байт на каждое значение тензора. Подробно о разных точностях и их компромиссах поговорим в разделе «Обучение со смешанной точностью» (часть 4); пока запомним, что форматы требуют разного объёма памяти, и это влияет на то, сколько занимают нужные нам данные.

Как быстро оценить потребление памяти по этим переменным? Простой способ — измерить его на практике.

Профилирование памяти

С помощью профилировщика PyTorch можно увидеть, как память выделяется в ходе обучения. Оказывается, использование памяти — не статичная величина: оно сильно меняется по ходу обучения и даже внутри одного шага.

Что на графике оригинала

График показывает несколько шагов обучения. Видно, как на прямом проходе растут активации, на обратном — накапливаются градиенты, а активации освобождаются; после первого шага появляются состояния оптимизатора. Как профилировать свою модель, описано в приложении A1 оригинала («Distributed Training Profiling»).

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

Как сказано, первый шаг выглядит иначе: активации быстро растут, а потом какое-то время держатся на одном уровне. Почему? На первом шаге кэширующий аллокатор PyTorch проделывает много подготовительной работы: заранее раскладывает выделения памяти, чтобы на следующих шагах не искать свободные блоки, и это их ускоряет (см. блог Зака). После первого шага появляются ещё и состояния оптимизатора, которые обычно сдвигают вверх потребление памяти на следующих шагах.

Замечали, что иногда обучение проходит первый шаг, а на следующих падает с OOM? Объяснение — как раз в состояниях оптимизатора, которые появляются после первого шага.

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

Память для весов, градиентов и состояний оптимизатора

Начнём с первых трёх пунктов списка: веса модели, градиенты и состояния оптимизатора. Их объём оценить довольно легко.

Для простой LLM-трансформера число параметров даётся следующей формулой:

$$N = h \cdot v + L \cdot (12 \cdot h^2 + 13 \cdot h) + 2 \cdot h$$

Мы не учитываем позиционные эмбеддинги, потому что не используем фиксированные позиционные эмбеддинги.

Здесь $h$ — скрытая размерность, $v$ — размер словаря, $L$ — число слоёв модели. Из формулы видно, что при большой скрытой размерности доминирует слагаемое с $h^2$: только оно растёт квадратично, когда мы увеличиваем модель.

Пример переводчика

Возьмём $h = 4096$, $v = 128\,256$, $L = 32$ (размеры, похожие на Llama 3 8B). Тогда $h \cdot v \approx 0{,}525$ млрд, $L \cdot (12h^2 + 13h) \approx 6{,}44$ млрд, итого $N \approx 6{,}97$ млрд. У настоящей Llama 3 8B параметров больше (8,03 млрд): у неё SwiGLU с более широким FFN и раздельные матрицы входных и выходных эмбеддингов, а формула описывает «простой» трансформер. Порядок величины формула даёт верно.

Память для параметров и градиентов — это просто число параметров, умноженное на число байт на параметр. В старом добром обучении в полной точности (FP32) и параметрам, и градиентам нужно по 4 байта, а оптимизатору Adam — ещё хранить момент и дисперсию, то есть ещё 8 байт на параметр (по 4 байта на каждое). Итого:

$$\begin{aligned} & m_{params} = 4 \cdot N \\ & m_{grad} = 4 \cdot N \\ & m_{opt} = (4+4) \cdot N \end{aligned}$$

Посмотрим, что изменится в меньшей точности. Ради стабильности (см. раздел о смешанной точности в части 4) обычно учат не целиком в низкой точности, а в смеси высокой и низкой, которую называют «смешанной точностью» (mixed precision). Сейчас по умолчанию большинство вычислений идёт в BF16 — 2 байта на параметр и на градиент, — а дополнительно хранится копия весов и градиентов в FP32; всего 12 байт на параметр. Кроме параметров и градиентов, нужно хранить состояния оптимизатора: для Adam это момент и дисперсия, которые ради численной стабильности обычно хранят в FP32, по 4 байта каждое.

Подробнее увидим ниже, когда дойдём до методов ZeRO.

Сводка:

$$\begin{aligned} & m_{params} = 2 \cdot N \\ & m_{grad} = 2 \cdot N \\ & m_{params\_fp32} = 4 \cdot N \\ & m_{opt} = (4+4) \cdot N \end{aligned}$$

Заметка

Некоторые библиотеки хранят градиенты в FP32, что требует ещё $m_{params\_fp32} = 4 \cdot N$ памяти. Так делают, например, в Nanotron: BF16 теряет точность на малых значениях, а стабильность для нас всегда важнее. Подробнее — в этом обсуждении DeepSpeed.

Заметка

Копию параметров в FP32 ($m_{params\_fp32}$) в литературе и в коде иногда называют «мастер-весами» (master weights).

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

Вот сколько примерно памяти нужно модели (полная и смешанная точность дают одинаковые итоговые числа):

Параметров в модели	FP32 или BF16 без накопления градиентов в FP32	BF16 с накоплением градиентов в FP32
1B	16 ГБ	20 ГБ
7B	112 ГБ	140 ГБ
70B	1120 ГБ	1400 ГБ
405B	6480 ГБ	8100 ГБ

Обучение в FP8 вместо BF16 ещё уменьшило бы память, но оно менее стабильно. Это очень активная область исследований (см. этот твит), и мы ещё вернёмся к ней.

Как видим, уже на 7B (!) веса, градиенты и состояния оптимизатора в сумме занимают очень много и превышают память типичной GPU (например, 80 ГБ у H100).

Примечание переводчика: ваша машина

У RTX 4080 из вашей конфигурации 32 ГБ. Полное обучение (не LoRA) модели 1B со смешанной точностью требует около 20 ГБ только на веса, градиенты и оптимизатор — плюс активации. Модель 7B без шардирования (ZeRO/FSDP из части 2) не поместится даже в сумму памяти всех четырёх карт (128 ГБ): нужно 140 ГБ без учёта активаций.

Но пока будем работать с моделями, которые помещаются в одну GPU, и посмотрим на последний большой пункт нашего бюджета памяти — активации.

Память для активаций

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

$$m_{act} = L \cdot seq \cdot bs \cdot h \cdot \left(34 + \frac{5 \cdot n_{heads} \cdot seq}{h}\right)$$

Здесь $L$ — число слоёв, $seq$ — длина последовательности, $bs$ — размер батча в примерах, $h$ — скрытая размерность модели, $n_{heads}$ — число голов.

Точный вывод этих чисел можно найти в исходной статье NVIDIA о перевычислении: по сути, нужно аккуратно учесть размеры всех промежуточных активаций между операциями в слое трансформера.

Интересное наблюдение: для заданной модели потребление памяти не постоянно — оно растёт линейно с размером батча и квадратично с длиной последовательности. Значит, именно память под активации «взрывается», когда мы увеличиваем батч или учим на более длинных последовательностях. По этой формуле можно посмотреть, как меняется память для разных длин последовательности, например для моделей Llama (bs=1).

Пример переводчика

Подставим $L = 32$, $h = 4096$, $n_{heads} = 32$, $bs = 1$ (размеры как у Llama 3 8B; результат — в байтах). При $seq = 1024$: $32 \cdot 1024 \cdot 4096 \cdot (34 + 40) \approx 9{,}9$ ГБ. При $seq = 4096$ скобка равна $34 + 160 = 194$, и выходит уже ≈104 ГБ. При $seq = 8192$ — ≈380 ГБ. Длина выросла в 8 раз, а память — почти в 40: так работает квадратичное слагаемое.

Графики оригинала рассказывают поразительную историю: для коротких последовательностей (или маленьких батчей) память под активации почти не заметна, но примерно с 2–4 тысяч токенов она начинает занимать очень много, тогда как память под параметры, градиенты и состояния оптимизатора (о них ниже) почти не зависит от длины последовательности и размера батча.

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

Можно ли как-то укротить этот «взрыв активаций»? Хороший вопрос, читатель!

Пора объяснить наш первый приём — перевычисление активаций (activation recomputation). Он поможет ограничить память под активации; сегодня это обязательный инструмент в наборе для обучения больших моделей.

Перевычисление активаций

Общая идея перевычисления активаций (его ещё называют gradient checkpointing или rematerialization) — отбросить часть активаций при прямом проходе, чтобы сэкономить память, и потратить немного лишних вычислений, чтобы пересчитать их на лету при обратном проходе. Без перевычисления мы храним каждое скрытое состояние между двумя обучаемыми операциями (например, FFN, LayerNorm и т. д.), чтобы использовать их на обратном проходе для вычисления градиентов. С перевычислением мы обычно храним активации лишь в нескольких ключевых точках архитектуры, а остальные отбрасываем и пересчитываем на лету при обратном проходе от ближайших сохранённых. По сути, мы повторно выполняем часть прямого прохода, обменивая память на вычисления.

flowchart LR
X["Вход слоя (сохранён)"] --> A["Внимание"] --> N["LayerNorm"] --> F["FFN"] --> Y["Выход слоя (сохранён)"]
Y -. "обратный проход: промежуточные активации пересчитываются от сохранённого входа" .-> X

Схема переводчика: при перевычислении сохраняются только «контрольные точки», остальное считается заново

Есть несколько стратегий выбора ключевых активаций для хранения.

• Полная (full). Сохраняем активации на стыке каждого слоя трансформера. Эту стратегию обычно называют «полной», потому что она требует прямого прохода через каждый слой, то есть по сути добавляет ещё один полный прямой проход во время обратного. Она экономит больше всего памяти, но дороже всех по вычислениям: обычно увеличивает объём вычислений и время до 30–40 %, что очень заметно.
• Выборочная (selective). Как правило, можно сделать лучше, чем «полная». Авторы статьи о перевычислении подробно проанализировали, какие активации растут сильнее всего и дешевле всего пересчитываются по числу операций с плавающей точкой (FLOPs). Оказалось, что в эту категорию попадают вычисления внимания, поэтому их обычно можно отбросить, а сохранять дорогие вычисления FFN. Для GPT-3 (175B) это означает сокращение памяти под активации на 70 % ценой 2,7 % вычислений.

В свежих моделях вроде DeepSeek-V3 используется выборочное перевычисление, причём память под активации внимания ещё меньше благодаря так называемому «многоголовому латентному вниманию» (Multi-Head Latent Attention, MLA).

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

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

Заметка

Когда вы измеряете, насколько эффективно ваша схема обучения использует GPU/TPU/ускоритель, обычно нужно учесть перевычисление при подсчёте общего числа FLOPs (операций с плавающей точкой) и сравнить его с теоретическим максимумом FLOPS (операций в секунду) ускорителя. FLOPs шага обучения с учётом перевычисления называют «аппаратными FLOPs» (hardware FLOPs) — это реальное число операций на ускорителе. Разделив их на длительность шага (в секундах), получаем фактически достигнутые FLOPS. А разделив достигнутые FLOPS на максимум ускорителя, получаем утилизацию аппаратных FLOPS (hardware FLOPS utilization, HFU).

Но в конечном счёте важно общее время обучения модели на заданном наборе данных. Поэтому, например, при сравнении разных GPU/TPU/ускорителей тот, у кого хватает памяти обойтись без перевычисления и поэтому выполняющий меньше операций (меньше аппаратных FLOPs), но обучающий быстрее, должен получать за это плюс, а не минус. Отсюда альтернатива — утилизация FLOPs модели (model FLOPS utilization, MFU): в отличие от HFU, она учитывает только операции, необходимые для прямого и обратного проходов через модель, без перевычисления. Эта величина характеризует скорее модель, чем конкретную реализацию обучения.

Большинство современных фреймворков обучения используют FlashAttention (подробнее — в части 4), в которую перевычисление активаций встроено изначально: оценки и матрицы внимания пересчитываются при обратном проходе, а не хранятся. Так что большинство пользователей FlashAttention уже применяют выборочное перевычисление.

Как вы уже поняли, перевычисление активаций немного увеличивает число операций, но заметно сокращает накладные расходы на доступ к памяти.

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

Теперь, зная о перевычислении, мы можем укротить память под активации, которую видели на графиках!

Однако активации по-прежнему линейно зависят от размера батча, а все профили на графиках выше были при bs=1, так что при переходе к большим батчам проблема может вернуться. К счастью, у нас есть второй инструмент — на помощь приходит накопление градиентов (gradient accumulation)!

Накопление градиентов

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

Назовём размер батча для одного прямого прохода размером микробатча ($mbs$), а общий размер батча между двумя шагами оптимизатора — глобальным размером батча ($gbs$). Если мы делаем один шаг оптимизатора на каждые восемь прямых и обратных проходов, глобальный размер батча в восемь раз больше размера микробатча.

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

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

$$bs = gbs = mbs \times grad\_acc$$

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

flowchart LR
M1["Микробатч 1: прямой + обратный"] --> ACC["Буфер градиентов: сумма"]
M2["Микробатч 2: прямой + обратный"] --> ACC
M3["…"] --> ACC
M8["Микробатч 8: прямой + обратный"] --> ACC
ACC --> AVG["Среднее"] --> OPT["Один шаг оптимизатора"]

Схема переводчика: накопление градиентов, grad_acc = 8

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

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

Недостаток в том, что для одного шага оптимизации нужно несколько последовательных прямых и обратных проходов, а это увеличивает вычислительные накладные расходы и замедляет обучение. Бесплатных обедов не бывает!

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

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

Профилирование вычислений и коммуникации на GPU

Профилировщик PyTorch позволяет отследить и визуализировать, что именно происходит на CPU и на GPU во время обучения. Он встроен в PyTorch. Вот как им пользоваться:

import torch

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,   # записывать события CPU
        torch.profiler.ProfilerActivity.CUDA,  # и события GPU
    ],
    # 1 шаг пропустить, 1 шаг «прогреть», 3 шага записать
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),
    # сохранить трассу для TensorBoard в папку ./log/profile
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/profile'),
    with_stack=True,  # запоминать стек вызовов Python
) as prof:
    for step in range(steps):
        train_step()  # ваш шаг обучения
        prof.step()   # сообщить профилировщику, что шаг закончился

Получится трасса, которую можно посмотреть в TensorBoard или в просмотрщике трасс Chrome. На трассе видно:

• поток CPU, асинхронно запускающий ядра на GPU;
• несколько потоков CUDA (CUDA streams), параллельно занятых вычислениями и коммуникацией;
• время выполнения ядер и выделение памяти.

На рисунке оригинала — пример такой трассы: поток CPU асинхронно запускает ядра на GPU, а вычислительные ядра и коммуникация идут параллельно в разных потоках CUDA.

Трасса помогает найти узкие места, например:

• последовательные вычисления и коммуникацию, которые можно было бы перекрыть;
• простой GPU в ожидании передачи данных;
• синхронизации CUDA и перемещение данных между CPU и GPU;
• накладные расходы на запуск ядер на GPU.

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

Теперь возьмём рабочую станцию побольше, с парой GPU, и начнём изучать первый приём масштабирования — параллелизм по данным (data parallelism). Как мы увидим, это просто параллельная версия накопления градиентов. Продолжение — в части 2, а перед ней полезно прочитать приложение ниже.

Приложение A0. Ускоренный курс параллельного программирования

На протяжении книги мы масштабируем обучение LLM с одной до сотен GPU. Для этого нужно обмениваться весами, градиентами и данными между всеми машинами и синхронизировать их. Для этого есть набор распределённых шаблонов — коллективные операции (collective operations). В этом разделе — короткий ускоренный курс по ним: Broadcast, AllReduce, Scatter и другие. Поехали!

Общая постановка: есть несколько независимых узлов — ядра CPU, GPU или вычислительные узлы. Каждый что-то вычисляет, а потом мы хотим передать результат или его часть другим узлам для следующего шага вычислений ($t+1$).

Может быть, нужно отправить результат с одного узла на все остальные или сложить промежуточные результаты всех узлов, чтобы получить общий итог. Обычно есть один узел с особым статусом — его называют корнем (root), — который служит источником или получателем в некоторых операциях. Начнём с одного из простейших примитивов — Broadcast.

Broadcast

Очень частая ситуация: данные есть на узле 1, и вы хотите поделиться ими со всеми остальными узлами, чтобы они что-то с ними посчитали. Операция Broadcast («широковещательная рассылка») делает ровно это.

flowchart LR
R["Узел 0 (root): [1,2,3,4,5]"] --> N1["Узел 1: [1,2,3,4,5]"]
R --> N2["Узел 2: [1,2,3,4,5]"]
R --> R2["Узел 0: [1,2,3,4,5]"]

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

Коллективные операции встроены в PyTorch, поэтому легко написать маленький пример, показывающий работу Broadcast. Сначала нужно инициализировать группу процессов вызовом dist.init_process_group: она настраивает бэкенд коммуникации (о NCCL поговорим позже), определяет, сколько всего воркеров (то есть узлов), и присваивает каждому номер — ранг (rank), который можно узнать через dist.get_rank. Наконец, она устанавливает соединение между воркерами.

Чтобы показать dist.broadcast, создадим на rank=0 тензор с ненулевыми значениями, а на остальных воркерах — тензоры из нулей. Затем разошлём тензор с rank=0 всем остальным рангам через dist.broadcast(tensor, src=0):

import torch
import torch.distributed as dist

def init_process():
    dist.init_process_group(backend='nccl')   # бэкенд для обмена между GPU
    torch.cuda.set_device(dist.get_rank())   # процесс с рангом i работает на GPU i

def example_broadcast():
    if dist.get_rank() == 0:
        tensor = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32).cuda()
    else:
        tensor = torch.zeros(5, dtype=torch.float32).cuda()
    print(f"Before broadcast on rank {dist.get_rank()}: {tensor}")
    dist.broadcast(tensor, src=0)            # разослать тензор ранга 0 всем
    print(f"After broadcast on rank {dist.get_rank()}: {tensor}")

init_process()
example_broadcast()

Скрипт запускается командой torchrun --nproc_per_node=3 dist_op.py (для этого нужны три GPU, или поменяйте nproc_per_node). Вывод должен быть таким:

Before broadcast on rank 0: tensor([1., 2., 3., 4., 5.], device='cuda:0')
Before broadcast on rank 1: tensor([0., 0., 0., 0., 0.], device='cuda:1')
Before broadcast on rank 2: tensor([0., 0., 0., 0., 0.], device='cuda:2')

After broadcast on rank 0: tensor([1., 2., 3., 4., 5.], device='cuda:0')
After broadcast on rank 1: tensor([1., 2., 3., 4., 5.], device='cuda:1')
After broadcast on rank 2: tensor([1., 2., 3., 4., 5.], device='cuda:2')

Отлично, всё работает как ожидалось. Учтите, что сообщения разных рангов могут печататься в другом порядке: мы не управляем тем, какой print выполнится первым (здесь мы упорядочили их для удобства чтения). Перейдём к шаблонам Reduce и AllReduce!

Примечание переводчика

В оригинальном коде приложения есть опечатки (dist.initi_process_group, вызов example_broadcats(), переменная rank без определения в example_reduce). В переводе они исправлены. На вашей машине из 4× RTX 4080 все примеры запускаются командой torchrun --nproc_per_node=4 dist_op.py; если NCCL капризничает без NVLink, попробуйте переменную окружения NCCL_P2P_DISABLE=1.

Reduce и AllReduce

Шаблоны свёртки (reduce) — одни из самых фундаментальных в распределённой обработке данных. Идея в том, чтобы объединить данные со всех узлов функцией f() — например, суммой или средним. В Reduce результат отправляется только корневому узлу, а в AllReduce он рассылается всем узлам.

Конечно, нет никакого волшебного «свободно парящего» узла, который сам выполнит такую операцию; обычно каждый узел делает часть вычислений, а узлы организованы в кольцо или дерево. Простой пример: нужно сложить числа на всех узлах, и узлы соединены кольцом. Первый узел отправляет своё число соседу, тот прибавляет к нему своё и передаёт дальше. Пройдя круг, первый узел получит общую сумму.

Вот код простой операции Reduce, складывающей тензоры. Операция задаётся параметром op=dist.ReduceOp.SUM (другие поддерживаемые операции — в документации PyTorch):

def example_reduce():
    rank = dist.get_rank()
    tensor = torch.tensor([rank + 1] * 5, dtype=torch.float32).cuda()
    print(f"Before reduce on rank {rank}: {tensor}")
    dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)   # сумма попадёт только на ранг 0
    print(f"After reduce on rank {rank}: {tensor}")

init_process()
example_reduce()

Обратите внимание: при Reduce обновляется только тензор на узле dst:

Before reduce on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
Before reduce on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')
Before reduce on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')

After reduce on rank 0: tensor([6., 6., 6., 6., 6.], device='cuda:0')
After reduce on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')
After reduce on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')

Аналогично выполняется AllReduce (здесь получателя указывать не нужно):

def example_all_reduce():
    rank = dist.get_rank()
    tensor = torch.tensor([rank + 1] * 5, dtype=torch.float32).cuda()
    print(f"Before all_reduce on rank {rank}: {tensor}")
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)      # сумма окажется на всех рангах
    print(f"After all_reduce on rank {rank}: {tensor}")

init_process()
example_all_reduce()

В этом случае результат есть на всех узлах:

Before all_reduce on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
Before all_reduce on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')
Before all_reduce on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')

After all_reduce on rank 0: tensor([6., 6., 6., 6., 6.], device='cuda:0')
After all_reduce on rank 1: tensor([6., 6., 6., 6., 6.], device='cuda:1')
After all_reduce on rank 2: tensor([6., 6., 6., 6., 6.], device='cuda:2')

Перейдём к следующей операции. Часто каждый узел сам по себе выполняет много сложных вычислений, и нужно поделиться итоговыми результатами между всеми узлами. Для этого служат Gather и AllGather. Посмотрим!

Gather и AllGather

Gather и AllGather похожи на Broadcast тем, что распределяют данные между узлами без изменений. Главное отличие: нет одного значения, которым один узел делится со всеми; у каждого узла свой кусок данных, и мы хотим собрать все данные на одном узле (Gather) или на всех узлах (AllGather). На рисунке оригинала пунктиром показаны данные, которые никуда не перемещаются (они уже на нужном узле).

flowchart LR
subgraph S1["До"]
  A0["GPU 0: a"]
  A1["GPU 1: b"]
  A2["GPU 2: c"]
end
subgraph S2["После AllGather"]
  B0["GPU 0: a b c"]
  B1["GPU 1: a b c"]
  B2["GPU 2: a b c"]
end
S1 --> S2

Схема переводчика: AllGather — каждый узел получает куски всех узлов

Для Gather нужно подготовить контейнер, куда будут сложены собранные тензоры, — в этом примере объект gather_list:

def example_gather():
    rank = dist.get_rank()
    tensor = torch.tensor([rank + 1] * 5, dtype=torch.float32).cuda()
    if rank == 0:
        # только на корне готовим место для тензоров всех рангов
        gather_list = [
            torch.zeros(5, dtype=torch.float32).cuda()
            for _ in range(dist.get_world_size())
        ]
    else:
        gather_list = None
    print(f"Before gather on rank {rank}: {tensor}")
    dist.gather(tensor, gather_list, dst=0)
    if rank == 0:
        print(f"After gather on rank 0: {gather_list}")

init_process()
example_gather()

И видим, что gather_list действительно содержит тензоры всех рангов:

Before gather on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
Before gather on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')
Before gather on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')

After gather on rank 0: [tensor([1., 1., 1., 1., 1.], device='cuda:0'),
                         tensor([2., 2., 2., 2., 2.], device='cuda:0'),
                         tensor([3., 3., 3., 3., 3.], device='cuda:0')]

Для AllGather нужно изменить только одно: место под результат понадобится на каждом узле:

def example_all_gather():
    rank = dist.get_rank()
    tensor = torch.tensor([rank + 1] * 5, dtype=torch.float32).cuda()
    # место под результат — на каждом ранге
    gather_list = [
        torch.zeros(5, dtype=torch.float32).cuda()
        for _ in range(dist.get_world_size())
    ]
    print(f"Before all_gather on rank {rank}: {tensor}")
    dist.all_gather(gather_list, tensor)
    print(f"After all_gather on rank {rank}: {gather_list}")

init_process()
example_all_gather()

Теперь у каждого узла есть все данные:

Before all_gather on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
Before all_gather on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')
Before all_gather on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')

After all_gather on rank 0: [tensor([1., 1., 1., 1., 1.], device='cuda:0'),
                             tensor([2., 2., 2., 2., 2.], device='cuda:0'),
                             tensor([3., 3., 3., 3., 3.], device='cuda:0')]
After all_gather on rank 1: [tensor([1., 1., 1., 1., 1.], device='cuda:1'),
                             tensor([2., 2., 2., 2., 2.], device='cuda:1'),
                             tensor([3., 3., 3., 3., 3.], device='cuda:1')]
After all_gather on rank 2: [tensor([1., 1., 1., 1., 1.], device='cuda:2'),
                             tensor([2., 2., 2., 2., 2.], device='cuda:2'),
                             tensor([3., 3., 3., 3., 3.], device='cuda:2')]

А что обратно Gather? Когда все данные на одном узле и их нужно разрезать и раздать узлам, возможно, с промежуточной обработкой. Для этого есть Scatter, а если над данными перед раздачей выполняется операция — ReduceScatter.

Scatter и ReduceScatter

Как следует из названия, Scatter («рассеять») берёт данные с одного узла и разносит их по всем узлам, отдавая каждому свой срез. Этим он отличается от Broadcast, который отправляет каждому узлу полную копию без нарезки, и логически обратен Gather.

ReduceScatter чуть сложнее. Как и в AllReduce, к данным со всех узлов применяется операция. Но, в отличие от AllReduce, где каждый узел получает весь выходной тензор, в ReduceScatter каждый узел получает только свой срез выходного тензора.

Scatter в коде записывается как противоположность Gather: вместо списка тензоров-приёмников мы готовим исходные данные как список тензоров, которые хотим раздать. Ещё нужно указать src:

def example_scatter():
    rank = dist.get_rank()
    if rank == 0:
        # на корне — список кусков: i-й кусок уйдёт рангу i
        scatter_list = [
            torch.tensor([i + 1] * 5, dtype=torch.float32).cuda()
            for i in range(dist.get_world_size())
        ]
        print(f"Rank 0: Tensor to scatter: {scatter_list}")
    else:
        scatter_list = None
    tensor = torch.zeros(5, dtype=torch.float32).cuda()
    print(f"Before scatter on rank {rank}: {tensor}")
    dist.scatter(tensor, scatter_list, src=0)
    print(f"After scatter on rank {rank}: {tensor}")

init_process()
example_scatter()

В итоге пустые тензоры заполняются содержимым scatter_list:

Rank 0: Tensor to scatter: [tensor([1., 1., 1., 1., 1.], device='cuda:0'),
                            tensor([2., 2., 2., 2., 2.], device='cuda:0'),
                            tensor([3., 3., 3., 3., 3.], device='cuda:0')]
Before scatter on rank 0: tensor([0., 0., 0., 0., 0.], device='cuda:0')
Before scatter on rank 1: tensor([0., 0., 0., 0., 0.], device='cuda:1')
Before scatter on rank 2: tensor([0., 0., 0., 0., 0.], device='cuda:2')

After scatter on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
After scatter on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')
After scatter on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')

Чтобы показать логику ReduceScatter, создадим данные поинтереснее. На каждом узле сделаем список двухэлементных векторов, зависящих от ранга узла через степень и сдвиг:

def example_reduce_scatter():
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    # j-й элемент списка: вектор [(rank+1)*1, (rank+1)*2] в степени j+1
    input_tensor = [
        torch.tensor([(rank + 1) * i for i in range(1, 3)], dtype=torch.float32).cuda() ** (j + 1)
        for j in range(world_size)
    ]
    output_tensor = torch.zeros(2, dtype=torch.float32).cuda()
    print(f"Before ReduceScatter on rank {rank}: {input_tensor}")
    dist.reduce_scatter(output_tensor, input_tensor, op=dist.ReduceOp.SUM)
    print(f"After ReduceScatter on rank {rank}: {output_tensor}")

init_process()
example_reduce_scatter()

По выводу видно, какие данные мы создали, и сразу виден ReduceScatter в действии: первый ранг получил сумму первых тензоров со всех узлов, второй ранг — сумму вторых и так далее:

Before ReduceScatter on rank 0: [tensor([1., 2.], device='cuda:0'),
                                 tensor([1., 4.], device='cuda:0'),
                                 tensor([1., 8.], device='cuda:0')]
Before ReduceScatter on rank 1: [tensor([2., 4.], device='cuda:1'),
                                 tensor([ 4., 16.], device='cuda:1'),
                                 tensor([ 8., 64.], device='cuda:1')]
Before ReduceScatter on rank 2: [tensor([3., 6.], device='cuda:2'),
                                 tensor([ 9., 36.], device='cuda:2'),
                                 tensor([ 27., 216.], device='cuda:2')]

After ReduceScatter on rank 0: tensor([ 6., 12.], device='cuda:0')
After ReduceScatter on rank 1: tensor([14., 56.], device='cuda:1')
After ReduceScatter on rank 2: tensor([ 36., 288.], device='cuda:2')

Теперь быстро посмотрим на распространённую реализацию AllReduce через ReduceScatter и AllGather — кольцевой AllReduce (Ring AllReduce).

Кольцевой AllReduce

Кольцевой AllReduce (Ring AllReduce) — реализация AllReduce, оптимизированная для масштабирования. Вместо того чтобы все устройства общались друг с другом напрямую (а это может создать узкие места), кольцевой AllReduce раскладывается на два ключевых шага: ReduceScatter и AllGather. Вот как это работает.

1. ReduceScatter
   • Каждое устройство делит свои данные (например, градиенты) на $N$ кусков ($N$ — число GPU) и отправляет один кусок соседу. Одновременно оно получает кусок от другого соседа.
   • Получив кусок, устройство прибавляет (reduce) к нему свой соответствующий кусок.
   • Это продолжается по кругу, пока у каждого устройства не окажется полностью свёрнутый кусок — сумма градиентов со всех устройств для этого куска.
2. AllGather
   • Теперь каждому устройству нужно собрать полностью свёрнутые куски с других устройств.
   • Каждое устройство отправляет свой свёрнутый кусок соседу и получает свёрнутый кусок от другого соседа.
   • Устройства продолжают пересылать полученные куски, пока у каждого не окажутся все свёрнутые куски, то есть полный суммарный градиент.

В оригинале это показано анимациями: 5 GPU, у каждой тензор длины 5. Первая анимация — шаг ReduceScatter, в конце которого каждая GPU получает свёрнутый результат для своего куска данных; вторая — шаг AllGather, в конце которого у каждой GPU полный результат AllReduce.

Вы могли заметить, что каждая из $N$ GPU отправляет и получает значения $N-1$ раз и на шаге ReduceScatter, и на шаге AllGather. За одну передачу каждая GPU отправляет $\frac{K}{N}$ значений, где $K$ — общее число значений в массиве, который суммируется по GPU. Поэтому общий объём данных, передаваемых на каждую GPU и с неё, равен $2 \times (N-1) \times \frac{K}{N}$. При большом $N$ (числе GPU) это примерно $2 \times K$, где $K$ — общее число параметров.

О AllReduce нужно помнить две вещи:

1. Стоимость коммуникации AllReduce при большом $N$ примерно $2 \times K$.
2. AllReduce можно разложить на ReduceScatter и следующий за ним AllGather. Стоимость каждой из этих двух операций — половина стоимости AllReduce, то есть примерно $K$.

Пример переводчика

У вас $N = 4$ GPU. Для модели 1B в BF16 градиенты занимают $K = 2$ ГБ. Кольцевой AllReduce передаёт с каждой карты $2 \cdot 3 \cdot \frac{2}{4} = 3$ ГБ. При реальной скорости PCIe 4.0 порядка 20–25 ГБ/с это примерно 0,12–0,15 с на каждый шаг — поэтому так важно перекрывать коммуникацию с вычислениями (см. часть 2).

Как видите, такая реализация эффективно использует даже ограниченную пропускную способность между узлами.

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

Barrier

Барьер (barrier) — простая операция синхронизации всех узлов. Барьер не снимается, пока до него не дойдут все узлы. Только после этого узлам разрешено продолжать вычисления.

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

import time

def example_barrier():
    rank = dist.get_rank()
    t_start = time.time()
    print(f"Rank {rank} sleeps {rank} seconds.")
    time.sleep(rank)  # имитируем разное время обработки
    dist.barrier()    # ждём, пока все ранги дойдут до этой строки
    print(f"Rank {rank} after barrier time delta: {time.time()-t_start:.4f}")

init_process()
example_barrier()

Видно, что хотя первый ранг вообще не спал, ему тоже понадобилось 2 секунды, чтобы пройти барьер:

Rank 0 sleeps 0 seconds.
Rank 1 sleeps 1 seconds.
Rank 2 sleeps 2 seconds.

Rank 0 after barrier time delta: 2.0025
Rank 1 after barrier time delta: 2.0025
Rank 2 after barrier time delta: 2.0024

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

Прежде чем перейти к практическим реализациям распределённого обучения, раскроем одну тайну: что такое NCCL?

NCCL

Обучая большие модели на многих GPU, иногда можно наткнуться на золото, но никель (nickel — созвучно NCCL 🥁) встретится обязательно! Что это такое?

Есть несколько библиотек, реализующих коллективную коммуникацию, которые поддерживает PyTorch: классический MPI (Message Passing Interface), Gloo от Meta и, наконец, NCCL (NVIDIA Collective Communications Library). Все они дают похожие коллективные операции, но оптимизированы под разное железо: NCCL создан для эффективного обмена GPU–GPU, а MPI и Gloo — для обмена CPU–CPU или CPU–GPU. У PyTorch есть хорошее руководство по выбору, но суть такова:

• обучение на GPU — используйте NCCL;
• обучение на CPU — используйте Gloo.

Есть и более тонкие моменты в дереве решений — их мы оставляем читателю изучить в упомянутом руководстве PyTorch.

Что пропущено

Приложения A1 (профилирование распределённого обучения), A2 (типичные масштабы обучения LLM) и A3 (математика перекрытия вычислений и коммуникации) мы не переводили: A3 по смыслу повторяет оценки из частей 2–3, а A1 и A2 удобнее читать в оригинале рядом с интерактивными графиками.

Словарик

EN	RU
forward / backward pass	прямой / обратный проход
batch size (bs), batch size tokens (bst)	размер батча в примерах / в токенах
micro-batch size (mbs), global batch size (gbs)	размер микробатча, глобальный размер батча
out-of-memory (OOM)	нехватка памяти
optimizer states	состояния оптимизатора
master weights	мастер-веса (копия весов в FP32)
mixed precision	смешанная точность
activation recomputation, gradient checkpointing	перевычисление активаций
gradient accumulation	накопление градиентов
HFU / MFU	утилизация аппаратных FLOPS / утилизация FLOPs модели
collective operations	коллективные операции
rank, root	ранг (номер процесса), корень
Broadcast, Reduce, AllReduce, Gather, AllGather, Scatter, ReduceScatter, Barrier	названия коллективных операций не переводятся
Ring AllReduce	кольцевой AllReduce