18B 파라미터 GPT 모델을 Single GPU로 학습하기 (Colossal-AI)

https://github.com/hpcaitech/ColossalAI 

GitHub - hpcaitech/ColossalAI: Colossal-AI: A Unified Deep Learning System for Big Model Era

Colossa-AI의 특징

• vanilla Pytorch 와 MS DeepSpeed와 같은 메인스트림 분산 솔루션보다 더 좋은 성능을 낸다.

Colossa-AI의 사용예

• RTX 2060 (6GB) : 1.5B 모델 학습 가능
• RTX 3090 (24GB) : 18B 모델 학습 가능

Collossa-AI vs DeepSpeed

• GPT-2 모델 Scaling 케이스
  • 동일한 GPU 환경에서 DeepSpeed에 비해 2~5배 더 큰 모델을 학습 가능함

Key Technology: Enhanced Heterogeneous Training

• Consumer-grade GPU를 이용하여 Large AI 모델을 학습할 때 가장 큰 문제는 GPU 메모리 용량 제한이다.
• ZeRO-Offload의 문제점
  • 고정된 메모리 레이아웃을 사용하여 CPU와 GPU 메모리에 모델 파라미터들을 분할하여 저장함 
• Gemini 메커니즘 (designed by Colossal-AI)
  • GPU와 CPU의 이기종 메모리를 효율적으로 관리 및 사용
  • 학습 프로세스 동안 CPU-GPU 메모리에 동적으로 텐서들을 분산시킴
  • Iteration 수에 따라 학습 단계를 2 단계로 나눔
    • Warmup Stage
      • 초기 Warmup Phase에서 메모리 정보를 모니터링함
    • Non-Warmup Stage
      • Warmup Phase에서 수집된 정보는 CPU-GPU 데이터 이동을 최소화하기 위해 사용됨
• Gemini 구현시 어려움
  • Non-Model 데아터의 메모리 사용량을 확보하기 어려움
    • Non-Model 데이터의 수명은 사용자에 의해 관리되지 않음
    • 현 DL 프레임워크는 사용자에게 Non-Model 데이터를 트래킹하기 위한 인터페이스를 제공하지 않음
  • CUDA 컨텍스트와 같이 Non-framework overhead를 고려해야할 필요성
• Colossal-AI의 구현
  • Warmup 단계에서 샘플링을 통해 CPU와 GPU 메모리의 사용량을 확보함
  • Non-model 데이터의 사용량은 최대 시스템 메모리 사용량과 모델 메모리 사용량, 두 순간을 비교하여 확보함
  • Model의 메모리 사용량은 메모리 매니저에 질의하여 알아냄 (검은색 Solid 선)
  • 모델의 모든 텐서는 메모리 매니저에 의해 관리됨
    • 각 텐서는 HOLD, COMPUTE, FREE와 같은 상태 정보가 마킹됨
    • Colossal-AI는 메모리 사용량을 동적으로 확보하여 텐서의 상태를 변화시키고 연속적으로 텐서의 위치를 조절함
    • CPU-GPU 메모리를 효율적으로 이용하면, 극단적으로 제한된 HW에서도 모델 용량를 최대화하면서 학습 속도의 균형을 맞추는 것이 가능함

Colossal-AI의 효율성

• GPT-3 학습시, 학습 시작에 필요한 최소 리소스는 NVIDIA(Megatron-LM)에 비해 절반만 필요함
• 같은 컴퓨팅 리소스를 사용할  경우,  NVIDIA에 비해 11% 학습속도가 빠름

 

 

원문

https://medium.com/@hpcaitech/train-18-billion-parameter-gpt-models-with-a-single-gpu-on-your-personal-computer-8793d08332dc

Train 18-billion-parameter GPT models with a single GPU on your personal computer!