Large Model Memory 효율적 사용하는 방법 (huggingface 기준)
- 최근 나오는 LLM 모델들의 경우 메모리 사용량이 많은데 메모리 사용을 효율적으로 사용 할 수 있는 방법들 기록
- 항상 성능과 효율의 적정선을 찾기 위해 노력해야함
- Model load 시에 8bit 로 로딩
- AutoModelForCausalLM.from_pretrained() 의 옵션으로 load_in_8bit=True
- 모델의 가중치와 활성화 함수의 출력이 8 bit 정밀도로 저장되어 처리
- 활성화 함수의 출력이란? => ReLU, Sigmoid 등의 활성화 함수가 입력 신호를 받아 처리한 결과
- 기존 학습된 32 bit(float) 형식에 비해 8 bit 정밀도로 로드시 메모리 사용량이 감소 => 데이터 전송 속도가 빨라지고 연산이 더 효율적으로 이뤄져 모델 훈련과 추론 속도 향상
- 32 bit(float) 형식에 비해 8 bit 의 정밀도가 낮아 모델의 정확도와 같은 성능은 떨어질 수 있음
- 하드웨어가 8bit 연산을 지원하는지 확인하고 사용해야함
- 모델 훈련 전 모델을 얼리기(freeze)
- param.requires_grad = False 를 이용
- 모델의 모든 파라미터에 대해 역전파 과정에서 그라디언트 계산을 수행하지 않도록 설정
- 이는 추후 어댑터(adapter)를 훈련할 때를 제외하고는 모델의 가중치가 업데이트되지 않도록 하는 데 사용
- 모델 훈련시 Gradient Checkpointing 적용
- model.gradient_checkpointing_enable() 함수를 이용
- 메모리 집약적 대규모 모델 훈련시 메모리 효율성을 크게 향상
- 기존 모델의 학습시에는 backpropagation 에서 gradient 계산을 하기 위해 신경망은 forward pass 동안의 각 레이어 출력인 activations를 저장하는데 이를 모두 저장하는것은 메모리를 많이 소비
- 모든 activations를 저장하는 대신 특정 레이어에서만 activations 를 저장하고 필요시 나머지를 다시 계산하여 전체적인 메모리 사용량을 줄임
- 이를 이용하면 메모리 절약으로 더 큰 배치크기로 훈련이 가능하거나 GPU 메모리 제약으로 Train 이 힘든 LM 도 훈련 가능
- 더 큰 배치크기로 훈련하면 한 번의 업데이트에 더 많은 데이터를 사용하여 학습 과정에서 더 안정적인 경사 하강을 가능하게 하여 안정성을 가져옴
- 하지만 activations 를 다시 계산해야하니 계산시간 증가로 훈련 시간은 증가할 수 있음
- LoRA(Low-Rank Adaptaion) / QLoRA(Quantized Low-Rank Adaptation)와 같은 PEFT 사용
- LoraConfig() 사용
- LoRA 방법은 모델의 기존 가중치는 고정한 뒤 저차원 행렬을 도입하여 기존 가중치에 적용 => full fine tuning 에 비해 시간 절약
- 전체 모델의 모든 가중치를 미세 조정하는 대신, 작은 수의 추가 파라미터를 사용하여 모델의 일부분만을 조정
- config 에 설정한 하이퍼 파라미터에 의존 (복잡도 조정 가능 => 모델 용량과 성능 균형조절 가능)
- 모델 훈련시 fp16 사용
- transformers.TrainingArguments() 옵션에 fp16=True 설정
- 훈련시 16비트 부동소수점(Half Precision)을 사용
- 메모리 사용량을 줄임 => FP16 은 FP32 부동소수점(Full Precision)에 비해 메모리 사용량 절만 사용
- 훈련 / 추론 속도를 높임 => 메모리의 대역폭이 감소하여 연산 속도가 빨라짐
- FP32 에 비해 FP16 은 낮은 정밀도로 정확도에 영향을 줄 수 있음 => FP16 과 FP32 를 혼합하여 사용 (Mixed Precision)
- 하드웨어가 FP16 연산을 지원하는지 확인하고 사용해야함
- 모델 훈련시 Gradient Accumulation 사용
- transformers.TrainingArguments() 옵션에 gradient_accumulation_steps 사용
- 모델 훈련 시 배치에 대한 gradient 를 즉시 업데이트 하지 않고, 여러 스탭에 걸쳐 누적된 gradient 를 업데이트
- 더 큰 배치 크기를 학습시킨 효과를 냄 => 큰 배치 크기는 gradient update의 분산을 줄임 => 안정적인 최적화
- 각 스탭에서 처리하는 실제 데이터 양이 작게 유지되어 GPU 메모리 사용량 감소
- DeepSpeed / Acceralate
- Accelerate는 자동으로 데이터 병렬 처리, 모델 병렬 처리를 관리하며, 간단한 설정을 통해 여러 GPU 또는 TPU에서 효율적으로 모델을 훈련할 수 있도록 지원
- DeepSpeed는 마이크로소프트에서 개발한 라이브러리 / 제로 Redundancy Optimizer(ZeRO), 모델 병렬 처리, 그리고 그라디언트 체크포인팅과 같은 기술을 통해 메모리 사용량을 현저히 줄임
- 작은 메모리를 가진 하드웨어에서도 대규모 모델을 훈련
- ZeRO는 모델의 파라미터, 그라디언트, 옵티마이저 상태를 효율적으로 관리하여, 기존 방식 대비 메모리 사용량을 크게 줄이면서도 훈련 속도를 향상
- 위 두 내용은 아래에서 추가 기술
# 관련코드
* https://github.com/wotres/deep-learning-study/blob/main/Large_Model_Memory_%EC%A0%88%EC%95%BD_%EA%B8%B0%EB%B2%95.ipynb