The Problem

Frontier MoE language models (GLM-5.2, DeepSeek) are documented in papers but nearly impossible to study locally—they require multi-GPU clusters, custom kernels, and distributed collectives. A conventional tiny Transformer doesn't answer architecture questions; it just shows a small decoder can learn TinyStories. This repo compresses the defining GLM-5.2 mechanisms (MLA, DSA, IndexShare, sparse MoE, shared MTP) into 50M parameters that fit on one consumer GPU, in plain PyTorch with no custom CUDA or frameworks obscuring the data flow.

What This Does

smol-glm-5.2 is a training-oriented miniature of the public GLM-5.2 architecture. The authoritative shape lives in configs/glm52_50m.json: 13 layers, layers 1–3 dense SwiGLU, layers 4–13 sparse MoE, with the FFFSSSFSSSFSS IndexShare schedule. The core machinery is in glm52/model.py (700 lines, 34 functions, 13 classes)—MLA with separate RoPE/NoPE channels, a trainable multi-head lightning indexer, sigmoid top-k routing with no auxiliary loss, and a parameter-shared multi-token predictor.

Supporting modules: glm52/optim.py implements Muon Split for matrices and AdamW for embeddings; glm52/monitoring.py handles atomic checkpoints and metrics; glm52/config.py defines GLM52Config with mlp_type and indexer_type variants. A live training dashboard (dashboard.py, dashboard/app.js) serves metrics over HTTP. Tests in tests/test_model.py cover forward/backward, BF16 autocast safety, and parameter budget.

How It Is Wired

Execution starts at main in dashboard.py:57, which reaches 42 functions. The import graph shows 12 internal modules with 11 edges and zero circular dependencies—clean structure. The hub is glm52/__init__.py (Ca 4, Ce 2, instability 0.33); glm52/model.py is the widest blast radius, with GLM52ForCausalLM called from 9 places and RMSNorm instantiated 9 times in __init__.

Tracing a run: main → save_checkpoint → path.parent.mkdir (filesystem write, 3 hops). main also calls load_state_dict (3×), save_checkpoint (3×), and set_status (4×). The model itself routes through forward → apply_interleaved_rope (4×) and forward_step (2× callers). glm52/monitoring.py owns all state persistence (_write_state called from 5 places); scripts/export_model.py does a sha256 for integrity. The dashboard do_GET handler routes to _send_bytes and _send_file (2× each).

How To Use It

git clone https://github.com/moses-y/smol-glm-5.2
cd smol-glm-5.2
pip install -r requirements.txt  # or: pip install -e .
python scripts/prepare_tinystories.py  # build tokenizer + dataset
python train.py --config configs/glm52_50m.json
python generate.py  # load trained model and sample
python dashboard.py  # live training dashboard at localhost

Configuration lives in configs/glm52_50m.json (full) and configs/glm52_smoke.json (smoke test). No environment variables are required. The README documents these commands verbatim.

Real-World Use

This fits as a teaching tool or architecture-validation harness. A researcher studying sparse MoE routing can modify glm52/model.py's _indexer_distillation_loss, run tests/test_model.py to verify gradient flow, then train on TinyStories to observe whether the indexer learns meaningful token selection. The dashboard gives live visibility into metrics without external tooling. For production, the export script (scripts/export_model.py) produces a checkpoint with SHA-256 verification for artifact integrity.

Code Health & Issues

Static analysis (deterministic, not opinion) found 7 medium findings across 3 kinds:

  • Medium - Deep nesting (×5) in glm52/monitoring.py, glm52/model.py, glm52/optim.py; max indentation depth 6. Fix: early returns/guard clauses.
  • Medium - Oversized file glm52/model.py at 700 lines. Fix: split by responsibility.
  • Medium - High branching density in dashboard/app.js; 33 branch points over 92 lines. Fix: table/strategy dispatch.

SDLC gaps from structure: High - no LICENSE (default all-rights-reserved blocks reuse); High - no lockfile (non-reproducible builds, torch@2.5 range reaches CVE-2025-32434); High - no CI (12 source files, no build/test gate); Medium - no Dependabot (2 manifests, no update bot); Medium - unpinned resolution (declared ranges reach critical advisories). datasets is 2 major versions behind (3.0 vs 5.0.1), pytest 1 behind (8.0 vs 9.1.1).

The Bottom Line

A genuinely readable architecture study—plain PyTorch, no framework obscuring the GLM-5.2 mechanisms, clean module graph with no cycles. The trade-off is operational: no CI, no lockfile, no license, and model.py needs splitting. Use it to learn or teach sparse MoE architecture; don't ship it without addressing the SDLC gaps.