A research-oriented Small Language Model (SLM) built entirely from scratch in PyTorch to explore modern transformer architectures under real-world compute constraints.
CodeMind is an experimental language model project focused on understanding and implementing modern LLM architectures from first principles rather than fine-tuning existing models.
Instead of modifying LLaMA or Mistral, this repository implements an entire transformer stack from scratch, integrating recent architectural ideas proposed by DeepSeek, Moonshot AI, and Muon into a unified training framework.
The project serves both as
- a research platform for experimenting with new architectures
- a production-quality implementation of a modern language model
- a learning resource for understanding large-scale model training
Implements DeepSeek's latent KV-cache compression mechanism for reducing inference memory while preserving attention quality.
- 4 Routed Experts
- Top-2 Routing
- 1 Shared Expert
- Gradient-compatible load balancing
Implementation inspired by Kimi (Moonshot AI), replacing additive residual connections with attention-based aggregation across transformer blocks.
Predicts multiple future tokens during training to improve latent representations and enable speculative decoding research.
Instead of using a single optimizer:
- Muon handles large weight matrices
- AdamW optimizes embeddings, norms, and other 1D parameters
CodeMind/
│
├── config/
├── model/
│ ├── attention/
│ ├── moe/
│ ├── transformer/
│ └── codemind.py
│
├── tokenizer/
├── dataset/
├── training/
├── inference/
├── scripts/
│
├── checkpoints/
├── assets/
├── technical_report.pdf
└── README.md
CodeMind combines multiple modern research ideas into one architecture.
Embedding
│
RoPE
│
Multi-Head Latent Attention
│
Block Attention Residual
│
Mixture of Experts
│
Multi-Token Prediction Head
│
Language Modeling Head
The initial pretraining checkpoint was trained completely from scratch.
| Item | Value |
|---|---|
| Parameters | ~1B |
| Active Parameters | ~400M |
| Hardware | NVIDIA A40 48GB |
| Dataset | GitHub + OpenWebMath + FineWeb-Edu |
| Tokens | 147 Million |
| Precision | bfloat16 |
| Training Time | ~21 Hours |
The model successfully converged without instability despite combining several modern architectural components.
Highlights:
- Stable MoE routing
- Stable Muon optimization
- Successful MLA implementation
- ~1940 Tokens/sec
- Training Loss: 10.5 → 3.1
Training curves are available below.
A complete write-up of the implementation and training methodology is available.
📄 Technical Report
The pretrained checkpoint is available on Hugging Face.
🤗 https://huggingface.co/B4K2xx/CodeMind-1B-v0.1
CodeMind is only the beginning.
Current research directions include
- Scaling MoE at small model sizes
- Muon vs AdamW optimization
- Long-context training
- Efficient KV-cache compression
- Alternative attention mechanisms
- Instruction tuning
- Reinforcement Learning
- Distributed training
This project draws inspiration from
- DeepSeek
- Moonshot AI (Kimi)
- Muon Optimizer
- PyTorch
- Hugging Face