아래 글은 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가지다.
- Batch: batch를 각 GPU로 쪼개서 각 GPU에서 학습하자
- Optimizer State: 해당 Batch를 위한 Optimizer만 가져오기
- Gradient: backward 위한 Gradient를 해당 batch만 쓰자
- Model weight(parameters): 모델조차 쪼개버리자, Model parallel
- Transformer 모델 기준 Head 단위로 모델을 쪼갠다고 이해하면 편하다.
- Transformer 모델 기준 Layer 단위로 모델을 쪼갠다고 이해하면 편하다.
4-1. Model Parallel
4-2. Pipeline Parallel
- 쪼개지 않고 하는 방법도 물론 있다.
- 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]
- 혹은 수동으로 아래와 같이 설치해줄 수도 있다.
정말 공식 가이드대로 저 위 한 줄만 하면 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하는 방식
- 여기서
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 🤔
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_2or--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"
--fp16 옵션이 필수!