오늘, Answer.AI의 첫 번째 프로젝트를 공개합니다: 정규 데스크톱 컴퓨터에서 두 개 이상의 표준 게이밍 GPU(RTX 3090 또는 4090)로 효율적으로 70b 크기의 대형 언어 모델을 처음으로 훈련할 수 있는 완전 오픈 소스 시스템입니다. 이 시스템은 FSDP와 QLoRA를 결합한 것으로, Answer.AI, Tim Dettmers(워싱턴 대학), Hugging Face의 Titus von Koeller 및 Sourab Mangrulkar의 협력 결과입니다.
이 시스템은 오픈 소스 커뮤니티가 더 나은 모델을 공개하는 데 도움이 될 것입니다. 오버 50만 다운로드를 기록한 매우 인기 있는 OpenHermes 모델과 데이터셋의 창조자인 Teknium은 다음과 같이 말했습니다:
“이 능력을 통해 우리는 거대한 모델을 지역적으로 새로운 높이로 끌어올릴 수 있으며, 수십억 개의 매개변수를 가진 거대한 모델들이 이제 작은 연구소에서도 접근 가능해졌습니다.”
Answer.AI에서는 이것을 첫 번째 프로젝트로 만들었습니다. 왜냐하면 이것은 우리의 핵심 가치인 모든 사람에게 유용한 AI를 이용할 수 있도록 돕는 것의 기반입니다. 다른 사람들의 모델을 사용하는 것만으로는 충분하지 않습니다. 우리는 모든 사람이 자신만의 맞춤형 모델을 만들 수 있기를 원합니다. 그래서 그들이 자신의 AI 시스템을 통제할 수 있도록 하고 싶습니다.
딥러닝 모델을 학습하는 데 사용되는 하드웨어는 두 가지 매우 다른 수준이 있습니다. H100s 및 A100s와 같은 데이터 센터급 하드웨어가 있으며, 이는 수십만 달러에 이릅니다. 그리고 게이밍 GPU가 장착된 데스크탑 컴퓨터도 있습니다. 예를 들어, 듀얼 4090을 장착한 데스크탑은 10,000달러 미만에 구입할 수 있습니다(사전 제작 시스템의 절반 가격 이하로 2차 손상 부품을 사용하여 조립할 수도 있습니다).
하지만 여기가 핵심입니다: 게이밍 GPU는 10배 이상 비싼 데이터 센터 GPU와 유사한 성능을 가지고 있습니다! 이 10배 저렴한(거의 똑같이 빠른) 카드를 사용하여 대형 언어 모델을 훈련할 수 있다면 좋을텐데, 그러나 그들은 훨씬 적은 메모리를 가지고 있기 때문에 그렇게 할 수 없습니다. 현재 사용 가능한 최고의 데이터 센터 카드는 80GB RAM을 가지고 있지만, 게이밍 카드는 최대 24GB RAM까지만 지원합니다. 가장 큰 모델만이 최상의 결과를 내기 때문에, 최고의 모델을 만드는 것은 대부분의 사람들에게는 크게 접근하기 어려웠습니다.
우리는 실제로 이에 대한 본질적인 이유가 없다는 것을 깨달았습니다. 초고속 하드웨어는 이미 존재하며, 사용될 준비가 되어 있습니다. 단지 모델과 데이터를 메모리 제약 조건을 충족하는 방식으로 공급할 방법이 필요할 뿐입니다. 당연한 질문은: 그렇다면 왜 이것이 이루어지지 않았을까요? 대규모 산업 연구소들은 이미 10배 비싼 하드웨어를 보유하고 있기 때문에, 이를 해결하려는 동기가 별로 없습니다.
큰 아이디어는 간단합니다: 이러한 더 저렴하고 낮은 메모리를 가진 게임용 GPU를 사용하여 최고의 오픈 소스 모델을 훈련하는 방법을 찾아보세요. 따라서 목표는 이겁니다: 게임용 GPU만을 사용하여 700억 개의 파라미터(70b) 모델을 훈련하는 것입니다. 이는 각 GPU의 메모리가 최대 24GB일 것을 의미합니다. 각 파라미터는 보통 16비트(2바이트)를 차지하기 때문에 가중치를 저장하는 데는 70*2=140GB가 필요하며, 활성화, 그래디언트 및 최적화 상태와 같은 다른 모든 데이터를 포함하지 않은 상태입니다!
Answer.AI는 매우 특이한 유형의 조직입니다 - 오늘날의 AI 연구 그룹보다는 영적으로 19세기 전기 실험실에 더 가까운 이윤을 추구하는 연구 및 개발 연구소입니다. 대규모 모델 훈련을 저렴하고 접근 가능하게 만드는 방법을 찾는 것은 지난해 NeurIPS에서 출범한 이후 Eric Ries와 Jeremy Howard가 조직이 할 수 있기를 희망했던 일종의 일입니다.
이 문제를 해결하는 것은 어렵습니다. 많은 별개의 라이브러리(e.g bitsandbytes, PEFT, Transformers, Accelerate, 그리고 PyTorch)와 컴퓨터 과학 및 수학 개념(e.g 이산화, 분산 컴퓨팅, GPU 프로그래밍, 선형 대수, 그래디언트 체크포인팅과 같은 SGD 개념)을 이해해야 하며, 그들이 모두 어떻게 상호작용하는지 알아야 합니다.
학계는 어려운 문제를 해결하는 뛰어난 사람들로 가득하지만, 이 특정 문제를 해결하지 못했습니다. 대학 연구자들이 이러한 종류의 작업에 시간을 할애하는 것을 정당화하기 어려운데, 기존 도구와 기술을 결합하는 것은 일반적으로 충분히 '새로운' 것으로 간주되지 않아 높은 영향력을 가진 저널에 게재되기 어렵습니다. 하지만 이것이 학계가 필요로 하는 통화입니다. 게다가, 학계는 일반적으로 자신의 분야에서 높은 전문성을 갖추기를 기대받기 때문에, 이렇게 많은 조각들을 하나의 해결책으로 통합하는 것이 어려운 과제가 됩니다.
또한, 물론 대형 기술 기업들도 어려운 문제를 해결하는 뛰어난 인재들로 가득합니다. 그러나 소비자용 GPU로 모델을 학습하는 이 문제는 그들이 해결해야 할 문제가 아닙니다 - 이미 비싼 대형 GPU를 구입했기 때문입니다! 많은 스타트업도 어려운 문제를 해결하는 뛰어난 인재들로 가득합니다! 그러나 Eric Ries가 설명한 대로, '오늘날의 금융 시장은 기업에게 단기 이익을 모든 것보다 우선시하도록 강요합니다'. 스타트업이 투자자들에게 자신들의 자금을 오픈 소스 소프트웨어와 공개 연구에 사용하는 이유를 정당화하는 것은 매우 어렵습니다.
학계, 대형 기술 기업 및 스타트업이 이 문제를 해결하지 않은 이유는 있었지만, 이러한 이유들이 바로 이 문제를 해결하기에 적합했던 이유들이었습니다. 회사에서 일하는 모든 사람들은 이 문제에 대해 작업해야 했던 시스템들을 구축해 왔기 때문에, 모든 조각들이 어떻게 맞아 떨어지는지 이해할 수 있었습니다. 소프트웨어와 AI의 기초를 깊이 이해하고, 재미있고 흥미로운 end-to-end 시스템을 해킹하는 것을 좋아하는 사람들이 Answer.AI에 끌리는 사람들이며, 그 반대도 마찬가지입니다.
함께 해결하려고 선택하는 문제들은 해결할 사람들에 의해 선택됩니다. 그래서 우리는 여러 아이디어를 결합하여 실용적인 해결책을 만드는 프로젝트를 선택하는 경향이 있습니다. 그리고 우리는 AI, 오픈 소스 소프트웨어로부터 장기적인 이점을 얻기 위한 특허를 가진 공익 회사이기 때문에, 공개 연구와 오픈 소스 소프트웨어가 우리의 미션과 직접적으로 일치합니다.
최근에 이를 실현하기 위한 첫 번째 중요한 단계를 거친 두 프로젝트가 출시되었습니다: QLoRA (Tim Dettmers et al가 개발), 그리고 FSDP (Meta의 PyTorch 팀가 개발).
QLoRA는 현대 신경망에서 중요한 두 가지 발전인 _양자화_와 _LoRA_의 뛰어난 조합입니다. 양자화는 신경망의 가중치를 저장할 때 16비트나 심지어 32비트를 사용하는 대신 4비트(또는 그 이하)를 사용하는 기술입니다. 4비트 숫자에는 16가지 가능한 값만 있지만 Dettmers와 Zettlemoyer가 보여준 바에 따르면, 오늘날 인기 있는 대형 언어 모델에서 충분할 수 있다고 합니다. Tim Dettmers는 이 4비트 '양자화' 모델을 bitsandbytes 라이브러리 덕분에 쉽게 만들 수 있었으며, 최근 Hugging Face가 이 라이브러리를 유지 및 문서화하는 데 도움을 주었는데, 이는 특히 Titus von Koeller의 노력 덕분입니다.
불행히도, 모델이 양자화되면 일반적인 방법으로 더 이상 훈련할 수 없습니다. 16가지 가능한 값만 사용하면 모델 훈련에 사용되는 경사 하강법은 거의 모든 곳에서 기울기가 거의 0이 되어, 양자화된 가중치에 업데이트를 할 수 없습니다. 이것은 양자화가 추론에만 사용될 수 있고 계속된 사전 훈련이나 미세 조정에 사용될 수 없다는 주요 문제입니다. 추론은 유용하고 중요하지만, 실제로는 모델을 '소비'하는 것뿐입니다. 그러나 우리는 모든 사람이 모델을 '창조'할 수 있기를 원합니다!
이 제한을 피하는 꿀팁은 LoRA - 'Large Language Models의 저랭크 적응'을 사용하는 것입니다. LoRA는 전체 대형 언어 모델을 전혀 훈련시키지 않고 대신 '어댑터'를 추가합니다. 이 어댑터는 매우 작은 행렬(일반적으로 전체 모델의 1% 미만)로 훈련되며 나머지 모델은 일정한 상태를 유지합니다. Stable Diffusion과 같은 모델을 다뤄본 적이 있다면 이러한 어댑터를 여러 번 본 적이 있을 것입니다. 이것이 이러한 모델이 일반적으로 공유되는 방식이며 그 이유로 이 모델들이 작고 빠르게 다운로드되는 이유입니다.
Tim은 LoRA를 양자화와 결합할 수 있다는 것을 깨달았습니다: 양자화된 기본 모델을 사용하고, 이 모델은 훈련 중에 전혀 변경되지 않으며 양자화되지 않은 훈련 가능한 LoRA 어댑터를 추가합니다. 이 조합은 _QLoRA_입니다. Tim의 팀은 이를 사용하여 GPU보다 큰(양자화되지 않은) 모델을 처음으로 훈련시킬 수 있었습니다: 48GB 카드에서 65b 모델(양자화되지 않은 경우 130GB)을 훈련시켰습니다.
Hugging Face stepped in once again here, creating the PEFT library, which made LoRA training far simpler, and also integrating it directly with bitsandbytes to allow anyone to use QLoRA with just a few lines of code. The Hugging Face team has been working tirelessly behind the scenes to ensure that the open source community can use these technologies to train their models. If you’ve ever used Transformers to load a 4-bit model using a single function argument, then you’ve got them to thank (and even if you haven’t, you’ve almost certainly used the work of folks that have built their model with this ecosystem).
QLoRA는 우리가 해결하려고 했던 문제를 완전히 해결하지는 못했지만, 24GB 카드에 70b 모델을 훈련시키는 것에 가장 가까이 다가갔습니다. 4비트(즉, 0.5바이트)로 양자화된 70b 모델은 70/2 = 35GB가 소요되며, 이는 사용하려는 24GB 게이밍 GPU보다 큽니다.
QLoRA에는 다른 제한 사항이 있습니다. 48GB 카드는 매우 비싸며, 65b 모델을 훈련시키는 것은 그런 카드에 맞게 딱 맞게 들어갑니다. 이것은 문제가 될 수 있습니다. 왜냐하면 훈련 중에 모델의 활성화2, 그래디언트 및 최적화 상태를 포함하여 다른 많은 것들을 저장해야하기 때문입니다. 모델 가중치를 로드한 후에 메모리가 많이 남아 있지 않으면, 훈련을 지원할 충분한 작업 메모리가 부족합니다.
예를 들어, 언어 모델의 장점 중 하나는 '채팅'하거나 이해하거나 긴 문서나 대화를 분석하는 데 사용할 수 있다는 것입니다. 그러한 긴 시퀀스를 처리할 수 있는 모델을 만들려면 훈련 중에 긴 시퀀스의 예시를 보여주어야 합니다. 훈련에 사용된 가장 긴 시퀀스를 '시퀀스 길이'라고 합니다. 48GB 카드에서 65b QLoRA 모델을 훈련할 때 짧은 시퀀스 길이 이외의 것을 사용하려고 하면 오류가 발생할 것입니다. 왜냐하면 시퀀스에 대한 모든 정보를 저장할 메모리가 충분하지 않기 때문에, 거의 모든 메모리가 모델 자체를 저장하는 데 사용됩니다.
또한, 모델이 한 번에 하나의 시퀀스만 볼 수 있다면, 교육 세트의 모든 데이터를 처리하는 데 매우 오랜 시간이 걸릴 것입니다. 따라서 우리는 한 번에 몇 개의 시퀀스를 함께 '배치(batch)'할 수 있도록 하고 싶습니다. 포함된 시퀀스의 수를 '배치 크기(batch size)'라고 합니다. 모델 가중치를 로드한 후 GPU에 공간이 매우 적은 경우, 매우 작은 배치 크기만 사용할 수 있어 굉장히 느린 학습 결과를 초래할 수 있습니다.
단일 소비자 GPU의 RAM 한계 문제에 대한 한 가지 명백한 해결책은 두 개 이상의 GPU를 사용하는 것입니다! 오픈 소스 커뮤니티에서 매우 흔한 접근 방식은 간단히 각 카드에 모델의 몇 개의 레이어를 배치하는 것입니다. 그런 다음 훈련을 위해 첫 번째 GPU에서 처음 몇 레이어를 실행한 다음 두 번째 GPU에서 다음 몇 레이어를 실행하고 이와 같이 계속합니다. 예를 들어, 70b (140GB) 모델은 8개의 24GB GPU에 분산되어 각각 17.5GB를 사용할 수 있습니다. Hugging Face Transformers에는 device_map='auto'라는 편리한 설정이 있습니다. 아마도 사용해 보았을 것입니다. 이 설정이 실제로 하는 일이 바로 이것입니다. 이 방법은 작동하지만 큰 단점이 있습니다. 한 번에 하나의 GPU만 활성화되며 다른 모든 GPU는 '차례'를 기다립니다. 이는 계산의 ⅞이 낭비된다는 것을 의미합니다.
'분산 데이터 병렬 (DDP)은 이전에 여러 GPU를 사용하여 모델을 효율적으로 학습하는 데 사용되는 골드 표준 방법이었습니다. 이를 위해서는 각 GPU에 전체 모델을 유지해야 합니다. 예를 들어, 작은 모델(예: 4GB RAM을 차지하는 2b 모델)의 경우 각 GPU에 전체 모델을 간단히 로드하고, 각 GPU가 병렬로 학습 예제를 처리하도록 할 수 있습니다. 예를 들어, 4개의 GPU가 있다면 학습 속도가 4배 빨라집니다. 그러나 모델이 GPU에 맞지 않고 학습 프로세스에 필요한 데이터를 위한 충분한 공간이 없는 경우 DDP는 작동하지 않습니다.
따라서 우리는 모델을 GPU들 사이에 분할할 수 있는 것이 필요하며(device_map='auto'와 같이) 동시에 그들을 병렬로 사용할 수 있는 것이 필요합니다(DPP와 같이). 이것이 Meta의 Fully Sharded Data Parallel (FSDP) 라이브러리가 등장하는 곳입니다. 이 라이브러리는 큰 모델을 'shard'로 나누어 그 파라미터들을 여러 GPU들 사이에 분산시킴으로써 모든 GPU들이 동시에 사용될 수 있도록 합니다. 훈련 중에 신경망의 한 레이어가 특정 GPU에서 계산될 때, 모든 필요한 shard들이 그곳으로 복사됩니다. 그런 다음 계산이 이루어지고, 마지막으로 복사된 부분들이 해당 GPU에서 삭제됩니다. 이 방법은 매우 비효율적으로 들릴 수 있지만, 실제로는 현재 레이어가 계산 중일 때 다음 레이어의 데이터를 동시에 복사함으로써 이 접근 방식이 DDP와 비교했을 때 속도 저하 없이 결과를 얻을 수 있습니다.
FSDP의 능력은 한 대의 GPU보다 큰 모델에 DDP의 성능을 가져다 줄 수 있다는 것을 밝혀냈습니다. 예를 들어, 70억(70 billion) 개의 파라미터가 있는 양자화되지 않은 모델은 각 파라미터가 16비트(즉, 2바이트)로 저장되기 때문에 140GB의 RAM을 차지합니다. 그러나 심지어 NVIDIA의 H100 카드(한 장에 약 40,000달러의 비용이 듭니다!)도 필요한 것을 충족시키지 못하며 80GB의 RAM을 가지고 있습니다. 그러나 FSDP를 사용하면 4개의 H100 GPU를 결합하여 총 320GB의 RAM을 얻을 수 있습니다.
(그런데, 그런 기계는 약 15만 달러 정도 들 거에요…)
Answer.AI의 핵심 가치는 유용한 AI를 보다 접근 가능하게 만드는 것입니다. 자신만의 고품질 맞춤형 모델을 만들기 위해 15만 달러는 분명히 접근 가능한 금액이 아닙니다! 그래서 우리가 먼저 시작한 프로젝트는 소비자용 게임 GPU를 사용하여 70b 모델을 효율적으로 훈련할 수 있도록 만드는 것이었습니다. 만약 QLoRA를 사용하여 모델 크기를 약 400% 줄일 수 있다면 (따라서 70b 모델이 35GB RAM에 맞게 될 것), 그리고 그 모델을 24GB 소비자용 카드 두 개 이상으로 분할하기 위해 FSDP를 사용한다면, 그것은 모델을 훈련하기에 충분한 RAM이 남을 것입니다.
제레미와 팀은 2023년 말에 FSDP와 QLoRA를 결합하는 아이디어에 대해 논의했습니다. 팀은 제레미를 티투스 폰 켈러와 연결시켰고, 제레미와 티투스는 두 라이브러리를 결합했을 때 발생하는 문제를 시도하고 탐색하며 이해하고 문서화하기 위해 함께 작업했습니다.
Answer.AI의 Johno Whitaker가 중요한 첫 번째 단계를 마련했습니다: 우리가 문제를 더 깊이 이해하고 해결책을 시험할 수 있도록 하는 간단한 독립형 테스트 스크립트입니다. 주요한 획기는 2024년 초에 Answer.AI의 Benjamin Warner와 Titus가 독립적으로 중요한 아이디어를 제시한 것입니다: 양자화된 매개변수를 선택 가능한 데이터 유형에 저장하는 것, 그 저장 데이터 유형이 모델의 '계산 유형'과 동일한 데이터 유형인 것입니다3.
문제를 해결한 후 FSDP를 사용하여 양자화된 모델로 첫 번째 데이터 일괄 처리를 성공적으로 훈련할 수 있었습니다! Benjamin과 Answer.AI의 Kerem Turgutlu는 이를 모두 포장하여 bitsandbytes의 pull request로 필요한 모든 테스트 및 리팩터링과 함께 제출할 수 있었습니다. bitsandbytes 프로젝트의 유지 보수자들에게 극도로 감사드립니다. 그들은 우리의 PR을 프로세스를 통해 안내하는 데 매우 민첩했습니다.
이 시점에서 우리는 다시 한 번 모든 것을 빨리 끝낼 것으로 생각했지만, 다시 한 번 작업의 복잡성을 과소평가했습니다! 우리가 깨달은 첫 번째 것은 여전히 단일 GPU보다 큰 양자화된 모델을 로드하는 것이 실제로 불가능했다는 것이었습니다. 왜냐하면 로딩 및 양자화 과정 자체가 전체 모델을 하나의 GPU에 넣어야 했기 때문입니다.
제레미는 몇 주 동안 Meta의 훌륭한 Llama-Recipes 프로젝트를 주의 깊게 연구했는데, 이는 그가 찾은 최고의 완벽한 FSDP 파인튜닝 구현이었고, bitsandbytes와 함께 작동하는 방식을 밀접히 추적하면서 Hugging Face의 PEFT, Transformers, 그리고 Accelerate 프로젝트와 함께, 모델을 파인튜닝하는 데 필요한 모든 단계를 수동으로 완료하는 최소한의 독립형 스크립트를 구축하는 데 성공했습니다.
벤자민은 몇 가지 조정을 통해 한 번에 한 층씩 로딩 및 이산화를 수행할 수 있다는 것을 깨달았으며, 이를 통해 전체 모델을 단일 GPU에 올려두지 않아도 된다는 것을 알게 되었습니다. 또한 PEFT 라이브러리가 양자화 상태를 CPU로 이동하지 않도록 방지하는 방법을 찾아냈습니다. 케렘은 LoRA 알고리즘의 사용자 정의 구현을 작성하여 벤자민의 변경 사항과 함께 작동할 수 있도록 했습니다.
그리고 그로 인해, 우리는 처음으로 듀얼 3090 게이밍 GPU에서 70b 모델을 섬세하게 조정할 수 있었습니다!
이 작업을 수행하기 위해 FSDP와 QLoRA뿐만 아니라, 지난 몇 년 동안 학계 및 오픈 소스 커뮤니티에서 개발된 다양한 똑똑한 기술들을 활용했습니다. 사용한 것들은 다음과 같습니다:
- Gradient checkpointing (also known as activation checkpointing) to avoid storing full gradients, instead saving activations at a number of ‘checkpoints’ throughout the model and then re-computing gradients by re-running the forward computation step as needed
- CPU offloading to store weights in CPU RAM rather than on the GPU when they are not in use, drastically reducing the GPU memory required. This technique isn’t very useful to the “GPU rich” using H100 GPUs, which have highly optimized ways to pass weights to each other. But for our use case, it’s absolutely necessary, since gaming GPUs and motherboards don’t have these systems
- Flash Attention 2 to efficiently calculate attention using a memory-optimized Cuda kernel.
FSDP와 QLoRA와 함께 작동하도록 하는 것은 항상 간단하지 않았습니다. 예를 들어, CPU 오프로딩을 사용할 때, Benjamin은 bitsandbytes에서 문제를 발견하고 수정했습니다. '오프로딩'된 가중치가 GPU로 다시 복사될 때마다, 자동으로 재양자화되어 사전 훈련된 모델을 무작위 가중치로 변환하는 문제가 있었습니다! 우리는 bitsandbytes에 풀 리퀘스트를 제출하여 이미 양자화된 매개변수를 추적하도록 했으므로 중복 계산을 피할 수 있었습니다.
이 모든 작업을 마친 후, 소비자용 GPU로 대규모 모델을 훈련할 수 있다는 것을 보게 되어 매우 기쁘게 생각했습니다. Jeremy는 원래 llama-recipes의 다양한 GPU 하드웨어 범위에서 자세한 벤치마킹을 실행했고, Kerem은 새 프로젝트를 위한 포괄적인 벤치마킹 시스템을 개발했습니다. 두 가지를 비교해 보니 우리는 여전히 희망했던 시퀀스 길이나 배치 크기를 사용할 수 없다는 것을 깨달았습니다 - 어떤 이유로 인해 예상보다 더 많은 메모리가 사용되고 있었습니다.
자세히 살펴보니, 우리의 FSDP/QLoRA 통합 때문이 아니었음을 알게 되었습니다. 사실 FSDP 없이도 bitsandbytes에서 seqlen을 증가시키면 메모리 사용량이 초선형적으로 증가하여, 양자화 없이도 더 높은 메모리 사용량으로 이어졌습니다! 이 문제를 발견한 사람들은 처음이 아닌 것으로 밝혀졌습니다. 아직 bitsandbytes 솔루션이 없지만(조사 중입니다), 이 문제로 우리는 흥미로운 발견을 하게 되었습니다...