Upload slimshard.py

slimshard.py

@@ -1461,18 +1461,26 @@ def _write_lm_head_calibrated_chunked(
         and "model.layers.0.input_layernorm.weight" in writer._entries
         and "model.layers.0.self_attn_raon_mhc.phi_pre" in writer._entries
     )
+    is_deepseek_v41 = (
+        "layers.0.hc_attn_fn" in writer._entries
+        or "model.layers.0.hc_attn_fn" in writer._entries
+    )
     has_generic_layers = any(LAYER_RE.match(k) for k in writer._entries)
 
     h_states_list: list["torch.Tensor"] = []
     target_ids: list[int] = []
     bg_states_list: list["torch.Tensor"] = []
 
-    if is_raon or has_generic_layers:
+    if is_raon or is_deepseek_v41 or has_generic_layers:
         writer._tmp.flush()
         raw_tokens = ["hello", '"', " is", ' "', "Hello", ",", " World", '!".']
         if im_end_id >= 0:
             raw_tokens.append("<|im_end|>")
+        elif int(arch.get("eos_token_id", -1)) >= 0:
+            raw_seq_eos = int(arch.get("eos_token_id", 1))
         raw_seq = [tok_id(x) for x in raw_tokens]
+        if im_end_id < 0 and int(arch.get("eos_token_id", -1)) >= 0:
+            raw_seq.append(int(arch["eos_token_id"]))
 
         chat_prefix_tokens = [
             "<|im_start|>", "user", "\n", "hello", "<|im_end|>", "\n",
@@ -1502,6 +1510,17 @@ def _write_lm_head_calibrated_chunked(
                     outs = _run_exact_reconstructed_forward(
                         writer, seqs_to_run, coords, basis_in, scale_in, arch, dev
                     )
+                elif is_deepseek_v41:
+                    outs = _run_deepseek_v41_forward(
+                        lambda k: writer.read_tensor_torch(k, dev),
+                        lambda k: k in writer._entries,
+                        seqs_to_run,
+                        coords,
+                        basis_in,
+                        scale_in,
+                        arch,
+                        dev,
+                    )
                 else:
                     outs = _run_generic_transformer_forward(
                         writer, seqs_to_run, coords, basis_in, scale_in, arch, dev
@@ -2204,10 +2223,23 @@ def recover_from_shard(
                 _notify(name, "SYNTHESIZE_CALIBRATED")
                 continue
 
-            tail_match = (
-                by_tail.get(".".join(name.split(".")[-2:]))
-                or by_tail.get(".".join(name.split(".")[-3:]))
-            )
+            tail_match = None
+            if name in _NORM_NAMES or "norm" in name:
+                hidden_sz = int(arch.get("hidden_size", 0))
+                if hidden_sz > 0 and name in _NORM_NAMES:
+                    tail_match = next(
+                        (
+                            n
+                            for n in sorted(present)
+                            if "norm" in n and tuple(f.get_slice(n).get_shape()) == (hidden_sz,)
+                        ),
+                        None,
+                    )
+            if tail_match is None:
+                tail_match = (
+                    by_tail.get(".".join(name.split(".")[-2:]))
+                    or by_tail.get(".".join(name.split(".")[-3:]))
+                )
             if tail_match is None and "norm" in name:
                 tail_match = next((n for n in present if "norm" in n), None)
             if tail_match is None and "bias" in name: