The Challenge
Frontier Models API costs are expensive, free models suck at reasoning, and the AI tech stack is way too expensive in the long-run. Small Language Models (SLMs) excel at their domain, often beating Frontier Models.
The Solution: Domain-Specific SLMs
I embarked on building three domain-specific models (SLMs) for long-term AI model solutions to power my data platforms. Google's 6th Gen and 4th Gen TPUs, available for free for 30 days under the TRC Program, were the perfect solution.
Data Sources
- FineWebEdu
- StackOverflow
- Wikipedia
- General Web Scraping
- CWE Corpus
- Security adversaries
- Compliance docs
- Regulatory texts
Setup & Architecture
To maximize TPU utilization, I used:
- MaxText/JAX framework
- FSDP (Fully Sharded Data Parallel)
- Orbax Checkpoints on GCS
- Automated scout + autopilot layer
The MoE architecture uses 2B active parameters per token, giving the knowledge of 9B while paying inference costs closer to a 2B dense model.
Training Process
The most painful aspect of using free TPUs is that Trillium VMs get preempted frequently. To automate this:
- Auto setup and launch training
- Scout script runs every 10 minutes
- Automatically resumes training when VMs get preempted
I initially made the mistake of saving all checkpoints, resulting in 40TB of training data and exploding Google Cloud Storage costs. Later, I learned to save only the last 1 or 2 checkpoints.
Workflow
- Setup
- Data Download
- Tokenize
- Train
- SFT + GRPO
- Deploy
Post-Training
- SFT + GRPO
- Raw Weights
- GGUF Formats
- Ollama integration
Timeline & Costs
- 40+ Days of training
- $5K-$10K USD in total costs
- Hidden costs and miscellaneous expenses
Post-mortem: what went wrong, and how we found it
Added August 2026, four months after the original post. The three models finished pretraining with clean loss curves. They were still not usable, and the reason took a full diagnostic pass to find. This section is the part of the story that matters more than the build log above it.
The loss curve was fine and the models were not. Every run descended smoothly and every checkpoint saved. What the loss could not show was that the training configuration had been quietly damaging the parts of the network that do not appear in it.
The root cause: weight decay applied to normalisation gains
MaxText's adamw_mask defaults to an empty list, which means weight decay is applied to every parameter, including RMSNorm gains and the embedding table. A normalisation gain and the projection it feeds are multiplicatively coupled: only the product W·diag(g) affects the function the network computes. So decay can shrink a gain toward zero, remove a layer's contribution, and leave the loss curve perfectly smooth while it happens.
This is a configuration default, not a bug in anyone's code. It is also invisible to every metric a normal training dashboard shows.
What we measured on the finished weights
| Measurement | Result | What it means |
|---|---|---|
| Dead embedding rows | 8,507 of 49,152 | 17% of the vocabulary has an input embedding that never trained |
| Effective weight spread | 837,534× | Ratio between the strongest and weakest effective layer weights. Healthy models sit in the tens |
| Special-token imbalance | ids 0–3 at ~0.005–0.05 | Input embeddings are dead while their output rows are live, so the model can emit tokens it cannot read and those ids must be suppressed at inference |
| In-domain perplexity (SQL) | 3.8 | The models genuinely learned their domains |
| General instruction following | poor | Fluent in domain, unreliable outside it |
The layer-31 spike, and why we did not "fix" it
One model carried an activation spike at layer 31 measuring 7,068× the residual stream. It looked like obvious damage, and the obvious move was to attenuate it. We ran the sweep before touching anything: removing the spike moved in-domain SQL perplexity from 5.10 to 406.4. The anomaly was load-bearing. The model had routed a large part of its function through it.
That result changed the conclusion. These models are not damaged in a way a repair could undo. They were shaped this way by the training configuration, and the shape is what they compute with. There is no patch. The fix has to happen before the run, which is the whole point of what we built next.
What we did with each model
- Kautilyaa (finance): recovered far enough to be useful with supervised fine-tuning and grounding work, evaluated at 0 hallucinated figures across 150 test cases, and published on HuggingFace under Apache-2.0 with the limitations documented in the model card.
- Bonacci (data engineering) and CyberAgent (security): competent in their domains and retained as internal research checkpoints. Most of the diagnostic work described above came out of them.
The decisions we would make differently
- Set the decay exclusion mask explicitly, and never inherit it from a framework default. This one line accounts for most of what went wrong.
- Watch a second curve. We now record the minimum RMSNorm gain across layers at every checkpoint alongside the loss. It is the series that moves when decay is eating the gains, and the loss is the series that does not.
- Checkpoint every 2,000 steps on spot capacity, not 10,000. A long interval meant preemptions restarted from zero and cost days.
- Compare against the base model on every evaluation. An absolute score looks acceptable until you see what an untouched open checkpoint scores on the same prompts.
- Keep a hand-written probe set. Automated evaluation passed on models that failed obviously the moment a human read the outputs.
What this became
Every failure above is now a rule that runs before a training job starts, in a system we built for exactly this reason. The decay-mask check blocks a run whose exclusion list is empty. The checkpoint-interval check blocks an interval too long for spot capacity. The corpus gates measure duplication, PII and template collapse before compute is spent, and the artifact gates inspect the delivered weights for the dead-token and spread signatures listed above.
We paid for that list with three 9B models and 40 days of TPU time. It is the reason we can tell you which rung you need in a week, and the reason we are comfortable telling you when the answer is that you do not need a model at all.
Key Takeaways
Owning your stack means owning the failures too. Building your own models gives you control and long-term cost savings, and it also gives you a class of failure that no API vendor exposes you to. The three models here cost 40 days and $5–10K to produce, and most of their value turned out to be the diagnosis rather than the weights.
A clean loss curve is not evidence that a model works. That is the single most expensive thing we learned, and it is why every rule in our build system cites the measurement that justifies it rather than a convention.
- own your stack ✎
← back to blog