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. Зв'яжіться з нами для безкоштовної оцінки вашого проєкту — допоможемо підібрати оптимальну конфігурацію та прискорити навчання без зайвих витрат.