File size: 4,445 Bytes
d10ba4c c6f07f0 d10ba4c c6f07f0 eebf5da c6f07f0 eebf5da c6f07f0 eebf5da c6f07f0 eebf5da c6f07f0 | 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 | ---
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
---
# Pebble-10M
Pebble-10M 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:** ~10,000,000 (10M)
- **Hidden Dimension:** 384
- **Layers:** 8 (6 Mamba2, 2 Attention)
- **Vocab Size:** 2,048 (Custom Byte-Level BPE)
- **Context Length:** 512
- **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
Pebble-10M performs above random chance on several commonsense and arithmetic benchmarks.
| Benchmark | Accuracy | Random Baseline |
|-----------|----------|------------------|
| PIQA | 58.43% | 50.00% |
| ARC-Easy | 37.29% | 25.00% |
| ARC-Challenge | 18.60% | 25.00% |
| HellaSwag | 26.81% | 25.00% |
| ArithMark-2.0 | 27.64% | 25.00% |
| ArithMark-3.0 | 32.80% | 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.
```bash
pip install transformers huggingface_hub torch
pip install causal-conv1d mamba-ssm
```
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-10M"
def main():
print("Loading Pebble 10M...")
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! 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
|