aakashlahoti commited on
Commit
bc6b5d0
·
0 Parent(s):

Mamba-3 public release

Browse files
Files changed (4) hide show
  1. .gitattributes +5 -0
  2. README.md +105 -0
  3. config.json +29 -0
  4. pytorch_model.bin +3 -0
.gitattributes ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ *.bin filter=lfs diff=lfs merge=lfs -text
2
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
3
+ *.pt filter=lfs diff=lfs merge=lfs -text
4
+ *.pth filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ library_name: mamba-ssm
4
+ tags:
5
+ - mamba
6
+ - mamba-3
7
+ - mimo
8
+ - state-space-model
9
+ - causal-lm
10
+ - language-modeling
11
+ ---
12
+
13
+ # Mamba-3 MIMO 1.50 B
14
+
15
+ `mamba3-mimo-1.5b` is a pretrained causal language model with 1.50 B
16
+ parameters. It is built from stacked blocks, each containing a Mamba-3
17
+ MIMO mixer followed by a gated MLP. It contains no attention layers and
18
+ is released with BF16 weights in the public `mamba_ssm` checkpoint format.
19
+
20
+ ## Model architecture
21
+
22
+ | Property | Value |
23
+ |---|---:|
24
+ | Parameters | 1.50 B |
25
+ | Layers | 24 |
26
+ | Model dimension | 2,048 |
27
+ | SSM state size | 128 |
28
+ | SSM head dimension | 64 |
29
+ | SSM heads | 64 |
30
+ | SSM groups | 1 |
31
+ | MIMO rank | 4 |
32
+ | Chunk size | 16 |
33
+ | Context length | 2,048 |
34
+
35
+ The model was pretrained on 100B tokens from
36
+ [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu).
37
+ It uses tied input and output embeddings and the `meta-llama/Llama-3.1-8B` tokenizer.
38
+
39
+ ## Installation
40
+
41
+ Install CUDA-enabled PyTorch first, followed by the latest Mamba source:
42
+
43
+ ```bash
44
+ pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation
45
+ ```
46
+
47
+ While this repository is private, authenticate with Hugging Face:
48
+
49
+ ```bash
50
+ hf auth login
51
+ ```
52
+
53
+ ## Usage
54
+
55
+ ```python
56
+ import torch
57
+ from transformers import AutoTokenizer
58
+ from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
59
+
60
+ model_id = "state-spaces/mamba3-mimo-1.5b"
61
+
62
+ tokenizer = AutoTokenizer.from_pretrained(
63
+ "meta-llama/Llama-3.1-8B",
64
+ )
65
+
66
+ model = MambaLMHeadModel.from_pretrained(
67
+ model_id,
68
+ device="cuda",
69
+ dtype=torch.bfloat16,
70
+ )
71
+ model.eval()
72
+
73
+ input_ids = tokenizer(
74
+ "Mamba-3 is",
75
+ return_tensors="pt",
76
+ ).input_ids.cuda()
77
+
78
+ with torch.inference_mode():
79
+ logits = model(input_ids).logits
80
+
81
+ print(logits.shape)
82
+ ```
83
+
84
+ ## References
85
+
86
+ - [Mamba-3: Improved Sequence Modeling using State Space Principles](https://arxiv.org/abs/2603.15569)
87
+ - [Official Mamba implementation](https://github.com/state-spaces/mamba)
88
+
89
+ ## Citation
90
+
91
+ If you use this model, please cite:
92
+
93
+ ```bibtex
94
+ @misc{lahoti2026mamba3improvedsequencemodeling,
95
+ title = {Mamba-3: Improved Sequence Modeling using State Space Principles},
96
+ author = {Aakash Lahoti and Kevin Y. Li and Berlin Chen and
97
+ Caitlin Wang and Aviv Bick and J. Zico Kolter and
98
+ Tri Dao and Albert Gu},
99
+ year = {2026},
100
+ eprint = {2603.15569},
101
+ archivePrefix = {arXiv},
102
+ primaryClass = {cs.LG},
103
+ url = {https://arxiv.org/abs/2603.15569}
104
+ }
105
+ ```
config.json ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "d_model": 2048,
3
+ "d_intermediate": 3824,
4
+ "n_layer": 24,
5
+ "vocab_size": 128256,
6
+ "ssm_cfg": {
7
+ "layer": "Mamba3",
8
+ "d_state": 128,
9
+ "expand": 2,
10
+ "headdim": 64,
11
+ "ngroups": 1,
12
+ "rope_fraction": 0.5,
13
+ "dt_min": 0.001,
14
+ "dt_max": 0.1,
15
+ "dt_init_floor": 0.0001,
16
+ "A_floor": 0.0001,
17
+ "chunk_size": 16,
18
+ "is_mimo": true,
19
+ "mimo_rank": 4,
20
+ "is_outproj_norm": false
21
+ },
22
+ "attn_layer_idx": [],
23
+ "attn_cfg": {},
24
+ "rms_norm": true,
25
+ "residual_in_fp32": true,
26
+ "fused_add_norm": false,
27
+ "pad_vocab_size_multiple": 16,
28
+ "tie_embeddings": true
29
+ }
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:139c09e9728d1f10ac7f4354d2ae29aaca41851a954aedd8f6728c05aeb6c744
3
+ size 2994280907