LLM 미세 조정 메모리 최적화 기법 TACO 소개
TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning
대규모 언어 모델(LLM)의 전체 파라미터 미세 조정 시 발생하는 막대한 옵티마이저 메모리 문제를 해결하기 위해 TACO라는 새로운 최적화 기법이 제안되었습니다.
- TACO는 옵티마이저 상태를 극도로 줄이는 동시에 1차 그래디언트를 유지하는 방식으로, 기존 대비 메모리 사용량을 최대 174배까지 낮추고 학습 효율성을 높였습니다.
- 이 기술은 이전에는 어려웠던 30~320억 개 파라미터 규모의 대형 모델을 단일 GPU 환경에서도 전체 파라미터를 활용해 미세 조정할 수 있게 하는 혁신적인 성능 향상을 제공합니다.
- TACO는 기존 최적화 기법과 비교하여 정확도와 실행 시간 면에서 유사한 수준을 유지하면서, 압도적인 메모리 절감 효과를 달성했다는 것이 핵심 장점입니다.
(LLMs)의 전체 은 막대한 옵티마이저 상태 메모리를 요구하여, 현대 에서 처리 가능한 모델 크기에 제한을 초래한다. 기존 접근 방식들은 옵티마이저 상태를 압축하거나 1차 그래디언트를 포기하는 등의 방법을 사용해왔다. 이에 연구진은 정확도와 계산 효율성을 희생하지 않으면서 메모리 오버헤드를 줄이는 새로운 최적화 기법인 TACO(Ternary Absolute-max Column-wise One-sparse Optimizer)를 제안했다.
TACO는 Muon 옵티마이저의 연산자 노름 하강 경사 관점을 따르지만, 이를 더욱 발전시킨 방식이다. 이 방법은 차원 정규화된 $1 o1$ 연산자 노름 하에서 정확한 최단 하강 방향을 계산하며, 2차원 행렬의 각 열에서 가장 큰 크기의 부호를 선택하는 방식으로 구현된다. 이를 통해 옵티마이저 상태 메모리를 거의 무시할 수 있는 수준으로 유지하면서도 1차 그래디언트를 보존한다.
실제 테스트 결과, TACO는 OPT-13B 모델을 기준으로 AdamW8bit 대비 지속적인 옵티마이저 상태를 $174 imes$ 감소(27.7 GB $ o$ 0.16 GB)시키고, 최대 학습 메모리를 $2.9 imes$ 줄여(80.6 GB $ o$ 27.5 GB), 유사한 정확도와 실행 시간을 달성했다. 이 기술 덕분에 단일 80 GB H100 GPU 환경에서도 30~32B 파라미터 규모의 모델 전체 파라미터를 활용하여 미세 조정이 가능해졌다.
용어 풀이
- 대규모 언어 모델
- 방대한 글을 학습해 문장을 이해하고 만들어 내는 AI 모델. ChatGPT, Claude, Gemini가 여기에 속해요.
- 파라미터
- 모델이 학습하면서 조정한 내부 숫자. 개수가 많을수록 대체로 모델이 크고 무거워요.
- 미세 조정
- 이미 학습된 모델을 특정 목적의 데이터로 조금 더 학습시키는 일.
- GPU
- 많은 계산을 한꺼번에 처리하는 칩. AI 학습과 실행에 가장 많이 써요.
- 가중치
- 학습으로 정해진 모델 내부의 숫자 값. 내려받는 모델 파일의 본체예요.
대규모 언어 모델을 미세 조정할 때 어떤 문제가 발생하나요?
LLM의 전체 파라미터를 미세 조정하려면 막대한 옵티마이저 상태 메모리가 필요해서, 현대 GPU가 처리할 수 있는 모델 크기에 제한이 생겨요.
TACO는 어떤 방식으로 메모리 사용량을 줄이나요?
옵티마이저 상태를 거의 무시할 수 있는 수준으로 유지하면서도 1차 그래디언트를 보존해요. 이는 차원 정규화된 $1 o1$ 연산자 노름 하에서 정확한 최단 하강 방향을 계산하는 방식으로 구현돼요.
TACO를 사용했을 때 구체적으로 어떤 성능 향상이 있었나요?
OPT-13B 모델을 기준으로 옵티마이저 상태가 기존 대비 174배 감소했고, 최대 학습 메모리도 줄어들었어요. 이 덕분에 단일 GPU 환경에서도 대형 모델 전체 파라미터 활용 미세 조정이 가능해졌어요.