Tian Yuandong 등의 새로운 작업: 메모리 병목 현상을 극복하고 4090 사전 훈련된 7B 대형 모델 허용

Meta FAIR 지난 달 참여한 연구 프로젝트 Tian Yuandong이 폭넓은 호평을 받았습니다. "MobileLLM: 장치 내 사용 사례를 위한 10억 미만 매개 변수 언어 모델 최적화"라는 논문에서 그들은 모바일 장치에서 대규모 언어 모델을 실행한다는 목표를 달성하기 위해 10억 개 미만의 매개 변수를 사용하여 작은 모델을 최적화하는 방법을 탐색하기 시작했습니다. .

3월 6일 Tian Yuandong 팀은 LLM 메모리 효율성 향상에 초점을 맞춘 최신 연구 결과를 발표했습니다. 연구팀에는 Tian Yuandong 자신 외에도 캘리포니아 공과대학(California Institute of Technology), 텍사스 대학교 오스틴 캠퍼스(University of Texas at Austin) 및 CMU의 연구원도 포함되어 있습니다. 본 연구의 목표는 LLM 메모리의 성능을 더욱 최적화하고 향후 기술 개발을 위한 지원과 지침을 제공하는 것입니다.

그들은 전체 매개변수 학습이 가능한 GaLore(Gradient Low-Rank Projection)라는 훈련 전략을 공동으로 제안했습니다. 이는 LoRA와 같은 일반적인 하위 순위 적응 방법에 비해 GaLore는 메모리 효율성이 더 높습니다.

이 연구에서는 모델 병렬 처리, 체크포인트 또는 오프로드 전략을 사용하지 않고도 7B 모델이 24GB 메모리(예: NVIDIA RTX 4090)를 갖춘 소비자 GPU에서 성공적으로 사전 훈련될 수 있음을 처음으로 보여줍니다.

논문 주소: https://arxiv.org/abs/2403.03507

논문 제목: GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

다음을 살펴보겠습니다. 기사의 주요 내용입니다.

현재 LLM(Large Language Model)은 여러 분야에서 뛰어난 잠재력을 보여주고 있지만 현실적인 문제도 직면해야 합니다. 즉, LLM을 사전 훈련하고 미세 조정하려면 많은 양의 컴퓨팅 리소스가 필요할 뿐만 아니라, 그러나 또한 많은 양의 메모리 지원이 필요합니다.

LLM의 메모리 요구 사항에는 수십억 개의 매개변수뿐만 아니라 저장소 자체보다 클 수 있는 기울기 및 최적화 상태(예: Adam의 기울기 모멘텀 및 분산)도 포함됩니다. 예를 들어 단일 배치 크기를 사용하여 처음부터 사전 훈련된 LLaMA 7B에는 최소 58GB의 메모리가 필요합니다(훈련 가능한 매개변수의 경우 14GB, Adam Optimizer 상태 및 가중치 기울기의 경우 42GB, 활성화의 경우 2GB). 이로 인해 24GB 메모리를 갖춘 NVIDIA RTX 4090과 같은 소비자급 GPU에서는 LLM 교육이 불가능해집니다.

위 문제를 해결하기 위해 연구자들은 사전 학습 및 미세 조정 중에 메모리 사용량을 줄이기 위한 다양한 최적화 기술을 계속 개발하고 있습니다.

이 방법은 최적화 상태에서 메모리 사용량을 65.5% 줄이면서 최대 19.7B 토큰이 포함된 C4 데이터 세트를 사용하는 LLaMA 1B 및 7B 아키텍처에서 사전 학습의 효율성과 성능을 유지하고 GLUE에서 작업에 대한 RoBERTa의 효율성과 성능. BF16 기준에 비해 8비트 GaLore는 최적화 메모리를 82.5%, 총 훈련 메모리를 63.3% 더 줄입니다.

이 연구를 본 네티즌들은 "이제 클라우드와 HPC를 잊어야 할 때입니다. GaLore를 사용하면 모든 AI4Science가 2,000달러짜리 소비자급 GPU에서 완성될 것입니다."라고 말했습니다.

Tian Yuandong은 다음과 같이 말했습니다. "GaLore를 사용하면 이제 24G 메모리를 갖춘 NVidia RTX 4090s에서 7B 모델을 사전 훈련하는 것이 가능해졌습니다.

LoRA와 같은 낮은 순위의 가중치 구조를 가정하지 않았지만 가중치 기울기가 자연스럽게 낮다는 것을 증명했습니다. 순위를 부여하므로 (다양한) 저차원 공간에 투영될 수 있습니다. 따라서 기울기, Adam 모멘텀 및 분산에 대한 메모리를 동시에 절약할 수 있습니다.

따라서 LoRA와 달리 GaLore는 훈련 역학을 변경하지 않으며 메모리를 많이 소모하는 워밍업 없이 7B 모델 사전 훈련을 시작하면 GaLore를 미세 조정에도 사용할 수 있어 LoRA와 비슷한 결과를 얻을 수 있습니다.

방법 소개

앞서 언급했듯이 GaLore는 전체 매개변수 학습을 허용하는 훈련 전략이지만 일반적인 낮은 순위 적응 방법(예: LoRA)보다 메모리 효율적입니다. GaLore의 핵심 아이디어는 가중치 행렬을 직접적으로 low-rank 형태로 근사시키려는 것이 아니라, 가중치 행렬 W의 그래디언트의 천천히 변화하는 low-rank 구조를 활용하는 것입니다.

이 기사에서는 먼저 경사 행렬 G가 훈련 과정에서 낮은 순위가 된다는 것을 이론적으로 증명했습니다. 이론을 바탕으로 이 기사에서는 GaLore를 사용하여 두 개의 투영 행렬 및 을 계산하여 경사 행렬 G를 투영합니다. 낮은 순위 형식 P^⊤GQ. 이 경우 구성요소 그라데이션 통계에 의존하는 최적화 상태의 메모리 비용을 크게 줄일 수 있습니다. 표 1에서 볼 수 있듯이 GaLore는 LoRA보다 메모리 효율성이 더 높습니다. 실제로 이는 LoRA에 비해 사전 훈련 중에 메모리를 최대 30%까지 줄일 수 있습니다.

이 기사는 GaLore가 사전 훈련 및 미세 조정에서 좋은 성능을 발휘한다는 것을 증명합니다. C4 데이터 세트에서 LLaMA 7B를 사전 훈련할 때 8비트 GaLore는 8비트 최적화 프로그램과 레이어별 가중치 업데이트 기술을 결합하여 최적화 프로그램 상태에서 10% 미만의 메모리 비용으로 전체 순위에 필적하는 성능을 달성합니다.

사전 훈련의 경우 GaLore는 ReLoRA와 같은 전체 등급 훈련을 요구하지 않고 훈련 과정 전반에 걸쳐 메모리를 낮게 유지한다는 점에 주목할 가치가 있습니다. GaLore의 메모리 효율성 덕분에 처음으로 LLaMA 7B는 값비싼 메모리 오프로딩 기술 없이도 24GB 메모리를 갖춘 단일 GPU(예: NVIDIA RTX 4090)에서 처음부터 훈련할 수 있습니다(그림 1).

GaLore는 그라디언트 투영 방법으로 최적화 프로그램 선택에 독립적이며 알고리즘 1과 같이 단 두 줄의 코드만으로 기존 최적화 프로그램에 쉽게 연결할 수 있습니다.

다음 그림은 GaLore를 Adam에 적용하는 알고리즘을 보여줍니다.

실험 및 결과

연구원들은 GaLore의 사전 훈련과 LLM의 미세 조정을 평가했습니다. 모든 실험은 NVIDIA A100 GPU에서 수행되었습니다.

성능을 평가하기 위해 연구원들은 GaLore를 적용하여 C4 데이터 세트에서 LLaMA를 기반으로 한 대규모 언어 모델을 훈련했습니다. C4 데이터 세트는 주로 언어 모델과 단어 표현을 사전 훈련하는 데 사용되는 Common Crawl 웹 크롤링 코퍼스의 거대하고 정제된 버전입니다. 실제 사전 훈련 시나리오를 가장 잘 시뮬레이션하기 위해 연구원들은 데이터를 복제하지 않고 모델 크기가 최대 70억 개에 달하는 매개변수를 사용하여 충분히 많은 양의 데이터를 훈련했습니다.

이 문서는 RMSNorm 및 SwiGLU 활성화와 함께 LLaMA3 기반 아키텍처를 사용하는 Lialin 등의 실험 설정을 따릅니다. 학습률을 제외한 각 모델 크기에 대해 동일한 하이퍼 매개변수 세트를 사용하고 모든 실험을 BF16 형식으로 실행하여 동일한 계산 예산으로 각 방법의 학습률을 조정하면서 메모리 사용량을 줄이고 최적의 성능을 보고했습니다.

또한 연구원들은 GaLore 및 LoRA의 메모리 효율적인 미세 조정을 위한 벤치마크로 GLUE 작업을 사용했습니다. GLUE는 감정 분석, 질문 답변, 텍스트 상관관계 등 다양한 작업에서 NLP 모델의 성능을 평가하기 위한 벤치마크입니다.

본 논문에서는 먼저 Adam 옵티마이저를 사용하여 GaLore와 기존 하위 순위 방법을 비교했으며 그 결과는 표 2에 나와 있습니다.

연구원들은 GaLore를 다양한 학습 알고리즘, 특히 메모리 효율적인 최적화 프로그램에 적용하여 메모리 사용량을 더욱 줄일 수 있음을 입증했습니다. 연구원들은 GaLore를 AdamW, 8비트 Adam 및 Adafactor 최적화 프로그램에 적용했습니다. 그들은 성능 저하를 피하기 위해 1차 통계적 Adafactor를 사용합니다.

실험을 통해 LLaMA 1B 아키텍처에서 10K 학습 단계로 평가하고 각 설정에 대한 학습 속도를 조정한 후 최고의 성능을 보고했습니다. 그림 3에 표시된 것처럼 아래 그래프는 GaLore가 AdamW, 8비트 Adam 및 Adafactor와 같은 널리 사용되는 최적화 프로그램과 함께 작동함을 보여줍니다. 게다가 아주 적은 수의 하이퍼파라미터를 도입해도 GaLore의 성능에는 영향을 미치지 않습니다.

표 4에서 볼 수 있듯이 GaLore는 대부분의 작업에서 더 적은 메모리 사용량으로 LoRA보다 높은 성능을 달성할 수 있습니다. 이는 GaLore가 LLM 사전 훈련 및 미세 조정을 위한 전체 스택 메모리 효율적인 훈련 전략으로 사용될 수 있음을 보여줍니다.

그림 4에서 볼 수 있듯이 BF16 벤치마크 및 8비트 Adam과 비교하면 8비트 GaLore는 훨씬 적은 메모리를 필요로 하며 LLaMA 7B를 사전 훈련할 때 22.0G 메모리만 필요하며 각 GPU의 토큰 배치 크기는 다음과 같습니다. 더 작습니다(최대 500개의 토큰).

더 자세한 기술적인 내용은 원본 논문을 읽어보세요.

위 내용은 Tian Yuandong 등의 새로운 작업: 메모리 병목 현상을 극복하고 4090 사전 훈련된 7B 대형 모델 허용의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!