MLVXN/125M-v6e-trillium

125M v6e Trillium

Best 125M param LLM trained on TPU v6e-1 within $30 budget (9.7h @ 3.09 USD/hr).

Architecture: 30 layers, 576 hidden, 9 heads GQA 3 KV, 1536 FFN SwiGLU, RoPE, RMSNorm, ~125M params — deep-thin per MobileLLM / SmolLM2-135M. Tokenizer: HuggingFaceTB/SmolLM2-135M (49152 vocab) Data mix pretrain 1.5B: 45% FineWeb-Edu, 20% DCLM, 15% Cosmopedia-v2 (synthetic), 10% FineMath+StackEdu, 5% Python-Edu, 5% Distil (teacher CoT Qwen2.5-1.5B + smoltalk) Synthetic: Cosmopedia v2 (Mixtral synthetic textbooks), FineMath synthetic reasoning Distil: SmolTalk traces, OpenHermes-2.5 (GPT-4 distil), UltraFeedback Schedule: WSD warmup 30M, peak 3e-3, decay to 3e-4, AdamW 0.9/0.95, wd 0.1, bf16, grad clip 1.0, seq 1024, batch 64 (65k tokens) Post-train: SFT 200M on smoltalk/openhermes (LR 2e-5) + DPO on UltraFeedback (beta 0.1, LR 5e-7) Hardware: TPU v6e-1 (TPU v6 lite, 1 chip via JAX), libtpu 0.0.44, JAX 0.11, Flax 0.12, Optax, Orbax Eval: HellaSwag, ARC, PIQA, Winogrande (see logs)

Checkpoints: /teamspace/studios/this_studio/content/125m-v6e/checkpoints* Logs: /teamspace/studios/this_studio/logs

Usage

from transformers import AutoTokenizer, AutoModelForCausalLM
tok = AutoTokenizer.from_pretrained("MLVXN/125M-v6e-trillium")
# Note: weights are Flax pickle in flax_params.pkl, convert to PyTorch via script if needed
# For JAX inference: load pickle + train.py LlamaFlax

Training cost: ~5.3h pretrain * 3.09 = ~16.4 USD + SFT/DPO