DeepSpeed

Что такое DeepSpeed

DeepSpeed — библиотека для распределённого обучения от Microsoft Research, оптимизированная для LLM. Ключевые компоненты: алгоритм ZeRO (Zero Redundancy Optimizer) для устранения дублирования состояния оптимизатора между GPU, поддержка смешанной точности, gradient checkpointing и multi-node training.

Хорошо интегрирован с Hugging Face Transformers через Trainer(deepspeed=...).

ZeRO: три стадии

Стандартное обучение дублирует полную копию модели на каждом GPU. ZeRO устраняет это:

При ZeRO-3 на 8 GPU каждый GPU хранит 1/8 весов, 1/8 градиентов, 1/8 оптимизатора — модель 70B умещается на 8× A100 40GB вместо 8× 80GB.

Быстрый старт

# ds_config.json
{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {"device": "cpu"},  # опционально: выгрузить на CPU
    "allgather_partitions": true,
    "reduce_scatter": true
  },
  "bf16": {"enabled": true},
  "gradient_accumulation_steps": 4,
  "train_micro_batch_size_per_gpu": 2
}

# Запуск через Hugging Face Trainer
from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    deepspeed="ds_config.json",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    bf16=True,
    output_dir="./output"
)
trainer = Trainer(model=model, args=args, ...)

# Запуск на нескольких GPU
deepspeed --num_gpus=8 train.py --deepspeed ds_config.json

DeepSpeed vs FSDP

Связанные термины

• ZeRO — оптимизатор, реализованный в DeepSpeed
• FSDP — альтернатива от PyTorch
• gradient checkpointing — совместно используется
• mixed precision — BF16/FP16 в DeepSpeed