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

Тысячи GPU, гудящих в идеальной гармонии. Именно это требуется для обучения самых мощных современных AI-моделей: симфония вычислительной мощности, которая до недавнего времени оставалась прерогативой элитных исследовательских лабораторий. Open source радикально изменил этот ландшафт, но не полностью. Да, вы можете скачать последние модели Llama или DeepSeek. Да, вы можете изучить их технические и экспериментальные отчеты. Но самая сложная часть, а именно training code, знания и техники, необходимые для координации GPU при обучении этих огромных систем, по-прежнему окутана сложностью и рассеяна по несвязанным статьям и зачастую закрытым codebase.

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

По мере роста кластеров, используемых для обучения таких моделей, появились такие техники, как data parallelism, tensor parallelism, pipeline parallelism и context parallelism, а также ZeRO и kernel fusion, призванные обеспечивать максимально высокую утилизацию GPU в каждый момент времени. Это существенно сокращает время обучения и позволяет использовать дорогостоящее оборудование с максимальной эффективностью. Эти техники распределенного обучения важны не только для обучения базовых моделей, но и стали критически необходимыми для fine-tuning больших моделей на специализированных данных, где часто достигаются наилучшие результаты. В этой книге мы последовательно разберем все эти методы, от самых простых до наиболее изощренных, сохраняя единую линию повествования, чтобы вам было понятно, откуда берется каждый подход.

Мы будем исходить из того, что у вас уже есть базовое представление о современных архитектурах LLM и вы в общих чертах понимаете, как обучаются модели глубокого обучения, но при этом вполне можете быть новичком в распределенном обучении. Если понадобится освежить основы, рекомендуем отличные курсы на DeepLearning.ai или руководства по PyTorch. Эту книгу можно рассматривать как вторую часть трилогии после нашей предыдущей статьи о подготовке данных для pretraining, а именно статьи FineWeb. Прочитав обе, вы получите почти весь базовый багаж знаний, необходимый для полноценного понимания того, как сегодня создаются высокопроизводительные LLM; останется лишь недостающий секретный ингредиент, связанный со смешиванием данных и выбором архитектуры, чтобы довести рецепт до совершенства (третья часть еще впереди).

Эта книга стоит на трех китах:

1. Короткие введения в теорию и основные понятия: Прежде чем погружаться в код и эксперименты, важно на высоком уровне понять, как работает каждый метод и в чем его сильные и слабые стороны. Например, вы увидите, какие части языковой модели сильнее всего расходуют память и на каких этапах обучения это происходит. Мы также разберем, как обходить ограничения по памяти с помощью распараллеливания модели и как повышать пропускную способность, масштабируясь на большее число GPU. После этого вам будет гораздо проще разобраться в следующем виджете, который оценивает разбивку памяти Transformer-модели:

(Не переживайте, если сейчас не до конца понимаете, что именно происходит в этом виджете. Сейчас мы во всем разберемся.)

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

2. Понятные реализации в коде: Теория - это одно, но при реальной реализации неизбежно вскрываются edge cases и важные детали. Поэтому там, где возможно, мы ссылаемся на референсные реализации. В зависимости от случая мы будем использовать два code reference:

• Репозиторий Picotron создан для обучения, поэтому концепции в нем обычно реализованы в виде коротких самодостаточных файлов.

• С другой стороны, для знакомства с production-ready кодом мы будем обращаться к реализациям Nanotron. Это production codebase для обучения моделей, используемый в Hugging Face.

3. Реальные бенчмарки эффективности обучения: То, как на практике масштабировать обучение LLM, зависит от вашей инфраструктуры: типа чипов, interconnect и множества других факторов, поэтому единого универсального рецепта здесь нет. Зато мы можем показать, как бенчмаркать разные конфигурации. Именно это мы и сделали на нашем кластере. Мы провели более 4100 распределенных экспериментов (а если считать test runs, то свыше 16 тысяч) на конфигурациях вплоть до 512 GPU, чтобы просканировать широкий набор distributed training layouts и размеров моделей.

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

Обзор на высоком уровне

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

1. Использование памяти: Это жесткое ограничение: если training step не помещается в память, обучение невозможно продолжить.
2. Эффективность вычислений: Мы хотим, чтобы оборудование проводило как можно больше времени за вычислениями, поэтому нужно сокращать время на передачу данных и ожидание работы других GPU.
3. Communication overhead: Мы хотим минимизировать communication overhead, поскольку он оставляет GPU в простое. Для этого мы постараемся максимально эффективно использовать bandwidth внутри узла (быстрый) и между узлами (более медленный), а также как можно сильнее перекрывать communication вычислениями.

Во многих местах мы увидим, что один из этих ресурсов (вычисления, коммуникация, память) можно обменивать на другой, например через recomputation или tensor parallelism. Найти правильный баланс между ними - ключ к успешному масштабированию обучения.

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

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

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

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

В общем виде это выглядит так:

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

Batch size (bs) - один из важных hyperparameter обучения модели; он влияет как на сходимость модели, так и на throughput.

Небольшой batch size может быть полезен на ранних этапах обучения, поскольку позволяет быстрее двигаться по ландшафту оптимизации и находить удачную траекторию. Однако на более поздних стадиях слишком маленький batch size делает градиенты шумными, и модель может так и не выйти на лучшее достижимое качество. Слишком большой batch size, наоборот, дает очень точные оценки градиентов, но часто менее эффективно расходует каждый training token, из-за чего сходимость замедляется и вычислительные ресурсы могут тратиться впустую. Хорошее раннее обсуждение этой темы можно найти в статье OpenAI о large batch training или в разделе 4.2 technical report MiniMax-01.

Batch size также влияет на то, сколько времени требуется, чтобы обучиться на заданном текстовом датасете: маленький batch size потребует больше optimizer steps для обработки того же числа примеров. Шаги optimizer дороги по compute time, поэтому общее время обучения возрастает по сравнению со случаем более крупного batch size. При этом важно помнить, что вокруг оптимального значения batch size его часто можно довольно широко варьировать без существенного влияния на качество модели: чувствительность финального качества к точному значению batch size обычно невелика в окрестности оптимума.

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

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

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

Для современных LLM типичный рабочий диапазон batch size находится примерно между 4 и 60 миллионами токенов на batch. И batch size, и размер training corpus стабильно росли на протяжении последних лет: Llama 1 обучалась с batch size около 4M токенов на корпусе в 1,4 триллиона токенов, тогда как DeepSeek обучалась с batch size около 60M токенов на 14 триллионах токенов.

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

Начнем с разбора того, что именно вообще приводит к OOM. Это даст нам полезную интуицию относительно memory requirements при обучении модели.

Использование памяти в transformer-моделях

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

• веса модели
• градиенты модели
• состояния optimizer
• активации, необходимые для вычисления градиентов

Все эти объекты хранятся в виде tensors, которые имеют разные shapes и precision. Shapes определяются такими hyperparameters, как batch size, длина последовательности, hidden dimensions модели, число attention heads, размер словаря и, как мы увидим позже, возможный model sharding. Precision относится к форматам вроде FP32, BF16 или FP8, которые требуют соответственно 4, 2 или 1 байт на хранение одного значения tensor. Подробно разные precision и связанные с ними trade-off мы обсудим в разделе "Mixed precision training"; пока достаточно помнить, что memory requirements для этих форматов различаются, а значит, меняется и объем памяти, необходимый для хранения нужных нам объектов.

Как же быстро оценить использование памяти по всем этим переменным? Один из простых способов - подойти эмпирически и просто измерить его.

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

С помощью PyTorch profiler можно понять, как именно выделяется память на протяжении обучения. Мы видим, что использование памяти - это не статичная величина: оно заметно меняется как по ходу обучения, так и внутри одного training step:

Очевидно, что первый шаг выглядит совсем не так, как последующие, но прежде чем к этому переходить, давайте посмотрим на общую анатомию одного шага. Сначала во время forward pass быстро растут activations, затем в ходе backward pass накапливаются gradients, а по мере распространения backward pass сохраненные activations, использовавшиеся для вычисления градиентов, постепенно освобождаются. Наконец, выполняется optimization, для которой нужны все gradients, после чего обновляются optimizer states перед следующим forward pass.

Как уже сказано, первый шаг выглядит иначе: activations быстро растут, а затем на какое-то время выходят на плато. Почему? На этом первом шаге PyTorch caching allocator выполняет значительный объем подготовительной работы, заранее организуя memory allocations, чтобы на последующих шагах не приходилось искать свободные блоки памяти, что ускоряет выполнение (см. блог Zach). После первого шага также появляются optimizer states, которые в дальнейшем обычно смещают общий профиль использования памяти.

Теперь, когда у нас появилось первое представление о памяти, давайте посмотрим, почему масштабирование обучения часто сводится к максимизации compute efficiency при одновременном удержании memory requirements всех этих сущностей (activations, parameters, gradients, optimizer states) в пределах memory constraints GPU.

Память для weights/grads/optimizer states

Начнем с первых трех пунктов нашего списка: весов модели, градиентов и optimizer states. Для них на самом деле довольно легко оценить необходимый объем памяти.

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

В этом уравнении h - hidden dimension, v - размер словаря, а L - число слоев модели. Обратите внимание, что при больших hidden dimensions доминировать начинает член h^2, поскольку только он растет квадратично по мере масштабирования параметров.

Memory requirements для параметров и градиентов определяются просто: нужно умножить число параметров на число байт на параметр. В классическом full precision (FP32) training и параметры, и градиенты требуют по 4 байта, а optimizer, если это Adam, дополнительно хранит momentum и variance, что добавляет еще 8 байт на параметр (по 4 байта на каждое состояние). Итого:

Теперь посмотрим, что меняется при использовании более низкой precision. По соображениям стабильности (см. позже раздел о mixed precision training) на практике мы часто используем не полностью low precision training, а смесь высокой и низкой precision, называемую "mixed precision". Сегодня стандартный вариант mixed precision training обычно использует BF16 для большинства вычислений, что требует 2 байта на параметр и градиент, а также хранит дополнительную копию весов модели и градиентов в FP32, что в сумме дает 12 байт на параметр. Помимо параметров и градиентов, нужно хранить optimizer states; для Adam это momentum и variance, которые ради численной стабильности обычно хранятся в FP32, то есть по 4 байта каждое.

Итоговая сводка такова:

Интересно, что mixed precision training само по себе не экономит память; оно лишь перераспределяет ее между тремя компонентами и, если мы накапливаем градиенты в FP32, фактически добавляет еще 4 байта по сравнению с full precision training. Тем не менее это выгодно, поскольку вычисление forward/backward pass в half precision, во-первых, позволяет использовать на GPU оптимизированные low precision operations, а значит, ускоряет вычисления, а во-вторых, снижает memory requirements для activations во время forward pass, которые, как мы видели на графике выше, составляют значительную долю общего потребления памяти.

Давайте в целом оценим, сколько памяти требуется для модели (если рассматривать полную картину, full precision и mixed precision здесь дают одинаковый порядок величин):

Как видно, уже на масштабе 7B(!) memory requirements для весов, градиентов и optimizer states становятся настолько велики, что превосходят объем памяти типичного GPU, например 80 GB у H100.

Но пока ограничимся моделями, которые помещаются в один GPU, и посмотрим на последний крупный вклад в memory budget: activations.

Память для activations

Memory requirements для activations считать несколько сложнее, чем для weights, gradients и optimizer states, в том числе потому, что они зависят от входов модели. Если вы не до конца помните, зачем вообще нужно хранить activations для backward pass, эта статья дает хорошее быстрое напоминание. После аккуратного разбора того, как именно вычисляется backward pass, можно оценить общий объем памяти, необходимый для activations в mixed precision. Получается следующее уравнение:

Здесь L - число слоев, seq - длина последовательности, bs - batch size в примерах, h - hidden dimension модели, а n_{heads} - число heads.

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

Интересно, что использование памяти для заданной модели не является статичной величиной: оно растет линейно по batch size и квадратично по длине последовательности. Это означает, что именно activation memory взрывается первой, когда мы увеличиваем batch size или обучаемся на более длинных последовательностях. Используя это уравнение, можно посмотреть, как меняется использование памяти при разных длинах последовательности, например для моделей Llama (bs=1):

Эти графики показывают довольно наглядную картину: для коротких последовательностей, или маленьких batch size, память на activations почти несущественна, но примерно начиная с 2-4k токенов она начинает занимать заметную долю памяти, тогда как использование памяти под parameters, gradients и optimizer states, как мы позже обсудим, в целом почти не зависит ни от длины последовательности, ни от batch size.

При большом числе входных токенов, то есть при больших batch size или длинных последовательностях, activations становятся безусловно крупнейшим источником memory pressure.

Можно ли укротить этот «взрыв activations»? Отличный вопрос.

Пора перейти к нашей первой технике - activation recomputation, которая помогает ограничить activation memory footprint. Сегодня это один из базовых инструментов в toolbox обучения больших моделей.

Activation recomputation

Общая идея activation recomputation, также известной как gradient checkpointing или rematerialization, состоит в том, чтобы выбрасывать часть activations во время forward pass ради экономии памяти, а затем тратить дополнительный compute на их повторное вычисление на лету во время backward pass. Без recomputation мы сохраняем каждое hidden state между двумя обучаемыми операциями, например feedforward, LayerNorm и так далее, чтобы использовать их в backward pass при вычислении gradients. При recomputation обычно сохраняются activations только в нескольких ключевых точках архитектуры, а остальные выбрасываются и затем пересчитываются во время backward pass из ближайших сохраненных activations. Иначе говоря, мы повторно выполняем часть forward pass, обменивая память на вычисления. В общем виде это выглядит так:

Существует несколько стратегий выбора ключевых activations, которые стоит сохранять:

• Full: Мы checkpoint'им activations на границе каждого слоя Transformer-модели. Обычно это называют стратегией «full», поскольку она требует фактически повторного forward pass через каждый слой, то есть почти полного дополнительного forward pass во время backward pass. Эта стратегия дает максимальную экономию памяти, но и обходится дороже всего по compute. На практике она обычно увеличивает compute cost и время на 30-40%, что весьма заметно.
• Selective: В общем случае можно сделать лучше, чем full. Авторы статьи о recomputation подробно проанализировали, какие activations растут сильнее всего и при этом имеют наименьшую стоимость повторного вычисления в терминах floating-point operations per second (FLOPS). Оказалось, что в эту категорию попадают attention-computations, поэтому их обычно можно выбрасывать, а checkpointing сосредоточить на более дорогих feedforward-вычислениях. Для модели GPT-3 (175B) это дает снижение activation memory на 70% при compute cost всего 2.7%.

Давайте посмотрим, насколько сильно стратегии recomputation уменьшают memory footprint на практике и как selective recomputation находит хороший баланс между экономией памяти и стоимостью повторного вычисления:

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

Сегодня большинство training frameworks используют FlashAttention (мы обсудим ее подробнее позже в книге), которая изначально встраивает activation recomputation в свою стратегию оптимизации, повторно вычисляя attention scores и matrices во время backward pass вместо их сохранения. Поэтому большинство пользователей FlashAttention уже фактически используют selective recomputation.

Как вы уже поняли, activation recomputation немного увеличивает число FLOPs из-за повторных вычислений, но при этом существенно снижает overhead, связанный с обращением к памяти.

Такой trade-off особенно выгоден на hardware с ограниченным объемом быстрой памяти, например на GPU, поскольку доступ к памяти обычно медленнее, чем выполнение вычислений. Поэтому, несмотря на дополнительные операции, итоговый эффект часто оказывается не только в меньшем memory footprint, но и в более быстром вычислении.

Теперь, когда мы разобрались с recomputation, мы уже можем держать под контролем activation memory usage, который видели на предыдущих графиках.

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

Gradient accumulation

Gradient accumulation - это очень прямолинейный способ избежать взрыва памяти: мы просто разбиваем batch на micro-batches. Затем последовательно выполняем forward и backward pass для каждого micro-batch, вычисляем gradients и, как и следует из названия, суммируем gradients по всем micro-batches перед шагом optimization. На практике optimization выполняется не по сумме, а по среднему значению gradients, чтобы результат не зависел от числа шагов gradient accumulation.

Будем называть batch size для каждого отдельного forward pass micro-batch size (mbs). Общий batch size между двумя соседними optimizer step будем называть global batch size (gbs). Если мы делаем один optimizer step на каждые восемь forward/backward pass, то global batch size будет в восемь раз больше micro-batch size.

То, что теперь называется global batch size, соответствует тому, что до этого ради простоты мы называли просто batch size; сейчас мы лишь уточняем терминологию, чтобы избежать неоднозначности.

При gradient accumulation global batch size вычисляется так:

Gradient accumulation позволяет нам фактически увеличивать batch size почти до бесконечности, при этом memory footprint остается постоянным. Кроме того, gradient accumulation хорошо сочетается с activation recomputation, что дает дополнительное снижение потребления памяти.

Gradient accumulation позволяет сократить activation memory, которая растет линейно с batch size, обрабатывая меньшие micro-batches последовательно. Это уменьшает объем одновременно хранимых activations и gradients, поскольку в памяти в каждый момент времени нужно удерживать только activations одного micro-batch, а значит, снижается общий activation memory footprint.

Но есть и минус: gradient accumulation требует нескольких последовательных forward/backward pass на один шаг optimization, что увеличивает compute overhead и замедляет training. Бесплатных обедов не бывает.

Если вы внимательно следили за рассуждением, то, вероятно, уже заметили, что forward/backward pass для разных micro-batches на самом деле можно выполнять параллельно. Они независимы друг от друга и отличаются только входными примерами. Похоже, пора расширять обучение за пределы одного GPU.

Но прежде чем это делать, давайте быстро посмотрим, как визуализировать compute и communication. Для этого устроим короткую экскурсию по одному из самых полезных инструментов в арсенале distributed training: profiler. Он крайне полезен для того, чтобы понимать и проверять, как именно взаимодействуют GPU-коммуникации и вычисления и где возникают bottlenecks.

Профилирование GPU compute и communication

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

Это создаст trace, который можно открыть в TensorBoard или в trace viewer от Chrome. В нем видно:

• как CPU thread асинхронно запускает kernels на GPU
• как несколько CUDA streams параллельно обрабатывают compute и communication
• времена выполнения kernels и allocation памяти

Такой trace помогает выявлять bottlenecks вроде:

• последовательных compute и communication, которые можно было бы перекрыть
• простоя GPU в ожидании передачи данных
• CUDA Sync и перемещений памяти между CPU и GPU
• накладных расходов на запуск kernels на GPU

Понимание этих паттернов критично для оптимизации производительности distributed training. Например, в trace сразу видно, корректно ли перекрывается синхронизация gradients с backward-вычислениями, о чем мы еще поговорим ниже.

Теперь давайте возьмем рабочую станцию покрупнее, с несколькими GPU, и перейдем к нашей первой технике масштабирования: data parallelism. Как мы увидим, по сути это параллельная версия gradient accumulation.

Data Parallelism

Идея data parallelism, или DP, состоит в том, чтобы скопировать модель на несколько GPU, эти копии мы будем называть "экземплярами модели", и параллельно запускать forward и backward pass на разных micro-batches данных на каждом GPU. Отсюда и название data parallelism. Скорее всего, вы уже встречали этот подход в простых примерах обучения, но в этом разделе мы разберем его гораздо глубже, так что не переключайтесь даже если общая идея вам уже знакома.

Если на каждом GPU используется свой micro-batch, gradients на разных GPU тоже будут разными. Чтобы экземпляры модели оставались синхронизированными, нужно усреднить gradients между всеми экземплярами с помощью операции "all-reduce". Эта операция выполняется во время backward pass, до шага optimizer.

Здесь мы впервые встречаемся с примитивом all-reduce, который отвечает за синхронизацию и communication между экземплярами модели и nodes.

Наивная реализация DP просто дождалась бы конца backward pass, чтобы собрать все gradients, а затем запустила бы all-reduce по всем DP-ranks для их синхронизации. Но такая последовательность "сначала compute, потом communication" - плохая идея, потому что в этот момент GPU будут простаивать, ожидая завершения communication, как показано на схеме выше.

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

Давайте рассмотрим три оптимизации, которые позволяют уйти далеко вперед по сравнению с этой наивной реализацией.

Первая оптимизация: overlap синхронизации gradients с backward pass

Главный недостаток наивного DP-подхода в том, что после backward pass, то есть compute, приходится ждать синхронизации gradients, то есть communication, и только потом обновлять параметры. Можно ли перекрыть эту communication нашими вычислениями? Да, можно.

Как видно на схеме выше, gradients для конкретного слоя, розовые блоки, можно собирать и суммировать еще до того, как будут вычислены gradients для более ранних слоев, то есть блоков слева. Например, как только завершился backward pass последнего слоя, самого правого блока, его gradients уже можно отправлять в all-reduce, пока backward-вычисления продолжаются для более ранних слоев.

В PyTorch этого можно добиться, если прикрепить к каждому параметру all-reduce hook. Тогда all-reduce будет запускаться сразу, как только gradient для конкретного параметра готов, пока gradients для остальных параметров еще продолжают вычисляться. В результате большая часть all-reduce операций перекрывается вычислением gradients, что повышает эффективность. Вот простая функция для регистрации такого hook:

Такое перекрытие compute и communication сокращает время ожидания синхронизации gradients по всей модели. Сама синхронизация начинает происходить, по крайней мере частично, параллельно backward pass в рамках того же training step, что заметно ускоряет data parallelism. Ниже показана полная реализация наивного DP с overlap синхронизации:

Это первый пример overlap между compute и communication, то есть одной из ключевых техник, без которой нельзя добиться высокой эффективности масштабирования. В этой книге мы еще не раз к ней вернемся. Но и это еще не предел.

Вторая оптимизация: bucketing gradients

Операции на GPU обычно работают эффективнее, когда выполняются на больших tensors, а не на множестве мелких. Это верно и для communication. Поэтому gradients выгодно объединять в "buckets" и запускать один all-reduce сразу для всех gradients внутри bucket, вместо того чтобы делать отдельный all-reduce для каждого gradient. Общая идея выглядит так:

Это похоже на упаковку вещей в коробки перед отправкой: отправить несколько крупных коробок выгоднее, чем множество мелких. Один all-reduce на bucket заметно снижает communication overhead и ускоряет обмен данными.

Ниже приведена реализация с bucketing:

Третья оптимизация: взаимодействие с gradient accumulation

Наконец, как мы уже видели, gradient accumulation работает так: мы выполняем несколько forward и backward pass, а параметры обновляем только после этого через optimizer.step(). Когда gradient accumulation сочетается с data parallelism, нужно внимательно выбрать момент синхронизации gradients.

В наивной версии all-reduce автоматически срабатывает после каждого backward pass во время accumulation. Это не лучший вариант: достаточно было бы сделать один reduce после последнего шага и тем самым получить тот же эффект с меньшими накладными расходами.

В PyTorch это обычно решается с помощью model.no_sync(), который отключает синхронизацию gradients на тех backward pass, где reduce не нужен.

Теперь посмотрим, что все это означает для global batch size.

Возвращаемся к global batch size

Теперь можно обновить формулу batch size, добавив в нее параметры data parallelism и gradient accumulation:

Здесь grad\_acc - это число шагов gradient accumulation, а dp - число параллельных экземпляров, участвующих в data parallelism.

Значит, при фиксированном целевом global batch size мы можем обменивать шаги gradient accumulation на процессы data parallelism, чтобы ускорять обучение.

На практике обычно стараются максимально наращивать размер data parallelism, то есть dp, а уже потом добавлять gradient accumulation, то есть grad\_acc, потому что data parallelism по своей природе параллелен, в отличие от последовательной gradient accumulation. Gradient accumulation накладывают сверху, когда одного масштабирования по data parallelism уже недостаточно для достижения целевого global batch size или просто заканчиваются GPU.

Возможность распределять обучение по разным примерам дает нам первое измерение параллелизации. Поэтому это и есть 1D parallelism; дальше мы постепенно добавим еще четыре измерения.

Где мы сейчас

Коротко подытожим, как собрать наш первый 1D training setup в виде рабочего чернового рецепта для data parallelism:

1. Сначала определяем лучший global batch size в токенах: либо по литературе, либо по собственным экспериментам на сходимость модели.
2. Затем выбираем длину последовательности для обучения, снова либо опираясь на литературу, либо по экспериментам. В общем случае 2-8k tokens сегодня обычно надежно работают на типичных evaluation benchmarks. В training recipes мы здесь не уходим, но команды часто увеличивают длину последовательности ближе к концу обучения, добавляя в смесь данные с более длинным контекстом.
3. После этого мы знаем batch size, то есть gbs. Далее находим максимальный local batch size, то есть mbs, который помещается в память одного GPU.
4. Наконец, определяем число доступных GPU для нашего целевого dp. Отношение gbs к dp подскажет, сколько шагов gradient accumulation еще потребуется для достижения нужного gbs.

Если отношение gradient accumulation оказывается меньше 1, то есть GPU у нас больше, чем необходимо, можно либо не использовать все GPU, либо попробовать больший gbs, либо проверить, ускорит ли обучение меньший mbs. В последнем случае мы сознательно ставим throughput выше индивидуальной compute-эффективности каждого GPU и используем меньший mbs, чем в принципе могли бы.

Теперь рассмотрим конкретный пример. Пусть мы хотим обучить современную модель с gbs в 4M tokens и длиной последовательности 4k. Тогда наш batch size составит 1024 примера, берем ближайшую степень двойки. Допустим, на одном GPU в память помещается только mbs=2, а для обучения у нас есть 128 GPU. Это значит, что при 4 шагах gradient accumulation мы получим нужные 1024 примера, то есть 4M tokens, на один training step. А что если внезапно становится доступно 512 GPU? Тогда тот же gbs можно получить, сохранив mbs=2 и уменьшив число шагов gradient accumulation до 1, а значит обучение пойдет быстрее.

Хотя data parallelism отлично перекрывает all-reduce синхронизацию gradients с backward-вычислениями и тем самым экономит время, на больших масштабах это преимущество начинает разрушаться. Почему? Потому что по мере роста числа GPU, до сотен и тысяч, резко растут накладные расходы на координацию, а требования к сети начинают перевешивать выигрыш. В результате каждая дополнительная GPU делает setup все менее и менее эффективным.

На практике это хорошо видно по benchmark-ам:

Как видно на графике, после определенного порога throughput начинает довольно заметно падать, тогда как использование памяти на один GPU остается постоянным и не меняется при добавлении новых DP-ranks.

Data parallelism был нашей первой и самой простой стратегией масштабирования обучения на большее число GPU. Этот подход похож на gradient accumulation, но в отличие от него распараллеливает forward и backward pass по micro-batches, тем самым повышая throughput.

Но внимательный читатель уже, вероятно, заметил важное допущение: все это работает только если хотя бы один forward pass для одного входного примера, то есть mbs=1, помещается в память GPU. А это бывает далеко не всегда. Как мы увидим ниже, крупные модели часто не влезают в один GPU даже при включенной activation recomputation:

Мы уже видели, что при достаточно большом масштабе data parallelism начинает упираться в накладные расходы на обмен данными. Есть ли другие варианты для более крупных моделей или больших batch size? К счастью, да: можно либо выносить часть tensors на CPU, либо разносить tensors с weights, gradients и optimizer states между несколькими GPU.

Есть два основных подхода к такому разбиению: parallelism, например tensor, context или pipeline parallelism, и sharding, например DeepSpeed ZeRO или PyTorch FSDP. Эти подходы в значительной степени ортогональны и на практике могут комбинироваться.

Парадигма sharding тесно связана с DP, поэтому сначала мы разберем именно ее, на примере метода ZeRO.

Zero Redundancy Optimizer (ZeRO)

В этом разделе мы познакомимся с DeepSpeed ZeRO - технологией оптимизации памяти, предназначенной для уменьшения избыточности памяти при обучении LLM.

Хотя data parallelism хорошо масштабирует обучение, наивное дублирование optimizer states, gradients и parameters на каждом DP-ранге создает заметную избыточность по памяти. ZeRO устраняет ее, разбивая optimizer states, gradients и parameters вдоль измерения data parallelism, но при этом все еще позволяя считать с полным набором параметров. Иногда это требует дополнительных communications между DP-рангами, которые, как мы вскоре увидим, могут как полностью перекрываться вычислениями, так и не перекрываться.

У этого подхода есть три возможные стадии оптимизации:

• ZeRO-1: разбиение optimizer states
• ZeRO-2: разбиение optimizer states и gradients
• ZeRO-3: разбиение optimizer states, gradients и parameters

Вы, возможно, заметили, что activations нет в списке сущностей, которые можно shard'ить. Причина в том, что каждая DP-реплика модели получает свой micro-batch, а значит, activations на каждом DP-ранге различаются. Они не дублируются между рангами, и потому shard'ить тут просто нечего.

Давайте посмотрим, какую экономию памяти дает каждая стадия ZeRO.

Снова о потреблении памяти

Ранее мы уже обсуждали, сколько памяти во время обычного обучения занимают optimizer states, gradients и parameters. Обозначим число параметров модели через \Psi - раньше мы использовали N, но здесь перейдем к обозначению из оригинальной статьи про ZeRO. В mixed precision training с optimizer Adam, подробнее о ней мы поговорим позже, затраты памяти на каждую из сущностей, которые нужно хранить, таковы:

• Параметры модели в half precision, то есть BF16/FP16: 2\Psi
• Градиенты модели в half precision, то есть BF16/FP16: 2\Psi
• Параметры модели в FP32 и optimizer states: 4\Psi + (4\Psi + 4\Psi)
• Градиенты модели в FP32: 4\Psi, опционально, только если мы хотим накапливать gradients в FP32

Если не накапливать gradients в FP32, получаем суммарное потребление памяти 2\Psi + 2\Psi + 12\Psi; если накапливать, то 2\Psi + 6\Psi + 12\Psi. Для простоты пока сосредоточимся на варианте без FP32 gradient accumulation.

Идея ZeRO в том, чтобы shard'ить эти объекты между DP-ranks: каждый node хранит только свою долю, а при необходимости эти доли собираются обратно. В итоге использование памяти делится на степень data parallelism N_d:

Здесь \Psi обозначает число параметров, k - memory multiplier для optimizer states, для Adam это k=12, как мы только что видели, а N_d - степень DP.

Теперь разберем, как работает каждая стадия ZeRO. Начнем с ZeRO-1.

ZeRO-1: Partitioning optimizer states

В vanilla DP все ranks после backward pass собирают одни и те же gradients и одновременно выполняют одинаковые optimizer steps. Выглядит как большой объем дублирующейся работы. Можно ли этого избежать и одновременно уменьшить использование памяти?

В ZeRO-1 optimizer states разбиваются на N_d равных частей, где N_d - степень DP. Это означает, что каждая реплика модели, размещенная на DP ranks, отслеживает только \frac{1}{N_d} optimizer states, а во время optimization step обновляется лишь \frac{1}{N_d} FP32-весов.

Однако во время forward pass каждой реплике нужны все параметры целиком. Поэтому после optimizer step приходится добавлять дополнительную операцию all-gather, это уже второй встречающийся нам тип collective communication primitive, чтобы каждая реплика модели снова получила полный набор обновленных весов.

Этим и объясняется формула памяти 2\Psi + 2\Psi + \frac{k\Psi}{N_d}, которую мы видели на предыдущем рисунке. Ниже кратко описана последовательность операций для одного training step:

1. Выполнить forward pass с одним и тем же полным набором BF16-параметров на каждой реплике, но на разных micro-batches.
2. Выполнить backward pass с тем же полным набором gradients на каждой реплике, снова на разных micro-batches.
3. Выполнить reduce-scatter для gradients, это еще один primitive, который мы скоро поясним.
4. Каждая реплика выполняет optimizer step на своих локальных optimizer states, то есть только на \frac{1}{N_d} от полного объема, и получает \frac{1}{N_d} обновленных FP32-параметров, которые затем можно преобразовать в \frac{1}{N_d} полного набора BF16-параметров.
5. Выполнить all-gather для BF16-параметров, чтобы вернуть недостающие срезы каждой реплике. Это новая операция в ZeRO, которой нет в vanilla DP.

Вы можете спросить, что именно представляет собой операция "reduce-scatter" и как все это выглядит на практике, поэтому давайте сделаем объяснение более наглядным с помощью схемы ниже. Мы пройдем по всем шагам цикла forward/backward pass:

С точки зрения практических communications, по сравнению с vanilla DP, ZeRO-1 заменяет all-reduce для gradients на reduce-scatter и добавляет all-gather по всем параметрам после optimizer step. Выглядит это так:

Если вы следили за предыдущим разбором vanilla DP, то помните, что communication all-reduce для gradients можно перекрывать с вычислениями backward pass. В ZeRO-1 также можно изучать, как эффективно перекрывать и вновь добавленный all-gather BF16-параметров. Здесь есть две основные стратегии:

• Во время optimizer step: можно запускать all-gather сразу после того, как optimizer обновил первый срез параметров. Это потенциально позволяет перекрыть communication с обновлением остальных параметров.
• Во время forward pass: можно перекрывать all-gather параметров каждого слоя с самим forward pass.

В ZeRO-1 optimizer states уже partitioned, а значит каждая реплика обновляет только \frac{1}{N_d} состояний. Внимательный читатель мог заметить, что и хранить все gradients на всех DP ranks вовсе не обязательно, ведь для optimization step нужен только их поднабор. Здесь и появляется ZeRO-2.

ZeRO-2: Adding gradient partitioning

Поскольку на каждой реплике нужен только тот gradient shard, который соответствует ее shard optimizer states, имеет смысл shard'ить и gradients по той же схеме. Тогда во время backward pass вместо all-reduce по gradients мы выполняем только reduce-scatter. В памяти теперь хранится лишь \frac{1}{N_d} gradients, которые действительно нужны, и это дает дополнительную экономию памяти по сравнению с ZeRO-1.

Теперь легко увидеть, что sharding gradients приводит нас к формуле 2\Psi + \frac{2\Psi+k\Psi}{N_d}, и при росте N_d мы можем использовать до 8 раз меньше памяти по сравнению с baseline. С точки зрения communication процесс почти такой же, как в ZeRO-1: разница лишь в том, что теперь данные передаются и память освобождается on the fly. В обоих случаях нужны reduce-scatter для gradients и all-gather по всем параметрам. Поэтому по своим communication-свойствам ZeRO-2 эквивалентен vanilla DP training.

Теперь, когда мы shard'им и gradients, закончили ли мы на этом, или можно пойти еще дальше? Встречайте ZeRO-3.

ZeRO-3: добавляем parameter partitioning (FSDP)

На третьей стадии мы расширяем описанный выше подход: теперь между DP-репликами shard'ятся не только optimizer states и gradients, но и сами параметры модели.

Как же на практике выполнять forward или backward pass, если параметры модели распределены? Довольно просто: мы собираем их по требованию, когда они нужны. Во время forward pass это выглядит так:

По мере выполнения forward pass и последовательного прохода по слоям мы по запросу поднимаем нужные параметры и сразу же выбрасываем их из памяти, как только они перестают быть нужны. Backward pass работает по тому же принципу, только в обратном направлении. Здесь мы получаем gradient shards:

Другая проблема состоит в том, что такие all-gather приходится выполнять непрерывно на протяжении всего forward и backward pass внутри training step, что дает 2\cdot \text{num\_layers} -1 дополнительных all-gather по сравнению с ZeRO-2. Каждая из них вносит небольшой overhead по base latency, что хорошо видно на следующем рисунке:

Во время forward pass мы выполняем all-gather параметров по мере необходимости, и это дает communication tax размером \Psi. Поскольку после использования в forward pass параметры сразу выбрасываются, в backward pass нужен еще один all-gather, то есть еще один communication tax размером \Psi. Наконец, для gradients нужен тот же reduce-scatter, что и в ZeRO-2, а это еще \Psi communication cost. Итого получаем суммарную communication cost в 3\Psi против 2\Psi у ZeRO-2.

На слух это может казаться большим communication overhead, но на практике не все так страшно, потому что communication параметров следующего слоя можно перекрывать с forward pass текущего слоя с помощью того, что называется prefetching. При prefetching мы выполняем all-gather весов для Layer n+1, пока идет forward pass для Layer n, и аналогично собираем веса для Layer n-1, пока выполняется backward pass для Layer n. Разумеется, такое перекрытие работает только пока мы не слишком сильно увеличиваем DP; грубое практическое правило состоит в том, что DP не стоит поднимать выше 512.

С точки зрения памяти наше уравнение теперь принимает окончательную форму \frac{2\Psi +2\Psi+k\Psi}{N_d}, а это означает, что по крайней мере для model-related parameters мы теоретически можем снижать memory usage почти неограниченно, если продолжаем увеличивать DP size. Но важно помнить, что это никак не помогает с intermediate activations; для них, как мы уже видели, нужны activation checkpointing и gradient accumulation.

Подведем промежуточный итог по DP и ZeRO. Мы увидели, что DP позволяет заметно повышать throughput обучения, просто масштабируя training добавлением новых реплик модели. А ZeRO позволяет обучать даже те модели, которые в обычном случае не поместились бы в один GPU, shard'я parameters, gradients и optimizer states между DP replicas и расплачиваясь за это сравнительно небольшим communication cost.

Однако и здесь есть пределы: DP работает только если хотя бы слой модели помещается в один GPU, а ZeRO умеет partition'ить только parameters, gradients и optimizer states, но не activation memory. Напомним из раздела про activation memory, что эта часть памяти растет с длиной последовательности и batch size. Формально можно было бы просто ограничивать их, но на практике никто не хочет, чтобы hardware заставляло обучаться только на коротких последовательностях.

Чтобы преодолеть это ограничение, пора перейти к новой, ортогональной оси parallelism - tensor parallelism (TP). В отличие от ZeRO-3, который опирается на активный обмен параметрами, TP предлагает shard'ить parameters, gradients, optimizer states и activations между устройствами без необходимости постоянно перегонять параметры модели между GPU.

Как? Как это вообще возможно? Давайте разберем этот на первый взгляд почти магический подход.

Tensor Parallelism

Итак, с помощью ZeRO мы уже научились shard'ить parameters, gradients и optimizer states модели, но упираемся в предел, когда activation memory начинает съедать весь бюджет памяти. Здесь и появляется tensor parallelism (TP), метод, который shard'ит веса, gradients, optimizer states, а также activations - и при этом не требует заранее собирать их все перед вычислением. Звучит почти как мечта. Для начала посмотрим, как TP работает на простых операциях matrix multiplication, или matmul.

Tensor parallelism использует математические свойства matrix multiplication, A \times B. Чтобы понять, как именно это работает, давайте рассмотрим два фундаментальных уравнения, которые и делают такую параллелизацию возможной:

Это означает, что произведение матриц можно вычислять либо по отдельности умножая каждый столбец B, либо по отдельности умножая строки и затем объединяя результат. В нейронной сети matrix multiplication чаще записывается в виде X \times W, где:

• X обозначает вход или activation values
• W обозначает веса Linear layer

На практике небольшой пример такой операции выглядит так:

Посмотрим, как можно распараллелить эту операцию. В tensor parallelism tensors разбиваются на N shards вдоль выбранного измерения и распределяются по N GPU. Матрицы можно делить либо по столбцам, либо по строкам, что приводит соответственно к row или column parallelism. Как мы увидим дальше, row и column sharding требуют разных communication primitives.

Первый вариант - использовать column-wise, также называемый column-linear, sharding: мы копируем полные входные матрицы на каждый worker, для этого нужна операция broadcast, а матрицу весов делим по столбцам. Затем входы умножаются на частичные матрицы весов, а в конце результаты объединяются через all-gather.

Ниже показана реализация column-wise tensor parallelism в коде:

Второй вариант называется row-wise, или row-linear, sharding. Как нетрудно догадаться, row-linear означает, что матрица весов делится на куски по строкам. Но тогда приходится делить и входы, а значит вместо broadcast, который использовался при column-linear sharding, нужна операция scatter, это уже четвертый distributed communication primitive. Результаты на каждом worker уже имеют правильную форму, но для получения финального результата их нужно просуммировать, поэтому в этом случае также требуется all-reduce:

Ниже показана реализация row-wise tensor parallelism:

Теперь, когда у нас есть базовые строительные блоки TP, посмотрим, как эффективно комбинировать их внутри transformer layer.

Tensor parallelism в transformer block

Чтобы перейти от игрушечного примера к реальной стратегии, давайте возьмем настоящий базовый блок модели. Transformer-модель состоит из двух основных блоков: feedforward multi-layer perceptron (MLP) и multi-head attention (MHA). Tensor parallelism можно применять к обоим.

Feedforward-часть удобно распараллеливать, комбинируя сначала column-linear, а затем row-linear разбиение. В такой схеме в forward pass это соответствует broadcast для копирования входа и all-reduce в конце. При реальном обучении broadcast обычно не нужен, потому что можно заранее гарантировать, что входы уже синхронизированы между TP-рангами. Такая схема эффективнее, чем вариант с row-linear, за которым следует column-linear, поскольку позволяет избежать промежуточного all-reduce между двумя этапами разбиения.

Теперь, когда мы нашли эффективную схему для feedforward-части transformer-блока, посмотрим на multi-head attention.

В случае attention можно в целом следовать той же логике: матрицы Query (Q), Key (K) и Value (V) делятся column-parallel способом, а выходную projection удобно трактовать как row-linear. Для multi-head attention такой column-parallel подход особенно естественен: каждый GPU считает attention для одного attention head или для подмножества heads. Та же идея работает и для multi-query attention (MQA), и для grouped query attention (GQA), где keys и values разделяются между разными queries.

Tensor parallelism так хорошо ложится и на Attention, и на MLP потому, что в обоих случаях есть естественно независимые измерения. Attention-блок можно распараллеливать по измерению num_attention_heads, поскольку каждый attention head работает независимо. Аналогично, MLP-блок можно распараллеливать по измерению hidden_dim, так как операции внутри feedforward-сети независимы вдоль этой оси.

Однако важно помнить, что степень tensor parallelism не должна превышать число attention heads, поскольку QKV projection мы shard'им именно по измерению num_attention_heads. В случае Grouped Query Attention (GQA) у нас есть num\_attention\_heads query heads, но только num\_kv\_heads key/value heads, причем num\_attention\_heads >= num\_kv\_heads. В этом случае все еще можно взять TP = num\_attention\_heads, но потребуется аккуратно следить за тем, чтобы K/V heads корректно синхронизировались между GPU. Например, у Llama-3 8B есть 32 query heads, но только 8 key/value heads. Теоретически степень TP можно было бы поднять до 32, но для этого потребовалась бы особенно аккуратная реализация синхронизации K/V heads между tensor-parallel workers.

Важно понимать и то, что tensor parallelism не является серебряной пулей. Мы добавили несколько distributed communication primitives прямо в вычислительный путь модели, а значит, их трудно полностью скрыть за computation, как это иногда удается в ZeRO. Поэтому итоговая производительность всегда оказывается компромиссом между выигрышем по вычислениям и памяти с одной стороны и дополнительным communication overhead с другой. Посмотрим на это наглядно:

Если посмотреть на timeline операций в tensor-parallel MLP, та же логика применима и к MHA, можно лучше понять возникающие компромиссы. В forward pass каждого decoder layer мы упираемся в точку синхронизации, связанную с all-reduce, и ее уже нельзя перекрыть computation. Этот явный communication overhead необходим, чтобы объединить частичные результаты с разных tensor-parallel рангов до того, как можно будет применить финальный LayerNorm.

Tensor parallelism действительно уменьшает activation memory для matrix multiplications, поскольку промежуточные activations shard'ятся между GPU. Однако для операций вроде LayerNorm нам все равно приходится собирать полные activations, а значит, мы не получаем максимальной возможной экономии памяти. Кроме того, TP вносит существенные требования к communication, и они сильно зависят от сетевой инфраструктуры. Неспособность полностью скрыть этот конкретный all-reduce за вычислениями означает, что он напрямую попадает в critical path forward propagation, то есть в последовательность операций, определяющую минимальное время завершения forward pass.

Давайте чуть внимательнее посмотрим на этот компромисс по мере роста степени TP:

Рост TP снижает throughput на один GPU, график слева, но позволяет обрабатывать более крупные batch size, график справа. Это наглядно показывает компромисс между вычислительной эффективностью и доступной памятью в distributed training.

На практике, как видно по левому графику выше, communication overhead у tensor parallelism особенно заметен при масштабировании выше 8 GPU. Пока TP укладывается в пределы одного узла, можно опираться на быстрые interconnect вроде NVLink. Но при выходе между узлами приходится использовать более медленные сетевые соединения. Мы видим заметное падение при переходе от TP=8 к TP=16 и еще более резкий спад при переходе от TP=16 к TP=32. При больших степенях parallelism communication overhead быстро становится доминирующей частью общего времени.

Тем не менее tensor parallelism дает важный выигрыш по памяти, распределяя между GPU параметры модели, gradients, optimizer states и, в определенной степени, activations. Посмотрим, как это выглядит на модели с 70B параметров:

Увеличение tensor parallelism уменьшает объем памяти, который на каждом GPU нужен под параметры модели, gradients и optimizer states, настолько, что крупную модель уже можно уместить в пределах одного узла из 8 GPU.

Можно ли выжать из этой техники еще больше? Layer normalization и dropout все еще требуют собирать полные activations на каждом GPU, а это частично съедает выигрыш по памяти. Значит, можно попробовать распараллелить и эти оставшиеся операции.

Именно этим мы и займемся дальше, рассмотрев небольшое, но очень естественное расширение tensor parallelism под названием sequence parallelism.

Sequence parallelism

Sequence parallelism, или SP, состоит в том, чтобы разделять activations и вычисления для тех частей модели, которые не покрываются tensor parallelism, например dropout и LayerNorm, но делать это вдоль измерения входной последовательности, а не hidden dimension.

Это нужно потому, что для корректной работы этим операциям нужен доступ ко всей hidden dimension. Например, LayerNorm требуется полная hidden dimension, чтобы вычислить mean и variance:

где \mu = \text{mean}(x) и \sigma^2 = \text{var}(x) вычисляются по hidden dimension h.

Из-за этого, хотя сами операции вычислительно дешевы, они все равно требуют заметного объема activation memory. Sequence parallelism позволяет распределить эту нагрузку по памяти между GPU, разбивая данные вдоль измерения sequence.

Следующая схема показывает, как мы переходим между tensor-parallel и sequence-parallel регионами с помощью разных collective operations, обозначенных как f и g. На практике мы движемся слева направо:

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

В tensor parallelism во время forward pass происходит следующее:

• f - это no-op, потому что activations уже дублируются между ranks.
• f* - это all-reduce, который синхронизирует activations и обеспечивает корректность.

А в backward pass:

• f* - это no-op, поскольку gradients уже дублируются между ranks.
• f - это all-reduce для синхронизации gradients.

Операции f и f* называются conjugate pairs, потому что они дополняют друг друга: в каждом проходе, когда одна из них является no-op, другая выполняет all-reduce, а в следующем проходе роли меняются местами.

Для sequence parallelism используются другие операции, обозначенные как g и g*. В частности, внутри SP-регионов мы избегаем all-reduce, потому что для этого пришлось бы собирать activations целиком, а значит peak memory usage вырос бы и убил сам смысл SP.

Так что же именно здесь происходит? Как сказала бы известная LLM, давайте разберем по шагам:

Начальный слой LayerNorm (SP-регион)

• На вход поступают tensors X1* и X2* размера (b,s/2,h), уже разделенные по измерению sequence.
• Каждый GPU независимо вычисляет LayerNorm на своем куске последовательности, получая Y1* и Y2*.

Первый переход (SP → TP)

• Операция g, то есть all-gather, объединяет Y1 и Y2, восстанавливая полную длину последовательности.
• Тем самым восстанавливается Y размера (b,s,h), поскольку column-linear layers нужен полный hidden dimension h.

Первый linear layer (TP-регион)

• A1 и A2 - это column-linear layers, поэтому они делят Y вдоль hidden dimension.
• GELU применяется независимо на каждом GPU.
• Z1* и Z2* имеют форму (b,s,h/2).

Второй linear layer (TP-регион)

• B1 и B2 - row-linear layers, поэтому они восстанавливают hidden dimension.
• W1 и W2 имеют форму (b,s,h) и должны быть просуммированы.

Финальный переход (TP → SP)

• Операция g*, то есть reduce-scatter, выполняет reduction для корректности предыдущего row-linear слоя и одновременно scatter по измерению sequence.
• W1* и W2* имеют форму (b,s/2,h).

Ключевое преимущество sequence parallelism состоит в том, что оно уменьшает максимальный размер activations, которые нам нужно хранить. При одном только tensor parallelism в разных точках приходилось хранить activations формы (b,s,h). При sequence parallelism максимальный размер activation уменьшается до \frac{b \cdot s \cdot h}{tp}, потому что мы всегда делим либо по sequence, либо по hidden dimension.

За всеми частями, которые по-разному shard'ятся в TP и TP+SP, следить непросто. Поверьте, нам самим это тоже дается не бесплатно, поэтому мы сделали небольшую таблицу, суммирующую, как меняется форма activations, то есть hidden_states, по hidden dimension h и sequence dimension s во время forward pass:

А для embedding layer картина такая:

Используя sequence parallelism, мы можем еще сильнее сократить activation memory, а значит - продвинуть и batch size, и длину последовательности дальше, чем это возможно с одним только tensor parallelism. Посмотрим, что это означает для нашего предыдущего примера с моделью 70B:

Мы снова сильно сократили максимальное использование памяти на один GPU, и теперь можем уместить последовательности длиной 16k токенов при TP+SP=16 - это уже заметное улучшение относительно vanilla TP. Да, TP=16 все еще несколько великоват, как мы видели в предыдущем разделе, но в следующем разделе посмотрим, как это улучшить.

Логичный вопрос: не приносит ли TP+SP больше communication overhead, чем vanilla TP? Ответ и да, и нет. В forward pass у vanilla TP было две all-reduce операции на каждый transformer block, а в SP на каждый block приходится две all-gather и две reduce-scatter операции. Формально получается, что SP делает вдвое больше communication operations, чем TP. Но поскольку all-reduce можно разложить на all-gather и reduce-scatter, см. раздел "Ring AllReduce" в appendix, по communication cost эти варианты фактически эквивалентны. Та же логика работает и для backward pass, так как мы просто используем сопряженные версии операций, no-op ↔ allreduce и allgather ↔ reducescatter.

Если вы внимательно следили за рассуждением, то заметили, что в каждом слое речь идет уже о четырех communication operations, двух для attention и двух для MLP. Вот как выглядит profiling MLP при использовании TP+SP:

Как и vanilla TP, TP+SP непросто перекрывать с compute, поэтому throughput здесь сильно зависит от communication bandwidth. И снова, как и в случае vanilla TP, TP+SP обычно выполняют только внутри одного узла, то есть стараются держать TP degree не выше числа GPU на узел, например TP≤8.

Теперь можно посмотреть на бенчмарках, как этот communication overhead становится все более болезненным по мере роста tensor parallelism. Измерим throughput и использование памяти при масштабировании TP вместе с SP для модели на 3B параметров и длины последовательности 4096:

И здесь снова проявляется тот же компромисс: слева - вычислительная эффективность, справа - доступная емкость памяти. Более высокие степени parallelism позволяют обрабатывать значительно более крупные batch size за счет снижения activation memory, но одновременно уменьшают throughput на один GPU, особенно после порога, соответствующего числу GPU на узел.

Кратко подведем итоги:

• Для обоих методов самое заметное падение производительности происходит при переходе от TP=8 к TP=16, потому что именно в этот момент мы выходим из режима обмена только внутри одного узла, NVLink, в режим межузловой communication, EFA.
• Экономия activation memory при использовании TP вместе с SP позволяет уместить существенно более крупные batches, чем при одном только TP.

Итак, мы увидели, как TP помогает shard'ить activations между несколькими GPU, разбивая attention и feedforward-операции вдоль hidden dimension, а SP естественным образом дополняет его, беря на себя оставшиеся операции и разделяя их вдоль измерения sequence.

И все же у TP+SP есть два предела. Во-первых, при дальнейшем росте длины последовательности activation memory в TP-регионе все равно начнет взрываться. Во-вторых, если модель слишком велика и не помещается даже при TP=8, нас ждет резкое замедление из-за межузловой connectivity.

Первую проблему можно атаковать с помощью context parallelism, а вторую - с помощью pipeline parallelism. Для начала разберем context parallelism.

Context Parallelism

С tensor parallelism + sequence parallelism можно заметно уменьшить memory requirements на один GPU, поскольку и веса модели, и activations распределяются между GPU. Однако при обучении на все более длинных последовательностях, например 128k токенов и выше, мы все равно можем выйти за пределы памяти одного узла, потому что внутри TP-региона нам все еще приходится обрабатывать полную последовательность.

Более того, даже если использовать full recomputation activations, а это влечет тяжелый compute overhead порядка 30%, нам все равно нужно хранить в памяти часть activations на границах слоев, а этот объем растет линейно с длиной последовательности. Давайте посмотрим, как здесь помогает context parallelism, или CP.

Базовая идея context parallelism похожа на sequence parallelism, то есть мы снова делим по длине последовательности, но теперь этот подход применяется и к тем модулям, где уже используется tensor parallelism. Иначе говоря, такие модули мы начинаем делить сразу по двум измерениям, тем самым дополнительно уменьшая эффект длины последовательности. После всего, что мы уже обсудили, эта идея должна казаться интуитивной, но здесь есть важный трюк, так что не теряйте нить.

В context parallelism, как и в sequence parallelism, мы делим вход по измерению sequence, но теперь распространяем это деление на всю модель, а не только на sequence-parallel регионы, как было в TP+SP.

Разбиение последовательности почти не влияет на большинство модулей, например MLP и LayerNorm, где каждый токен обрабатывается независимо. Кроме того, здесь не нужна столь дорогая communication, как в TP, поскольку делятся только входы, а не матрицы весов. Как и в data parallelism, после вычисления gradients запускается all-reduce для синхронизации gradients внутри CP group.

Однако есть одно важное исключение: нам нужно особенно внимательно посмотреть на attention blocks. В attention-модуле каждый токен должен иметь доступ к key/value pairs от всех остальных токенов последовательности, либо, в случае causal attention, как минимум ко всем предыдущим токенам.

Поскольку context parallelism делит входы по измерению sequence между GPU, attention-модулю потребуется полноценная communication между GPU для обмена необходимыми key/value data.

Если делать это наивно, звучит очень дорого. Можно ли сделать дешевле и быстрее? К счастью, да: существует ключевая техника, позволяющая эффективно обрабатывать communication key/value pairs, - Ring Attention.

Ring Attention

В этой реализации механизма attention каждый GPU сначала инициирует асинхронную communication operation, чтобы отправить свои key/value pairs другим GPU. Пока он ждет данные от соседей, он вычисляет attention score для той части данных, которая уже находится у него в памяти. В идеале следующий набор key/value приходит с другого GPU до завершения этих вычислений, и тогда GPU может сразу перейти к следующему раунду compute.

Для наглядности предположим, что у нас есть четыре GPU и вход из четырех токенов. Изначально входная последовательность равномерно разделена по измерению sequence, поэтому на каждом GPU находится по одному токену вместе с соответствующими Q/K/V. Пусть Q1, K1 и V1 - это query, key и value первого токена, расположенного на первом GPU. Вычисление attention займет четыре временных шага. На каждом шаге каждый GPU выполняет три последовательные операции:

1. Неблокирующим образом отправляет текущие keys и values на следующую машину, на всех шагах кроме последнего, чтобы можно было перейти к следующей операции, не дожидаясь завершения этой.
2. Локально вычисляет attention score на текущих keys и values, обычно в виде Softmax(\frac{QK^T}{\sqrt{d}}) * V.
3. Ждет получения keys и values с предыдущего GPU и затем возвращается к шагу 1, где текущими keys и values становятся только что полученные значения.

Чтобы завершить вычисление attention, мы повторяем эти три шага четыре раза.

Весь процесс для четырех GPU показан в следующей анимации:

По этой анимации, вероятно, уже очевидно, почему авторы назвали этот подход Ring Attention!

Однако здесь есть одна серьезная проблема: наивная реализация Ring Attention приводит к сильному дисбалансу между GPU из-за формы каузальной attention matrix. Посмотрим на вычисление softmax для attention score matrix с каузальной attention mask:

Softmax вычисляется построчно, а значит, как только GPU получил все токены для строки, он может ее посчитать. Видно, что GPU 1 может сделать это сразу, поскольку начинает с токенов 1-4 и не должен получать информацию от других GPU. А вот GPU 2 придется ждать второго раунда, чтобы получить токены 1-4 и тем самым собрать все значения для токенов 1-8. Кроме того, GPU 1 выполняет заметно меньше работы, чем остальные GPU.

Посмотрим, можно ли лучше сбалансировать вычисления.

Zig-Zag Ring Attention - сбалансированная реализация вычислений

Нужен более удачный способ распределять входные последовательности. Для этого можно назначать токены GPU не строго по порядку, а слегка перемешивать их так, чтобы на каждом GPU было хорошее сочетание ранних и поздних токенов. Такой подход называется Zig-Zag Attention. В этой новой раскладке attention mask показывает более равномерное распределение вычислений: если посчитать число закрашенных клеток, станет видно, что нагрузка теперь сбалансирована между всеми GPU.

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

В целом computation и communication здесь можно перекрывать двумя способами: либо выполнить общий all-gather и одновременно собрать все keys и values на каждом GPU, по аналогии с ZeRO-3, либо получать их по мере необходимости с каждого GPU на каждый GPU.

Ключевые различия между этими двумя реализациями лежат в communication patterns и использовании памяти:

1. Реализация через all-gather:

• Все GPU одновременно собирают полные пары key/value со всех остальных GPU.
• Требует больше временной памяти, поскольку каждому GPU нужно хранить все K/V пары сразу.
• Communication происходит за один шаг, но с более высоким memory overhead.

2. Реализация через all-to-all (ring):

• GPU обмениваются K/V парами по кольцевому паттерну, по одному chunk за раз.
• Более экономна по памяти, так как каждому GPU нужно временно хранить только один дополнительный chunk.
• Communication распределена по времени и перекрывается вычислениями, хотя платой за это становится некоторая дополнительная базовая latency из-за нескольких шагов обмена.

Подход all-to-all обычно дает лучшую эффективность по памяти ценой немного более сложного communication pattern, тогда как подход all-gather проще, но требует больше временной памяти во время attention computation.

Теперь мы увидели, как с помощью TP можно разрезать модель внутри одного узла, чтобы справляться с большими моделями, и как CP помогает укротить взрыв activations при длинных последовательностях.

Однако мы уже знаем, что TP плохо масштабируется между узлами. Что делать, если веса модели не помещаются в пределах одного узла? Здесь на помощь приходит pipeline parallelism - наше четвертое измерение parallelism.

Pipeline Parallelism

В разделе "Tensor Parallelism" мы увидели, что попытка масштабировать tensor parallelism выше числа GPU на одном узле, обычно это 4 или 8, вынуждает нас использовать сетевую communication с более низкой bandwidth, что может заметно ухудшать производительность. Эффект такой inter-node communication хорошо виден на операции all-reduce, если замерить ее на нашем кластере на нескольких узлах, в каждом узле здесь по 8 GPU:

Sequence и context parallelism помогают при длинных последовательностях, но мало что дают, если корневая причина memory issues состоит не в длине sequence, а в размере самой модели. Для крупных моделей, 70B+ параметров, одного только объема весов уже достаточно, чтобы выйти за пределы 4-8 GPU на одном узле. Здесь на помощь приходит еще одно измерение parallelism: pipeline parallelism, или PP.

Pipeline parallelism - техника простая по идее, но очень мощная: мы просто делим слои модели между несколькими GPU. Например, если у нас 8 GPU, можно разместить слои 1-4 на GPU 1, слои 5-8 на GPU 2 и так далее. В таком случае каждому GPU нужно хранить и обрабатывать только часть слоев модели, а значит memory requirements на один GPU существенно снижаются. Посмотрим, как это влияет на использование памяти для модели на 8B параметров:

На рисунке выше видно нечто интересное: хотя параметры модели аккуратно распределены по GPU, activation memory на каждом GPU остается прежней. Иначе говоря, эта техника не экономит activation memory.

Здесь появляется новый тип communication pattern: вместо передачи параметров, как в ZeRO-3 внутри data parallelism, мы теперь последовательно передаем activation tensors между GPU по "pipeline". Концептуально это просто, но эффективная реализация на практике довольно непроста. Перейдем к деталям.

Разбиение слоев по узлам - all forward, all backward

Для начала предположим, что мы просто разложили слои по нескольким устройствам: например, первый GPU берет первые несколько слоев, второй - следующую часть модели и так далее. Тогда forward pass через модель сводится к последовательной передаче batch данных вдоль модели, то есть к последовательному использованию каждого compute device.

Здесь сразу появляется первое важное преимущество: требуемая interconnect bandwidth остается довольно низкой, поскольку мы пересылаем activations умеренного размера лишь в нескольких точках вдоль глубины модели. По сравнению, например, с TP, где communications происходят по нескольку раз внутри каждого слоя, это может давать очень большую разницу.

Но, возможно, вы уже чувствуете, в чем будут проблемы: “последовательно” и “по очереди” звучит не очень-то эффективно в мире parallel computations, особенно после всего, что мы обсуждали про overlap computation и communication.

Именно так. Главная проблема pipeline parallelism - как эффективно обойти последовательную природу PP, чтобы GPU были заняты все время и не возникала ситуация, где один GPU считает, а остальные ждут. Вот как выглядит утилизация GPU при наивном и простом forward и backward pass через модель, здесь числа обозначают слои модели:

Оставшееся время простоя выделено серым и обычно называется “bubble”. После всех наших усилий по оптимизации throughput на это больно смотреть.

Эффективность pipeline setup можно количественно оценить, посмотрев, сколько времени теряется из-за bubble. Пусть t_f и t_b - времена forward и backward pass соответственно, измеренные для одного micro-batch и одной стадии pipeline. Часто принимают простое приближение t_b \approx 2 \times t_f, как на графике выше. При идеальной parallelization общее время было бы t_{id}=t_f + t_b. Но в этом примере из-за pipeline bubble появляется дополнительное время t_{pb}=(p-1)*(t_f+t_b), где p - степень pipeline parallelism, то есть число GPU. Это и есть время, которое каждый GPU проводит в ожидании, пока считают другие GPU.

Отношение дополнительного bubble time к идеальному времени можно вычислить так:

По мере добавления стадий bubble time растет, а utilization падает. Как видно, в наивной реализации bubble может быть очень большой.

К счастью, существует несколько схем pipeline parallelism, которые уменьшают размер bubble.

Достанем первый инструмент из нашего набора и разобьем batch на более мелкие порции, которые можно обрабатывать параллельно или почти параллельно, как мы уже делали в DP. Тогда, пока второй GPU занят обработкой micro-batch 1, первый GPU уже может начать micro-batch 2. Вот как выглядит schedule с восемью micro-batches:

Эта схема называется all forward, all backward (AFAB), потому что сначала выполняются все forward passes, а затем все backward passes. Ее преимущество в том, что этапы forward и backward в целом остаются последовательными, то есть мы сохраняем привычную организацию training code. Это одна из самых простых реализаций PP.

Полную реализацию AFAB pipeline можно посмотреть в Picotron:

Оценим bubble в этом примере. Отличие от самого первого случая в том, что теперь идеальное время обработки m micro-batches равно t_{id} = m*(t_f+t_b):

Как видно, часть неэффективности pipeline stages можно компенсировать, добавляя больше micro-batches и уменьшая размер bubble в m раз.

Однако не меньше, чем bubble, мешает память, необходимая для хранения всех activations. Нам приходится держать их в памяти до начала backward stage, и в таких реализациях PP это быстро приводит к memory explosion. Можно ли сделать лучше и избежать этой проблемы?

Поскольку memory explosion вызван activations, которые мы сохраняем для backward pass, давайте посмотрим, можно ли начать backward еще во время forward-части вычислений. Это позволит как можно раньше освобождать activations, нужные для backward.

One forward, one backward и схемы Llama 3.1

Эта схема называется one forward, one backward (1F1B), потому что в ее среднем, устойчивом режиме по очереди выполняются один forward и один backward pass. Общая идея состоит в том, чтобы начинать backward как можно раньше. Схема выглядит так:

Если внимательно посчитать, окажется, что bubble все еще имеет тот же размер, так что training efficiency почти не улучшается. Однако теперь нужно хранить activations только для p micro-batches, где p - степень pipeline parallelism, вместо m, где m - число micro-batches. Это уменьшает тот activation memory explosion, который был в схеме AFAB. А значит, мы можем добавить больше micro-batches, и это уже действительно уменьшит bubble.

Крупная сложность этого setup, хорошо заметная на рисунке выше, в том, что forward и backward passes больше не идут в чисто последовательном порядке, а выполняются параллельно на разных устройствах и переплетаются между собой. Это значит, что переключение между forward и backward придется планировать независимо на каждом устройстве, а не внутри одного простого общего training loop, как обычно.

Это одна из причин, по которым реализация pipeline parallelism обычно требует довольно серьезных изменений и в training code, и в modeling code.

Полную реализацию 1F1B тоже можно найти в Picotron:

Посмотрим, как схема 1F1B pipeline parallelism масштабируется на практике на бенчмарках нашего кластера:

Слева, когда число micro-batches равно или меньше степени PP минус один (m = p - 1), хорошо видно, насколько вредным может быть pipeline bubble: производительность низкая и даже падает при увеличении PP. На графике справа показано, что использование числа micro-batches, значительно большего, чем степень PP (m = 32 \gg p - 1), улучшает производительность при малых степенях PP, хотя при очень больших степенях PP ограничения все равно остаются. На практике бесконечно увеличивать число micro-batches нельзя, чтобы сохранять соотношение m \gg p - 1, потому что нас в конечном счете ограничивает целевой global batch size. Поэтому при максимально возможном числе micro-batches и дальнейшем росте PP размер bubble в итоге неизбежно растет в соответствии с r_{bubble} = \frac{p - 1}{m}.

Интересно, что при малом числе micro-batches производительность падает всего на 14% при переходе от одного узла (p = 8) к двум узлам (p = 16). Это намного лучше, чем масштабирование tensor parallelism, где в похожих cross-node сценариях обычно наблюдается около 43% деградации. Именно такое поведение при выходе на более медленную inter-node network делает pipeline parallelism особенно привлекательным для распределенного обучения на нескольких узлах.

Хотя 1F1B существенно снижает наш activation memory footprint, на последнем графике видно, что pipeline bubble остается крупным bottleneck по эффективности. Пока размер bubble по-прежнему пропорционален числу pipeline stages, значительная часть GPU compute простаивает. Можно ли придумать еще более умную схему, чтобы уменьшить это потерянное время?

Interleaving stages

Схема 1F1B улучшила использование памяти, но почти не повлияла на размер idle bubble. Можно ли продвинуться дальше?

Оказывается, да, если мы готовы добавить несколько дополнительных communication operations. Пора поговорить про interleaved stages.

До сих пор мы наивно разрезали модель вдоль глубины, например размещая слои 1-4 на первом GPU, а слои 5-8 на втором. Но есть и другие способы разрезать слои, например держать нечетные слои (1, 3, 5, 7) на первом GPU, а четные (2, 4, 6, 8) - на втором.

В общем виде это можно понимать как своего рода "циклический pipeline", где micro-batch по кругу перемещается от одного GPU к следующему, проходя forward pass через модель. Посмотрим, как это работает:

Здесь требуются дополнительные communications, потому что модель проходит через каждый GPU несколько раз в рамках того же вычисления, которое раньше занимало всего один проход. Однако каждый forward и backward pass делится на фактор v, где v - число stages, или model chunks, на один GPU, поскольку мы можем лучше переплетать forward и backward passes:

Итак, теперь мы можем уменьшать bubble, добавляя micro-batches и interleaved stages, но важно помнить, что количественно объем communication тоже растет в v раз, так что это компромисс. На следующем графике показаны несколько конфигураций для PP setup с p=8: особый случай m=1, v=1 соответствует наивному pipeline parallelism, конфигурации с v=1 - это AFAB или 1F1B, а варианты с v \neq 1 - interleaved-конфигурации.

Scheduling здесь тоже становится сложнее, потому что на данном GPU и в данный момент нужно решать, что приоритетнее: более ранние micro-batches, проходящие через более поздние слои, то есть как можно быстрее замыкать forward и backward loops (подход "depth-first", который старается как можно быстрее выводить batches из модели), или более поздние micro-batches, проходящие через ранние слои (подход "breadth-first", который старается как можно полнее заполнить pipeline). Этот выбор подробно разбирается в статье "Breadth-Fist Pipeline Parallelism".

Теперь у вас есть все, чтобы понять подход pipeline parallelism в Llama 3.1, где используется 1F1B setup с interleaved stages и настраиваемым приоритетом между depth-first и breadth-first:

Но и это еще не предел возможных pipeline schedules. Недавно были предложены методы, которые позволяют свести bubble почти к нулю. Эти техники, например, использовались в реализации DeepSeek-V3/R1. Заинтригованы? Давайте перед выходом из мира pipeline parallelism быстро посмотрим и на эти почти магические схемы.

Zero bubble и DualPipe

Недавно появились еще более изощренные способы уменьшать bubble, вплоть до режима, близкого к "zero bubble" - например, pipeline-подход DualPipe, использованный в DeepSeek-V3/R1. Секрет в том, чтобы дробить операции еще на более мелкие части и переплетать их максимально эффективно.

Кратко посмотрим, как это работает, опираясь на zero bubble work от Sea AI Lab, который предшествовал DualPipe. Базовое наблюдение здесь в том, что backward pass через matrix multiplication на самом деле состоит из двух отдельных операций: backward для входов (B) и backward для весов (W).

Выход B, то есть backward pass для входов, необходим для backward pass нижележащих слоев. А вот backward pass для весов, W, для этого не нужен и обычно должен быть выполнен лишь до optimizer step. Это видно на следующей диаграмме из статьи Zero Bubble:

Это означает, что W можно гибко планировать в любой момент после соответствующего B той же стадии. Благодаря этому W можно стратегически размещать так, чтобы заполнять pipeline bubbles. Схема ZB-H2 в правом верхнем углу - пример теоретического schedule с zero bubble, использующего это более тонкое разложение.

DualPipe от DeepSeek, представленный в техническом отчете V3 , расширяет этот разложенный подход на дополнительный случай двух потоков, распространяющихся с обоих концов измерения PP. Эти потоки переплетаются между собой, чтобы еще сильнее сократить idle time на GPU. Эта схема показана на следующем графе scheduling - как видно, она еще сложнее предыдущих:

Вообще, полная оптимизация таких сложных schedules требует аккуратно измерять длительность разных мелкозернистых операций и решать задачу Integer Linear Programming (ILP), чтобы минимизировать итоговый bubble time. Обсуждение используемых эвристик и алгоритмов можно найти, например, в статье Zero Bubble. Поэтому схемы zero bubble и DualPipe слишком сложны, чтобы приводить здесь код, но общее представление о лежащих в основе идеях у вас уже есть.

На этом завершается наш тур по миру pipeline schedules и bubbles.

Теперь пора перейти к последнему методу parallelism, который мы подробно разберем и который позволяет эффективно обучать большие модели: expert parallelism.

Expert Parallelism

Это последний метод parallelism, который мы обсудим. Если вы пока мало знакомы с моделями Mixture of Experts (MoE), возможно, сначала стоит прочитать этот более короткий блог-пост, который мы публиковали раньше: он поможет лучше понять архитектуру MoE в целом.

Парадигма Mixture of Experts в последнее время получила большую популярность благодаря моделям вроде GPT-4, Mixtral и DeepSeek-V3/R1. Базовая идея в том, что вместо одного feedforward module на слой можно иметь несколько параллельных модулей и маршрутизировать через них токены, чтобы они обрабатывались по-разному.

Устройство MoE-слоев делает очень удобной реализацию parallelism по измерению experts - то, что мы называем expert parallelism (EP). Поскольку feedforward layers полностью независимы, можно просто разместить feedforward layer каждого expert на отдельном worker. По сравнению с TP этот подход гораздо легче, потому что нам не нужно дробить matrix multiplication; достаточно лишь отправить hidden states токена к нужному expert.

На практике EP обычно используется вместе с другими формами parallelism, например с data parallelism. Причина в том, что EP затрагивает только MoE-слои и не shard'ит входные токены, в отличие от context parallelism, которая shard'ит токены вдоль sequence length. Это значит, что если использовать только EP, наши GPU будут выполнять избыточные вычисления во всех не-MoE блоках. Комбинируя EP с DP, можно эффективно shard'ить и experts, и входные batches между GPU, как показано на упрощенной схеме ниже:

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

На практике для эффективной работы EP есть несколько важных трюков, и они тесно связаны с дизайном модели. Например, в DeepSeek-V3 в router введено ограничение, которое гарантирует, что каждый токен отправляется не более чем на M узлов, в их случае на 4, чтобы удерживать токены в пределах одного узла и уменьшать communication overhead. Хотя expert parallelism существует уже давно, новый импульс этот подход получил именно сейчас, вместе с ростом популярности MoE-архитектур.

Мы планируем скоро добавить более полный пример EP в Picotron/Nanotron.

5D Parallelism в двух словах

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

1. Data parallelism (DP) - вдоль batch dimension
2. Tensor parallelism (TP) - вдоль hidden dimension
3. Sequence и context parallelism (SP/CP) - вдоль sequence dimension
4. Pipeline parallelism (PP) - вдоль слоев модели
5. Expert parallelism (EP) - вдоль измерения experts

а также три стратегии ZeRO, которые можно комбинировать с data parallelism для уменьшения потребления памяти:

1. ZeRO-1 - sharding optimizer states между DP replicas
2. ZeRO-2 - sharding optimizer states и gradients между DP replicas
3. ZeRO-3 - sharding optimizer states, gradients и parameters между DP replicas

На этом этапе естественно возникает вопрос: как все эти стратегии parallelism и ZeRO соотносятся друг с другом и как они взаимодействуют? Иначе говоря, какие из них можно эффективно комбинировать, а какие лучше держать отдельно?

Давайте разберем сходства и взаимодействия. Начнем с прямого сравнения pipeline parallelism и ZeRO-3, поскольку между ними есть как очень близкие сходства, так и важные различия.

И pipeline parallelism, и ZeRO-3 - это способы partition'ить веса модели между несколькими GPU и организовывать communication/computation вдоль оси глубины модели, например, в ZeRO-3 мы prefetch'им следующий слой во время вычислений. Это означает, что в обоих случаях на каждом устройстве выполняются операции целого слоя, в отличие, например, от TP или EP, где вычисления идут на более мелких подслойных единицах.

Тем не менее между PP и ZeRO-3 есть несколько крупных различий:

Как видно, ZeRO-3 и PP решают одну и ту же задачу, но разными способами, и выбор между ними зависит от того, хотите ли вы сосредоточить communication на передаче weights или activations. Хотя их можно комбинировать, на практике это делается нечасто, потому что для амортизации communication cost приходится заметно увеличивать global batch size, а это создает trade-off между global batch size, размером модели, network bandwidth и эффективностью обучения. Если вы все же комбинируете эти подходы, ZeRO-3 стоит настраивать так, чтобы weights оставались в памяти на протяжении серии PP micro-batches и не создавали лишний communication overhead.

С другой стороны, ZeRO-1 и ZeRO-2, которые работают с optimizer states и gradients, легко сочетаются с pipeline parallelism и естественно его дополняют. Эти комбинации не порождают каких-то особых новых проблем. Например, при обучении DeepSeek-v3 использовалась комбинация PP и ZeRO-1.

Tensor parallelism, вместе с sequence parallelism, естественно дополняет и pipeline parallelism, и ZeRO-3, потому что опирается на дистрибутивное свойство matrix multiplication, позволяющее shard'ить weights и activations, вычислять их независимо, а затем объединять.

Главная причина, по которой не хочется полагаться только на TP, состоит в двух практических ограничениях, которые мы уже обсуждали. Во-первых, его communication operations находятся на critical path вычислений, поэтому после некоторой точки communication overhead начинает доминировать и TP масштабируется плохо. Во-вторых, в отличие от ZeRO и PP, которые model-agnostic, TP требует очень аккуратной работы с activation sharding, иногда по hidden dimension, в TP-регионе, а иногда по sequence dimension, в SP-регионе. Из-за этого его сложнее правильно реализовать, и он требует model-specific знаний, чтобы гарантировать корректные sharding patterns по всей модели.

Как следствие, при комбинировании стратегий parallelism TP обычно оставляют для быстрых intra-node communications, тогда как ZeRO-3 или PP можно использовать для групп parallelism, выходящих на более медленные inter-node communications, потому что их communication patterns либо требуют меньше bandwidth, как в PP, либо проще перекрываются вычислениями, как в ZeRO-3. Главная задача при комбинировании этих техник - эффективно организовать GPU в группы для каждого измерения parallelism так, чтобы максимизировать throughput и минимизировать communication overhead, не забывая об ограничениях TP по масштабированию. Например, группы GPU, которые общаются в рамках TP, стоит держать внутри узлов.

Context parallelism и expert parallelism тоже помогают shard'ить activations и могут рассматриваться как дополнение к TP. CP работает с длинными последовательностями, а EP позволяет распределенно обучать Mixture of Experts модели, и эти подходы можно комбинировать без каких-то особых проблем.

CP специально нацелена на задачу обучения с очень длинными последовательностями, shard'я activations по измерению sequence между GPU. Хотя большинство модулей, например MLP и LayerNorm, могут независимо работать с такими sharded sequences, attention blocks требуют communication, потому что каждому токену нужен доступ к keys/values всей последовательности. Как мы видели в разделе про CP, это эффективно решается через паттерны Ring Attention, которые перекрывают computation и communication. CP особенно ценна при переходе к экстремальным длинам последовательности, 128k+ токенов, где даже с full activation recomputation memory requirements для attention были бы неприемлемы для одного GPU.

Expert parallelism специально решает задачу обучения MoE-моделей, shard'я специализированных "experts" между GPU и динамически маршрутизируя токены к нужным experts во время вычислений. Ключевые communication operations в EP - это "all-to-all" операции, которые отправляют токены их назначенным experts и затем собирают результаты обратно. Хотя это добавляет определенный communication overhead, взамен мы получаем возможность сильно наращивать model capacity, поскольку каждый токен во время inference, и во время training тоже, обрабатывается лишь небольшой долей от общего числа параметров. В distributed training и inference partitioning experts между GPU становится особенно актуальным, когда число experts в модели становится большим.

Область действия и фокус

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

• Tensor parallelism, вместе с sequence parallelism, влияет на вычисления по всей модели, shard'я и weights, и activations.
• Context parallelism в первую очередь влияет на attention layers, поскольку именно там нужна cross-sequence communication, тогда как остальные слои работают независимо на sharded sequences.
• Expert parallelism в первую очередь затрагивает MoE layers, которые заменяют стандартные MLP blocks, оставляя attention layers и другие компоненты без изменений.
• Pipeline parallelism и ZeRO не привязаны к какому-то одному подмодулю или компоненту, за исключением того, что в pipeline parallelism нужно балансировать модули и слои, поэтому первый и последний слои нередко обрабатываются отдельно из-за дополнительных embedding layers.

Сводим все вместе

А что если собрать все рассмотренные техники в одну общую схему, объединяющую их все сразу? Да, мы готовы к такому вызову.

В этой итоговой диаграмме показаны activations и модули одного transformer layer в его MoE-варианте. Также здесь изображены разные направления parallelism и communication operations, которые мы обсуждали в предыдущих разделах.

Кроме того, можно дать и полный обзор экономии памяти для каждой из этих стратегий. Мы построим графики для разных длин sequence, а также для selective recomputation, сверху, и full recomputation, снизу, чтобы показать, как все это взаимодействует с activations:

Завершим этот раздел высокоуровневым обзором всех этих техник, их базовых идей и главных bottlenecks:

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

Как подобрать лучшую конфигурацию обучения

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

Мы уже слегка коснулись этого в предыдущем разделе, но теперь давайте более подробно разберем возможный процесс принятия решения, шаг за шагом. Разумеется, в реальной жизни все равно придется провести несколько экспериментов, чтобы найти действительно оптимальный setup для вашего compute cluster с учетом его физических свойств, network bandwidth, числа GPU на узел, памяти на GPU и так далее.

Шаг 1: уместить training step в память

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

1. Сценарий GPU-rich 🤑 - когда у вас в распоряжении много GPU:

• Для моделей меньше 10B параметров можно использовать одну технику parallelism, например tensor parallelism или ZeRO-3/DP с full recompute на 8 GPU.
• Для моделей в диапазоне 10B-100B параметров, которым требуется больше 8 GPU, есть несколько вариантов:
• Комбинировать tensor parallelism, TP=8, с pipeline parallelism
• Комбинировать tensor parallelism, TP=8, с data parallelism, ZeRO-3
• Использовать только ZeRO-3, то есть чистый data parallelism
• На масштабе 512+ GPU чистый data parallelism, ZeRO-3, начинает терять эффективность из-за communication cost, и тогда часто лучше комбинировать DP либо с tensor, либо с pipeline parallelism.
• На масштабе 1024+ GPU разумной рекомендацией может быть tensor parallelism, TP=8, в сочетании с data parallelism, ZeRO-2, и pipeline parallelism.

Особые случаи:

• Для очень длинных последовательностей вам, скорее всего, захочется добавить context parallelism между узлами.
• Для архитектур Mixture of Experts будет выгодно использовать expert parallelism между узлами.

2. Сценарий GPU-poor 😭 - когда GPU-ресурсов мало:

• Можно включить full activation recomputation, обменивая часть compute на память, пусть и ценой более медленного обучения.
• Можно увеличить gradient accumulation, чтобы обрабатывать большие batches при ограниченной памяти.

Теперь, когда первый экземпляр модели уже можно обучать, нужно убедиться, что batch size выбран правильно.

Шаг 2: достичь целевого global batch size

В зависимости от того, к чему привел нас шаг 1 с точки зрения micro-batch size и DP, текущий batch size может оказаться либо слишком маленьким, либо слишком большим. Теперь нужно попасть в целевой batch size.

Чтобы увеличить текущий global batch size:

• Можно увеличить data parallelism или число шагов gradient accumulation.
• Для длинных последовательностей можно использовать context parallelism.

Чтобы уменьшить текущий global batch size:

• Можно уменьшить data parallelism в пользу других стратегий parallelization.
• Для длинных последовательностей можно уменьшить context parallelism.

Хорошо, теперь модель уже работает в нужной нам общей конфигурации с точки зрения размера модели и batch size. Но обучаем ли мы ее максимально быстро? Последний шаг - оптимизация throughput.

Шаг 3: оптимизация training throughput

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

• Увеличивать tensor parallelism, используя быструю intra-node bandwidth, пока не приблизимся к числу GPU в узле, чтобы затем уменьшить другие формы parallelism.
• Увеличивать data parallelism вместе с ZeRO-3, сохраняя целевой batch size.
• Когда communication в data parallelism начинает становиться bottleneck, переключаться на pipeline parallelism.
• Поочередно пробовать увеличивать разные виды parallelism.
• Экспериментировать с micro-batch size, mbs, чтобы найти оптимальный баланс между max global batch size, размером модели, compute и communication.

Бенчмаркинг тысяч конфигураций

Теперь, когда пошаговый процесс описан, давайте посмотрим, как реализовать такой поиск в реальной жизни.

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

Мы сами прогнали бенчмарки на нескольких тысячах распределенных конфигураций, охватив все размеры моделей, о которых говорили, а также очень большое число конфигураций кластера, а именно от 1 до 64 узлов по 8xH100, чтобы получить результаты, которые использовались в этой книге.

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

Все приведенные ниже бенчмарки проводились с длиной последовательности 4096 и global batch size в 1M токенов. Мы собрали лучшие конфигурации для каждой модели и размера кластера и изобразили их на следующих heatmap:

Из этой высокоуровневой визуализации можно вынести несколько важных наблюдений:

• Во-первых, с ростом числа узлов, то есть при усилении parallelism, мы наблюдаем снижение эффективности. Особенно заметно это на маленьких моделях, у которых отношение compute к размеру модели ниже. Обычно это можно было бы компенсировать увеличением batch size, но здесь нас ограничивает global batch size в 1M.
• Во-вторых, у больших моделей возникает другая проблема. По мере роста размера модели memory requirements резко увеличиваются. Поэтому на малом числе узлов возникают два сценария: либо модель вообще не помещается, либо помещается, но работает неэффективно, потому что почти упирается в пределы памяти GPU, как, например, в случае модели на 80B параметров при обучении на 4 узлах.
• Наконец, наши бенчмарки показывают, насколько сильно производительность зависит от качества реализации. Когда мы впервые реализовали обе стратегии, tensor parallelism обгонял pipeline parallelism. После оптимизации нашего PP-кода более быстрым стал PP. Сейчас, улучшая communication overlap в реализации TP, мы ожидаем, что лидерство снова вернется к TP.

Что мы поняли из бенчмаркинга

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

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

• Процессы PyTorch иногда не могли корректно завершиться и очистить ресурсы.
• Менеджер задач Slurm принудительно завершал jobs, что приводило к отказам узлов.
• Простые бенчмарки, которые должны были занимать минуты, растягивались на часы.
• Некоторые jobs зависали бесконечно.

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

• Минимизации времени перезапуска кластера и оптимизации idle time
• Анализа подробных NCCL debug logs
• Понимания паттернов потребления памяти и поведения CUDA memory allocator
• Улучшения производительности pipeline parallelism на multi-node setup

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

Воспроизводить теоретические результаты в реальной жизни непросто, особенно если учесть ограниченную доступность production training code. Через open source-проекты вроде Nanotron и Picotron мы надеемся сделать техники распределенного обучения доступнее, а также совместно развивать простые и эффективные codebase, которые помогают исследователям и инженерам извлекать максимум из их аппаратных ресурсов.

На этом завершается наш очень глубокий разбор методов распределения в 5D parallelism.

Если сделать шаг назад, окажется, что до сих пор наше обсуждение часто опиралось на критически важное допущение: будто computation и communication можно эффективно перекрывать на GPU без влияния на computation throughput. В реальности все тоньше. При использовании распространенных communication primitives вроде NCCL send/recv возникает скрытая конкуренция между ресурсами вычислений и коммуникации, потому что communication kernels обычно используют те же streaming multiprocessors GPU, что и вычисления. Это снижает throughput, когда communication перекрывается с computation. Чтобы по-настоящему оптимизировать распределенное обучение, нужно глубже разобраться в архитектуре самих GPU.

Пора приглушить свет и включить режим CUDA.

Погружение в GPU - fusing, threading и mixing

До сих пор мы говорили в основном о высокоуровневой организации операций модели. Мы переставляли вычисления между разными accelerator, учитывая общие ограничения памяти и высокоуровневый scheduling compute units.

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

В этом разделе мы углубимся в детали архитектуры GPU. Основной фокус будет на архитектуре GPU от NVIDIA, но общие идеи, как это часто бывает, можно перенести и на похожие accelerator.

Сначала мы кратко объясним, как устроены GPU, а затем разберем революцию FlashAttention, эффективный scheduling workloads на GPU и то, как на GPU можно эффективно использовать разные численные точности.

Краткое введение в GPU

В целом GPU устроены очень иерархично. Со стороны вычислений GPU состоит из массива compute units, называемых streaming multiprocessors (SMs). Каждый SM содержит и управляет набором streaming processors, которые также называют cores. Например, GPU NVIDIA H100 имеет 132 SM и по 128 cores в каждом SM, то есть всего 16 896 cores (подробности можно посмотреть в документации по tensor cores), и каждый из них способен одновременно обрабатывать несколько threads.

Подсистема памяти тоже сильно иерархична и включает несколько уровней cache и memory. Registers - самые маленькие единицы, и они приватны для thread во время выполнения. Shared memory и L1 cache разделяются между thread, работающими на одном SM. Выше находится L2 cache, общая для всех SM, а затем global memory - самая большая память на GPU, например заявленные 80 GB у H100, но при этом и самая медленная по доступу.

Цель при использовании GPU - запускать как можно больше workloads параллельно на доступных cores, используя преимущества этой иерархической организации compute/memory resources.

Фрагмент кода, который выполняется на core GPU, называется kernel. Его можно написать на высоком уровне, например на CUDA или Triton, а затем он компилируется в Parallel Thread Execution (PTX) - низкоуровневый assembly, используемый GPU от NVIDIA.

Чтобы запустить kernel, нужен еще и host code, который исполняется на CPU/host и отвечает за подготовку выделения памяти, загрузку данных и кода:

Обычно kernels планируются следующим образом:

• Threads объединяются в warps, каждая из которых содержит 32 threads. Все threads внутри warp синхронизированы и одновременно исполняют инструкции, но над разными частями данных.
• Warps объединяются в более крупные blocks более гибкого размера, например по 512 или 1 024 threads, причем каждый block назначается одному SM. Один SM может параллельно исполнять несколько blocks. Однако в зависимости от доступных ресурсов не все blocks будут запущены немедленно: часть может ждать, пока освободятся ресурсы.

Главное, что стоит вынести отсюда: существуют разнообразные ограничения по размерам и размещению - объемы разных типов памяти, число одновременно выполняющихся blocks и число threads в warps, - и их нужно учитывать, если вы хотите использовать архитектуру GPU максимально эффективно.

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

Улучшение производительности с помощью kernels

Если вам нужно добавить новую операцию без готового оптимизированного kernel или ускорить уже существующую функцию PyTorch, может показаться, что самый прямой путь - писать kernel с нуля. Но создание действительно быстрых CUDA-kernel требует большого опыта и сопровождается крутой кривой входа. На практике разумнее начать с torch.compile: он динамически оптимизирует код PyTorch, перехватывает ваши операции и генерирует низкоуровневые высокопроизводительные kernels на Triton.

Допустим, вы хотите написать kernel для функции активации Exponential Linear Unit (ELU):

Начать можно с простой реализации на PyTorch, а затем просто добавить сверху декоратор @torch.compile:

Как видно на следующем графике, разница в производительности между compiled- и non-compiled-версиями весьма заметна, особенно если учесть, что мы лишь добавили один декоратор (N здесь обозначает число столбцов):

Если этого прироста все же недостаточно, следующим шагом может стать написание Triton-kernel. Начать удобно с kernel, который сгенерировал @torch.compile. Для этого достаточно выставить переменную окружения TORCH_LOGS в "output_code":

После запуска Python-скрипта с декоратором @torch.compile будет сгенерирован и выведен соответствующий Triton-kernel. В нашем случае он выглядит так:

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

Здесь tl.program_id(0) задает уникальный идентификатор блока, по которому мы определяем, какой фрагмент данных он будет обрабатывать. Из этого идентификатора вычисляется block_start, то есть начальный индекс участка, а block_indices задает диапазон индексов внутри него. Маска valid_mask гарантирует, что будут обработаны только индексы в пределах num_elements, после чего данные безопасно читаются через tl.load. Затем применяется ELU, которая по-разному обрабатывает отрицательные и неотрицательные значения, а результат записывается обратно в память с помощью tl.store.

Если прогнать сгенерированный kernel через triton.testing.Benchmark, получится такая производительность:

Этот standalone-kernel даже показывает более высокую производительность на небольших размерах, чем @torch.compile, хотя, скорее всего, это просто артефакт, связанный со временем компиляции torch.compile. Но главный вывод в другом: вовсе не обязательно начинать с чистого листа. Гораздо практичнее взять уже сгенерированный kernel и сосредоточиться на его доработке, что экономит массу времени.

Даже в Triton не всегда удается выжать из устройства потолок производительности: язык ограничен в работе с низкоуровневыми деталями вроде shared memory и планирования внутри streaming multiprocessors. Triton умеет управлять блоками и их распределением по SM, но если нужен еще более тонкий контроль, придется писать kernels прямо на CUDA, где доступны все эти низкоуровневые механизмы.

На уровне CUDA существует множество приемов, которые помогают сделать kernels эффективнее. Мы разберем лишь некоторые из них: оптимизацию паттернов доступа к памяти для снижения задержек, использование shared memory для часто переиспользуемых данных и управление загрузкой threads, чтобы минимизировать простои.

Прежде чем углубляться в примеры на CUDA, кратко подытожим уже встречавшиеся нам инструменты, с помощью которых можно писать kernel-код для GPU:

1. PyTorch: просто, но медленно
2. @torch.compile: просто и быстро, но не слишком гибко
3. Triton: сложнее, быстрее и гибче
4. CUDA: сложнее всего, зато быстрее всего и максимально гибко, если сделать все правильно

Начнем с одного из самых частых сценариев использования CUDA: оптимизации доступа к памяти. Global memory в GPU, то есть самая большая область памяти, о которой мы говорили выше, отличается высокой задержкой и сравнительно низкой пропускной способностью по сравнению с кэшем. Для многих приложений это превращается в главный bottleneck, поэтому эффективная работа с global memory может радикально повысить производительность.

Memory coalescing

Чтобы по-настоящему использовать пропускную способность global memory, нужно понимать, как она устроена. В CUDA-устройствах global memory реализована на базе DRAM.

Memory coalescing использует тот факт, что DRAM отдает данные пакетами при обращении к памяти. Когда считывается один адрес, параллельно читается и последовательность соседних адресов, включая запрошенный, после чего эти данные быстрым burst-передачей отправляются процессору. В CUDA coalescing позволяет максимально эффективно использовать этот механизм: threads внутри warp, то есть группы из 32 threads, исполняющих одну и ту же инструкцию синхронно, должны обращаться к последовательным адресам в памяти. Например, если thread 0 читает адрес M, thread 1 - M + 1, thread 2 - M + 2 и так далее, GPU объединяет эти обращения в один крупный и эффективный запрос к DRAM вместо того, чтобы обрабатывать каждое обращение по отдельности.

Рассмотрим на примере матричного умножения. В самой простой реализации каждый thread вычисляет один элемент выходной матрицы:

Вот отличная визуализация этого kernel из прекрасного поста в блоге Simon Boehm:

Однако при профилировании этого kernel с помощью инструмента вроде ncu сразу видны проблемы: низкая пропускная способность памяти и некоалесцированные обращения к памяти:

Причина в том, что в этом kernel два threads из одного блока с идентификаторами (0, 0) и (1, 0), которые в итоге попадут в один warp, читают одну и ту же колонку матрицы B, но разные строки матрицы A. Поскольку элементы матрицы хранятся в row-major порядке, то есть элементы строки лежат по соседним адресам, как показано на рисунке ниже, в первой итерации при i = 0 thread (0, 0) загрузит A_{0,0}, а thread (1, 0) - A_{1,0}. Эти элементы находятся далеко друг от друга в памяти, и такое несоответствие повторяется на каждой итерации, не давая обращениям к памяти склеиться в coalesced-доступ.

Чтобы ускорить этот kernel, можно изменить способ вычисления координат x и y следующим образом:

Вместо двумерного блока мы переходим к одномерному и по-новому определяем x и y. Теперь threads внутри одного warp, у которых близкие значения threadIdx.x, будут иметь одинаковый x, но разные y. Это означает, что они читают одну и ту же строку матрицы A, но разные столбцы матрицы B. Для матриц в row-major представлении такой доступ уже можно коалесцировать.

Если снова запрофилировать новый kernel, предупреждение о некоалесцированных обращениях исчезает, а пропускная способность памяти GPU вырастает примерно в 10 раз.

Вдобавок к этому время выполнения kernel уменьшается примерно в 10 раз. Неплохо.

Теперь перейдем к еще одной технике, которая постоянно всплывает в литературе: tiling.

Tiling

Tiling - это техника, использующая shared memory для оптимизации паттернов доступа к памяти. Как мы уже говорили, shared memory на GPU - это небольшая, но очень быстрая область памяти, доступная всем threads внутри блока. Она позволяет переиспользовать данные между несколькими threads и тем самым реже обращаться к более медленной global memory.

Например, при матричном умножении каждому thread внутри блока могут понадобиться элементы из двух матриц, скажем A и B. Если каждый thread будет независимо загружать нужные ему строку и столбец из global memory, получится масса избыточных загрузок, потому что threads внутри одного блока часто работают с пересекающимися данными. Вместо этого tiling позволяет один раз загрузить блок, или tile, матриц A и B в shared memory, после чего все threads этого блока переиспользуют один и тот же набор данных.

В подходе с tiling на каждой итерации все threads блока совместно загружают в shared memory два tile: один из матрицы A, другой из матрицы B. Конкретно, они загружают tile матрицы A размера BLOCK_SIZE_M на BLOCK_SIZE_K и tile матрицы B размера BLOCK_SIZE_K на BLOCK_SIZE_N. После этого threads выполняют умножение этих tile, и вычисления идут быстро, поскольку все нужные данные уже находятся под рукой. Результаты складываются в accumulation-матрицу с промежуточными значениями. На каждой итерации частичный результат от текущих tile добавляется в эту матрицу накопления, и процесс продолжается, пока не будут обработаны все tile обеих матриц.

Посмотрим на ключевые части реализации, которые здесь важно понять:

Каждый thread начинает с загрузки по одному элементу из матриц A и B в shared memory. В этом сценарии добиться coalesced-доступа довольно просто: если использовать threadIdx.x как локальный индекс столбца (localCol), то threads внутри одного warp будут обращаться к соседним элементам обеих матриц. После того как все threads блока закончат загрузку своих элементов в shared memory, что обеспечивается вызовом __syncthreads(), они переходят к вычислению скалярного произведения двух tile. Когда threads пройдут по всем tile, по горизонтали для A и по вертикали для B, итоговая сумма записывается в соответствующую позицию матрицы C.

При benchmarking этого kernel через ncu мы увидели, что пропускная способность памяти выросла до 410 Gb/s, а время выполнения сократилось примерно на 43%, что соответствует производительности около 6.6 TFLOPS.

Thread coarsening

Tiling заметно ускорил наш kernel. Однако если посмотреть на warp states, то есть на распределение тактов по различным состояниям выполнения, мы увидим следующее:

Расшифровку этих криптических состояний можно найти в руководстве NVIDIA по профилированию kernel, в разделе "Warp Stall Reasons". Там указано, что smsp__pcsamp_warps_issue_stalled_mio_throttle означает следующее: warp простаивал в ожидании, пока очередь инструкций MIO (memory input/output) перестанет быть переполненной. Такое часто происходит при экстремальной нагрузке на MIO pipeline, куда входят специальные математические инструкции, динамические ветвления и инструкции shared memory. Если проблема вызвана доступом к shared memory, снизить давление на pipeline иногда помогает меньшее число, но более широких загрузок.

Иными словами, warps здесь простаивают, ожидая завершения обращений к shared memory. Один из способов ослабить проблему - применить thread coarsening, то есть объединить несколько threads в один укрупненный thread. Тогда число обращений к shared memory существенно уменьшится, потому что каждый такой укрупненный thread сможет вычислять сразу несколько выходных элементов.

Теперь коротко разберем еще один важный момент, который стоит учитывать при написании и оптимизации custom-kernel: минимизацию control divergence.

Minimizing control divergence

Streaming multiprocessor устроен так, чтобы исполнять все threads внутри warp по модели Single Instruction, Multiple Data (SIMD). Это значит, что в каждый момент времени одна и та же инструкция выбирается и выполняется одновременно для всех threads данного warp. Threads работают с разными фрагментами данных, но следуют одной инструкции, откуда и название SIMD. Главный плюс этого подхода - эффективность: управляющая логика, отвечающая за выборку и диспетчеризацию инструкций, делится между несколькими вычислительными блоками, что снижает аппаратные накладные расходы и освобождает больше ресурсов под арифметику.

Control divergence возникает, когда threads внутри одного warp расходятся по разным веткам исполнения. Например, если условие вроде if приводит к тому, что часть threads идет по одной ветке кода, а часть - по другой, warp вынужден исполнять эти ветки последовательно, а значит часть threads будет простаивать, ожидая остальных. Чтобы минимизировать control divergence, kernels проектируют так, чтобы threads внутри warp по возможности проходили один и тот же путь исполнения. Добиваются этого перестройкой кода для уменьшения ветвлений, подбором структур данных, которые ведут threads по похожим траекториям, и приемами вроде predication.

Мы разобрали несколько ключевых идей, которые помогают писать custom-kernels и улучшать производительность и memory footprint GPU-операций. Но прежде чем перейти к реальному примеру, осталось затронуть еще одну важную концепцию: fusion kernel.

Fused kernels

Мы уже не раз упоминали, что работа CPU и GPU может идти асинхронно. В частности, host-код на CPU может неблокирующим образом ставить задачи на GPU.

Это полезно не только для overlap между коммуникацией и вычислениями, который мы много раз обсуждали, но и для более общей цели: как можно реже гонять управление туда-сюда между host-кодом и отдельными GPU-kernel. Идею очень наглядно иллюстрирует Horace He в этих схемах:

Последовательность kernels, в которой данные постоянно курсируют между global memory и вычислительными блоками

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

Как избежать этого бессмысленного хождения туда-сюда, показанного слева? Лучший способ - сделать GPU как можно более автономным. Для этого в один kernel упаковывают как можно больше последовательных вычислительных операций. Такой kernel и называется fused kernel, как на схеме справа.

Fused kernels особенно эффективны и сравнительно просты в написании для цепочек point-wise операций, которые выполняются независимо для каждого входного token. В такой ситуации нет смысла отправлять промежуточные значения обратно в global memory, а потом снова переносить их в память SM и запускать новый kernel. Гораздо эффективнее держать все локально, пока не будут выполнены все вычисления.

В Transformer-модели такой подход к fusion можно применять всякий раз, когда у нас идет последовательность point-wise операций, например в вычислениях внутри слоев LayerNorm.

Теперь у нас есть все необходимое, чтобы оценить настоящий шедевр kernel engineering: FlashAttention.

FlashAttention

FlashAttention был предложен Tri Dao как способ ускорить вычисления attention с помощью custom CUDA-kernel, сделав их одновременно быстрее и экономнее по памяти. Идея FlashAttention в том, чтобы рационально использовать разные уровни памяти GPU и как можно меньше зависеть от самого медленного из них: global memory.

Global memory в современных GPU часто построена на технологии High Bandwidth Memory (HBM). Несмотря на название, в иерархии памяти GPU она все равно медленнее, чем SRAM. Термин HBM будет важен дальше, когда мы перейдем к деталям реализации FlashAttention.

В базовой реализации механизма attention приходится много гонять данные между памятью и вычислителями. Нужно материализовать матрицу S, где S = QK^T, то есть attention scores, а также матрицу P, где P = \text{softmax}(S), то есть нормализованные веса attention, прямо в HBM. Это означает, что промежуточные результаты сначала отправляются в HBM, а затем возвращаются в SRAM для следующих вычислений:

Поскольку пропускная способность HBM заметно ниже, это создает тяжелый bottleneck в вычислении attention. Можно ли лучше? Tri Dao отвечает утвердительно.

Ключевая идея состоит в том, чтобы вычислять матрицу S небольшими фрагментами, которые помещаются в сравнительно небольшую shared memory внутри SM. Но можно пойти еще дальше и вовсе не материализовывать огромную матрицу S, а сохранять только статистики, необходимые для вычисления нормировочного множителя в softmax. Тогда часть выхода O можно посчитать напрямую в SRAM за один проход, не гоняя промежуточные результаты туда-сюда. В результате мы не только эффективно используем shared memory, но и снимаем memory bottleneck, связанный с материализацией одной из самых крупных activation-матриц модели при больших длинах контекста: матрицы attention.

Идея FlashAttention снимает столько bottlenecks в обучении моделей, что этот подход очень быстро стал стандартным способом считать attention практически во всех transformer-моделях. В частности:

• не материализуя матрицу S, мы снижаем memory burden механизма attention;
• кроме того, мы сильно ослабляем наивное влияние квадратичной цены attention, то есть O(S^2).

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

После FlashAttention-1 та же лаборатория выпустила еще две улучшенные версии: FlashAttention-2 и FlashAttention-3. По сравнению с первой версией прогресс в них связан уже не столько с общей идеей attention, сколько с тем, как еще точнее подогнать низкоуровневую реализацию под GPU: во-первых, максимально сократить число non-matmul операций, во-вторых, аккуратно разложить workload между warps и thread blocks, это особенно важно для FlashAttention-2, и, в-третьих, тонко оптимизировать поддержку FP8 и Tensor Core под архитектуру Hopper, то есть H100, в FlashAttention-3.

FlashAttention - это образцовая демонстрация того, какой прорыв в производительности возможен, если проектировать алгоритм с учетом внутренней memory/compute-архитектуры современных GPU-ускорителей.

Техники, которые мы разобрали в этом разделе, требовали либо изменений в modeling code, либо написания custom-kernels для отдельных операций, чтобы ускорить обучение.

В последней части нашего низкоуровневого погружения в compute-операции мы рассмотрим набор методов, которые не зависят от конкретного modeling code. Они применимы практически к любой модели и настолько широко используются, что давно стали отраслевым стандартом. Следующий пункт программы - mixed precision training.

Обучение в mixed precision

В разных разделах этой книги мы уже говорили о форматах пониженной точности и о том, как они влияют на потребности в памяти для хранения activations, parameters и optimizer states. Теперь пора углубиться в детали этих форматов и лучше понять их компромиссы, преимущества и ограничения.

Mixed precision training, как следует из названия, состоит в смешивании разных точностей в процессе обучения. Численная точность по умолчанию для tensors в PyTorch - это single-precision floating-point format, также называемый FP32 или float32, то есть каждое число занимает 32 бита, или 4 байта. Биты, которыми представляется число, делятся на три части:

• Знак: первый бит определяет знак числа
• Экспонента: задает порядок величины
• Мантисса: задает значащие цифры

Принцип floating-point чисел легко проиллюстрировать через научную запись, например - 5.734 \times 10^{7}, где сначала идет sign, затем mantissa и exponent. Благодаря этому числа можно представлять в очень широком диапазоне величин с адаптивной точностью. Хотя float32 является форматом по умолчанию, в PyTorch доступен целый набор floating-point форматов:

Сокращение общего числа бит дается не бесплатно, но мы в какой-то мере можем выбирать, чем за это платить: можно жертвовать битами либо в mantissa, либо в exponent. Именно поэтому существуют два формата float8 (FP8), названные по своим exponent и mantissa, чтобы можно было гибко выбирать наиболее подходящий вариант. Посмотрим на диапазон чисел, доступный для каждого формата:

Видно, что float32 покрывает 80 порядков величины, тогда как float16 сильно жертвует диапазоном, а bfloat16 сохраняет его полностью. Два формата float8 сокращают диапазон еще сильнее: e5m2 способен удерживать диапазон float16, а у e4m3 диапазон еще меньше.

Почему же одни форматы способны сохранять полный диапазон, а другие нет? Давайте посмотрим на их разрешение, изобразив 10 000 точек между 1 и 2. Каждая точка будет округлена к ближайшему представимому числу в каждом формате:

Здесь видно, что хотя bfloat16 и сохраняет диапазон float32, в отличие от float16, он делает это ценой более низкой точности. В случае float8 ситуация еще жестче: e4m3 способен представить только 7 чисел, а e5m2 только 3 числа на интервале [1,2].

Обычная метрика для измерения разрешения формата - это epsilon: первое представимое число после 1.00. Видно, что для float32 верхней оценкой будет 10^{-4}, хотя реальное значение равно 1.19^{-7}. Для float16 это примерно 10^{-3}, а для bfloat еще в 10 раз больше.

Идея mixed precision training состоит в том, чтобы использовать некоторые из этих низкоточных форматов для отдельных вычислений, при этом сохраняя качество обучения уровня full precision. На практике полностью отказаться от float32 нельзя, и часть вычислений обычно все равно приходится выполнять в полной точности.

Теперь посмотрим на обучение моделей в 16 битах, а затем разберемся, можно ли пойти еще дальше - вплоть до 8 бит.

Обучение в FP16 и BF16

К сожалению, наивный перевод всех tensors и операций в float16 не работает: результатом обычно становятся расходящиеся loss. Однако в оригинальной статье про mixed precision training были предложены три приема, позволяющие приблизить поведение к float32 training:

1. FP32 copy of weights: С FP16 weights возможны две проблемы. Во время обучения некоторые веса могут становиться очень маленькими и округляться до 0. Но даже если сами weights не близки к 0, очень маленькие updates могут из-за разницы порядков величины уходить в underflow при сложении. Как только weights стали равны 0, они уже останутся нулями до конца обучения, потому что через них больше не проходит gradient signal.
2. Loss scaling: Похожая проблема возникает и с gradients, поскольку они часто заметно меньше 1 и потому подвержены underflow. Простая и эффективная стратегия - масштабировать loss перед backward pass и затем снимать масштаб с gradients после backward pass. Это гарантирует отсутствие underflow во время backward, а на само обучение не влияет, потому что scale снимается до дальнейшей обработки gradients, например clipping, и до optimizer step.
3. Accumulation: Наконец, при некоторых арифметических операциях в 16-битной точности, таких как усреднение или суммирование, тоже возможны underflow и overflow. Решение - накапливать промежуточные результаты в FP32 и только в конце приводить итог обратно к 16-битной точности.

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

И естественно, как любопытный читатель, уже, возможно, слегка зависимый от максимизации throughput, вы можете спросить: можно ли пойти еще дальше и быстрее, чем 16-битная точность?

Возможно.

Предобучение в FP8

Даже если нам удается идеально перекрывать communication и computation, в какой-то момент мы все равно упираемся в низкоуровневый теоретический предел FLOPS самого железа, то есть в эффективность каждой отдельной операции на оборудовании. Здесь критически важной становится численная точность. Например, на GPU NVIDIA H100 матричные умножения в FP8, то есть GEMM operations, достигают вдвое большего теоретического FLOPS, чем BF16, что делает обучение в пониженной точности очень привлекательным направлением дальнейшей оптимизации.

Недавние работы, включая FP8-LM, torchao и DeepSeek-V3, показали потенциал FP8 training для large-scale моделей. Однако FP8 pretraining вносит серьезную проблему: stability. При пониженной точности численная нестабильность часто приводит к расходимости loss, из-за чего трудно добиться точности уровня более высокоточного обучения.

Мы знаем, что при фиксированном размере модели нестабильность растет с увеличением learning rate, поэтому FP8 pretraining оказывается особенно непростой задачей.

Вот пример типичной расходящейся loss curve для FP8 training:

О первом публично описанном успешном very large scale training в FP8 mixed precision было сообщено в техническом отчете DeepSeek-V3. Авторы внимательно проанализировали каждую операцию forward pass (Fprop), а также backward pass для activations (Dgrad) и weights (Wgrad). Как и в BF16 mixed precision training, некоторые aggregations и master weights оставляются в более высокой точности, тогда как сами операции выполняются в FP8.

Чтобы перейти от высокой точности, например FP32 или BF16, к более низкой, например FP16 или FP8, с меньшим диапазоном, нужно нормализовать диапазон значений activations, например вычисляя их абсолютный максимум. В DeepSeek-V3 была дополнительно предложена особая схема quantization, где диапазоны нормализуются по tile: 1x128 для inputs/activations и 128x128 для weights и scale elements. Благодаря этому нормализация меньше страдает от outlier values в activations. Авторы также предложили несколько дополнительных приемов для дальнейшего уменьшения memory и communication footprint; о них можно прочитать в разделе 3.3 технического отчета.

Вот краткая сводка нескольких известных подходов к FP8 training:

В целом FP8 по состоянию на начало 2025 года остается экспериментальной техникой, и методы все еще быстро развиваются. Но с учетом очевидных преимуществ вполне вероятно, что со временем FP8 станет стандартом и вскоре заменит BF16 mixed precision. Если хотите посмотреть open source-реализацию техник FP8 training, обратитесь к этому PR в Nanotron.

Если заглянуть чуть дальше в будущее, то Blackwell, следующее поколение чипов NVIDIA, уже анонсировано с поддержкой FP4 training, что обещает еще сильнее ускорить обучение, но почти наверняка принесет и новый вызов для training stability.

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

Заключение

Поздравляем, вы добрались до конца. Мы прошли большой путь: начали с того, как обучать простую модель на одном GPU, и дошли до освоения сложных техник, которые используются для эффективного обучения огромных языковых моделей вроде Llama-405B и DeepSeek-V3 на тысячах GPU. Теперь вы уже можете с относительной легкостью читать диаграммы вроде 4D parallel setup у Llama-3:

Оркестрация больших GPU-кластеров для эффективного обучения LLM - непростая задача, но теперь вы увидели, как оптимизировать computations и communications между GPU так, чтобы они работали с максимальной utilization. Как мы убедились, это требует выбирать правильную стратегию parallelization для конкретной модели и размера кластера, перекрывать communication и computation там, где это возможно, и писать custom kernels, учитывающие устройство железа, чтобы выполнять операции на GPU максимально быстро.

Возможно, вам все еще кажется, что это довольно нишевое знание и касается лишь узкого круга людей, занимающихся pretraining LLM. Исторически это, возможно, и было верно, но по мере того как быстро растут и сообщество AI builders, и размеры моделей, столь же быстро растет и число людей, использующих distributed techniques для inference, fine-tuning и training. Из-за этого distributed training setups становятся все более обычным делом, а значит, углубление в эту область может оказаться очень своевременным.

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

Что дальше?

Теперь у вас есть хорошее общее представление об основных концепциях distributed training, но по многим инструментам и техникам мы лишь слегка коснулись поверхности. Углубляться дальше можно множеством способов, но вот несколько шагов, которые мы бы рекомендовали:

• Внимательно прочитать несколько ключевых или совсем свежих статей. Большой список наиболее важных работ, блог-постов и книг, которые мы знаем, вы найдете в разделе References.
• Попробовать реализовать какой-нибудь алгоритм с нуля самостоятельно. Очень часто метод по-настоящему "щелкает" только после собственной реализации.
• Погрузиться в один из широко используемых framework и начать вносить вклад: исправлять баги, отвечать на issues или реализовывать новые функции. Это лучший способ войти в любую область ML.

Мы надеемся, что эта книга поможет вам уверенно начать работать с distributed training и что именно вы будете обучать следующее поколение выдающихся моделей под ровный гул вашего GPU cluster. Пусть сила open source и open science всегда будет с вами.

Благодарности

Мы благодарим Elie за тщательные ревью и создание аудиокомпонентов с помощью NotebookLM. Отдельное спасибо Hynek за оптимизацию производительности frontend. Также благодарим Simon за решение ряда проблем на hub.

Страница обсуждения

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

Ссылки

Ключевые статьи по масштабированию LLM

Фреймворки для обучения

Отладка

Техники распределения

Железо

Прочее

Приложение

A0: Краткий курс по параллельному программированию

На протяжении этой книги мы масштабировали обучение LLM от одного GPU до сотен GPU. Для этого нужны обмен данными и синхронизация weights, gradients и самих данных между всеми машинами. Существует набор distributed patterns, которые как раз это и обеспечивают; они называются collective operations. В этом разделе мы коротко разберем такие операции, как Broadcast, AllReduce, Scatter и другие.

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

Например, нам может понадобиться отправить результат с одного node всем остальным nodes или суммировать промежуточные результаты со всех nodes, чтобы получить итог. Обычно есть один node с особым статусом, играющий центральную роль; здесь он обозначен как root и служит источником или целью для некоторых операций. Начнем с одного из самых простых примитивов: операции Broadcast.

Broadcast

Очень распространенный паттерн - когда на одном node есть данные и вы хотите поделиться ими со всеми остальными nodes, чтобы те могли что-то на них посчитать. Именно это и делает операция Broadcast:

Collective operations встроены в PyTorch, поэтому мы легко можем написать маленький пример, показывающий, как работает broadcast. Сначала нужно инициализировать process group с помощью dist.initi_process_group, что настраивает communication backend; о NCCL мы поговорим позже. Затем определяется, сколько существует workers, то есть nodes, и каждому назначается rank, который можно получить через dist.get_rank. Наконец, между workers устанавливается соединение.

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

Этот скрипт можно запустить с помощью torchrun --nproc_per_node=3 dist_op.py. Для этого понадобятся три GPU, либо измените nproc_per_node соответствующим образом. На выходе вы должны увидеть следующее:

Отлично, все работает как ожидается. Обратите внимание, что сообщения от разных rank могут печататься не по порядку, потому что мы не контролируем, какой print выполнится первым; здесь они просто приведены в удобном для чтения порядке. Теперь перейдем к паттернам Reduce и AllReduce.

Reduce & AllReduce

Паттерны Reduce относятся к самым базовым конструкциям распределенной обработки данных. Идея в том, что вы хотите объединить данные, находящиеся на каждом node, с помощью функции f(), которая, например, может выполнять суммирование или усреднение. В парадигме Reduce результат отправляется только на root-node, тогда как в случае AllReduce он рассылается всем nodes:

Разумеется, никакого магического "летающего" node, который сам по себе выполняет эту операцию, не существует. Обычно каждый node делает часть работы, а сами nodes организованы, например, в кольцо или дерево. Простейший пример: пусть нам нужно просуммировать числа, расположенные на разных nodes, соединенных в ring-топологию. Первый node отправляет свое число соседу, тот прибавляет к нему свое и пересылает дальше. После одного полного прохода по кольцу первый node получает итоговую сумму.

Вот код для простой операции Reduce, суммирующей tensors. Операция задается через op=dist.ReduceOp.SUM; подробнее о поддерживаемых операциях можно прочитать в документации PyTorch:

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

Аналогично можно выполнить AllReduce. В этом случае destination указывать не нужно:

В этом случае результат оказывается доступен на всех nodes:

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

Gather & AllGather

Gather и AllGather во многом похожи на Broadcast: они тоже позволяют распределять данные между nodes без их изменения. Главное отличие в том, что здесь нет одного значения, которое нужно разослать с одного node всем остальным. Вместо этого у каждого node есть свой собственный фрагмент данных, и мы хотим либо собрать все данные на одном node, как в Gather, либо собрать их на всех nodes, как в AllGather. Картинка здесь действительно говорит лучше слов:

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

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

Как и ожидалось, gather_list действительно содержит tensors всех ranks:

Чтобы получить пример с AllGather, достаточно поменять только одно: теперь placeholder для результатов нужен на каждом node:

Здесь уже видно, что каждый node теперь получил весь набор данных:

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

Scatter & ReduceScatter

Как и следует из названия, задача Scatter - взять данные на одном node и "рассыпать" их по всем nodes, то есть раздать каждому свой срез данных. Поэтому Scatter отличается от Broadcast, где каждый node получает полную копию без разбиения, и логически является обратной операцией к Gather.

Паттерн ReduceScatter устроен чуть сложнее. Как и в случае AllReduce, вы сначала применяете операцию к данным со всех nodes. Но в отличие от AllReduce, где каждый node получает полный выходной tensor, в ReduceScatter каждый node получает только свой фрагмент выходного tensor. Следующая схема показывает разницу между этими операциями:

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

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

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

Вывод print наглядно показывает созданный нами паттерн данных. Одновременно хорошо видно и сам ReduceScatter в действии: первый rank получил сумму первых tensors со всех nodes, второй rank - сумму вторых tensors и так далее:

Теперь коротко посмотрим на распространенную реализацию AllReduce, которая строится из ReduceScatter и AllGather: Ring AllReduce.

Ring AllReduce

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

1. ReduceScatter
• Каждое устройство делит свои данные, например gradients, на N фрагментов, где N - число GPU, и отправляет один фрагмент соседу. Одновременно оно получает фрагмент от другого соседа.
• Получив фрагмент, устройство добавляет к нему свой соответствующий фрагмент, то есть выполняет reduce.
• Этот процесс продолжается по кольцу, пока у каждого устройства не окажется полностью редуцированный фрагмент, представляющий сумму gradients по всем устройствам для данного куска.
2. AllGather
• Теперь каждому устройству нужно собрать у себя полностью редуцированные фрагменты с остальных устройств.
• Каждое устройство отправляет свой редуцированный фрагмент соседу и получает редуцированный фрагмент от другого соседа.
• Устройства продолжают пересылать полученные фрагменты по кольцу, пока у каждого не окажется полный набор редуцированных кусков, то есть полный суммарный gradient.

Проиллюстрируем это на следующих gif, где у нас есть 5 GPU, и на каждом лежит tensor длины 5. Первая анимация показывает шаг ReduceScatter: в конце каждый GPU получает редуцированный результат для своего конкретного фрагмента данных, отмеченного оранжевым прямоугольником.

Следующая анимация показывает шаг AllGather, после которого каждый GPU получает полный результат операции AllReduce:

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

Для AllReduce полезно держать в голове две вещи:

1. Communication cost для AllReduce при большом N, то есть числе GPU, составляет примерно 2 \times K.
2. Операцию AllReduce можно разложить на ReduceScatter и затем AllGather. Communication cost этих двух операций в сумме вдвое меньше, то есть примерно K.

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

Теперь вы увидели основные строительные блоки distributed operations. Но прежде чем смотреть, как они работают в реальном обучении, стоит познакомиться еще с одной специальной операцией синхронизации: Barrier.

Barrier

Barrier - это простая операция для синхронизации всех nodes. Барьер не снимается, пока его не достигнут все nodes. Только после этого им разрешается продолжать дальнейшие вычисления:

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

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

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

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

NCCL

При обучении больших моделей на множестве GPU иногда удается наткнуться на золото, но на никель, то есть NCCL, вы точно наткнетесь всегда. Так что же это такое?

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

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

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

A1: Профилирование распределенного обучения

Kernels

Для начала предположим, что нужные нам kernels уже встроены в PyTorch. В качестве простого примера возьмем функцию layer normalization, реализованную в PyTorch как torch.nn.functional.layer_norm. Существует несколько способов профилировать kernel, лежащий под этой функцией. Самый прямолинейный - использовать модуль Python time. Но поскольку CUDA-операции асинхронны, такой замер будет учитывать в основном накладные расходы на запуск kernel из Python, а не реальное время выполнения самого kernel.

Чтобы корректно измерить время, можно использовать torch.cuda.Event и явно вызвать torch.cuda.synchronize(), чтобы дождаться завершения kernel. Это показано в следующем примере:

Более эффективный способ профилирования - использовать PyTorch profiler, как мы уже обсуждали выше. Например, можно взять такой код:

Он выведет агрегированные результаты профилирования, отсортированные по суммарному времени CUDA, и результат будет выглядеть примерно так:

Как мы уже упоминали, trace можно отдельно открыть и изучить в chrome://tracing/.

Если приблизить картину, можно увидеть поток операций, который возникает при вызове layer_norm:

Последовательность начинается на CPU, в верхней части trace, с aten::layer_norm, затем переходит в aten::native_layer_norm, после чего управление доходит до cudaLaunchKernel. Дальше исполнение перемещается на GPU, где вызывается kernel vectorized_layer_norm_kernel.

PyTorch profiler дает быстрый обзор производительности, но NVIDIA Nsight Compute CLI, то есть ncu, позволяет заглянуть гораздо глубже. Он показывает подробные времена выполнения и использование памяти для каждого kernel. Запуск выглядит очень просто:

Здесь layer_norm.py - это простой файл, который вызывает функцию layer normalization. Команда сгенерирует лог, но для удобной визуализации результатов лучше сразу указать выходной файл:

Если затем открыть файл output.ncu-rep в Nsight Compute, вы увидите примерно такую картину: предупреждения об утилизации compute и памяти, а также подсказки о том, как улучшить баланс между ними и добиться более высокой occupancy:

CPP extension

Если kernel, который вы хотите профилировать, еще не встроен в PyTorch, можно воспользоваться модулем cpp_extension, чтобы легко скомпилировать и запустить custom CUDA-код. Процесс довольно прямой: создаете свой CUDA-kernel в файле .cu, а затем загружаете его в Python через функцию load из модуля cpp_extension.

Для простого kernel сложения add файл .cu может выглядеть так:

А вот Python-файл, который этот kernel загружает:

С этим подходом custom CUDA-kernel можно профилировать точно так же, как мы только что делали это с помощью PyTorch profiler или инструментов NVIDIA.

A2: Типичные масштабы в обучении LLM

Теперь полезно прикинуть типичные порядки величин в обучении LLM. Когда мы говорим о памяти или compute, мы часто считаем просто число "элементов", то есть отдельных чисел в tensors. Чтобы получить реальный объем памяти в байтах, это число нужно умножить на размер одного значения, например 2 байта для BF16 или 4 байта для FP32.

Вот несколько быстрых ориентиров:

• Входные tokens: В каждом batch обрабатывается seq \cdot mbs tokens, где mbs - micro-batch size, а seq - длина последовательности.
• Activations (hidden states): Для одного слоя tensor hidden state имеет размер seq \cdot mbs \cdot h элементов.
• Веса модели и gradients: Каждая матрица весов в модели, например linear layer, содержит примерно h^2 элементов. Размер gradients совпадает с размером весов.
• Состояния optimizer: Для каждой матрицы весов из h^2 элементов optimizer вроде Adam при mixed precision training хранит моменты и дисперсии в FP32, то есть 2 \times 2 h^2, а также master weights в FP32, то есть еще 2 h^2. В сумме получается порядка 6 h^2 optimizer states на одну матрицу весов.
• Общее число параметров модели: Каждый transformer block содержит:
  • Параметры attention:
    • QKV projections: 3h^2 параметров
    • Выходная projection: h^2 параметров
  • Параметры MLP с Gated Linear Units (GLU):
    • Gate и up projections: 8h^2 параметров, то есть 2 матрицы размера h \times 4h
    • Down projection: 4h^2 параметров, то есть 1 матрица размера 4h \times h
  • Итого на один блок: 16h^2 с GLU MLP или 12h^2 без GLU
  • Для полной модели: 16h^2 \cdot num\_layers при наличии GLU
  • Дополнительные параметры:
    • Входные embeddings: vocab\_size \cdot h
    • LM head: vocab\_size \cdot h, если он не связан с входными embeddings
    • Позиционные embeddings, если они используются: max\_seq\_len \cdot h
• Compute в forward и backward pass (FLOPS): Очень грубая оценка FLOPS для forward pass равна 2 \cdot num\_tokens \cdot num\_params. Для backward pass compute будет вдвое больше: 4 \cdot num\_tokens \cdot num\_params.

📝 Примечание

Более точная формула FLOPs для пары forward+backward выглядит так: 6 \cdot seq\_len \cdot num\_params + 12 \cdot num\_layers \cdot h \cdot seq\_len^2. Она учитывает квадратичное масштабирование attention по всей последовательности, но для упрощения математики мы предполагаем, что seq\_len^2 << h.

A3: Математика overlap между compute и communication

Используя формулы из предыдущего раздела, можно оценить, в каких режимах вычисления и communication действительно могут перекрываться в распределенном обучении. В качестве первого примера рассмотрим data parallelism, то есть ZeRO-0.

Анализ communication в data parallelism

Полный размер gradients, которые нужно передать, равен:

• Gradients = Parameters ≈ num\_layers \cdot 16h^2

Во время backward pass эти gradients передаются bucket-ами, по умолчанию размером 25 MB. Время communication для all-reduce одного bucket равно:

Время вычислений для backward pass равно:

Чтобы overlap был эффективным, требуется:

Это отношение помогает понять, станет ли communication bottleneck во время обучения. Когда оно меньше 1, communication можно полностью спрятать за compute.

Анализ communication в ZeRO-3 (FSDP)

В ZeRO-3 параметры и gradients шардируются между GPU. Разберем паттерн communication для модели, где каждый transformer block содержит 16h^2 параметров:

• Для каждого transformer block в forward pass:
  • All-gather параметров: 16h^2/DP байт на rank
• Для каждого transformer block в backward pass:
  • All-gather параметров: 16h^2/DP байт на rank
  • Reduce-scatter gradients: 16h^2/DP байт на rank
• Итоговый communication на блок: 3 \cdot 16h^2/DP байт
• Итоговый communication для полной модели: 3 \cdot num\_layers \cdot 16h^2/DP байт

Время communication для all-gather операций равно:

Время compute для forward pass одного decoder layer равно:

Чтобы overlap между compute и communication был эффективным, нужно:

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

Анализ communication в TP

В tensor parallelism activations шардируются между GPU внутри TP-регионов, например внутри MLP-блока. Разберем соответствующий паттерн communication:

• Для каждой column-linear операции в forward pass:
  • All-gather activations: seq \cdot mbs \cdot h/TP байт на rank
• Для каждой column-linear операции в backward pass:
  • Reduce-scatter gradients: seq \cdot mbs \cdot h/TP байт на rank
• Для row-linear операций картина обратная. В каждом transformer block есть 2 column-linear и 2 row-linear операции.
• Итоговый communication на блок: 8 \cdot seq \cdot mbs \cdot h/TP байт
• Итоговый communication для полной модели: 8 \cdot num\_layers \cdot seq \cdot mbs \cdot h/TP байт

Возьмем TP-регион внутри слоя и проверим, можно ли перекрыть all-gather communication вычислением следующего linear layer. Время communication для all-gather операций равно:

А время compute для следующего linear layer, содержащего h^2 параметров, равно:

Для эффективного overlap хотим, чтобы communication занимала меньше времени, чем compute:

Это отношение показывает, удастся ли скрыть all-gather communication за вычислением следующего linear layer. Любопытно, что оно зависит только от hidden size h и степени tensor parallelism tp, но не зависит ни от длины последовательности, ни от batch size.

Анализ communication в PP

В pipeline parallelism activations и gradients передаются между pipeline stages. Разберем соответствующий паттерн communication:

• Для каждого micro-batch в forward pass:
  • Получение и отправка activations: 2 \cdot seq \cdot mbs \cdot h байт
• Для каждого micro-batch в backward pass:
  • Получение и отправка gradients: 2 \cdot seq \cdot mbs \cdot h байт
• Итоговый communication на один micro-batch: 4 \cdot seq \cdot mbs \cdot h байт
• Для шагов gradient accumulation, то есть gas, суммарный communication равен 4 \cdot gas \cdot seq \cdot mbs \cdot h байт

Проверим, можно ли перекрыть communication activations и gradients вычислением следующего transformer block. Время compute для transformer blocks на следующей pipeline stage равно:

А время communication для P2P-передачи равно:

Для эффективного overlap хотим:

Как и в случае TP, это отношение не зависит от длины последовательности и batch size. Оно определяется hidden size h, числом слоев на следующей pipeline stage и отношением compute-мощности к P2P-bandwidth данного железа.