Xiao Xiao는 Aofeisi
Qubit에서 왔습니다 | 공개 계정 QbitAI
대형 모델에 "작은 부품을 추가"하기만 하면 추론 속도가 즉시 2배 증가합니다!
추가 모델을 훈련하거나 컴퓨팅 하드웨어를 최적화할 필요가 없으며 단일 A100을 몇 시간 만에 미세 조정할 수 있습니다.
이 새로운 연구는 메두사(Medusa)라고 불리며 프린스턴, UIUC, CMU 및 코네티컷 대학교에서 나온 것이며 FlashAttention의 저자인 Tri Dao도 그 중 하나입니다.
현재 Berkeley의 70억 매개변수 "Lama" Vicuna 에 성공적으로 배포되었으며 향후 다른 대형 모델도 지원할 예정이며 이미 GitHub 인기 목록에 표시되었습니다.
하지만 실제로 이 방법이 도입되기 전에는 업계에 대규모 모델 추론 가속화 방법이 없었고, 주류는 DeepMind에서 출시한 추측 샘플링(추측 디코딩)이었습니다.
이 방법과 비교할 때 메두사의 차이점은 무엇입니까?
추측 샘플링의 두 가지 "버그"
대규모 모델 추론 속도를 높이려면 먼저 속도를 "제한"하는 것이 무엇인지 알아야 합니다.
계산량의 증가에 비해 대형 모델의 추론 속도는 메모리 대역폭 (메모리 바운드) 의 영향을 받기 쉽습니다 .
이는 대형 모델이 캐시 용량을 훨씬 초과하는 거대한 매개변수를 갖기 때문입니다. 따라서 추론 중에 외부 메모리(비디오 메모리)에서 캐시로 가중치를 한 번 읽어야 합니다. 이 프로세스는 메모리 대역폭에 의해 제한되며 일반적으로 매우 느립니다. .
따라서 모델이 일괄 추론을 할 때 한 번에 100개의 토큰을 처리하는 것과 1개의 토큰을 처리하는 것 사이의 시간에는 큰 차이가 없습니다.
이 기능을 바탕으로 딥마인드는 지난해 11월 추측 샘플링 이라는 마법 같은 연산을 고안해냈다.
더 작은 모델(초안 모델)을 훈련하고 큰 모델에 대한 "후보 단어" 배치를 미리 생성합니다. 큰 모델이 스스로 "생각"하고 생성하도록 하는 대신 직접 "선택"하면 됩니다.
소형 모델은 대형 모델보다 몇 배 빠르게 생성되기 때문에, 대형 모델이 소형 모델에서 이미 사용 가능한 단어를 '사용 가능'하다고 느끼면 다시 천천히 생성할 필요 없이 바로 사용할 수 있습니다.
이 프로세스는 입력 방법의 연관된 단어 후보와 약간 비슷합니다 . 우리(대형 모델)가 다음 단어에 사용할 항목을 결정하기 전에 입력 방법(소형 모델)은 먼저 몇 가지 대안을 나열합니다.
마음에 드는 것이 있으면 하나만 선택하여 사용하고, 생성된 것 중 어느 것도 좋지 않다고 생각되면 그냥 통과시키고 다시 시도하세요.
이 추측적 샘플링 방법은 실제로 놀라운 결과를 얻었으며 M2 Ultra에서 340억 개의 매개변수 LLaMA 대형 모델을 높은 정확도로 쉽게 실행할 수도 있습니다.
그러나 이 접근 방식에는 두 가지 문제가 있습니다.
한편으로는 큰 모델 에 대한 "후보 단어"를 생성하는 작은 초안 모델을 찾는 것이 그리 쉽지 않습니다 .
이 소형 모델은 단순히 사용할 수 있는 생성 모델이 아니며, 통합된 인터페이스 및 긴밀한 확률 분포와 같은 요구 사항 외에도 생성 품질이 대형 모델보다 훨씬 나쁠 수는 없습니다.
Meta에서 출시한 LLaMA 같은 모델이면 괜찮을 것 같습니다. 수백억 개의 매개변수를 포함하는 대형 모델 버전과 수십억 개의 매개변수를 포함하는 소형 모델 버전이 있습니다. 작은 버전을 초안 모델로 사용할 수 있습니다.
그러나 다른 대규모 오픈 소스 모델의 경우 이 방법은 그다지 적용 가능하지 않으며, 작은 모델을 직접 구축하고 훈련하려면 시간이 더 많이 걸릴 뿐만 아니라 생성 효과도 기대에 미치지 못할 수 있습니다.
반면, 이중 모델의 조합은 후속 시스템 튜닝을 더욱 복잡하게 만듭니다 .
시스템인 대형 모델 자체에 비하면 새로 추가된 초안 모델은 또 다른 시스템을 도입한 것과 같기 때문이다.
이로 인해 추가 네트워크 전송 및 고려해야 할 다양한 하드웨어 조건을 포함하여 모델 배포가 더욱 복잡해지고 계산 최적화 중에 난이도가 더욱 높아집니다.
이러한 문제를 해결하기 위해 메두사가 등장했습니다.
작은 모델은 필요하지 않습니다. "헤드" 몇 개만 추가하면 됩니다.
메두사 (머리가 여러 개인 괴물인 메두사)는 대규모 모델 추론을 가속화하는 새로운 방법입니다.
추측적 샘플링과 비교하여 Transformer 대형 모델에 여러 개의 디코딩 헤드를 직접 추가하는 것을 선택하며 각 헤드는 단일 레이어 피드포워드 네트워크입니다.
이러한 추가 디코딩 헤드를 사용하면 대형 모델이 하나씩 "치약을 짜는" 대신 한 번에 여러 단어를 직접 생성할 수 있습니다.
생성 정확도도 좋은데, '다음 단어의 다음 단어'를 예측할 때 메두사의 정확도는 60%에 달하며 아직 최적화 중이다.
그 후, 추론 가속을 달성하기 위해 트리 기반 주의 메커니즘을 사용하여 이러한 단어를 병렬로 검증합니다.
Medusa를 기반으로 Vicuna의 70억, 130억, 330억 매개변수 대형 모델 추론 속도가 1.9배 이상 향상되었습니다 .
연구진은 70억 개의 매개변수 모델에 대해 다양한 작업에 대한 가속 효과도 테스트했는데, 코드 생성 속도가 최대 2.15배 증가한 것으로 나타났습니다.
가장 중요한 것은 Medusa를 사용한 후에 전체 대형 모델을 재교육할 필요가 없다는 것입니다.
반면, 대형 모델의 매개변수를 동결하는 것만으로도 대형 모델과 함께 훈련할 수 있으며, 단일 GPU로도 가능합니다.
추가적인 모델이 추가되지 않기 때문에 분산 추론에도 친화적입니다.
작가에 대해
이 연구에는 두 명의 공동 저자가 있습니다.
공저자인 Cai Tianle는 프린스턴대학교 박사과정 학생으로, 최적화, 표현학습, 건축설계 등의 연구 관심분야를 갖고 있으며 북경대학교 수리과학대학원에서 응용수학과 컴퓨터공학을 복수전공했습니다.
공동저자는 University of Illinois at Urbana-Champaign(UIUC)의 박사과정 학생인 Yuhong (Jesse) Li입니다. 그의 연구 방향은 효율적인 기계 학습이며, 베이징 우편 통신 대학에서 학사 학위를 취득했습니다.
또한 이번 연구에는 FlashAttention의 저자이자 스탠포드 출신 박사인 Tri Dao도 참여했습니다.
FlashAttention은 Attention 속도를 높이고 메모리 사용량을 줄일 수 있는 방법으로, PyTorch의 표준 Attention 구현과 비교하여 최대 9배 더 빠를 수 있습니다.
GitHub 주소:
https://github.com/FasterDecoding/Medusa
연구 주소:
https://sites.google.com/view/medusa-llm