# Jamba: AI21이 공개한 Mamba 기반의 MoE 공개 모델 (OpenLLM)

**URL:** <https://discuss.pytorch.kr/t/jamba-ai21-mamba-moe-openllm/3888>\
**Category:** 읽을거리&정보공유\
**Tags:** open-weights, mamba, mixture-of-experts, jamba, mamba-ssm, ai21labs\
**Created:** [3월 29, 2024, 11:00오전 UTC](https://discuss.pytorch.kr/t/jamba-ai21-mamba-moe-openllm/3888 "2024-03-29T11:00:15Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![9bow](https://discuss.pytorch.kr/user_avatar/discuss.pytorch.kr/9bow/32/16301_2.png) [@9bow](https://discuss.pytorch.kr/u/9bow)\
**Post date:** [3월 29, 2024, 11:00오전 UTC](https://discuss.pytorch.kr/t/jamba-ai21-mamba-moe-openllm/3888/1 "2024-03-29T11:00:16Z")

</div>

![Jamba: AI21이 공개한 Mamba 기반의 MoE 공개 모델 (OpenLLM)](https://discuss.pytorch.kr/uploads/default/original/2X/0/0ca9a07d229c3d3c25ed38f666f127f9899da59e.jpeg)

AI21이 새로 공개한 Jamba는 기존의 트랜스포머 아키텍처의 한계를 극복하기 위해 설계된 새로운 구조의 SSM(Structured State Space) 모델, Mamba를 기반으로 합니다. Jamba는 높은 성능과 효율성을 자랑하며, 특히 다른 기존 모델들 대비 긴 Context 크기를 가지면서도 3배 가까운 처리량 증가를 보이고 있습니다.

* * *

# Jamba: AI21이 공개한 Mamba 기반의 MoE 공개 모델 (OpenLLM)

## 소개

 ![Jamba 소개: Transformer + Mamba](https://discuss.pytorch.kr/uploads/default/original/2X/a/a4a0d883e9a76e6ba43905584efe8dc9967875f2.png)

Jamba는 Mamba, 즉 새로운 구조의 SSM 아키텍처를 기반으로 한 첫 번째 생산 등급 모델입니다. 이 모델은 트랜스포머 아키텍처의 한계를 극복하려는 시도에서 출발했으나, Mamba만의 한계도 있었습니다. Jamba는 트랜스포머와 SSM의 장점을 모두 살려, 기존 모델들과 비교해 뛰어난 성능을 보여주는 동시에, 긴 문맥에서의 처리량을 3배 가까이 향상시키는 등 비용 효율성과 접근성을 우선시하는 모델입니다.

## 주요 특징

- 새로운 SSM-트랜스포머 하이브리드 아키텍처를 기반으로 구축된 최초의 프로덕션 수준의 맘바 기반 모델

- Mixtral 8x7B 대비 긴 컨텍스트에서 3배의 처리량 제공

- 대규모 256K 컨텍스트 창에 대한 액세스의 대중화

- 단일 GPU에서 최대 140K 컨텍스트까지 지원하는 동급 유일 모델

- Apache 2.0 라이선스로 모델 가중치를 공개한 공개 LLM (OpenLLM)

- [Hugging Face에서 사용 가능](https://huggingface.co/ai21labs/Jamba-v0.1)하며, [NVIDIA API 카탈로그](http://ai.nvidia.com/)에도 곧 추가 예정

### 하이브리드 아키텍처의 뛰어난 성능

 ![Jamba의 성능 비교](https://discuss.pytorch.kr/uploads/default/original/2X/8/8eb9945e585effaa714afa4a4fecde052f5be80a.jpeg)

Jamba는 비슷한 크기의 다른 모델들과 비교했을 때, 더 뛰어나거나 비슷한 성능을 보입니다. 추론(reasoning) 관련 벤치마크들에서 좋은 결과를 보이고 있습니다.

 ![Jamba의 모델 구조](https://discuss.pytorch.kr/uploads/default/original/2X/a/a962ad163d8c1ebbb33f584cceda005a1e8e054e.png)

Jamba는 SSM과 트랜스포머 아키텍처의 장점을 결합한 세계 최초의 프로덕션급의 Mamba 기반 모델입니다. 이 하이브리드 구조는 트랜스포머의 강력한 언어 이해 능력과 SSM의 효율적인 메모리 관리 및 처리 속도를 모두 활용합니다. 결과적으로, Jamba는 기존의 언어 모델들이 가지고 있던 메모리 사용량 증가와 처리 속도 저하 문제를 크게 개선합니다.

### 대규모 컨텍스트 윈도우 및 비용-효율성

 ![Jamba와 다른 모델들 간의 SingleGPU에서의 최대 Context 크기 비교](https://discuss.pytorch.kr/uploads/default/original/2X/e/e0cf769767327c9649d51419303bda5085a04058.jpeg)

Jamba는 256K의 컨텍스트 윈도우를 제공함으로써, 매우 긴 문서나 대화를 효율적으로 처리할 수 있습니다. 이는 AI 개발자들이 보다 복잡한 자연어 처리 작업을 수행할 수 있게 해주며, 기존 모델들이 다룰 수 없었던 긴 문맥의 이해에 기여합니다.

또한, 긴 컨텍스트에서 기존의 트랜스포머 기반 모델들과 비교해 약 3배 높은 처리량을 달성합니다. 이는 Jamba가 단일 GPU에서도 140K 컨텍스트를 처리할 수 있음을 의미하며, 이는 비용 효율적이며 접근하기 쉬운 배포를 가능하게 합니다.

### MoE(Mixture of Experts) 활용

 ![Jamba의 처리량: 기존 모델 대비 3배 가량의 높은 처리량](https://discuss.pytorch.kr/uploads/default/original/2X/9/9006a0525c241278cc31d7bc1d2586c52d929800.jpeg)

Jamba는 MoE 레이어를 통해 추론 시 사용 가능한 52B 파라미터 중 단 12B만을 활용합니다. 이로 인해 모델의 활성화 파라미터가 더 효율적으로 사용되며, 동일한 크기의 트랜스포머 전용 모델보다 더 나은 성능을 보입니다.

## 사용 방법

Jamba 모델을 사용하기 위해서는 Hugging Face의 `transformers` 라이브러리가 필요합니다. 다음은 Python을 사용하여 Jamba 모델을 불러오고, 간단한 텍스트 생성을 수행하는 예시 코드입니다. 이 코드를 실행하기 전에, `transformers` 라이브러리와 함께 `mamba-ssm`, `causal-conv1d` 라이브러리를 설치해야 합니다. 이는 Jamba의 최적화된 Mamba 구현을 사용하기 위함입니다.

### 필요한 라이브러리 설치

먼저, 필요한 라이브러리를 설치합니다:

```bash
pip install transformers>=4.39.0
pip install mamba-ssm causal-conv1d>=1.2.0

```

_⚠주의: Hugging Face의 예제 코드에는 transformers 4.39.0 이상을 설치하라고 하지만, 아직 `Jamba-v0` 모듈이 추가되지 않은 것 같습니다. 에러가 발생하면 아래와 같이 직접 transformers GitHub을 복제하여 4.40.0dev 이상의 버전을 설치하시면 됩니다._

```bash
git clone https://github.com/huggingface/transformers
cd transformers
pip install -e .

```

### 모델과 토크나이저 불러오기

`transformers` 라이브러리를 사용하여 Jamba 모델과 토크나이저를 불러옵니다. 그리고 간단한 입력 텍스트에 대한 텍스트 생성을 수행합니다:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

# 모델과 토크나이저 초기화
model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")

# 생성할 텍스트의 시작 부분
input_text = "PyTorch Korea User Group is "

# 토크나이저를 사용하여 입력 텍스트를 토큰화
input_ids = tokenizer(input_text, return_tensors='pt').input_ids

# 모델을 사용하여 텍스트 생성
outputs = model.generate(input_ids, max_length=50, num_return_sequences=1)

# 생성된 텍스트 출력
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

```

위 코드는 `"PyTorch Korea User Group is "` 이라는 입력 텍스트에 대해 모델이 텍스트를 생성하도록 합니다. `max_length`는 생성할 텍스트의 최대 길이를 지정하고, `num_return_sequences`는 생성할 텍스트 시퀀스의 수를 지정합니다. 생성된 텍스트는 `tokenizer.decode`를 통해 디코딩하여 출력합니다.

## 더 읽어보기

### Jmaba 홈페이지

> **[Jamba LLMs: The Best Long Context Models for Secure Enterprise Deployment | AI21](https://www.ai21.com/jamba/)**
>
> Explore AI21's Jamba models - a family of long-context, hyper-efficient open LLMs built for the enterprise and available for secure self-hosted deployments, such as on-premise and VPC.

### Jmaba 공개 글

> **[Introducing Jamba: AI21's Groundbreaking SSM-Transformer Model](https://www.ai21.com/blog/announcing-jamba/)**
>
> Debuting the first production-grade Mamba-based model delivering best-in-class quality and performance.

### 모델 가중치

> **[ai21labs/Jamba-v0.1 · Hugging Face](https://huggingface.co/ai21labs/Jamba-v0.1)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

* * *

_이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있으니 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다._ 🤗

[:pytorch:파이토치 한국 사용자 모임🇰🇷](https://pytorch.kr/)이 정리한 이 글이 유용하셨나요? [회원으로 가입](https://discuss.pytorch.kr/signup)하시면 주요 글들을 이메일💌로 보내드립니다! (기본은 Weekly지만 [Daily로 변경도 가능](https://discuss.pytorch.kr/my/preferences/emails)합니다.)

🎁 아래↘쪽에 좋아요❤를 눌러주시면 힘이 됩니다~ 🤩