[yongggg's] DeepSpeed + Zero 둘러보기

안녕하세요 이번 장에서는 효율적인 Deep Learning training 툴인 Deepspeed에 대해 간단히 살펴보겠습니다.

간략하게 핵심부터 말씀드리자면, GPU에 모델 학습 parameter를 복제하는 기존의 비효율 시스템을 효율적으로 구동할 수 있도록 만든 framework라고 할 수 있습니다. 아래는 DeepSpeed + Zero와 그렇지 않은 모듈의 GPU 사용량 및 여러 장점들을 나타냅니다.

지금부터 짧게 DeepSpeed + Zero에 대해 설명하겠습니다.

1. Introduction

현재 딥러닝 trend가 LLM으로 들어오면서, 해당 모델의 크기는 점점 커지며 좋은 데이터, 큰 모델의 시너지는 굉장히 powerful하다는 것이 입증되고 있다. 모델의 두 가지 조건 중에서 큰 모델의 학습은 점점 더 많은 Hardware 사양을 요구했지만, 기존의 방법대로 학습할 수 없을 만큼 모델의 크기가 늘어났다. 이 문제를 해결하기 위해서 Microsoft에서는 DeepSpeed라는 딥러닝 최적화 library를 공개했다.

 

DeepSpeed는 규모, 속도, 비용 및 사용성을 개선하여 대규모 모델 학습을 크게 발전시켜 1,000억 개의 파라미터를 학습할 수 있는 능력을 제공한다. 그 중, ZeRO(Zero Redundancy Optimizer)는 모델 및 데이터 병렬화에 필요한 리소스를 크게 줄이는 동시에 학습할 수 있는 파라미터의 수를 크게 늘릴 수 있는 새로운 병렬 최적화 도구이다.

 

2. Zero

Microsoft에서는 1.5억개의 파라미터인 GPT-2 모델만을 GPU에 로드했을 때는 3GB의 메모리를 차지하지만, 32GB GPU에서도 학습이 불가능하다는 것에 의문을 가졌습니다. 그 결과 나머제 메모리는 복제된 parameter에 의해 낭비된다는 것을 알았습니다.

 

기존 병렬 처리 학습 방법에는 크게 DataParallel(DP)과 ModelParallel(MP)가 있다. DP는 모델 state를 복제하는 방법으로 계산/통신 효율이 좋지만, 메모리 효율성이 떨어지고, MP는 모델 state를 분할하여 메모리 효율성이 좋지만, 계산/통신 효율성이 떨어질 수 있다. 또 이 두 가지 방법은 훈련 프로세스에 걸쳐 필요한 모든 모델 상태를 정적으로 유지한다. Zero는 MP의 메모리 효율성을 달성하면서 DP의 계산/통신 효율성을 달성한다.

 

ZeRO-DP는 optimization states, gradients 및 parameters 파티셔닝에 해당하는 세 가지 주요 최적화 단계가 있다. 아래 그림에서 누적으로 stage를 활성화할 경우

1. Optimizer State Partitioning(Pos): 메모리 4배 감소, DP와 동일한 통신 볼륨
2. Gradient Partitioning 추가(Pos+g): 메모리 8배 감소, DP와 동일한 통신 볼륨
3. Parameter Partitioning 추가(Pos+g+p): 메모리 감소는 DP의 Nd(GPU 개수)에 따라 선형적이다. 
   ex) 64개 GPU(Nd = 64개)에 걸쳐 분할하면 메모리가 64배 감소한다. 통신 볼륨은 50% 정도 증가한다.

ZeRO-1

Optimizer stage Partitioning을 구현하여 1,000억개의 parameter를 가진 모델을 지원할 수 있다.

ZeRO-2

기존 ZeRO-1에서 Optimizer Stage Partitionin외에도 gradient, activation memor, fragmented memory의 메모리 공간을 줄여 메모리를 최적화할 수 있는 범위를 확장했다.

 

Model state memory

• ZeRO-2는 Optimizer stage Partitioning($P_{os}$) 와 Gradients Partitioning을 누적하여 적용($P_{os+g}$)함으로써 GPU당 메모리 소비를 2배 더 줄인다.
• 이는 기존 데이터 병렬 처리와 비교하면, 동일한 통신 볼륨으로 최대 8배의 메모리를 절약할 수 있다.

Activation memory

• activation partitioning을 통해 기존 모델 병렬화 방식에서 activation memory의 복제를 제거하는 기술을 도입한다.
• 상황에 따라 activation memory를 cpu로 보낸다.(모델 사이즈가 크거나, 메모리가 제한되는 경우)

Fragmented Memory

• 다양한 tensor의 수명 주기로 인해 학습 중에 메모리 단편화가 발생한다.
• 이로 인해, 사용 가능 메모리가 충분하더라도 메모리 할당에 실패할 때가 있는데, ZeRO-2에서는 tensor의 서로 다른 수명을 기반으로 메모리를 관리하여 메모리 단편환를 방지한다.
  이 코드는 DeepSpeed(학습 최적화 라이브러리)와 배포 됐고, PyTorch와 호환 가능하다.

Speed

• ZeRO-2는 ZeRO-1에 비해 향상된 메모리 효율성으로 처리량 증가, 빠른 학습이 가능하다. 
• Megtron-LM은 tensor-slicing model parallelism을 사용한다. Megatron-LM의 tensor-slicing model parallelism에 DeepSpeed Zero-2의 data parallelism을 추가하면, DeepSpeed가 Megatron-LM보다 열배정도 더 빠르다.

Democratizing large model training

• ZeRO-2는 일반적으로 모델 리팩토링을 필요로하는 Model Parallelism없이 최대 130억 개의 파라미터를 효율적으로 학습할 수 있다.
• 일반적인 Data Parallelism(ex; PyTorch Distribued Data Parallel)은 14억 개의 parameter model로도 메모리가 부족하지만, ZeRO-1은 최대 60억 개의 parameter를 지원한다.

DeepSpeed + Zero Config

다음은 흔히 사용하는 DeepSpeed + Zero의 run config이다. 여기서 batch 정도만 바꾸어 사용해도 무리 없이 사용 가능하다.

만일, 인자를 customizing하고 싶다면, https://www.deepspeed.ai/docs/config-json/ 이 공식 문서를 참고해보자.

DeepSpeed Configuration JSON

zero-2

json = {
    "fp16": {
        "enabled": true, 
        "loss_scale": 0, 
        "loss_scale_window": 1000, 
        "initial_scale_power": 16, 
        "hysteresis": 2, 
        "min_loss_scale": 1
    }, 
    "optimizer": {
        "type": "AdamW", 
        "params": {
            "lr": 5e-05, 
            "betas": [0.9, 0.999], 
            "eps": 1e-08, 
            "weight_decay": 0.0
        }
    }, 
    "scheduler": {
        "type": "WarmupLR", 
        "params": {
            "warmup_min_lr": 0, 
            "warmup_max_lr": 5e-05, 
            "warmup_num_steps": 0
        }
    }, 
    "zero_optimization": {
        "stage": 2, 
        "allgather_partitions": true, 
        "allgather_bucket_size": 2.000000e+08, 
        "overlap_comm": true, 
        "reduce_scatter": true, 
        "reduce_bucket_size": 2.000000e+08, 
        "contiguous_gradients": true, 
        "cpu_offload": true
    }, 
    "gradient_accumulation_steps": 1, 
    "gradient_clipping": 1.0, 
    "train_batch_size": 32, 
    "train_micro_batch_size_per_gpu": 32
}

 

zero-3

json = {
    "fp16": {
        "enabled": true, 
        "loss_scale": 0, 
        "loss_scale_window": 1000, 
        "initial_scale_power": 16, 
        "hysteresis": 2, 
        "min_loss_scale": 1
    }, 
    "optimizer": {
        "type": "AdamW", 
        "params": {
            "lr": 5e-05, 
            "betas": [0.9, 0.999], 
            "eps": 1e-08, 
            "weight_decay": 0.0
        }
    }, 
    "scheduler": {
        "type": "WarmupLR", 
        "params": {
            "warmup_min_lr": 0, 
            "warmup_max_lr": 5e-05, 
            "warmup_num_steps": 0
        }
    }, 
    "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": 1.000000e+14, 
        "reduce_bucket_size": 5.898240e+05, 
        "stage3_prefetch_bucket_size": 5.308416e+05, 
        "stage3_param_persistence_threshold": 7.680000e+03, 
        "stage3_max_live_parameters": 1.000000e+09, 
        "stage3_max_reuse_distance": 1.000000e+09, 
        "stage3_gather_fp16_weights_on_model_save": true
    }, 
    "gradient_accumulation_steps": 1, 
    "gradient_clipping": 1.0, 
    "steps_per_print": 2.000000e+03, 
    "train_batch_size": 32, 
    "train_micro_batch_size_per_gpu": 32, 
    "wall_clock_breakdown": false
}