'Cuda out of Memory' Error 해결하기

[에러] 'Cuda out of Memory' Error 해결하기

 

 

안녕하세요 머킹입니다.

오늘은 모델 학습 중에 메모리를 관리하는 방법에 대해서 알아보겠습니다.

저는 종종 Cuda out of Memory를 마주하게 되는데 여러분들은 어떠신가요?

---

GPU 메모리 관리의 중요성과 최적화 기법

GPU 메모리 관리에 대한 이해와 적절한 최적화는 대형 모델을 학습하는 데 있어서 매우 중요합니다.

특히, PyTorch를 사용한 딥러닝 작업에서는 GPU 메모리 부족으로 인한 오류(CUDA out of memory)가

빈번하게 발생할 수 있습니다.

 

 

1. PYTORCH_CUDA_ALLOC_CONF 환경 변수 설정

PyTorch는 GPU 메모리를 관리하기 위해 PYTORCH_CUDA_ALLOC_CONF라는 환경 변수를 제공합니다. 이 변수는 메모리 할당 방식에 영향을 주어, 메모리 파편화를 줄이고 메모리 효율성을 높이는 데 도움을 줍니다. 특히, 대형 모델을 학습할 때 메모리 파편화로 인해 자주 발생하는 메모리 부족 문제를 해결할 수 있습니다.

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

 

2. Gradient Accumulation Steps

gradient_accumulation_steps는 작은 배치 크기를 사용할 때 학습을 안정적으로 진행하기 위해 사용되는 방법입니다. 배치 크기를 줄이면서도 모델의 성능을 유지하기 위해, 그래디언트를 여러 번 누적한 후 한 번에 업데이트하는 방식입니다. 이를 통해 메모리 사용량을 줄이면서도 효과적인 큰 배치 크기에서 학습하는 것과 동일한 효과를 얻을 수 있습니다.

training_args = TrainingArguments(
    gradient_accumulation_steps=16,  # 그래디언트 누적 단계를 증가시켜 메모리 사용을 최적화
)

 

 

3. Layer Freezing (레이어 동결)

레이어 동결은 학습이 필요하지 않은 모델의 일부 레이어를 동결하여, 이들 레이어가 학습 중에 메모리를 사용하지 않도록 하는 방법입니다. 예를 들어, 모델의 초기 레이어는 이미 충분히 학습된 상태일 수 있으며, 이런 경우 해당 레이어를 동결하여 학습할 파라미터의 수를 줄이면 메모리 사용을 크게 줄일 수 있습니다.

 

for name, param in model.named_parameters():
    if "layer" in name and int(name.split(".")[2]) < 6:  # 예: 처음 6개의 레이어를 동결
        param.requires_grad = False

 

---

추가적으로 사용할 수 있는 메모리 관리 방법

코드에서 사용된 방법 외에도 GPU 메모리 부족 문제를 해결할 수 있는 다양한 최적화 기법이 있습니다. 아래는 추가적으로 고려할 수 있는 방법들입니다:

1. Batch Size 줄이기

배치 크기를 줄이면 한 번에 처리하는 데이터의 양이 감소하여 메모리 사용량이 줄어듭니다. 작은 배치 크기와 gradient_accumulation을 함께 사용하면, 성능을 유지하면서도 메모리 사용을 최적화할 수 있습니다.

2. Mixed Precision Training (FP16/FP32 혼합 정밀도 학습)

FP16 정밀도를 사용하면 메모리 사용량을 줄일 수 있지만, 특정 하드웨어와의 호환성 문제로 인해 오류가 발생할 수 있습니다. FP16 대신 bf16 정밀도를 사용하는 것도 고려해볼 수 있습니다.

3. Model Pruning (모델 가지치기)

모델 가지치기는 학습 중에 중요하지 않은 파라미터를 제거하여 모델을 경량화하는 방법입니다. 이렇게 하면 학습에 필요한 메모리와 계산 비용을 줄일 수 있습니다.

4. Distillation (지식 증류)

지식 증류는 큰 모델에서 작은 모델로 지식을 전이하여 작은 모델이 더 적은 메모리로도 높은 성능을 발휘할 수 있게 하는 방법입니다. 이 방법은 메모리 효율성을 높이면서도 성능을 유지하는 데 효과적입니다.

5. 더 작은 모델 사용

학습하려는 작업의 복잡도에 따라 더 작은 모델을 사용하는 것도 좋은 방법입니다. 모델이 작을수록 필요한 메모리가 줄어들어, 메모리 부족 문제를 해결할 수 있습니다.

 

---

1. LoRA (Low-Rank Adaptation)

LoRA는 학습 가능한 파라미터 수를 줄이기 위해 모델의 일부 파라미터 행렬을 저랭크(low-rank) 근사로 분해하는 방법입니다. 이를 통해 대규모 모델의 일부만을 효율적으로 학습할 수 있으며, 메모리 사용량을 줄이면서도 성능 저하를 최소화할 수 있습니다.

LoRA의 핵심 아이디어는 원래의 파라미터 행렬을 학습하는 대신, 저랭크 행렬 두 개를 학습함으로써 전체 파라미터의 수를 줄이고, 메모리 사용량과 계산 비용을 절감하는 것입니다. 이 방법은 특히 대형 언어 모델에서 효과적입니다.

---

2. Adapter Layers

Adapter Layers는 기존의 사전 학습된 모델 위에 소규모의 학습 가능한 모듈(어댑터 레이어)을 추가하는 방법입니다. 원래 모델의 파라미터를 동결한 상태에서, 어댑터 레이어만을 학습하게 됩니다. 이 방법은 새로운 작업에 모델을 적응시키는 데 효과적이며, 메모리 사용량을 크게 줄일 수 있습니다.

어댑터 레이어의 크기는 매우 작기 때문에, 원래 모델의 학습 가능한 파라미터 수에 비해 메모리 사용량이 크게 줄어듭니다. 또한, 다양한 작업에 대해 원래 모델을 재사용할 수 있는 장점이 있습니다.

---

3. Quantization (양자화)

Quantization은 모델의 파라미터를 더 낮은 비트 정밀도로 변환하여 메모리 사용량을 줄이는 방법입니다. 예를 들어, 32비트 부동 소수점 숫자 대신 8비트 정밀도를 사용하면, 모델의 메모리 사용량을 크게 줄일 수 있습니다.

양자화는 모델의 크기를 줄이면서도 정확도를 유지할 수 있도록 설계되었습니다. 이 방법은 특히 하드웨어 가속기(GPU, TPU 등)에서의 효율성을 극대화하는 데 유리합니다.

---

4. Distillation (지식 증류)

Distillation은 큰 모델(교사 모델)의 지식을 더 작은 모델(학생 모델)에 전이하는 방법입니다. 큰 모델에서 학습한 정보를 작은 모델로 압축함으로써, 메모리와 계산 비용을 줄이면서도 높은 성능을 유지할 수 있습니다.

지식 증류는 특히 경량 모델을 학습할 때 유용하며, 큰 모델과 거의 비슷한 성능을 유지하면서도 메모리 사용량을 대폭 줄일 수 있는 방법입니다.

---

5. Pruning (모델 가지치기)

Pruning은 중요하지 않거나 불필요한 파라미터를 제거하여 모델을 경량화하는 방법입니다. 파라미터의 일부를 가지치기함으로써, 모델의 크기와 복잡도를 줄일 수 있습니다.

모델 가지치기는 학습 후 또는 학습 중에 수행될 수 있으며, 모델의 효율성을 높이는 데 매우 효과적입니다. 다만, 너무 많은 파라미터를 제거하면 성능 저하가 발생할 수 있으므로, 가지치기 전략을 신중히 선택해야 합니다.

---

6. Elastic Weight Consolidation (EWC)

**Elastic Weight Consolidation (EWC)**는 모델이 새로운 작업을 학습할 때, 이전 작업에서 중요했던 파라미터를 보호하여 성능 저하를 방지하는 방법입니다. 이를 통해 모델이 여러 작업에 걸쳐 학습할 때 메모리 사용량을 줄이면서도 성능을 유지할 수 있습니다.

EWC는 모델이 새로운 작업을 학습할 때, 이전 작업에서 중요한 파라미터의 가중치를 보존하는 방식으로 작동합니다. 이 방법은 특히 연속 학습 시에 유용합니다.

---

7. Sparse Tuning

Sparse Tuning은 모델의 파라미터 중 일부만을 선택적으로 학습하는 방법입니다. 이는 모델의 전체 파라미터를 모두 학습하는 대신, 중요한 파라미터만을 학습함으로써 메모리와 계산 자원을 절약할 수 있습니다.

이 방법은 특히 대형 모델에서 효과적이며, 메모리 사용량을 줄이면서도 높은 성능을 유지할 수 있습니다.