README.md

---
language: en
license: apache-2.0
tags:
- pytorch
- moe
- bitnet
- 1-bit
- 4-bit
- scratch
- turbowarp
---

# SSMoELM-Base

**Scratch Small MoE Language Model** — a 1-bit sparse MoE language model designed to run inside [Scratch](https://scratch.mit.edu) / [TurboWarp](https://turbowarp.org).

- **47M total / 25.8M active parameters** (top-2 sparse routing)
- **12.1 MB** packed weights (1-bit routed experts, 4-bit attention & embedding)

> **Note:** The HuggingFace model card may display ~12M parameters and an "8-bit" quantization badge. Both are artifacts of reading the packed `model.safetensors` directly: weights are stored as `uint8` arrays (bit-packed 1-bit or nibble-packed 4-bit values), which HF counts as fewer elements and infers as 8-bit storage. The actual model has **47M parameters** quantized to **1-bit and 4-bit**.
- Trained from scratch on English web text for **900M tokens**

> "Scratch" carries two meanings: built *for Scratch*, trained *from scratch*.

---

## Model Details

| | |
|---|---|
| Architecture | Decoder-only Transformer + Sparse MoE FFN |
| Total params | 47.04M |
| Active params | 25.80M (per forward pass) |
| d_model | 768 |
| Layers | 6 |
| Attention | GQA — 12 heads, kv_heads=3, head_dim=64 |
| Positional encoding | RoPE |
| Normalization | RMSNorm |
| Activation | SwiGLU |
| MoE | 8 routed experts + 1 shared expert, top-2 routing |
| d_ff (per expert) | 256 |
| Vocabulary | 8,192 (BPE, byte-fallback, English-optimized) |
| Context length | 2,048 tokens |
| Training tokens | 900M |

---

## Quantization Scheme

QAT (BitNet b1.58 style): forward pass uses quantized values, gradients flow through master bf16 weights via straight-through estimator (STE).

| Component | Bits | Notes |
|---|---|---|
| Routed expert gate/up/down | 1-bit | sparse, 2/8 active per token |
| Shared expert gate/up/down | 4-bit | runs on every token |
| Attention Q, K (all layers) | 4-bit | precision-sensitive |
| Attention V, O (layers 0, 5) | 4-bit | boundary layers |
| Attention V, O (layers 1–4) | 1-bit | inner layers |
| Embedding | 4-bit | tied with lm_head |
| Router, RMSNorm | bf16 | not quantized |

Packed format per tensor:
- **1-bit**: `{key}__scale` (fp16, per-row) + `{key}__bin` (uint8 packbits)
- **4-bit**: `{key}__scale` (fp16, per-row) + `{key}__int4` (uint8 nibble-packed)
- Metadata: `{key}__shape` (int32) + `{key}__cols` (int32)

Keys use `/` as separator in safetensors (replace with `.` for model loading).

---

## Training

| | |
|---|---|
| Dataset | FineWeb-Edu-score-2 (60%) + FineWeb (40%) |
| Tokens | 900M |
| Optimizer | AdamW, betas=(0.9, 0.95), weight_decay=0.1 |
| Learning rate | 3e-4 → 1e-5 cosine; router lr = 1e-4 |
| Warmup | 5% of total steps |
| Gradient clipping | 1.0 |
| Batch size | 64K tokens effective (bs=2, grad_accum=16, ctx=2048) |
| Hardware | Apple M2 (MLX) |
| Throughput | ~3,000 tokens/sec |

---

## Benchmark Results (0-shot, 500 samples)

| Task | Shot | Metric | Samples | Random | Score |
|---|---|---|---|---|---|
| HellaSwag | 0-shot | acc_norm | 500 | 25% | 33.4% |
| LAMBADA | 0-shot | acc | 500 | N/A | 13.8% |
| PIQA | 0-shot | acc_norm | 500 | 50% | 53.2% |
| WinoGrande | 0-shot | acc | 500 | 50% | 49.6% |
| ARC-Easy | 0-shot | acc_norm | 500 | 25% | 35.0% |
| ARC-Challenge | 0-shot | acc_norm | 500 | 25% | 21.0% |
| BoolQ | 0-shot | acc | 500 | 50% | 36.2% |
| MMLU (57 tasks avg) | 0-shot | acc | up to 500/task (total 12,173) | 25% | 23.4% |

---

## Expert Routing Statistics

Measured on 136 tokens (8 diverse text samples), top-2 routing. Uniform load = 12.5%.

| Layer | E0 | E1 | E2 | E3 | E4 | E5 | E6 | E7 | CV |
|---|---|---|---|---|---|---|---|---|---|
| 0 | 8.8% | 7.7% | 11.0% | 15.1% | **18.8%** | 18.0% | 11.0% | 9.6% | 0.32 |
| 1 | 12.5% | 11.0% | 11.4% | 13.2% | **16.5%** | 9.9% | 12.9% | 12.5% | 0.15 |
| 2 | 10.7% | **18.0%** | 16.2% | 10.3% | 9.6% | 11.4% | 11.4% | 12.5% | 0.22 |
| 3 | 10.7% | 6.2% | 14.3% | 8.5% | 11.8% | 7.7% | **22.8%** | 18.0% | 0.42 |
| 4 | 12.1% | **16.5%** | 10.3% | 10.7% | 14.0% | 18.0% | 8.8% | 9.6% | 0.25 |
| 5 | **18.0%** | 15.1% | 8.8% | 12.9% | 9.6% | 9.6% | 15.1% | 11.0% | 0.25 |

CV = coefficient of variation (lower = more balanced). No expert collapse observed.

---

## Tokenizer

- BPE, vocabulary size = 8,192
- Byte fallback enabled (no `<unk>`)
- ASCII/English-optimized segmentation

### Special Tokens

| Token | ID | Role |
|---|---|---|
| `<bos>` | 0 | sequence start |
| `<eos>` | 1 | end of sequence |
| `<pad>` | 2 | padding |
| `<\|system\|>` | 3 | system turn |
| `<\|user\|>` | 4 | user turn |
| `<\|assistant\|>` | 5 | assistant turn |
| `<\|eot\|>` | 6 | end of turn |

---

## Usage

Download `inference.py` from this repo. Requires: `torch`, `safetensors`, `tokenizers`.

```bash
pip install torch safetensors tokenizers
```

```python
from inference import load_packed_model
from tokenizers import Tokenizer

model = load_packed_model("model.safetensors")  # 12 MB, no dequantization
tok   = Tokenizer.from_file("tokenizer.json")

ids = [0] + tok.encode("The quick brown fox").ids  # 0 = <bos>
out = model.generate(ids, max_new_tokens=100, temperature=0.8)
print(tok.decode(out))
```

CLI:
```bash
python inference.py \
  --ckpt model.safetensors \
  --prompt "The quick brown fox" \
  --max-tokens 100 \
  --temperature 0.8
```

> **Memory:** Weights stay in packed uint8 format (12.2 MB). Unpacking is done on-the-fly per layer during forward — peak RAM usage is ~18 MB (12 MB stored + ~6 MB for the largest layer unpacked).

---

## License

Apache 2.0