Upload slimshard.py

slimshard.py

@@ -145,10 +145,95 @@ def _is_head_key(name: str) -> bool:
     return name in _HEAD_NAMES or name.endswith(("lm_head.weight", "embed_out.weight"))
 
 
-def infer_weight_map_from_shard(shard_path: str | Path, arch: dict) -> dict[str, str]:
+def _resolve_shard_paths(shard: str | Path | list[str | Path] | tuple[str | Path, ...]) -> list[Path]:
+    """Resolve one shard path, comma/semicolon-separated shard paths, or a directory into a list of Paths."""
+    if isinstance(shard, (list, tuple)):
+        raw_items = [str(x).strip() for x in shard if str(x).strip()]
+    else:
+        s = str(shard).strip()
+        if "," in s or ";" in s:
+            raw_items = [p.strip() for p in re.split(r"[,;]", s) if p.strip()]
+        else:
+            raw_items = [s] if s else []
+
+    out: list[Path] = []
+    seen: set[Path] = set()
+    for item in raw_items:
+        p = Path(item).expanduser()
+        if p.is_dir():
+            numbered = sorted(p.glob("model-*-of-*.safetensors"))
+            cands = numbered if numbered else [
+                f for f in sorted(p.glob("*.safetensors"))
+                if not f.name.endswith("-recovered.safetensors")
+            ]
+            for c in cands:
+                rp = c.resolve()
+                if rp not in seen and c.exists():
+                    seen.add(rp)
+                    out.append(c)
+        elif p.exists():
+            rp = p.resolve()
+            if rp not in seen:
+                seen.add(rp)
+                out.append(p)
+    if not out:
+        raise FileNotFoundError(f"No valid .safetensors shard found for: {shard!r}")
+    return out
+
+
+class _MultiShardReader:
+    """Unified safe_open reader across one or multiple .safetensors shards."""
+
+    def __init__(
+        self,
+        shard_paths: str | Path | list[str | Path] | tuple[str | Path, ...],
+        anchor_layers: set[int] | None = None,
+    ) -> None:
+        self.shard_paths = _resolve_shard_paths(shard_paths)
+        self.anchor_layers = set(anchor_layers) if anchor_layers else None
+        self._handles = []
+        self._key_to_handle: dict[str, object] = {}
+        self.key_to_shard_name: dict[str, str] = {}
+        self._meta: dict[str, str] = {}
+
+    def __enter__(self) -> "_MultiShardReader":
+        for p in self.shard_paths:
+            h = safe_open(str(p), framework="pt", device="cpu")
+            self._handles.append(h)
+            m = h.metadata() or {}
+            for mk, mv in m.items():
+                self._meta.setdefault(mk, mv)
+            for k in h.keys():
+                if self.anchor_layers is not None:
+                    lm = LAYER_RE.match(k)
+                    if lm and int(lm.group(2)) not in self.anchor_layers:
+                        continue
+                if k not in self._key_to_handle:
+                    self._key_to_handle[k] = h
+                    self.key_to_shard_name[k] = p.name
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self._handles.clear()
+        self._key_to_handle.clear()
+
+    def metadata(self) -> dict[str, str]:
+        return self._meta
+
+    def keys(self) -> list[str]:
+        return list(self._key_to_handle.keys())
+
+    def get_slice(self, name: str):
+        return self._key_to_handle[name].get_slice(name)
+
+    def get_tensor(self, name: str):
+        return self._key_to_handle[name].get_tensor(name)
+
+
+def infer_weight_map_from_shard(shard_path: str | Path | list[str | Path], arch: dict) -> dict[str, str]:
     """Infer a full model ``weight_map`` when ``model.safetensors.index.json`` is absent."""
-    shard_path = Path(shard_path)
-    sh_name = shard_path.name
+    paths = _resolve_shard_paths(shard_path)
+    sh_name = paths[0].name
     num_layers = int(arch.get("num_hidden_layers", 0))
     num_experts = int(arch.get("num_experts") or 0)
     first_dense = int(arch.get("first_k_dense_replace", 0))
@@ -156,7 +241,7 @@ def infer_weight_map_from_shard(shard_path: str | Path, arch: dict) -> dict[str,
     mla_set = {int(x) - 1 for x in rda_cfg.get("mla_layers", [])}
     rda_set = {int(x) - 1 for x in rda_cfg.get("rda_layers", [])}
 
-    with safe_open(str(shard_path), framework="pt", device="cpu") as f:
+    with _MultiShardReader(paths) as f:
         present = sorted(f.keys())
 
     if num_layers <= 0:
@@ -1712,15 +1797,466 @@ def _is_router_gate_tensor(name: str, orig_experts: int, shape: tuple[int, ...])
     ) and ".experts." not in name
 
 
+def _parse_anchor_layers(anchor_layers: list[int] | set[int] | str | None) -> set[int] | None:
+    """Parse optional anchor layer specification (e.g. ``'0,4,20,39'`` or ``[0, 4, 20, 39]``)."""
+    if anchor_layers is None:
+        return None
+    if isinstance(anchor_layers, (set, list, tuple)):
+        out = {int(x) for x in anchor_layers}
+        return out if out else None
+    s = str(anchor_layers).strip()
+    if not s:
+        return None
+    out: set[int] = set()
+    for part in re.split(r"[,;\s]+", s):
+        part = part.strip()
+        if not part:
+            continue
+        if "-" in part and not part.startswith("-"):
+            a_str, b_str = part.split("-", 1)
+            for v in range(int(a_str), int(b_str) + 1):
+                out.add(v)
+        else:
+            out.add(int(part))
+    return out if out else None
+
+
+def parse_hf_repo_url(url_or_repo: str) -> tuple[str, str]:
+    """Parse a Hugging Face repo URL or ``org/repo`` string into ``(repo_id, revision)``.
+
+    Examples::
+
+        https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/tree/main -> ('deepseek-ai/DeepSeek-V4-Pro', 'main')
+        https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash-1017  -> ('deepseek-ai/DeepSeek-V4.1-Flash-1017', 'main')
+        deepseek-ai/DeepSeek-V4-Pro                                  -> ('deepseek-ai/DeepSeek-V4-Pro', 'main')
+    """
+    s = str(url_or_repo).strip()
+    s = re.sub(r"^https?://(www\.)?huggingface\.co/", "", s)
+    s = re.sub(r"^hf\.co/", "", s)
+    s = s.strip("/")
+    parts = [p for p in s.split("/") if p]
+    if len(parts) < 2:
+        raise ValueError(
+            f"Invalid Hugging Face link or repo ID: {url_or_repo!r} (expected 'org/repo' or 'https://huggingface.co/org/repo/tree/main')"
+        )
+    repo_id = f"{parts[0]}/{parts[1]}"
+    revision = "main"
+    if len(parts) >= 4 and parts[2] in ("tree", "blob", "resolve"):
+        revision = parts[3]
+    return repo_id, revision
+
+
+def _hf_resolve_url(repo_id: str, revision: str, filename: str) -> str:
+    return f"https://huggingface.co/{repo_id}/resolve/{revision}/{filename.lstrip('/')}"
+
+
+def _http_get_bytes(url: str, byte_range: tuple[int, int] | None = None, timeout: int = 120) -> bytes:
+    import urllib.request
+
+    headers = {"User-Agent": "HPC-Quantize-SlimShard/2.0"}
+    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+    if byte_range is not None:
+        headers["Range"] = f"bytes={byte_range[0]}-{byte_range[1]}"
+    req = urllib.request.Request(url, headers=headers)
+    with urllib.request.urlopen(req, timeout=timeout) as resp:
+        return resp.read()
+
+
+def plan_minimum_shards_from_hf(
+    url_or_repo: str,
+    target_params_b: float | None = None,
+    max_experts: int | None = None,
+    num_anchor_layers: int = 3,
+    anchor_layers: list[int] | set[int] | str | None = None,
+) -> dict:
+    """Analyze a Hugging Face repo link and determine the minimum shards & anchor layers needed for full reconstruction."""
+    repo_id, revision = parse_hf_repo_url(url_or_repo)
+    cfg_bytes = _http_get_bytes(_hf_resolve_url(repo_id, revision, "config.json"))
+    raw_cfg = json.loads(cfg_bytes.decode("utf-8"))
+    arch = _normalize_arch_config(raw_cfg)
+
+    try:
+        idx_bytes = _http_get_bytes(_hf_resolve_url(repo_id, revision, "model.safetensors.index.json"))
+        idx_obj = json.loads(idx_bytes.decode("utf-8"))
+        raw_weight_map: dict[str, str] = idx_obj.get("weight_map", {})
+        total_size_bytes = int(idx_obj.get("metadata", {}).get("total_size", 0))
+    except Exception:
+        raw_weight_map = {"model.embed_tokens.weight": "model.safetensors", "lm_head.weight": "model.safetensors"}
+        total_size_bytes = 0
+
+    all_repo_shards = sorted(set(raw_weight_map.values()))
+    # Strip optional vision/engram/mtp branches and prune experts to target budget
+    filtered_wm, kept_experts, est_params_b = _filter_weight_map_for_target(
+        raw_weight_map, set(), arch, target_params_b=target_params_b, max_experts=max_experts
+    )
+
+    total_layers = int(arch.get("num_hidden_layers") or 0)
+    layer_to_shards: dict[int, set[str]] = {}
+    shard_to_layers: dict[str, set[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)
+    if total_layers <= 0 and layer_to_shards:
+        total_layers = max(layer_to_shards) + 1
+
+    selected_shards: list[str] = []
+    shard_roles: dict[str, list[str]] = {}
+
+    def _add_shard(sh: str, reason: str) -> None:
+        if sh not in selected_shards:
+            selected_shards.append(sh)
+        roles = shard_roles.setdefault(sh, [])
+        if reason not in roles:
+            roles.append(reason)
+
+    # 1. Input embeddings + Output LM head + Final Norm + Output HC
+    for k, sh in filtered_wm.items():
+        if _is_embed_key(k):
+            _add_shard(sh, f"Embedding ({k})")
+        elif _is_head_key(k):
+            _add_shard(sh, f"LM Head ({k})")
+        elif k in _NORM_NAMES or (not LAYER_RE.match(k) and ("norm" in k or "hc_" in k)):
+            _add_shard(sh, f"Final Norm/Head ({k})")
+
+    # 2. Choose multi-layer depth anchors across [0 .. total_layers - 1]
+    explicit_anchors = _parse_anchor_layers(anchor_layers)
+    first_dense = int(arch.get("first_k_dense_replace", 0))
+    if explicit_anchors:
+        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()
+            for g_l in grid:
+                # Prefer a nearby layer that lives in a single shard or already-selected shard
+                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)):
+                    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)))
+                    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)
+        target_anchors.sort()
+
+    for lidx in target_anchors:
+        for sh in sorted(layer_to_shards.get(lidx, set())):
+            _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
+    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 = {
+        (EXPERT_RE.match(k).group(4), EXPERT_RE.match(k).group(5), EXPERT_RE.match(k).group(6))
+        for k in covered_keys
+        if EXPERT_RE.match(k)
+    }
+    for k, sh in filtered_wm.items():
+        em = EXPERT_RE.match(k)
+        if em:
+            sig = (em.group(4), em.group(5), em.group(6))
+            if sig not in covered_exp_projs:
+                _add_shard(sh, f"Expert template e={em.group(4)}")
+                covered_exp_projs.add(sig)
+            continue
+        lm = LAYER_RE.match(k)
+        if lm:
+            suf = lm.group(3)
+            if suf not in covered_suffixes:
+                _add_shard(sh, f"Layer template ({suf})")
+                covered_suffixes.add(suf)
+
+    # Compute exact surgical tensors needed from selected_shards
+    anchor_set = set(target_anchors)
+    surgical_tensors: dict[str, str] = {}
+    for k, sh in filtered_wm.items():
+        if sh not in selected_shards:
+            continue
+        lm = LAYER_RE.match(k)
+        if lm:
+            if int(lm.group(2)) in anchor_set:
+                surgical_tensors[k] = sh
+        else:
+            surgical_tensors[k] = sh
+
+    return {
+        "repo_id": repo_id,
+        "revision": revision,
+        "model_type": arch.get("model_type", "unknown"),
+        "architectures": arch.get("architectures", []),
+        "num_hidden_layers": total_layers,
+        "orig_num_experts": int(arch.get("num_experts") or 0),
+        "kept_experts": kept_experts if kept_experts is not None else int(arch.get("num_experts") or 0),
+        "estimated_params_b": round(est_params_b, 2),
+        "total_repo_shards": len(all_repo_shards),
+        "total_repo_size_gb": round(total_size_bytes / 1.0e9, 2) if total_size_bytes else 0.0,
+        "minimum_shards": selected_shards,
+        "shard_roles": shard_roles,
+        "anchor_layers": target_anchors,
+        "surgical_tensor_count": len(surgical_tensors),
+        "surgical_tensors": surgical_tensors,
+        "raw_config": raw_cfg,
+        "raw_index": {"metadata": {"total_size": total_size_bytes}, "weight_map": raw_weight_map},
+    }
+
+
+def fetch_minimum_shards_from_hf(
+    url_or_repo: str,
+    output_dir: str | Path,
+    target_params_b: float | None = None,
+    max_experts: int | None = None,
+    num_anchor_layers: int = 3,
+    anchor_layers: list[int] | set[int] | str | None = None,
+    mode: str = "surgical",
+    verbose: bool = True,
+    progress_cb=None,
+    stop_event=None,
+) -> dict:
+    """Retrieve the minimum shards (or surgical HTTP-Range seed shard) from a Hugging Face link.
+
+    * ``mode='surgical'`` (default): Uses HTTP Range requests to download only
+      ``embed_tokens``, ``lm_head``, ``norm``, and the selected multi-layer
+      anchors (with ``kept_experts`` experts) from the minimum shards into
+      ``<output_dir>/minimal_seed.safetensors``.
+    * ``mode='full_shards'``: Downloads the complete ``.safetensors`` files for
+      the minimum shard set into ``<output_dir>/``.
+    """
+    out_dir = Path(output_dir).expanduser().resolve()
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    plan = plan_minimum_shards_from_hf(
+        url_or_repo,
+        target_params_b=target_params_b,
+        max_experts=max_experts,
+        num_anchor_layers=num_anchor_layers,
+        anchor_layers=anchor_layers,
+    )
+    repo_id = plan["repo_id"]
+    revision = plan["revision"]
+
+    # 1. Save config.json, model.safetensors.index.json, and companion tokenizer files
+    (out_dir / "config.json").write_text(json.dumps(plan["raw_config"], indent=2))
+    (out_dir / "model.safetensors.index.json").write_text(json.dumps(plan["raw_index"], indent=2))
+    for aux in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja", "special_tokens_map.json"):
+        dst = out_dir / aux
+        if dst.exists() and dst.stat().st_size > 0:
+            continue
+        try:
+            raw_aux = _http_get_bytes(_hf_resolve_url(repo_id, revision, aux), timeout=60)
+            dst.write_bytes(raw_aux)
+        except Exception:
+            pass
+
+    min_shards: list[str] = plan["minimum_shards"]
+    stats = {
+        "repo_id": repo_id,
+        "minimum_shards": min_shards,
+        "total_repo_shards": plan["total_repo_shards"],
+        "anchor_layers": plan["anchor_layers"],
+        "bytes_downloaded": 0,
+        "tensors_fetched": 0,
+        "shard_paths": [],
+    }
+
+    if mode == "full_shards":
+        import urllib.request
+
+        total_sh = len(min_shards)
+        downloaded_paths: list[str] = []
+        for idx_s, sh_name in enumerate(min_shards, 1):
+            if stop_event is not None and stop_event.is_set():
+                raise InterruptedError("Download cancelled by user.")
+            dst_sh = out_dir / sh_name
+            url = _hf_resolve_url(repo_id, revision, sh_name)
+            headers = {"User-Agent": "HPC-Quantize-SlimShard/2.0"}
+            token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+            if token:
+                headers["Authorization"] = f"Bearer {token}"
+            req = urllib.request.Request(url, headers=headers)
+            with urllib.request.urlopen(req, timeout=300) as resp, dst_sh.open("wb") as fw:
+                while True:
+                    if stop_event is not None and stop_event.is_set():
+                        raise InterruptedError("Download cancelled by user.")
+                    chunk = resp.read(4 * 1024 * 1024)
+                    if not chunk:
+                        break
+                    fw.write(chunk)
+                    stats["bytes_downloaded"] += len(chunk)
+                    if progress_cb is not None:
+                        progress_cb(
+                            idx_s - 1,
+                            total_sh,
+                            f"{sh_name} ({stats['bytes_downloaded'] / 1e9:.2f} GB)",
+                            "DOWNLOAD_SHARD",
+                            {"present": idx_s, "checkfree_interpolated": 0, "expert_interpolated": 0, "synthesized": 0},
+                        )
+            downloaded_paths.append(str(dst_sh))
+            if progress_cb is not None:
+                progress_cb(
+                    idx_s,
+                    total_sh,
+                    sh_name,
+                    "DOWNLOAD_SHARD",
+                    {"present": idx_s, "checkfree_interpolated": 0, "expert_interpolated": 0, "synthesized": 0},
+                )
+        stats["shard_paths"] = downloaded_paths
+        return stats
+
+    # Surgical HTTP-Range mode: fetch headers for each minimum shard, then range-fetch required tensors
+    surgical_map: dict[str, str] = plan["surgical_tensors"]
+    by_shard: dict[str, list[str]] = {}
+    for tname, sh_name in surgical_map.items():
+        by_shard.setdefault(sh_name, []).append(tname)
+
+    seed_path = out_dir / "minimal_seed.safetensors"
+    writer = SafetensorsWriter(
+        seed_path,
+        metadata={
+            "format": "pt",
+            "hpc_source_repo": repo_id,
+            "hpc_minimum_shards": ",".join(min_shards),
+            "hpc_anchor_layers": ",".join(str(x) for x in plan["anchor_layers"]),
+        },
+    )
+
+    total_t = len(surgical_map)
+    done_t = 0
+    try:
+        for sh_name in min_shards:
+            t_list = by_shard.get(sh_name, [])
+            if not t_list:
+                continue
+            local_sh = out_dir / sh_name
+            if local_sh.exists():
+                with safe_open(str(local_sh), framework="pt", device="cpu") as lf:
+                    lkeys = set(lf.keys())
+                    for tname in t_list:
+                        if stop_event is not None and stop_event.is_set():
+                            raise InterruptedError("Download cancelled by user.")
+                        if tname in lkeys:
+                            t = lf.get_tensor(tname)
+                            sl = lf.get_slice(tname)
+                            writer.add(tname, _payload_bytes(t), _safetensors_dtype(sl.get_dtype()), tuple(t.shape))
+                            done_t += 1
+                            stats["tensors_fetched"] += 1
+                            if progress_cb is not None and (done_t % 8 == 0 or done_t == total_t):
+                                progress_cb(
+                                    done_t,
+                                    total_t,
+                                    tname,
+                                    "LOCAL_SHARD_HIT",
+                                    {"present": done_t, "checkfree_interpolated": 0, "expert_interpolated": 0, "synthesized": 0},
+                                )
+                continue
+
+            url = _hf_resolve_url(repo_id, revision, sh_name)
+            hdr_len_b = _http_get_bytes(url, byte_range=(0, 7), timeout=60)
+            hdr_len = struct.unpack("<Q", hdr_len_b)[0]
+            hdr_json_b = _http_get_bytes(url, byte_range=(8, 8 + hdr_len - 1), timeout=120)
+            sh_hdr = json.loads(hdr_json_b.decode("utf-8"))
+            data_base = 8 + hdr_len
+
+            # Sort tensors by file offset and coalesce adjacent byte ranges (up to 32 MB) for fast HTTP streaming
+            req_items = []
+            for tname in t_list:
+                if tname not in sh_hdr:
+                    continue
+                info = sh_hdr[tname]
+                ofs = info["data_offsets"]
+                req_items.append((int(ofs[0]), int(ofs[1]), tname, info["dtype"], tuple(int(x) for x in info["shape"])))
+            req_items.sort(key=lambda x: x[0])
+
+            batches: list[list[tuple[int, int, str, str, tuple[int, ...]]]] = []
+            for item in req_items:
+                if (
+                    batches
+                    and item[0] == batches[-1][-1][1]
+                    and (item[1] - batches[-1][0][0]) <= 32 * 1024 * 1024
+                ):
+                    batches[-1].append(item)
+                else:
+                    batches.append([item])
+
+            for batch in batches:
+                if stop_event is not None and stop_event.is_set():
+                    raise InterruptedError("Download cancelled by user.")
+                b_start = data_base + batch[0][0]
+                b_end = data_base + batch[-1][1] - 1
+                if progress_cb is not None:
+                    progress_cb(
+                        done_t,
+                        total_t,
+                        f"{batch[0][2]} ({sh_name})",
+                        f"HTTP_RANGE ({(b_end - b_start + 1) / 1e6:.1f} MB)",
+                        {"present": done_t, "checkfree_interpolated": 0, "expert_interpolated": 0, "synthesized": 0},
+                    )
+                blob = _http_get_bytes(url, byte_range=(b_start, b_end), timeout=600)
+                stats["bytes_downloaded"] += len(blob)
+                base_rel = batch[0][0]
+                for s_off, e_off, tname, dt_str, shp in batch:
+                    sub = blob[s_off - base_rel : e_off - base_rel]
+                    writer.add(tname, sub, dt_str, shp)
+                    done_t += 1
+                    stats["tensors_fetched"] += 1
+                if progress_cb is not None:
+                    progress_cb(
+                        done_t,
+                        total_t,
+                        batch[-1][2],
+                        "HTTP_RANGE_OK",
+                        {"present": done_t, "checkfree_interpolated": 0, "expert_interpolated": 0, "synthesized": 0},
+                    )
+    except Exception:
+        writer.close()
+        raise
+
+    writer.close()
+    stats["shard_paths"] = [str(seed_path)]
+    if verbose:
+        print(
+            f"Fetched surgical seed shard -> {seed_path} "
+            f"({stats['tensors_fetched']} tensors from {len(min_shards)}/{plan['total_repo_shards']} shards, "
+            f"{stats['bytes_downloaded'] / 1e9:.2f} GB downloaded)"
+        )
+    return stats
+
+
 def inspect_recovery_plan(
-    shard: str | Path,
+    shard: str | Path | list[str | Path],
     index: str | Path | None = None,
     config: 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:
-    """Inspect a shard + index/config and return a full CheckFree recovery plan for the GUI."""
-    shard_path = Path(shard)
+    """Inspect one or more shards + index/config and return a full CheckFree recovery plan for the GUI."""
+    shard_paths = _resolve_shard_paths(shard)
+    shard_path = shard_paths[0]
     cfg_path = Path(config) if config else shard_path.parent / "config.json"
     if not cfg_path.exists() and index:
         cfg_path = Path(index).parent / "config.json"
@@ -1731,9 +2267,10 @@ def inspect_recovery_plan(
     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)
 
-    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 {}
         is_slimshard = meta.get("format") == SHARD_FORMAT
         present = set(f.keys())