GPU 메모리 부족 문제, LoRA, 양자화, FlashAttention으로 해결!
CUDA out of memory 오류는 AI 개발의 흔한 문제로, GPU의 HBM(High Bandwidth Memory) 부족이 원인임
모델의 가중치(Model Weights), 옵티마이저 상태(Optimizer States), 활성화(Activations)가 HBM을 소비하며, 배치 크기(Batch Size)에 따라 사용량이 동적으로 변동
LoRA(Low-Rank Adaptation), 모델 양자화(Model Quantization), FlashAttention 등의 기술을 활용하여 HBM 사용량 감소 가능
단일 GPU의 한계를 넘어 데이터 병렬 처리(Data Parallelism), 모델 병렬 처리(Model Parallelism), FSDP(Fully Sharded Data Parallelism) 등 멀티 GPU 전략 제시
16GB, 24GB, 40GB+ HBM 사용 사례를 제시하며, GPU 메모리(GPU Memory) 사용량에 따른 적절한 모델 및 배치 크기 선택의 중요성을 강조
HBM(High Bandwidth Memory) 사용량 분석
본문에서는 GPU의 HBM(High Bandwidth Memory) 사용량을 모델 가중치(Model Weights), 옵티마이저 상태(Optimizer States), 활성화(Activations) 세 가지 주요 요소로 분석한다.
모델 크기(Model Size): 모델 파라미터(Parameter) 수와 정밀도(Precision)에 비례하며, bfloat16(bfloat16) 사용 시 float32(float32) 대비 메모리 절약
옵티마이저 상태(Optimizer States) 및 기울기(Gradients): 모델 학습 과정에서 필요한 데이터로, AdamW(AdamW) 등 옵티마이저 종류에 따라 메모리 사용량 상이
활성화(Activations) 및 배치 데이터(Batch Data): 배치 크기(Batch Size)에 따라 동적으로 변동하며, 큰 배치 크기는 학습 속도 향상에 기여하나 메모리 사용량 증가
실제 프레임워크(Framework)에서는 임시 버퍼(Temporary Buffer) 및 커널 실행(Kernel Launch) 등으로 인해 최대 30%의 추가 오버헤드(Overhead) 발생에 유의해야 한다.
LoRA(Low-Rank Adaptation)를 활용한 메모리 절감
LoRA(Low-Rank Adaptation)는 모델의 모든 파라미터를 학습하는 대신, 작은 어댑터 레이어(Adapter Layer)만 학습하여 메모리 사용량을 줄이는 PEFT(Parameter-Efficient Fine-Tuning) 기법이다.
기본 모델(Base Model) 고정: 기존 모델의 가중치(Weight)는 고정하고, 새로운 파라미터(Parameter)를 추가하여 학습
메모리 절감 효과: 기울기(Gradients) 및 옵티마이저 상태(Optimizer States)는 작은 어댑터 레이어에 대해서만 계산되므로, 전체 모델 학습 대비 메모리 사용량 대폭 감소
예시: 40억 개 파라미터(Parameter) 모델에 LoRA 적용 시, 2000만 개 파라미터만 학습하여 HBM 사용량 8GB 수준으로 유지
LoRA는 모델 크기가 크고, HBM 용량이 제한적인 환경에서 효율적인 미세 조정(Fine-tuning)을 가능하게 한다.
모델 양자화(Model Quantization)를 통한 메모리 최적화
모델 양자화(Model Quantization)는 모델 가중치(Weight)의 정밀도를 낮춰 HBM 사용량을 줄이는 기법으로, int8(8-bit integer) 또는 int4(4-bit integer)와 같은 낮은 정밀도(Precision) 형식을 사용한다.
bfloat16(bfloat16) vs. int8(8-bit): bfloat16은 float32(float32)와 유사한 동적 범위를 가지면서 메모리 사용량을 절반으로 줄임
int4(4-bit)의 효과: int8(8-bit) 대비 메모리 사용량을 50% 추가 절감하여, 총 4배의 메모리 절약 효과를 얻을 수 있음
주의사항: 낮은 정밀도 사용 시, 수치적 불안정(Numerical Instability)으로 인해 NaN(Not a Number) 발생 가능성 존재
LoRA와 양자화를 결합한 QLoRA(QLoRA)는 4-bit 양자화된 모델과 bfloat16(bfloat16) LoRA 어댑터를 사용하여 최대 메모리 효율(Memory Efficiency)을 달성한다.
FlashAttention을 활용한 학습 속도 향상
FlashAttention은 트랜스포머(Transformer) 모델의 어텐션 메커니즘(Attention Mechanism)에서 발생하는 메모리 병목 현상을 해결하여 학습 속도를 향상시키는 알고리즘이다.
어텐션 행렬(Attention Matrix) 문제: 어텐션 메커니즘(Attention Mechanism)은 중간 결과를 저장하기 위해 많은 HBM을 사용
FlashAttention의 해결책: 어텐션 계산 순서를 재정렬하여 HBM 사용량 감소 및 계산 속도 향상을 동시에 달성
구현 용이성: 대부분의 경우, 코드 한 줄(One-line change)만 변경하여 FlashAttention을 활성화 가능
MedGemma 모델 예시: `attn_implementation="sdpa"` 설정을 통해 FlashAttention을 간편하게 적용
FlashAttention은 단일 GPU 및 멀티 GPU 환경 모두에서 학습 효율(Training Efficiency)을 극대화한다.
멀티 GPU 환경에서의 확장 전략
단일 GPU의 HBM 제약 극복을 위해 데이터 병렬 처리(Data Parallelism), 모델 병렬 처리(Model Parallelism), FSDP(Fully Sharded Data Parallelism) 등 다양한 멀티 GPU 전략을 활용한다.
데이터 병렬 처리(Data Parallelism): 각 GPU에 전체 모델 복제 후, 데이터를 분할하여 병렬 처리
모델 병렬 처리(Model Parallelism): 모델을 여러 GPU에 분산 배치하여, 모델 크기 제약 극복
FSDP(Fully Sharded Data Parallelism): 모델, 옵티마이저 상태, 기울기(Gradients)를 GPU 간에 분산하여, 각 GPU의 HBM 사용량 최소화
FSDP는 대규모 모델 학습 시 최대 메모리 효율(Memory Efficiency)을 제공하며, 여러 개의 소형 GPU를 활용하여 비용 효율적인 학습 환경(Cost-effective Training Environment)을 구축할 수 있다.
