mamba2-10m-slimpajama-500m/README.md

---
license: apache-2.0
library_name: transformers
tags:
  - mamba2
  - mamba
  - vllm
  - state-space-model
  - causal-language-modeling
  - slimpajama
  - hpml
datasets:
  - cerebras/SlimPajama-627B
pipeline_tag: text-generation
---

# Mamba2 10M SlimPajama 500M

This is the vLLM/Transformers-compatible export of the Mamba2 baseline from our HPML Simamba discretization experiments. It is the same checkpoint as `soumil1/mamba2-10m-slimpajama-500m`, converted from the training repository's `mamba_ssm` checkpoint layout to the standard Hugging Face `Mamba2ForCausalLM` layout that vLLM can identify natively.

The model is a research checkpoint, not an instruction-tuned assistant.

## Checkpoint

| Field | Value |
| --- | --- |
| Variant | Mamba2 baseline |
| Parameters | 9,472,400 unique parameters with tied input/output embeddings |
| Training data | SlimPajama subset, `data/slimpajama_500m_50m/train.bin` |
| Validation data | Held-out SlimPajama subset, `data/slimpajama_500m_50m/val.bin` |
| Tokenizer | `EleutherAI/gpt-neox-20b` tokenizer files are included |
| Context length | 128 tokens during training |
| Best checkpoint step | 122,000 |
| Best validation loss | 4.86253547668457 |
| W&B run | https://wandb.ai/ssb2234-columbia/simamba/runs/dm74y180 |

## vLLM

The repository uses:

```json
{
  "model_type": "mamba2",
  "architectures": ["Mamba2ForCausalLM"]
}
```

Serve with vLLM:

```bash
vllm serve soumil1/mamba2-10m-slimpajama-500m \
  --dtype float32 \
  --max-model-len 128
```

For offline inference:

```python
from vllm import LLM, SamplingParams

llm = LLM(
    model="soumil1/mamba2-10m-slimpajama-500m",
    dtype="float32",
    max_model_len=128,
)
outputs = llm.generate(
    ["State space models are"],
    SamplingParams(max_tokens=32, temperature=0.8),
)
print(outputs[0].outputs[0].text)
```

## Transformers

```python
import torch
from transformers import AutoTokenizer, Mamba2ForCausalLM

repo_id = "soumil1/mamba2-10m-slimpajama-500m"

tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = Mamba2ForCausalLM.from_pretrained(repo_id, torch_dtype=torch.float32).cuda()
model.eval()

input_ids = tokenizer("State space models are", return_tensors="pt").input_ids.cuda()
with torch.no_grad():
    logits = model(input_ids).logits
print(logits.shape)
```

## Architecture

```json
{
  "hidden_size": 160,
  "num_hidden_layers": 8,
  "vocab_size": 50280,
  "state_size": 64,
  "conv_kernel": 4,
  "expand": 2,
  "head_dim": 32,
  "num_heads": 10,
  "n_groups": 1,
  "chunk_size": 16,
  "tie_word_embeddings": true
}
```

## Limitations

This is a small pretraining checkpoint trained for controlled architecture comparison. It was trained on a limited SlimPajama subset with a short sequence length, so generation quality and downstream task performance should be interpreted cautiously.