DeepSpeed 깊이 파헤치기
멀티 GPU 하드웨어에서 최신 AI 모델을 훈련하는 일을 처음 시작하면 꽤 막막하게 느껴질 수 있습니다. 그러니 핵심 개념과 Microsoft DeepSpeed 라이브러리를 시작하는 방법을 하나씩 살펴봅시다. 다음은 짧은 DeepSpeed 튜토리얼입니다. 이 주제는 깊은(하하) 주제이므로 겉핥기 정도만 다루겠습니다.
DeepSpeed의 Core 기술은 ZeRO라고 불리는 Zero Redundancy Optimizer로, 모델 상태를 여러 GPU에 분할하여 대규모 모델의 학습을 가능하게 합니다. ZeRO는 여러 단계로 작동합니다.
- ZeRO-1: 옵티마이저 상태를 여러 GPU에 분할합니다
- ZeRO-2: 그래디언트와 옵티마이저 상태를 분할합니다
- ZeRO-3: 모델 매개변수, 그래디언트, 옵티마이저 상태를 분할합니다
ZeRO-1: 옵티마이저 상태 분할
ZeRO-1은 Adam의 모멘텀과 분산 같은 옵티마이저 상태를 여러 GPU에 분할하는 데 초점을 맞춥니다. 이 구성에서는 각 GPU가 옵티마이저 상태 중 자신의 몫만 유지하므로, Adam 옵티마이저를 사용할 때 메모리 사용량이 약 4배 줄어듭니다. 이 단계는 통신 오버헤드가 거의 없으며, 옵티마이저 상태가 주된 메모리 병목인 상황에서 특히 효과적입니다.
ZeRO-2: 그래디언트 분할
ZeRO-2는 ZeRO-1의 최적화를 바탕으로 분할 대상을 GPU 간 그래디언트까지 확장합니다. 각 GPU는 자신에게 지정된 매개변수의 그래디언트만 저장하므로, 표준 데이터 병렬 학습에 비해 메모리 사용량이 약 8배 줄어듭니다. 이 단계는 어느 정도의 통신 오버헤드가 생기지만, 메모리 절감과 성능 사이에서 훌륭한 균형을 제공합니다.
ZeRO-3: 매개변수 분할
ZeRO-3는 가장 공격적인 최적화 단계로, ZeRO-1과 ZeRO-2의 모든 이점을 포함하면서 모델 매개변수까지 여러 GPU에 분할합니다. 이 구성에서는 각 GPU가 자신에게 할당된 매개변수만 유지하므로, 메모리 사용량이 16배 이상 줄어듭니다. 통신 오버헤드는 더 커지지만, ZeRO-3는 그렇지 않으면 GPU 메모리에 들어가지 않는 초대형 모델을 학습하는 데 필수적이며, 수천억 개의 매개변수를 가진 모델의 학습을 가능하게 합니다.
설치
DeepSpeed는 PyPI에서 설치합니다.
pip install deepspeed
최적의 성능을 위해서는 하드웨어에 맞게 소스에서 설치하는 것을 권장합니다.
git clone https://github.com/microsoft/DeepSpeed.git
cd DeepSpeed
pip install -e .
훈련하기 전에 DeepSpeed의 내장 도구를 사용하여 GPU VRAM 메모리 요구량을 추정하십시오.
from transformers import AutoModel
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live
model = AutoModel.from_pretrained("meta-llama/Llama-3.2-1B")
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=1, num_nodes=1)
이 도구는 여러 구성에 대한 예상 메모리 요구량을 출력합니다. 출력은 다음과 같습니다.
Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 8 GPUs per node.
SW: Model with 2851M total params.
per CPU | per GPU | Options
127.48GB | 5.31GB | offload_optimizer=cpu
127.48GB | 15.93GB | offload_optimizer=none
구성 파일
DeepSpeed는 보통 ds_config.json이라는 이름의 JSON 파일로 구성합니다. 다음은 기본적인 ZeRO-3 구성입니다.
{
"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,
"stage3_gather_16bit_weights_on_model_save": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto"
}
구성에서 가장 중요한 부분은 ZeRO 적용 방식을 제어하는 zero_optimization 섹션입니다. 이 예제에서는 옵티마이저 상태와 매개변수 상태를 CPU 메모리로 오프로드한 상태에서 ZeRO-3를 사용합니다. 오프로드는 offload_optimizer 섹션과 offload_param 섹션으로 지정합니다. 이 구성은 모델 학습을 여러 GPU에 어떻게 분산하고 그 결과를 어떻게 다시 모을지를 기술합니다.
이제 train.py 파일에서 HuggingFace Trainer와 함께 DeepSpeed를 사용하려면, --deepspeed 플래그와 함께 DeepSpeed 구성 JSON 파일을 가리키는 구성 파일 경로를 추가하기만 하면 됩니다.
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir="output",
deepspeed="ds_config.json", # DeepSpeed config file
fp16=True,
# Put other training arguments here
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset,
# Put other trainer arguments here
)
trainer.train()
학습 시작하기
여러 GPU가 장착된 단일 서버를 사용한다면 다음 명령으로 학습을 시작할 수 있습니다.
deepspeed --num_gpus=8 train.py \
--num_nodes=1 \
--num_gpus=8 \
--deepspeed ds_config.json \
--other_training_args
다중 노드 환경을 사용한다면 다음 명령으로 학습을 시작할 수 있습니다.
deepspeed --hostfile=hostfile \
--num_nodes=2 \
--num_gpus=8 \
train.py \
--deepspeed ds_config.json \
--other_training_args
여기서 로컬 파일 hostfile에는 호스트 이름 목록과 호스트별 GPU 개수가 들어 있습니다.
worker-1 slots=8
worker-2 slots=8
DeepSpeed를 처음 사용할 때는 더 단순한 구성부터 시작하는 것이 가장 좋습니다. 먼저 ZeRO-2 최적화를 사용하고, 추가적인 메모리 절감이 필요하다고 판단될 때에만 더 공격적인 ZeRO-3로 넘어가십시오. GPU 메모리 제약을 관리할 때 CPU 오프로드가 도움이 될 수 있지만, 추가 오버헤드를 유발하므로 GPU 메모리만으로는 부족한 경우에만 활성화해야 합니다. 그레이디언트 체크포인팅은 추가 계산 시간을 대가로 메모리 사용량을 줄이는 또 다른 방법을 제공합니다.
배치 크기 역시 신중한 조정이 필요한 중요한 고려 사항입니다. 작은 배치 크기부터 시작하여 메모리 사용량을 면밀히 모니터링하면서 점진적으로 늘려, 사용하는 하드웨어에 맞는 최적의 균형점을 찾으십시오. 마지막으로, 가능하면 fp16 또는 bf16을 사용하는 혼합 정밀도 학습을 활용하십시오. 대부분의 경우 모델 품질을 희생하지 않고도 상당한 성능 향상을 얻을 수 있습니다.
문제 해결
멀티 GPU 구성에서 모델을 실행하는 것은 하드웨어 장애와 소프트웨어 장애 모두로 인해 매우 까다로울 수 있습니다. 가장 흔한 장애 유형(반드시 마주치게 될 것들)은 다음과 같습니다.
- NaN 손실 값.
- 메모리 부족 오류.
- CPU 오프로드 실패.
- 무작위 하드웨어 장애.
NaN 손실 값은 bf16으로 사전 학습된 모델을 fp16 정밀도로 사용할 때 흔히 발생하는 문제입니다. 이런 경우 bf16 또는 fp32 정밀도로 전환하면 대개 문제가 해결됩니다.
OOM 오류가 발생하면 GPU 메모리 사용량을 줄이기 위해 CPU 오프로드를 활성화하거나 ZeRO 최적화 단계를 높여 볼 수 있습니다. 다만 오프로드 기능을 사용할 때 프로세스가 아예 강제 종료된다면 CPU 메모리가 부족하다는 뜻일 수 있으므로 주의하십시오.
학습 속도 저하와 같은 성능 문제는 불필요한 CPU 오프로드가 원인인 경우가 많습니다. GPU 메모리가 모델과 학습 데이터를 처리하기에 충분하다면, CPU 오프로드를 비활성화하는 것이 일반적으로 학습 시간을 단축합니다.
GPU 메모리 사용량을 최적화할 때는 여러 ZeRO 단계를 활성화하기 전에 그레이디언트 체크포인팅부터 구현하십시오. fp16 또는 bf16을 사용하는 혼합 정밀도 학습은 모델 품질을 유지하면서도 메모리 요구량을 크게 줄일 수 있습니다. 학습 중에는 GPU 메모리 nvidia-smi를 정기적으로 모니터링하는 것이 필수적입니다. CPU 또는 NVMe 오프로딩은 GPU 메모리만으로는 학습에 필요한 양을 감당할 수 없을 때에만 고려해야 합니다.
CPU 오프로딩을 구현할 때는 옵티마이저 상태부터 오프로딩하십시오. 이는 성능에 미치는 영향이 비교적 작기 때문입니다. 매개변수 오프로딩은 상당한 오버헤드를 유발하므로 반드시 필요한 경우에만 사용해야 합니다. 오프로딩을 효과적으로 처리할 수 있도록 시스템에 충분한 CPU RAM(일반적으로 모델 크기의 2~4배)이 있는지 확인하십시오. CPU RAM 용량을 초과하는 모델의 경우, 최후의 수단으로 NVMe 오프로드 기능의 사용을 고려하십시오.
효과적인 배치 크기 튜닝은 작은 크기에서 시작하여 학습이 안정됨에 따라 점차 늘려 가는 방식으로 진행합니다. 그래디언트 누적을 사용하면 메모리 요구량을 비례해서 늘리지 않고도 사실상 더 큰 배치 크기를 얻을 수 있습니다. 손실 스케일링을 주의 깊게 모니터링하는 것이 필수적이며, 수렴 패턴에 따라 조정해야 합니다. 일부 학습 시나리오에서는 가용 메모리 상황에 맞춰 조정되는 동적 배치 크기 설정이 도움이 될 수 있습니다.
파이프라인 병렬성은 모델 계층을 여러 GPU에 나누어 배치함으로써 모델 분산의 또 다른 차원을 제공합니다. 이 방식은 GPU당 필요한 메모리를 효과적으로 줄여 주며, 단일 장치에 들어가지 않는 매우 깊은 모델도 학습할 수 있게 해 줍니다. 파이프라인 병렬성은 다양한 ZeRO 단계와 결합하여 특정 하드웨어 환경에 맞는 최적의 학습 구성을 구현할 수 있습니다.
활성화 체크포인팅은 활성화 값을 저장하는 대신 역전파 과정에서 다시 계산함으로써 연산량과 메모리 사용량 사이의 영리한 절충안을 제공합니다. 이 기법은 학습 시의 메모리 사용량을 크게 줄일 수 있습니다
현시점에서 멀티 GPU 학습을 설정하는 일은 간단하지 않으며 많은 숙고가 필요합니다. 이는 과학이라기보다는 훨씬 더 기예에 가까우며, 많은 시행착오를 수반합니다.