DS tweaks

slimshard.py

@@ -3051,40 +3051,10 @@ def export_safetensors_to_gguf(
             gname = f"blk.{lidx}.{mapped_proj}"
             work_items.append(("experts", (lidx, proj), members, gname))
 
-        # Stream quantized/converted tensor payloads into a temporary data file and record GGUF directory
-        tmp_data = tempfile.NamedTemporaryFile(delete=False, dir=str(out_path.parent))
-        gguf_tensors: list[tuple[str, tuple[int, ...], int, int]] = []
-        data_offset = 0
-        payload_cache: dict[tuple, tuple[int, tuple[int, ...], int, int]] = {}
-        stats = {"tensors": 0, "q2_k": 0, "q4_0": 0, "bf16": 0, "f32": 0}
-
-        if is_dsv4:
-            last_l = max(int(arch.get("num_hidden_layers", 40)) - 1, 0)
-            fn_k = f"layers.{last_l}.hc_ffn_fn"
-            bs_k = f"layers.{last_l}.hc_ffn_base"
-            sc_k = f"layers.{last_l}.hc_ffn_scale"
-            if fn_k in keys_set and bs_k in keys_set and sc_k in keys_set:
-                fn_t = f.get_tensor(fn_k).float()[:hc_mult].contiguous()
-                bs_t = f.get_tensor(bs_k).float()[:hc_mult].contiguous()
-                sc_t = f.get_tensor(sc_k).float()[:1].contiguous()
-                for g_hc_name, hc_t in (
-                    ("output_hc_fn.weight", fn_t),
-                    ("output_hc_base.weight", bs_t),
-                    ("output_hc_scale.weight", sc_t),
-                ):
-                    pad = _gguf_align(data_offset) - data_offset
-                    if pad > 0:
-                        tmp_data.write(b"\x00" * pad)
-                        data_offset += pad
-                    raw_hc = hc_t.cpu().numpy().astype("<f4").tobytes()
-                    gguf_tensors.append((g_hc_name, tuple(reversed(hc_t.shape)), _GGML_F32, data_offset))
-                    tmp_data.write(raw_hc)
-                    data_offset += len(raw_hc)
-                    stats["tensors"] += 1
-                    stats["f32"] += 1
-                tmp_data.flush()
-
-        total_items = len(work_items)
+        # Update effective expert count in arch metadata if stacked experts were pruned
+        if expert_groups:
+            first_group = next(iter(expert_groups.values()))
+            arch["num_experts"] = len(first_group)
 
         def _dequant_single_to_f32_torch(w_key: str, s_key: str | None) -> "torch.Tensor":
             wt = f.get_tensor(w_key).to(dev)
@@ -3113,7 +3083,6 @@ def export_safetensors_to_gguf(
                 if not is_emb and row_w % QK_K == 0 and n_el % QK_K == 0:
                     return "Q2_K"
                 return "BF16"
-            # mixed
             if is_emb:
                 return "Q4_0" if (row_w % QK4_0 == 0 and n_el % QK4_0 == 0) else "BF16"
             if row_w % QK_K == 0 and n_el % QK_K == 0:
@@ -3122,6 +3091,18 @@ def export_safetensors_to_gguf(
                 return "Q4_0"
             return "BF16"
 
+        def _target_type_and_bytes(tgt: str, torch_shape: tuple[int, ...]) -> tuple[int, int]:
+            n_el = int(np.prod(torch_shape)) if torch_shape else 1
+            if tgt == "F32":
+                return _GGML_F32, n_el * 4
+            if tgt == "BF16":
+                return _GGML_BF16, n_el * 2
+            if tgt == "Q4_0":
+                return _GGML_Q4_0, (n_el // QK4_0) * Q4_0_BLOCK_BYTES
+            if tgt == "Q2_K":
+                return _GGML_Q2_K, (n_el // QK_K) * 84
+            raise ValueError(tgt)
+
         def _encode_tensor_f32(t_f32: "torch.Tensor", tgt: str) -> tuple[int, bytes]:
             row_w = int(t_f32.shape[-1])
             if tgt == "F32":
@@ -3145,157 +3126,174 @@ def export_safetensors_to_gguf(
                 return _GGML_Q2_K, blocks_to_bytes(res.scales, res.qs, res.dm)
             raise ValueError(tgt)
 
-        for idx_w, item in enumerate(work_items, 1):
-            if stop_event is not None and stop_event.is_set():
-                tmp_data.close()
-                os.unlink(tmp_data.name)
-                raise InterruptedError("GGUF export cancelled by user.")
+        # Precompute exact GGUF tensor directory and offsets in O(1) so we write in a single pass
+        gguf_tensors: list[tuple[str, tuple[int, ...], int, int]] = []
+        planned_exec: list[tuple] = []
+        data_offset = 0
 
+        if is_dsv4:
+            last_l = max(int(arch.get("num_hidden_layers", 40)) - 1, 0)
+            fn_k = f"layers.{last_l}.hc_ffn_fn"
+            bs_k = f"layers.{last_l}.hc_ffn_base"
+            sc_k = f"layers.{last_l}.hc_ffn_scale"
+            if fn_k in keys_set and bs_k in keys_set and sc_k in keys_set:
+                fn_t = f.get_tensor(fn_k).float()[:hc_mult].contiguous()
+                bs_t = f.get_tensor(bs_k).float()[:hc_mult].contiguous()
+                sc_t = f.get_tensor(sc_k).float()[:1].contiguous()
+                for g_hc_name, hc_t in (
+                    ("output_hc_fn.weight", fn_t),
+                    ("output_hc_base.weight", bs_t),
+                    ("output_hc_scale.weight", sc_t),
+                ):
+                    data_offset = _gguf_align(data_offset)
+                    raw_hc = hc_t.cpu().numpy().astype("<f4").tobytes()
+                    g_dims = tuple(reversed(hc_t.shape))
+                    gguf_tensors.append((g_hc_name, g_dims, _GGML_F32, data_offset))
+                    planned_exec.append(("raw_bytes", g_hc_name, _GGML_F32, data_offset, raw_hc))
+                    data_offset += len(raw_hc)
+
+        for item in work_items:
             kind = item[0]
+            data_offset = _gguf_align(data_offset)
             if kind == "single":
                 _, w_key, s_key, gname = item
-                sl = f.get_slice(w_key)
-                t_shape = tuple(sl.get_shape())
+                t_shape = tuple(f.get_slice(w_key).get_shape())
                 if is_dsv4 and gname.endswith(".attn_output_a.weight") and len(t_shape) == 2 and t_shape[0] == o_groups * o_lora_rank:
                     t_shape = (o_groups, o_lora_rank, t_shape[1])
                 g_dims = tuple(reversed(t_shape)) if t_shape else (1,)
                 tgt = _choose_gguf_target(gname, t_shape or (1,))
-
-                # Fast signature for identical anchor layer reuse
-                lm = LAYER_RE.match(w_key)
-                cache_key = None
-                if lm:
-                    suf = lm.group(3)
-                    w_Sample = _payload_bytes(f.get_tensor(w_key).flatten()[:256])
-                    cache_key = ("single", suf, t_shape, tgt, w_Sample)
-                if cache_key is not None and cache_key in payload_cache:
-                    out_type, g_dims, src_off, src_len = payload_cache[cache_key]
-                    raw_bytes = os.pread(tmp_data.fileno(), src_len, src_off)
-                elif len(t_shape) == 2 and t_shape[0] > 4096 and s_key is None and tgt in ("BF16", "Q4_0"):
-                    # Stream large embedding/head matrices in 4096-row chunks to keep memory < 100 MB
-                    pad = _gguf_align(data_offset) - data_offset
-                    if pad > 0:
-                        tmp_data.write(b"\x00" * pad)
-                        data_offset += pad
-                    cur_off = data_offset
-                    written_len = 0
-                    out_type = _GGML_BF16 if tgt == "BF16" else _GGML_Q4_0
-                    for r0 in range(0, t_shape[0], 4096):
-                        r1 = min(r0 + 4096, t_shape[0])
-                        sub_t = sl[r0:r1].to(dev)
-                        if tgt == "BF16":
-                            b_part = sub_t.to(torch.bfloat16).view(torch.uint16).cpu().numpy().astype("<u2").tobytes()
-                        else:
-                            sub_np = sub_t.float().cpu().numpy().astype(np.float32)
-                            b_part, _, _ = quantize_tensor_q4_0(sub_np, cfg=hpc_cfg)
-                        tmp_data.write(b_part)
-                        written_len += len(b_part)
-                        del sub_t
-                    tmp_data.flush()
-                    data_offset += written_len
-                    gguf_tensors.append((gname, g_dims, out_type, cur_off))
-                    stats["tensors"] += 1
-                    stats["bf16" if out_type == _GGML_BF16 else "q4_0"] += 1
-                    if progress_cb is not None:
-                        progress_cb(idx_w, total_items, gname, tgt, stats)
-                    continue
-                else:
-                    t_f32 = _dequant_single_to_f32_torch(w_key, s_key)
-                    if not t_shape:
-                        t_f32 = t_f32.reshape(1)
-                    out_type, raw_bytes = _encode_tensor_f32(t_f32, tgt)
-                    del t_f32
+                out_type, out_sz = _target_type_and_bytes(tgt, t_shape or (1,))
+                gguf_tensors.append((gname, g_dims, out_type, data_offset))
+                planned_exec.append(("single", w_key, s_key, gname, t_shape, g_dims, tgt, out_type, out_sz, data_offset))
+                data_offset += out_sz
             else:
                 _, (lidx, proj), members, gname = item
                 first_w = members[0][1]
                 e_shape = tuple(f.get_slice(first_w).get_shape())
-                # If packed FP4 I8, logical in_dim is 2 * stored in_dim
                 dt_e = _safetensors_dtype(f.get_slice(first_w).get_dtype())
                 if dt_e == "I8" and members[0][2] is not None:
                     e_shape = (e_shape[0], e_shape[1] * 2)
                 full_shape = (len(members), e_shape[0], e_shape[1])
                 g_dims = (e_shape[1], e_shape[0], len(members))
                 tgt = _choose_gguf_target(gname, full_shape)
+                out_type, out_sz = _target_type_and_bytes(tgt, full_shape)
+                gguf_tensors.append((gname, g_dims, out_type, data_offset))
+                planned_exec.append(("experts", (lidx, proj), members, gname, full_shape, g_dims, tgt, out_type, out_sz, data_offset))
+                data_offset += out_sz
+
+        kv_pairs = _build_gguf_kv_pairs(
+            arch,
+            tok_path if tok_path.exists() else None,
+            chat_path if chat_path.exists() else None,
+            quant_format,
+            st_path.stem,
+        )
+
+        payload_cache: dict[tuple, bytes] = {}
+        stats = {"tensors": 0, "q2_k": 0, "q4_0": 0, "bf16": 0, "f32": 0}
+        total_items = len(planned_exec)
+
+        with out_path.open("wb") as fout:
+            fout.write(struct.pack("<IIQQ", _GGUF_MAGIC, 3, len(gguf_tensors), len(kv_pairs)))
+            for k, vt, raw_v in kv_pairs:
+                fout.write(_pack_gguf_str(k))
+                fout.write(struct.pack("<I", vt))
+                fout.write(raw_v)
+            for gname, g_dims, out_type, off in gguf_tensors:
+                fout.write(_pack_gguf_str(gname))
+                fout.write(struct.pack("<I", len(g_dims)))
+                for d in g_dims:
+                    fout.write(struct.pack("<Q", int(d)))
+                fout.write(struct.pack("<I", int(out_type)))
+                fout.write(struct.pack("<Q", int(off)))
+            data_start = _gguf_align(fout.tell())
+            if data_start > fout.tell():
+                fout.write(b"\x00" * (data_start - fout.tell()))
+
+            for idx_w, pitem in enumerate(planned_exec, 1):
+                if stop_event is not None and stop_event.is_set():
+                    raise InterruptedError("GGUF export cancelled by user.")
+
+                pkind = pitem[0]
+                if pkind == "raw_bytes":
+                    _, gname, out_type, cur_off, raw_bytes = pitem
+                    tgt = "F32"
+                    target_pos = data_start + cur_off
+                    if target_pos > fout.tell():
+                        fout.write(b"\x00" * (target_pos - fout.tell()))
+                    fout.write(raw_bytes)
+                elif pkind == "single":
+                    _, w_key, s_key, gname, t_shape, g_dims, tgt, out_type, out_sz, cur_off = pitem
+                    target_pos = data_start + cur_off
+                    if target_pos > fout.tell():
+                        fout.write(b"\x00" * (target_pos - fout.tell()))
+
+                    lm = LAYER_RE.match(w_key)
+                    cache_key = None
+                    if lm:
+                        suf = lm.group(3)
+                        w_sample = _payload_bytes(f.get_tensor(w_key).flatten()[:256])
+                        cache_key = ("single", suf, t_shape, tgt, w_sample)
+                    if cache_key is not None and cache_key in payload_cache:
+                        fout.write(payload_cache[cache_key])
+                    elif len(t_shape) == 2 and t_shape[0] > 4096 and s_key is None and tgt in ("BF16", "Q4_0"):
+                        sl = f.get_slice(w_key)
+                        for r0 in range(0, t_shape[0], 4096):
+                            r1 = min(r0 + 4096, t_shape[0])
+                            sub_t = sl[r0:r1].to(dev)
+                            if tgt == "BF16":
+                                b_part = sub_t.to(torch.bfloat16).view(torch.uint16).cpu().numpy().astype("<u2").tobytes()
+                            else:
+                                sub_np = sub_t.float().cpu().numpy().astype(np.float32)
+                                b_part, _, _ = quantize_tensor_q4_0(sub_np, cfg=hpc_cfg)
+                            fout.write(b_part)
+                            del sub_t
+                    else:
+                        t_f32 = _dequant_single_to_f32_torch(w_key, s_key)
+                        if not t_shape:
+                            t_f32 = t_f32.reshape(1)
+                        elif tuple(t_f32.shape) != t_shape:
+                            t_f32 = t_f32.reshape(t_shape)
+                        _, raw_bytes = _encode_tensor_f32(t_f32, tgt)
+                        del t_f32
+                        fout.write(raw_bytes)
+                        if cache_key is not None:
+                            payload_cache[cache_key] = raw_bytes
+                else:
+                    _, (lidx, proj), members, gname, full_shape, g_dims, tgt, out_type, out_sz, cur_off = pitem
+                    target_pos = data_start + cur_off
+                    if target_pos > fout.tell():
+                        fout.write(b"\x00" * (target_pos - fout.tell()))
+
+                    first_w = members[0][1]
+                    w_sample = _payload_bytes(f.get_tensor(first_w).flatten()[:256])
+                    cache_key = ("experts", proj, full_shape, tgt, w_sample)
+                    if cache_key in payload_cache:
+                        fout.write(payload_cache[cache_key])
+                    else:
+                        chunks = []
+                        for _, wk, sk in members:
+                            ef32 = _dequant_single_to_f32_torch(wk, sk)
+                            _, eb = _encode_tensor_f32(ef32, tgt)
+                            chunks.append(eb)
+                            del ef32
+                        raw_bytes = b"".join(chunks)
+                        fout.write(raw_bytes)
+                        payload_cache[cache_key] = raw_bytes
 
-                w_sample = _payload_bytes(f.get_tensor(first_w).flatten()[:256])
-                cache_key = ("experts", proj, full_shape, tgt, w_sample)
-                if cache_key in payload_cache:
-                    out_type, g_dims, src_off, src_len = payload_cache[cache_key]
-                    raw_bytes = os.pread(tmp_data.fileno(), src_len, src_off)
+                stats["tensors"] += 1
+                if out_type == _GGML_Q2_K:
+                    stats["q2_k"] += 1
+                elif out_type == _GGML_Q4_0:
+                    stats["q4_0"] += 1
+                elif out_type == _GGML_BF16:
+                    stats["bf16"] += 1
                 else:
-                    chunks = []
-                    for _, wk, sk in members:
-                        ef32 = _dequant_single_to_f32_torch(wk, sk)
-                        _, eb = _encode_tensor_f32(ef32, tgt)
-                        chunks.append(eb)
-                        del ef32
-                    raw_bytes = b"".join(chunks)
-                    out_type = (
-                        _GGML_Q2_K
-                        if tgt == "Q2_K"
-                        else (_GGML_Q4_0 if tgt == "Q4_0" else (_GGML_BF16 if tgt == "BF16" else _GGML_F32))
-                    )
+                    stats["f32"] += 1
 
-            pad = _gguf_align(data_offset) - data_offset
-            if pad > 0:
-                tmp_data.write(b"\x00" * pad)
-                data_offset += pad
-            cur_off = data_offset
-            tmp_data.write(raw_bytes)
-            tmp_data.flush()
-            data_offset += len(raw_bytes)
-
-            if cache_key is not None and cache_key not in payload_cache:
-                payload_cache[cache_key] = (out_type, g_dims, cur_off, len(raw_bytes))
-
-            gguf_tensors.append((gname, g_dims, out_type, cur_off))
-            stats["tensors"] += 1
-            if out_type == _GGML_Q2_K:
-                stats["q2_k"] += 1
-            elif out_type == _GGML_Q4_0:
-                stats["q4_0"] += 1
-            elif out_type == _GGML_BF16:
-                stats["bf16"] += 1
-            else:
-                stats["f32"] += 1
-
-            if progress_cb is not None and (idx_w % 8 == 0 or idx_w == total_items):
-                progress_cb(idx_w, total_items, gname, tgt, stats)
-
-    # Update effective expert count in arch metadata if stacked experts were pruned
-    if expert_groups:
-        first_group = next(iter(expert_groups.values()))
-        arch["num_experts"] = len(first_group)
-
-    kv_pairs = _build_gguf_kv_pairs(
-        arch,
-        tok_path if tok_path.exists() else None,
-        chat_path if chat_path.exists() else None,
-        quant_format,
-        st_path.stem,
-    )
+                if progress_cb is not None and (idx_w % 8 == 0 or idx_w == total_items):
+                    progress_cb(idx_w, total_items, gname, tgt, stats)
 
-    with out_path.open("wb") as fout:
-        fout.write(struct.pack("<IIQQ", _GGUF_MAGIC, 3, len(gguf_tensors), len(kv_pairs)))
-        for k, vt, raw_v in kv_pairs:
-            fout.write(_pack_gguf_str(k))
-            fout.write(struct.pack("<I", vt))
-            fout.write(raw_v)
-        for gname, g_dims, out_type, off in gguf_tensors:
-            fout.write(_pack_gguf_str(gname))
-            fout.write(struct.pack("<I", len(g_dims)))
-            for d in g_dims:
-                fout.write(struct.pack("<Q", int(d)))
-            fout.write(struct.pack("<I", int(out_type)))
-            fout.write(struct.pack("<Q", int(off)))
-        data_start = _gguf_align(fout.tell())
-        if data_start > fout.tell():
-            fout.write(b"\x00" * (data_start - fout.tell()))
-        tmp_data.seek(0)
-        shutil.copyfileobj(tmp_data, fout, length=16 * 1024 * 1024)
-
-    tmp_data.close()
-    os.unlink(tmp_data.name)
     stats["output_bytes"] = out_path.stat().st_size
     if verbose:
         print(