Налаштування DeepSpeed для розподіленого навчання LLM

Налаштування DeepSpeed для розподіленого навчання LLM

DeepSpeed — бібліотека Microsoft для ефективного навчання великих мовних моделей. Її ключові можливості: технологія ZeRO (Zero Redundancy Optimizer), яка усуває надмірність при зберіганні стану оптимізатора та параметрів моделі між GPU, а також підтримка mixed precision, gradient checkpointing та pipeline parallelism.

ZeRO: ключова інновація DeepSpeed

ZeRO Stage 1 — шардування стану оптимізатора (Adam states) між GPU. На 8 GPU: споживання пам'яті оптимізатором знижується у 8 разів.

ZeRO Stage 2 — додається шардування градієнтів. Сумарне зниження пам'яті: ~ 8x для стану оптимізатора + ~ 8x для градієнтів.

ZeRO Stage 3 - повне шардування: параметри, градієнти, стан оптимізатора. Дозволяє навчати моделі, які в сумі не поміщаються навіть на всі кластери GPU. При цьому параметри gather і scatter при кожному forward/backward pass overhead на комунікацію вище, ніж Stage 2.

ZeRO-Infinity — offloading параметрів на CPU RAM та NVMe SSD. Дозволяє навчати моделі з трильйонами параметрів на обмеженій кількості GPU за рахунок пропускної спроможності PCIe/NVMe.

Конфігурація DeepSpeed

{
  "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  // или true для A100/H100
  },
  "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

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()

Запуск:

deepspeed --num_gpus=8 train.py --deepspeed ds_config.json
# Или через torchrun:
torchrun --nproc_per_node=8 train.py --deepspeed ds_config.json

ZeRO Stage 3: навчання дуже великих моделей

Для моделей 30B+ параметрів на обмеженому кластері:

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  }
}

Оцінка конфігурації

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
)
# Выводит оценку RAM и GPU памяти для разных конфигураций

Практичні показники

DeepSpeed у поєднанні з gradient checkpointing та activation recomputation дозволяє навчати моделі в 3-5x більшого розміру на тому ж залізі порівняно з наївною DDP реалізацією.