KaptaanLM-Navigator-50M

KaptaanLM-Navigator-50M: Master Fleet Repository

KaptaanLM-Navigator (52.72M parameters) is the 50M foundation model of the KaptaanLM architecture fleet, designed and optimized for Google Cloud TPU v5e-8 accelerators using distributed JAX/Flax NamedSharding and activation rematerialization.

• GitHub Repository: https://github.com/rudy-07/KaptaanLM
• Architecture: 16 layers, d_model=512, 16 attention heads, 4 KV heads (4:1 GQA), d_ffn=1376 (SwiGLU), 16,384 Tied BPE vocabulary
• Context Length: 2,048 tokens (512 sliding window attention)
• Total Parameters: 52,718,336 (52.72M)
• Non-Embedding Parameters: 44,335,104 (44.34M / 84.11% reasoning capacity)
• Vocabulary: 16,384 BPE (Tied Embeddings)

---

Artifacts in this Repository

---

TPU v5e-8 Pilot Run Dynamics (1.0 Billion Tokens)

• Target Hardware: Google Cloud TPU v5e-8 (8 TensorCore chips, 128 GB HBM, pure DP=8 NamedSharding)
• Execution: 1,000,000,000 tokens processed across 3,815 steps in 2.24 hours
• Global Batch Size: 128 sequences = 262,144 tokens per step
• Loss Descent: Monotonic descent from 10.1630 down to 2.0172 during cosine cooldown to 7.5e-5
• Validation Perplexity: Held-out 131k-token perplexity reduced by 68.3% (98.77 -> 31.27)
• Memory Safety: Peak HBM held at 2.85 GB / 15.75 GB per core (>81% safety margin) via nn.remat(KaptaanBlock)

---

Planned Foundation Scale: Navigator-50M (30B Tokens)

• Target Token Budget: 30.0 Billion tokens delivering ~676x non-embedding parameter token density.
• Expected Outcomes: Surpassing 10%–15% GSM8K pass rate and 45%+ ARC-Easy at 52M parameters.

Citation

@software{kaptaanlm2026,
  author = {Shekhar, Rudransh},
  title = {KaptaanLM: Conversational and Reasoning Small Language Models (10M - 100M -> MoE)},
  url = {https://github.com/rudy-07/KaptaanLM},
  year = {2026}
}