Jamba 아키텍처 — blocks-and-layers 하이브리드

Jamba 아키텍처 — blocks-and-layers 하이브리드

Jamba가 왜 Transformer와 Mamba를 한 모델에 담을 수 있는지, 그 구조부터 보는 게 이해가 빨라요. 핵심은 blocks-and-layers 방식을 쓴다는 점입니다.

출처: https://huggingface.co/docs/transformers/model_doc/jamba

동작 원리

각 Jamba 블록에는 어텐션 레이어 또는 Mamba 레이어 중 하나가 들어가고, 그 뒤에 다층 퍼셉트론(MLP)이 이어져요. 이 조합을 반복하면서 전체 8개 레이어 중 1개만 Transformer 레이어가 되도록 비율을 유지합니다. MoE 레이어를 섞어 모델 용량을 키우는 것도 Jamba의 특징이에요.

HF Transformers에서 쓰기

Hugging Face Transformers에서 바로 사용할 수 있어요. Mamba 커널을 쓰려면 관련 패키지를 먼저 설치합니다.

pip install mamba-ssm causal-conv1d>=1.2.0

파이프라인으로 간단히 텍스트 생성을 돌려볼 수 있어요.

from transformers import pipeline

pipe = pipeline(
    task="text-generation",
    model="ai21labs/AI21-Jamba-Mini-1.6",
    device=0,
)
pipe("Plants create energy through a process known as")

큰 모델 로드 시 주의

AI21-Jamba-Large-1.6처럼 큰 모델은 여러 GPU에 나눠 담는 게 일반적이에요.

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/AI21-Jamba-Large-1.6",
    device_map="auto",
    attn_implementation="sdpa",
)
tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Large-1.6")

attn_implementation="sdpa"로 메모리 효율을 높이고, 더 빠르게 쓰려면 flash_attention_2를 고를 수도 있어요.

더 알아보기

  • https://huggingface.co/ai21labs — 공식 체크포인트 저장소
  • https://huggingface.co/docs/transformers/main/en/model_doc/jamba — 모델 문서