Initial commit: GLM-5.2 Baby from scratch with WSD schedule

.gitignore

@@ -0,0 +1,32 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Project specific ignores
+# Ignore the large dataset binaries
+data/*.bin
+data_phase3/*.bin
+
+# Ignore checkpoints
+out_glm5/*.pt
+*.pt
+
+# IDE / Editor
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# VSCode workspace / agents
+.agents/

README.md

@@ -0,0 +1,61 @@
+# 🍼 GLM-5.2 Baby (120M) - From Scratch
+
+This repository contains a from-scratch implementation and pretraining script for a "baby" version (~120M parameters) of **GLM-5.2** (GLM MoE DSA). It is heavily inspired by Andrej Karpathy's `nanoGPT` and aims to be highly educational.
+
+The model is small enough to train on a single consumer laptop GPU (e.g., RTX 4050 6GB VRAM) but includes all the cutting-edge architectural innovations of modern frontier models.
+
+## ✨ Architectural Features Implemented
+
+This isn't just a standard Transformer. It implements three major innovations from recent frontier models (like DeepSeek-V3 and GLM-5):
+
+1. **MLA (Multi-Latent Attention):** Compresses the attention mechanism using LoRA-style projections to drastically save VRAM during training and inference.
+2. **DSA (DeepSeek Sparse Attention):** Selects only the most relevant tokens to attend to via a learned indexer, rather than attending to the entire context uniformly.
+3. **MoE (Mixture of Experts):** Uses a fine-grained sigmoid-routed mixture of experts alongside a shared expert, activating only a subset of parameters per token.
+
+## 🚀 Training Features
+
+The training loop (`train_glm5.py`) is highly optimized for limited hardware (6GB VRAM) while maximizing throughput (~4,900 tokens/sec on an RTX 4050):
+- **Mixed Precision:** Uses `bfloat16` and TF32 Tensor Cores.
+- **Gradient Checkpointing:** Recomputes forward passes during backprop to save ~40% VRAM.
+- **Gradient Accumulation:** Achieves large effective batch sizes on a single GPU.
+- **WSD (Warmup-Stable-Decay) Learning Rate Schedule:** 
+  Supports multi-phase training by holding the learning rate at peak for a "stable" exploration phase before initiating a steep cosine decay. (Controlled via `--stable_iters`).
+
+## 📚 Educational Guide
+
+If you are new to LLM pretraining, learning rates, loss curves, and scaling laws, check out the included beginner guide:
+👉 **[LLM Training Guide for Beginners](llm_training_guide.md)**
+
+## 🛠️ Usage
+
+### Installation
+```bash
+pip install -r requirements.txt
+```
+
+### Training
+To train the model on a single GPU with the WSD schedule (holding LR stable for 217,000 steps):
+
+```bash
+python train_glm5.py \
+    --data_dir ./data \
+    --batch_size 6 \
+    --gradient_accumulation_steps 3 \
+    --max_iters 110000 \
+    --lr_decay_iters 260000 \
+    --warmup_iters 1500 \
+    --stable_iters 217000 \
+    --eval_interval 2000 \
+    --eval_iters 200 \
+    --log_interval 100
+```
+
+### Generation / Sampling
+To sample text from your best trained checkpoint:
+
+```bash
+python train_glm5.py --eval_only --ckpt out_glm5/ckpt_best.pt --prompt "The future of AI is"
+```
+
+## ⚖️ License
+MIT License

data/meta.json

@@ -0,0 +1,16 @@
+{
+  "tokenizer": "gpt2",
+  "vocab_size": 50257,
+  "eot_token": 50256,
+  "dtype": "uint16",
+  "train_tokens": 1979887628,
+  "val_tokens": 19913769,
+  "total_tokens": 1999801397,
+  "sources": [
+    "fineweb_edu",
+    "wikipedia",
+    "gutenberg",
+    "starcoder"
+  ],
+  "val_ratio": 0.01
+}

dataprep_phase3.py

@@ -0,0 +1,188 @@
+"""
+Pretraining Data Pipeline -- Phase 3 (0.4B Tokens)
+===================================================
+
+Streams 400 Million tokens of domain-diverse text:
+  - Wikipedia (wikimedia/wikipedia)  -- 200M tokens (50%)
+  - Gutenberg (emozilla/pg19)        -- 120M tokens (30%)
+  - Python Code (bigcode/the-stack)  --  80M tokens (20%)
+
+Outputs to `./data_phase3` to avoid touching active training in `./data`.
+
+Usage:
+  python dataprep_phase3.py
+"""
+
+import argparse
+import json
+import os
+import time
+
+import numpy as np
+import tiktoken
+from tqdm import tqdm
+
+DEFAULT_OUTPUT_DIR = "./data_phase3"
+DEFAULT_TOKENIZER = "gpt2"
+DEFAULT_TOTAL_TOKENS = 400_000_000  # 0.4B tokens
+DEFAULT_VAL_RATIO = 0.01
+DEFAULT_WRITE_CHUNK = 1_000_000
+
+SOURCE_ALLOCATIONS = {
+    "wikipedia": 0.50,  # 200M tokens
+    "gutenberg": 0.30,  # 120M tokens
+    "starcoder": 0.20,  # 80M tokens
+}
+
+GPT2_EOT = 50256
+
+
+def iter_wikipedia():
+    from datasets import load_dataset
+    ds = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)
+    for example in ds:
+        text = example.get("text", "")
+        if text:
+            yield text
+
+
+def iter_gutenberg():
+    from datasets import load_dataset
+    ds = load_dataset("emozilla/pg19", split="train", streaming=True)
+    for example in ds:
+        text = example.get("text", "")
+        if text:
+            yield text
+
+
+def iter_starcoder():
+    from datasets import load_dataset
+    try:
+        ds = load_dataset("bigcode/the-stack-smol", data_dir="data/python", split="train", streaming=True)
+        for example in ds:
+            text = example.get("content", "")
+            if text:
+                yield text
+    except Exception:
+        ds = load_dataset("codeparrot/codeparrot-clean", split="train", streaming=True)
+        for example in ds:
+            text = example.get("content", "")
+            if text:
+                yield text
+
+
+SOURCE_ITERATORS = {
+    "wikipedia": iter_wikipedia,
+    "gutenberg": iter_gutenberg,
+    "starcoder": iter_starcoder,
+}
+
+
+def tokenize_and_write(source_name, text_iterator, target_tokens, encoder, eot_token, train_file, val_file, val_ratio, write_chunk, dtype):
+    train_buffer, val_buffer = [], []
+    train_total, val_total = 0, 0
+    doc_count = 0
+    rng = np.random.default_rng(seed=42 + hash(source_name) % 10000)
+
+    pbar = tqdm(
+        total=target_tokens,
+        unit="tok",
+        unit_scale=True,
+        desc=f"  {source_name}",
+        bar_format="  {desc}: {percentage:3.0f}% |{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]",
+    )
+
+    for text in text_iterator:
+        tokens = encoder.encode_ordinary(text)
+        tokens.append(eot_token)
+
+        is_val = rng.random() < val_ratio
+
+        if is_val:
+            val_buffer.extend(tokens)
+            val_total += len(tokens)
+            if len(val_buffer) >= write_chunk:
+                val_file.write(np.array(val_buffer, dtype=dtype).tobytes())
+                val_buffer = []
+        else:
+            train_buffer.extend(tokens)
+            train_total += len(tokens)
+            if len(train_buffer) >= write_chunk:
+                train_file.write(np.array(train_buffer, dtype=dtype).tobytes())
+                train_buffer = []
+
+        doc_count += 1
+        pbar.update(len(tokens))
+
+        if train_total + val_total >= target_tokens:
+            break
+
+    if train_buffer:
+        train_file.write(np.array(train_buffer, dtype=dtype).tobytes())
+    if val_buffer:
+        val_file.write(np.array(val_buffer, dtype=dtype).tobytes())
+
+    pbar.close()
+    print(f"    {source_name}: {doc_count:,} docs | train {train_total:,} + val {val_total:,} tokens")
+    return train_total, val_total
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Pretraining Data Pipeline -- Phase 3")
+    parser.add_argument("--output_dir", type=str, default=DEFAULT_OUTPUT_DIR)
+    parser.add_argument("--total_tokens", type=int, default=DEFAULT_TOTAL_TOKENS)
+    args = parser.parse_args()
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    train_path = os.path.join(args.output_dir, "train.bin")
+    val_path = os.path.join(args.output_dir, "val.bin")
+    meta_path = os.path.join(args.output_dir, "meta.json")
+
+    encoder = tiktoken.get_encoding(DEFAULT_TOKENIZER)
+    vocab_size = encoder.n_vocab
+    dtype = np.uint16 if vocab_size <= 65535 else np.uint32
+
+    source_targets = {k: int(args.total_tokens * v) for k, v in SOURCE_ALLOCATIONS.items()}
+
+    print(f"\n{'='*70}")
+    print(f"  Phase 3 Data Pipeline (0.4B Tokens Target)")
+    print(f"{'='*70}")
+    print(f"  Output Dir: {os.path.abspath(args.output_dir)}")
+    for name, target in source_targets.items():
+        print(f"    {name:15s}: {target:,} tokens")
+    print(f"{'='*70}\n")
+
+    total_train, total_val = 0, 0
+    t0 = time.time()
+
+    with open(train_path, "wb") as train_file, open(val_path, "wb") as val_file:
+        for source_name, target in source_targets.items():
+            print(f"\n  [{source_name}] Streaming {target:,} tokens...")
+            try:
+                tr, va = tokenize_and_write(
+                    source_name, SOURCE_ITERATORS[source_name](), target,
+                    encoder, GPT2_EOT, train_file, val_file, DEFAULT_VAL_RATIO,
+                    DEFAULT_WRITE_CHUNK, dtype
+                )
+                total_train += tr
+                total_val += va
+            except Exception as e:
+                print(f"  [WARN] Error on {source_name}: {e}")
+
+    meta = {
+        "tokenizer": DEFAULT_TOKENIZER,
+        "vocab_size": vocab_size,
+        "dtype": "uint16" if dtype == np.uint16 else "uint32",
+        "train_tokens": total_train,
+        "val_tokens": total_val,
+        "total_tokens": total_train + total_val,
+        "sources": list(source_targets.keys()),
+    }
+    with open(meta_path, "w") as f:
+        json.dump(meta, f, indent=2)
+
+    print(f"\n  [DONE] Phase 3 data ready at {args.output_dir}! Total: {total_train + total_val:,} tokens")
+
+
+if __name__ == "__main__":
+    main()

dataprep_pretrain.py

@@ -0,0 +1,385 @@
+"""
+Pretraining Data Pipeline for Baby GLM-5.2
+============================================
+
+Streams, tokenizes, and writes 3.3B tokens to binary files for pretraining.
+
+Data Sources:
+  1. FineWeb-Edu       -- 2.0B tokens (high-quality web text)
+  2. Wikipedia EN      -- 0.7B tokens (encyclopedic knowledge)
+  3. Project Gutenberg -- 0.4B tokens (literary text)
+  4. StarCoder Python  -- 0.2B tokens (code, optional)
+
+Output:
+  data/train.bin   -- ~3.267B tokens, binary uint16 memmap
+  data/val.bin     -- ~0.033B tokens, binary uint16 memmap
+  data/meta.json   -- tokenizer info, vocab size, token counts
+
+Usage:
+  python dataprep_pretrain.py                    # Full 3.3B token run
+  python dataprep_pretrain.py --total_tokens 10000000  # Quick 10M test
+  python dataprep_pretrain.py --no_code          # Skip code data
+
+Estimated runtime: 2-4 hours on CPU with good internet (full run)
+Estimated disk: ~7 GB for output .bin files
+"""
+
+import argparse
+import json
+import os
+import time
+
+import numpy as np
+import tiktoken
+from tqdm import tqdm
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+DEFAULT_OUTPUT_DIR = "./data"
+DEFAULT_TOKENIZER = "gpt2"
+DEFAULT_TOTAL_TOKENS = 3_300_000_000  # 3.3B tokens (Chinchilla+ for 120M params)
+DEFAULT_VAL_RATIO = 0.01  # 1% validation split
+DEFAULT_WRITE_CHUNK = 1_000_000  # Flush to disk every 1M tokens
+
+# Data source allocation (fraction of total tokens)
+SOURCE_ALLOCATIONS = {
+    "fineweb_edu": 0.606,   # ~2.0B tokens
+    "wikipedia": 0.212,     # ~0.7B tokens
+    "gutenberg": 0.121,     # ~0.4B tokens
+    "starcoder": 0.061,     # ~0.2B tokens (optional)
+}
+
+# End-of-text token for GPT-2 tokenizer
+GPT2_EOT = 50256
+
+
+# =============================================================================
+# Dataset Iterators
+# =============================================================================
+# Each iterator yields raw text strings from a HuggingFace dataset stream.
+# Streaming means we never hold the full dataset in memory.
+# =============================================================================
+
+
+def iter_fineweb_edu():
+    """Stream high-quality educational web text from FineWeb-Edu."""
+    from datasets import load_dataset
+    ds = load_dataset(
+        "HuggingFaceFW/fineweb-edu",
+        name="sample-10BT",
+        split="train",
+        streaming=True,
+    )
+    for example in ds:
+        text = example.get("text", "")
+        if text:
+            yield text
+
+
+def iter_wikipedia():
+    """Stream English Wikipedia articles (script-free wikimedia version)."""
+    from datasets import load_dataset
+    ds = load_dataset(
+        "wikimedia/wikipedia",
+        "20231101.en",
+        split="train",
+        streaming=True,
+    )
+    for example in ds:
+        text = example.get("text", "")
+        if text:
+            yield text
+
+
+def iter_gutenberg():
+    """Stream public domain books from Project Gutenberg (PG-19 subset)."""
+    from datasets import load_dataset
+    ds = load_dataset(
+        "emozilla/pg19",
+        split="train",
+        streaming=True,
+    )
+    for example in ds:
+        # pg19 uses "text" for the full book text, with "short_book_title" as metadata
+        text = example.get("text", "")
+        if text:
+            yield text
+
+
+def iter_starcoder():
+    """
+    Stream Python code from The Stack v2 (non-gated subset).
+
+    If you want the full StarCoder dataset instead, first authenticate:
+        huggingface-cli login
+    Then change 'bigcode/the-stack-v2-train-smol-ids' below to 'bigcode/starcoderdata'.
+    """
+    from datasets import load_dataset
+    try:
+        # Try the non-gated smol subset first
+        ds = load_dataset(
+            "bigcode/the-stack-smol",
+            data_dir="data/python",
+            split="train",
+            streaming=True,
+        )
+        for example in ds:
+            text = example.get("content", "")
+            if text:
+                yield text
+    except Exception:
+        # Fallback: use codeparrot's cleaned Python dataset
+        ds = load_dataset(
+            "codeparrot/codeparrot-clean",
+            split="train",
+            streaming=True,
+        )
+        for example in ds:
+            text = example.get("content", "")
+            if text:
+                yield text
+
+
+# Map source names to their iterators
+SOURCE_ITERATORS = {
+    "fineweb_edu": iter_fineweb_edu,
+    "wikipedia": iter_wikipedia,
+    "gutenberg": iter_gutenberg,
+    "starcoder": iter_starcoder,
+}
+
+
+# =============================================================================
+# Tokenization & Writing
+# =============================================================================
+
+
+def tokenize_and_write(
+    source_name,
+    text_iterator,
+    target_tokens,
+    encoder,
+    eot_token,
+    train_file,
+    val_file,
+    val_ratio,
+    write_chunk,
+    dtype,
+):
+    """
+    Tokenize text from an iterator and write tokens to train/val binary files.
+
+    Each document is separated by an EOT token. Documents are randomly assigned
+    to val split with probability val_ratio.
+
+    Args:
+        source_name: Name of the data source (for logging)
+        text_iterator: Iterator yielding text strings
+        target_tokens: Number of tokens to collect from this source
+        encoder: tiktoken encoder
+        eot_token: End-of-text token ID
+        train_file: Open file handle for train.bin
+        val_file: Open file handle for val.bin
+        val_ratio: Fraction of documents for validation
+        write_chunk: Buffer size before flushing to disk
+        dtype: numpy dtype for token storage (uint16 or uint32)
+
+    Returns:
+        (train_tokens_written, val_tokens_written)
+    """
+    train_buffer = []
+    val_buffer = []
+    train_total = 0
+    val_total = 0
+    doc_count = 0
+
+    rng = np.random.default_rng(seed=42 + hash(source_name) % 10000)
+
+    pbar = tqdm(
+        total=target_tokens,
+        unit="tok",
+        unit_scale=True,
+        desc=f"  {source_name}",
+        bar_format="  {desc}: {percentage:3.0f}% |{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]",
+    )
+
+    for text in text_iterator:
+        # Tokenize the document
+        tokens = encoder.encode_ordinary(text)
+        tokens.append(eot_token)  # Separate documents with EOT
+
+        # Randomly assign entire documents to train or val
+        is_val = rng.random() < val_ratio
+
+        if is_val:
+            val_buffer.extend(tokens)
+            val_total += len(tokens)
+            # Flush val buffer
+            if len(val_buffer) >= write_chunk:
+                val_file.write(np.array(val_buffer, dtype=dtype).tobytes())
+                val_buffer = []
+        else:
+            train_buffer.extend(tokens)
+            train_total += len(tokens)
+            # Flush train buffer
+            if len(train_buffer) >= write_chunk:
+                train_file.write(np.array(train_buffer, dtype=dtype).tobytes())
+                train_buffer = []
+
+        doc_count += 1
+        pbar.update(len(tokens))
+
+        # Check if we've reached the target
+        if train_total + val_total >= target_tokens:
+            break
+
+    # Flush remaining buffers
+    if train_buffer:
+        train_file.write(np.array(train_buffer, dtype=dtype).tobytes())
+    if val_buffer:
+        val_file.write(np.array(val_buffer, dtype=dtype).tobytes())
+
+    pbar.close()
+    print(f"    {source_name}: {doc_count:,} docs | "
+          f"train {train_total:,} + val {val_total:,} = {train_total + val_total:,} tokens")
+
+    return train_total, val_total
+
+
+# =============================================================================
+# Main Pipeline
+# =============================================================================
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Pretraining Data Pipeline for Baby GLM-5.2",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--output_dir", type=str, default=DEFAULT_OUTPUT_DIR,
+                        help="Output directory for .bin and .json files")
+    parser.add_argument("--total_tokens", type=int, default=DEFAULT_TOTAL_TOKENS,
+                        help="Total tokens to collect across all sources")
+    parser.add_argument("--val_ratio", type=float, default=DEFAULT_VAL_RATIO,
+                        help="Fraction of documents for validation split")
+    parser.add_argument("--write_chunk", type=int, default=DEFAULT_WRITE_CHUNK,
+                        help="Buffer size (tokens) before flushing to disk")
+    parser.add_argument("--no_code", action="store_true",
+                        help="Exclude code data (StarCoder)")
+    args = parser.parse_args()
+
+    # --- Setup ---
+    os.makedirs(args.output_dir, exist_ok=True)
+    train_path = os.path.join(args.output_dir, "train.bin")
+    val_path = os.path.join(args.output_dir, "val.bin")
+    meta_path = os.path.join(args.output_dir, "meta.json")
+
+    # --- Tokenizer ---
+    encoder = tiktoken.get_encoding(DEFAULT_TOKENIZER)
+    vocab_size = encoder.n_vocab  # 50257 for GPT-2
+    eot_token = GPT2_EOT
+
+    # Determine dtype: uint16 if vocab fits, uint32 otherwise
+    if vocab_size <= 65535:
+        dtype = np.uint16
+        dtype_str = "uint16"
+    else:
+        dtype = np.uint32
+        dtype_str = "uint32"
+
+    # --- Compute per-source token targets ---
+    include_code = not args.no_code
+    active_sources = {k: v for k, v in SOURCE_ALLOCATIONS.items()
+                      if k != "starcoder" or include_code}
+
+    # Renormalize allocations if code is excluded
+    total_alloc = sum(active_sources.values())
+    source_targets = {
+        k: int(args.total_tokens * v / total_alloc)
+        for k, v in active_sources.items()
+    }
+
+    # --- Print Plan ---
+    print(f"\n{'='*70}")
+    print(f"  Pretraining Data Pipeline for Baby GLM-5.2")
+    print(f"{'='*70}")
+    print(f"  Tokenizer:     {DEFAULT_TOKENIZER} (vocab_size={vocab_size})")
+    print(f"  Token dtype:   {dtype_str}")
+    print(f"  Total target:  {args.total_tokens:,} tokens")
+    print(f"  Val ratio:     {args.val_ratio:.1%}")
+    print(f"  Output dir:    {os.path.abspath(args.output_dir)}")
+    print(f"  Code data:     {'Yes' if include_code else 'No'}")
+    print(f"\n  Source Allocation:")
+    for name, target in source_targets.items():
+        print(f"    {name:20s}  {target:>14,} tokens  ({target/args.total_tokens:.1%})")
+    print(f"{'='*70}\n")
+
+    # --- Process Each Source ---
+    total_train = 0
+    total_val = 0
+    t0 = time.time()
+
+    with open(train_path, "wb") as train_file, open(val_path, "wb") as val_file:
+        for source_name, target in source_targets.items():
+            print(f"\n  [{source_name}] Streaming {target:,} tokens...")
+            iterator_fn = SOURCE_ITERATORS[source_name]
+
+            try:
+                train_written, val_written = tokenize_and_write(
+                    source_name=source_name,
+                    text_iterator=iterator_fn(),
+                    target_tokens=target,
+                    encoder=encoder,
+                    eot_token=eot_token,
+                    train_file=train_file,
+                    val_file=val_file,
+                    val_ratio=args.val_ratio,
+                    write_chunk=args.write_chunk,
+                    dtype=dtype,
+                )
+                total_train += train_written
+                total_val += val_written
+            except Exception as e:
+                print(f"  [WARN] Error streaming {source_name}: {e}")
+                print(f"         Skipping this source and continuing...")
+                continue
+
+    elapsed = time.time() - t0
+
+    # --- Write Metadata ---
+    meta = {
+        "tokenizer": DEFAULT_TOKENIZER,
+        "vocab_size": vocab_size,
+        "eot_token": eot_token,
+        "dtype": dtype_str,
+        "train_tokens": total_train,
+        "val_tokens": total_val,
+        "total_tokens": total_train + total_val,
+        "sources": list(source_targets.keys()),
+        "val_ratio": args.val_ratio,
+    }
+    with open(meta_path, "w") as f:
+        json.dump(meta, f, indent=2)
+
+    # --- Summary ---
+    train_size_gb = os.path.getsize(train_path) / 1e9
+    val_size_gb = os.path.getsize(val_path) / 1e9
+
+    print(f"\n{'='*70}")
+    print(f"  [DONE] Data preparation complete!")
+    print(f"{'='*70}")
+    print(f"  Time elapsed:    {elapsed/3600:.1f} hours ({elapsed:.0f}s)")
+    print(f"  Train tokens:    {total_train:,}")
+    print(f"  Val tokens:      {total_val:,}")
+    print(f"  Total tokens:    {total_train + total_val:,}")
+    print(f"  train.bin:       {train_size_gb:.2f} GB")
+    print(f"  val.bin:         {val_size_gb:.2f} GB")
+    print(f"  meta.json:       {meta_path}")
+    print(f"\n  Next step: python train_glm5.py --data_dir {args.output_dir}")
+    print(f"{'='*70}\n")
+
+
+if __name__ == "__main__":
+    main()

llm_training_guide.md

@@ -0,0 +1,251 @@
+# 🧠 The Beginner's Guide to Understanding Your LLM Training Run
+
+*Tailored for your GLM-5.2 Baby 120M pretraining on RTX 4050*
+
+---
+
+## Part 1: What Do All These Numbers Mean?
+
+When you see a log line like this:
+
+```
+step 52000 | train 4.1230 | val 4.1071 | lr 5.51e-04
+```
+
+Here's what each piece tells you:
+
+### 📉 Loss (train & val)
+
+**Loss** = how wrong the model is. Lower = better.
+
+Your model is trying to predict the next word (token) in a sequence. The loss measures how surprised the model is by the correct answer. Think of it like a quiz score — but inverted: **4.10 is better than 4.30**.
+
+| Term | What it means | Your value |
+|------|--------------|------------|
+| **Training loss** | Error on data the model is actively learning from | `4.1230` |
+| **Validation loss** | Error on data the model has **never seen** (the real test) | `4.1071` |
+
+> [!IMPORTANT]
+> **Always trust validation loss over training loss.** Training loss can go down just because the model memorizes data. Validation loss tells you if it's actually *learning patterns*.
+
+#### What's "good" for a 120M model?
+
+- **Loss > 5.0** → The model is basically guessing randomly
+- **Loss ~4.0–4.5** → It's learning basic grammar and common patterns ← **You are here**
+- **Loss ~3.5–4.0** → It understands sentence structure well
+- **Loss ~3.0–3.5** → Coherent paragraphs, factual fragments
+- **Loss < 3.0** → Very strong for this model size (hard to achieve with 120M params)
+
+### 📐 Learning Rate (lr)
+
+```
+lr 5.51e-04  →  this means 0.000551
+```
+
+**Learning rate** = how big of a step the model takes when adjusting its weights after each batch.
+
+- **Too high** → Model overshoots and loss spikes or goes to NaN (crashes)
+- **Too low** → Model barely changes and training takes forever
+- **Just right** → Steady decrease in loss
+
+Your training uses a **cosine schedule with warmup**. Here's what that looks like:
+
+```
+Learning Rate over time:
+
+ 6e-4 |    /‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾\
+      |   /                              \
+      |  /                                \
+      | /                                  \
+ 6e-5 |/                                    \___
+      +----+--------+--------+--------+--------+
+      0   1.5k    65k     130k     195k     260k
+       warmup              steps →
+
+      Phase: [WARM] [——— COSINE DECAY ———————]  [MIN]
+```
+
+**What this means for YOUR run:**
+- Steps 0–1,500: LR ramped up from 0 → `6e-4` (warmup)
+- Steps 1,500–260,000: LR slowly decays from `6e-4` → `6e-5` following a cosine curve
+- At step 52,000 you're at `5.51e-04` — still very close to the peak!
+- **The big improvements come later** when the LR drops significantly (after ~130k steps)
+
+### ⚡ Tokens per Second (tok/s)
+
+```
+4,259 tok/s
+```
+
+This is your training speed — how many tokens (≈words) the model processes per second. Higher = faster training. Your RTX 4050 is pushing ~4,000-5,000 tok/s, which is solid for a laptop GPU.
+
+### 💾 VRAM
+
+```
+VRAM 4.78GB
+```
+
+How much GPU memory you're using. Your GPU has 6GB total, so 4.78GB means you have a ~1.2GB safety buffer. If this ever hits 6GB → Out of Memory crash (OOM).
+
+---
+
+## Part 2: How to Tell if Training is Going Well
+
+### ✅ Signs of Healthy Training
+
+1. **Val loss is trending down over thousands of steps** (not every single eval, but the overall trend)
+2. **Train loss and val loss are close together** (no big gap)
+3. **No NaN or Inf in loss** (that would mean the training exploded)
+4. **No sudden loss spikes** that don't recover
+
+### ⚠️ Warning Signs
+
+| Symptom | What it means | What to do |
+|---------|--------------|------------|
+| Val loss goes **up** while train loss goes **down** | **Overfitting** — memorizing instead of learning | Stop training, use more data, or add regularization |
+| Val loss **flat for 20,000+ steps** | **Plateau** — model may be stuck | Often resolves when LR decays; be patient |
+| Loss suddenly shoots to **100+** or **NaN** | **Training instability** | Reduce learning rate, check data for corruption |
+| Train loss **very noisy** (swings of ±1.0) | **Batch size too small** | Increase `gradient_accumulation_steps` |
+
+### 📊 Your Run's Health Check
+
+```
+Val Loss Trajectory:
+  Step 20k: 4.31 ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓░░░
+  Step 24k: 4.27 ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓░░░░
+  Step 28k: 4.22 ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓░░░░░
+  Step 30k: 4.19 ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓░░░░░░
+  Step 36k: 4.16 ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓░░░░░░░
+  Step 40k: 4.12 ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓░░░░░░░░  ← plateau started here
+  Step 50k: 4.18 ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓░░░░░��░  ← bouncing around
+  Step 52k: 4.11 ▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓░░░░░░░░  ← NEW BEST! Plateau broken 🎉
+```
+
+**Verdict: Your training is healthy.** The 10k-step plateau (40k–50k) was normal — it just broke through at step 52k with a new best of `4.1071`.
+
+---
+
+## Part 3: Key Concepts You Should Know
+
+### 🎯 Chinchilla Scaling Laws (The "20:1 Rule")
+
+DeepMind discovered in 2022 that for optimal training, you should use roughly **20 tokens of data per parameter**.
+
+**Your model:**
+- Parameters: ~120M
+- Chinchilla-optimal data: 120M × 20 = **2.4B tokens**
+- Your planned data: **2.4B tokens** ✅ (2.0B Phase 2 + 0.4B Phase 3)
+
+You're right on target! This means your model should be trained to near its full potential by the time you finish.
+
+### 🔄 Gradient Accumulation (Why `batch_size 6 × grad_accum 3`)
+
+Your GPU can only fit 6 sequences in VRAM at once. But training works better with larger "effective" batches (more stable gradients). **Gradient accumulation** is a trick:
+
+```
+Step 1: Process 6 sequences → compute gradients (DON'T update weights yet)
+Step 2: Process 6 more sequences → add gradients to step 1's
+Step 3: Process 6 more sequences → add gradients again
+→ NOW update weights using all 18 sequences' worth of gradients!
+```
+
+Effective batch = `6 × 3 = 18 sequences × 512 tokens = 9,216 tokens per weight update`
+
+### 🧊 Why Mixed Precision (bfloat16) Matters
+
+Normally, numbers in the model use 32 bits (float32). **bfloat16** uses only 16 bits:
+- **Pro:** Uses ~half the VRAM, runs ~2× faster on Tensor Cores
+- **Con:** Slightly less precise math
+- **Net result:** Massive win. This is why you can fit a 120M model in 6GB
+
+### 🏔️ Gradient Checkpointing
+
+Normally, the GPU stores ALL intermediate calculations during the forward pass (to use during backpropagation). With gradient checkpointing:
+- GPU **throws away** intermediate results to save VRAM
+- During backpropagation, it **recomputes** them on the fly
+- **Trade:** ~30% slower training, but ~40% less VRAM
+
+This is enabled in your run and is essential for fitting in 6GB.
+
+---
+
+## Part 4: Your Cosine Schedule — Why Patience Pays Off
+
+Here's something critical to understand about your specific run:
+
+```
+Your LR decay spans 260,000 steps, but this run goes to 110,000.
+
+At step 52,000:
+  - You're 20% through the cosine schedule
+  - LR has only dropped from 6.00e-4 → 5.51e-4 (an 8% decrease)
+  - The model is still taking BIG learning steps
+
+At step 110,000 (end of this run):
+  - You'll be 42% through the cosine schedule
+  - LR will be around ~4.4e-4
+  - Model will be learning at a moderate pace
+
+The BIGGEST gains happen in Phase 3 (steps 110k→260k):
+  - LR drops dramatically from 4.4e-4 → 6e-5
+  - This is where the model "settles in" and polishes its knowledge
+  - Think of it as: early training = rough sketch, late training = fine details
+```
+
+> [!TIP]
+> **This is why you saw a plateau.** The LR is still high enough that the model is "bouncing around" in the loss landscape. As LR decreases, these oscillations will shrink and the loss will drop more smoothly.
+
+---
+
+## Part 5: 📚 Must-Read Resources (Ordered by Difficulty)
+
+### 🟢 Beginner — Start Here
+
+| # | Resource | What You'll Learn | Link |
+|---|----------|-------------------|------|
+| 1 | **Karpathy: "Let's build GPT: from scratch"** (YouTube, ~2hrs) | How transformers work, attention, the training loop — all coded live | [YouTube](https://www.youtube.com/watch?v=kCc8FmEb1nY) |
+| 2 | **Karpathy: "A Recipe for Training Neural Networks"** (Blog) | THE guide for debugging training. When loss is weird, read this. | [karpathy.github.io](https://karpathy.github.io/2019/04/25/recipe/) |
+| 3 | **3Blue1Brown: "But what is a neural network?"** (YouTube) | Visual intuition for how neural nets learn | [YouTube](https://www.youtube.com/watch?v=aircAruvnKk) |
+
+### 🟡 Intermediate — After You've Done the Above
+
+| # | Resource | What You'll Learn | Link |
+|---|----------|-------------------|------|
+| 4 | **Karpathy: "Let's reproduce GPT-2 (124M)"** (YouTube, ~4hrs) | EXACTLY what you're doing — pretraining a GPT from scratch, optimized | [YouTube](https://www.youtube.com/watch?v=l8pRSuU81PU) |
+| 5 | **Sebastian Raschka: "Build a Large Language Model (From Scratch)"** (Book + GitHub) | Full pipeline: tokenization → training → finetuning, with code | [GitHub](https://github.com/rasbt/LLMs-from-scratch) |
+| 6 | **Chinchilla Paper (Summary)** | Why 20 tokens/param matters, scaling laws | Search: "Chinchilla scaling laws explained" |
+
+### 🔴 Advanced — When You Want to Go Deeper
+
+| # | Resource | What You'll Learn | Link |
+|---|----------|-------------------|------|
+| 7 | **EleutherAI: LM Evaluation Harness** | How to properly benchmark your model beyond just val loss | [GitHub](https://github.com/EleutherAI/lm-evaluation-harness) |
+| 8 | **DeepSeek-V3 Technical Report** | The MLA + MoE architecture your model is based on | [arXiv](https://arxiv.org/abs/2412.19437) |
+| 9 | **Sebastian Raschka: "Ahead of AI" newsletter** | Weekly updates on LLM research and training techniques | [Substack](https://magazine.sebastianraschka.com/) |
+
+---
+
+## Part 6: Quick Glossary
+
+| Term | Plain English |
+|------|--------------|
+| **Token** | A piece of a word. "training" → ["train", "ing"]. Your model sees tokens, not words. |
+| **Epoch** | One full pass through all training data. You won't complete 1 full epoch in this run. |
+| **Perplexity** | `e^loss` — another way to express loss. Your val loss 4.11 = perplexity ~61. Means "the model is choosing between ~61 equally likely next tokens." |
+| **Overfitting** | Model memorizes training data instead of learning general patterns. Val loss goes up while train loss keeps going down. |
+| **Underfitting** | Model hasn't learned enough yet. Both losses are still high. |
+| **Cosine decay** | LR schedule that follows a cosine curve from high → low. Industry standard. |
+| **AdamW** | The optimizer (the algorithm that updates weights). It's what everyone uses for LLMs. |
+| **Gradient clipping** | Caps the size of gradient updates to prevent explosions. Your clip = 1.0. |
+| **MoE** | Mixture of Experts — only some "expert" sub-networks activate per token, making the model bigger without proportionally more compute. |
+| **MLA** | Multi-Latent Attention — compresses the attention mechanism using LoRA-style projections to save VRAM. |
+| **DSA** | DeepSeek Sparse Attention — selects only the most relevant tokens to attend to, instead of attending to everything. |
+
+---
+
+> [!NOTE]
+> **The single most important resource for you right now is Karpathy's ["Let's reproduce GPT-2 (124M)"](https://www.youtube.com/watch?v=l8pRSuU81PU) video.** It covers the exact same workflow you're doing — pretraining a ~124M parameter model from scratch on a single GPU with the same optimizer, LR schedule, and training loop design. Your GLM-5.2 script is heavily inspired by nanoGPT, so it will feel very familiar.
+
+---
+
+*Your training is going well. Val loss 4.1071 at step 52k is solid progress. Keep it running!* 🚀

requirements.txt

@@ -0,0 +1,16 @@
+# Requirements for "Let's Reproduce GLM-5.2"
+# Install with: pip install -r requirements.txt
+
+# Core
+torch>=2.1.0
+numpy>=1.24.0
+
+# Tokenization (GPT-2 BPE tokenizer)
+tiktoken>=0.5.0
+
+# Data Preparation (used by dataprep_pretrain.py)
+datasets>=2.14.0
+transformers>=4.30.0
+
+# Progress bars
+tqdm>=4.65.0

train_glm5.py

@@ -0,0 +1,1380 @@
+"""
+Let's Reproduce GLM-5.2 (GLM MoE DSA) — From Scratch!
+=======================================================
+
+A baby version of GLM-5.2 (zhipu-ai / zai-org / GLM-5), trained from scratch.
+
+GLM-5.2 combines three cutting-edge innovations:
+  1. MLA (Multi-Latent Attention) — LoRA-compressed Q and KV projections
+  2. DSA (DeepSeek Sparse Attention) — top-k token selection via a learned indexer
+  3. MoE (Mixture of Experts) — sigmoid-routed fine-grained experts + shared expert
+
+This script implements ALL of these from scratch in a single file,
+scaled down to ~120M parameters for training on a single consumer GPU.
+
+Architecture Reference:
+  HuggingFace transformers — models/glm_moe_dsa/modeling_glm_moe_dsa.py
+
+Paper References:
+  - DeepSeek-V3 (MLA + MoE): https://arxiv.org/abs/2412.19437
+  - DeepSeek Sparse Attention: https://arxiv.org/abs/2603.12201
+  - GLM: https://github.com/THUDM/GLM
+
+Inspired by Andrej Karpathy's "Let's reproduce GPT-2" and nanoGPT.
+
+Usage:
+    python train_glm5.py                                # Train with defaults (RTX 4050 friendly)
+    python train_glm5.py --batch_size 2                  # Smaller batch for less VRAM
+    python train_glm5.py --compile                       # Use torch.compile (faster, needs warmup)
+    python train_glm5.py --eval_only --ckpt out/ckpt.pt  # Generate from a checkpoint
+    python train_glm5.py --no_gradient_checkpointing     # Disable grad checkpointing (needs more VRAM)
+"""
+
+import argparse
+import json
+import math
+import os
+import time
+from dataclasses import dataclass
+
+import numpy as np
+import tiktoken
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+# =============================================================================
+# Section 1: Model Configuration
+# =============================================================================
+# The full GLM-5.2 has 78 layers, 6144 hidden, 256 experts — far too large.
+# We scale everything down to ~120M total params while preserving EVERY
+# architectural innovation. Think of this as "baby GLM-5.2".
+# =============================================================================
+
+
+@dataclass
+class GLM5Config:
+    """
+    Configuration for baby GLM-5.2.
+    Default values give a ~159M total param model (~82M non-embedding),
+    trainable on a single RTX 4050 (6GB VRAM) with gradient checkpointing + bf16.
+
+    The full GLM-5.2 values are shown in comments for reference.
+    """
+
+    # --- Vocabulary & Embedding ---
+    vocab_size: int = 50304  # GPT-2 tokenizer (50257) padded to nearest 128  [full: 154880]
+
+    # --- Core Dimensions ---
+    hidden_size: int = 768  # Model width (d_model)                           [full: 6144]
+    num_hidden_layers: int = 12  # Total decoder layers                       [full: 78]
+
+    # --- Multi-Latent Attention (MLA) ---
+    num_attention_heads: int = 12  # Number of query heads                    [full: 64]
+    q_lora_rank: int = 384  # Query LoRA bottleneck                           [full: 2048]
+    kv_lora_rank: int = 128  # Key/Value LoRA bottleneck                      [full: 512]
+    qk_nope_head_dim: int = 32  # Non-rotary Q/K head dim                    [full: 192]
+    qk_rope_head_dim: int = 32  # Rotary Q/K head dim                        [full: 64]
+    v_head_dim: int = 64  # Value head dim                                    [full: 256]
+
+    # --- Dense MLP ---
+    intermediate_size: int = 2048  # Dense FFN intermediate dim               [full: 12288]
+    hidden_act: str = "silu"  # Activation function
+
+    # --- Mixture of Experts (MoE) ---
+    first_k_dense_replace: int = 3  # First K layers use dense MLP           [full: 3]
+    moe_intermediate_size: int = 256  # Per-expert FFN intermediate dim       [full: 2048]
+    n_routed_experts: int = 8  # Number of routed experts                     [full: 256]
+    num_experts_per_tok: int = 2  # Top-k experts per token                   [full: 8]
+    n_shared_experts: int = 1  # Always-active shared experts                 [full: 1]
+    n_group: int = 1  # Expert groups for routing                             [full: 1]
+    topk_group: int = 1  # Top groups selected                                [full: 1]
+    routed_scaling_factor: float = 2.5  # Expert weight scaling               [full: 2.5]
+    norm_topk_prob: bool = True  # Normalize routing probabilities
+
+    # --- DeepSeek Sparse Attention (DSA) ---
+    index_topk: int = 256  # Top tokens selected by DSA indexer               [full: 2048]
+    index_head_dim: int = 64  # Head dim in DSA indexer                       [full: 128]
+    index_n_heads: int = 12  # Heads in DSA indexer                           [full: 32]
+
+    # --- Positional Encoding ---
+    max_position_embeddings: int = 4096  # Max context length                 [full: 202752]
+    rope_theta: float = 10000.0  # RoPE base frequency
+
+    # --- Regularization & Precision ---
+    rms_norm_eps: float = 1e-5  # RMSNorm epsilon
+    attention_dropout: float = 0.0  # Attention dropout
+
+    # --- Initialization ---
+    initializer_range: float = 0.02  # Std dev for weight init
+
+    # --- Weight Tying ---
+    # if False, Total model param = 158.8 M params. If True param cout changes to around 120 M
+    tie_word_embeddings: bool = True  # Tie embed + lm_head                  [full: False]
+
+    def __post_init__(self):
+        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        # Per-layer MLP type: first K layers = dense, rest = MoE (sparse)
+        n_dense = min(self.first_k_dense_replace, self.num_hidden_layers)
+        self.mlp_layer_types = ["dense"] * n_dense + ["sparse"] * (self.num_hidden_layers - n_dense)
+        # DSA indexer pattern: alternating "full" (run indexer) / "shared" (reuse previous)
+        # Full GLM-5.2 uses a freq/offset schedule; we simplify to alternating.
+        self.indexer_types = [
+            "full" if i % 2 == 0 else "shared" for i in range(self.num_hidden_layers)
+        ]
+
+
+# =============================================================================
+# Section 2: Architecture Components
+# =============================================================================
+# We build each component bottom-up, with detailed comments explaining
+# WHY each design choice was made in GLM-5.2.
+# =============================================================================
+
+# ---------------------------------------------------------------------------
+# 2a: RMSNorm — Root Mean Square Layer Normalization
+# ---------------------------------------------------------------------------
+
+
+class RMSNorm(nn.Module):
+    """
+    RMSNorm (Root Mean Square Layer Normalization).
+
+    Unlike LayerNorm, RMSNorm does NOT center activations (no mean subtraction).
+    This is cheaper and works just as well for LLMs.
+
+    Formula: output = x / sqrt(mean(x²) + eps) * weight
+    """
+
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x):
+        input_dtype = x.dtype
+        x = x.float()  # Always compute in float32 for numerical stability
+        variance = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return self.weight * x.to(input_dtype)
+
+
+# ---------------------------------------------------------------------------
+# 2b: Rotary Position Embedding (RoPE) — Interleaved variant
+# ---------------------------------------------------------------------------
+
+
+class RotaryEmbedding(nn.Module):
+    """
+    Standard Rotary Position Embedding (RoPE).
+
+    Computes cos/sin tables for position encoding. The actual rotation is
+    applied by apply_rotary_pos_emb_interleave() — see below.
+    """
+
+    def __init__(self, dim, max_position_embeddings=4096, base=10000.0, device=None):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float, device=device) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        """
+        Args:
+            x: [B, T, D] — only used for dtype/device reference
+            position_ids: [B, T]
+        Returns:
+            cos, sin: each [B, T, dim]
+        """
+        inv_freq = self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
+        pos = position_ids[:, None, :].float()
+        freqs = (inv_freq.float() @ pos.float()).transpose(1, 2)  # [B, T, dim/2]
+        emb = torch.cat((freqs, freqs), dim=-1)  # [B, T, dim]
+        return emb.cos().to(x.dtype), emb.sin().to(x.dtype)
+
+
+def apply_rotary_pos_emb_interleave(q, k, cos, sin, unsqueeze_dim=1):
+    """
+    Apply INTERLEAVED Rotary Position Embedding.
+
+    GLM-5.2 (and DeepSeek) uses interleaved RoPE pairs: (x0,x1), (x2,x3), ...
+    Each pair is rotated by a single frequency.
+
+    This is DIFFERENT from standard LLaMA-style RoPE which splits the
+    first/second half of the head dimension. The interleaved version avoids
+    memory-shuffling copies from 'rotate_half'.
+
+    ┌──────────────────────────────────────────────────────────────────┐
+    │  Standard RoPE:     [x0..x_d/2 | x_d/2..x_d] → rotate halves  │
+    │  Interleaved RoPE:  [x0,x1 | x2,x3 | ...] → rotate pairs      │
+    └──────────────────────────────────────────────────────────────────┘
+    """
+    # cos/sin come as cat(freqs, freqs) → take the first half
+    cos = cos[..., : cos.shape[-1] // 2].unsqueeze(unsqueeze_dim)
+    sin = sin[..., : sin.shape[-1] // 2].unsqueeze(unsqueeze_dim)
+
+    # Split into even and odd indexed elements (the interleaved pairs)
+    q1, q2 = q[..., 0::2], q[..., 1::2]
+    k1, k2 = k[..., 0::2], k[..., 1::2]
+
+    # Apply 2D rotation to each (even, odd) pair
+    q_embed = torch.cat([q1 * cos - q2 * sin, q2 * cos + q1 * sin], dim=-1)
+    k_embed = torch.cat([k1 * cos - k2 * sin, k2 * cos + k1 * sin], dim=-1)
+    return q_embed, k_embed
+
+
+# ---------------------------------------------------------------------------
+# 2c: Gated MLP (SwiGLU)
+# ---------------------------------------------------------------------------
+
+
+class GatedMLP(nn.Module):
+    """
+    SwiGLU MLP: down_proj( SiLU(gate_proj(x)) ⊙ up_proj(x) )
+
+    The gating mechanism (SiLU on gate, element-wise multiply with up)
+    consistently outperforms vanilla ReLU/GELU FFNs in modern LLMs.
+    Used in LLaMA, DeepSeek, GLM, Gemma, Qwen, and many others.
+    """
+
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+# ---------------------------------------------------------------------------
+# 2d: Top-K Expert Router (Sigmoid-based, DeepSeek-style)
+# ---------------------------------------------------------------------------
+
+
+class TopKRouter(nn.Module):
+    """
+    Sigmoid-based top-k expert router with bias correction.
+
+    Key differences from the traditional softmax MoE router:
+    ┌──────────────────────────────────────────────────────────────────┐
+    │  1. SIGMOID scoring (not softmax) — experts scored independently │
+    │  2. Correction bias — loaded from checkpoint, helps balance load │
+    │  3. Group routing — select top groups, then experts within them  │
+    │  4. Normalize + scale — weights normalized then scaled by 2.5x  │
+    └──────────────────────────────────────────────────────────────────┘
+
+    The sigmoid approach prevents "expert collapse" where softmax routing
+    causes only a few experts to receive all the tokens.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.n_routed_experts
+        self.hidden_dim = config.hidden_size
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.norm_topk_prob = config.norm_topk_prob
+
+        # Router weight: one logit per expert
+        self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim))
+        # Correction bias: pretrained load-balancing signal (zeros for training from scratch)
+        self.register_buffer("e_score_correction_bias", torch.zeros(self.num_experts))
+
+    def forward(self, x):
+        x_flat = x.view(-1, self.hidden_dim)
+
+        # Step 1: Sigmoid scoring (NOT softmax!)
+        # Each expert gets an independent 0-1 probability
+        router_logits = F.linear(x_flat.float(), self.weight.float())
+        scores = router_logits.sigmoid()
+        scores_for_choice = scores + self.e_score_correction_bias
+
+        # Step 2: Group-based routing
+        # With n_group=1 (our config), this is standard top-k.
+        # With n_group>1 (full GLM-5.2), first select best groups, then pick
+        # top experts only from those groups — prevents cross-group interference.
+        group_scores = (
+            scores_for_choice.view(-1, self.n_group, self.num_experts // self.n_group)
+            .topk(2, dim=-1)[0]
+            .sum(dim=-1)
+        )
+        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_idx, 1)
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(-1, self.n_group, self.num_experts // self.n_group)
+            .reshape(-1, self.num_experts)
+        )
+        scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
+
+        # Step 3: Select top-k experts per token
+        topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
+        topk_weights = scores.gather(1, topk_indices)
+
+        # Step 4: Normalize probabilities and scale
+        if self.norm_topk_prob:
+            topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20)
+        topk_weights = topk_weights * self.routed_scaling_factor
+
+        return topk_weights, topk_indices
+
+
+# ---------------------------------------------------------------------------
+# 2e: MoE Expert Collection (Batched 3D Tensors)
+# ---------------------------------------------------------------------------
+
+
+class MoEExperts(nn.Module):
+    """
+    Collection of expert MLPs stored as batched 3D parameter tensors.
+
+    Instead of N separate nn.Linear modules, we store ALL expert weights
+    in single tensors. This enables efficient batched dispatch.
+
+      gate_up_proj: [num_experts, 2*intermediate, hidden]
+      down_proj:    [num_experts, hidden, intermediate]
+
+    Each expert computes: SiLU(gate(x)) * up(x) → down → output
+
+    NOTE: This is the naive loop implementation. Production systems (DeepSeek,
+    GLM-5.2) use custom CUDA kernels for grouped GEMM — orders of magnitude faster.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.n_routed_experts
+        self.hidden_dim = config.hidden_size
+        self.intermediate_dim = config.moe_intermediate_size
+
+        # Fused gate+up projection: [E, 2*I, D]
+        self.gate_up_proj = nn.Parameter(
+            torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)
+        )
+        # Down projection: [E, D, I]
+        self.down_proj = nn.Parameter(
+            torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)
+        )
+
+    def forward(self, x, topk_indices, topk_weights):
+        """
+        Route tokens to their selected experts and accumulate outputs.
+
+        Args:
+            x: [num_tokens, hidden_dim]
+            topk_indices: [num_tokens, top_k] — which experts each token uses
+            topk_weights: [num_tokens, top_k] — routing weights
+        """
+        final = torch.zeros_like(x)
+
+        # Build per-expert assignment mask
+        with torch.no_grad():
+            expert_mask = F.one_hot(topk_indices, self.num_experts)  # [tokens, top_k, E]
+            expert_mask = expert_mask.permute(2, 1, 0)  # [E, top_k, tokens]
+            expert_hit = (expert_mask.sum(dim=(-1, -2)) > 0).nonzero()
+
+        # Process each active expert
+        for idx in expert_hit:
+            e = idx[0]
+            top_k_pos, token_idx = torch.where(expert_mask[e])
+            current = x[token_idx]
+
+            # SwiGLU: SiLU(gate) * up → down
+            gate, up = F.linear(current, self.gate_up_proj[e]).chunk(2, dim=-1)
+            hidden = F.silu(gate) * up
+            out = F.linear(hidden, self.down_proj[e])
+
+            # Weight by routing probability and accumulate
+            out = out * topk_weights[token_idx, top_k_pos, None]
+            final.index_add_(0, token_idx, out.to(final.dtype))
+
+        return final
+
+
+# ---------------------------------------------------------------------------
+# 2f: Full MoE Block (Router + Routed Experts + Shared Expert)
+# ---------------------------------------------------------------------------
+
+
+class MoEBlock(nn.Module):
+    """
+    Full Mixture-of-Experts block.
+
+    Output = Routed_Experts(x) + Shared_Expert(x)
+
+    The shared expert ALWAYS processes all tokens — it provides a stable
+    "backbone" of computation. The routed experts add specialized capacity
+    for different types of tokens/patterns.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.gate = TopKRouter(config)
+        self.experts = MoEExperts(config)
+        # Shared expert: always-on, processes every token unconditionally
+        self.shared_experts = GatedMLP(
+            config.hidden_size,
+            config.moe_intermediate_size * config.n_shared_experts,
+        )
+
+    def forward(self, x):
+        residual = x
+        orig_shape = x.shape
+        topk_weights, topk_indices = self.gate(x)
+        x = x.view(-1, x.shape[-1])
+        x = self.experts(x, topk_indices, topk_weights).view(*orig_shape)
+        x = x + self.shared_experts(residual)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# 2g: DeepSeek Sparse Attention (DSA) Indexer
+# ---------------------------------------------------------------------------
+
+
+class DSAIndexer(nn.Module):
+    """
+    DeepSeek Sparse Attention (DSA) Indexer.
+
+    THE key innovation of DSA: instead of attending to ALL past tokens (O(n²)),
+    the indexer selects only the top-k most relevant tokens per query position.
+    This makes long-context attention tractable (O(n·k) where k << n).
+
+    Architecture:
+    ┌─────────────────────────────────────────────────────────────────────┐
+    │  1. Separate Q/K projections (NOT shared with main MLA attention)  │
+    │  2. Multi-head dot-product scoring with ReLU (not softmax!)        │
+    │  3. Learned per-head importance weights for aggregation            │
+    │  4. Returns top-k token indices for the main attention to use      │
+    └─────────────────────────────────────────────────────────────────────┘
+
+    NOTE: The @torch.no_grad() decorator matches the official implementation.
+    The indexer doesn't backpropagate gradients — in production GLM-5.2, it's
+    pre-trained separately. For training from scratch, the random-but-causal
+    token selection acts as attention regularization. The model learns to be
+    robust to approximate attention through its main MLA weights.
+    """
+
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.n_heads = config.index_n_heads
+        self.head_dim = config.index_head_dim
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.index_topk = config.index_topk
+        self.q_lora_rank = config.q_lora_rank
+
+        # The indexer has its OWN projections — completely separate from main attention!
+        self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+        self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6)
+        # Learned per-head importance: "how much should we trust each head's score?"
+        self.weights_proj = nn.Linear(self.hidden_size, self.n_heads, bias=False)
+        self.softmax_scale = self.head_dim**-0.5
+
+    @torch.no_grad()
+    def forward(self, hidden_states, q_resid, cos, sin, position_ids):
+        """
+        Select top-k most relevant tokens for each query position.
+
+        Args:
+            hidden_states: [B, S, hidden_size] — input to this layer
+            q_resid: [B, S, q_lora_rank] — query residual from MLA's q_a_layernorm
+            cos, sin: position embeddings
+            position_ids: [B, S]
+
+        Returns:
+            topk_indices: [B, S, topk] — indices of selected tokens (int32)
+        """
+        B, S, _ = hidden_states.shape
+
+        # --- Query: project from q_resid (shared with main attention's LoRA output) ---
+        q = self.wq_b(q_resid).view(B, S, self.n_heads, self.head_dim)
+        q_rot, q_pass = q.split([self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
+
+        # --- Key: project from hidden states (fresh, independent projection) ---
+        k = self.k_norm(self.wk(hidden_states)).unsqueeze(2)  # [B, S, 1, head_dim]
+        k_rot, k_pass = k.split([self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
+
+        # --- Apply interleaved RoPE to both Q and K ---
+        q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin, unsqueeze_dim=2)
+        q = torch.cat([q_rot, q_pass], dim=-1)  # [B, S, n_heads, head_dim]
+        k = torch.cat([k_rot, k_pass], dim=-1).squeeze(2)  # [B, S, head_dim]
+
+        # --- Multi-head relevance scoring ---
+        # Each head independently scores every (query, key) pair
+        # q: [B, S, n_heads, D] @ k^T: [B, 1, D, S] → [B, S, n_heads, S]
+        scores = (
+            torch.matmul(q.float(), k.transpose(-1, -2).float().unsqueeze(1)) * self.softmax_scale
+        )
+        scores = F.relu(scores)  # ReLU! Not softmax. This creates naturally sparse scores.
+
+        # --- Weighted head aggregation ---
+        # Learn which heads' opinions matter more, then combine
+        weights = self.weights_proj(hidden_states.to(self.weights_proj.weight.dtype)).float()
+        weights = weights * (self.n_heads**-0.5)
+        # [B, S, 1, n_heads] @ [B, S, n_heads, S] → [B, S, 1, S] → squeeze → [B, S, S]
+        index_scores = torch.matmul(weights.unsqueeze(-2), scores).squeeze(-2)
+
+        # --- Enforce causality: can't select future tokens! ---
+        key_positions = torch.arange(S, device=hidden_states.device)
+        causal = key_positions[None, None, :] > position_ids[:, :, None]
+        index_scores = index_scores.masked_fill(causal, float("-inf"))
+
+        # --- Select top-k most relevant tokens ---
+        topk = min(self.index_topk, S)
+        return index_scores.topk(topk, dim=-1).indices.to(torch.int32)
+
+
+# ---------------------------------------------------------------------------
+# 2h: Multi-Latent Attention (MLA) + DSA Integration
+# ---------------------------------------------------------------------------
+
+
+class MultiLatentAttention(nn.Module):
+    """
+    Multi-Latent Attention (MLA) with DeepSeek Sparse Attention (DSA).
+
+    MLA compresses queries and key-values through LoRA-style bottlenecks.
+    This dramatically reduces KV-cache size during inference.
+
+    ┌─────────────────────────────────────────────────────────────────────────┐
+    │  Query path:                                                           │
+    │    x → q_a_proj (compress) → RMSNorm → q_b_proj (expand per-head)     │
+    │    → split into [q_nope, q_rope] → apply RoPE to q_rope               │
+    │                                                                        │
+    │  KV path:                                                              │
+    │    x → kv_a_proj (compress to [kv_latent + k_rope])                    │
+    │    → kv_latent → RMSNorm → kv_b_proj (expand per-head)                │
+    │    → split into [k_nope, value]                                        │
+    │    → k_rope gets RoPE and is broadcast to all heads                    │
+    │                                                                        │
+    │  Then: q = [q_nope, q_rope], k = [k_nope, k_rope]                     │
+    │  Standard scaled dot-product attention with DSA sparse masking          │
+    └─────────────────────────────────────────────────────────────────────────┘
+
+    Cross-layer DSA sharing:
+    - "full" layers run the indexer to compute fresh top-k indices
+    - "shared" layers reuse the previous full layer's indices (saves compute)
+    """
+
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.q_lora_rank = config.q_lora_rank
+        self.kv_lora_rank = config.kv_lora_rank
+        self.qk_nope_head_dim = config.qk_nope_head_dim
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.qk_head_dim = config.qk_head_dim  # nope + rope
+        self.v_head_dim = config.v_head_dim
+
+        # === Query LoRA compression ===
+        # hidden → compress → normalize → expand to per-head queries
+        self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False)
+        self.q_a_layernorm = RMSNorm(config.q_lora_rank, eps=config.rms_norm_eps)
+        self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)
+
+        # === KV LoRA compression ===
+        # hidden → compress to [kv_latent, k_rope_shared]
+        self.kv_a_proj_with_mqa = nn.Linear(
+            config.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=False,
+        )
+        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
+        # kv_latent → expand to per-head [k_nope, value]
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+            bias=False,
+        )
+
+        # === Output projection ===
+        self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, config.hidden_size, bias=False)
+
+        # === Attention scaling ===
+        self.scaling = self.qk_head_dim ** (-0.5)
+
+        # === DSA: indexer or shared ===
+        self.skip_topk = config.indexer_types[layer_idx] == "shared"
+        self.indexer = None if self.skip_topk else DSAIndexer(config, layer_idx)
+
+    def forward(self, x, cos, sin, position_ids, prev_topk_indices=None):
+        B, T, _ = x.shape
+
+        # ============ Query Path ============
+        # x → compress(768→384) → RMSNorm → expand(384→12*64=768)
+        q_resid = self.q_a_layernorm(self.q_a_proj(x))  # [B, T, q_lora_rank=384]
+        q = self.q_b_proj(q_resid)  # [B, T, num_heads * qk_head_dim]
+        q = q.view(B, T, self.num_heads, self.qk_head_dim).transpose(1, 2)  # [B, H, T, qk_head_dim]
+        q_nope, q_rope = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+        # ============ KV Path ============
+        # x → compress(768→160) → split [kv_latent(128), k_rope(32)]
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        k_compressed, k_rope = compressed_kv.split(
+            [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+        )
+        # kv_latent → RMSNorm → expand(128→12*96=1152) → split [k_nope(32), v(64)]
+        kv = self.kv_b_proj(self.kv_a_layernorm(k_compressed))
+        kv = kv.view(B, T, self.num_heads, self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2)
+        k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        # ============ Apply Interleaved RoPE ============
+        # k_rope is shared across heads (MQA-style for the rope component)
+        k_rope = k_rope.view(B, 1, T, self.qk_rope_head_dim)
+        q_rope, k_rope = apply_rotary_pos_emb_interleave(q_rope, k_rope, cos, sin)
+        k_rope = k_rope.expand(B, self.num_heads, T, -1)  # broadcast to all heads
+
+        # Concatenate nope + rope components for final Q and K
+        q = torch.cat([q_nope, q_rope], dim=-1)  # [B, H, T, qk_head_dim=64]
+        k = torch.cat([k_nope, k_rope], dim=-1)  # [B, H, T, qk_head_dim=64]
+
+        # ============ DSA: Sparse Token Selection ============
+        if self.indexer is not None:
+            # "Full" layer: run indexer to get fresh top-k indices
+            topk_indices = self.indexer(x, q_resid, cos, sin, position_ids)
+        else:
+            # "Shared" layer: reuse previous full layer's indices
+            assert prev_topk_indices is not None, (
+                f"Layer {self.layer_idx} is 'shared' DSA but got no previous top-k indices!"
+            )
+            topk_indices = prev_topk_indices
+
+        # ============ Attention Computation ============
+        attn_weights = torch.matmul(q, k.transpose(-1, -2)) * self.scaling  # [B, H, T, T]
+
+        # Causal mask: prevent attending to future positions
+        causal_mask = torch.triu(
+            torch.full((T, T), torch.finfo(q.dtype).min, device=x.device, dtype=q.dtype),
+            diagonal=1,
+        )
+        attn_weights = attn_weights + causal_mask[None, None, :, :]
+
+        # DSA sparse mask: ONLY attend to the indexer's top-k selected tokens
+        # index_mask[b, t, t'] = True → position t' is NOT selected → mask it out
+        index_mask = torch.ones(B, T, T, device=x.device, dtype=torch.bool)
+        index_mask.scatter_(-1, topk_indices.long(), False)  # Set selected positions to False (unmasked)
+        attn_weights = attn_weights.masked_fill(
+            index_mask.unsqueeze(1),  # [B, 1, T, T] — broadcast across heads
+            torch.finfo(q.dtype).min,
+        )
+
+        # Softmax + weighted sum of values
+        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
+        attn_output = torch.matmul(attn_weights, v)  # [B, H, T, v_head_dim]
+
+        # Reshape and project output
+        attn_output = attn_output.transpose(1, 2).reshape(B, T, self.num_heads * self.v_head_dim)
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, topk_indices
+
+
+# ---------------------------------------------------------------------------
+# 2i: Transformer Decoder Layer
+# ---------------------------------------------------------------------------
+
+
+class DecoderLayer(nn.Module):
+    """
+    Pre-norm Transformer decoder layer.
+
+    Structure:
+        x → LayerNorm → MLA Attention → +residual → LayerNorm → MLP/MoE → +residual
+
+    Layers 0..first_k_dense_replace use dense SwiGLU MLP.
+    Remaining layers use Mixture-of-Experts (MoE).
+    """
+
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.self_attn = MultiLatentAttention(config, layer_idx)
+
+        # Choose MLP type based on layer position
+        if config.mlp_layer_types[layer_idx] == "sparse":
+            self.mlp = MoEBlock(config)
+        else:
+            self.mlp = GatedMLP(config.hidden_size, config.intermediate_size)
+
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(self, x, cos, sin, position_ids, prev_topk_indices=None):
+        # Pre-norm → Attention → Residual
+        residual = x
+        x = self.input_layernorm(x)
+        attn_out, topk_indices = self.self_attn(x, cos, sin, position_ids, prev_topk_indices)
+        x = residual + attn_out
+
+        # Pre-norm → MLP/MoE → Residual
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = residual + self.mlp(x)
+
+        return x, topk_indices
+
+
+# ---------------------------------------------------------------------------
+# 2j: Full Model (Base + CausalLM head)
+# ---------------------------------------------------------------------------
+
+
+class GLM5Model(nn.Module):
+    """GLM-5.2 base model: token embeddings → N decoder layers → final RMSNorm."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = RotaryEmbedding(
+            config.qk_rope_head_dim,
+            max_position_embeddings=config.max_position_embeddings,
+            base=config.rope_theta,
+        )
+        self.gradient_checkpointing = False
+
+    def forward(self, input_ids):
+        B, T = input_ids.shape
+        assert T <= self.config.max_position_embeddings, (
+            f"Sequence length {T} > max_position_embeddings {self.config.max_position_embeddings}"
+        )
+
+        x = self.embed_tokens(input_ids)
+
+        # Compute position embeddings once (shared across all layers)
+        position_ids = torch.arange(T, device=input_ids.device).unsqueeze(0).expand(B, -1)
+        cos, sin = self.rotary_emb(x, position_ids)
+
+        # Forward through decoder layers
+        # Each layer returns (hidden_states, topk_indices)
+        # topk_indices propagate from "full" DSA layers to "shared" layers
+        topk_indices = None
+        for layer in self.layers:
+            if self.gradient_checkpointing and self.training:
+                # gradient_checkpointing.checkpoint does not support None inputs.
+                # Pass a sentinel zero-tensor when topk_indices is None (layer 0 full-indexer layers),
+                # and detect it inside with a flag. Simpler: just skip checkpointing for the very
+                # first "full" layer (layer 0) which has no prev indices to receive.
+                if topk_indices is None:
+                    # Layer 0 is always a "full" DSA layer — run normally, then checkpoint the rest.
+                    x, topk_indices = layer(x, cos, sin, position_ids, None)
+                else:
+                    x, topk_indices = torch.utils.checkpoint.checkpoint(
+                        layer, x, cos, sin, position_ids, topk_indices,
+                        use_reentrant=False,
+                    )
+            else:
+                x, topk_indices = layer(x, cos, sin, position_ids, topk_indices)
+
+        return self.norm(x)
+
+
+class GLM5ForCausalLM(nn.Module):
+    """
+    GLM-5.2 for Causal Language Modeling.
+    = GLM5Model (base) + Linear lm_head (vocab projection).
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.model = GLM5Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+
+        # Initialize all weights
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        """Initialize weights following GLM-5.2 conventions."""
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, RMSNorm):
+            nn.init.ones_(module.weight)
+        elif isinstance(module, MoEExperts):
+            nn.init.normal_(module.gate_up_proj, mean=0.0, std=std)
+            nn.init.normal_(module.down_proj, mean=0.0, std=std)
+        elif isinstance(module, TopKRouter):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+
+    def forward(self, input_ids, targets=None):
+        """
+        Args:
+            input_ids: [B, T] token indices
+            targets: [B, T] target token indices (shifted by 1 in get_batch)
+
+        Returns:
+            logits: [B, T, vocab_size]
+            loss: scalar if targets provided, else None
+        """
+        hidden_states = self.model(input_ids)
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=-1,
+            )
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=200):
+        """
+        Simple autoregressive generation with temperature + top-k sampling.
+
+        No KV-cache for simplicity — recomputes the full context each step.
+        This is slower but simpler, and matches Karpathy's nanoGPT style.
+        """
+        self.eval()
+        for _ in range(max_new_tokens):
+            # Crop to max context length if needed
+            idx_cond = (
+                idx
+                if idx.size(1) <= self.config.max_position_embeddings
+                else idx[:, -self.config.max_position_embeddings :]
+            )
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = float("-inf")
+
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx
+
+    def param_count(self):
+        """Return a detailed parameter count breakdown."""
+        total = sum(p.numel() for p in self.parameters())
+        embed_params = self.model.embed_tokens.weight.numel()
+        head_params = self.lm_head.weight.numel() if not self.config.tie_word_embeddings else 0
+        non_embed = total - embed_params - head_params
+
+        # MoE active params per token
+        moe_layers = sum(1 for t in self.config.mlp_layer_types if t == "sparse")
+        if moe_layers > 0:
+            expert_params_per_layer = (
+                2 * self.config.moe_intermediate_size * self.config.hidden_size
+                + self.config.hidden_size * self.config.moe_intermediate_size
+            )
+            total_expert_params = expert_params_per_layer * self.config.n_routed_experts * moe_layers
+            active_expert_params = expert_params_per_layer * self.config.num_experts_per_tok * moe_layers
+            active_ratio = self.config.num_experts_per_tok / self.config.n_routed_experts
+        else:
+            total_expert_params = 0
+            active_expert_params = 0
+            active_ratio = 1.0
+
+        active_params = total - total_expert_params + active_expert_params
+        return {
+            "total": total,
+            "non_embedding": non_embed,
+            "active_per_token": active_params,
+            "moe_active_ratio": active_ratio,
+        }
+
+
+# =============================================================================
+# Section 3: Data Loading
+# =============================================================================
+# Loads pre-tokenized binary data produced by dataprep_pretrain.py.
+# Data format: train.bin / val.bin (uint16 memmap) + meta.json.
+#
+# Run dataprep_pretrain.py first to prepare the data:
+#   python dataprep_pretrain.py              # Full 3.3B tokens
+#   python dataprep_pretrain.py --total_tokens 10000000  # Quick 10M test
+# =============================================================================
+
+
+def load_pretrain_data(data_dir):
+    """
+    Load pre-tokenized binary data from data_dir.
+
+    Expects:
+      data_dir/train.bin  -- binary token file (uint16 or uint32)
+      data_dir/val.bin    -- binary token file (uint16 or uint32)
+      data_dir/meta.json  -- metadata (vocab_size, dtype, token counts)
+
+    Returns:
+      train_data: np.memmap of training tokens
+      val_data: np.memmap of validation tokens
+
+    Raises:
+      FileNotFoundError if data files are missing.
+    """
+    train_path = os.path.join(data_dir, "train.bin")
+    val_path = os.path.join(data_dir, "val.bin")
+    meta_path = os.path.join(data_dir, "meta.json")
+
+    # --- Validate files exist ---
+    for path, name in [(train_path, "train.bin"), (val_path, "val.bin"), (meta_path, "meta.json")]:
+        if not os.path.exists(path):
+            raise FileNotFoundError(
+                f"  [ERR] {name} not found at: {path}\n"
+                f"        Run dataprep_pretrain.py first to prepare the data:\n"
+                f"          python dataprep_pretrain.py\n"
+                f"        Or for a quick test:\n"
+                f"          python dataprep_pretrain.py --total_tokens 10000000"
+            )
+
+    # --- Load metadata ---
+    with open(meta_path, "r") as f:
+        meta = json.load(f)
+
+    dtype_str = meta.get("dtype", "uint16")
+    dtype = np.uint16 if dtype_str == "uint16" else np.uint32
+    train_tokens = meta.get("train_tokens", 0)
+    val_tokens = meta.get("val_tokens", 0)
+
+    print(f"  [Data]")
+    print(f"    Tokenizer:   {meta.get('tokenizer', 'unknown')}")
+    print(f"    Vocab size:  {meta.get('vocab_size', 'unknown')}")
+    print(f"    Dtype:       {dtype_str}")
+    print(f"    Train:       {train_tokens:,} tokens")
+    print(f"    Val:         {val_tokens:,} tokens")
+    print(f"    Total:       {train_tokens + val_tokens:,} tokens")
+    print(f"    Sources:     {', '.join(meta.get('sources', []))}")
+
+    # --- Memory-map the binary files ---
+    # memmap reads directly from disk without loading into RAM.
+    # This is critical for 3.3B tokens (~7GB) on a 6GB VRAM machine.
+    train_data = np.memmap(train_path, dtype=dtype, mode="r")
+    val_data = np.memmap(val_path, dtype=dtype, mode="r")
+
+    return train_data, val_data
+
+
+def get_batch(split, train_data, val_data, block_size, batch_size, device):
+    """
+    Sample a random batch of token sequences from memmap data.
+
+    Uses .copy() on numpy slices to avoid torch tensor issues with
+    non-writable memmap arrays.
+    """
+    data = train_data if split == "train" else val_data
+    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
+    x = torch.stack([torch.from_numpy(data[i : i + block_size].astype(np.int64).copy()) for i in ix])
+    y = torch.stack([torch.from_numpy(data[i + 1 : i + 1 + block_size].astype(np.int64).copy()) for i in ix])
+    return x.to(device), y.to(device)
+
+
+# =============================================================================
+# Section 4: Training
+# =============================================================================
+# AdamW optimizer with cosine LR schedule, gradient accumulation,
+# mixed precision, gradient checkpointing, periodic eval + checkpointing.
+# =============================================================================
+
+
+def get_lr(it, warmup_iters, lr_decay_iters, learning_rate, min_lr, stable_iters=0):
+    """WSD (Warmup-Stable-Decay) learning rate schedule.
+
+    Phases:
+      1. Warmup:  steps [0, warmup_iters)           — linear ramp 0 → learning_rate
+      2. Stable:  steps [warmup_iters, stable_iters) — constant at learning_rate
+      3. Decay:   steps [stable_iters, lr_decay_iters] — cosine decay → min_lr
+
+    If stable_iters <= warmup_iters (default: 0), this reduces to the standard
+    cosine schedule with warmup (backward-compatible).
+    """
+    # Phase 1: Linear warmup
+    if it < warmup_iters:
+        return learning_rate * (it + 1) / warmup_iters
+    # Phase 3 ended: hold at min_lr
+    if it > lr_decay_iters:
+        return min_lr
+    # Phase 2: Stable (constant LR) — only if stable_iters is set
+    if stable_iters > warmup_iters and it < stable_iters:
+        return learning_rate
+    # Phase 3: Cosine decay from learning_rate → min_lr
+    decay_start = max(stable_iters, warmup_iters)
+    decay_ratio = (it - decay_start) / (lr_decay_iters - decay_start)
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    return min_lr + coeff * (learning_rate - min_lr)
+
+
+@torch.no_grad()
+def estimate_loss(model, train_data, val_data, eval_iters, block_size, batch_size, device, ctx):
+    """Estimate loss on train and val splits (averaged over eval_iters batches)."""
+    model.eval()
+    out = {}
+    for split in ["train", "val"]:
+        losses = []
+        for _ in range(eval_iters):
+            X, Y = get_batch(split, train_data, val_data, block_size, batch_size, device)
+            with ctx:
+                _, loss = model(X, Y)
+            losses.append(loss.item())
+        out[split] = np.mean(losses)
+    model.train()
+    return out
+
+
+def train(args):
+    """Main training function."""
+    # --- Device Setup ---
+    device = args.device
+    if device == "auto":
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    print(f"\n{'='*70}")
+    print(f"  >> Let's Reproduce GLM-5.2 (GLM MoE DSA) -- From Scratch!")
+    print(f"{'='*70}")
+    print(f"  Device: {device}")
+    if device == "cuda":
+        torch.set_float32_matmul_precision("high")
+        print(f"  GPU:    {torch.cuda.get_device_name()}")
+        print(f"  VRAM:   {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+        print(f"  [OK] TF32 Tensor Cores enabled")
+
+    # --- Data ---
+    train_data, val_data = load_pretrain_data(args.data_dir)
+
+    # --- Model ---
+    config = GLM5Config()
+    model = GLM5ForCausalLM(config)
+    counts = model.param_count()
+
+    print(f"\n  [Model Architecture]")
+    print(f"    Hidden size:        {config.hidden_size}")
+    print(f"    Layers:             {config.num_hidden_layers} "
+          f"({config.first_k_dense_replace} dense + "
+          f"{config.num_hidden_layers - config.first_k_dense_replace} MoE)")
+    print(f"    Attention heads:    {config.num_attention_heads}")
+    print(f"    Q LoRA rank:        {config.q_lora_rank}  ->  qk_head_dim: {config.qk_head_dim} "
+          f"(nope:{config.qk_nope_head_dim} + rope:{config.qk_rope_head_dim})")
+    print(f"    KV LoRA rank:       {config.kv_lora_rank}  ->  v_head_dim: {config.v_head_dim}")
+    print(f"    Experts:            {config.n_routed_experts} routed "
+          f"(top-{config.num_experts_per_tok}) + {config.n_shared_experts} shared")
+    print(f"    DSA index_topk:     {config.index_topk}")
+    print(f"    Indexer pattern:    {''.join('F' if t == 'full' else 'S' for t in config.indexer_types)}")
+    print(f"\n  [Parameters]")
+    print(f"    Total:              {counts['total']:>12,}")
+    print(f"    Non-embedding:      {counts['non_embedding']:>12,}")
+    print(f"    Active per token:   {counts['active_per_token']:>12,}  "
+          f"({counts['moe_active_ratio']:.0%} of experts active)")
+    print(f"    VRAM (est. train):  ~{counts['total'] * 16 / 1e9:.1f} GB  "
+          f"(weights + optimizer + gradients)")
+
+    model = model.to(device)
+
+    # --- Gradient Checkpointing ---
+    if args.gradient_checkpointing:
+        model.model.gradient_checkpointing = True
+        print(f"\n  [OK] Gradient checkpointing: ON (saves ~40% VRAM, ~30% slower)")
+
+    # --- Mixed Precision ---
+    if device == "cuda" and torch.cuda.is_bf16_supported():
+        dtype = torch.bfloat16
+        ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
+        print(f"  [OK] Mixed precision: bfloat16")
+    elif device == "cuda":
+        dtype = torch.float16
+        ctx = torch.amp.autocast(device_type="cuda", dtype=torch.float16)
+        print(f"  [OK] Mixed precision: float16")
+    else:
+        dtype = torch.float32
+        ctx = torch.amp.autocast(device_type="cpu", enabled=False)
+        print(f"  [WARN] No mixed precision (CPU mode)")
+
+    # GradScaler only needed for float16 (bfloat16 doesn't need scaling)
+    # On CPU, GradScaler must be disabled entirely (no CUDA streams available)
+    scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16 and device == "cuda"))
+
+    # --- Dry Run: verify forward pass and check VRAM ---
+    print(f"\n  Verifying forward pass...")
+    try:
+        with torch.no_grad():
+            dummy = torch.randint(0, config.vocab_size, (1, args.block_size), device=device)
+            with ctx:
+                _, test_loss = model(dummy, dummy)
+            print(f"  [OK] Forward pass OK (dummy loss={test_loss.item():.4f})")
+            if device == "cuda":
+                print(f"  [OK] VRAM after forward: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")
+                torch.cuda.reset_peak_memory_stats()
+    except torch.cuda.OutOfMemoryError:
+        print(f"  [ERR] OOM during forward pass! Try reducing --batch_size or --block_size")
+        return
+
+    # --- torch.compile ---
+    raw_model = model  # Keep a reference to the un-compiled model for saving
+    if args.compile and device == "cuda":
+        print(f"  [OK] Compiling model with torch.compile (first step will be slow)...")
+        model = torch.compile(model)
+
+    # --- Optimizer ---
+    # Separate weight decay: only for 2D+ params (weight matrices), not biases/norms
+    decay_params = []
+    no_decay_params = []
+    for name, param in raw_model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if param.dim() >= 2:
+            decay_params.append(param)
+        else:
+            no_decay_params.append(param)
+
+    optimizer = torch.optim.AdamW(
+        [
+            {"params": decay_params, "weight_decay": args.weight_decay},
+            {"params": no_decay_params, "weight_decay": 0.0},
+        ],
+        lr=args.learning_rate,
+        betas=(args.beta1, args.beta2),
+        fused=(device == "cuda"),
+    )
+
+    tokens_per_step = args.batch_size * args.block_size * args.gradient_accumulation_steps
+    print(f"\n  [Training Configuration]")
+    print(f"    Batch size:         {args.batch_size} x {args.gradient_accumulation_steps} "
+          f"grad accum = {args.batch_size * args.gradient_accumulation_steps} effective")
+    print(f"    Sequence length:    {args.block_size}")
+    print(f"    Tokens per step:    {tokens_per_step:,}")
+    print(f"    Max iterations:     {args.max_iters:,}")
+    print(f"    Total tokens:       ~{tokens_per_step * args.max_iters:,}")
+    print(f"    Learning rate:      {args.learning_rate} -> {args.min_lr} (cosine)")
+    print(f"    Warmup:             {args.warmup_iters} steps")
+    print(f"{'='*70}\n")
+
+    # --- Resume from Checkpoint (if requested or if ckpt.pt exists) ---
+    start_iter = 0
+    best_val_loss = float("inf")
+    ckpt_path = os.path.join(args.out_dir, "ckpt.pt")
+
+    if args.resume and os.path.exists(ckpt_path):
+        print(f"  [RESUME] Loading checkpoint from {ckpt_path}...")
+        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+        raw_model.load_state_dict(ckpt["model"])
+        if "optimizer" in ckpt:
+            optimizer.load_state_dict(ckpt["optimizer"])
+        start_iter = ckpt.get("iter_num", 0) + 1
+        best_val_loss = ckpt.get("best_val_loss", float("inf"))
+        print(f"  [RESUME] Resuming from step {start_iter} (best val loss: {best_val_loss:.4f})")
+
+    # --- Training Loop ---
+    os.makedirs(args.out_dir, exist_ok=True)
+    t0 = time.time()
+    tokens_processed = 0
+
+    for iter_num in range(start_iter, args.max_iters):
+        # Update learning rate (WSD schedule: warmup → stable → cosine decay)
+        lr = get_lr(iter_num, args.warmup_iters, args.lr_decay_iters, args.learning_rate, args.min_lr, args.stable_iters)
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = lr
+
+        # --- Periodic Evaluation ---
+        if iter_num % args.eval_interval == 0:
+            losses = estimate_loss(
+                model, train_data, val_data,
+                args.eval_iters, args.block_size, args.batch_size, device, ctx,
+            )
+            print(
+                f"  step {iter_num:>5d} | "
+                f"train {losses['train']:.4f} | val {losses['val']:.4f} | "
+                f"lr {lr:.2e}"
+            )
+
+            # Save latest checkpoint at every eval interval so progress is never lost
+            ckpt = {
+                "model": raw_model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "config": config,
+                "iter_num": iter_num,
+                "best_val_loss": best_val_loss,
+            }
+            torch.save(ckpt, ckpt_path)
+            print(f"           [SAVED] latest checkpoint to {ckpt_path} (step {iter_num})")
+
+            # Save separate best checkpoint when val loss improves
+            if losses["val"] < best_val_loss:
+                best_val_loss = losses["val"]
+                ckpt["best_val_loss"] = best_val_loss
+                best_ckpt_path = os.path.join(args.out_dir, "ckpt_best.pt")
+                torch.save(ckpt, best_ckpt_path)
+                print(f"           [SAVED] BEST checkpoint to {best_ckpt_path} (val_loss={best_val_loss:.4f})")
+
+        # --- Gradient Accumulation Loop ---
+        optimizer.zero_grad(set_to_none=True)
+        for micro_step in range(args.gradient_accumulation_steps):
+            X, Y = get_batch("train", train_data, val_data, args.block_size, args.batch_size, device)
+            with ctx:
+                _, loss = model(X, Y)
+                loss = loss / args.gradient_accumulation_steps
+            scaler.scale(loss).backward()
+            tokens_processed += X.numel()
+
+        # Gradient clipping
+        if args.grad_clip > 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(raw_model.parameters(), args.grad_clip)
+
+        scaler.step(optimizer)
+        scaler.update()
+
+        # --- Logging ---
+        if iter_num > 0 and iter_num % args.log_interval == 0:
+            dt = time.time() - t0
+            tps = tokens_processed / dt if dt > 0 else 0
+            lossf = loss.item() * args.gradient_accumulation_steps
+            vram = ""
+            if device == "cuda":
+                vram = f" | VRAM {torch.cuda.max_memory_allocated() / 1e9:.2f}GB"
+            print(f"  step {iter_num:>5d} | loss {lossf:.4f} | lr {lr:.2e} | {tps:,.0f} tok/s{vram}")
+            t0 = time.time()
+            tokens_processed = 0
+
+    print(f"\n{'='*70}")
+    print(f"  [DONE] Training complete! Best val loss: {best_val_loss:.4f}")
+    print(f"  Checkpoint saved to: {os.path.join(args.out_dir, 'ckpt.pt')}")
+    print(f"{'='*70}")
+
+
+# =============================================================================
+# Section 5: Text Generation / Sampling
+# =============================================================================
+
+
+def sample(args):
+    """Generate text from a trained checkpoint."""
+    device = args.device
+    if device == "auto":
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Load checkpoint
+    ckpt_path = args.ckpt or os.path.join(args.out_dir, "ckpt.pt")
+    if not os.path.exists(ckpt_path):
+        print(f"  [ERR] Checkpoint not found at {ckpt_path}")
+        print(f"    Train first with: python train_glm5.py")
+        return
+
+    print(f"  Loading checkpoint from {ckpt_path}...")
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    config = ckpt["config"]
+
+    model = GLM5ForCausalLM(config)
+    model.load_state_dict(ckpt["model"])
+    model = model.to(device)
+    model.eval()
+
+    # Encode prompt
+    enc = tiktoken.get_encoding("gpt2")
+    prompt = args.prompt or "\n"
+    tokens = enc.encode(prompt)
+    idx = torch.tensor([tokens], dtype=torch.long, device=device)
+
+    print(f"\n  Prompt: {prompt!r}")
+    print(f"  {'-'*60}")
+
+    # Generate — use bfloat16 on CUDA, float32 on CPU (autocast doesn't support bf16 on CPU)
+    if device == "cuda" and torch.cuda.is_bf16_supported():
+        gen_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
+    elif device == "cuda":
+        gen_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.float16)
+    else:
+        import contextlib
+        gen_ctx = contextlib.nullcontext()
+    with gen_ctx:
+        output = model.generate(
+            idx,
+            max_new_tokens=args.max_new_tokens,
+            temperature=args.temperature,
+            top_k=args.top_k,
+        )
+
+    generated = enc.decode(output[0].tolist())
+    print(generated)
+    print(f"  {'-'*60}\n")
+
+
+# =============================================================================
+# Section 6: Entry Point
+# =============================================================================
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Let's Reproduce GLM-5.2 (GLM MoE DSA) -- From Scratch!",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    # --- Mode ---
+    parser.add_argument("--eval_only", action="store_true", help="Generate text only (no training)")
+    parser.add_argument("--resume", action="store_true", help="Resume training from existing out_dir/ckpt.pt")
+    parser.add_argument("--ckpt", type=str, default=None, help="Checkpoint path for generation")
+
+    # --- Training ---
+    parser.add_argument("--out_dir", type=str, default="out_glm5", help="Output directory for checkpoints")
+    parser.add_argument("--data_dir", type=str, default="./data", help="Data directory containing train.bin, val.bin, meta.json")
+    parser.add_argument("--max_iters", type=int, default=800000, help="Training iterations (800K for 3.3B tokens)")
+    parser.add_argument("--batch_size", type=int, default=4, help="Micro batch size per step")
+    parser.add_argument("--block_size", type=int, default=512, help="Context/sequence length")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=4, help="Gradient accumulation steps")
+    parser.add_argument("--learning_rate", type=float, default=6e-4, help="Peak learning rate")
+    parser.add_argument("--min_lr", type=float, default=6e-5, help="Minimum learning rate (end of cosine)")
+    parser.add_argument("--warmup_iters", type=int, default=2000, help="LR warmup iterations")
+    parser.add_argument("--lr_decay_iters", type=int, default=800000, help="Cosine decay length (match max_iters)")
+    parser.add_argument("--stable_iters", type=int, default=0, help="WSD: keep LR at peak until this step, then cosine decay (0=standard cosine)")
+    parser.add_argument("--weight_decay", type=float, default=0.1, help="Weight decay")
+    parser.add_argument("--beta1", type=float, default=0.9, help="AdamW beta1")
+    parser.add_argument("--beta2", type=float, default=0.95, help="AdamW beta2")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping (0=disable)")
+
+    # --- Efficiency ---
+    parser.add_argument(
+        "--gradient_checkpointing", action="store_true", default=True,
+        help="Enable gradient checkpointing (default: on, saves VRAM)",
+    )
+    parser.add_argument(
+        "--no_gradient_checkpointing", action="store_false", dest="gradient_checkpointing",
+        help="Disable gradient checkpointing",
+    )
+    parser.add_argument("--compile", action="store_true", help="Use torch.compile (faster, needs warmup)")
+    parser.add_argument("--device", type=str, default="auto", help="Device: auto, cuda, cpu")
+
+    # --- Evaluation ---
+    parser.add_argument("--eval_interval", type=int, default=2000, help="Evaluate every N steps")
+    parser.add_argument("--eval_iters", type=int, default=50, help="Batches per evaluation")
+    parser.add_argument("--log_interval", type=int, default=100, help="Log loss every N steps")
+
+    # --- Generation ---
+    parser.add_argument("--prompt", type=str, default=None, help="Prompt for text generation")
+    parser.add_argument("--max_new_tokens", type=int, default=500, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature")
+    parser.add_argument("--top_k", type=int, default=200, help="Top-k sampling")
+
+    args = parser.parse_args()
+
+    if args.eval_only:
+        sample(args)
+    else:
+        train(args)
+        # Generate a sample after training completes
+        print("\n  >> Generating sample text from the trained model...\n")
+        args.prompt = args.prompt or "First Citizen:\n"
+        sample(args)
+
+
+if __name__ == "__main__":
+    main()