DeepSpeed 파인튜닝 시 CUDA OOM 문제 해결: 메모리 효율 극대화 전략
대규모 모델 파인튜닝 시 DeepSpeed를 사용하면서 CUDA OOM (Out Of Memory) 에러를 겪는 개발자를 위한 가이드입니다. 이 글에서는 ZeRO, 데이터 병렬 처리, Gradient Accumulation 등 DeepSpeed의 핵심 기능을 활용하여 GPU 메모리 사용량을 극적으로 줄이는 전략을 소개하고, 실제 설정 예시와 팁을 제공합니다. 작은 변경만으로도 OOM 문제를 해결하고, 더 큰 모델 또는 더 큰 배치 사이즈로 파인튜닝이 가능하게 됩니다.
1. The Challenge / Context
최근 몇 년 동안 NLP와 컴퓨터 비전 분야에서 모델의 크기가 기하급수적으로 증가했습니다. 이러한 거대 모델은 뛰어난 성능을 제공하지만, 파인튜닝 과정에서 엄청난 컴퓨팅 자원을 요구하며, 특히 GPU 메모리 부족 문제 (CUDA OOM)를 빈번하게 발생시킵니다. DeepSpeed는 이러한 문제를 해결하기 위한 강력한 솔루션이지만, 올바르게 설정하지 않으면 여전히 OOM 에러에 직면할 수 있습니다. 특히 자원이 제한적인 환경에서 개인 개발자나 소규모 팀은 이러한 문제를 해결하는 데 어려움을 겪습니다. 이 글은 DeepSpeed의 다양한 기능을 효과적으로 활용하여 CUDA OOM 문제를 극복하고, 메모리 효율적인 파인튜닝을 가능하게 하는 실질적인 전략을 제시합니다.
2. Deep Dive: DeepSpeed
DeepSpeed는 Microsoft에서 개발한 딥러닝 최적화 라이브러리입니다. 주로 대규모 모델 학습을 위해 설계되었으며, 다양한 메모리 최적화 및 병렬 처리 기술을 제공합니다. DeepSpeed의 핵심 기능은 다음과 같습니다.
- ZeRO (Zero Redundancy Optimizer): 모델 파라미터, 그래디언트, 옵티마이저 상태를 여러 GPU에 분산시켜 메모리 사용량을 획기적으로 줄입니다. ZeRO는 Stage 1, Stage 2, Stage 3의 세 가지 레벨로 나뉘며, 각 레벨은 메모리 효율성과 통신 오버헤드 간의 균형을 제공합니다.
- 데이터 병렬 처리 (Data Parallelism): 데이터를 여러 GPU에 분할하여 각 GPU에서 독립적으로 학습을 수행합니다. 이를 통해 전체 학습 시간을 단축할 수 있습니다. DeepSpeed는 데이터 병렬 처리를 효율적으로 관리하고, 통신 오버헤드를 최소화합니다.
- Gradient Accumulation: 작은 배치 사이즈로 여러 번의 forward/backward pass를 수행한 후 그래디언트를 누적하여 업데이트합니다. 이는 큰 배치 사이즈를 사용하는 것과 유사한 효과를 내면서도 GPU 메모리 사용량을 줄일 수 있습니다.
- Mixed Precision Training (FP16): half-precision (FP16) 부동 소수점 연산을 사용하여 메모리 사용량을 줄이고 학습 속도를 향상시킵니다.
- Offload: CPU 또는 NVMe SSD로 모델 파라미터, 그래디언트, 옵티마이저 상태를 오프로드하여 GPU 메모리 부담을 줄입니다.
DeepSpeed는 PyTorch와 통합되어 있으며, 사용하기 쉽고 유연합니다. 하지만 효과적인 사용을 위해서는 각 기능의 작동 방식과 장단점을 이해해야 합니다.
3. Step-by-Step Guide / Implementation
이제 DeepSpeed를 사용하여 CUDA OOM 문제를 해결하고 메모리 효율적인 파인튜닝을 수행하는 단계별 가이드를 살펴보겠습니다.
Step 1: DeepSpeed 설치 및 환경 설정
먼저 DeepSpeed를 설치해야 합니다. pip를 사용하여 간단하게 설치할 수 있습니다.
pip install deepspeedCUDA 및 PyTorch가 올바르게 설정되었는지 확인하십시오. 또한, DeepSpeed는 분산 학습을 위해 NCCL과 같은 통신 라이브러리가 필요합니다. 필요한 드라이버와 라이브러리가 설치되어 있는지 확인하세요.
Step 2: DeepSpeed 설정 파일 (JSON) 작성
DeepSpeed는 JSON 형식의 설정 파일을 사용하여 다양한 옵션을 구성합니다. 다음은 기본적인 설정 파일 예시입니다.
{
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"gradient_accumulation_steps": "auto",
"optimizer": {
"type": "AdamW",
"params": {
"lr": 1e-5,
"weight_decay": 0.01
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": 0,
"warmup_max_lr": 1e-5,
"warmup_num_steps": 1000
}
},
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16,
"hysteresis": 2,
"min_loss_scale": 1
},
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto"
},
"gradient_clipping": 1.0,
"steps_per_print": 2000,
"wall_clock_breakdown": false
}각 옵션에 대한 설명은 다음과 같습니다.
- train_batch_size: 전체 배치 사이즈입니다. "auto"로 설정하면 DeepSpeed가 자동으로 최적의 값을 찾습니다.
- train_micro_batch_size_per_gpu: 각 GPU에서 사용하는 마이크로 배치 사이즈입니다. OOM 에러가 발생하는 경우 이 값을 줄여야 합니다.
- gradient_accumulation_steps: 그래디언트 누적 단계 수입니다. 마이크로 배치 사이즈를 줄이는 대신 이 값을 늘려 큰 배치 사이즈와 유사한 효과를 낼 수 있습니다.
- fp16: Mixed Precision Training을 활성화합니다. "auto"로 설정하면 DeepSpeed가 자동으로 활성화 여부를 결정합니다.
- zero_optimization: ZeRO 최적화를 설정합니다. Stage 2 또는 Stage 3를 사용하는 것이 좋습니다. offload_optimizer 및 offload_param을 사용하여 CPU 또는 NVMe SSD로 파라미터를 오프로드할 수 있습니다.
Step 3: DeepSpeedEngine 초기화 및 학습 루프 수정
PyTorch 학습 스크립트에서 DeepSpeedEngine을 초기화하고 학습 루프를 수정해야 합니다.
import deepspeed
import torch
# 모델, 옵티마이저, 데이터 로더 초기화
model = ...
optimizer = ...
train_dataloader = ...
# DeepSpeedEngine 초기화
model, optimizer, train_dataloader, _ = deepspeed.initialize(
config_params=config_json,
model=model,
optimizer=optimizer,
train_dataloader=train_dataloader
)
# 학습 루프
for epoch in range(num_epochs):
for step, batch in enumerate(train_dataloader):
# 데이터 준비
inputs = batch[0].to(model.device)
labels = batch[1].to(model.device)
# Forward pass
outputs = model(inputs)
loss = criterion(outputs, labels)
# Backward pass 및 파라미터 업데이트
model.backward(loss)
model.step()
deepspeed.initialize() 함수는 모델, 옵티마이저, 데이터 로더를 DeepSpeedEngine으로 래핑합니다. model.backward() 및 model.step() 함수를 사용하여 backward pass 및 파라미터 업데이트를 수행합니다. 모델을 GPU로 옮기는 과정 (model.to(device)) 은 DeepSpeed에서 자동으로 처리하므로 필요하지 않습니다. DeepSpeedEngine은 자동으로 GPU 메모리 관리를 수행하며, OOM 에러를 방지합니다.
Step 4: ZeRO Stage 조정
ZeRO Stage는 메모리 효율성에 큰 영향을 미칩니다. 일반적으로 Stage 2 또는 Stage 3를 사용하는 것이 좋습니다. Stage 3는 가장 높은 메모리 효율성을 제공하지만, 통신 오버헤드가 더 클 수 있습니다. 따라서, 모델 크기와 GPU 메모리 용량에 따라 적절한 Stage를 선택해야 합니다. 설정 파일에서 zero_optimization.stage 값을 변경하여 ZeRO Stage를 조정할 수 있습니다.
"zero_optimization": {
"stage": 3,
...
}Step 5: Gradient Accumulation Step 조정
GPU 메모리가 부족한 경우 gradient_accumulation_steps 값을 늘려 마이크로 배치 사이즈를 줄일 수 있습니다. 예를 들어, gradient_accumulation_steps를 2로 설정하면 마이크로 배치 사이즈를 절반으로 줄일 수 있습니다. 설정 파일에서 다음과 같이 값을 변경합니다.
"gradient_accumulation_steps": 2이 값을 늘리면 학습 시간이 늘어나지만, GPU 메모리 사용량을 줄일 수 있습니다.
Step 6: Offload 활성화
GPU 메모리가 여전히 부족한 경우, 모델 파라미터 및 옵티마이저 상태를 CPU 또는 NVMe SSD로 오프로드할 수 있습니다. 설정 파일에서 offload_optimizer 및 offload_param 옵션을 사용하여 오프로드를 활성화합니다.
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
...
}device 옵션을 "cpu" 또는 "nvme"로 설정하여 오프로드 장치를 지정할 수 있습니다. NVMe SSD를 사용하는 경우, 학습 속도를 향상시킬 수 있습니다.
4. Real-world Use Case / Example
개인적으로 GPT-2 Large 모델을 단일 24GB GPU에서 파인튜닝해야 하는 프로젝트가 있었습니다. 초기 설정에서는 CUDA OOM 에러가 지속적으로 발생했습니다. DeepSpeed를 사용하여 위에서 설명한 단계를 적용한 결과, 다음과 같은 효과를 얻을 수 있었습니다.
- ZeRO Stage 2 활성화: 메모리 사용량을 약 30% 줄였습니다.
- Gradient Accumulation Step 4로 설정: 마이크로 배치 사이즈를 줄이고, 메모리 사용량을 추가로 20% 줄였습니다.
- FP16 활성화: 메모리 사용량을 절반으로 줄였습니다.
결과적으로 OOM 에러 없이 GPT-2 Large 모델을 성공적으로 파인튜닝할 수 있었고, 실험 결과는 만족스러웠습니다. 중요한 점은, DeepSpeed의 여러 기능들을 조합하여 적절한 설정을 찾는 것이 중요하다는 것입니다.
5. Pros & Cons / Critical Analysis
- Pros:
- 메모리 효율성 극대화: DeepSpeed는 ZeRO, Gradient Accumulation, Offload 등의 기능을 통해 GPU 메모리 사용량을 획기적으로 줄일 수 있습니다.
- 대규모 모델 학습 가능: DeepSpeed를 사용하면 단일 GPU 또는 소규모 클러스터에서도 거대 모델을 학습할 수 있습니다.
- 사용 용이성: PyTorch와 통합되어 있으며, 간단한 설정 파일과 몇 줄의 코드 변경만으로 DeepSpeed를 적용할 수 있습니다.
- Cons:
- 설정 복잡성: DeepSpeed의 다양한 옵션을 올바르게 구성하는 데 시간이 걸릴 수 있습니다. 최적의 설정을 찾기 위해 실험과 튜닝이 필요합니다.
- 통신 오버헤드: ZeRO Stage 3와 같은 일부 기능은 통신 오버헤드를 증가시킬 수 있습니다.
- 디버깅 어려움: DeepSpeed 환경에서 발생하는 에러는 일반적인 PyTorch 환경보다 디버깅이 더 어려울 수 있습니다. DeepSpeed 관련 문서와 커뮤니티 지원을 적극적으로 활용해야 합니다.
6. FAQ
- Q: DeepSpeed를 사용하지 않고는 대규모 모델을 파인튜닝할 수 없나요?
A: 반드시 그런 것은 아닙니다. 그러나 DeepSpeed는 메모리 효율성을 극대화하여 더 큰 모델 또는 더 큰 배치 사이즈로 파인튜닝할 수 있도록 도와줍니다. 다른 방법으로는 모델 병렬 처리 (Model Parallelism) 등이 있지만, DeepSpeed가 더 간단하고 효과적인 경우가 많습니다. - Q: ZeRO Stage를 어떻게 선택해야 하나요?
A: 일반적으로 Stage 2 또는 Stage 3를 사용하는 것이 좋습니다. Stage 3는 가장 높은 메모리 효율성을 제공하지만, 통신 오버헤드가 더 클 수 있습니다. 모델 크기와 GPU 메모리 용량에 따라 적절한 Stage를 선택해야 합니다. 작은 모델에는 Stage 1이 적합할 수 있습니다. - Q: Offload를 사용하면 학습 속도가 느려지나요?
A: Offload를 사용하면 학습 속도가 느려질 수 있습니다. 특히 CPU로 오프로드하는 경우, 데이터 전송 속도가 GPU 메모리보다 느리기 때문에 성능 저하가 발생할 수 있습니다. NVMe SSD를 사용하는 경우, 성능 저하를 최소화할 수 있습니다. - Q: DeepSpeed와 Horovod의 차이점은 무엇인가요?
A: Horovod는 분산 학습을 위한 프레임워크이지만, DeepSpeed는 메모리 최적화 및 대규모 모델 학습에 특화되어 있습니다. DeepSpeed는 ZeRO, Gradient Accumulation, Offload 등의 기능을 통해 GPU 메모리 사용량을 획기적으로 줄일 수 있습니다. Horovod는 주로 데이터 병렬 처리에 사용됩니다.
7. Conclusion
DeepSpeed는 CUDA OOM 문제를 해결하고 메모리 효율적인 파인튜닝을 가능하게 하는 강력한 도구입니다. 이 글에서 소개한 전략과 팁을 활용하여 DeepSpeed를 효과적으로 사용하고, 더 큰 모델 또는 더 큰 배치 사이즈로 파인튜닝을 수행할 수 있기를 바랍니다. 지금 바로 DeepSpeed를 설치하고, 당신의 모델을 더 크게 만들어 보세요! DeepSpeed 공식 문서를 참조하여 더 자세한 정보를 얻을 수 있습니다.
