:D :D

slimshard.py

@@ -66,12 +66,13 @@ from .q4_0 import Q4_0_BLOCK_BYTES, QK4_0, dequantize_q4_0_bytes, quantize_tenso
 
 SHARD_FORMAT = "hpc-slimshard-v1"
 EXPERT_RE = re.compile(
-    r"^((?:[a-zA-Z0-9_]+\.)*(?:layers|h|blocks))\.(\d+)\.((?:moe\.|mlp\.|block_sparse_moe\.|feed_forward\.)?experts)\.(\d+)\.([a-zA-Z0-9_]+)\.(weight|bias|weight_scale_inv|weight_scale|scales)$"
+    r"^((?:[a-zA-Z0-9_]+\.)*(?:layers|h|blocks))\.(\d+)\.((?:moe\.|mlp\.|ffn\.|block_sparse_moe\.|feed_forward\.)?experts)\.(\d+)\.([a-zA-Z0-9_]+)\.(weight|bias|weight_scale_inv|weight_scale|scales|scale)$"
 )
 LAYER_RE = re.compile(r"^((?:[a-zA-Z0-9_]+\.)*(?:layers|h|blocks))\.(\d+)\.(.+)$")
 
 _EMBED_NAMES = (
     "model.embed_tokens.weight",
+    "embed.weight",
     "transformer.wte.weight",
     "gpt_neox.embed_in.weight",
     "model.decoder.embed_tokens.weight",
@@ -81,6 +82,7 @@ _EMBED_NAMES = (
 )
 _HEAD_NAMES = (
     "lm_head.weight",
+    "head.weight",
     "embed_out.weight",
     "output.weight",
     "language_model.lm_head.weight",
@@ -88,6 +90,7 @@ _HEAD_NAMES = (
 )
 _NORM_NAMES = (
     "model.norm.weight",
+    "norm.weight",
     "transformer.ln_f.weight",
     "gpt_neox.final_layer_norm.weight",
     "model.decoder.final_layer_norm.weight",
@@ -241,6 +244,8 @@ _TORCH_DTYPE = {
     "float8_e5m2": "F8_E5M2",
     "float8_e4m3fnuz": "F8_E4M3",
     "float8_e5m2fnuz": "F8_E5M2",
+    "float8_e8m0fnu": "U8",
+    "F8_E8M0": "U8",
 }
 _NP_DTYPE = {
     "F32": np.float32,
@@ -1235,10 +1240,86 @@ def _write_lm_head_calibrated_chunked(
     writer.add_chunked_bf16(name, (vocab, hidden), gen())
 
 
+def _filter_weight_map_for_target(
+    weight_map: dict[str, str],
+    present: set[str],
+    arch: dict,
+    target_params_b: float | None = None,
+    max_experts: int | None = None,
+) -> tuple[dict[str, str], int | None, float]:
+    """Filter ``weight_map`` for single-shard recovery and optional parameter-budget compression.
+
+    * Strips external multi-hundred-GB ``.engram.`` tables, ``mtp.`` draft layers,
+      and ``vision.`` / ``aligner.`` / ``image_`` branches if absent from ``present``.
+    * Computes ``kept_experts`` when ``target_params_b`` (e.g. ``30.0``) or
+      ``max_experts`` is specified, pruning routed experts ``>= kept_experts``.
+    * Returns ``(filtered_weight_map, kept_experts, estimated_params_b)``.
+    """
+    has_vision = any(k.startswith(("vision.", "aligner.", "image_")) for k in present)
+    has_mtp = any(k.startswith("mtp.") for k in present)
+    has_engram = any(".engram." in k for k in present)
+
+    orig_experts = int(arch.get("num_experts") or 0)
+    num_layers = int(arch.get("num_hidden_layers") or 25)
+    hidden = int(arch.get("hidden_size") or 2560)
+    moe_inter = int(arch.get("moe_intermediate_size") or arch.get("intermediate_size") or 1088)
+    vocab = int(arch.get("vocab_size") or 128000)
+    top_k = int(arch.get("num_experts_per_token") or arch.get("num_experts_per_tok") or 6)
+
+    # Estimate dense/non-routed parameters (embed + head + attention + shared experts)
+    tie_emb = bool(arch.get("tie_word_embeddings", False))
+    embed_params = vocab * hidden * (1 if tie_emb else 2)
+    attn_shared_per_layer = (4 * hidden * hidden) + (3 * moe_inter * hidden)
+    dense_params = float(embed_params + num_layers * attn_shared_per_layer)
+    params_per_expert = float(3 * moe_inter * hidden)
+
+    kept_experts: int | None = None
+    if max_experts is not None and max_experts > 0 and orig_experts > 0:
+        kept_experts = min(orig_experts, max(top_k, int(max_experts)))
+    elif target_params_b is not None and target_params_b > 0.0 and orig_experts > 0 and params_per_expert > 0:
+        target_p = float(target_params_b) * 1.0e9
+        avail_for_routed = max(0.0, target_p - dense_params)
+        moe_layers = max(1, num_layers - int(arch.get("first_k_dense_replace", 0)))
+        est_e = int(round(avail_for_routed / (moe_layers * params_per_expert)))
+        kept_experts = min(orig_experts, max(top_k, est_e))
+
+    filtered: dict[str, str] = {}
+    for k, sh in weight_map.items():
+        if not has_engram and ".engram." in k:
+            continue
+        if not has_mtp and k.startswith("mtp."):
+            continue
+        if not has_vision and (k.startswith(("vision.", "aligner.", "image_")) or k.endswith(".bias_vl")):
+            continue
+        em = EXPERT_RE.match(k)
+        if em and kept_experts is not None:
+            eidx = int(em.group(4))
+            if eidx >= kept_experts:
+                continue
+        filtered[k] = sh
+
+    moe_layers = max(1, num_layers - int(arch.get("first_k_dense_replace", 0)))
+    eff_experts = kept_experts if kept_experts is not None else orig_experts
+    est_total_b = (dense_params + (moe_layers * eff_experts * params_per_expert if eff_experts > 0 else 0.0)) / 1.0e9
+    return filtered, kept_experts, est_total_b
+
+
+def _is_router_gate_tensor(name: str, orig_experts: int, shape: tuple[int, ...]) -> bool:
+    if orig_experts <= 0 or not shape or shape[0] != orig_experts:
+        return False
+    return (
+        ".gate.weight" in name
+        or ".gate.bias" in name
+        or name.endswith("e_score_correction_bias")
+    ) and ".experts." not in name
+
+
 def inspect_recovery_plan(
     shard: str | Path,
     index: str | Path | None = None,
     config: str | Path | None = None,
+    target_params_b: float | None = None,
+    max_experts: int | None = None,
 ) -> dict:
     """Inspect a shard + index/config and return a full CheckFree recovery plan for the GUI."""
     shard_path = Path(shard)
@@ -1250,9 +1331,9 @@ def inspect_recovery_plan(
 
     index_path = Path(index) if index else shard_path.parent / "model.safetensors.index.json"
     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)
 
     with safe_open(str(shard_path), framework="pt", device="cpu") as f:
         meta = f.metadata() or {}
@@ -1265,6 +1346,10 @@ def inspect_recovery_plan(
             present_shapes[k] = tuple(sl.get_shape())
             present_dtypes[k] = _safetensors_dtype(sl.get_dtype())
 
+    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
+    )
+
     surviving_layers = sorted(
         {int(LAYER_RE.match(k).group(2)) for k in present if LAYER_RE.match(k)}
     )