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

license, library_name, tags, datasets, pipeline_tag

license	library_name	tags	datasets	pipeline_tag
apache-2.0	transformers	mamba2 mamba vllm state-space-model causal-language-modeling slimpajama hpml	cerebras/SlimPajama-627B	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:

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

Serve with vLLM:

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

For offline inference:

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

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

{
  "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.