mamba

Mamba는 state-spaces 그룹이 개발한 오픈소스 연구 코드베이스로, 시퀀스 모델링을 위한 일련의 상태 공간 모델(SSM) 아키텍처를 구현합니다. 이는 "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", "Transformers are SSMs"(Mamba-2), "Mamba-3: Improved Sequence Modeling using State Space Principles" 논문과 함께 제공됩니다. 이 프로젝트는 특히 언어 모델링과 같은 정보 밀도가 높은 데이터에서 트랜스포머에 대한 하위 이차(subquadratic) 대안으로 자리매김하며, FlashAttention의 정신을 따른 하드웨어 인지 구현을 통해 이전의 구조화된 상태 공간 연구(S4)를 기반으로 구축되었습니다. 저장소가 제공하는 것: - 선택적 SSM 레이어: 핵심 선택적 스캔 연산으로, 소스는 ops/selective_scan_interface.py에 있습니다. - Mamba 블록: 선택적 SSM을 감싸는 주요 아키텍처 모듈(modules/mamba_simple.py)로, d_model, d_state, d_conv, expand 매개변수를 구성할 수 있는 바로 사용 가능한 PyTorch 모듈입니다. - Mamba-2 블록: modules/mamba2.py에 구현되어 있으며 modules/mamba2_simple.py에 더 간단한 변형이 있고, modules/ssd_minimal.py에 최소한의 SSD(구조화된 상태 공간 이중성) 참조 모듈이 있습니다. - Mamba-3 블록: modules/mamba3.py에 구현되어 있으며, d_state, headdim, mimo_rank를 사용한 MIMO 모드, chunk_size, 선택적 출력 프로젝션 정규화, bfloat16 dtype 등의 옵션을 제공합니다. - 언어 모델 예제: models/mixer_seq_simple.py에 전체 백본과 LM 헤드가 있으며, 생성 스크립트에서 사용됩니다. 설치 및 빌드 옵션: 패키지는 pip(mamba-ssm)로 설치되며 Linux, Python 3.10+ 및 PyTorch 1.12+가 필요합니다. 기본적으로 설치 시 selective_scan_cuda 확장을 컴파일하지 않고 캐시된 CUDA 휠을 가져오지 않습니다. 선택적 추가 기능과 환경 플래그가 이 동작을 변경합니다: causal-conv1d 추가 기능은 해당 의존성을 추가하고, MAMBA_FORCE_BUILD=TRUE는 로컬 빌드를 강제하며, MAMBA_KEEP_CUDA_BUILD=TRUE는 CUDA 선택적 스캔 확장을 선택하여 먼저 일치하는 사전 빌드된 CUDA/HIP 휠을 시도하고, 없으면 로컬에서 컴파일합니다. 두 플래그를 결합하면 로컬 CUDA 컴파일이 강제됩니다. README는 CUDA 빌드에 --no-build-isolation이 필요하여 pip가 기존 CUDA 지원 PyTorch를 사용하도록 한다고 설명합니다. CUDA 빌드 및 GPU 실행에는 NVIDIA GPU와 CUDA 11.6+가 추가로 필요합니다. AMD 카드의 경우 README는 ROCm 6.0 패치 단계를 문서화합니다(ROCm 6.1부터는 필요 없음). 사전 학습 모델: 가중치는 Hugging Face의 state-spaces 조직 아래 게시되어 있으며, mamba-130m, mamba-370m, mamba-790m, mamba-1.4b, mamba-2.8b, mamba2-130m, mamba2-370m, mamba2-780m, mamba2-1.3b, mamba2-2.7b, transformerpp-2.7b, mamba2attn-2.7b가 포함되고, 이들은 Pile의 300B 토큰으로 학습되었습니다. 또한 SlimPajama의 600B 토큰으로 학습된 mamba-2.8b-slimpj도 있습니다. README는 이들을 명령어 튜닝이나 기타 다운스트림 수정이 없는 기본 모델로 설명하며, 성능이 유사한 데이터로 학습된 다른 아키텍처와 비슷하거나 더 나을 것으로 기대되지만 더 크거나 미세 조정된 모델과는 일치하지 않을 것이라고 명시합니다. 모델 차원은 GPT-3 스타일 확장을 따릅니다(예: 24개 레이어와 768 모델 차원의 130M 매개변수부터 64개 레이어와 2560 차원의 2.8B까지). 평가 및 추론 도구: 제로샷 평가는 lm-evaluation-harness 라이브러리(pip install lm-eval==0.4.2)를 통해 수행되며, lambada_openai, hellaswag, piqa, arc_easy, arc_challenge, winogrande, openbookqa, boolq, race, truthfulqa_mc2, mmlu와 같은 작업을 포함하는 예제 명령이 제공됩니다. README는 평가 노이즈로 인해 작업별 결과가 보고된 값과 0.1~0.3 차이가 날 수 있다고 설명합니다. 생성 벤치마크 스크립트(benchmarks/benchmark_generation_mamba_simple.py)는 Hugging Face Hub에서 모델을 자동 로드하고, 사용자 프롬프트에 대한 완성 텍스트를 생성하며, top-p, min-p, top-k, temperature, 반복 페널티, 배치 크기 옵션으로 추론 속도를 측정합니다. 문제 해결 노트: README는 두 가지 실용적인 문제를 강조합니다. 첫째, 정밀도: 모델은 매개변수를 float32로 유지하고 필요에 따라 절반 정밀도로 캐스팅하는 PyTorch AMP로 학습되었습니다. SSM은 순환 역학에 민감하므로 float16으로 매개변수를 저장하는 프레임워크(예: DeepSpeed)는 불안정성을 유발할 수 있으며, fp32 매개변수 저장이 첫 번째 해결책으로 제안됩니다. 둘째, 초기화: 일부 구성 요소는 S4 스타일 초기화를 상속합니다(예: 선형 투영 편향을 통한 델타 매개변수의 목표 범위). nn.Linear 편향을 0으로 만드는 초기화 후 훅이 있는 프레임워크는 이를 보존하기 위해 사용자 지정 논리가 필요할 수 있습니다. 라이선스 및 인용: README는 Mamba, Mamba-2, Mamba-3 논문에 대한 BibTeX 항목을 제공하며, 코드베이스를 사용할 때 이를 인용하도록 요청합니다.