The Problem

Training a modern transformer‑style language model usually requires dozens of gigabytes of GPU memory and a large codebase that hides the experimental components. Researchers who want to experiment with the Kimi K3 attention mechanisms must either re‑implement the tricks from scratch or work with a heavyweight repository that cannot run on a single 8 GB GPU.

What This Does

smol‑kimi‑k3 delivers a 49 M‑parameter K3‑inspired model that fits on one 8 GB GPU. The core implementation lives in the k3/ package:

  • k3/model.py – defines the K3ForCausalLM class, its forward and step logic.
  • k3/config.py – holds K3Config and helpers (head_size, attention_type).
  • k3/optim.py – custom optimizer utilities (zeroth_power_via_newton_schulz, step).
  • k3/monitoring.py – runtime metrics (_now, _write_state).

Training utilities are in train.py (batch preparation, learning‑rate scaling, checkpointing, and a generate_sample helper). A minimal live dashboard is provided by dashboard.py (HTTP handler, file serving) together with the static assets in dashboard/. The scripts/prepare_tinystories.py script tokenizes the TinyStories dataset and writes the split files used by training.

How It Is Wired

Entry point – the only explicit main function is in dashboard.py (line 56). Starting the dashboard runs:

dashboard.py → main
   ├─ encode_split → writes split files (output_path.open)
   ├─ set_status / load_state_dict (model checkpoint handling)
   └─ K3ForCausalLM (instantiated 5 times across the repo)

K3ForCausalLM is the central hub: it is created in train.py, generate.py, the dashboard handler, and the test suite. Each instance calls:

  • __init__ → constructs layers (RMSNorm, CausalDepthwiseConv1d, AttentionResidual, SiTUGLU).
  • forward → invokes _shape_heads (4 calls per forward) and the custom KDA/MLA kernels.

The call graph shows the widest blast radius on K3ForCausalLM (5 distinct callers) and on internal helpers _write_state (5 callers) and _now (4 callers). Changing these functions will affect most of the codebase.

train.py orchestrates the training loop:

  • get_batch → reads tokenized data files.
  • evaluate → runs a forward pass and computes loss.
  • save_checkpoint → writes model state (state_dict) to disk.
  • generate_sample → calls K3ForCausalLM.forward and then generate.py’s main for inference.

tests/test_model.py validates the same pathways, exercising tiny_config (a minimal K3Config) and exercising forward/backward passes, causal masking, and cached decoding.

No circular imports were detected; the import graph is shallow (11 modules, 9 edges). Deep nesting (max indentation depth 6) appears in k3/optim.py, dashboard.py, and train.py, which can make the control flow harder to follow.

How To Use It

# Clone the repo
git clone https://github.com/moses-y/smol-kimi-k3
cd smol-kimi-k3

# Install Python dependencies
pip install -r requirements.txt   # or: pip install -e .  (pyproject.toml defines the package)

# Prepare the TinyStories tokenized split (requires internet access)
python scripts/prepare_tinystories.py

# Start the live training dashboard (optional)
python dashboard.py &   # serves on localhost, default port 8000

# Train the model (single‑GPU)
python train.py          # uses default config from configs/k3_50m.json

# Generate a sample after training
python generate.py

The repository supplies configs/k3_50m.json as the default model hyper‑parameters; no environment variables are required beyond a CUDA‑capable GPU.

Real‑World Use

A data‑science team can embed the model in a microservice that calls K3ForCausalLM.forward for inference, while developers experiment with the KDA/MLA blocks by editing k3/model.py and re‑running train.py. The lightweight dashboard provides live loss curves without needing external monitoring tools.

Code Health & Issues

  • High – Missing LICENSE – no LICENSE file; reuse is legally undefined.
  • High – No lockfile – pyproject.toml lacks a corresponding lockfile; builds are non‑reproducible.
  • High – No CI pipeline – repository contains no GitHub Actions, Travis, etc.; changes are not automatically verified.
  • Medium – No Dependabot/Renovate – dependency updates are manual; known vulnerabilities may linger.
  • Medium – Version ranges include vulnerable releases – e.g., torch>=2.5 covers a CVE‑critical version; without a lockfile the exact version is ambiguous.

Additional observation from the measured analysis: deep nesting (max indentation depth 6) in k3/optim.py, dashboard.py, and train.py raises cognitive load; refactoring into smaller helper functions would improve readability.

The Bottom Line

smol‑kimi‑k3 provides a compact, reproducible implementation of Kimi K3’s novel attention mechanisms that can be trained on a modest GPU. The code is functional but lacks essential production hygiene (license, lockfile, CI). It is best suited for research prototyping and educational exploration rather than for deployment without further engineering.