DeepSpeed ZeRO-3 GPU 메모리 에러 디버깅 마스터: 고급 메모리 프로파일링 및 분산 학습 최적화
DeepSpeed ZeRO-3는 거대한 모델 학습을 가능하게 하지만, GPU 메모리 에러는 흔한 문제입니다. 이 글에서는 고급 메모리 프로파일링 기술과 분산 학습 최적화 전략을 통해 ZeRO-3 환경에서 메모리 에러를 진단하고 해결하는 실질적인 방법을 제시합니다. 문제 해결 시간을 단축하고 GPU 활용률을 극대화하여 더 큰 모델을 더 빠르게 학습시킬 수 있습니다.
1. The Challenge / Context
최근 거대 언어 모델 (LLM)의 발전으로 DeepSpeed ZeRO와 같은 분산 학습 기술의 중요성이 더욱 커지고 있습니다. ZeRO-3는 모델, 옵티마이저 상태, 기울기 등을 여러 GPU에 분산하여 단일 GPU 메모리 제한을 극복하도록 설계되었지만, 복잡성 증가로 인해 메모리 부족 (OOM) 에러가 자주 발생합니다. 이러한 에러는 모델 개발 주기를 늦추고, 디버깅에 상당한 시간을 소모하게 만듭니다. 특히, ZeRO-3 구성이 잘못되었거나 메모리 누수가 있는 경우, 디버깅은 더욱 어려워집니다. GPU 메모리 사용량에 대한 심층적인 이해와 효과적인 최적화 전략이 필수적입니다.
2. Deep Dive: DeepSpeed ZeRO-3 메모리 분산 전략
DeepSpeed ZeRO (Zero Redundancy Optimizer)는 데이터 병렬 처리의 일종으로, 모델의 메모리 발자국을 줄여 더 큰 모델을 학습할 수 있도록 돕습니다. ZeRO-3는 다음과 같은 방식으로 메모리를 분산합니다.
- Sharded Model Parameters (모델 파라미터 분할): 모델의 파라미터를 모든 GPU에 복제하는 대신, GPU마다 파라미터의 일부를 저장합니다.
- Sharded Optimizer States (옵티마이저 상태 분할): Adam, SGD와 같은 옵티마이저 상태 (예: momentum, variance)도 GPU마다 분산 저장됩니다.
- Sharded Gradients (기울기 분할): 기울기 역시 모든 GPU에 복제되지 않고 분산 저장됩니다.
- Data Parallelism (데이터 병렬 처리): 각 GPU는 데이터의 서로 다른 배치(batch)를 처리합니다.
ZeRO-3는 모델 크기와 학습 효율성 간의 균형을 맞추기 위해 설계되었습니다. 핵심은 데이터, 모델 및 옵티마이저 상태를 GPU에 효율적으로 분산하여 각 GPU의 메모리 압박을 줄이는 것입니다. 하지만 올바른 설정과 디버깅 기술 없이는 OOM 에러를 피하기 어렵습니다. 특히 데이터 로딩, 활성화 함수 메모리 관리, 배치 크기 설정 등 다양한 요소가 메모리 사용량에 영향을 미칩니다.
3. Step-by-Step Guide / Implementation
이제 ZeRO-3 환경에서 GPU 메모리 에러를 디버깅하고 최적화하는 단계별 가이드를 살펴보겠습니다.
Step 1: 초기 설정 및 환경 준비
가장 먼저 DeepSpeed와 필요한 라이브러리가 올바르게 설치되었는지 확인합니다. deepspeed --version 명령어를 통해 설치된 DeepSpeed 버전을 확인하고 최신 버전으로 업데이트하는 것이 좋습니다.
pip install deepspeed
pip install torch # PyTorch는 필수
pip install psutil # 메모리 프로파일링에 유용
pip install pandas # 데이터 분석에 유용 (선택 사항)
Step 2: DeepSpeed 설정 파일 (JSON) 분석
DeepSpeed 설정 파일 (`ds_config.json`)은 학습에 사용되는 다양한 옵션을 정의합니다. 중요한 메모리 관련 설정을 확인해야 합니다.
{
"train_batch_size": 32,
"train_micro_batch_size_per_gpu": 4,
"gradient_accumulation_steps": 8,
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"reduce_bucket_size": 5e8,
"stage3_prefetch_bucket_size": 5e8,
"stage3_param_persistence_threshold": 1e4,
"sub_group_size": 1e9,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 32,
"hysteresis": 2,
"min_loss_scale": 1
},
"gradient_clipping": 1.0
}
- `train_batch_size`: 전체 배치 크기.
- `train_micro_batch_size_per_gpu`: GPU당 배치 크기. 이 값을 줄이면 GPU 메모리 사용량을 줄일 수 있습니다.
- `gradient_accumulation_steps`: 기울기 누적 단계 수. 이 값을 늘리면 가상 배치 크기가 증가하지만 메모리 사용량도 증가합니다. `train_batch_size`는 `train_micro_batch_size_per_gpu * gradient_accumulation_steps * world_size`와 같아야 합니다 (여기서 world_size는 총 GPU 수).
- `zero_optimization.stage`: ZeRO 단계 (1, 2, 또는 3). 3단계가 가장 메모리 효율적입니다.
- `offload_optimizer` & `offload_param`: 옵티마이저 상태 및 모델 파라미터를 CPU로 오프로드할지 여부를 결정합니다. True로 설정하면 GPU 메모리를 절약할 수 있지만 학습 속도가 느려질 수 있습니다.
- `fp16.enabled`: FP16 (반정밀도 부동 소수점) 학습 활성화 여부. True로 설정하면 메모리 사용량을 절약할 수 있습니다.
- `reduce_bucket_size`, `stage3_prefetch_bucket_size`: 이 파라미터들은 통신 버킷 크기를 제어합니다. 통신 오버헤드를 줄이기 위해 조정할 수 있습니다.
- `gradient_clipping`: 기울기 클리핑 임계값. 기울기가 너무 커지는 것을 방지하여 학습 안정성을 높입니다.
Step 3: 메모리 프로파일링
GPU 메모리 사용량을 자세히 분석하는 것이 중요합니다. PyTorch의 `torch.cuda.memory_summary()` 또는 DeepSpeed의 자체 메모리 통계를 활용할 수 있습니다.
import torch
import deepspeed
def print_gpu_memory(rank=0):
if rank == 0:
print(torch.cuda.memory_summary())
# DeepSpeed 엔진 초기화 후 (예시)
model_engine, optimizer, _, _ = deepspeed.initialize(
config_params=ds_config,
model=model,
optimizer=optimizer
)
# 학습 루프 내에서 메모리 사용량 출력
for step, batch in enumerate(train_dataloader):
# ... 학습 코드 ...
if step % 10 == 0: # 10 스텝마다 메모리 사용량 출력
print_gpu_memory(model_engine.local_rank)
`torch.cuda.memory_summary()`는 캐시된 메모리, 할당된 메모리, 사용 가능한 메모리 등 자세한 정보를 제공합니다. DeepSpeed 엔진의 `local_rank`를 사용하여 각 GPU의 메모리 사용량을 출력할 수 있습니다. TensorBoard를 사용하여 학습 중 메모리 사용량을 시각적으로 추적할 수도 있습니다 (DeepSpeed integrates with TensorBoard).
Step 4: 메모리 누수 감지
메모리 누수는 학습이 진행됨에 따라 GPU 메모리 사용량이 지속적으로 증가하는 현상을 말합니다. PyTorch의 `torch.autograd.set_detect_anomaly(True)`를 사용하여 자동 미분 과정에서 발생하는 문제를 식별할 수 있습니다.
import torch
torch.autograd.set_detect_anomaly(True)
# 학습 루프
for step, batch in enumerate(train_dataloader):
try:
# ... 학습 코드 ...
except Exception as e:
print(f"오류 발생: {e}")
break # 오류 발생 시 학습 중단
오류가 발생하면 역전파 과정에서 문제가 발생한 연산을 알려줍니다. 또한, 주기적으로 `torch.cuda.empty_cache()`를 호출하여 사용하지 않는 메모리를 해제할 수 있습니다. 하지만 과도한 캐시 비우기는 성능 저하를 야기할 수 있으므로 신중하게 사용해야 합니다.
Step 5: 모델 구조 및 배치 크기 조정
모델 구조 자체가 메모리 사용량에 큰 영향을 미칩니다. 레이어 수를 줄이거나, hidden size를 줄이거나, attention 메커니즘을 효율적인 버전으로 변경하는 것을 고려할 수 있습니다. 또한, `train_micro_batch_size_per_gpu`를 줄이는 것이 가장 간단한 방법 중 하나입니다. 메모리 제한 내에서 가능한 한 큰 배치 크기를 사용하는 것이 좋습니다. Gradient Accumulation Steps을 활용하면 작은 GPU 메모리를 가진 환경에서도 큰 배치 크기로 학습하는 효과를 얻을 수 있습니다.
Step 6: 오프로딩 (Offloading) 전략 활용
GPU 메모리가 부족한 경우, 옵티마이저 상태 및 모델 파라미터를 CPU 또는 NVMe SSD로 오프로드하는 것을 고려할 수 있습니다. DeepSpeed는 `offload_optimizer` 및 `offload_param` 옵션을 제공합니다. CPU 오프로드는 일반적으로 GPU 오프로드보다 느리지만, GPU 메모리가 매우 제한적인 경우 유용합니다. NVMe SSD 오프로드는 CPU 오프로드와 GPU 오프로드 사이의 절충안이 될 수 있습니다. DeepSpeed는 NVMe 오프로드를 위한 최적화된 구현을 제공합니다.
{
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "nvme",
"nvme_path": "/path/to/nvme",
"pin_memory": true
},
"offload_param": {
"device": "nvme",
"nvme_path": "/path/to/nvme",
"pin_memory": true
}
}
}
4. Real-world Use Case / Example
최근 175B 파라미터 규모의 LLM을 학습하면서 ZeRO-3 OOM 에러를 겪었습니다. 초기 설정에서 `train_micro_batch_size_per_gpu`를 너무 높게 설정했고, `gradient_accumulation_steps`는 충분히 활용하지 않았습니다. 또한, 불필요한 텐서가 메모리에 남아있는 메모리 누수 현상도 발견했습니다. 위에 설명된 단계를 따라, 먼저 `train_micro_batch_size_per_gpu`를 줄이고 `gradient_accumulation_steps`를 늘렸습니다. 그리고 메모리 프로파일링을 통해 메모리 누수 부분을 찾아 수정했습니다. 마지막으로, 옵티마이저 상태를 NVMe SSD로 오프로드하여 GPU 메모리 압박을 더욱 줄였습니다. 결과적으로, OOM 에러 없이 안정적으로 학습을 진행할 수 있었고, GPU 활용률도 향상되었습니다. 이전에는 몇 시간마다 발생하던 OOM 에러 때문에 학습이 중단되었지만, 최적화 후에는 일주일 이상 문제 없이 학습을 진행할 수 있었습니다.
5. Pros & Cons / Critical Analysis
- Pros:
- 더 큰 모델 학습 가능: ZeRO-3는 단일 GPU 메모리 제한을 극복하여 훨씬 큰 모델을 학습할 수 있게 해줍니다.
- 향상된 GPU 활용률: 적절한 최적화를 통해 GPU 활용률을 높여 학습 속도를 향상시킬 수 있습니다.
- 유연한 구성: 다양한 설정 옵션을 통해 특정 하드웨어 및 모델에 맞게 최적화할 수 있습니다.
- Cons:
- 복잡한 설정: ZeRO-3는 설정이 복잡하고, 최적의 성능을 위해서는 많은 실험이 필요합니다.
- 디버깅 어려움: 메모리 에러 발생 시, 디버깅이 어려울 수 있습니다.
- 통신 오버헤드: 파라미터 분산으로 인해 통신 오버헤드가 발생할 수 있으며, 이는 학습 속도에 영향을 미칠 수 있습니다.
6. FAQ
- Q: ZeRO-1, ZeRO-2, ZeRO-3의 차이점은 무엇인가요?
A: ZeRO-1은 옵티마이저 상태만 분산하고, ZeRO-2는 옵티마이저 상태와 기울기를 분산하며, ZeRO-3는 모델 파라미터, 옵티마이저 상태, 기울기를 모두 분산합니다. ZeRO-3가 가장 메모리 효율적이지만, 통신 오버헤드가 가장 높습니다. - Q: `gradient_accumulation_steps`를 늘리면 어떤 효과가 있나요?
A: `gradient_accumulation_steps`를 늘리면 가상 배치 크기가 증가하는 효과가 있습니다. GPU 메모리가 제한적인 경우, 작은 배치 크기로 학습하면서도 큰 배치 크기의 효과를 얻을 수 있습니다. 하지만 학습 스텝당 forward/backward 계산 횟수가 증가하므로 학습 시간이 늘어날 수 있습니다. - Q: FP16 학습이 항상 좋은가요?
A: FP16 학습은 메모리 사용량을 줄이고 학습 속도를 향상시킬 수 있지만, 수치 안정성 문제가 발생할 수 있습니다. Loss scaling을 적절히 조정하여 이러한 문제를 완화할 수 있습니다. - Q: `torch.cuda.empty_cache()`를 언제 사용해야 하나요?
A: `torch.cuda.empty_cache()`는 사용하지 않는 캐시된 메모리를 해제하는 데 유용합니다. 메모리 누수가 의심될 때 또는 OOM 에러가 발생할 가능성이 높을 때 사용할 수 있습니다. 하지만 과도한 캐시 비우기는 성능 저하를 야기할 수 있으므로 신중하게 사용해야 합니다. - Q: DeepSpeed config 파일을 어떻게 테스트하나요?
A: DeepSpeed에는 config 파일의 유효성을 검사하는 기능이 있습니다. `deepspeed --check_config ds_config.json` 명령어를 사용하여 config 파일에 오류가 있는지 확인할 수 있습니다.
7. Conclusion
DeepSpeed ZeRO-3는 거대한 모델 학습을 위한 강력한 도구이지만, 효과적으로 사용하려면 깊이 있는 이해와 디버깅 기술이 필요합니다. 이 글에서 제시된 단계별 가이드와 최적화 전략을 통해 ZeRO-3 환경에서 GPU 메모리 에러를 극복하고, 더 큰 모델을 더 빠르게 학습시킬 수 있기를 바랍니다. 지금 바로 코드를 적용하고, 자신의 모델에 맞는 최적의 설정을 찾아보세요! DeepSpeed 공식 문서도 잊지 말고 참고하세요. DeepSpeed Configuration Documentation
