DeepSpeed ZeRO

DeepSpeed ZeRO (Zero Redundancy Optimizer) menghilangkan redundansi memori dalam pelatihan terdistribusi dengan membagi (sharding) status optimizer, gradien, dan parameter di antara GPU. ZeRO memiliki tiga tahap, di mana setiap tahap membagi lebih banyak status daripada tahap sebelumnya. DeepSpeed juga mendukung offloading ke memori CPU atau NVMe untuk penghematan lebih lanjut. Setiap tahap tambahan dan tingkat offload mengurangi memori puncak, dengan biaya komunikasi antar-GPU yang lebih banyak.

                  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 memecah gradien dan status optimizer dengan overhead komunikasi yang lebih rendah daripada ZeRO-3. Gunakan ZeRO-3 hanya jika model Anda tidak muat di seluruh GPU dengan ZeRO-2.

Instalasi

Install DeepSpeed dari PyPI, atau install Transformers dengan extra deepspeed.

pip install deepspeed
pip install transformers[deepspeed]

Jika Anda mengalami error instalasi terkait CUDA, periksa dokumentasi DeepSpeed CUDA. Installing from source adalah opsi yang lebih andal karena sesuai dengan perangkat keras Anda secara persis dan mencakup fitur yang belum tersedia di rilis PyPI.

Konfigurasi

Trainer mengintegrasikan DeepSpeed melalui argumen deepspeed, yang menerima file konfigurasi JSON. Sebagai alternatif, gunakan file konfigurasi Accelerate alih-alih TrainingArguments.

TrainingArguments

File konfigurasi Accelerate

Gunakan "auto" di konfigurasi Anda untuk nilai yang ingin diisi oleh DeepSpeed dari TrainingArguments. Jika Anda ingin secara eksplisit menentukan nilai, pastikan Anda menggunakan nilai yang sama untuk argumen DeepSpeed dan TrainingArguments.

Lihat referensi DeepSpeed Configuration JSON untuk daftar lengkap opsi konfigurasi DeepSpeed.

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

Konfigurasi DeepSpeed	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`

Lewatkan config ke argumen deepspeed.

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

Tahapan ZeRO

Pilih konfigurasi tahap ZeRO untuk digunakan sebagai titik awal.

ZeRO-1

ZeRO-2

ZeRO-3

Field-field berikut penting untuk menyesuaikan pelatihan.

zero_optimization mengatur tahap ZeRO.
```
{ "zero_optimization": { "stage": 3 } }
```
Atur argumen batch size dan gradient accumulation ke "auto". Jika Anda mengatur nilai-nilai ini secara manual yang tidak sesuai dengan TrainingArguments, pelatihan akan terus berjalan secara diam-diam dengan nilai yang salah.
```
{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}
```
bf16 mengatur presisi pelatihan. Atur menjadi "auto" agar mencerminkan flag bf16 di TrainingArguments.
```
{ "bf16": { "enabled": "auto" } }
```
stage3_gather_16bit_weights_on_model_save melakukan all-gather di semua GPU sebelum menyimpan, merekonstruksi tensor penuh dari shard-nya. Ini adalah argumen ZeRO-3.
```
{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```
Set overlap_comm ke true untuk menyembunyikan latensi all-reduce di balik backward pass. allgather_bucket_size dan reduce_bucket_size menukar kecepatan komunikasi dengan memori GPU. Nilai yang lebih rendah menggunakan memori lebih sedikit tetapi memperlambat komunikasi.
```
{
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}
```
offload_optimizer memindahkan optimizer ke memori CPU. Untuk menghemat memori lebih banyak, pindahkan juga parameter model dengan offload_param (hanya ZeRO-3). Atur pin_memory ke true untuk mempercepat transfer CPU-GPU, tetapi ini mengunci RAM yang tidak tersedia untuk proses lain.
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}
```

optimizer dan scheduler secara default menggunakan optimizer dan scheduler yang dikonfigurasi di TrainingArguments. Atur ke "auto" agar DeepSpeed membaca nilai dari TrainingArguments kecuali Anda memerlukan optimizer native DeepSpeed seperti LAMB.

{
    "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" }
    }
}

Jika Anda melakukan offloading optimizer, atur zero_force_ds_cpu_optimizer ke false untuk menggunakan CPU Adam optimizer dari DeepSpeed.

{
    "zero_force_ds_cpu_optimizer": false
}

Checkpoint

DeepSpeed menyimpan checkpoint dalam format terpecah (sharded) yang tidak dapat dimuat secara langsung dengan from_pretrained(). Atur load_best_model_at_end() ke True agar Trainer melacak dan memuat ulang checkpoint terbaik di akhir pelatihan.

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

Menetapkan save_only_model=True akan melewati penyimpanan state optimizer penuh, yang berarti Anda tidak dapat memuat ulang model terbaik di akhir pelatihan. Juga, tetapkan stage3_gather_16bit_weights_on_model_save: true untuk merekonstruksi bobot penuh dari shard-nya. Ini diperlukan untuk menyimpan artefak model 16-bit yang terpadu atau state dict 16-bit dengan ZeRO-3. Transformers akan melempar error ketika save_only_model=True dikombinasikan dengan load_best_model_at_end=True.

Untuk melanjutkan dengan konfigurasi paralelisasi yang berbeda, lihat panduan Universal Checkpointing dari DeepSpeed.

Langkah selanjutnya

Baca bab Zero Redundancy Optimizer dari The Ultra-Scale Playbook untuk mempelajari lebih lanjut tentang cara kerja ZeRO.
Baca makalah ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, Democratizing Billion-Scale Model Training, dan Breaking the GPU Memory Wall for Extreme Scale Deep Learning.

Update on GitHub