Upload slimshard.py

slimshard.py

@@ -1061,7 +1061,7 @@ def _run_deepseek_v41_forward(
         del emb_w
 
     # Cache layer weights if multiple layers are identical (e.g. single-shard anchor recovery)
-    cached_layer_ signature = None
+    cached_layer_signature = None
     cached_layer_data = None
 
     for lidx in range(num_layers):
@@ -1270,6 +1270,19 @@ def _synthesize_embed_tokens_checkfree(
             t = f_handle.get_tensor(cand).to(dev, dtype=torch.float32)
             if t.ndim == 2 and t.shape[1] == hidden:
                 in_mats.append(t[: min(384, t.shape[0])])
+    if not in_mats:
+        for cand in sorted(present):
+            if cand.endswith(".weight") and ".experts." not in cand:
+                sl = f_handle.get_slice(cand)
+                shp = tuple(sl.get_shape())
+                if len(shp) == 2 and shp[1] == hidden:
+                    t = f_handle.get_tensor(cand).to(dev)
+                    sc_key = cand[:-7] + ".scale"
+                    if str(t.dtype).endswith("e4m3fn") and sc_key in present:
+                        t = _ds4_dequant_fp8_32x32(t, f_handle.get_tensor(sc_key).to(dev))
+                    in_mats.append(t.float()[: min(384, t.shape[0])])
+                    if len(in_mats) >= 4:
+                        break
     k_dim = min(128, hidden)
     if in_mats:
         cat_in = torch.cat(in_mats, dim=0)
@@ -1411,13 +1424,19 @@ def _write_lm_head_calibrated_chunked(
             (
                 "self_attn.o_proj.weight",
                 "attn.o_proj.weight",
+                "attn.wo_b.weight",
                 "attention.wo.weight",
                 "shared_experts.down_proj.weight",
+                "shared_experts.w2.weight",
                 "mlp.down_proj.weight",
                 "feed_forward.w2.weight",
             )
         ):
-            t = f_handle.get_tensor(cand).to(dev, dtype=torch.float32)
+            t = f_handle.get_tensor(cand).to(dev)
+            sc_key = cand[:-7] + ".scale"
+            if str(t.dtype).endswith("e4m3fn") and sc_key in present:
+                t = _ds4_dequant_fp8_32x32(t, f_handle.get_tensor(sc_key).to(dev))
+            t = t.float()
             if t.ndim == 2 and t.shape[0] == hidden:
                 out_mats.append(t[:, : min(384, t.shape[1])].T)
     if out_mats: