Create convert_to_7b.py

convert_to_7b.py

@@ -0,0 +1,175 @@
+# ==============================================================================
+# COPYRIGHT (C) 2026 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
+# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
+# ==============================================================================
+#
+# grow_to_7b.py
+# -------------
+# Grow a trained JiRackNative checkpoint into a larger model by:
+#   (A) widening the FFN of every existing layer to a target intermediate size
+#       (new neurons: input side random, output side zero  -> alive, not dead)
+#   (B) inserting NEW transformer layers initialized as IDENTITY:
+#         out_proj.weight = 0 , ffn_w2.weight = 0   -> block computes x -> x
+#         q/k/v/ffn_w1/ffn_w3 = small random         -> gradient flows, learns
+#         norm1/norm2 = 1
+#       New layers are interleaved among the trained ones (not appended in one
+#       block), which trains more stably.
+#
+# Why not duplicate layers? Copied layers share identical gradients and stay
+# locked together -> no real added capacity. Identity-init new layers do add
+# capacity while preserving the function at initialization (no loss spike).
+#
+# After running:
+#   1) set NUM_LAYERS = <--layers> and INTERMEDIATE_SIZE = <--intermediate>
+#      in JiRackNative_3b.py
+#   2) load the new checkpoint and continue training (normal LR; a short warmup
+#      helps the new layers settle)
+#
+# Example (3B -> ~7B, 36 layers, intermediate 16384):
+#   python grow_to_7b.py --in jirack_3b.safetensors --out jirack_7b.safetensors \
+#          --layers 36 --intermediate 16384
+# ==============================================================================
+
+import argparse
+import glob
+import math
+import os
+import torch
+from safetensors.torch import load_file, save_file
+
+INIT_STD = 0.02
+torch.manual_seed(0)
+
+# suffix -> how to initialize a NEW layer's tensor
+#   "rand" small gaussian | "zero" zeros | "one" ones
+NEW_LAYER_INIT = {
+    "q_proj.weight":  "rand",
+    "k_proj.weight":  "rand",
+    "v_proj.weight":  "rand",
+    "out_proj.weight": "zero",   # -> attention residual = 0 at init
+    "ffn_w1.weight":  "rand",
+    "ffn_w3.weight":  "rand",
+    "ffn_w2.weight":  "zero",    # -> ffn residual = 0 at init
+    "norm1.weight":   "one",
+    "norm2.weight":   "one",
+}
+
+
+def load_state(path):
+    if os.path.isdir(path):
+        sd = {}
+        for f in sorted(glob.glob(os.path.join(path, "*.safetensors"))):
+            sd.update(load_file(f))
+        if not sd:
+            raise FileNotFoundError(f"No .safetensors in {path}")
+        return sd
+    return load_file(path)
+
+
+def detect_layers(sd):
+    idx = set()
+    for k in sd:
+        if k.startswith("blocks."):
+            idx.add(int(k.split(".")[1]))
+    return sorted(idx)
+
+
+def widen_ffn(block, target_I):
+    """In-place widen one block's FFN tensors to target intermediate size."""
+    w1 = block["ffn_w1.weight"]
+    cur_I, H = w1.shape
+    if target_I <= cur_I:
+        return block
+    add = target_I - cur_I
+    dt, dev = w1.dtype, w1.device
+    for name in ("ffn_w1.weight", "ffn_w3.weight"):        # [I, H]: new rows random
+        w = block[name]
+        new_rows = torch.randn(add, H, dtype=dt, device=dev) * INIT_STD
+        block[name] = torch.cat([w, new_rows], dim=0)
+    w2 = block["ffn_w2.weight"]                            # [H, I]: new cols zero
+    new_cols = torch.zeros(w2.shape[0], add, dtype=dt, device=dev)
+    block["ffn_w2.weight"] = torch.cat([w2, new_cols], dim=1)
+    return block
+
+
+def make_identity_block(template):
+    """Build a fresh identity-initialized block from {suffix: tensor} template."""
+    out = {}
+    for suffix, ref in template.items():
+        rule = NEW_LAYER_INIT[suffix]
+        if rule == "zero":
+            out[suffix] = torch.zeros_like(ref)
+        elif rule == "one":
+            out[suffix] = torch.ones_like(ref)
+        else:  # rand
+            out[suffix] = torch.randn_like(ref) * INIT_STD
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--in", dest="inp", required=True)
+    ap.add_argument("--out", dest="out", required=True)
+    ap.add_argument("--layers", type=int, required=True, help="target NUM_LAYERS")
+    ap.add_argument("--intermediate", type=int, default=0,
+                    help="target INTERMEDIATE_SIZE (0 = keep current)")
+    args = ap.parse_args()
+
+    print(f"Loading {args.inp} ...")
+    sd = load_state(args.inp)
+    old_idx = detect_layers(sd)
+    old_L = len(old_idx)
+    if args.layers < old_L:
+        raise ValueError(f"--layers ({args.layers}) < existing layers ({old_L})")
+
+    # --- pull each existing block into a dict of {suffix: tensor} ---
+    def block_of(j):
+        pfx = f"blocks.{j}."
+        return {k[len(pfx):]: v for k, v in sd.items() if k.startswith(pfx)}
+
+    orig_blocks = [block_of(j) for j in old_idx]
+
+    # --- (A) widen FFN of every existing block ---
+    tgt_I = args.intermediate or orig_blocks[0]["ffn_w1.weight"].shape[0]
+    for b in orig_blocks:
+        widen_ffn(b, tgt_I)
+    template = orig_blocks[0]  # correct shapes AFTER widening, for identity blocks
+
+    # --- (B) decide where trained layers sit; fill gaps with identity ---
+    L_new = args.layers
+    step = L_new / old_L
+    positions = [int(math.floor(i * step)) for i in range(old_L)]  # strictly increasing
+    pos_to_orig = {p: i for i, p in enumerate(positions)}
+
+    new_sd = {}
+    # copy non-layer tensors unchanged (token_emb, ln_f, lm_head, buffers)
+    for k, v in sd.items():
+        if not k.startswith("blocks."):
+            new_sd[k] = v
+
+    n_identity = 0
+    for n in range(L_new):
+        if n in pos_to_orig:
+            block = orig_blocks[pos_to_orig[n]]
+        else:
+            block = make_identity_block(template)
+            n_identity += 1
+        for suffix, tensor in block.items():
+            new_sd[f"blocks.{n}.{suffix}"] = tensor
+
+    # --- report ---
+    total = sum(t.numel() for t in new_sd.values())
+    print(f"Existing layers : {old_L}  -> new layers: {L_new} "
+          f"({n_identity} identity-initialized)")
+    print(f"Intermediate    : {tgt_I}")
+    print(f"Trained layers placed at indices: {positions}")
+    print(f"Total parameters: {total/1e9:.3f} B")
+
+    print(f"Saving {args.out} ...")
+    save_file(new_sd, args.out)
+    print("Done. New layers start as identity (no loss spike) and learn from ~step 2.")
+    print("Remember to update NUM_LAYERS and INTERMEDIATE_SIZE in the model file.")
+
+
+if __name__ == "__main__":
+    main()