Tiny-Talker-55M (Tri-Checkpoint Bundle)
Tiny-Talker-55M is a custom 55.1-million parameter causal language model built in JAX/XLA, engineered for high-throughput execution on TPU v6e hardware within SRAM (VMEM) limits.
This repository releases three artifacts:
base_model_sft.pkl: Pretrained on 2.0B tokens (FineWeb-Edu & Cosmopedia-v2) and SFT aligned on SmolTalk.dora_adapters.pkl: Standalone Weight-Decomposed Low-Rank Adaptation (DoRA, Rank 16) parameters ($A, B, m$).merged_model.pkl: Fully merged model weights for inference.
Download Model Weights
- Download merged model (~110 MiB) — Ready for inference with DoRA already merged.
- Download base SFT model (~110 MiB) — Original model before DoRA adaptation.
- Download DoRA adapters (~2.34 MiB) — Standalone adapter weights; requires the base model.
Architectural Specifications
| Parameter | Specification |
|---|---|
| Total Parameters | 55.06 Million |
| Model Type | Decoder-only Transformer (LLaMA/Gemma style) |
| Layers ($L$) | 6 |
| Hidden Size ($D$) | 768 |
| Attention Heads ($NH$) | 6 (Head Dim: 128) |
| MLP Dimension ($FF$) | 2048 (SwiGLU activation) |
| Position Embeddings | Rotary Position Embeddings (RoPE) |
| Normalization | Pre-RMSNorm ($\epsilon = 10^{-6}$) |
| Vocabulary Size | 16,384 (Byte-Level BPE) |
| Context Length | 1,024 tokens |
Artifact Manifest
| File | Approximate Size | Description |
|---|---|---|
merged_model.pkl |
~110 MiB | Weights with DoRA adaptation merged into base parameters. |
base_model_sft.pkl |
~110 MiB | Pristine post-SFT checkpoint. |
dora_adapters.pkl |
~2.34 MiB | Direction matrices ($A, B$) and magnitude vectors ($m$). |
tokenizer.json |
~1.8 MiB | Byte-level BPE tokenizer (16,384 vocab). |
config.json |
<1 KiB | Model topology and generation special token mappings. |
Reported Hardware Performance (Google Cloud TPU v6e-1)
- Static Weights (INT8 Quantized): 52.60 MiB
- On-Chip VMEM Working Set: 73.67 MiB / 128.00 MiB (57.6% SRAM footprint)
- Training Throughput: ~0.73 Million tokens/second in-SRAM
- Prefill Speed: 1.99 Million tokens/second
Quickstart: Loading the Merged Weights
These are custom JAX weights. Generation requires the matching model implementation; this example only loads weights and prepares input tokens. Only unpickle files from a trusted source.
import pickle
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download
REPO_ID = "niko3x/tiny-talker-55m"
tok_path = hf_hub_download(REPO_ID, "tokenizer.json")
weights_path = hf_hub_download(REPO_ID, "merged_model.pkl")
tokenizer = Tokenizer.from_file(tok_path)
with open(weights_path, "rb") as f:
params = pickle.load(f)
prompt = "Explain why the sky is blue in simple terms."
US, EOS, AS = [
tokenizer.token_to_id(s)
for s in ["<|user|>", "<|eos|>", "<|assistant|>"]
]
assert all(t is not None for t in (US, EOS, AS)), "Missing special tokens"
input_ids = [US] + tokenizer.encode(prompt).ids + [EOS, AS]
Limitations
- Parameter Capacity: At 55M parameters, the model cannot reliably solve multi-step math or deep symbolic logic.
- Hallucination Profile: Outputs reflect synthetic pretraining biases and may generate plausible-sounding technical inaccuracies outside common prompt patterns.
- Downloads last month
- -