Upload slimshard.py

slimshard.py

@@ -2445,7 +2445,7 @@ def export_hf_bundle(
 # CheckFree Single-Shard Recovery
 # ---------------------------------------------------------------------------
 def recover_from_shard(
-    shard: str | Path,
+    shard: str | Path | list[str | Path],
     index: str | Path | None,
     output: str | Path,
     config: str | Path | None = None,
@@ -2456,16 +2456,18 @@ def recover_from_shard(
     export_hf_dir: str | Path | None = None,
     target_params_b: float | None = None,
     max_experts: int | None = None,
+    anchor_layers: list[int] | set[int] | str | None = None,
 ) -> dict:
-    """Regenerate a full multi-shard HF model from **one** ``.safetensors`` shard.
+    """Regenerate a full multi-shard HF model from one or more ``.safetensors`` shards.
 
-    Strictly uses only ``shard`` (plus ``index``, ``config``, and ``tokenizer.json``
-    metadata). Does **not** open any other ``.safetensors`` shard.
+    Supports both single-shard recovery and multi-shard / multi-layer anchor
+    recovery (e.g. from the minimum shard set retrieved from Hugging Face).
     """
     if safe_open is None or torch is None:
         raise RuntimeError("safetensors + torch are required for recovery")
 
-    shard_path = Path(shard)
+    shard_paths = _resolve_shard_paths(shard)
+    shard_path = shard_paths[0]
     index_path = Path(index) if index else shard_path.parent / "model.safetensors.index.json"
     cfg_path = Path(config) if config else (index_path.parent if index_path.exists() else shard_path.parent) / "config.json"
     if not cfg_path.exists():
@@ -2476,7 +2478,7 @@ def recover_from_shard(
     if index_path.exists():
         raw_weight_map: dict[str, str] = json.loads(index_path.read_text())["weight_map"]
     else:
-        raw_weight_map = infer_weight_map_from_shard(shard_path, arch)
+        raw_weight_map = infer_weight_map_from_shard(shard_paths, arch)
 
     tok_path = cfg_path.parent / "tokenizer.json"
     if not tok_path.exists():
@@ -2498,9 +2500,10 @@ def recover_from_shard(
             return "rda"
         return "generic"
 
-    with safe_open(str(shard_path), framework="pt", device="cpu") as f:
+    parsed_anchors = _parse_anchor_layers(anchor_layers)
+    with _MultiShardReader(shard_paths, anchor_layers=parsed_anchors) as f:
         meta = f.metadata() or {}
-        if meta.get("format") == SHARD_FORMAT:
+        if len(shard_paths) == 1 and meta.get("format") == SHARD_FORMAT:
             res = regenerate(shard_path, output, verbose=verbose, progress_cb=progress_cb, stop_event=stop_event)
             return {
                 "present": res["tensors"],
@@ -2742,7 +2745,10 @@ def recover_from_shard(
                     if kept_experts is not None and _is_router_gate_tensor(name, orig_experts, tuple(ta.shape)):
                         ta = ta[:kept_experts].contiguous()
                     if la == lb or alpha <= 1.0e-6 or dt_t not in ("F32", "F16", "BF16"):
-                        writer.add(name, _payload_bytes(ta), dt_t, tuple(ta.shape))
+                        t_sel = ta if (la == lb or alpha < 0.5) else f.get_tensor(nb)
+                        if kept_experts is not None and _is_router_gate_tensor(name, orig_experts, tuple(t_sel.shape)):
+                            t_sel = t_sel[:kept_experts].contiguous()
+                        writer.add(name, _payload_bytes(t_sel), dt_t, tuple(t_sel.shape))
                     else:
                         tb = f.get_tensor(nb)
                         if kept_experts is not None and _is_router_gate_tensor(name, orig_experts, tuple(tb.shape)):