Create convert_to_14b.py

convert_to_14b.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+# ==============================================================================
+# COPYRIGHT (C) 2026 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
+# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
+# ==============================================================================
+#
+# size_planner.py
+# ---------------
+# Compute exact parameter counts for a JiRack config, and suggest
+# (num_layers, intermediate_size) combinations that hit a target size.
+# No torch needed - pure arithmetic, so you can plan before touching a GPU.
+#
+# Examples:
+#   python size_planner.py --hidden 3072 --layers 36 --intermediate 16384
+#   python size_planner.py --hidden 3072 --target-b 14      # suggest combos
+# ==============================================================================
+
+import argparse
+
+
+def param_count(hidden, layers, intermediate, vocab=128256,
+                num_heads=24, num_kv_heads=8, tied_embeddings=False):
+    head_dim = hidden // num_heads
+    kv_dim = num_kv_heads * head_dim
+
+    attn = (hidden * hidden          # q_proj
+            + hidden * kv_dim        # k_proj
+            + hidden * kv_dim        # v_proj
+            + hidden * hidden)       # out_proj
+    ffn = 3 * intermediate * hidden  # w1, w3, w2 (SwiGLU)
+    per_layer = attn + ffn
+
+    emb = vocab * hidden * (1 if tied_embeddings else 2)  # token_emb (+ lm_head)
+    total = emb + layers * per_layer
+
+    return {
+        "total": total,
+        "embeddings": emb,
+        "per_layer": per_layer,
+        "attn_per_layer": attn,
+        "ffn_per_layer": ffn,
+        "layers_total": layers * per_layer,
+    }
+
+
+def intermediate_for_target(hidden, layers, target_params, vocab=128256,
+                            num_heads=24, num_kv_heads=8, tied_embeddings=False):
+    """Solve for the intermediate size that hits target_params at fixed layers."""
+    head_dim = hidden // num_heads
+    kv_dim = num_kv_heads * head_dim
+    attn = 2 * hidden * hidden + 2 * hidden * kv_dim
+    emb = vocab * hidden * (1 if tied_embeddings else 2)
+    # target = emb + layers*(attn + 3*I*hidden)  ->  solve I
+    budget_per_layer = (target_params - emb) / layers
+    ffn_budget = budget_per_layer - attn
+    if ffn_budget <= 0:
+        return None
+    inter = ffn_budget / (3 * hidden)
+    # round to a multiple of 256 for clean shapes
+    return int(round(inter / 256) * 256)
+
+
+def fmt(n):
+    return f"{n/1e9:.3f} B" if n >= 1e9 else f"{n/1e6:.1f} M"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--hidden", type=int, default=3072)
+    ap.add_argument("--vocab", type=int, default=128256)
+    ap.add_argument("--heads", type=int, default=24)
+    ap.add_argument("--kv-heads", type=int, default=8)
+    ap.add_argument("--tied", action="store_true", help="tie lm_head to token_emb")
+    ap.add_argument("--layers", type=int, help="report a single exact config")
+    ap.add_argument("--intermediate", type=int, help="report a single exact config")
+    ap.add_argument("--target-b", type=float, help="suggest combos for this size (in B)")
+    args = ap.parse_args()
+
+    if args.layers and args.intermediate:
+        r = param_count(args.hidden, args.layers, args.intermediate, args.vocab,
+                        args.heads, args.kv_heads, args.tied)
+        print(f"hidden={args.hidden} layers={args.layers} "
+              f"intermediate={args.intermediate} vocab={args.vocab}")
+        print(f"  embeddings      : {fmt(r['embeddings'])} "
+              f"({100*r['embeddings']/r['total']:.1f}% of total)")
+        print(f"  per layer       : {fmt(r['per_layer'])} "
+              f"(attn {fmt(r['attn_per_layer'])} + ffn {fmt(r['ffn_per_layer'])})")
+        print(f"  all layers      : {fmt(r['layers_total'])}")
+        print(f"  TOTAL           : {fmt(r['total'])}")
+        ratio = args.intermediate / args.hidden
+        print(f"  ffn expansion   : {ratio:.2f}x hidden "
+              f"({'normal' if 2.0 <= ratio <= 3.0 else 'unusual - very fat' if ratio > 4 else 'narrow'})")
+        return
+
+    if args.target_b:
+        target = args.target_b * 1e9
+        print(f"Configs hitting ~{args.target_b}B  (hidden={args.hidden}, "
+              f"vocab={args.vocab}):\n")
+        print(f"  {'layers':>6}  {'intermediate':>12}  {'ffn_ratio':>9}  {'total':>9}")
+        print("  " + "-" * 46)
+        for L in range(24, 81, 4):
+            I = intermediate_for_target(args.hidden, L, target, args.vocab,
+                                        args.heads, args.kv_heads, args.tied)
+            if I is None or I < args.hidden:  # skip infeasible / too-narrow FFN
+                continue
+            r = param_count(args.hidden, L, I, args.vocab,
+                            args.heads, args.kv_heads, args.tied)
+            ratio = I / args.hidden
+            flag = "  <- fat FFN" if ratio > 4 else ("  <- balanced" if ratio <= 3 else "")
+            print(f"  {L:>6}  {I:>12}  {ratio:>7.2f}x  {r['total']/1e9:>7.2f}B{flag}")
+        return
+
+    ap.error("provide either --layers with --intermediate, or --target-b")
+
+
+if __name__ == "__main__":
+    main()