Small models for the physical world: how we trained Dhīr from scratch
Why we built our own foundation models, where small models win, and how we pre-trained Dhīr-2B and Dhīr-5B MoE from random weights on Google Cloud TPU v6e.
Why train our own models at all?
The next decade of AI will not live only in chat windows. It will live inside machines: drones, robots, vehicles, plants, grids and the systems that defend them. Those machines need intelligence they can carry with them, trust completely, and run on their own terms.
There are good open models to fine-tune, so starting from random weights is a deliberate choice. We made it for three reasons.
- The base model sets the ceiling. What a model can learn later depends on what it saw in pre-training. Engineering knowledge, physical reasoning and Indian languages have to be in the foundation, not bolted on afterwards.
- Sovereignty you can audit. Defence, government and critical-infrastructure customers need to know exactly what went into a model. Training from scratch means we can account for every dataset, every stage and every weight.
- Control of the whole stack. Owning the data pipeline and the training recipe lets us shape each model for the hardware it will run on, and build the next one faster.
Why small models, and why now?
Large language models are outstanding generalists. For writing code, open-ended research or long-form drafting, a frontier model in the cloud is often the right tool, and we use them ourselves.
Physical systems have different constraints. Answers are needed in milliseconds, not seconds. Networks are unreliable, contested or switched off on purpose. Data is too sensitive to leave the site. Power, memory and cost per device are fixed by the hardware. A model that needs a distant data centre to think is the wrong shape for that work.
The timing is right too. Capable accelerators are now shipping in vehicles, drones, robots and industrial gateways. The hardware has arrived at the edge. It needs models built for it.
Large models in the cloud
- Broad knowledge across every topic
- Best for coding and open-ended research
- Need a fast, reliable connection
- Data leaves the device
Small models at the edge
- Run on the device, in milliseconds
- Work offline and in contested networks
- Data stays where it is created
- Cheap to fine-tune for one site, fleet or platform
Where will Dhīr run?
Every sector below shares one requirement: the model has to live where the work happens. Choose one to see what that means in practice.
What Dhīr does there
Why it has to run at the edge
What were we aiming for?
We set four objectives for the first Dhīr models before a single training step ran. Every decision in the rest of this post traces back to one of them.
Models at the 2B scale that fit on a single accelerator or edge box, with a mixture-of-experts variant that adds knowledge without adding per-token cost.
A data schedule that moves from broad knowledge toward maths, code and reasoning, and a 16k-token context for reading full manuals and logs.
Ten Indian languages in the pre-training mix, in native script and in the romanised forms people actually type.
High hardware utilisation on Google Cloud TPUs, every dataset traceable, and every phase evaluated before the next one begins.
What did we train?
Two decoder-only transformers, both pre-trained from random weights, with the same depth and the same training plan. They differ in how their parameters are used.
Dhīr-2B is dense: every parameter works on every token. Dhīr-5B MoE replaces the feed-forward block in each layer with 12 experts and a router that sends each token to 2 of them. It holds 5.09 billion parameters but uses only 1.8 billion per token, so it carries more knowledge for close to the per-token compute of the dense model.
Dhīr-2B
Dense decoder
- Total parameters
- 2.03B
- Active per token
- 2.03B
- Layers / width
- 28 / 2,048
- Feed-forward
- Dense
- Context
- 8k in pre-training, 16k after
- Tokens seen
- 6T+
- Hardware
- Google Cloud TPU v6e
Dhīr-5B MoE
Mixture-of-experts decoder
- Total parameters
- 5.09B
- Active per token
- 1.8B
- Layers / width
- 28 / 2,560
- Feed-forward
- 12 experts, top-2 routing
- Context
- 8k in pre-training, 16k after
- Tokens seen
- 6T+
- Hardware
- Google Cloud TPU v6e
How did we use the TPUs?
Both runs used Google Cloud TPU v6e (Trillium) with MaxText, Google’s JAX training framework, bf16 compute, and Grain for deterministic data loading.
Two networks connect the chips. Inside a pod, chips are wired directly to each other through the inter-chip interconnect (ICI), a high-bandwidth mesh built for collective operations. Between pods, traffic goes over the data-centre network (DCN), which is slower. The rule we designed around was simple: keep the heavy, frequent traffic on ICI, and send only what must cross pods over DCN.
Weights, gradients and optimiser state are fully sharded across the pod’s chips. Each chip gathers the parameters it needs over ICI just in time, so no chip holds a full copy of the model.
Pods act as data-parallel replicas and exchange only gradient synchronisation over DCN, so the slower network carries the least traffic and adding pods adds throughput.
For Dhīr-5B MoE, tokens are routed to their 2 of 12 experts with sparse grouped-matrix kernels (megablox), which skip the work for experts a token does not use.
How hard did the chips work?
Model FLOPs utilisation (MFU) is the share of the hardware’s peak arithmetic that goes into the model’s own forward and backward passes. It is the standard way labs report training efficiency, and it is harder to keep high on small models, because each chip has less work per step to hide overhead behind.
Model FLOPs utilisation on TPU v6e
Against a peak of 918 TFLOP/s per chip in bf16. The markers show published figures from much larger dense runs. Select a bar for detail.
What did the models learn from?
Each model read over six trillion tokens of English text, code, maths, multilingual web text and ten Indian languages, in three stages.
Within pre-training, the mixture changes from phase to phase. Early phases build broad knowledge from web text. Later phases lean into maths, reasoning and code, and bring in ZenteiQ’s own synthetic corpus.
Mixture by pre-training phase
Share of tokens by source. Dhīr-2B ran two phases; Dhīr-5B MoE ran three.
Overall share by category
Share of all tokens each model saw, with both models side by side. The remainder is multilingual web text.
Why build the Indian-language data ourselves?
Because the open web will not supply it. In one sample of 19,200 documents from a major multilingual crawl, Indian languages made up 0.39%. We curated this part of the corpus per language, from sources built for each one.
How did we set the recipe?
Batch size, sequence length and learning rate decide how much a model learns from each token, so we tuned them as one decision rather than three.
Why 8,192-token sequences?
Long enough to learn document structure, whole code files and multi-step reasoning, while keeping attention cost manageable for the bulk of training. Longer context was added at the end, where it is far cheaper.
Why grow the batch, and why stop at 38%?
Every model has a critical batch size. Below it, a larger batch speeds up learning almost for free; above it, extra tokens per step buy very little, and the compute goes into throughput rather than learning. The critical batch also grows as the model improves.
So Dhīr-2B started at 256 sequences per step, about 2.1 million tokens, and moved to 1,024 sequences, about 8.4 million tokens, at step 660,000. In our ablations, a larger batch still reached 40% MFU, but it sat past the critical point. We chose the batch that made the best use of every token and ran at 38%.
What shape is the learning rate?
A short warm-up, then a long stretch at a high, stable rate while the model takes in most of its data, followed by a decay as each stage hands over to the next. Mid-training restarts from a lower peak and decays to near zero, so the final checkpoint settles cleanly.
How did the training run go?
This is Dhīr-2B’s real pre-training log: 1,250,000 steps and 6.19 trillion tokens across two phases, before mid-training and long-context training. Drag across the chart, or use the slider, to read any point.
Dhīr-2B pre-training
Training loss or held-out bits per byte. Lower is better.
What happens in the first few thousand steps?
Starting from random weights, loss falls about fiftyfold in the first 2,000 steps, so steeply that it starts above the top of the chart, as the model learns token frequencies and basic structure, then settles into the slow, steady decline that makes up most of pre-training.
Why does the loss step down at phase 2?
At step 972,000 the training loss drops sharply. That is the data, not the model. Phase 2 shifted the mixture toward maths and reasoning, which is more predictable text, so the loss level moved with it. Hugging Face’s SmolLM training guide makes the same point: changing the data mixture changes loss in ways that say little about capability on their own.
So which number do we trust?
Held-out bits per byte, measured on the same fixed test set at every checkpoint. Switch to it and the step disappears: the model improves smoothly through both phases.
How do we know it worked?
Every checkpoint goes through the same evaluation suite, and a phase only hands over to the next once the results hold up.
A fixed held-out set spanning 20 domains, including English, code, biomedical text and Indian languages. It never changes, so every checkpoint is compared on equal terms.
Multiple-choice benchmarks, run once the model is far enough along for them to carry signal. Early in training they sit at chance and tell you nothing.
Cross-lingual suites such as XStoryCloze and XNLI, read side by side with an English control, so gains in one language are not hiding losses in another.
Any two checkpoints can be compared domain by domain, so a regression in one area is caught before it reaches the next phase.
How do we pick the checkpoint that moves on?
With brahmai-eval, the evaluation harness we built in-house for Dhīr and released as open source. Give it a checkpoint and it runs the model on 228 tasks across English understanding, code, maths, Indian languages, Indic generation and tool use, in one command. vLLM does the inference; the EleutherAI lm-evaluation-harness and the official BFCL checkers do the scoring.
Every checkpoint we save goes through it. At the end of each phase we compare the candidates side by side and hand the best one to the next stage of training, not simply the latest.
Checkpoints are written at regular intervals throughout each phase.
Bits per byte on the fixed 20-domain set, for a smooth signal at every step.
The benchmark suites that suit the stage, from English and maths to Indic and code.
The strongest checkpoint across domains starts the next phase.
What did we learn?
Measure the model, not the batch
Training loss moves with the data mixture. A fixed held-out set beside every run is what tells a better model apart from an easier batch, and it is how we decide when a phase is done.
Batch size is a learning decision
The batch that maximises chip utilisation is not always the one that maximises learning per token. Giving up two points of MFU to stay at the critical batch was the right trade.
Hand over the best checkpoint, not the last
Evaluating every checkpoint with brahmai-eval and choosing the strongest across domains gave each new phase the best possible starting point.
Keep the heavy traffic on ICI
Designing the sharding so that only gradient synchronisation crosses pods is what let both runs scale across pods without the network becoming the bottleneck.
Indian-language data has to be built, not found
Curating it per language, in native and romanised script, is part of the training work, not a preprocessing step.