DeepSpeed ZeRO

DeepSpeed ZeRO (Zero Redundancy Optimizer) loại bỏ sự dư thừa bộ nhớ trong huấn luyện phân tán bằng cách chia nhỏ trạng thái optimizer, gradient và tham số trên các GPU. ZeRO có ba giai đoạn, mỗi giai đoạn chia nhỏ nhiều trạng thái hơn giai đoạn trước. DeepSpeed cũng hỗ trợ offload sang bộ nhớ CPU hoặc NVMe để tiết kiệm thêm. Mỗi giai đoạn và mức offload bổ sung đều làm giảm bộ nhớ đỉnh, với cái giá là tăng lượng giao tiếp giữa các GPU.

                  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 chia nhỏ gradient và trạng thái optimizer với chi phí truyền thông thấp hơn so với ZeRO-3. Chỉ sử dụng ZeRO-3 khi mô hình của bạn không vừa trên các GPU với ZeRO-2.

Cài đặt

Cài đặt DeepSpeed từ PyPI, hoặc cài đặt Transformers với phần mở rộng deepspeed.

pip install deepspeed
pip install transformers[deepspeed]

Nếu bạn gặp lỗi cài đặt liên quan đến CUDA, hãy kiểm tra tài liệu DeepSpeed CUDA. Cài đặt từ mã nguồn là lựa chọn đáng tin cậy hơn vì nó khớp chính xác với phần cứng của bạn và bao gồm các tính năng chưa có trong bản phát hành PyPI.

Cấu hình

Trainer tích hợp DeepSpeed thông qua đối số deepspeed, nhận một tệp cấu hình JSON. Thay vào đó, hãy sử dụng một tệp cấu hình Accelerate thay vì TrainingArguments.

TrainingArguments

Tệp cấu hình Accelerate

Sử dụng "auto" trong cấu hình của bạn cho các giá trị mà bạn muốn DeepSpeed tự điền từ TrainingArguments. Nếu bạn muốn chỉ định rõ một giá trị, hãy đảm bảo sử dụng giá trị giống nhau cho cả đối số DeepSpeed và TrainingArguments.

Xem tài liệu tham khảo DeepSpeed Configuration JSON để biết danh sách đầy đủ các tùy chọn cấu hình DeepSpeed.

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

Cấu hình 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`

Truyền config vào đối số 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

Các giai đoạn của ZeRO

Chọn một cấu hình giai đoạn ZeRO để sử dụng làm điểm khởi đầu.

ZeRO-1

ZeRO-2

ZeRO-3

Các trường sau đây quan trọng để tùy chỉnh quá trình huấn luyện.

zero_optimization thiết lập giai đoạn ZeRO.
```
{ "zero_optimization": { "stage": 3 } }
```
Đặt các đối số kích thước lô (batch size) và tích lũy gradient thành "auto". Nếu bạn tự động đặt các giá trị này không khớp với TrainingArguments, quá trình huấn luyện sẽ tiếp tục một cách im lặng với các giá trị sai.
```
{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}
```
bf16 thiết lập độ chính xác của quá trình huấn luyện. Đặt nó thành "auto" để nó phản ánh cờ bf16 trong TrainingArguments.
```
{ "bf16": { "enabled": "auto" } }
```
stage3_gather_16bit_weights_on_model_save thực hiện all-gather trên tất cả các GPU trước khi lưu, tái tạo các tensor đầy đủ từ các shard của chúng. Đây là một đối số của ZeRO-3.
```
{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```
Đặt overlap_comm thành true để che độ trễ all-reduce phía sau bước truyền ngược. allgather_bucket_size và reduce_bucket_size đánh đổi tốc độ truyền thông lấy bộ nhớ GPU. Giá trị thấp hơn sử dụng ít bộ nhớ hơn nhưng làm chậm truyền thông.
```
{
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}
```
offload_optimizer chuyển bộ tối ưu sang bộ nhớ CPU. Để tiết kiệm thêm bộ nhớ, hãy chuyển cả tham số mô hình bằng offload_param (chỉ ZeRO-3). Đặt pin_memory thành true để tăng tốc chuyển dữ liệu giữa CPU và GPU, nhưng điều này sẽ khóa RAM khiến các tiến trình khác không thể sử dụng.
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}
```

optimizer và scheduler mặc định là optimizer và scheduler được cấu hình trong TrainingArguments. Đặt thành "auto" để DeepSpeed đọc các giá trị từ TrainingArguments trừ khi bạn cần một optimizer bản địa của DeepSpeed như 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" }
    }
}

Nếu bạn đang offload optimizer, hãy đặt zero_force_ds_cpu_optimizer thành false để sử dụng CPU Adam optimizer của DeepSpeed.

{
    "zero_force_ds_cpu_optimizer": false
}

Điểm kiểm tra

DeepSpeed lưu checkpoint ở định dạng phân mảnh, không thể tải trực tiếp bằng from_pretrained(). Đặt load_best_model_at_end() thành True để Trainer theo dõi và tải lại checkpoint tốt nhất vào cuối quá trình huấn luyệ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")

Đặt save_only_model=True sẽ bỏ qua việc lưu trạng thái optimizer đầy đủ, điều này có nghĩa là bạn không thể tải lại mô hình tốt nhất ở cuối quá trình huấn luyện. Ngoài ra, hãy đặt stage3_gather_16bit_weights_on_model_save: true để tái tạo trọng số đầy đủ từ các shard của chúng. Điều này là bắt buộc để lưu một artifact mô hình 16-bit đã hợp nhất hoặc state dict 16-bit với ZeRO-3. Transformers sẽ gây ra lỗi khi save_only_model=True được kết hợp với load_best_model_at_end=True.

Để tiếp tục huấn luyện với các cấu hình song song khác nhau, hãy xem hướng dẫn Universal Checkpointing của DeepSpeed.

Các bước tiếp theo

Đọc chương Zero Redundancy Optimizer trong The Ultra-Scale Playbook để tìm hiểu thêm về cách ZeRO hoạt động.
Đọc các bài báo về ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, Democratizing Billion-Scale Model Training và Breaking the GPU Memory Wall for Extreme Scale Deep Learning.

Update on GitHub