train_glm5.py

"""
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.
# Everything is scaled 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; this simplifies to alternating.
        self.indexer_types = [
            "full" if i % 2 == 0 else "shared" for i in range(self.num_hidden_layers)
        ]


# =============================================================================
# Section 2: Architecture Components
# =============================================================================
# Each component is built 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 (the current 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, ALL expert weights are stored
    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 each head's score be trusted?"
        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 scripts/dataprep_pretrain.py.
# Data format: train.bin / val.bin (uint16 memmap) + meta.json.
#
# Run scripts/dataprep_pretrain.py first to prepare the data:
#   python scripts/dataprep_pretrain.py              # Full 3.3B tokens
#   python scripts/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 scripts/dataprep_pretrain.py first to prepare the data:\n"
                f"          python scripts/dataprep_pretrain.py\n"
                f"        Or for a quick test:\n"
                f"          python scripts/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()