# DEEPSPEED 介紹

## :memo: Where do I start?
- 聯絡窗口 Email us : [e-mail] 王小姐

## 訓練深度學習模式時的顯存消耗可分為兩大部分
- 模型狀態(model states)。對於大型模型來說，大部分顯存消耗都是被模型狀態佔用的，主要包含三個部分：最佳化器的狀態(Optimizer States)、梯度(Gradients)、參數(Parameters)。三者簡稱為OPG。
- 殘餘狀態(residual states)。包括激活函數、臨時緩衝區和不可用的記憶體碎片。

## Deepspeed Zero
ZeRO主要包括三個等級的最佳化：ZeRO-1，ZeRO-2和ZeRO-3，每個等級都在前一個等級的基礎上進行了進一步的最佳化。
- ZeRO-0: Distributed Data Parallel, DDP
- ZeRO-1：優化器狀態分區（Optimizer State Partitioning） 在傳統的資料並行訓練中，每個設備都會保存一份完整的模型參數和優化器狀態，這會導致記憶體的大量浪費。ZeRO-1透過在裝置間分散式儲存優化器狀態，從而顯著減少了記憶體的使用。
- ZeRO-2：優化器狀態分區+ 梯度分區（Optimizer State Partitioning + Gradient Partitioning） 在ZeRO-1的基礎上，ZeRO-2進一步將梯度也進行了分區儲存。這樣，每個設備只需要保存一部分的優化器狀態和梯度，從而進一步減少了記憶體的使用。這使得我們可以在相同的硬體上訓練更大的模型，或者在更小的硬體上訓練相同的模型。
- ZeRO-3：優化器狀態分區+ 梯度分區+ 參數分區（Optimizer State Partitioning + Gradient Partitioning + Parameter Partitioning） ZeRO-3是ZeRO的最高等級的最佳化。除了優化器狀態和梯度，ZeRO-3還將模型參數進行了分區儲存。這樣，每個設備只需要保存一部分的優化器狀態、梯度和模型參數，從而最大限度地減少了記憶體的使用

## Deepspeed Zero configure
- ds_zero2_config.json

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "zero_allow_untested_optimizer": true,
  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
        "device": "cpu",
        "pin_memory": true
    },      
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

- ds_zero3_config.json

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "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
  },
  "activation_checkpointing": {
    "partition_activations": true,
    "cpu_checkpointing": true,
    "contiguous_memory_optimization": true,
    "number_checkpoints": 4
  }
}

## 如何選擇不同的Zero stage和offload
- 由左至右，越來越慢

Stage 0 (DDP) > Stage 1 > Stage 2 > Stage 2 + offload > Stage 3 > Stage 3 + offloads

- 由左至右，所需GPU顯存越來越少

Stage 0 (DDP) < Stage 1 < Stage 2 < Stage 2 + offload < Stage 3 < Stage 3 + offloads

## 調參步驟
1. 將batch_size設定為1，透過梯度累積實現任意的有效batch_size
2. 如果OOM則，設定--gradient_checkpointing 1 (HF - Trainer)，或者model.gradient_checkpointing_enable()
3. 如果OOM則，試試ZeRO stage 2
4. 如果OOM則，試試ZeRO stage 2 +offload_optimizer
5. 如果OOM則，試試ZeRO stage 3
6. 如果OOM則，嘗試offload_param到CPU
7. 如果OOM則，嘗試offload_optimizer到CPU
8. 如果OOM則，請嘗試降低一些預設參數。例如使用generate時，減少beam search的搜尋範圍
9. 如果OOM則，使用混合精度訓練，在Ampere的GPU上使用bf16，在舊版GPU上使用fp16
10. 如果仍然OOM，則使用ZeRO-Infinity ，使用offload_param和offload_optimizer到NVME
11. 一旦使用batch_size=1時，沒有導致OOM，測量此時的有效吞吐量，然後盡可能增加batch_size
12. 開始最佳化參數，可以關閉offload參數，或降低ZeRO stage，然後調整batch_size，然後繼續測量吞吐量，直到效能比較滿意（調參可以增加66%的效能）

## Deepspeed Zero 補充說明

模型狀態分區