xtrain/crates/xtrain-model/src/config.rs

//! Tiny-transformer hyperparameters. Host-only (no GPU), always compiled.

/// Architecture config for [`crate::TinyTransformer`]. Keep it tiny — T5 is a
/// correctness bring-up, not a real training run.
#[derive(Debug, Clone, Copy)]
pub struct Config {
    /// Vocabulary size (char-level in the bring-up).
    pub vocab: usize,
    /// Model / residual width. Must equal `n_heads * head_dim`.
    pub dim: usize,
    /// Number of decoder blocks.
    pub n_layers: usize,
    /// Number of attention (query) heads.
    pub n_heads: usize,
    /// Number of key/value heads (Phase T15, GQA). Each KV head is shared by a
    /// group of `n_heads / num_kv_heads` query heads (repeat_kv). Must divide
    /// `n_heads`. `num_kv_heads == n_heads` (the default) = MHA, bit-identical to
    /// the pre-T15 path; `num_kv_heads < n_heads` = real grouped-query attention,
    /// shrinking the K/V projections to `num_kv_heads * head_dim` and exported as a
    /// real `num_key_value_heads`.
    pub num_kv_heads: usize,
    /// Per-head dimension (`dim / n_heads`).
    pub head_dim: usize,
    /// SwiGLU hidden width (gate/up project to this, down projects back).
    pub ffn_hidden: usize,
    /// RMSNorm epsilon.
    pub eps: f32,
    /// RoPE base frequency (theta).
    pub rope_theta: f32,
    /// Dropout probability `p` (Phase T18). Applied at the attention/MLP sub-block
    /// outputs (before each residual add) at TRAINING time, with inverted scaling
    /// `1/(1-p)`; disabled (identity) at eval. Default `0.0` = no dropout, and the
    /// forward graph is then bit-identical to the pre-T18 path.
    pub dropout: f32,
}

impl Config {
    /// A minimal config used by the bring-up / overfit test.
    pub fn tiny() -> Self {
        let n_heads = 2;
        let head_dim = 16;
        Config {
            vocab: 0, // set by the caller from the char vocab
            dim: n_heads * head_dim,
            n_layers: 2,
            n_heads,
            num_kv_heads: n_heads, // default = MHA
            head_dim,
            ffn_hidden: 64,
            eps: 1e-5,
            rope_theta: 10000.0,
            dropout: 0.0,
        }
    }

    /// Build a config from the architecture knobs, deriving `dim = n_heads *
    /// head_dim`. The scaling-run entry (`bin/train`) passes these from CLI so the
    /// model size is a tunable ladder rung (v1 = dim256/8L, v2/v3 scale further),
    /// instead of a hardcoded tiny config. `eps`/`rope_theta` keep the engine
    /// defaults (also what the xserv export reconciles against).
    pub fn from_arch(
        vocab: usize,
        n_heads: usize,
        head_dim: usize,
        n_layers: usize,
        ffn_hidden: usize,
    ) -> Self {
        Config {
            vocab,
            dim: n_heads * head_dim,
            n_layers,
            n_heads,
            num_kv_heads: n_heads, // default = MHA; set via with_kv_heads for GQA
            head_dim,
            ffn_hidden,
            eps: 1e-5,
            rope_theta: 10000.0,
            dropout: 0.0,
        }
    }

    /// Set the number of K/V heads (Phase T15, GQA). Builder-style so existing
    /// `from_arch` call sites stay MHA unless they opt in. Asserts `num_kv_heads`
    /// divides `n_heads`.
    pub fn with_kv_heads(mut self, num_kv_heads: usize) -> Self {
        assert!(num_kv_heads > 0, "num_kv_heads must be > 0");
        assert_eq!(
            self.n_heads % num_kv_heads,
            0,
            "n_heads {} not divisible by num_kv_heads {num_kv_heads}",
            self.n_heads
        );
        self.num_kv_heads = num_kv_heads;
        self
    }

    /// KV projection width (`num_kv_heads * head_dim`). For GQA this is smaller than
    /// `dim`; for MHA it equals `dim`.
    pub fn kv_dim(&self) -> usize {
        self.num_kv_heads * self.head_dim
    }

    /// Transformer-core parameter count: everything except the token embedding and
    /// the LM head (the two `vocab × dim` tables). This is the figure the scaling
    /// ladder is sized against — the 50257-vocab embed+lm_head adds a fixed ~25M on
    /// top that does not reflect model capacity. `num_params() = core + 2·vocab·dim`.
    pub fn core_params(&self) -> usize {
        self.num_params() - 2 * self.vocab * self.dim
    }

    /// Total learnable parameter count (for logging / sanity).
    pub fn num_params(&self) -> usize {
        let per_layer = 2 * self.dim                       // 2 rmsnorm gammas
            + 2 * self.head_dim                             // q/k per-head norm gammas
            + self.dim * self.dim                           // q proj [dim,dim]
            + 2 * self.dim * self.kv_dim()                  // k/v proj [dim,kv_dim] (GQA: smaller)
            + self.dim * self.dim                           // out proj
            + 2 * self.dim * self.ffn_hidden                // gate/up proj
            + self.ffn_hidden * self.dim; // down proj
        self.vocab * self.dim          // embedding
            + self.n_layers * per_layer
            + self.dim                 // final norm
            + self.dim * self.vocab // lm head
    }
}