AGILLM4_include_NAT_in_param_estimates

AGILLM-4.md

@@ -22,9 +22,9 @@ Implemented presets:
 
 | Preset | Shape | Approx params | Use |
 | --- | --- | ---: | --- |
-| `agillm4_floor` | d=1280, L=28, H=20, rank=160 | ~1.05B | minimum AGILLM-4 run |
-| `agillm4_main` | d=1536, L=32, H=24, rank=192 | ~1.5B | main target |
-| `agillm4_big` | d=1792, L=36, H=28, rank=224 | ~2.1B | stretch target after memory works |
+| `agillm4_floor` | d=1280, L=28, H=20, rank=160 | ~1.21B | minimum AGILLM-4 run |
+| `agillm4_main` | d=1536, L=32, H=24, rank=192 | ~1.70B | main target |
+| `agillm4_big` | d=1792, L=36, H=28, rank=224 | ~2.40B | stretch target after memory works |
 
 Default recommendation: train `agillm4_main` if B200/B300 availability is good. On a 24GB 4090, start with `agillm4_floor` so the run is still larger than AGILLM-3 while leaving enough VRAM for AR+SAT+NAT.
 

README.md

@@ -14,7 +14,7 @@ AGILLM-4 is the next training target after AGILLM-3. The current code is a
 production-oriented starting point, copied from the proven single-file trainer
 and extended for:
 
-- >1B parameter floor preset (`agillm4_floor`) and ~1.5B main preset (`agillm4_main`)
+- >1B parameter floor preset (`agillm4_floor`) and ~1.7B main preset (`agillm4_main`) with AR+SAT+NAT heads
 - 100 tokens per parameter target ratio, above the AGILLM-3 training ratio
 - longer block-size work on 24GB, B200, and B300 class GPUs
 - AR+SAT+NAT training, with sequential backward to reduce peak VRAM

estimate_agillm4_params.py

@@ -25,8 +25,9 @@ def estimate(vocab: int, d: int, layers: int, heads: int, rank: int, tie_weights
     block = attn + ff + norms
     core = embed + layers * block + 2 * d
     ar = 0 if tie_weights else vocab * d + vocab
+    nat = vocab * d + vocab
     sat = vocab * d + vocab + 2 * d + 2
-    total = core + ar + sat
+    total = core + ar + sat + nat
     return {
         "vocab": vocab,
         "d_model": d,
@@ -38,6 +39,7 @@ def estimate(vocab: int, d: int, layers: int, heads: int, rank: int, tie_weights
         "block_params_each": block,
         "core_params": core,
         "ar_head_params": ar,
+        "nat_head_params": nat,
         "sat_head_params": sat,
         "total_params": total,
         "tokens_at_100_to_1": total * 100,

nB300_agillm4.py

@@ -832,8 +832,8 @@ PRESETS: Dict[str, Dict[str, int]] = {
     "base18":    dict(d=768, layers=18, heads=24, rank=96),
     "large":     dict(d=1024, layers=24, heads=16, rank=128),
     # AGILLM-4 tiers. These are intentionally above the ~700M AGILLM-3 size.
-    # Approx dense parameter count with the current untied embedding+AR+SAT heads:
-    # agillm4_floor ~= 1.05B, agillm4_main ~= 1.50B, agillm4_big ~= 2.1B.
+    # Approx dense parameter count with the current untied embedding+AR+SAT+NAT heads:
+    # agillm4_floor ~= 1.21B, agillm4_main ~= 1.70B, agillm4_big ~= 2.40B.
     "agillm4_floor": dict(d=1280, layers=28, heads=20, rank=160),
     "agillm4_main":  dict(d=1536, layers=32, heads=24, rank=192),
     "agillm4_big":   dict(d=1792, layers=36, heads=28, rank=224),