train: fix SwiGLU width in build_config variants (was 2x)

train/train_ounce100m.py

@@ -33,19 +33,21 @@ BASE = dict(hidden=576, layers=22, heads=9, kv_heads=3, ffn=1536, vocab=49152,
 
 
 def build_config(seq_len, layers=None, hidden=None):
+    """The frozen shape, and the two variants test T1 needs (20L at the same width, and a wider/shallower
+    alternative). SwiGLU width is int(2*4h/3) rounded to a 64-multiple -- 1536 at h=576, which is the
+    value the frozen 106,194,240 count was derived from. Getting this factor wrong silently changes the
+    parameter count and makes T1 compare two different models."""
     h = hidden or BASE["hidden"]
     L = layers or BASE["layers"]
-    inter = BASE["ffn"] if hidden is None and layers is None else int(round(8 / 3 * h)) * 2 // 64 * 64
+    inter = BASE["ffn"] if hidden is None and layers is None else (int(round(8 * h / 3)) // 64) * 64
+    kv = BASE["kv_heads"] if hidden is None else max(1, h // 64)
     cfg = LlamaConfig(vocab_size=BASE["vocab"], hidden_size=h, intermediate_size=inter,
                       num_hidden_layers=L, num_attention_heads=BASE["heads"],
-                      num_key_value_heads=BASE["kv_heads"] if hidden is None else max(1, h // 64),
+                      num_key_value_heads=kv,
                       max_position_embeddings=seq_len, rope_theta=BASE["rope_theta"],
                       rms_norm_eps=BASE["rms_eps"], tie_word_embeddings=BASE["tie"],
                       attention_dropout=0.0, hidden_act="silu", mlp_bias=False,
                       pad_token_id=BASE["vocab"] - 1, eos_token_id=0)
-    # GQA ratio is part of the frozen shape; 9 heads with 3 kv heads, never MHA at this budget.
-    if hidden is None:
-        cfg.num_key_value_heads = BASE["kv_heads"]
     return cfg