DeepSpeed ZeRO

DeepSpeed ZeRO

DeepSpeed ZeRO (Zero Redundancy Optimizer)는 옵티마이저 상태, 그라디언트, 파라미터를 GPU 간에 샤딩해서 분산 학습에서 메모리 중복을 제거해요. ZeRO에는 세 단계가 있고, 각 단계는 이전 단계보다 더 많은 상태를 샤딩해요. DeepSpeed는 추가 절약을 위해 CPU 또는 NVMe 메모리로의 offloading도 지원해요. 단계와 offload 수준을 더할수록 최고(peak) 메모리는 줄어들지만, 그 대가로 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는 ZeRO-3보다 낮은 통신 오버헤드로 그라디언트와 옵티마이저 상태를 샤딩해요. ZeRO-2로 모델이 GPU들에 들어가지 않을 때만 ZeRO-3를 사용하세요.

출처: 문서

본문

설치

PyPI에서 DeepSpeed를 설치하거나, deepspeed extra와 함께 Transformers를 설치해요.

pip install deepspeed
pip install transformers[deepspeed]

CUDA 관련 설치 에러가 발생하면 DeepSpeed CUDA 문서를 확인해요. 소스에서 설치하는 것이 더 신뢰할 수 있는 옵션이에요. 정확한 하드웨어에 맞춰지고 PyPI 릴리스에 아직 없는 기능도 포함되기 때문이에요.

구성

Trainer는 JSON config 파일을 받는 deepspeed 인자를 통해 DeepSpeed를 통합해요. 또는 TrainingArguments 대신 Accelerate config file을 사용할 수도 있어요.

DeepSpeed가 TrainingArguments에서 채워주길 원하는 값에는 config에서 "auto"를 사용해요. 값을 명시적으로 지정하려면 DeepSpeed 인자와 TrainingArguments 양쪽에 동일한 값을 사용해야 해요.

[!NOTE] DeepSpeed config 옵션의 전체 목록은 DeepSpeed Configuration JSON 참조를 확인하세요.

{
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto" }
    },
    "fp16": { "enabled": "auto" }
}
DeepSpeed config 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

config를 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

[!NOTE] Accelerate는 TrainingArguments의 deepspeed 인자를 무시합니다.

accelerate config 명령을 실행하고 하드웨어와 학습 설정에 대한 질문에 답해서 캐시에 default_config.yaml 파일을 만들어요.

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 4

Trainer 기반 스크립트로 accelerate launch를 실행해요.

accelerate launch --config_file deepspeed_config.yaml train.py

ZeRO 단계

시작점으로 사용할 ZeRO 단계 config를 선택해요.

{
    "bf16": { "enabled": "auto" },
    "zero_optimization": { "stage": 1 },
    "gradient_clipping": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto"
}

{
    "bf16": { "enabled": "auto" },
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_clipping": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto"
}

[!WARNING] ZeRO-3는 초기화 중에 파라미터를 샤딩합니다. 모델을 로드하기 전에 TrainingArguments를 먼저 인스턴스화해야 합니다. DeepSpeed가 구성되기 전에 모델이 이미 각 GPU에 있다면 메모리가 절약되지 않습니다.

{
    "bf16": { "enabled": "auto" },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_gather_16bit_weights_on_model_save": true,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    },
    "gradient_clipping": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto"
}

학습을 커스터마이즈할 때 중요한 필드는 다음과 같아요.

  • zero_optimization은 ZeRO 단계를 설정해요.

    { "zero_optimization": { "stage": 3 } }

  • 배치 크기와 그라디언트 누적 인자를 "auto"로 설정해요. TrainingArguments와 일치하지 않는 값으로 수동 설정하면 학습이 잘못된 값으로 조용히 계속돼요.

    {
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}

  • bf16은 학습 정밀도를 설정해요. TrainingArguments의 bf16 플래그를 반영하도록 "auto"로 설정해요.

    { "bf16": { "enabled": "auto" } }

  • stage3_gather_16bit_weights_on_model_save는 저장 전에 모든 GPU에 걸쳐 all-gather를 수행해서 샤드에서 전체 텐서를 재구성해요. 이는 ZeRO-3 인자예요.

    {
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}

  • overlap_comm을 true로 설정하면 all-reduce 지연 시간을 backward pass 뒤에 숨겨요(hide). allgather_bucket_size와 reduce_bucket_size는 통신 속도와 GPU 메모리를 맞바꿔요. 값이 낮을수록 메모리는 덜 쓰지만 통신은 느려져요.

    {
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}

  • offload_optimizer는 옵티마이저를 CPU 메모리로 오프로드해요. 더 많은 메모리를 절약하려면 offload_param으로 모델 파라미터도 오프로드해요 (ZeRO-3에서만). pin_memory를 true로 설정하면 CPU-GPU 전송을 빠르게 하지만, 다른 프로세스가 사용할 수 없는 RAM을 잠가요.

    {
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}

  • optimizer와 scheduler는 기본적으로 TrainingArguments에 구성된 옵티마이저와 스케줄러를 사용해요. LAMB 같은 DeepSpeed 네이티브 옵티마이저가 필요하지 않다면 "auto"로 설정해서 DeepSpeed가 TrainingArguments에서 값을 읽도록 해요.

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

    옵티마이저를 오프로드한다면 DeepSpeed의 CPU Adam 옵티마이저를 사용하도록 zero_force_ds_cpu_optimizer를 false로 설정해요.

    {
    "zero_force_ds_cpu_optimizer": false
}

체크포인트

DeepSpeed는 from_pretrained()으로 직접 로드할 수 없는 샤딩된 형식으로 체크포인트를 저장해요. load_best_model_at_end()를 True로 설정하면 Trainer가 학습이 끝날 때 최상의 체크포인트를 추적하고 다시 로드해요.

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를 설정하면 전체 옵티마이저 상태 저장을 건너뛰므로 학습이 끝날 때 최상의 모델을 다시 로드할 수 없어요. 또한 stage3_gather_16bit_weights_on_model_save: true를 설정해서 샤드에서 전체 가중치를 재구성해야 해요. 이는 ZeRO-3로 통합된 16-bit 모델 아티팩트 또는 16-bit state dict를 저장하는 데 필요해요. Transformers는 save_only_model=True가 load_best_model_at_end=True와 결합되면 에러를 발생시켜요.

[!TIP] 서로 다른 병렬 구성 간에 재개(resume)하려면 DeepSpeed의 Universal Checkpointing 가이드를 참고하세요.

더 알아보기 (Learn more)