Learning project: train and re-architect a small LLM entirely on a Jetson Orin Nano Super (8 GB unified memory) — capability, then three independent efficiency axes. Every notebook is fully self-contained, top-to-bottom runnable, and heavily annotated with the why behind every knob.
| Phase | Axis | What it does |
|---|---|---|
| 01 | Capability | Distill + align a personal assistant |
| 02 | Architectural efficiency | Bound the KV-cache via attention surgery (SWA/GLA) |
| 03 | Sparsity | Only activate a fraction of params per token (MoE) |
| 04 | Numerical + retrieval efficiency | Compress the KV-cache and give the assistant real memory (RAG) |
Distill a personal assistant from Qwen2.5-0.5B in three stages:
- Cognitive distillation — LoRA (r=32) trained on small slices of three teacher datasets (DeepSeek-R1 reasoning traces, code Q&A, chat structure).
- Personal alignment — Stage 1 is merged into the base, then a second LoRA (r=16) learns your handwritten assistant behaviors (
STAGE2_PAIRS). - Merge + sanity check — both adapters baked into one plain FP16 checkpoint:
final_jetson_assistant_0.5b/.
Key lessons: FP16 LoRA instead of QLoRA (no bitsandbytes on ARM64), unified-memory budgeting, gradient checkpointing + accumulation, sequential adapter stacking via merge_and_unload().
Trained artifacts (Hugging Face): stage1 adapter · stage2 adapter · merged model
Upcycle the Phase 1 output into a flat-KV-cache hybrid architecture:
- Load the Phase 1 model (falls back to base Qwen if you haven't run notebook 1).
- Replace every attention layer, alternating Sliding Window Attention (O(W) cache) and Gated Linear Attention (O(1) cache, learnable per-head decay, chunk-parallel recurrence) — warm-started from the trained Q/K/V/O projections.
- Short continual-pre-training (CPT) pass to stabilize the new attention layers, with checkpointing + resume so an interrupted run doesn't cost the whole pass.
- Save + convert to GGUF for llama.cpp / Ollama. GLA layers have no native llama.cpp kernel yet, so they fall back to sliding-window softmax attention at GGUF inference time — the conversion strips the GLA-only
log_decayparameters (no GGUF equivalent) before quantizing.
Key lessons: attention-module surgery on a Hugging Face model, linear-attention recurrence, why warm-starting preserves competence, GGUF quantization trade-offs, and why device_map="auto" can silently run layers on CPU under Jetson's unified-memory model.
Trained artifacts (Hugging Face): CPT adapter · hybrid model + GGUF
Two independent, contrasting approaches to sparse MoE, both device-agnostic (auto-detect cuda/mps/cpu, Jetson-specific tuning applied only when actually on a Jetson):
Franken_MoE/franken_moe_pipeline.ipynb — stitch pretrained specialists, no training
Combine Qwen2.5-0.5B-Instruct and Qwen2.5-Coder-0.5B into a sparse Qwen2MoeForCausalLM via mergekit-moe: clone each specialist's MLP blocks into expert slots, calibrate a router from domain prompts, quantize to GGUF Q4_K_M, benchmark and evaluate with llama.cpp.
Real hardware-driven trade-off, not a hypothetical one: mergekit's best router-calibration mode (gate_mode: "hidden", a real forward pass) OOM-killed on this Jetson's 8 GB unified memory — confirmed via dmesg. Switched to "cheap_embed" (embedding-average calibration, no forward pass) to actually complete, and mergekit itself flagged the result as degenerate ("ALL layers have degenerate routing parameters"). The notebook's own qualitative eval confirms it: the merged model's outputs are genuine gibberish. This is the honest result of a memory-constrained router calibration, not a bug — the mechanical pipeline (merge → quantize → build → benchmark) runs cleanly end-to-end and produces real numbers (~60 tok/s generation, ~900MB Q4_K_M). On a machine with more headroom, switching back to "hidden" is a one-line change and a straightforward quality upgrade.
Also fixed along the way: Q4_K_M isn't a direct conversion --outtype (needs llama-quantize, built from source, run as a second pass after an F16 conversion); Qwen2.5's tied embeddings need an explicit lm_head.weight materialized before GGUF conversion (llama.cpp's Qwen2Moe loader, unlike plain Qwen2's, has no tied-embedding fallback); nvcc often isn't on PATH in a Jupyter kernel's environment even when CUDA is installed; and llama-cli needs -st (single-turn), not just -no-cnv, or it hangs in an interactive loop when scripted. llama.cpp itself is built once into a shared, gitignored cache (.cache/llama_cpp_shared/) at the repo root and reused by both this notebook and notebook 02, instead of every notebook compiling its own copy.
Mini_MoE/mini_moe_from_scratch.ipynb — train a real MoE transformer from scratch
A ~45M-parameter decoder-only transformer, every component written out in PyTorch: causal self-attention with RoPE, a top-k router with a real Switch-Transformer-style auxiliary load-balancing loss (not just routing logic — the mechanism that actually prevents expert collapse), SwiGLU experts, mixed-precision training (torch.amp.autocast, bf16 on Ampere+) under a 3.5 GB VRAM budget. Trains a small byte-level BPE tokenizer (vocab_size=8192) on this project's own local dataset/reasoning//dataset/code/ text rather than inheriting Qwen's ~152k-token vocabulary, which alone would blow the parameter budget.
Where Franken-MoE's router is calibrated from prompts, this one's is genuinely learned end-to-end — slower and the base model is far less capable in absolute terms, but the routing behavior is real, not approximated. Ends with a routing-by-domain plot comparing code vs. reasoning prompts, reported honestly rather than assumed.
Implements Google Research's TurboQuant (arXiv:2504.19874, Zandieh/Daliri/Hadian/Mirrokni) from the paper's actual math — not a library import — and validates the reimplementation against the paper's own reported distortion numbers before applying it to anything. Random-rotate each vector so its coordinates become independent and identically distributed (the marginal of one coordinate of a uniform point on a sphere), then fit an optimal scalar quantizer to that known, fixed distribution — no calibration data, no per-tensor statistics.
kv_cache_quant.ipynb — compress notebook 02's KV-cache
Applies the MSE-optimal quantizer to the hybrid model's SWA-layer attention. A real finding surfaced along the way: notebook 02's SWAAttention/GLAAttention never actually wire up incremental past_key_values — every forward() recomputes over the full sequence, so there's no live cache buffer to swap for a compressed one yet. This notebook measures the honest thing instead: the real fidelity cost of quantized K/V (round-tripped through quantize/dequantize inline, at bit-widths 1 through 4 plus fractional 2.5/3.5 via split-coordinate allocation) against the theoretical memory savings a proper incremental cache would realize — 97.5% token agreement with the FP16 baseline at 4 bits, dropping sharply below that. Also checks the obvious follow-up question rather than assuming: is quantizing GLA's O(1) state worth it too? (No — it's already ~128x smaller than the SWA cache and doesn't grow with context length.)
vector_search_rag.ipynb — compressed retrieval, wired into the Phase 1 assistant
Implements the paper's two-stage TurboQuant_prod construction (MSE quantizer + a Quantized Johnson-Lindenstrauss residual correction) for unbiased inner-product estimation directly from the compressed representation — no dequantization needed to rank search results. Builds a compressed index over this project's own docs (README, progress reports, session history) with sentence-transformers/all-MiniLM-L6-v2, then wires retrieval into jetson-assistant-0.5b as RAG, so the tiny Phase 1 model can answer questions about its own training project instead of hallucinating. Reports recall@5 against exact brute-force search (80% mean at 4 bits) and is honest about the actual demo results: retrieval genuinely finds the right passages, but a 0.5B model doesn't reliably extract the correct fact from them even when handed the right context — RAG fixes "no information available," not "perfect reading comprehension." Point CORPUS_ROOT at your own notes to index something more personal.
JetPack 7.2 (CUDA 13.2) with CUDA-enabled PyTorch (torch==2.13.0, generic PyPI +cu130 wheel — runs on Orin via PTX JIT), plus transformers, datasets, peft, trl, accelerate. All training fits in 8 GB; the long pole is Phase 1 Stage 1 (~2–4 h). 03_MoE_arch/Franken_MoE additionally needs cmake + a C++ toolchain (to build llama.cpp) and pulls mergekit into its own isolated venv — see that notebook for why. 04_TurboQuant additionally needs sentence-transformers (installs cleanly into the shared venv — no version conflicts, unlike mergekit).