Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading

이 논문은 Transformer 모델 학습 시 GPU 메모리 부족 문제를 해결하기 위해 순전파, 역전파, 업데이트 단계의 메모리 사용량 변동을 활용하여 옵티마이저 상태를 CPU 와 GPU 간에 동적으로 분배하는 'Deep Optimizer States' 기법을 제안함으로써, 기존 최첨단 방법 대비 2.5 배 빠른 학습 속도를 달성했다고 요약할 수 있습니다.

핵심

1. 문제 상황: 도서관이 너무 작아졌어요! (메모리 벽)

상상해 보세요. 여러분은 **수십억 권의 책 (AI 모델의 지식)**을 한 번에 읽고 이해해야 하는 **초고속 도서관 사서 (GPU)**가 되었습니다.

• 현재의 문제: 책이 너무 많아서 책장 (GPU 메모리) 에 다 담을 수가 없습니다.
• 기존 해결책 (DeepSpeed 등): 책장에는 중요한 책만 두고, 나머지 책들은 **창고 (컴퓨터의 일반 메모리/CPU)**에 보관합니다.
  • 사서가 책을 읽을 때, 창고에서 책을 가져와야 하므로 시간이 많이 걸립니다.
  • 특히 사서가 "이 책을 어떻게 이해했지? 다음엔 어떻게 수정할까?"라고 **학습 노트 (옵티마이저 상태)**를 정리할 때, 창고까지 다녀와야 해서 속도가 매우 느려집니다.
  • 마치 사서가 책장 앞에 서서 창고로 뛰어갔다 오기를 반복하는 것과 같아, 책장 앞에서는 아무것도 못 하고 기다리는 시간이 길어집니다.

2. 새로운 아이디어: "교차 오프로딩" (Deep Optimizer States)

이 논문은 **"창고와 책장 사이를 오가는 시간을 줄이고, 사서와 창고 관리자가 동시에 일하게 하자"**는 아이디어를 제시합니다.

🌟 핵심 비유: "동시 작업과 교차 이동"

기존 방식은 순서대로 일했습니다.

1. 창고에서 책 가져오기 (기다림)
2. 책 읽기 (작업)
3. 창고로 노트 가져가기 (기다림)
4. 노트 정리하기 (작업)
5. 다시 책장으로 가져오기 (기다림)
   → 사서가 일할 때 창고 관리자는 쉬고, 창고 관리자가 일할 때 사서는 쉬는 '비효율'이 발생했습니다.

Deep Optimizer States는 다음과 같이 바꿉니다:

1. 책장 (GPU) 에는 '학습 노트'의 일부만 두고, 나머지는 창고 (CPU) 에 둡니다.
2. 사서 (GPU) 가 노트를 정리하는 동안, 창고 관리자 (CPU) 는 다음에 필요한 노트를 미리 책장으로 가져옵니다.
3. 동시에, 창고 관리자가 노트를 정리하는 동안, 사서는 이미 정리된 노트를 받아서 다음 책을 읽을 준비를 합니다.

이처럼 GPU 와 CPU 가 서로의 일을 방해하지 않고, 동시에 움직이면서 데이터를 주고받는 방식을 **'교차 오프로딩 (Interleaved Offloading)'**이라고 합니다.

3. 이 기술의 3 가지 비밀 무기

이 시스템이 어떻게 그렇게 빨라질 수 있었는지 3 가지 비법을 소개합니다.

① "빈 공간을 활용한 스마트한 배치"

• 상황: 사서가 책을 읽는 동안은 책장 (GPU 메모리) 이 꽉 차지만, 책을 다 읽고 노트를 정리할 때는 책장이 비어있는 공간이 생깁니다.
• 해결: 이 빈 공간을 이용해 노트 (학습 데이터) 를 잠시 책장에 올려둡니다. 창고에서 가져오는 시간을 아껴서, 사서가 바로 작업을 시작할 수 있게 합니다.

② "데이터를 미리 가져오는 (Prefetching) 기술"

• 상황: 창고에서 책을 가져오는 데 시간이 걸립니다.
• 해결: 사서가 지금 책을 읽고 있을 때, 다음에 읽을 책은 이미 창고에서 책장 앞으로 이동 중이게 합니다. 사서가 책을 다 읽는 순간, 다음 책은 이미 준비되어 있어 기다릴 필요가 없습니다.

③ "데이터 변환의 효율화"

• 상황: 창고 (CPU) 는 큰 숫자 (FP32) 로 계산하고, 책장 (GPU) 은 작은 숫자 (FP16) 로 계산합니다. 둘 사이를 오갈 때 숫자를 변환하는 데 시간이 걸립니다.
• 해결: 숫자 변환 작업을 책장 (GPU) 에서 바로 처리하도록 합니다. 창고로 가서 변환했다가 다시 오는 수고를 덜어주어, 데이터 이동 속도를 획기적으로 높였습니다.

4. 결과: 얼마나 빨라졌나요?

이 기술을 적용한 결과, 기존 방식보다 최대 2.5 배 더 빠른 훈련 속도를 달성했습니다.

• 비유: 예전에는 사서가 창고로 왕복하는 데 1 시간 걸려서 하루에 10 권의 책만 정리했다면, 이제는 동시에 일을 처리해서 하루에 25 권의 책을 정리할 수 있게 된 것입니다.
• 의미: 이 기술 덕분에, 비싼 고성능 컴퓨터 (HPC) 가 없어도 일반 기업이나 연구실에서도 거대한 AI 모델을 빠르게 훈련시킬 수 있게 되었습니다.

5. 결론

이 논문은 **"메모리가 부족해서 창고에 데이터를 두는 것은 어쩔 수 없지만, 창고와 책장 사이를 오가는 시간을 줄이고, 두 곳에서 동시에 일하게 하면 속도를 2 배 이상 높일 수 있다"**는 것을 증명했습니다.

앞으로 더 큰 AI 모델이 등장하더라도, 이 스마트한 데이터 관리 기술을 통해 우리가 더 빠르고 저렴하게 AI 를 개발할 수 있는 길이 열렸습니다.

기술 요약

1. 문제 정의 (Problem Statement)

대규모 언어 모델 (LLM) 과 트랜스포머 모델의 파라미터 수가 수천억 개에서 조 (Trillion) 단위까지 급증함에 따라, GPU 메모리만으로는 모델 파라미터, 옵티마이저 상태 (Optimizer State), 그래디언트, 활성화 값 등을 모두 저장하기 어려워지는 '메모리 벽 (Memory Wall)' 문제가 발생하고 있습니다.

• 기존 접근법의 한계: 메모리 부족을 해결하기 위해 DeepSpeed Offload, TwinFlow, Zero-Infinity 등의 기술이 옵티마이저 상태 (FP32 파라미터, 모멘텀, 분산) 를 호스트 (CPU) 메모리로 오프로딩 (Offloading) 하는 방식을 사용합니다.
• 주요 병목 현상:
  1. PCIe 대역폭 부족: GPU 와 CPU 간의 데이터 이동 (H2D, D2H) 이 PCIe 링크를 통해 이루어지는데, 이 대역폭 (약 25~50 GB/s) 이 계산 속도에 비해 매우 느립니다.
  2. CPU 계산 속도 저하: GPU 에 비해 CPU 의 계산 성능이 현저히 낮습니다 (예: H100 GPU 는 초당 1000 억 파라미터 업데이트, CPU 는 80 억).
  3. 비효율적인 리소스 활용: 기존 하이브리드 방식 (TwinFlow 등) 은 옵티마이저의 일부를 GPU 에 고정 (Static) 하고 나머지를 CPU 에서 처리하지만, Forward/Backward Pass 동안 GPU 메모리가 비어있거나 PCIe 링크가 유휴 상태일 때에도 이를 활용하지 못해 전체적인 학습 속도가 느려집니다.

2. 방법론 (Methodology)

저자들은 옵티마이저 업데이트 단계에서 발생하는 GPU 메모리 사용량의 변동과 PCIe 링크의 비활용 시간을 역이용하여, CPU 와 GPU 간의 작업을 교차 (Interleaved) 하여 동적으로 스케줄링하는 새로운 기법인 Deep Optimizer States를 제안합니다.

핵심 설계 원칙

1. 교차 오프로딩 (Interleaved Offloading):
   • 옵티마이저 상태를 작은 '서브그룹 (Subgroup)' 단위로 쪼개고, 이 중 일부를 GPU 에서, 나머지를 CPU 에서 업데이트합니다.
   • 고정된 비율이 아닌, 성능 모델을 기반으로 동적 비율로 업데이트를 분배합니다.
2. 중첩 실행 (Overlapping Execution):
   • CPU 가 서브그룹 A 의 업데이트를 계산하는 동안, GPU 는 서브그룹 B 의 업데이트를 수행하고, PCIe 를 통해 서브그룹 C 의 데이터를 비동기적으로 전송 (Prefetch/Flush) 합니다.
   • 이를 통해 CPU 계산, GPU 계산, PCIe 데이터 전송을 동시에 수행하여 대기 시간을 최소화합니다.
3. 효율적인 그라디언트 관리:
   • 활성화 값 (Activations) 이 해제된 GPU 메모리를 활용하여, GPU 에서 업데이트할 서브그룹의 그라디언트를 저장합니다.
4. 정밀도 변환 최적화:
   • 기존 방식은 FP16 그라디언트를 CPU 로 보낸 후 FP32 로 변환하는 과정에서 병목이 발생했습니다. Deep Optimizer States 는 GPU 에서 FP16 → FP32 변환을 수행한 후 FP32 상태로 전송하여, 메모리 할당 및 변환 오버헤드를 줄이고 PCIe 대역폭 효율을 극대화합니다.

성능 모델 (Performance Model)

• CPU 업데이트 속도 ($U_c$), GPU 업데이트 속도 ($U_g$), PCIe 대역폭 ($B$), 정밀도 변환 속도 ($D_c$) 등을 고려하여 **최적의 "업데이트 스트라이드 (Update Stride)"**를 계산합니다.
• 이 모델을 통해 "몇 개의 CPU 업데이트 후 GPU 업데이트를 수행해야 PCIe 와 계산 자원을 최대화할 수 있는지"를 결정합니다.

3. 주요 기여 (Key Contributions)

1. 심층 분석: 옵티마이저 오프로딩 시 학습 반복 (Iteration) 동안의 GPU 메모리 사용량 변동과 PCIe 링크의 비활용 구간을 정량적으로 분석했습니다.
2. 새로운 아키텍처 설계:
   • CPU 와 GPU 간의 옵티마이저 업데이트를 교차 실행하는 설계 원칙 제시.
   • 그라디언트 이동 및 정밀도 변환을 최적화하는 메커니즘 도입.
3. 성능 모델 및 알고리즘: 동적 오프로딩 비율을 결정하는 성능 모델과 이를 구현하는 스케줄링 알고리즘을 제안했습니다.
4. 실제 구현 및 통합: DeepSpeed 및 Megatron-LM 런타임에 통합된 오픈소스 미들웨어 (Deep Optimizer States) 를 구현했습니다.
5. 광범위한 평가: 다양한 모델 크기 (7B~20B) 와 하드웨어 구성에서 기존 최첨단 기법 대비 성능을 검증했습니다.

4. 실험 결과 (Results)

ALCF 의 JLSE 테스트베드 (4×H100 GPU, 192 CPU 코어) 를 사용하여 실험한 결과는 다음과 같습니다.

• 학습 속도 향상:
  • 기존 DeepSpeed ZeRO-3 (전체 오프로딩) 대비 최대 2.5 배 빠른 반복 (Iteration) 시간을 달성했습니다.
  • 옵티마이저 업데이트 단계만 보면 최대 3 배의 속도 향상을 보였습니다.
• 리소스 활용도:
  • GPU 사용률: 100% 에 근접 (기존 방식은 오프로딩 중 대기 상태로 낮음).
  • PCIe 대역폭: 최대 40% 수준 활용 (기존 방식은 10% 미만).
  • CPU 사용률: 약 60~70% 수준으로 유지되면서도 전체 TFLOPS 는 2.5 배 증가 (약 75 TFLOPS 달성).
• 모델 크기 및 설정별 성능:
  • 20B 파라미터 모델에서 7B 모델의 학습 시간과 유사한 수준으로 단축되었습니다.
  • GPU 메모리에 일부 옵티마이저를 고정 (Static) 하는 TwinFlow 방식보다도 1.7~2.3 배 빠른 성능을 보였습니다.
  • 마이크로배치 크기 증가 및 데이터 병렬화 (Data Parallelism) 확대 시에도 효율적인 확장성을 입증했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 메모리 제약 해결: 제한된 GPU 메모리 환경에서도 대규모 LLM 을 효율적으로 학습할 수 있는 새로운 패러다임을 제시합니다.
• 하드웨어 효율성 극대화: CPU 와 GPU 의 계산 능력, 그리고 PCIe 대역폭을 동시에 활용하여 '기다림'을 '계산/전송'으로 변환함으로써 시스템 전체의 처리량을 극대화합니다.
• 미래 지향성: 차세대 시스템 (예: Grace Hopper, C2C 인터커넥트 등) 에서 CPU-GPU 간 대역폭이 더욱 증가할 경우, 이러한 동적 교차 오프로딩 기법의 중요성과 효과는 더욱 커질 것으로 예상됩니다.

이 연구는 단순히 메모리 용량을 늘리는 것을 넘어, 기존 하드웨어 자원을 어떻게 지능적으로 조율하여 학습 속도를 높일 수 있는지에 대한 중요한 통찰을 제공합니다.