train: train_ounce100m -- readback counts as verification; GQA kv must divide heads; pargs reports before optional probe

train/train_ounce100m.py

@@ -32,6 +32,14 @@ BASE = dict(hidden=576, layers=22, heads=9, kv_heads=3, ffn=1536, vocab=49152,
             rope_theta=10000.0, rms_eps=1e-5, tie=True)
 
 
+def _kv_dividing(heads, hidden):
+    """GQA requires num_key_value_heads to divide num_attention_heads; SDPA raises otherwise
+    ("Number of heads in key and value must divide the number of heads in Q"). Variant shapes are built
+    from hidden/64, which need not divide 9 -- take the largest divisor of `heads` that fits."""
+    per_head = max(1, hidden // 64)
+    return max(d for d in range(1, min(per_head, heads) + 1) if heads % d == 0)
+
+
 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
@@ -40,7 +48,7 @@ def build_config(seq_len, layers=None, hidden=None):
     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 * h / 3)) // 64) * 64
-    kv = BASE["kv_heads"] if hidden is None else max(1, h // 64)
+    kv = BASE["kv_heads"] if hidden is None else _kv_dividing(BASE["heads"], h)
     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=kv,