DeepSpeed ZeRO

DeepSpeed ZeRO (Zero Redundancy Optimizer), dağıtık eğitimde bellek gereksizliğini ortadan kaldırmak için optimizör durumlarını, gradyanları ve parametreleri GPU'lar arasında parçalara ayırır. ZeRO'nun üç aşaması vardır ve her aşama, bir öncekinden daha fazla durumu parçalar. DeepSpeed, ek tasarruflar için CPU veya NVMe belleğe aktarımı da destekler. Her ek aşama ve aktarma seviyesi, GPU'lar arası iletişimin artması pahasına tepe belleği azaltır.

                  params        grads       opt states
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-1          │██████████│ │██████████│ │███░░░░░░░│  GPU 0
                │██████████│ │██████████│ │░░░███░░░░│  GPU 1
                │██████████│ │██████████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-2          │██████████│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │██████████│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │██████████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-3          │███░░░░░░░│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │░░░███░░░░│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │░░░░░░████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
  █ resident    ░ held on another GPU

ZeRO-2, ZeRO-3'e kıyasla daha düşük iletişim maliyetiyle gradyanları ve optimizör durumlarını parçalar. ZeRO-3'ü yalnızca modeliniz ZeRO-2 ile GPU'lar arasında sığmıyorsa kullanın.

Kurulum

DeepSpeed'i PyPI'den yükleyin veya Transformers'ı deepspeed ek paketiyle yükleyin.

pip install deepspeed
pip install transformers[deepspeed]

CUDA ile ilgili kurulum hatalarıyla karşılaşırsanız, DeepSpeed CUDA belgelerine göz atın. Kaynaktan kurulum, donanımınızla tam olarak eşleştiği ve PyPI sürümünde henüz bulunmayan özellikleri içerdiği için daha güvenilir bir seçenektir.

Yapılandır

Trainer, deepspeed argümanı aracılığıyla DeepSpeed'i entegre eder; bu argüman bir JSON yapılandırma dosyası kabul eder. Alternatif olarak, TrainingArguments yerine bir Accelerate yapılandırma dosyası kullanabilirsiniz.

TrainingArguments

Accelerate yapılandırma dosyası

"auto" değerini, DeepSpeed'in TrainingArguments içinden doldurmasını istediğiniz değerler için yapılandırmanızda kullanın. Bir değeri açıkça belirtmek istiyorsanız, hem DeepSpeed argümanı hem de TrainingArguments için aynı değeri kullandığınızdan emin olun.

DeepSpeed yapılandırma seçeneklerinin tam listesi için DeepSpeed Configuration JSON referansına bakın.

{
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto" }
    },
    "fp16": { "enabled": "auto" }
}

DeepSpeed yapılandırması	TrainingArguments
`train_micro_batch_size_per_gpu`	`per_device_train_batch_size`
`gradient_accumulation_steps`	`gradient_accumulation_steps`
`optimizer.params.lr`	`learning_rate`
`fp16.enabled`	`fp16`

Yapılandırmayı deepspeed argümanına iletin.

from transformers import TrainingArguments

args = TrainingArguments(
    deepspeed="path/to/deepspeed_config.json",
    ...
)

# DeepSpeed launcher
deepspeed --num_gpus 4 train.py

# torchrun
torchrun --nproc_per_node 4 train.py

# Accelerate
accelerate launch --num_processes 4 train.py

ZeRO aşamaları

Başlangıç noktası olarak kullanılacak bir ZeRO aşaması yapılandırması seçin.

ZeRO-1

ZeRO-2

ZeRO-3

Aşağıdaki alanlar, eğitimi özelleştirmek için önemlidir.

zero_optimization, ZeRO aşamasını ayarlar.
```
{ "zero_optimization": { "stage": 3 } }
```
Toplu işlem boyutunu ve gradyan biriktirme argümanlarını "auto" olarak ayarlayın. Bu değerleri TrainingArguments ile çelişen değerlere manuel olarak ayarlarsanız, eğitim yanlış değerlerle sessizce devam eder.
```
{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}
```
bf16, eğitim hassasiyetini belirler. TrainingArguments içindeki bf16 bayrağını yansıtması için "auto" olarak ayarlayın.
```
{ "bf16": { "enabled": "auto" } }
```
stage3_gather_16bit_weights_on_model_save, kaydetmeden önce tüm GPU'lar arasında bir all-gather işlemini gerçekleştirir ve tam tensörleri parçalarından yeniden oluşturur. Bu bir ZeRO-3 argümanıdır.
```
{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```
Tüm all-reduce gecikmesini geriye doğru geçişin arkasına gizlemek için overlap_comm değerini true olarak ayarlayın. allgather_bucket_size ve reduce_bucket_size parametreleri, iletişim hızını GPU belleğiyle takas eder. Daha düşük değerler daha az bellek kullanır ancak iletişimi yavaşlatır.
```
{
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}
```
offload_optimizer, optimizatörü CPU belleğine taşır. Daha fazla bellek tasarrufu için model parametrelerini de offload_param ile taşıyın (yalnızca ZeRO-3). CPU-GPU aktarmalarını hızlandırmak için pin_memory değerini true olarak ayarlayın; ancak bu, diğer süreçlerin kullanamayacağı RAM'i kilitleyecektir.
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}
```

optimizer ve scheduler, TrainingArguments içinde yapılandırılan optimizör ve zamanlayıcıya varsayılan olarak ayarlanır. LAMB gibi DeepSpeed yerel bir optimizöre ihtiyacınız olmadıkça, DeepSpeed'in değerleri TrainingArguments içinden okuması için "auto" olarak ayarlayın.

{
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": { "total_num_steps": "auto", "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto" }
    }
}

Optimizatörü dışa aktarıyorsanız, DeepSpeed'in CPU Adam optimizatörünü kullanmak için zero_force_ds_cpu_optimizer değerini false olarak ayarlayın.

{
    "zero_force_ds_cpu_optimizer": false
}

Kontrol Noktaları

DeepSpeed, kontrol noktalarını doğrudan from_pretrained() ile yüklenemeyen bir parçalı formatta kaydeder. Trainer'ın eğitim sonunda en iyi kontrol noktasını izleyip yeniden yüklemesi için load_best_model_at_end() parametresini True olarak ayarlayın.

from transformers import TrainingArguments, Trainer

args = TrainingArguments(
    deepspeed="ds_config_zero3.json",
    load_best_model_at_end=True,
    ...
)
# after training, save a normal transformers checkpoint
trainer.save_model("./best-model")

save_only_model=True ayarlandığında tam optimizör durumu kaydedilmez; bu da eğitim sonunda en iyi modelin yeniden yüklenemeyeceği anlamına gelir. Ayrıca, tam ağırlıkları parçalarından yeniden oluşturmak için stage3_gather_16bit_weights_on_model_save: true değerini de ayarlayın. Bu, ZeRO-3 ile birleştirilmiş 16-bit model artefaktının veya 16-bit durum sözlüğünün kaydedilmesi için gereklidir. save_only_model=True değeri load_best_model_at_end=True değeriyle birleştirildiğinde Transformers bir hata üretir.

Farklı paralellik yapılandırmaları arasında devam etmek için DeepSpeed'in Universal Checkpointing kılavuzuna bakın.

Sonraki adımlar

ZeRO'nun nasıl çalıştığı hakkında daha fazla bilgi edinmek için The Ultra-Scale Playbook'taki Zero Redundancy Optimizer bölümünü okuyun.
ZeRO makalelerini okuyun: Memory Optimizations Toward Training Trillion Parameter Models, Democratizing Billion-Scale Model Training ve Breaking the GPU Memory Wall for Extreme Scale Deep Learning.

Update on GitHub