DeepSpeed 파인튜닝 시 CUDA OOM 오류 디버깅 마스터: 메모리 프로파일링, 최적화 기법, 그리고 코드 예제
DeepSpeed를 이용한 대규모 모델 파인튜닝은 뛰어난 성능을 제공하지만, 종종 CUDA Out-of-Memory (OOM) 오류라는 난관에 부딪히게 됩니다. 이 글에서는 DeepSpeed 환경에서 OOM 오류를 효과적으로 디버깅하고 해결하기 위한 실용적인 접근 방식, 메모리 프로파일링 도구 활용법, 그리고 다양한 최적화 기법을 코드 예제와 함께 자세히 다룹니다. 이제 더 이상 OOM 오류에 좌절하지 마세요!
1. The Challenge / Context
최근 몇 년간, 자연어 처리(NLP) 모델의 크기는 기하급수적으로 증가했습니다. BERT, GPT-3와 같은 거대 모델들은 놀라운 성능을 보여주지만, 이를 파인튜닝하기 위해서는 막대한 컴퓨팅 자원이 필요합니다. DeepSpeed는 이러한 대규모 모델을 효율적으로 학습시키기 위한 강력한 프레임워크이지만, GPU 메모리 부족으로 인한 CUDA OOM 오류는 여전히 흔하게 발생합니다. 특히, 고가의 GPU 자원을 효율적으로 활용해야 하는 상황에서 OOM 오류는 연구 및 개발 생산성을 저하시키는 주요 요인이 됩니다. 따라서 OOM 오류를 신속하게 진단하고 해결하는 능력은 DeepSpeed를 사용하는 모든 개발자에게 필수적입니다.
2. Deep Dive: 메모리 프로파일링 (Memory Profiling)
OOM 오류를 해결하기 위한 첫 번째 단계는 메모리 사용량을 정확하게 프로파일링하는 것입니다. NVIDIA의 Nsight Systems 및 torch.cuda.memory_summary()는 메모리 사용량을 분석하는 데 유용한 도구입니다. Nsight Systems는 CPU, GPU 활동을 모두 프로파일링할 수 있는 강력한 도구이며, torch.cuda.memory_summary()는 PyTorch 환경에서 간단하게 GPU 메모리 사용량을 확인할 수 있습니다. DeepSpeed는 자체적으로 메모리 프로파일링 기능(deepspeed.profiling.flops_profiler)을 제공하지만, 더 자세한 분석을 위해서는 외부 도구를 함께 사용하는 것이 좋습니다.
3. Step-by-Step Guide / Implementation
Step 1: torch.cuda.memory_summary()를 이용한 간단한 메모리 사용량 확인
학습 코드 중간중간에 torch.cuda.memory_summary()를 삽입하여 특정 시점의 GPU 메모리 사용량을 확인합니다. 특히, OOM 오류가 발생하는 지점 주변에 집중적으로 삽입하여 메모리 사용량이 급증하는 부분을 찾아냅니다.
import torch
# ... (학습 코드)
print(torch.cuda.memory_summary())
# ... (학습 코드)
Step 2: Nsight Systems를 이용한 상세 메모리 프로파일링
Nsight Systems를 사용하여 전체 학습 과정을 프로파일링합니다. Nsight Systems GUI를 통해 CPU, GPU 메모리 사용량, CUDA 커널 실행 시간 등 다양한 정보를 시각적으로 확인할 수 있습니다. OOM 오류가 발생하는 시점의 메모리 할당 패턴을 분석하여 메모리 누수 또는 비효율적인 메모리 사용 부분을 식별합니다.
Nsight Systems를 실행하는 방법은 다음과 같습니다.
nsys profile -o profile_output.qdrep python your_training_script.py --deepspeed_config ds_config.json
프로파일링이 완료되면 profile_output.qdrep 파일을 Nsight Systems GUI에서 열어 분석합니다.
Step 3: Gradient Accumulation Steps 조정
Gradient accumulation은 배치 크기를 늘리는 효과를 내면서도 GPU 메모리 제약 하에서 학습을 가능하게 하는 기법입니다. Gradient accumulation steps를 늘리면, 더 많은 미니배치를 처리한 후 가중치를 업데이트하므로, 각 업데이트에 필요한 GPU 메모리를 줄일 수 있습니다.
DeepSpeed 설정 파일 (ds_config.json)에서 gradient_accumulation_steps 값을 조정합니다.
{
"train_batch_size": 32,
"gradient_accumulation_steps": 4,
"fp16": {
"enabled": true
}
}
위 설정은 실제 배치 크기를 32 * 4 = 128로 설정하는 것과 같습니다. gradient_accumulation_steps 값을 늘리면 OOM 오류를 해결할 수 있지만, 학습 속도가 느려질 수 있으므로 적절한 값을 찾아야 합니다.
Step 4: Mixed Precision Training (FP16) 활용
Mixed precision training은 FP32 (32비트 부동 소수점) 대신 FP16 (16비트 부동 소수점)을 사용하여 메모리 사용량을 줄이는 기법입니다. DeepSpeed는 FP16 training을 간편하게 지원합니다. DeepSpeed 설정 파일에서 FP16을 활성화합니다.
{
"fp16": {
"enabled": true,
"loss_scale": 0,
"initial_scale_power": 32,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
}
}
FP16 training은 일반적으로 학습 속도를 향상시키고 메모리 사용량을 줄이지만, 수치적 안정성 문제가 발생할 수 있습니다. 따라서 DeepSpeed에서 제공하는 Loss Scaling 기법을 함께 사용하여 수치적 안정성을 확보해야 합니다.
Step 5: Activation Checkpointing (Gradient Checkpointing) 적용
Activation checkpointing은 순전파 과정에서 중간 활성화 값을 저장하지 않고, 역전파 과정에서 다시 계산하여 메모리 사용량을 줄이는 기법입니다. DeepSpeed는 activation checkpointing을 지원하며, 모델의 특정 레이어에 적용할 수 있습니다.
import torch
import deepspeed
model = ... # your model
model = deepspeed.checkpointing.make_activation_checkpointing(model)
# or, apply to specific layers
# from deepspeed.checkpointing import checkpoint
# def forward(self, x):
# x = self.layer1(x)
# x = checkpoint(self.layer2, x)
# x = self.layer3(x)
# return x
Activation checkpointing은 메모리 사용량을 줄이지만, 추가적인 계산 비용이 발생하므로, 학습 속도가 느려질 수 있습니다. 따라서, 메모리 제약과 학습 속도 간의 균형을 고려하여 적절하게 적용해야 합니다.
Step 6: ZeRO 최적화 적용
ZeRO (Zero Redundancy Optimizer)는 모델 파라미터, gradients, optimizer states를 여러 GPU에 분산시켜 메모리 사용량을 획기적으로 줄이는 DeepSpeed의 핵심 기능입니다. ZeRO는 여러 단계로 구성되어 있으며, 각 단계별로 메모리 절약 효과가 다릅니다. 일반적으로 ZeRO-2 또는 ZeRO-3를 많이 사용합니다.
DeepSpeed 설정 파일에서 ZeRO를 활성화합니다.
{
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"reduce_scatter": true,
"contiguous_gradients": true,
"allgather_partitions": true
}
}
위 설정은 ZeRO-2를 활성화하고, optimizer states와 parameters를 CPU로 오프로드하는 것을 의미합니다. CPU 오프로드는 GPU 메모리 사용량을 줄이지만, 학습 속도가 느려질 수 있습니다. 필요에 따라 GPU 오프로드를 고려할 수 있습니다.
4. Real-world Use Case / Example
최근 BERT-large 모델을 특정 도메인 데이터셋으로 파인튜닝하는 프로젝트에서 OOM 오류가 빈번하게 발생했습니다. 초기에는 배치 크기를 줄이고 FP16 training을 적용했지만, 여전히 OOM 오류가 발생했습니다. Nsight Systems를 이용하여 메모리 프로파일링을 수행한 결과, Transformer 레이어의 attention 연산에서 많은 메모리가 사용되는 것을 확인했습니다. Activation checkpointing을 Transformer 레이어에 적용하고 ZeRO-2를 활성화한 결과, OOM 오류 없이 파인튜닝을 완료할 수 있었습니다. 뿐만 아니라, FP16 training과 ZeRO 최적화 덕분에 학습 속도도 20% 향상되었습니다.
5. Pros & Cons / Critical Analysis
- Pros:
- DeepSpeed는 대규모 모델 학습을 위한 강력한 도구입니다.
- ZeRO 최적화를 통해 GPU 메모리 제약을 극복할 수 있습니다.
- FP16 training, Activation Checkpointing 등의 기법을 통해 메모리 사용량을 추가적으로 줄일 수 있습니다.
- 상세한 메모리 프로파일링 도구를 통해 OOM 오류의 원인을 정확하게 파악할 수 있습니다.
- Cons:
- DeepSpeed 설정이 복잡할 수 있습니다.
- FP16 training은 수치적 안정성 문제를 야기할 수 있습니다.
- Activation checkpointing은 추가적인 계산 비용을 발생시킬 수 있습니다.
- CPU 오프로드는 학습 속도를 저하시킬 수 있습니다.
- 최적의 성능을 얻기 위해서는 다양한 설정을 실험적으로 조정해야 합니다.
6. FAQ
- Q: DeepSpeed를 처음 사용하는데, 어떤 설정을 먼저 시도해야 할까요?
A: 먼저fp16.enabled=true를 설정하여 FP16 training을 활성화하고,gradient_accumulation_steps를 적절히 조절하여 배치 크기를 늘려보세요. 그 다음 ZeRO-2 또는 ZeRO-3를 활성화하는 것을 추천합니다. - Q: Nsight Systems 외에 다른 메모리 프로파일링 도구가 있나요?
A: PyTorch Profiler (torch.profiler)도 사용할 수 있습니다. 하지만, DeepSpeed 환경에서는 Nsight Systems가 더 상세한 정보를 제공하므로 더 유용합니다. - Q: ZeRO-3는 항상 ZeRO-2보다 좋은가요?
A: ZeRO-3는 ZeRO-2보다 더 많은 메모리를 절약할 수 있지만, 통신 오버헤드가 더 커서 학습 속도가 느려질 수 있습니다. GPU 개수가 많고, 메모리 제약이 매우 심각한 경우에 ZeRO-3를 사용하는 것이 좋습니다.
7. Conclusion
DeepSpeed는 대규모 모델 학습을 위한 필수적인 도구이지만, OOM 오류는 피할 수 없는 과제입니다. 이 글에서 제시된 메모리 프로파일링, 최적화 기법들을 활용하여 OOM 오류를 효과적으로 디버깅하고 해결함으로써, DeepSpeed를 더욱 효율적으로 활용할 수 있을 것입니다. 지금 바로 코드를 적용해보고, DeepSpeed 커뮤니티에 참여하여 경험을 공유하고 함께 성장해 나가세요!
