mamba 다운로드 - mamba 소스 코드 다운로드

mamba

기타 소스코드

v2.2.2

다운로드

맘바

Mamba

Mamba: 선택적 상태 공간을 사용한 선형-시간 시퀀스 모델링
Albert Gu*, Tri Dao*
논문: https://arxiv.org/abs/2312.00752

Mamba-2

변환기는 SSM입니다: 일반화 모델 및 효율적인 알고리즘
구조화된 상태 공간 이중성을 통해
트라이 다오*, Albert Gu*
논문: https://arxiv.org/abs/2405.21060

에 대한

Mamba는 이전의 하위 2차 모델이 Transformers에 비해 부족한 언어 모델링과 같은 정보 밀도가 높은 데이터에 대해 유망한 성능을 보여주는 새로운 상태 공간 모델 아키텍처입니다. 이는 FlashAttention의 정신으로 효율적인 하드웨어 인식 설계 및 구현을 통해 구조화된 상태 공간 모델의 발전 방향을 기반으로 합니다.

설치

[옵션] pip install causal-conv1d>=1.4.0 : Mamba 블록 내부에서 사용되는 간단한 인과 Conv1d 레이어를 효율적으로 구현합니다.
pip install mamba-ssm : 핵심 Mamba 패키지입니다.
pip install mamba-ssm[causal-conv1d] : 핵심 Mamba 패키지와 causal-conv1d를 설치합니다.
pip install mamba-ssm[dev] : 핵심 Mamba 패키지 및 개발 종속성을 설치합니다.

pip install . 이 저장소에서.

pip PyTorch 버전에 대해 불만을 표시하는 경우 --no-build-isolation pip 에 전달해 보세요.

기타 요구 사항:

리눅스
엔비디아 GPU
파이토치 1.12+
쿠다 11.6+

AMD 카드의 경우 아래의 추가 전제 조건을 참조하세요.

용법

우리는 Mamba 모델과 함께 여러 수준의 인터페이스를 공개합니다.

선택적 SSM

Mamba는 논문의 초점인 선택적 SSM 계층을 기반으로 합니다(섹션 3, 알고리즘 2).

출처: ops/selective_scan_interface.py.

맘바 블록

이 저장소의 기본 모듈은 선택적 SSM을 래핑하는 Mamba 아키텍처 블록입니다.

출처: 모듈/mamba_simple.py.

용법:

 import torch
from mamba_ssm import Mamba

batch , length , dim = 2 , 64 , 16
x = torch . randn ( batch , length , dim ). to ( "cuda" )
model = Mamba (
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model = dim , # Model dimension d_model
    d_state = 16 ,  # SSM state expansion factor
    d_conv = 4 ,    # Local convolution width
    expand = 2 ,    # Block expansion factor
). to ( "cuda" )
y = model ( x )
assert y . shape == x . shape

맘바-2

Mamba-2 블록은 module/mamba2.py에서 구현됩니다.

더 간단한 버전은 module/mamba2_simple.py에 있습니다.

사용법은 Mamba(-1)와 유사합니다.

 from mamba_ssm import Mamba2
model = Mamba2 (
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model = dim , # Model dimension d_model
    d_state = 64 ,  # SSM state expansion factor, typically 64 or 128
    d_conv = 4 ,    # Local convolution width
    expand = 2 ,    # Block expansion factor
). to ( "cuda" )
y = model ( x )
assert y . shape == x . shape

SSD

"이산" 및 "연속" SSM 버전 간 변환이 포함된 내부 SSD 모듈(Mamba-2 문서의 목록 1)의 최소 버전은 module/ssd_minimal.py에 있습니다.

맘바 언어 모델

마지막으로 완전한 언어 모델의 예를 제공합니다. 즉, 심층 시퀀스 모델 백본(반복하는 Mamba 블록 포함) + 언어 모델 헤드입니다.

출처: models/mixer_seq_simple.py.

이는 Mamba를 엔드투엔드 신경망에 통합하는 방법의 예입니다. 이 예는 아래 생성 스크립트에서 사용됩니다.

사전 훈련된 모델

사전 학습된 모델은 Hugging Face에 업로드됩니다: 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 토큰으로 훈련됨) 및 mamba-2.8b-slimpj (SlimPajama 데이터 세트에서 600B 토큰으로 훈련됨).

모델은 아래 생성 스크립트에 의해 자동 다운로드됩니다.

이러한 모델은 Pile에서 훈련되었으며 GPT-3에 설명된 표준 모델 차원과 많은 오픈 소스 모델을 따릅니다.

매개변수	레이어	모델이 어두워요.
1억 3천만	24	768
370M	48	1024
790M	48	1536년
1.4B	48	2048년
2.8B	64	2560

(Mamba의 레이어 수는 비슷한 크기의 Transformer의 레이어 수의 두 배입니다. 이는 Transformer의 각 "레이어"(MHA 블록 + MLP 블록)에 두 개의 Mamba 블록이 필요하기 때문입니다.)

참고: 이는 어떠한 형태의 다운스트림 수정(명령 튜닝 등) 없이 300B 토큰에 대해서만 훈련된 기본 모델입니다. 성능은 유사한 데이터로 훈련된 다른 아키텍처와 비슷하거나 더 좋을 것으로 예상되지만 더 크거나 미세 조정된 모델과 일치하지는 않습니다.

평가

모델의 제로샷 평가(논문의 표 3에 해당)를 실행하기 위해 lm-evaluation-harness 라이브러리를 사용합니다.

pip install lm-eval==0.4.2 로 lm-evaluation-harness 설치합니다.
다음을 사용하여 평가를 실행합니다(lm-evaluation-harness 저장소에서 추가 문서 참조).

lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-130m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
python evals/lm_harness_eval.py --model hf --model_args pretrained=EleutherAI/pythia-160m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande --device cuda --batch_size 64

블로그 게시물에 보고된 mamba-2.8b-slimpj 모델의 결과를 재현하려면:

lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks boolq,piqa,hellaswag,winogrande,arc_easy,arc_challenge,openbookqa,race,truthfulqa_mc2 --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks mmlu --num_fewshot 5 --device cuda --batch_size 256

Mamba-2 모델에 대한 평가를 실행하려면 모델 이름을 바꾸면 됩니다.

lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/transformerpp-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2attn-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256

각 작업의 결과는 평가 과정의 노이즈로 인해 보고된 값과 0.1~0.3 정도 다를 수 있습니다.

추론

스크립트 벤치마킹/benchmark_세대_mamba_simple.py

Hugging Face Hub에서 모델을 자동 로드합니다.
사용자가 지정한 프롬프트의 완성을 생성합니다.
이 세대의 추론 속도를 벤치마킹합니다.

기타 구성 가능한 옵션에는 top-p(핵 샘플링) 확률 및 소프트맥스 온도가 포함됩니다.

예

다양한 샘플링 전략을 사용하여 생성 지연 시간(예: 배치 크기 = 1)을 테스트하려면 다음을 수행하세요.

python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba-2.8b " --prompt " My cat wrote all this CUDA code for a new language model and " --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name " EleutherAI/pythia-2.8b " --prompt " My cat wrote all this CUDA code for a new language model and " --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba-2.8b " --prompt " My cat wrote all this CUDA code for a new language model and " --minp 0.05 --topk 0 --temperature 0.7 --repetition-penalty 1.2

무작위 프롬프트로 생성 처리량을 테스트하려면(예: 대규모 배치 크기):

python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba-2.8b " --batch 64
python benchmarks/benchmark_generation_mamba_simple.py --model-name " EleutherAI/pythia-2.8b " --batch 64

Mamba-2를 사용하면 모델 이름만 변경하면 됩니다.

python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba2-2.7b " --prompt " My cat wrote all this CUDA code for a new language model and " --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

문제 해결

정도

우리 모델은 혼합 정밀도를 위해 PyTorch AMP를 사용하여 훈련되었습니다. AMP는 모델 매개변수를 float32에 유지하고 필요한 경우 절반 정밀도로 캐스팅합니다. 반면 DeepSpeed와 같은 다른 프레임워크는 float16에 매개변수를 저장하고 필요한 경우(예: 최적화 프로그램 축적을 위해) 업캐스트합니다.

우리는 SSM이 반복되는 역학에 민감하기 때문에 주요 모델 매개변수에 대해 더 높은 정밀도가 필요할 수 있음을 관찰했습니다. 불안정한 경우 첫 번째 단계로 fp32(예: AMP)에 매개변수를 저장하는 프레임워크를 사용해 보세요.

초기화

모델의 일부 부분에는 S4 모델에 대한 이전 작업에서 상속된 초기화가 있습니다. 예를 들어, $델타$ 매개변수는 선형 투영의 편향을 초기화하여 목표 범위를 갖습니다. 그러나 일부 프레임워크에는 초기화 후 후크가 있을 수 있습니다(예: nn.Linear 모듈의 모든 바이어스 항을 0으로 설정). 이 경우 훈련 프레임워크에 특정한 사용자 지정 논리(예: 이 줄은 트레이너에서 다시 초기화를 끄지만 다른 프레임워크에서는 작동하지 않음)를 추가해야 할 수 있습니다.

AMD 카드에 대한 추가 전제 조건

ROCm 패치 중

ROCm 6.0을 사용하는 경우 다음 단계를 실행하여 컴파일 중 오류를 방지하세요. ROCm 6.1 이상에서는 이는 필요하지 않습니다.

ROCm 설치 디렉터리를 찾으세요. 이는 일반적으로 /opt/rocm/ 에 있지만 설치에 따라 다를 수 있습니다.
패치를 적용합니다. 권한 문제가 발생할 경우 sudo 사용하여 실행하세요.
```
 patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h < rocm_patch/rocm6_0.patch 
```

소환

이 코드베이스를 사용하거나 우리 작업이 가치 있다고 생각하는 경우 Mamba를 인용해 주세요.

 @article{mamba,
  title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
  author={Gu, Albert and Dao, Tri},
  journal={arXiv preprint arXiv:2312.00752},
  year={2023}
}

@inproceedings{mamba2,
  title={Transformers are {SSM}s: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
  author={Dao, Tri and Gu, Albert},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2024}
}

확장하다

추가 정보