[논문 리뷰] Jamba-1.5: Hybrid Transformer-Mamba Models at Scale

Jamba Team, Barak Lenz외 59명|arXiv (Cornell University)|2024. 08. 22.

Fluid Dynamics and Vibration Analysis인용 수 4

한 줄 요약

Jamba-1.5는 하이브리드 트랜스포머-앰바 미스처 오브 익스퍼트 아키텍처를 기반으로 한 두 가지 지침 퍼티드 대규모 언어 모델—Jamba-1.5-Mini(12B 활성 파라미터)와 Jamba-1.5-Large(94B 활성 파라미터)—를 소개한다. 이 모델은 256K 토큰의 컨텍스트 길이를 지원하며, 상태 기반의 성능을 달성했고, 낮은 메모리 사용량(9GB KV 캐시)과 높은 처리량을 기록했다. 이 모델은 새로운 양자화 기술인 ExpertsInt8를 통해 효율적인 추론이 최적화되었으며, 품질 손실 없이 8x80GB GPU에서의 배포가 가능하다.

ABSTRACT

We present Jamba-1.5, new instruction-tuned large language models based on our Jamba architecture. Jamba is a hybrid Transformer-Mamba mixture of experts architecture, providing high throughput and low memory usage across context lengths, while retaining the same or better quality as Transformer models. We release two model sizes: Jamba-1.5-Large, with 94B active parameters, and Jamba-1.5-Mini, with 12B active parameters. Both models are fine-tuned for a variety of conversational and instruction-following capabilties, and have an effective context length of 256K tokens, the largest amongst open-weight models. To support cost-effective inference, we introduce ExpertsInt8, a novel quantization technique that allows fitting Jamba-1.5-Large on a machine with 8 80GB GPUs when processing 256K-token contexts without loss of quality. When evaluated on a battery of academic and chatbot benchmarks, Jamba-1.5 models achieve excellent results while providing high throughput and outperforming other open-weight models on long-context benchmarks. The model weights for both sizes are publicly available under the Jamba Open Model License and we release ExpertsInt8 as open source.

연구 동기 및 목표

• 장기 컨텍스트 입력을 처리할 수 있는 대규모, 효율적이고 고성능의 지시 퍼티드 언어 모델 개발.
• 품질 손실 없이 대규모 모델의 메모리 사용량을 줄이고 추론 효율성을 향상시키기.
• 신규 양자화 및 아키텍처 혁신을 통해 Jamba-1.5-Large와 같은 대규모 모델을 표준 하드웨어(8x80GB GPU)에 배포할 수 있도록 하기.
• 하이브리드 트랜스포머-앰바-MoE 아키텍처가 순수 트랜스포머 또는 앰바 모델보다 스케일링 시 더 우수한 성능을 보임을 입증하기, 특히 장기 컨텍스트 환경에서.
• OECD AI 원칙에 부합하고 투명성을 확보한 Jamba Open Model License 하에 오픈 웨이트 모델 공개하기.

제안 방법

• Jamba-1.5 아키텍처는 트랜스포머 레이어, 앰바 상태공간 레이어, 16개의 전문가와 상위 2개 게이팅을 갖춘 미스처 오브 익스퍼트(MoE) 모듈을 통합한다.
• 각 블록은 주로 1:7 비율의 어텐션 레이어 대 앰바 레이어를 사용하며, MoE는 매 2개 레이어마다 적용되며, 히든 차원은 8192이다.
• ExpertsInt8 양자화 기술이 도입되었으며, MoE 레이어의 85%에 해당하는 가중치를 대상으로 하여 품질 저하 없이 8비트 양자화를 구현한다.
• 모델은 1000억 토큰에 걸쳐 훈련되었으며, 다양한 벤치마크에서 지시 따르기 및 대화 기능 향상을 위한 피니어팅이 수행되었다.
• 앰바의 효율적인 시퀀스 모델링과 MoE의 희소 활성화를 활용함으로써, 높은 처리량과 낮은 KV 캐시 사용량(256K 토큰 기준 9GB)을 달성한다.
• 평가에서는 학술 벤치마크(MMLU, GSM8K 등), 챗봇 벤치마크, RULER과 같은 장기 컨텍스트 작업을 수행하였으며, LLaMA-3, Mistral-Large-2, Mixtral 등과의 비교 분석을 포함한다.

(a) 350M.

실험 결과

연구 질문

• RQ1 하이브리드 트랜스포머-앰바-MoE 아키텍처가 활성 파라미터 수 94B에 도달하면서도 모델 품질을 유지하거나 향상시킬 수 있는가?
• RQ2 스케일링된 환경에서 앰바와 트랜스포머 레이어의 조합이 순수 앰바-2 또는 트랜스포머 전용 모델보다 우수한 성능을 내는가?
• RQ3 ExpertsInt8 양자화가 MoE 레이어에 대해 전체 8비트 양자화를 가능하게 하며, 품질 저하 없이 8x80GB GPU에서의 배포를 가능하게 하는가?
• RQ4 다른 오픈 웨이트 모델과 비교해, 특히 256K 토큰 컨텍스트 길이에서 장기 컨텍스트 작업에서의 성능은 어떠한가?
• RQ5 Jamba-1.5 모델은 다국어 작업, 특히 자원이 부족한 언어에서 얼마나 일반화되는가?

주요 결과

• Jamba-1.5-Large는 94B 활성 파라미터를 확보했으며, ExpertsInt8 양자화를 통해 8x80GB GPU에서 단일 머신에서 256K 컨텍스트 길이로 배포되었고, 품질 손실 없이 구현되었다.
• 256K 토큰 기준으로 KV 캐시 메모리 사용량을 9GB로 줄여 LLaMA-3.1-70B 및 Mistral-Large-2 대비 10배 감소시켰다.
• RULER 벤치마크에서 Jamba-1.5 모델은 유일하게 효과적 컨텍스트 길이 256K 토큰을 지원하는 오픈 웨이트 모델이었다.
• Jamba-1.5-Mini는 다국어 MMLU에서 64.3%를 기록하여 LLaMA-3.1-8B와 Gemma-9B를 초월했으며, Jamba-1.5-Large는 동일한 벤치마크에서 73.94%를 달성했다.
• 장기 컨텍스트 평가에서 Jamba-1.5-Large는 LLaMA-3.1-70B 및 Mistral-Large-2를 압도하며 뛰어난 처리량과 지연 시간 성능을 보였다.
• ExpertsInt8 양자화가 MoE 레이어에 효과적으로 작용하며, 모델 품질을 유지하면서 소비자 수준 하드웨어에서의 효율적 배포를 가능하게 했다.

(b) 1.3B.