Delete convert_to_7b.py

convert_to_7b.py

@@ -1,175 +0,0 @@
-# ==============================================================================
-# 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()