Offload


Contents

  1. Offload란?
  2. Offload의 종류
  3. Offload 기술의 핵심 장점
  4. Code (feat. DeepSpeed)
  5. Summary


1. Offload란?

일부 계산을 GPU에서 CPU 또는 NVMe로 옮겨서 메모리를 절약하는 기술!

필요성? 일반적으로 GPU 메모리는 한정적!

\(\rightarrow\) 대형 모델을 학습할 때 메모리 부족이 문제가 될 수 있음!


딥러닝에서 큰 메모리를 차지하는 요소 3가지

  1. 모델 파라미터 (Weights)
  2. 그래디언트 (Gradients)
  3. 옵티마이저 상태 (Optimizer States, e.g., Adam의 1차, 2차 모멘트)


Offload: 이 중 일부를 CPU 또는 NVMe로 옮겨서 메모리를 아끼는 기술


2. Offload의 종류

어떤 데이터를 오프로드(이동)하느냐에 따라 방식이 다름!


(1) Optimizer Offload

  • “옵티마이저 상태”를 GPU → CPU로 이동
  • 메모리 절약 효과: 중간

  • 예시) Adam 옵티마이저를 사용시, 1차, 2차 모멘트 값(m, v)이 GPU에서 많은 메모리를 차지

    \(\rightarrow\) 이를 CPU로 옮기면 GPU 메모리 사용량이 줄어든다

    ( 단점: CPU에서 연산이 수행되므로 속도가 조금 느려질 수 있음 )


(2) Parameter + Optimizer Offload

  • “모델 파라미터 & 옵티마이저 상태”를 GPU → CPU로 이동
  • 메모리 절약 효과:

  • 예시) 모델의 가중치(Weights)도 CPU 메모리에 저장하고, 연산이 필요할 때 GPU로 불러와서 계산하는 방식.

    ( 단점: 계산할 때마다 CPU ↔ GPU 간 데이터 이동이 필요! 속도가 더 느려질 수 있음. )


(3) NVMe Offload (SSD Offload)

  • “옵티마이저 상태”를 GPU → NVMe (SSD)로 이동
  • 메모리 절약 효과: 최대

  • 예시) CPU 메모리도 부족할 경우, SSD 같은 저장장치에 데이터를 저장하고 필요할 때 불러오는 방식.

    ( GPU → CPU보다 더 느리지만, 아주아주 큰 모델을 학습할 때 유용함. )


3. Offload 기술의 핵심 장점

  1. GPU 메모리를 효과적으로 절약→ 더 큰 모델을 학습 가능
  2. **기존 DDP(Distributed Data Parallel) 대비 **메모리 사용량이 줄어든다
  3. Zero Redundancy Optimizer (ZeRO) 와 결합해서 더욱 강력한 메모리 최적화 가능


4. Code (feat. DeepSpeed)

ds_config = {
    "zero_optimization": {
        "stage": 2,  # ZeRO Stage 2 사용
        "offload_optimizer": {
            "device": "cpu"  # 옵티마이저 상태를 CPU로 이동
        }
    }
}
  • 설명: 옵티마이저 상태를 CPU로 오프로드해서 GPU 메모리를 절약할 수 있어.
  • device: "nvme" : SSD(NVMe)로 옵티마이저 상태를 이동


5. Summary

  • Offload는 GPU 메모리 부족 문제를 해결하기 위한 기술!
    • 옵티마이저 상태, 모델 파라미터 등을 CPU나 NVMe로 이동해서 GPU 부담을 줄임
  • 속도는 다소 느려질 수 있지만, 큰 모델 학습이 가능해짐
  • ZeRO와 함께 사용하면 더욱 강력한 메모리 절약 효과를 얻을 수 있음