Upload slimshard.py

slimshard.py

@@ -1431,7 +1431,9 @@ def inspect_recovery_plan(
         "architectures": arch.get("architectures", []),
         "hidden_size": arch.get("hidden_size", 0),
         "num_hidden_layers": arch.get("num_hidden_layers", len(all_layers)),
-        "num_experts": arch.get("num_experts", 0),
+        "num_experts": kept_experts if kept_experts is not None else arch.get("num_experts", 0),
+        "orig_num_experts": arch.get("num_experts", 0),
+        "estimated_params_b": round(est_params_b, 2),
         "vocab_size": arch.get("vocab_size", 0),
         "surviving_layers": surviving_layers,
         "missing_layers": missing_layers,
@@ -1447,6 +1449,10 @@ def export_hf_bundle(
     safetensors_path: str | Path,
     source_dir: str | Path,
     bundle_dir: str | Path,
+    kept_experts: int | None = None,
+    has_engram: bool = False,
+    has_mtp: bool = False,
+    has_vision: bool = False,
 ) -> Path:
     """Create a complete Hugging Face model directory around ``safetensors_path``."""
     st_path = Path(safetensors_path).resolve()
@@ -1461,7 +1467,6 @@ def export_hf_bundle(
     if not target_model.exists():
         os.symlink(st_path, target_model)
 
-    # Remove any multi-shard index in bundle_dir so transformers loads model.safetensors directly
     idx_out = out_dir / "model.safetensors.index.json"
     if idx_out.exists() or idx_out.is_symlink():
         idx_out.unlink()
@@ -1472,6 +1477,29 @@ def export_hf_bundle(
         if item.name.startswith("."):
             continue
         dst = out_dir / item.name
+        if item.name == "config.json" and (kept_experts is not None or not has_engram or not has_mtp):
+            try:
+                cfg_obj = json.loads(item.read_text())
+                for sub in (cfg_obj, cfg_obj.get("text_config") if isinstance(cfg_obj.get("text_config"), dict) else None):
+                    if not isinstance(sub, dict):
+                        continue
+                    if kept_experts is not None:
+                        for ek in ("num_experts", "n_routed_experts", "num_local_experts", "dspark_n_routed_experts"):
+                            if ek in sub:
+                                sub[ek] = int(kept_experts)
+                    if not has_engram:
+                        if "engram_layer_ids" in sub:
+                            sub["engram_layer_ids"] = []
+                        if "engram_num_embeddings" in sub:
+                            sub["engram_num_embeddings"] = []
+                    if not has_mtp and "num_nextn_predict_layers" in sub:
+                        sub["num_nextn_predict_layers"] = 0
+                if dst.is_symlink() or dst.exists():
+                    dst.unlink()
+                dst.write_text(json.dumps(cfg_obj, indent=2))
+                continue
+            except Exception:
+                pass
         if dst.exists() or dst.is_symlink():
             continue
         os.symlink(item.resolve(), dst)
@@ -1491,6 +1519,8 @@ def recover_from_shard(
     progress_cb=None,
     stop_event=None,
     export_hf_dir: str | Path | None = None,
+    target_params_b: float | None = None,
+    max_experts: int | None = None,
 ) -> dict:
     """Regenerate a full multi-shard HF model from **one** ``.safetensors`` shard.
 
@@ -1509,9 +1539,9 @@ def recover_from_shard(
     arch = _normalize_arch_config(raw_arch)
 
     if index_path.exists():
-        weight_map: dict[str, str] = json.loads(index_path.read_text())["weight_map"]
+        raw_weight_map: dict[str, str] = json.loads(index_path.read_text())["weight_map"]
     else:
-        weight_map = infer_weight_map_from_shard(shard_path, arch)
+        raw_weight_map = infer_weight_map_from_shard(shard_path, arch)
 
     tok_path = cfg_path.parent / "tokenizer.json"
     if not tok_path.exists():
@@ -1520,6 +1550,7 @@ def recover_from_shard(
     use_cuda = (device in ("auto", "cuda")) and torch.cuda.is_available()
     dev = torch.device("cuda:0" if use_cuda else "cpu")
     total_layers = int(arch.get("num_hidden_layers", 25))
+    orig_experts = int(arch.get("num_experts") or 0)
 
     rda_cfg = arch.get("rda_config", {})
     mla_set = {int(x) - 1 for x in rda_cfg.get("mla_layers", [])}
@@ -1548,6 +1579,12 @@ def recover_from_shard(
             }
 
         present = set(f.keys())
+        weight_map, kept_experts, est_params_b = _filter_weight_map_for_target(
+            raw_weight_map, present, arch, target_params_b=target_params_b, max_experts=max_experts
+        )
+        has_engram = any(".engram." in k for k in present)
+        has_mtp = any(k.startswith("mtp.") for k in present)
+        has_vision = any(k.startswith(("vision.", "aligner.", "image_")) for k in present)
         present_dtype: dict[str, str] = {}
         expert_anchors: dict[tuple[str, str, str, int], list[int]] = {}
         kind_suffix_anchors: dict[tuple[str, str, str], list[int]] = {}