Create convert_to_3b.py

convert_to_3b.py

@@ -0,0 +1,111 @@
+# ==============================================================================
+# COPYRIGHT (C) 2026 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
+# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
+# ==============================================================================
+#
+# expand_3b.py
+# ------------
+# Widen a trained JiRackNative_3b checkpoint's FFN from INTERMEDIATE=4096 (~2B)
+# to INTERMEDIATE=8192 (~3B) so the NEW neurons are actually alive.
+#
+# Because hidden_size is unchanged (only the FFN width grows), this is a CLEAN,
+# function-preserving expansion:
+#   * ffn_w1 / ffn_w3  (shape [I, H]) : new ROWS  = small random  (input side)
+#   * ffn_w2           (shape [H, I]) : new COLS  = ZERO           (output side)
+# At init the new neurons contribute 0 to the output (loss is unchanged), but
+# their input weights are non-zero, so gradient flows and they start learning
+# from ~step 2. No dead neurons, no loss spike.  (Net2Net, Chen et al. 2016.)
+#
+# After running this:
+#   1) set INTERMEDIATE_SIZE = 8192 in JiRackNative_3b.py
+#   2) load the new checkpoint and continue training (keep your normal LR)
+#
+# Usage:
+#   python expand_3b.py --in jirack_2b.safetensors --out jirack_3b.safetensors
+#   python expand_3b.py --in ckpt/ --out jirack_3b.safetensors   # dir of shards
+# ==============================================================================
+
+import argparse
+import glob
+import os
+import torch
+from safetensors.torch import load_file, save_file
+
+OLD_I = 4096          # current intermediate size (~2B)
+NEW_I = 8192          # target intermediate size (~3B)
+INIT_STD = 0.02       # std for the new input-side weights
+torch.manual_seed(0)  # reproducible expansion
+
+
+def expand_ffn_in(w):
+    """ffn_w1 / ffn_w3, shape [I, H]: add new ROWS (new neurons' input) = random."""
+    I, H = w.shape
+    if I >= NEW_I:
+        return w
+    new_rows = torch.randn(NEW_I - I, H, dtype=w.dtype, device=w.device) * INIT_STD
+    return torch.cat([w, new_rows], dim=0)
+
+
+def expand_ffn_out(w):
+    """ffn_w2, shape [H, I]: add new COLS (new neurons' output) = ZERO."""
+    H, I = w.shape
+    if I >= NEW_I:
+        return w
+    new_cols = torch.zeros(H, NEW_I - I, dtype=w.dtype, device=w.device)
+    return torch.cat([w, new_cols], dim=1)
+
+
+def load_state(path):
+    if os.path.isdir(path):
+        sd = {}
+        files = sorted(glob.glob(os.path.join(path, "*.safetensors")))
+        if not files:
+            raise FileNotFoundError(f"No .safetensors in {path}")
+        for f in files:
+            sd.update(load_file(f))
+        return sd
+    return load_file(path)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--in", dest="inp", required=True,
+                    help="input .safetensors file OR directory of shards")
+    ap.add_argument("--out", dest="out", required=True,
+                    help="output .safetensors path")
+    args = ap.parse_args()
+
+    print(f"Loading checkpoint from {args.inp} ...")
+    sd = load_state(args.inp)
+
+    new_sd = {}
+    n_w1 = n_w3 = n_w2 = 0
+
+    for k, v in sd.items():
+        if k.endswith("ffn_w1.weight"):
+            new_sd[k] = expand_ffn_in(v);  n_w1 += 1
+        elif k.endswith("ffn_w3.weight"):
+            new_sd[k] = expand_ffn_in(v);  n_w3 += 1
+        elif k.endswith("ffn_w2.weight"):
+            new_sd[k] = expand_ffn_out(v); n_w2 += 1
+        else:
+            new_sd[k] = v  # embeddings, attention, norms, lm_head unchanged
+
+    print(f"Expanded: ffn_w1={n_w1}  ffn_w3={n_w3}  ffn_w2={n_w2}  "
+          f"({OLD_I} -> {NEW_I})")
+
+    # --- sanity check: confirm new neurons are "alive" (nonzero input side) ---
+    sample = next(k for k in new_sd if k.endswith("ffn_w1.weight"))
+    w1 = new_sd[sample]
+    new_block = w1[OLD_I:NEW_I]
+    live = (new_block.abs().sum(dim=1) > 0).sum().item()
+    print(f"Sanity ({sample}): {live}/{NEW_I - OLD_I} new neurons have "
+          f"nonzero input weights (should be all).")
+
+    print(f"Saving to {args.out} ...")
+    save_file(new_sd, args.out)
+    print("Done. Now set INTERMEDIATE_SIZE = 8192 and continue training.")
+
+
+if __name__ == "__main__":
+    main()