---
license: mit
tags:
- text-generation
- custom-architecture
- pytorch
- bfloat16
---

# SAN x-small

Simple Attention Network — a ~31M-parameter custom LLM, a faithful PyTorch port of the
`needle` architecture (arXiv:2607.18363). Trained locally on an RTX 4070 Ti 12 GB.

This repo contains the **released weights** for the 12-layer configuration.

## Files
| file | what |
|---|---|
| `pytorch_model.bin` | model `state_dict` (load via `san_model.SimpleAttentionNetwork.load_state_dict`) |
| `san_latest.pt` | raw training checkpoint (step 137,260): `{step, loss, model_state_dict, optimizer_state_dict, config}` |
| `config.json` | architecture hyperparameters |
| `tokenizer/` | SmolLM2-135M tokenizer (vocab 49152) |
| `san_model.py`, `san_triton.py` | model definition (needed to load the weights) |
| `load_example.py` | minimal load + forward example |

## Load
```python
from san_model import SimpleAttentionNetwork, SANConfig
import torch
cfg = SANConfig(num_layers=12)
model = SimpleAttentionNetwork(cfg).eval()
sd = torch.load("pytorch_model.bin", map_location="cpu")
model.load_state_dict(sd)
```
See `load_example.py`. Full training/eval code: GitHub
[`kenpeter/x-small`](https://github.com/kenpeter/x-small).

## Architecture (released config)
- params ~31M · d_model 384 · 12 layers · 8 heads / 4 KV heads (GQA)
- vocab 49152 (SmolLM2-135M) · context 2048 · bfloat16
- components: ZCRMSNorm, HadamardMLP, GQA + RoPE + learned gate, Engram n-gram KV memory,
  Multi-Lane Hyper-Connections (MHC, Sinkhorn routing), MTP head, tied embeddings

## Training & eval
- Curriculum over 16 domains; target 100w (1,000,000) steps; eff batch 64.
- Domains: easy — code_easy, math_easy, reformat_easy, synth_easy, web_easy; medium — code_medium, math_medium, synth_medium, web_medium; hard — code_gold, code_hard, gold_hard, math_hard, synth_hard, web_gold, web_hard.
- Training token budget (nominal, full 64 × 2048 batches): 131,072 tokens/step; ~17.99B at step 137,260; ~131.07B at 1,000,000 steps. Loss-masked curriculum tokens are not subtracted.
- Hardware limit: single RTX 4070 Ti 12GB; micro-batch 8 × grad-accum 8 (eff 64) with gradient checkpointing, memory-efficient block-wise attention, and expandable CUDA segments to stay under the 12GB ceiling.
- Last training step **137,260** (of 1,000,000).
- Last eval (step 123,141) perplexity: **mean 23.74** (code 10.25 / math 29.62 / prose 31.34).
  Perplexity is the usable quality signal at this scale; sample text loops regardless.
- The flat training loss near ~2.0 is curriculum-mix masking, not model stall.

## Benchmarks
- Quality trend (`benchmarks/ppl_trend.png`, data in `benchmarks/ppl_trend.csv`): code/math/prose perplexity across the 12-layer evals (steps 63,663 → 110,938 → 123,141).
- Throughput (`benchmarks/throughput.png`, data in `benchmarks/tput_release.csv`): released 12-layer weights, eager forward+backward with grad-checkpointing, bf16, seq 2048, RTX 4070 Ti 12GB. B=8 OOM because the GPU was shared with another training job at bench time; the release run used B=8 on a dedicated GPU.

![perplexity trend](benchmarks/ppl_trend.png)
![throughput](benchmarks/throughput.png)

## License
MIT (code). Weights released for research use.