Upload 2 files

slimshard.py

@@ -1895,12 +1895,14 @@ def plan_minimum_shards_from_hf(
     total_layers = int(arch.get("num_hidden_layers") or 0)
     layer_to_shards: dict[int, set[str]] = {}
     shard_to_layers: dict[str, set[int]] = {}
+    layer_suffix_count: dict[int, int] = {}
     for k, sh in filtered_wm.items():
         lm = LAYER_RE.match(k)
         if lm:
             lidx = int(lm.group(2))
             layer_to_shards.setdefault(lidx, set()).add(sh)
             shard_to_layers.setdefault(sh, set()).add(lidx)
+            layer_suffix_count[lidx] = layer_suffix_count.get(lidx, 0) + 1
     if total_layers <= 0 and layer_to_shards:
         total_layers = max(layer_to_shards) + 1
 
@@ -1930,38 +1932,29 @@ def plan_minimum_shards_from_hf(
         target_anchors = sorted(l for l in explicit_anchors if l in layer_to_shards)
     else:
         n_anc = max(1, int(num_anchor_layers))
-        already_layers = sorted(
-            {
-                l
-                for sh in selected_shards
-                for l in shard_to_layers.get(sh, set())
-                if l >= first_dense
-            }
-        )
         target_anchors = []
         if first_dense > 0 and 0 in layer_to_shards:
             target_anchors.append(0)
         moe_start = min(first_dense, max(total_layers - 1, 0))
         if total_layers > 0:
-            grid = np.linspace(moe_start, max(total_layers - 1, moe_start), max(n_anc, 2), dtype=int).tolist()
+            grid = np.linspace(moe_start, max(total_layers - 1, moe_start), n_anc, dtype=int).tolist()
             for g_l in grid:
-                # Prefer a nearby layer that lives in a single shard or already-selected shard
+                # Prefer a nearby layer that lives in fewest new shards and has maximal suffixes (e.g. compressor+indexer)
                 best_l = int(g_l)
-                best_cost = (999, 999)
-                for cand_l in range(max(moe_start, int(g_l) - 3), min(total_layers, int(g_l) + 4)):
+                best_cost = (999, 999, 999)
+                for cand_l in range(max(moe_start, int(g_l) - 5), min(total_layers, int(g_l) + 6)):
                     if cand_l not in layer_to_shards:
                         continue
                     shs = layer_to_shards[cand_l]
                     new_shs = len([s for s in shs if s not in selected_shards])
-                    cost = (new_shs, abs(cand_l - int(g_l)))
+                    cost = (new_shs, -layer_suffix_count.get(cand_l, 0), abs(cand_l - int(g_l)))
                     if cost < best_cost:
                         best_cost = cost
                         best_l = cand_l
                 if best_l not in target_anchors:
                     target_anchors.append(best_l)
-        for al in already_layers:
-            if len(target_anchors) < n_anc and al not in target_anchors:
-                target_anchors.append(al)
+                for sh in sorted(layer_to_shards.get(best_l, set())):
+                    _add_shard(sh, f"Anchor Layer {best_l}")
         target_anchors.sort()
 
     for lidx in target_anchors:
@@ -1969,6 +1962,7 @@ def plan_minimum_shards_from_hf(
             _add_shard(sh, f"Anchor Layer {lidx}")
 
     # 3. Ensure every non-layer and layer template suffix in filtered_wm has at least one anchor in selected_shards
+    extra_template_keys: set[str] = set()
     covered_keys = {k for k, sh in filtered_wm.items() if sh in selected_shards}
     covered_suffixes = {LAYER_RE.match(k).group(3) for k in covered_keys if LAYER_RE.match(k)}
     covered_exp_projs = {
@@ -1983,6 +1977,7 @@ def plan_minimum_shards_from_hf(
             if sig not in covered_exp_projs:
                 _add_shard(sh, f"Expert template e={em.group(4)}")
                 covered_exp_projs.add(sig)
+                extra_template_keys.add(k)
             continue
         lm = LAYER_RE.match(k)
         if lm:
@@ -1990,6 +1985,7 @@ def plan_minimum_shards_from_hf(
             if suf not in covered_suffixes:
                 _add_shard(sh, f"Layer template ({suf})")
                 covered_suffixes.add(suf)
+                extra_template_keys.add(k)
 
     # Compute exact surgical tensors needed from selected_shards
     anchor_set = set(target_anchors)
@@ -1999,7 +1995,7 @@ def plan_minimum_shards_from_hf(
             continue
         lm = LAYER_RE.match(k)
         if lm:
-            if int(lm.group(2)) in anchor_set:
+            if int(lm.group(2)) in anchor_set or k in extra_template_keys:
                 surgical_tensors[k] = sh
         else:
             surgical_tensors[k] = sh
@@ -2566,7 +2562,7 @@ def recover_from_shard(
 
         writer = SafetensorsWriter(
             output,
-            metadata={"format": "pt", "hpc_checkfree_source": str(shard_path.name)},
+            metadata={"format": "pt", "hpc_checkfree_source": ",".join(p.name for p in shard_paths)},
         )
         stats = {
             "present": 0,
@@ -3868,16 +3864,26 @@ def main() -> None:
     i = sub.add_parser("info", help="print SlimShard metadata")
     i.add_argument("--input", required=True)
 
-    rec = sub.add_parser("recover", help="CheckFree single-shard model recovery")
-    rec.add_argument("--shard", required=True, help="single surviving HF .safetensors shard")
+    rec = sub.add_parser("recover", help="CheckFree single/multi-shard model recovery")
+    rec.add_argument("--shard", required=True, help="surviving HF .safetensors shard(s) (comma-separated or directory)")
     rec.add_argument("--index", default=None, help="model.safetensors.index.json (auto-detected or inferred if omitted)")
     rec.add_argument("--output", required=True, help="full model .safetensors")
     rec.add_argument("--config", default=None, help="config.json (defaults to index/shard dir)")
     rec.add_argument("--export-hf-dir", default=None, help="optional ready-to-load HF model directory")
     rec.add_argument("--target-params-b", type=float, default=None, help="optional target parameter budget in billions (e.g. 30.0)")
     rec.add_argument("--max-experts", type=int, default=None, help="optional max routed experts per layer to keep")
+    rec.add_argument("--anchor-layers", default=None, help="optional comma-separated anchor layer indices (e.g. 0,4,20,39)")
     rec.add_argument("--device", default="auto", choices=("auto", "cuda", "cpu"))
 
+    fhf = sub.add_parser("fetch-hf", help="retrieve minimum shards needed for full reconstruction from a Hugging Face link")
+    fhf.add_argument("--repo", required=True, help="Hugging Face repo URL or org/repo (e.g. https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/tree/main)")
+    fhf.add_argument("--output-dir", required=True, help="directory to store retrieved minimum shards and config/tokenizer")
+    fhf.add_argument("--target-params-b", type=float, default=None, help="optional target parameter budget in billions")
+    fhf.add_argument("--max-experts", type=int, default=None, help="optional max routed experts per layer")
+    fhf.add_argument("--num-anchor-layers", type=int, default=3, help="number of depth anchor layers to retrieve")
+    fhf.add_argument("--anchor-layers", default=None, help="explicit comma-separated anchor layers (e.g. 0,14,39)")
+    fhf.add_argument("--mode", choices=("surgical", "full_shards"), default="surgical")
+
     eg = sub.add_parser("export-gguf", help="convert a SafeTensors model (FP4/FP8/BF16) to GGUF")
     eg.add_argument("--input", required=True, help="input .safetensors file")
     eg.add_argument("--output", required=True, help="output .gguf file")
@@ -3913,6 +3919,18 @@ def main() -> None:
             export_hf_dir=args.export_hf_dir,
             target_params_b=args.target_params_b,
             max_experts=args.max_experts,
+            anchor_layers=args.anchor_layers,
+        )
+    elif args.cmd == "fetch-hf":
+        fetch_minimum_shards_from_hf(
+            args.repo,
+            args.output_dir,
+            target_params_b=args.target_params_b,
+            max_experts=args.max_experts,
+            num_anchor_layers=args.num_anchor_layers,
+            anchor_layers=args.anchor_layers,
+            mode=args.mode,
+            verbose=True,
         )
     elif args.cmd == "export-gguf":
         export_safetensors_to_gguf(

slimshard_gui.py

@@ -1010,7 +1010,7 @@ class SlimShardGUI:
             return
         mode = self.op_mode_var.get()
         out = self.output_var.get().strip()
-        if not out:
+        if mode != "fetch-hf" and not out:
             messagebox.showwarning("SlimShard", "Please specify an output .safetensors path.")
             return
 
@@ -1054,10 +1054,40 @@ class SlimShardGUI:
                         export_hf_dir=exp_dir,
                         target_params_b=tp if tp > 0.0 else None,
                         max_experts=me if me > 0 else None,
+                        anchor_layers=self.anchor_layers_var.get().strip() or None,
                     )
                     if exp_dir:
                         self._queue.put(("set_test_dir", exp_dir))
                     self._queue.put(("done", ("CheckFree Recovery Complete", res)))
+                elif mode == "fetch-hf":
+                    url = self.hf_url_var.get().strip()
+                    mdir = self.model_dir_var.get().strip() or str(Path.cwd() / "hf_min_shards")
+                    try:
+                        tp = float(self.target_params_b_var.get())
+                    except Exception:
+                        tp = 0.0
+                    try:
+                        me = int(self.max_experts_var.get())
+                    except Exception:
+                        me = 0
+                    try:
+                        na = int(self.num_anchor_layers_var.get())
+                    except Exception:
+                        na = 3
+                    res = fetch_minimum_shards_from_hf(
+                        url,
+                        mdir,
+                        target_params_b=tp if tp > 0.0 else None,
+                        max_experts=me if me > 0 else None,
+                        num_anchor_layers=na,
+                        anchor_layers=self.anchor_layers_var.get().strip() or None,
+                        mode=self.hf_fetch_mode_var.get(),
+                        verbose=False,
+                        progress_cb=progress_cb,
+                        stop_event=self._stop_event,
+                    )
+                    self._queue.put(("hf_fetched", (mdir, res)))
+                    self._queue.put(("done", ("HF Minimum Shards Retrieved", res)))
                 elif mode == "build":
                     mdir = self.model_dir_var.get().strip()
                     cfg_obj = HPCConfig(device=self.device_var.get(), verbose=False).validate()
@@ -1346,6 +1376,14 @@ class SlimShardGUI:
                         self.tree.item(tname, values=vals)
                 elif kind == "set_test_dir":
                     self.test_model_dir_var.set(str(payload))
+                elif kind == "hf_fetched":
+                    mdir, res = payload
+                    sh_paths = res.get("shard_paths", [])
+                    if sh_paths:
+                        self._set_shard_and_companions(Path(sh_paths[0]))
+                        self.shard_var.set(", ".join(str(p) for p in sh_paths))
+                    self.op_mode_var.set("recover")
+                    self._inspect_plan()
                 elif kind == "done":
                     title, stats = payload
                     elapsed = time.time() - (self._start_time or time.time())