AI21이 새로 공개한 Jamba는 기존의 트랜스포머 아키텍처의 한계를 극복하기 위해 설계된 새로운 구조의 SSM(Structured State Space) 모델, Mamba를 기반으로 합니다. Jamba는 높은 성능과 효율성을 자랑하며, 특히 다른 기존 모델들 대비 긴 Context 크기를 가지면서도 3배 가까운 처리량 증가를 보이고 있습니다.
Jamba: AI21이 공개한 Mamba 기반의 MoE 공개 모델 (OpenLLM)
소개
Jamba는 Mamba, 즉 새로운 구조의 SSM 아키텍처를 기반으로 한 첫 번째 생산 등급 모델입니다. 이 모델은 트랜스포머 아키텍처의 한계를 극복하려는 시도에서 출발했으나, Mamba만의 한계도 있었습니다. Jamba는 트랜스포머와 SSM의 장점을 모두 살려, 기존 모델들과 비교해 뛰어난 성능을 보여주는 동시에, 긴 문맥에서의 처리량을 3배 가까이 향상시키는 등 비용 효율성과 접근성을 우선시하는 모델입니다.
주요 특징
-
새로운 SSM-트랜스포머 하이브리드 아키텍처를 기반으로 구축된 최초의 프로덕션 수준의 맘바 기반 모델
-
Mixtral 8x7B 대비 긴 컨텍스트에서 3배의 처리량 제공
-
대규모 256K 컨텍스트 창에 대한 액세스의 대중화
-
단일 GPU에서 최대 140K 컨텍스트까지 지원하는 동급 유일 모델
-
Apache 2.0 라이선스로 모델 가중치를 공개한 공개 LLM (OpenLLM)
-
Hugging Face에서 사용 가능하며, NVIDIA API 카탈로그에도 곧 추가 예정
하이브리드 아키텍처의 뛰어난 성능
Jamba는 비슷한 크기의 다른 모델들과 비교했을 때, 더 뛰어나거나 비슷한 성능을 보입니다. 추론(reasoning) 관련 벤치마크들에서 좋은 결과를 보이고 있습니다.
Jamba는 SSM과 트랜스포머 아키텍처의 장점을 결합한 세계 최초의 프로덕션급의 Mamba 기반 모델입니다. 이 하이브리드 구조는 트랜스포머의 강력한 언어 이해 능력과 SSM의 효율적인 메모리 관리 및 처리 속도를 모두 활용합니다. 결과적으로, Jamba는 기존의 언어 모델들이 가지고 있던 메모리 사용량 증가와 처리 속도 저하 문제를 크게 개선합니다.
대규모 컨텍스트 윈도우 및 비용-효율성
Jamba는 256K의 컨텍스트 윈도우를 제공함으로써, 매우 긴 문서나 대화를 효율적으로 처리할 수 있습니다. 이는 AI 개발자들이 보다 복잡한 자연어 처리 작업을 수행할 수 있게 해주며, 기존 모델들이 다룰 수 없었던 긴 문맥의 이해에 기여합니다.
또한, 긴 컨텍스트에서 기존의 트랜스포머 기반 모델들과 비교해 약 3배 높은 처리량을 달성합니다. 이는 Jamba가 단일 GPU에서도 140K 컨텍스트를 처리할 수 있음을 의미하며, 이는 비용 효율적이며 접근하기 쉬운 배포를 가능하게 합니다.
MoE(Mixture of Experts) 활용
Jamba는 MoE 레이어를 통해 추론 시 사용 가능한 52B 파라미터 중 단 12B만을 활용합니다. 이로 인해 모델의 활성화 파라미터가 더 효율적으로 사용되며, 동일한 크기의 트랜스포머 전용 모델보다 더 나은 성능을 보입니다.
사용 방법
Jamba 모델을 사용하기 위해서는 Hugging Face의 transformers 라이브러리가 필요합니다. 다음은 Python을 사용하여 Jamba 모델을 불러오고, 간단한 텍스트 생성을 수행하는 예시 코드입니다. 이 코드를 실행하기 전에, transformers 라이브러리와 함께 mamba-ssm, causal-conv1d 라이브러리를 설치해야 합니다. 이는 Jamba의 최적화된 Mamba 구현을 사용하기 위함입니다.
필요한 라이브러리 설치
먼저, 필요한 라이브러리를 설치합니다:
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 이상의 버전을 설치하시면 됩니다.
git clone https://github.com/huggingface/transformers
cd transformers
pip install -e .
모델과 토크나이저 불러오기
transformers 라이브러리를 사용하여 Jamba 모델과 토크나이저를 불러옵니다. 그리고 간단한 입력 텍스트에 대한 텍스트 생성을 수행합니다:
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 홈페이지
Jmaba 공개 글
모델 가중치
이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있으니 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. 
파이토치 한국 사용자 모임
이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일
로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)
아래
쪽에 좋아요
를 눌러주시면 힘이 됩니다~ 