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

• CPU
• Single GPU
• 1 Node, Multi GPU
• Multi Node, Multi GPU
• TPU
• TPU Pods

모델 크기 한계

• 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

• 이것과 관련해서 당근마켓 블로그에 잘 정리된 글이 있다.

  🔥PyTorch Multi-GPU 학습 제대로 하기

  딥러닝은 기본적으로 GPU 에서 학습을 합니다. Deep Neural Network는 기본적으로 매트릭스 연산을 하기 때문에 GPU를 사용할 경우 연산처리 속도가 상당히 빨라집니다. 딥러닝이 발전함에 따라 점점 네트워크의 크기도 커졌습니다. 비전 분야에서 Neural Network를 깊게 쌓는 방식이 성공하면서부터 그 이후의 딥러닝 연구에서는 대부분 큰...

  https://medium.com/daangn/pytorch-multi-gpu-%ED%95%99%EC%8A%B5-%EC%A0%9C%EB%8C%80%EB%A1%9C-%ED%95%98%EA%B8%B0-27270617936b

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

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

• 결국 대규모 모델 학습을 위해서 쪼갤 수 있는건 크게 4가지다.
  1. Batch: batch를 각 GPU로 쪼개서 각 GPU에서 학습하자

     

  2. Optimizer State: 해당 Batch를 위한 Optimizer만 가져오기
  3. Gradient: backward 위한 Gradient를 해당 batch만 쓰자
  4. Model weight(parameters): 모델조차 쪼개버리자, Model parallel

     4-1. Model Parallel

     

     • Transformer 모델 기준 Head 단위로 모델을 쪼갠다고 이해하면 편하다.

     4-2. Pipeline Parallel

     

     • Transformer 모델 기준 Layer 단위로 모델을 쪼갠다고 이해하면 편하다.
• 쪼개지 않고 하는 방법도 물론 있다.
  • Zero-2 (aka Zero-Offload), up to 13 Billion on 1 GPU

    

    • 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]

• 혹은 수동으로 아래와 같이 설치해줄 수도 있다.

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 파일을 아래와 같이 만들어서 쓰기만 해도 성능 ++

{
  "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 제한 못한다!

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하는 방식
  • 여기서 device 를 cpu 아니라 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나 컴퓨터 상태에 따라서 달라진다.
• 따라서 https://github.com/microsoft/DeepSpeed/issues/998 이슈에 달린 답 처럼 로컬에서 벤치마크 돌려본 뒤에 값을 정해주는게 성능 향상에 유리하다.

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 🤔

설치

• 설치는 간단하게 아래와 같이 할 수 있다.

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 옵션이 필수!