File size: 4,628 Bytes
88801ff 3c5b617 88801ff 3c5b617 9abce57 3c5b617 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | ---
license: apache-2.0
language:
- en
pipeline_tag: text-generation
tags:
- pebble
- language-model
- base-model
- small-language-model
- pytorch
- safetensors
- custom-code
- mamba2
- hybrid
library_name: transformers
---
# Pebble-25M
Pebble-25M is a compact, hybrid autoregressive language model. It combines the efficiency of state-space models with the proven performance of attention layers, optimized using a custom Muon + AdamW optimizer split.
## Model Details
- **Architecture:** Hybrid Mamba2 / Transformer
- **Block Pattern:** 3 Mamba2 blocks : 1 Attention block (repeating)
- **Parameters:** ~24,500,000 (25M)
- **Hidden Dimension:** 608
- **Layers:** 8 (6 Mamba2, 2 Attention)
- **Vocab Size:** 2,048 (Custom Byte-Level BPE)
- **Context Length:** 2048
- **Training Tokens:** ~25,000,000,000 (~25 Billion)
- **Optimizer:** Muon (for 2D hidden weights) + AdamW (for embeddings, norms, and scalars)
- **Precision:** fp32 master weights with bf16 autocast
## Dataset Sources
The model was trained on a 25B token subset of the following datasets:
| Dataset | Token Allocation | Share |
|---|---:|---:|
| FineWeb-Edu | 7.50 billion | 30% |
| DCLM | 5.00 billion | 20% |
| Cosmopedia-v2 | 3.75 billion | 15% |
| FineMath-4+ | 3.75 billion | 15% |
| FinePhrase | 3.00 billion | 12% |
| NPset | 2.00 billion | 8% |
## Benchmarks
| Benchmark | **Pebble-25M** | Pebble-25M Chat | Pebble-10M | BananaMind-2-Mini | Random |
|---|---:|---:|---:|---:|---:|
| PIQA | 59.25% | 53.37% | 58.43% | **59.63%** | 50.00% |
| ARC-Easy | 38.17% | 26.68% | 37.29% | **39.86%** | 25.00% |
| ARC-Challenge | 18.60% | 19.62% | 18.60% | **25.68%** | 25.00% |
| HellaSwag | 27.62% | 25.63% | 26.81% | **29.72%** | 25.00% |
| ArithMark-2.0 | 27.60% | 26.20% | **27.64%** | 27.52% | 25.00% |
| ArithMark-3.0 | 33.80% | 28.80% | 32.80% | **34.90%** | 25.00% |
### Evaluation Notes
- PIQA, ARC-Easy, ARC-Challenge, and HellaSwag were evaluated on their respective test splits.
- ArithMark-2.0 was evaluated on its train split due to the lack of a suitable test split.
- ArithMark-3.0 was evaluated on its train split due to the lack of a suitable test split.
- Results were obtained using zero-shot multiple-choice evaluation.
- No task-specific fine-tuning was performed.
## Usage
To run the model for text generation, you will need to install the required dependencies. The included Mamba2 implementation relies on CUDA/Triton kernels and is intended to run on a CUDA-enabled GPU. Ampere-class GPUs or newer are recommended.
> **Note:** The model uses custom architecture code, so you must pass `trust_remote_code=True` when loading both the tokenizer and the model.
### Installation
```bash
pip install transformers huggingface_hub torch
pip install causal-conv1d mamba-ssm
```
### Generation
Here is a simple Python script to load the model and generate text interactively:
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "basically-ai/Pebble-25M"
def main():
print("Loading Pebble 25M...")
tokenizer = AutoTokenizer.from_pretrained(
MODEL_ID,
trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
trust_remote_code=True,
dtype=torch.float32,
).to("cuda")
model.eval()
print(
f"Model loaded successfully! "
f"VRAM usage: {torch.cuda.memory_allocated() / 1e9:.2f} GB"
)
print("Type 'quit' or 'exit' to stop.\n")
while True:
prompt = input("You: ")
if prompt.lower() in ["quit", "exit"]:
break
# Tokenize the prompt
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
# Generate text
print("Pebble: ", end="", flush=True)
with torch.inference_mode():
outputs = model.generate(
**inputs,
max_new_tokens=100, # How many tokens to generate
do_sample=True, # Use sampling (more creative)
temperature=0.7, # Controls randomness
top_k=50, # Consider top 50 tokens
top_p=0.95, # Nucleus sampling
repetition_penalty=1.2, # Prevent repeating words
)
# Decode and print (skip the prompt part)
generated_text = tokenizer.decode(
outputs[0][inputs["input_ids"].shape[1]:],
skip_special_tokens=True,
)
print(generated_text)
print()
if __name__ == "__main__":
main()
```
## License
Apache 2.0 |