Перейти к содержанию

Автоматический перевод

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

Масштабирование больших языковых моделей: стратегии multi-GPU и multi-node, которые реально работают на практике

Сегодняшние LLM не помещаются на один GPU. Модели на 70B параметров нужно около 140GB только под веса в FP16 — почти вдвое больше, чем помещается в A100. Обучение или инференс таких моделей означает, что работу нужно разделять между несколькими GPU, и неверное разбиение быстро сжигает большую часть вашего вычислительного бюджета.

Это практический разбор стратегий параллелизма, которые действительно работают в проде, на основе Ultra-Scale Playbook от Hugging Face.

Предварительные требования

Максимум пользы вы получите, если уже уверенно понимаете:

  • Backpropagation, градиенты и оптимизаторы вроде AdamW.
  • Устройство Transformer: attention и feed-forward сети.
  • Базовые вещи в PyTorch: nn.Module, DataLoader, стандартный training loop.

Почему масштабирование важно

Есть четыре причины, по которым одного GPU становится недостаточно. Модели на 70B нужно примерно 140GB в FP16 — почти вдвое больше, чем 80GB у A100. Даже на восьми A100 обучение 13B-модели с нуля занимает недели. Длинные контексты (32K токенов и выше) упираются в память одного GPU еще до того, как вы начнете делать что-то реально полезное. А при продовой нагрузке именно распределенный inference не дает tail latency выйти из-под контроля.

1. Техники параллелизма, простым языком

1.1 Data parallelism (DP)

Самое простое разбиение. Каждый GPU хранит полную модель и обрабатывает свой фрагмент батча. После backprop все GPU делают all-reduce градиентов (усреднение), а затем каждый обновляет свою копию модели. Одинаковые модели, разные данные, синхронизированные веса.

Используйте это, когда модель уверенно помещается на одном GPU и вам нужно просто прогонять больше батчей в секунду. Стоимость настройки почти нулевая: DDP в PyTorch — это обертка в одну строку. Ограничение в памяти никуда не девается: каждый GPU по-прежнему хранит полную модель, состояния оптимизатора и градиенты, так что вы покупаете throughput, а не capacity.

Параллелизм по данным

Инструменты: PyTorch DDP, Horovod.

1.2 Fully sharded data parallelism (FSDP)

DP, но с учетом памяти. Параметры, градиенты и состояния оптимизатора шардингом распределяются между GPU. Во время forward pass каждый GPU собирает у остальных нужные ему параметры, считает, а затем сбрасывает полученные шарды, чтобы освободить память. В backward pass это повторяется, после чего градиенты редуцируются так, чтобы каждый GPU обновлял только свой шард.

Это первый слой, к которому стоит переходить, когда модель уже переросла один GPU (обычно все, что больше ~10B параметров) и вы хотите продолжать обучение на одной машине без переписывания кода. На практике FSDP позволяет обучать модели в 4–8× больше тех, что помещаются на одном GPU.

FSDP

[!NOTE] Стадии ZeRO, кратко FSDP часто описывают в терминах ZeRO (Zero Redundancy Optimizer):

  • Stage 1: шардинг только состояний оптимизатора (~4× экономии памяти).
  • Stage 2: шардинг градиентов + состояний оптимизатора (~8× экономии памяти).
  • Stage 3: шардинг параметров + градиентов + состояний оптимизатора (линейное масштабирование по числу GPU N).

В PyTorch FSDP по умолчанию соответствует поведению Stage 3.

Включение FSDP в PyTorch

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

# 1. Wrap your model
model = MyLLM()
model = FSDP(model)

# 2. Train as usual
output = model(input)
loss = output.sum()
loss.backward()
optimizer.step()

Инструменты: PyTorch FSDP, DeepSpeed ZeRO-3.

1.3 Tensor parallelism (TP)

Разбиение отдельных слоев между GPU. Берете матрицу весов, режете ее по столбцам (или строкам) и отдаете каждому GPU свой кусок. Каждое устройство вычисляет свою часть выхода; затем all-reduce или конкатенация собирает результат обратно перед следующим слоем. Так происходит на каждом слое.

TP имеет смысл, когда отдельные слои слишком большие даже после FSDP — например, огромные attention-матрицы или широкие FFN. Он также предполагает быстрый interconnect внутри узла: NVLink или NVSwitch, а не PCIe. Между узлами all-reduce на каждом слое становится bottleneck, и выигрыш исчезает. Типичная sweet spot — TP степени 2–8 внутри одной машины.

Тензорный параллелизм

Инструменты: Megatron-LM, TensorRT-LLM, ColossalAI.

1.4 Pipeline parallelism (PP)

Разбиение модели по вертикали, на группы слоев. Слои 1–10 живут на GPU 1, слои 11–20 — на GPU 2 и так далее. Затем по pipeline прогоняются micro-batch, чтобы каждая стадия была занята: GPU 1 заканчивает batch 1 и передает его GPU 2, а сам сразу начинает batch 2. Когда в полете достаточно micro-batch, каждое устройство постоянно чем-то занято.

Используйте PP, когда модель настолько глубокая, что даже FSDP не позволяет ее уместить, или когда нужно растянуть вычисления на несколько узлов, и bottleneck уже в межузловой пропускной способности. Неприятная часть — pipeline "bubbles", то есть простои стадий в начале и конце каждого batch. Их уменьшают, запуская много маленьких micro-batch вместо нескольких больших.

Конвейерный параллелизм

Инструменты: DeepSpeed PP, Megatron-LM, GPipe.

1.5 Context parallelism (CP)

Для очень длинных последовательностей. Контекст в 64K токенов делится, например, между четырьмя GPU по 16K токенов на каждый. Каждый GPU выполняет self-attention на своем локальном куске, затем GPU обмениваются keys и values, чтобы вычислить части attention между кусками. Объединенный результат совпадает с тем, что вы получили бы, выполнив весь контекст на одном устройстве, но без такого счета за память.

Это тот рычаг, который нужен, когда bottleneck — не размер модели, а длина контекста: анализ длинных документов, рассуждение на уровне книги, генерация кода по большому репозиторию. Именно CP делает обучение на 100K+ токенов возможным на железе, которое иначе уперлось бы в 8K.

Параллелизм по контексту

Инструменты: Picotron, Nanotron.

1.6 Expert parallelism (Mixture of Experts)

Специализация. Плотные FFN-слои заменяются на N экспертных подсетей (8, 64, иногда больше). Небольшая gating-сеть выбирает top-k экспертов (обычно top-2) для каждого токена. Только эти эксперты обрабатывают токен; остальные простаивают. Разные эксперты могут жить на разных GPU, поэтому модель может быть огромной по общему числу параметров, а вычисления на токен остаются небольшими.

У Mixtral-8x7B суммарно 56B параметров, но на токен активны только ~13B; Grok и DeepSeek-V2 используют тот же прием. Цена — сложность на стороне обучения: балансировка нагрузки между экспертами сама по себе отдельная инженерная задача, а нестабильность роутинга уже не раз приводила к расходимости в MoE-обучении.

Mixture of Experts

Инструменты: Picotron, Nanotron.

Быстрое сравнение: какой параллелизм использовать?

Technique What it splits Best for Memory savings Communication cost
Data Parallelism (DP) Data batches Models that fit on 1 GPU None (copies model) Low (only gradients)
FSDP Model + optimizer + gradients Models too big for 1 GPU High (4–8×) Medium
Tensor Parallelism (TP) Individual layers Huge layers, fast GPUs Medium High (per layer)
Pipeline Parallelism (PP) Layer groups (stages) Very deep models Medium Low (between stages)
Context Parallelism (CP) Sequence length Long contexts (64K+ tokens) High (for activations) Medium
Expert Parallelism (MoE) Experts in MoE layers Massive sparse models None (more params, less FLOPs) Medium

Разумный вариант по умолчанию: начинайте с FSDP. Добавляйте TP, если отдельные слои все еще слишком большие. Добавляйте PP, когда нужно масштабироваться на несколько узлов. Добавляйте CP, когда bottleneck — длина контекста.

2. Практические стратегии обучения

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

2.1 Одна машина, 2–8 GPU

Сначала используйте чистый FSDP — PyTorch FSDP или DeepSpeed ZeRO-2/ZeRO-3 — через Hugging Face accelerate или torchrun. Если отдельные attention- или FFN-слои все еще слишком большие даже после шардинга, добавьте TP=2.

Пара замечаний по конкретному железу. Пользовательским GPU (RTX 4090 и подобным) на PCIe стоит держаться TP=1 или максимум TP=2; interconnect не вытягивает больше. Серверные GPU (A100, H100) с NVLink нормально справляются с TP=2 до TP=4. А на восьми GPU в одном сервере чистый FSDP часто справляется с моделями до 70B вообще без TP.

2.2 Небольшой кластер, 2–16 узлов (≤128 GPU)

Здесь нужен 2D или 3D parallelism: TP плюс FSDP, при необходимости еще и PP. Практически рабочая схема такая:

  1. TP внутри каждого узла (TP=4 или TP=8 с NVLink).
  2. FSDP между узлами для data parallelism.
  3. Добавляйте PP, если модель настолько глубокая, что даже FSDP ее не вмещает, разрезая ее по вертикали между узлами.

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

Если вы добавляете PP, ставьте число micro-batch как минимум 4× от степени pipeline. Если меньше, bubbles начинают съедать throughput.

2.3 Большой кластер (сотни и тысячи GPU)

Здесь уже имеет смысл 4D parallelism (DP × TP × PP × CP). Сопоставьте каждое измерение с топологией вашего железа и используйте Megatron-LM или Nanotron — они поддерживают 4D из коробки, а писать все это самостоятельно — отдельный проект.

Реалистично это нужно только тогда, когда вы pretraining 70B+ моделей с окнами контекста 32K+ с нуля. Большинству задач fine-tuning, даже для больших моделей, это не требуется.

Конкретный пример. Обучение 70B-модели с контекстом 32K на 512 GPU:

  • TP=8 внутри каждого 8-GPU узла.
  • PP=4 на четыре узла.
  • CP=4 для длинного контекста.
  • DP=4 для throughput.
  • 8 × 4 × 4 × 4 = 512 GPU.

Эффективность масштабирования на такой конфигурации обычно оказывается в районе 70–80% при хорошем InfiniBand, а при аккуратной настройке можно дойти до ~85%. Все, что ниже, — это уже прямые потери денег.

3. Инструменты, которые стоит изучить

Короткий cheat sheet, чтобы выбрать подходящий инструмент:

Tool When to use Learning curve Best for
Hugging Face Accelerate Any distributed training with minimal code changes ★☆☆☆☆ Beginners, quick prototypes
PyTorch FSDP Medium-to-large models (1–30B) on a single node ★★☆☆☆ The common case
DeepSpeed ZeRO Multi-node training with good documentation ★★★☆☆ Production training
Megatron-LM Very large models (70B+), 3D/4D parallelism ★★★★☆ Production at scale
Nanotron Learning and research on modern parallelism ★★★☆☆ Education, experimentation
vLLM Fast inference with PagedAttention and KV caching ★★☆☆☆ Serving in production
TensorRT-LLM Maximum inference speed on NVIDIA GPUs ★★★★☆ Production inference optimization

Минимальная FSDP-конфигурация для accelerate

compute_environment: LOCAL_MACHINE
distributed_type: FSDP
fsdp_config:
    fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
    fsdp_backward_prefetch: BACKWARD_PRE
    fsdp_state_dict_type: SHARDED_STATE_DICT
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
use_cpu: false

Если вы только начинаете, я бы выбрал Hugging Face Accelerate, чтобы быстро что-то запустить, а затем перешел бы на PyTorch FSDP или DeepSpeed, когда понадобится более тонкий контроль.

4. Фреймворк для принятия решений

Дерево решений по масштабированию

Схема выше отражает то, как я бы действовал на практике: сначала FSDP, затем TP, если слои слишком большие, PP — для глубины на нескольких узлах, CP — для длинного контекста. Усложняйте систему только после того, как более простой подход действительно перестал справляться.

5. Ultra-Scale cheat sheet

Команда Hugging Face собрала одностраничную визуальную выжимку, которая покрывает большую часть того, что описано выше:

Шпаргалка по Ultra-Scale LLM

Заключение

FSDP закрывает большую часть того, с чем вы столкнетесь. TP помогает внутри узла, когда отдельные слои все еще не помещаются. PP распределяет модель между узлами, когда ограничение — глубина. CP нужен, когда память заканчивается из-за длины контекста. Общий принцип для всех этих подходов один: подбирайте стратегию параллелизма под реальную топологию доступного железа и добавляйте новое измерение только тогда, когда более простое уже исчерпало себя.

Дополнительное чтение: