Масштабирование LLM с помощью мульти-GPU и мультиузельного параллелизма

Автоматический перевод Эта статья была автоматически переведена с оригинальной английской версии.

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

Это практический разбор основных стратегий параллелизма и ограничений, стоящих за ними, подготовленный с опорой на Ultra-Scale Playbook от Hugging Face. Цель — показать, что даёт каждое разбиение, какими данными обмениваются ранги и когда становятся необходимы комбинации стратегий.

Кратко. Реплицированный параллелизм по данным повышает throughput обучения, когда одна реплика помещается в память. Полностью шардированный параллелизм по данным партиционирует состояние модели, но добавляет all-gather параметров и reduce-scatter градиентов. Тензорный, конвейерный, контекстный и экспертный параллелизм разделяют математику слоёв, глубину модели, последовательность и слои mixture-of-experts (MoE). Комбинируйте их только после того, как определите ограничение, связанное с памятью или коммуникацией.

Предполагается, что вы уверенно работаете с backpropagation, слоями Transformer и стандартным обучающим циклом PyTorch.

Начните с двух бюджетов памяти

У обучения и инференса разный memory footprint.

training peak ≈ parameters
              + gradients
              + optimizer state
              + saved activations
              + temporary buffers
              + communication buffers
              + allocator headroom

inference peak ≈ resident weights
               + key/value (KV) cache
               + runtime workspace
               + communication buffers
               + allocator headroom

У модели со 70 миллиардами параметров нижняя граница только для весов в BF16 составляет 140 ГБ в десятичном исчислении. Для обучения это число мало о чём говорит: градиенты, состояние оптимизатора, master weights и активации могут занимать гораздо больше памяти. Размер памяти для serving также нельзя оценивать только по весам: важны политика кэша, длина последовательности, конкурентность батчей и квантизация.

Профилируйте точную архитектуру, precision, длину последовательности, micro-batch, оптимизатор, политику checkpointing и рантайм. Записывайте пиковую allocated- и reserved-память, токены в секунду, время в ядрах и время, проведённое в collectives.

Масштабирование распределённой системы начинается с измеренного боттлнекаМасштабирование распределённой системы начинается с измеренного боттлнека

Каждое измерение параллелизма создаёт компромисс

Для каждой стратегии ответьте на три вопроса: какое измерение тензора разделяется, какое состояние реплицируется и какой collective оказывается на критическом пути.

СтратегияЧто разделяетсяОсновное облегчениеДобавляемая коммуникация
Реплицированный параллелизм по даннымбатчthroughput обученияgradient all-reduce
Полностью шардированный параллелизм по даннымпараметры, градиенты и состояние оптимизатора между участниками группы data-parallel (DP)память под состояние моделиall-gather параметров, reduce-scatter градиентов
Тензорный параллелизмматричные или attention-измерения внутри слоёввеса и активации слоёвcollectives внутри блоков Transformer
Конвейерный параллелизмгруппы слоёвглубина модели и состояние отдельных стадийpoint-to-point передача активаций и пузыри планирования
Контекстный параллелизмизмерение последовательностипамять под активации длинных последовательностейобмен key/value или attention между участниками sequence-группы
Экспертный параллелизмэксперты MoE и маршрутизируемые токеныёмкость экспертов на каждом рангеdispatch и combine токенов, обычно all-to-all

Уменьшение потребления памяти не является фиксированным множителем. Оно зависит от степени шардинга, того, что остаётся реплицированным, временного несшардированного состояния, политики работы с активациями, padding, дисбаланса и буферов.

Реплицированный параллелизм по данным: throughput без увеличения ёмкости

Реплицированное обучение с параллелизмом по даннымРеплицированное обучение с параллелизмом по данным

Реплицированный параллелизм по данным, обычно называемый distributed data parallel, хранит полную обучающую реплику на каждом ранге. В документации PyTorch по DistributedDataParallel описана эта модель с репликацией и синхронизацией градиентов. Каждый ранг обрабатывает отдельный micro-batch, а градиенты синхронизируются перед шагом оптимизатора.

Используйте этот подход, когда полное состояние обучения помещается в память с безопасным запасом, а global batch можно увеличить или скорректировать gradient accumulation. Его главные преимущества — простая семантика и зрелая реализация, которая перекрывает вычисления backward с редукцией градиентов по бакетам.

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

Полностью шардированный параллелизм по данным: память под состояние в обмен на collectives

Выполнение полностью шардированного параллелизма по даннымВыполнение полностью шардированного параллелизма по данным

Полностью шардированный параллелизм по данным хранит шарды параметров, градиентов и состояния оптимизатора в группе. В статье о ZeRO описан этот паттерн шардинга состояния, а API FSDP2 в PyTorch его реализует. Параметры слоя собираются через all-gather для вычисления, после чего их можно снова зашардировать. Градиенты возвращаются владельцам через reduce-scatter.

В актуальной документации PyTorch различаются API fully_shard в полностью шардированном параллелизме по данным версии 2 (FSDP2) и старый wrapper FullyShardedDataParallel. FSDP2 группирует коммуникацию по модулям, к которым применяется fully_shard, и рекомендует применять его снизу вверх, чтобы группы слоёв могли перекрывать коммуникацию и вычисления.

from torch.distributed.fsdp import fully_shard
from torch.optim import AdamW

# Apply bottom-up: each block becomes a communication group.
for block in model.transformer.blocks:
    fully_shard(block)

# Shard remaining root parameters such as embeddings and output projection.
fully_shard(model)

# Construct the optimizer after parameters have become sharded distributed tensors (DTensors).
optimizer = AdamW(model.parameters(), lr=learning_rate)

Это структурная схема, а не полноценный launcher. Device meshes, mixed precision, checkpointing, инициализация, состояние оптимизатора и распределённые checkpoints должны соответствовать обучающему стеку.

Шардинг привлекателен, когда состояние модели — главное ограничение, а вычисления слоёв способны скрыть достаточный объём collective-трафика. Для небольших моделей, медленных соединений, маленьких слоёв или конфигураций, в которых shard-группа пересекает неподходящую топологическую границу, это может быть невыгодным компромиссом.

Тензорный параллелизм: разделение математики слоёв

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

Тензорный параллелизм партиционирует линейную алгебру внутри слоя. Примеры — column-parallel и row-parallel проекции. В руководстве NVIDIA по параллелизму описано это разбиение на уровне слоёв. Для частичных результатов требуются collectives внутри блоков Transformer, поэтому latency и bandwidth многократно влияют на forward и backward pass.

Используйте его, когда слой или его активации не помещаются в память либо когда матричные умножения достаточно велики, чтобы partitioned kernels превзошли один ранг. Размещайте tensor-parallel группу в самом быстром доступном домене коммуникации, а затем измеряйте результат. При большом tensor-parallel degree каждая локальная матрица может стать настолько маленькой, что эффективность ядер упадёт, а накладные расходы на collectives вырастут.

Sequence parallelism часто комбинируют с tensor parallelism, чтобы не реплицировать часть вычислений над активациями. Это отдельный механизм, не совпадающий с context parallelism по всей входной последовательности модели.

Конвейерный параллелизм: разделение глубины и времени планирования

Стадии конвейерного параллелизма и micro-batchСтадии конвейерного параллелизма и micro-batch

Конвейерный параллелизм размещает разные группы слоёв на отдельных стадиях и передаёт между ними активации. Micro-batch позволяют стадиям работать одновременно. В статье о GPipe этот schedule используется для гигантских нейросетей.

Он уменьшает объём состояния модели на каждой стадии и может снизить объём коммуникации через медленную границу по сравнению с per-layer tensor collectives. Цена — пузыри, передача активаций, дисбаланс стадий, более сложное планирование, а также более трудные recovery и checkpointing.

Для простого сбалансированного schedule в стиле GPipe с p стадиями и m micro-batch идеализированная доля пузыря в forward составляет примерно:

(p - 1) / (m + p - 1)

В реальных schedule могут использоваться one-forward/one-backward, interleaving или zero-bubble варианты, а неравномерная стоимость слоёв может полностью изменить формулу. Определяйте границы стадий по измеренным времени и потреблению памяти, а не по одинаковому количеству слоёв.

Контекстный параллелизм: разделение активаций длинных последовательностей

Обмен attention при контекстном параллелизмеОбмен attention при контекстном параллелизме

Контекстный параллелизм распределяет измерение последовательности. В документации NVIDIA по context parallelism описаны разбиение последовательности и обмен key/value, необходимый для attention. Каждый ранг владеет своим shard последовательности, а attention обменивается данными, необходимыми для сохранения семантики полного контекста. В реализациях могут использоваться point-to-point кольца, all-gather, all-to-all или иерархические комбинации.

Это снижает потребление памяти активациями при обучении на длинном контексте, но реплицирует веса между участниками context-группы и добавляет коммуникацию в attention. Выигрыш зависит от типа attention, causal masking, длины последовательности, recomputation и того, как context-группы сочетаются с tensor- и data-parallel группами.

Не выбирайте эту стратегию по универсальному порогу в 8K, 32K или 100K. Профилируйте память активаций и коммуникацию в attention для фактической архитектуры.

Экспертный параллелизм: только для архитектуры MoE

Маршрутизация токенов при экспертном параллелизмеМаршрутизация токенов при экспертном параллелизме

Экспертный параллелизм распределяет экспертов в слоях mixture-of-experts. В руководстве NVIDIA по параллелизму описаны это размещение экспертов и его сочетание с другими измерениями параллелизма. Роутер отправляет представления токенов выбранным экспертам и объединяет их результаты. Для каждого токена вычисления выполняют только выбранные эксперты, однако суммарные веса экспертов всё равно требуют хранения и размещения при serving.

Экспертный параллелизм не является переключателем оптимизации для dense-модели. Это часть архитектуры MoE. Среди ключевых проблем — балансировка нагрузки, ограничения ёмкости, all-to-all токенов, отброшенные или дополненные padding токены, auxiliary losses и дисбаланс при сбоях. Отслеживайте число токенов на эксперта, энтропию маршрутизации, переполнение ёмкости, время коммуникации и качество в зависимости от маршрута.

Соберите конфигурацию с учётом топологии

Системы обучения больших dense-моделей без expert-parallel группы обычно используют произведение размеров групп data-parallel (DP), tensor-parallel (TP), pipeline-parallel (PP) и context-parallel (CP):

world size = DP × TP × PP × CP

Если expert parallelism (EP) является независимой группой, руководство NVIDIA по параллелизму вычисляет общее число как:

total GPUs = TP × PP × CP × EP × DP

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

Стройте конфигурацию в следующем порядке:

  1. Нарисуйте домены коммуникации: соединения GPU-GPU, коммутаторы, границы non-uniform memory access (NUMA), межузловую сеть, oversubscription и путь к хранилищу.
  2. Разместите частые latency-sensitive collectives, обычно TP, в самом быстром подходящем домене.
  3. Выберите группы FSDP или реплицированного DP с учётом оставшейся ёмкости и bandwidth.
  4. Добавьте PP, если размещение по глубине или трафик между доменами дают выигрыш; сбалансируйте измеренные время и память стадий.
  5. Добавляйте CP только из-за ограничения по последовательности, а expert parallelism (EP) — только при наличии соответствующей экспертной топологии модели.
  6. Проверьте делимость числа heads, hidden dimensions, слоёв, экспертов, batch и sequence для выбранного mesh.
  7. Проведите бенчмарк нескольких допустимых mesh. Эвристики с учётом топологии выбирают кандидатов, но не победителя.

Два кластера с одинаковым числом GPU могут предпочитать разные конфигурации из-за различий в bandwidth каналов, иерархии коммутаторов, подключении CPU и сетевой конкуренции.

Для обучения и serving нужны разные решения

При инференсе обычно нет градиентов и состояния оптимизатора, поэтому обучающие конфигурации в стиле FSDP не переносятся автоматически.

Для serving задайте следующие вопросы:

  • Помещаются ли веса, KV cache, workspace и целевая конкурентность в одну реплику?
  • Что лучше обслуживает throughput: больше независимых реплик или шардинг одной реплики?
  • Снижает ли TP нагрузку на веса и кэш каждого ранга настолько, чтобы окупить коммуникацию на уровне слоёв?
  • Эффективно ли PP поддерживается для этой модели и request scheduler?
  • Как prefill и decode по-разному нагружают вычисления, memory bandwidth и interconnect?
  • Что происходит с tail latency, когда у запросов различаются длины промптов и выходов?

Проводите бенчмарк полного сервера вместе с scheduler, квантизацией, распределением контекста, политикой batching и формой трафика. Число токенов в секунду при обучении не позволяет предсказать time to first token или inter-token latency при serving.

Честно измеряйте масштабирование конфигурации

Для каждого кандидата записывайте:

  • модель, код, рантайм, kernels и идентификатор топологии
  • global и local batch, распределение длины последовательности и число токенов
  • пиковое потребление памяти по категориям, если оно доступно
  • полезные токены в секунду и утилизацию model floating-point operation (FLOP), если она рассчитывается согласованно
  • время, проведённое в all-reduce, all-gather, reduce-scatter, all-to-all и point-to-point операциях
  • задержки при подаче данных, время checkpointing, поведение после перезапуска и распределение straggler
  • loss обучения или parity выходов serving с baseline

Целенаправленно сравнивайте weak и strong scaling. Strong scaling сохраняет общий объём работы постоянным при увеличении числа рангов. При weak scaling объём работы на ранг остаётся постоянным, поэтому общий объём работы растёт вместе с числом рангов. Процент с названием «эффективность масштабирования» бессмысленен без указания знаменателя и baseline.

Заключение

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

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

Ссылки