DeepSpeed для LLM: ZeRO, offloading, мульти-GPU
Обучение LLM с числом параметров >10B упирается в лимиты VRAM. Даже на 8× A100 80GB модель 30B не помещается без шардирования. DeepSpeed от Microsoft — стандарт индустрии для распределённого обучения: ZeRO снижает потребление памяти в 3–8 раз, а offloading на CPU/NVMe позволяет обучать модели до 175B на ограниченном кластере. В этой статье разберём, как подобрать конфигурацию под вашу задачу.
Проблемы, которые решаем
Переполнение VRAM. LLM с 30B+ параметров не помещается даже на 8× A100 80GB. ZeRO Stage 3 шардирует не только оптимизатор, но и параметры, градиенты — экономия до 3x памяти. Коммуникационные узкие места: all-reduce в DDP блокирует GPU. DeepSpeed использует overlap коммуникации и reduce-scatter — p99 latency падает на 40%. Сложность конфигурации: неверный выбор stage или batch size приводит к OOM или низкой утилизации GPU. Мы подбираем параметры под ваше железо и модель.
Как мы это делаем
Мы используем итоговые конфиги, протестированные на кластерах A100/H100. Пример для ZeRO Stage 2 с BF16:
{ "zero_optimization": { "stage": 2, "allgather_partitions": true, "allgather_bucket_size": 2e8, "overlap_comm": true, "reduce_scatter": true, "reduce_bucket_size": 2e8, "contiguous_gradients": true }, "fp16": { "enabled": true, "loss_scale": 0, "loss_scale_window": 1000, "initial_scale_power": 16, "hysteresis": 2, "min_loss_scale": 1 }, "bf16": { "enabled": false }, "gradient_accumulation_steps": 4, "gradient_clipping": 1.0, "train_batch_size": "auto", "train_micro_batch_size_per_gpu": 4, "wall_clock_breakdown": false } Интеграция с Hugging Face Transformers — стандартный Trainer:
from transformers import TrainingArguments, Trainer training_args = TrainingArguments( output_dir="./results", deepspeed="ds_config_zero2.json", per_device_train_batch_size=4, gradient_accumulation_steps=4, fp16=True, num_train_epochs=3, logging_steps=100, save_steps=1000, ) trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, tokenizer=tokenizer, ) trainer.train() Запуск:
deeepspeed --num_gpus=8 train.py --deepspeed ds_config.json # Или через torchrun: torchrun --nproc_per_node=8 train.py --deepspeed ds_config.json Как ZeRO Stage 2 ускоряет обучение?
Стадия 2 шардирует состояние оптимизатора и градиенты — на 8 GPU память на каждый узел снижается в 8x. Коммуникация через reduce-scatter вместо all-reduce сокращает объём пересылаемых данных на 30%. В итоге throughput на LLaMA 7B достигает 7000 tokens/s на 8× A100. Это в 4 раза быстрее стандартной DDP.
Сравнение ZeRO stages
| Stage | Шардирование | Память на GPU | Коммуникация | Когда использовать |
|---|---|---|---|---|
| 1 | Оптимизатор | Умеренная | Стандартная | Модели <6B |
| 2 | Оптимизатор + градиенты | Низкая | Reduce-scatter | Модели <30B |
| 3 | Все (параметры, градиенты, оптимизатор) | Минимальная | All-gather | Модели >30B |
Практические показатели
| Конфигурация | Модель | Кластер | Throughput |
|---|---|---|---|
| ZeRO-2, BF16 | LLaMA 7B | 8× A100 80GB | ~7000 tokens/s |
| ZeRO-2, BF16 | LLaMA 13B | 8× A100 80GB | ~3500 tokens/s |
| ZeRO-3, BF16 | LLaMA 30B | 8× A100 80GB | ~1200 tokens/s |
| ZeRO-3 + Offload | LLaMA 65B | 8× A100 80GB + 512GB RAM | ~400 tokens/s |
Когда offloading оправдан?
Offloading на CPU через NVMe добавляет latency, но позволяет обучать модели 65B+ на 8 картах. Если VRAM хватает, лучше обойтись Stage 2 или Stage 3 без offload. Типичный кейс: кластер с 8× A100 40GB для модели 30B — Stage 3 с offload на CPU даёт стабильные 1200 tokens/s. Без offload модель просто не влезет.
Типичные ошибки при настройке DeepSpeed
- Неправильный выбор stage: Stage 3 на малых моделях даёт оверхед из-за all-gather.
- Слишком большой micro batch: OOM, хотя конфиг правильный. Рекомендуем начинать с 1 и увеличивать.
- Игнорирование gradient accumulation: для batch size 256 на 8 GPU достаточно 32 шагов накопления.
- Отсутствие overlap коммуникации: снижает утилизацию GPU. Включаем
overlap_comm: true.
Как настроить DeepSpeed для вашей модели?
Процесс работы включает пять этапов. Аналитика: выясняем вашу модель, железо, требования к скорости и памяти. Проектирование: подбираем stage, batch size, градиентные шаги. Реализация: пишем конфиг, интегрируем с вашим кодом. Тест: проверяем на небольшом датасете, снимаем метрики. Деплой: разворачиваем на кластере, оптимизируем при необходимости.
Для предварительной оценки памяти используем встроенный инструмент DeepSpeed:
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf") estimate_zero3_model_states_mem_needs_all_live( model, num_gpus_per_node=8, num_nodes=1 ) Что входит в работу и сроки
В рамках настройки вы получаете: готовый конфиг DeepSpeed под вашу модель и железо; код интеграции с Hugging Face Transformers или другим фреймворком; отчёт с замерами throughput, памяти и советами по масштабированию; консультации в течение 30 дней после внедрения; документацию по запуску и мониторингу. Сроки — от 3 до 10 рабочих дней в зависимости от сложности. Стоимость рассчитывается индивидуально, но обычно проект окупается за счёт сокращения времени обучения на 30-50%.
Мы выполнили более 50 проектов по обучению и дообучению LLM: от GPT-2 до LLaMA 65B, на кластерах от 4 до 64 GPU. Свяжитесь с нами для бесплатной оценки вашего проекта — поможем подобрать оптимальную конфигурацию и ускорить обучение без лишних затрат.