Parallelism and memory: почему одна GPU перестает быть единицей мышления
На маленькой модели вопрос звучит просто: “помещается ли batch в память?” На больших моделях это уже не один вопрос. В памяти лежат веса, градиенты, optimizer states, activations, временные buffers и KV/attention-related tensors. Часть этого можно шардировать, часть пересчитывать, часть разрезать по слоям или матрицам. Поэтому distributed training превращается в инженерное планирование: что именно не помещается, какую цену мы готовы заплатить коммуникацией, и где будет следующий bottleneck.
Полезный mental model: parallelism - это не набор “трюков”, а способы распределить память, вычисления и коммуникацию между GPU. Data parallelism распределяет данные, но реплицирует модель. Tensor parallelism режет матричные операции. Pipeline parallelism режет слои. FSDP/ZeRO режут model states. Activation checkpointing не распределяет ничего, а меняет память на дополнительный compute.
Сначала считаем память
Без memory accounting легко выбрать неправильную технологию. Например, если модель не помещается из-за optimizer states, помогает sharding. Если взрываются activations из-за длинного контекста, FSDP может не решить проблему без checkpointing или sequence/context parallelism. Если bottleneck уже не HBM, а interconnect, новый уровень sharding может сделать обучение медленнее.
| Компонент | От чего растет | Что обычно помогает | Важная оговорка |
|---|---|---|---|
| Weights / parameters | Число параметров и precision. | FSDP/ZeRO-3, tensor parallelism, quantized states для некоторых режимов. | Во время compute шардированные параметры часто надо all-gather-ить. |
| Gradients | Число trainable параметров и precision. | DDP buckets, ZeRO-2/3, FSDP reduce-scatter. | Communication может стать bottleneck при маленьком compute per step. |
| Optimizer states | Adam обычно хранит дополнительные моменты; часто это больше, чем сами bf16 weights. | ZeRO-1/2/3, FSDP optimizer state sharding. | SGD/Adam/8-bit optimizers имеют разные memory profiles. |
| Activations | Batch size, sequence length, hidden size, depth, attention pattern. | Activation checkpointing, sequence/context parallelism, меньший micro-batch. | Checkpointing снижает memory, но добавляет recompute и может снизить throughput. |
| Temporary buffers | Kernels, collectives, fragmentation, runtime. | Profiling, allocator tuning, более стабильные shapes. | Эта часть часто всплывает только в реальном запуске. |
Основные виды parallelism
| Подход | Что распределяет | Где силен | Чем платим |
|---|---|---|---|
| Data Parallel / DDP | Данные. Модель полностью копируется на каждый rank. | Простой baseline, хорошо работает когда модель помещается на GPU. | All-reduce gradients и полная репликация weights/optimizer states. |
| FSDP / ZeRO | Parameters, gradients, optimizer states. | Когда model states не помещаются или съедают слишком много HBM. | All-gather/reduce-scatter, сложнее checkpoint/resume, чувствительно к wrapping/sharding policy. |
| Tensor Parallelism | Большие матричные операции внутри слоя. | Очень большие Transformer layers, когда слой сам слишком тяжелый. | Требует быстрый interconnect; часто держат внутри node с NVLink/NVSwitch. |
| Pipeline Parallelism | Слои модели между устройствами. | Когда модель глубоко не помещается целиком и можно разбить stack. | Pipeline bubbles, micro-batches, сложнее балансировать разные stages. |
| Sequence / Context Parallelism | Части sequence/activation-heavy работы. | Длинный контекст, где activations/attention становятся главной болью. | Сложная коммуникация и поддержка не во всех training stacks. |
| Activation Checkpointing | Ничего не распределяет; сохраняет меньше activations. | Когда memory bottleneck в activations. | Recompute во время backward, иногда заметная потеря throughput. |
В реальных training runs редко выбирают ровно один метод. Часто план выглядит как композиция: FSDP или ZeRO для model states, tensor parallelism внутри node для больших матриц, pipeline parallelism по слоям, gradient accumulation для effective batch, activation checkpointing для длинного контекста. Поэтому хороший ответ на собесе не “я бы включил FSDP”, а “я сначала посчитаю память, потом выберу комбинацию под topology и throughput target”.
Почему topology важнее красивой схемы
Communication не одинаковая везде. GPU внутри одной node могут общаться через NVLink/NVSwitch, между nodes - через InfiniBand или RoCE, а storage/checkpoint traffic идет своим путем. Поэтому tensor parallelism, который часто требует частой коммуникации внутри слоя, обычно стараются держать там, где bandwidth максимальный. Pipeline или data parallel traffic легче растянуть между nodes, но там появляются bubbles, stragglers and checkpoint/restart costs.
Как читать production case studies
Мини-decision tree
- Модель помещается, но хочется быстрее. Начни с DDP, проверь data pipeline, batch size, mixed precision and profiler traces.
- Не помещаются optimizer states / gradients. Смотри FSDP или ZeRO stages; сравни memory saving с extra communication.
- Не помещается один layer или матрица слишком большая. Смотри tensor parallelism, но проверь interconnect.
- Модель слишком глубокая для одной GPU/node. Смотри pipeline parallelism, micro-batch count and bubble overhead.
- Длинный context раздувает activations. Смотри activation checkpointing, smaller micro-batch, sequence/context parallelism.
- Все помещается, но throughput плохой. Ищи bottleneck: dataloader, communication overlap, stragglers, checkpoint I/O, kernel efficiency, network.
Что важно уметь объяснить на собеседовании
Сильный кандидат не обязан помнить все флаги Megatron или DeepSpeed, но должен уметь разложить систему. Если тебе дают 70B model, 64 GPU, sequence length 8k and OOM, нужно спросить: precision? optimizer? micro-batch? activation checkpointing? sharding? topology? где именно peak memory? Если тебе говорят, что FSDP стал медленнее DDP, нужно объяснить communication overhead и размер модели/workload, а не удивляться.
Memory budget simulator
Ниже интерактив не пытается дать “правильную конфигурацию”. Он показывает порядок величин: как model size, sequence length, micro-batch, precision, optimizer, DP/TP/PP, FSDP and activation checkpointing двигают HBM per GPU, communication pressure, pipeline bubbles and checkpoint I/O risk. Именно так полезно думать перед запуском реального profiler/benchmark.
Загрузка интерактивного виджета...
Материалы
Официальные docs
Production cases
Papers / reports
Original tensor model parallelism grounding.
Shows how data, tensor and pipeline parallelism are composed at large scale.
Recent large-scale training systems paper with practical scaling and failure context.