Масштабирование 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 позволяют стадиям работать одновременно. В статье о 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 варианты, а неравномерная стоимость слоёв может полностью изменить формулу. Определяйте границы стадий по измеренным времени и потреблению памяти, а не по одинаковому количеству слоёв.
Контекстный параллелизм: разделение активаций длинных последовательностей
Контекстный параллелизм распределяет измерение последовательности. В документации 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, поддерживаемый вашим фреймворком, вместо перемножения неподдерживаемых конфигураций.
Стройте конфигурацию в следующем порядке:
- Нарисуйте домены коммуникации: соединения GPU-GPU, коммутаторы, границы non-uniform memory access (NUMA), межузловую сеть, oversubscription и путь к хранилищу.
- Разместите частые latency-sensitive collectives, обычно TP, в самом быстром подходящем домене.
- Выберите группы FSDP или реплицированного DP с учётом оставшейся ёмкости и bandwidth.
- Добавьте PP, если размещение по глубине или трафик между доменами дают выигрыш; сбалансируйте измеренные время и память стадий.
- Добавляйте CP только из-за ограничения по последовательности, а expert parallelism (EP) — только при наличии соответствующей экспертной топологии модели.
- Проверьте делимость числа heads, hidden dimensions, слоёв, экспертов, batch и sequence для выбранного mesh.
- Проведите бенчмарк нескольких допустимых 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 и профилируйте их. Лучшая конфигурация — та, которая помещается в память с запасом и минимизирует открытую коммуникацию для конкретной выполняемой задачи.
Ссылки
- Документация PyTorch по FSDP2
fully_shard— шардинг параметров на уровне модулей, all-gather, reduce-scatter и поведение DTensor. - Документация PyTorch по DistributedDataParallel — обучение реплицированной модели с синхронизированными градиентами.
- Стратегии параллелизма NVIDIA Megatron Core — измерения data, tensor, pipeline, context, expert и fully sharded параллелизма.
- Ultra-Scale Playbook от Hugging Face — интерактивный обзор крупномасштабного параллелизма моделей.
- Статья DeepSpeed о ZeRO — шардинг состояния оптимизатора, градиентов и параметров для снижения потребления памяти.
- Статья о GPipe — обучение с pipeline, micro-batch и поэтапным выполнением.
- Документация NERSC по параллелизму — определения strong и weak scaling на основе фиксированного общего объёма работы или объёма работы на ранг.