🤗

Huggingface + DeepSpeed + FairScale

Tags
NLP
MLDL Framework
Published
Published May 16, 2021
💡
아래 글은 PyTorch ≥ 1.8.0, Transformers ≥ 4.6.0 기준으로 작성함

PyTorch로 큰 모델을 학습해보자?

Huggingface에서 제공하는 내장된 Trainer를 사용할 경우 가능한 학습 Device는 다음과 같다.
  • CPU
  • Single GPU
  • 1 Node, Multi GPU
  • Multi Node, Multi GPU
  • TPU
  • TPU Pods
여기서 가장 자주 쓰는게 당연히 SingleGPU, 1Node-MultiGPU, 그리고 간혹 TPU를 쓴다.
한편 단일 GPU나 TPU를 통한 학습을 진행하거나 혹은 DDP를 통해 1Node-MultiGPU를 학습한다면 VRam의 한계로 인해 학습 가능한 모델 최대 크기의 한게가 있다.

모델 크기 한계

  • V100 32G 기준 약 1.3Billion params 모델이 올라갈 수 있다.
    • Model + Optimizer + batch 1 정도..
  • 하지만 LM 학습시 Batch size 역시 성능에 더 큰 영향을 주기 때문에 너무 작은 Batch size를 억지로 학습하는 것은 최종 결과물이 만족스럽지가 않다.
  • TPU v3-8 기준 16GB내에 맞춰야 하는 이슈가 있다.
    • GPT-2 기준 Batch size 8정도가 아슬하다.

2019년까지의 분산학습: Data Parallel, Distributed Data Parallel, Apex

  • 다만 이 글은 단일 Node의 Multi GPU에서의 학습 상황을 고려하고 있기 때문에, 최근의 모델에서 필요로 하는 MultiNode등의 학습은 어렵다.

좀더 큰 사이즈의 학습을 위해: ZeRO, FairScale

  • 결국 대규모 모델 학습을 위해서 쪼갤 수 있는건 크게 4가지다.
      1. Batch: batch를 각 GPU로 쪼개서 각 GPU에서 학습하자
        1. notion image
      1. Optimizer State: 해당 Batch를 위한 Optimizer만 가져오기
      1. Gradient: backward 위한 Gradient를 해당 batch만 쓰자
      1. Model weight(parameters): 모델조차 쪼개버리자, Model parallel
        1. 4-1. Model Parallel
          notion image
          • Transformer 모델 기준 Head 단위로 모델을 쪼갠다고 이해하면 편하다.
          4-2. Pipeline Parallel
          notion image
          • Transformer 모델 기준 Layer 단위로 모델을 쪼갠다고 이해하면 편하다.
  • 쪼개지 않고 하는 방법도 물론 있다.
    • Zero-2 (aka Zero-Offload), up to 13 Billion on 1 GPU
      • notion image
      • Single GPU + CPU Ram
      • GPU에서 Forward & Backwrd 후 Gradient → CPU(ram)으로 이동
      • CPU에서 Model(params) update → GPU로 Copy
    • Zero-3, up to 40 Billion on 1 GPU & 40 Trillion on 512 GPU
      • Zero-3 = Zero-2 + Parameter(model) partitioning

Huggingface + DeepSpeed(ZeRO) 👍

설치

  • 설치는 간단하다.
pip install transformers[deepspeed]
  • 혹은 수동으로 아래와 같이 설치해줄 수도 있다.
💡
정말 공식 가이드대로 저 위 한 줄만 하면 DeepSpeed의 모든 것을 쓸 수 있을까? → ❌ CPU Fused ADAM등, CPU Offload를 full로 사용하는 등 여러 고급기능을 사용하려면 CUDA를 비롯해 C++ build 환경을 갖춘 상태에서 아래 공식 가이드를 따라 커스텀 빌드를 해야 한다. 공식 링크: https://www.deepspeed.ai/tutorials/advanced-install/
git clone https://github.com/microsoft/DeepSpeed/ cd DeepSpeed rm -rf build TORCH_CUDA_ARCH_LIST="6.1;8.6" DS_BUILD_OPS=1 pip install . \ --global-option="build_ext" --global-option="-j8" --no-cache -v \ --disable-pip-version-check 2>&1 | tee build.log # TORCH_CUDA_ARCH_LIST는 사용하는 GPU 환경에 맞춰 쓰자. # A100=8.0, RTX_titan=7.5, RTX 3090=8.6, ... 이렇게 맞춰주면 된다. # 위 코드는 아마(?) 6.1부터 8.6까지 모두에 대해 빌드하는 듯?

PyTorch distributed → DeepSpeed

  • torch.distributed.launch 를..
python -m torch.distributed.launch --nproc_per_node=2 your_program.py <normal cl args>
  • deepspeed 로 바꾸자!
deepspeed --num_gpus=2 your_program.py <normal cl args> --deepspeed ds_config.json
  • 만약 --num_gpus 안쓰면, 보이는 모든 GPU를 갖다 쓴다.

간단한 샘플

deepspeed examples/pytorch/translation/run_translation.py \ --deepspeed tests/deepspeed/ds_config_zero3.json \ --model_name_or_path t5-small --per_device_train_batch_size 1 \ --output_dir output_dir --overwrite_output_dir --fp16 \ --do_train --max_train_samples 500 --num_train_epochs 1 \ --dataset_name wmt16 --dataset_config "ro-en" \ --source_lang en --target_lang ro
  • Huggingface의 examples/pytorch/translation/run_translation.py 를 DeepSpeed로.

DeepSpeed on 단일 GPU

  • ZeRO Offload를 사용하기 위한 경우
  • DeepSpeed(ZeRO)의 메모리 관리(파편화 방지)로 보다 큰 Batch size 사용
  • 간단한 ds_config.json 파일을 아래와 같이 만들어서 쓰기만 해도 성능 ++
💡
아래 추천 코드는 아직 Stage 2, 즉 ZeRO-2 기반 코드! zero-3은 아직(20210517) 성능 평가가 되지 않아서인듯.
{ "zero_optimization": { "stage": 2, "allgather_partitions": true, "allgather_bucket_size": 2e8, "reduce_scatter": true, "reduce_bucket_size": 2e8, "overlap_comm": true, "contiguous_gradients": true, "cpu_offload": true } }
  • 위 코드 보니까... ZeRO Optimize하고 Vram 최적화, 그리고 CPU Offload를 추가하는 수준인 듯함

DeepSpeed에서는 CUDA_VISIBLE_DEVICES 로 GPU 제한 못한다!

대신 --include localhost:GPU번호 로 해야함.
deepspeed --include localhost:1 examples/pytorch/translation/run_translation.py ...
  • 위 코드는 로컬의 1번(2번째) GPU를 쓴다는 뜻.

Jupyter Notebook에서 띄우기(1 GPU only)

  • OS environ으로 넣어주고.. deepspeed 항목으로 Trainer에 넣어주면 된다는 것
# DeepSpeed requires a distributed environment even when only one process is used. # This emulates a launcher in the notebook import os os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '9994' # modify if RuntimeError: Address already in use os.environ['RANK'] = "0" os.environ['LOCAL_RANK'] = "0" os.environ['WORLD_SIZE'] = "1" # Now proceed as normal, plus pass the deepspeed config file training_args = TrainingArguments(..., deepspeed="ds_config_zero3.json") trainer = Trainer(...) trainer.train()

Jupyter Notebook + Multi GPU = ❌

  • 노트북 셀에서 아래와 같이 ds_config_zero3.json 파일을 만들수야 있지만....
  • 그냥 json파일 따로 만들어서 아래 내용 잘 넣고, deepspeed 커맨드로 실행하는게 맞다.
%%bash cat <<'EOT' > ds_config_zero3.json { "fp16": { "enabled": "auto", "loss_scale": 0, "loss_scale_window": 1000, "initial_scale_power": 16, "hysteresis": 2, "min_loss_scale": 1 }, "optimizer": { "type": "AdamW", "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" } }, "scheduler": { "type": "WarmupLR", "params": { "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto" } }, "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": 1e14, "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_fp16_weights_on_model_save": true }, "gradient_accumulation_steps": "auto", "gradient_clipping": "auto", "steps_per_print": 2000, "train_batch_size": "auto", "train_micro_batch_size_per_gpu": "auto", "wall_clock_breakdown": false } EOT

ZeRO-2에서 많은 기능을 활성화한 ds_config.json 파일

  • 아래 파일에서 auto 는 이후 Transformers Trainer에서 자동으로 값을 잡아주도록 세팅하는 것
  • 공식 docs에서는 웬만하면 auto 쓰라고 권장함.
{ "fp16": { "enabled": "auto", "loss_scale": 0, "loss_scale_window": 1000, "initial_scale_power": 16, "hysteresis": 2, "min_loss_scale": 1 }, "optimizer": { "type": "AdamW", "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" } }, "scheduler": { "type": "WarmupLR", "params": { "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto" } }, "zero_optimization": { "stage": 2, "allgather_partitions": true, "allgather_bucket_size": 2e8, "overlap_comm": true, "reduce_scatter": true, "reduce_bucket_size": 2e8, "contiguous_gradients": true, "cpu_offload": true }, "gradient_accumulation_steps": "auto", "gradient_clipping": "auto", "train_batch_size": "auto", "train_micro_batch_size_per_gpu": "auto", }

ZeRO-3을 위한 간단한 세팅 파일

  • ZeRO stage = 3인 세팅 파일
  • Optimizer, Params(model)도 Offload하는 방식
    • 여기서 devicecpu 아니라 nvme 로 바꿔줄수도 있다!
  • pin_memory 를 사용하면 고정된 Memory 주소 사용을 통해 속도 향상이 있다.
    • (이정도는 해야 offload같은걸 하겠지..?)
{ "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": 1e14, "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_fp16_weights_on_model_save": true } }
  • stage3_gather_fp16_weights_on_model_save 를 쓰면 fp16으로 모델을 저장
    • 속도가 매우 느려질수 있음
    • 하지만 이걸 해줘야... 학습 뻑날 때 ckpt부터 resume 학습이 가능함

NVME Offload

  • ZeRO-3은 nvme SSD에 offload하는 기능을 추가했다.
  • 아래와 같이 local device의 nvme에 offload할 수 있다!
  • 물론 nvme아니라 HDD나 일반 SSD에도 offload는 가능하지만... 속도가 1/10이하로 떨어짐
    • NVME: 3.5GB/s, SSD: 0.5GB/s, HDD: 0.1-0.2GB
{ "zero_optimization": { "stage": 3, "offload_optimizer": { "device": "nvme", "nvme_path": "/local_nvme", "pin_memory": true, "buffer_count": 4, "fast_init": false }, "offload_param": { "device": "nvme", "nvme_path": "/local_nvme", "pin_memory": true, "buffer_count": 5, "buffer_size": 1e8, "max_in_cpu": 1e9 } "aio": { "block_size": 262144, "queue_depth": 32, "thread_count": 1, "single_submit": false, "overlap_events": true } "overlap_comm": true, "contiguous_gradients": true, "sub_group_size": 1e14, "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_fp16_weights_on_model_save": true }, }
  • aio 값은 nvme나 컴퓨터 상태에 따라서 달라진다.

ZeRO-3을 위한 '온전한' ds_config.json 파일

  • 모든걸 auto로 맡겨버리자!
{ "fp16": { "enabled": "auto", "loss_scale": 0, "loss_scale_window": 1000, "initial_scale_power": 16, "hysteresis": 2, "min_loss_scale": 1 }, "optimizer": { "type": "AdamW", "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" } }, "scheduler": { "type": "WarmupLR", "params": { "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto" } }, "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": 1e14, "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_fp16_weights_on_model_save": true }, "gradient_accumulation_steps": "auto", "gradient_clipping": "auto", "steps_per_print": 2000, "train_batch_size": "auto", "train_micro_batch_size_per_gpu": "auto", "wall_clock_breakdown": false }

ZeRO-2 vs ZeRO-3

  • ZeRO-2가 약간 더 빠르다고 함
  • ZeRO-3에서 Model 자르는 것 때문에 모델 Scatter과정에서 속도 저하가 있음
    • ZeRO-2에서는 GPU에 model params가 있지만, ZeRO-3에서는 Model params도 Offload 한다.
  • 엄청 많은 GPU쓸거 아니면 ZeRO-2써도 된다고..?

Checkpoint Save & Load

  • DeepSpeed는 기본적으로 FP32로 모델 ckpt 저장 (커스텀 형식)
  • global_step*/*optim_states.pt 형식으로 저장
  • Deepspeed로 학습 & DeepSpeed로 load하면 전혀 문제 없음
  • 근데 PyTorch만으로 load하려고 할 때 문제가 된다. → fp16 pytorch_model.bin 파일로 저장 해야 함
    • ZeRO-2에서는 알아서 저장 잘 해준다.
    • ZeRO-3에서는 앞서 언급한 것 처럼, "stage3_gather_fp16_weights_on_model_save": true 옵션이 되어있어야 저장됨(단, 느림.)
    • 근데 이렇게 하는 것보다.. DeepSpeed로 저장된걸 FP32로 추출하는게 낫다.
      • 만약 output_dir/checkpoint-1/ 폴더 안에 DeepSpeed 파일들이 있다면..
      • python zero_to_fp32.py global_step1 pytorch_model.bin

그 외에...

  • Gradient Clipping, Accumulation도 있지만 다들 잘 안쓰고...
  • AMP를 Apex나 Native로 쓸수 있지만 굳이....? fp16으로 다 하는게 나을 듯
  • ZeRO-Infinity도 새로 나왔지만.. ZeRO-3에 대충 다 통합되는 느낌. 따로 설정 안해도 괜찮음.
  • Trillion size model 이야기도 있지만.... 설마......;;
  • deepspeed 커맨드가 에러 없이 죽는다 = OS가 Memory 배정해주다가 뻗었다 = CPU Offload 대신 NVME offload로 실험해보자
  • 1Bit-adam은 pip install~로 못쓴다. 결국 source install 해줘야 함.

Huggingface + FairScale 🤔

💡
ZeRO/DeepSpeed 작성하다보니 FairScale 써야하나? 싶은 의문이...

설치

  • 설치는 간단하게 아래와 같이 할 수 있다.
pip install transformers[fairscale]

가장 간단한 사용 예시

python -m torch.distributed.launch --nproc_per_node=2 examples/pytorch/translation/run_translation.py \ --model_name_or_path t5-small --per_device_train_batch_size 1 \ --output_dir output_dir --overwrite_output_dir \ --do_train --max_train_samples 500 --num_train_epochs 1 \ --dataset_name wmt16 --dataset_config "ro-en" \ --source_lang en --target_lang ro \ --fp16 --sharded_ddp simple
  • PyTorch의 내장 Distributed와 동일한 Launcher를 사용한다.
  • TPU 지원 안함!
  • --fp16 지원 됨.
  • --sharded_ddp simple 옵션 켜면 Vram 여유가 늘어서 → 더 큰 Batch size 사용 가능
    • --sharded_ddp zero_dp_2 or --sharded_ddp zero_dp_3 로 ZeRO 사용도 가능하다.
    • 근데 두개를 굳이 같이 쓸 이유가 있을까....???
      • DeepSpeed Launcher 대신 PyTorch Distributed만 쓰지만, ZeRO Optimizer를 쓰고 싶다, 인 경우일까?

GPU 2대로 하는 간단한 예시

python -m torch.distributed.launch --nproc_per_node=2 examples/pytorch/translation/run_translation.py \ --model_name_or_path t5-small --per_device_train_batch_size 1 \ --output_dir output_dir --overwrite_output_dir \ --do_train --max_train_samples 500 --num_train_epochs 1 \ --dataset_name wmt16 --dataset_config "ro-en" \ --source_lang en --target_lang ro \ --fp16 --sharded_ddp zero_dp_2
  • CPU Offload를 활성화 시키려면 아래와 같이 --shared_ddp 에 추가해주면 된다.
    • --sharded_ddp "zero_dp_2 cpu_offload"
    • 단 CPU Offload는 --fp16 옵션이 필수!