Upload scripts/extract_300.py with huggingface_hub

scripts/extract_300.py

@@ -0,0 +1,73 @@
+"""Extrae los 300 expertos seleccionados por capa del 397B-FP8 -> modelo 300exp-FP8.
+Prunea experts + gate (router 512->300), mantiene shared_expert/atencion/vision, quita MTP.
+"""
+import json, os, re, time
+import numpy as np
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+MODEL="/model"; OUT="/model300"
+os.makedirs(OUT, exist_ok=True)
+sel=np.load("/expert_selection.npy")  # [L,300]
+L,N=sel.shape
+print(f"seleccion [{L},{N}]", flush=True)
+idx=json.load(open(f"{MODEL}/model.safetensors.index.json"))
+wmap=idx["weight_map"]
+files=sorted(set(wmap.values()))
+handles={f:safe_open(os.path.join(MODEL,f),framework="pt",device="cpu") for f in files}
+def get(k): return handles[wmap[k]].get_tensor(k)
+
+# mapa: experto original -> nuevo indice, por capa
+remap={l:{int(o):n for n,o in enumerate(sel[l])} for l in range(L)}
+
+new={}; t0=time.time()
+exp_re=re.compile(r"^model\.language_model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(.+)$")
+gate_re=re.compile(r"^model\.language_model\.layers\.(\d+)\.mlp\.gate\.weight$")
+for ki,k in enumerate(wmap.keys()):
+    if k.startswith("mtp."):
+        continue  # quitar MTP
+    me=exp_re.match(k)
+    if me:
+        l,e,rest=int(me.group(1)),int(me.group(2)),me.group(3)
+        if e in remap[l]:
+            ne=remap[l][e]
+            new[f"model.language_model.layers.{l}.mlp.experts.{ne}.{rest}"]=get(k)
+        continue
+    mg=gate_re.match(k)
+    if mg:
+        l=int(mg.group(1))
+        w=get(k)  # [512, hidden]
+        new[k]=w[sel[l]].contiguous()  # [300, hidden]
+        continue
+    # resto (atencion, shared_expert, norms, embeds, vision, lm_head): copiar
+    new[k]=get(k)
+    if ki%5000==0: print(f"  {ki} tensores, {time.time()-t0:.0f}s", flush=True)
+
+print(f"nuevo state: {len(new)} tensores. guardando shards...", flush=True)
+# guardar en shards ~5GB
+items=list(new.items()); shard=[]; sz=0; si=0; sidx={}
+SHARD_MAX=5e9
+def flush_shard(shard,si):
+    fn=f"model-{si:05d}.safetensors"
+    save_file(dict(shard), os.path.join(OUT,fn), metadata={"format":"pt"})
+    for kk,_ in shard: sidx[kk]=fn
+    return fn
+for k,v in items:
+    shard.append((k,v)); sz+=v.numel()*v.element_size()
+    if sz>=SHARD_MAX:
+        flush_shard(shard,si); print(f"  shard {si} ({sz/1e9:.1f}GB)",flush=True); shard=[]; sz=0; si+=1
+if shard: flush_shard(shard,si)
+# index + config
+json.dump({"metadata":{},"weight_map":sidx}, open(os.path.join(OUT,"model.safetensors.index.json"),"w"))
+cfg=json.load(open(f"{MODEL}/config.json"))
+cfg["text_config"]["num_experts"]=N
+cfg["text_config"]["mtp_num_hidden_layers"]=0
+json.dump(cfg, open(os.path.join(OUT,"config.json"),"w"), indent=2)
+# copiar tokenizer/processor configs
+import shutil
+for f in os.listdir(MODEL):
+    if f.endswith(".json") and "index" not in f and f!="config.json" or "token" in f.lower() or "processor" in f.lower() or "merges" in f or "vocab" in f:
+        try: shutil.copy(os.path.join(MODEL,f), os.path.join(OUT,f))
+        except: pass
+print(f"DONE: 300exp guardado en {OUT}, {si+1} shards, {time.time()-t0:.0f}s", flush=True)