From 15f1e526c71e2116bbd471f60ccc74c36813649e Mon Sep 17 00:00:00 2001
From: Gahow Wang <gahow.wang@gmail.com>
Date: Mon, 15 Jun 2026 18:34:39 +0800
Subject: [PATCH] train: parameterize model size (scaling ladder)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn) so the model
size is a tunable rung instead of a hardcoded tiny config, and Config::core_params()
(num_params minus the two vocab×dim tables) — the figure the ladder is sized
against (the 50257-vocab embed+lm_head adds a fixed ~25M that is not capacity).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 crates/xtrain-model/src/config.rs | 32 +++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/crates/xtrain-model/src/config.rs b/crates/xtrain-model/src/config.rs
index d5aac59..3eba07a 100644
--- a/crates/xtrain-model/src/config.rs
+++ b/crates/xtrain-model/src/config.rs
@@ -39,6 +39,38 @@ impl Config {
         }
     }
 
+    /// 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,
+            head_dim,
+            ffn_hidden,
+            eps: 1e-5,
+            rope_theta: 10000.0,
+        }
+    }
+
+    /// 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