💡
아래 글은 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
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
- Zero-2 (aka Zero-Offload), up to 13 Billion on 1 GPU
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_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
옵션이 필수!
- 단 CPU Offload는