File size: 5,935 Bytes
fdc6474
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
#!/usr/bin/env python3
"""Parse GGUF headers from HuggingFace via HTTP range requests (no full download).

Produces a JSON tensor map per shard: name, shape, ggml type, absolute byte
range within the shard file. This lets us (a) recover unsloth's per-tensor
quant-type assignment (their importance map), and (b) later download exactly
the expert tensors we need for the hybrid checkpoint build.
"""
import json
import struct
import sys
import urllib.request

GGML_TYPES = {
    0: "F32", 1: "F16", 2: "Q4_0", 3: "Q4_1", 6: "Q5_0", 7: "Q5_1", 8: "Q8_0",
    9: "Q8_1", 10: "Q2_K", 11: "Q3_K", 12: "Q4_K", 13: "Q5_K", 14: "Q6_K",
    15: "Q8_K", 16: "IQ2_XXS", 17: "IQ2_XS", 18: "IQ3_XXS", 19: "IQ1_S",
    20: "IQ4_NL", 21: "IQ3_S", 22: "IQ2_S", 23: "IQ4_XS", 24: "I8", 25: "I16",
    26: "I32", 27: "I64", 28: "F64", 29: "IQ1_M", 30: "BF16",
    34: "TQ1_0", 35: "TQ2_0", 39: "MXFP4",
}
# bytes per block / elements per block
GGML_BLOCK = {
    "F32": (4, 1), "F16": (2, 1), "BF16": (2, 1),
    "Q8_0": (34, 32), "Q4_0": (18, 32), "Q4_1": (20, 32),
    "Q5_0": (22, 32), "Q5_1": (24, 32),
    "Q2_K": (84, 256), "Q3_K": (110, 256), "Q4_K": (144, 256),
    "Q5_K": (176, 256), "Q6_K": (210, 256),
    "IQ1_S": (50, 256), "IQ1_M": (56, 256),
    "IQ2_XXS": (66, 256), "IQ2_XS": (74, 256), "IQ2_S": (82, 256),
    "IQ3_XXS": (98, 256), "IQ3_S": (110, 256),
    "IQ4_NL": (18, 32), "IQ4_XS": (136, 256),
    "I8": (1, 1), "I16": (2, 1), "I32": (4, 1), "I64": (8, 1), "F64": (8, 1),
    "MXFP4": (17, 32),
}


def fetch_range(url: str, start: int, length: int) -> bytes:
    req = urllib.request.Request(url)
    req.add_header("Range", f"bytes={start}-{start + length - 1}")
    with urllib.request.urlopen(req) as r:
        return r.read()


class Cursor:
    def __init__(self, buf: bytes):
        self.buf = buf
        self.pos = 0

    def need(self, n):
        if self.pos + n > len(self.buf):
            raise EOFError(f"header larger than fetched window ({len(self.buf)} bytes)")

    def u32(self):
        self.need(4); v = struct.unpack_from("<I", self.buf, self.pos)[0]; self.pos += 4; return v

    def u64(self):
        self.need(8); v = struct.unpack_from("<Q", self.buf, self.pos)[0]; self.pos += 8; return v

    def i64(self):
        self.need(8); v = struct.unpack_from("<q", self.buf, self.pos)[0]; self.pos += 8; return v

    def f32(self):
        self.need(4); v = struct.unpack_from("<f", self.buf, self.pos)[0]; self.pos += 4; return v

    def f64(self):
        self.need(8); v = struct.unpack_from("<d", self.buf, self.pos)[0]; self.pos += 8; return v

    def s(self):
        n = self.u64(); self.need(n)
        v = self.buf[self.pos:self.pos + n].decode("utf-8", errors="replace"); self.pos += n
        return v

    def value(self, t):
        if t == 0: self.need(1); v = self.buf[self.pos]; self.pos += 1; return v
        if t == 1: self.need(1); v = struct.unpack_from("<b", self.buf, self.pos)[0]; self.pos += 1; return v
        if t == 2: self.need(2); v = struct.unpack_from("<H", self.buf, self.pos)[0]; self.pos += 2; return v
        if t == 3: self.need(2); v = struct.unpack_from("<h", self.buf, self.pos)[0]; self.pos += 2; return v
        if t == 4: return self.u32()
        if t == 5: self.need(4); v = struct.unpack_from("<i", self.buf, self.pos)[0]; self.pos += 4; return v
        if t == 6: return self.f32()
        if t == 7: self.need(1); v = self.buf[self.pos] != 0; self.pos += 1; return v
        if t == 8: return self.s()
        if t == 9:  # array
            et = self.u32(); n = self.u64()
            return [self.value(et) for _ in range(n)]
        if t == 10: return self.u64()
        if t == 11: return self.i64()
        if t == 12: return self.f64()
        raise ValueError(f"bad gguf kv type {t}")


def parse_header(url: str, window: int = 4 << 20, want_kv: bool = False):
    buf = fetch_range(url, 0, window)
    while True:
        try:
            return _parse(buf, want_kv)
        except EOFError:
            window *= 4
            if window > (1 << 30):
                raise
            buf = fetch_range(url, 0, window)


def _parse(buf: bytes, want_kv: bool):
    c = Cursor(buf)
    magic = c.u32()
    assert magic == 0x46554747, f"not gguf: {magic:#x}"
    version = c.u32()
    assert version in (2, 3), version
    n_tensors = c.u64()
    n_kv = c.u64()
    kvs = {}
    alignment = 32
    for _ in range(n_kv):
        k = c.s()
        t = c.u32()
        v = c.value(t)
        if k == "general.alignment":
            alignment = v
        if want_kv and not (isinstance(v, list) and len(v) > 64):
            kvs[k] = v
    tensors = []
    for _ in range(n_tensors):
        name = c.s()
        nd = c.u32()
        dims = [c.u64() for _ in range(nd)]
        ty = GGML_TYPES.get(c.u32(), "?")
        off = c.u64()
        tensors.append({"name": name, "dims": dims, "type": ty, "off": off})
    data_start = (c.pos + alignment - 1) // alignment * alignment
    for t in tensors:
        bpb, epb = GGML_BLOCK[t["type"]]
        nelem = 1
        for d in t["dims"]:
            nelem *= d
        nbytes = nelem // epb * bpb
        t["abs_off"] = data_start + t["off"]
        t["nbytes"] = nbytes
    return {"n_tensors": n_tensors, "data_start": data_start, "tensors": tensors, "kv": kvs}


if __name__ == "__main__":
    base = sys.argv[1]      # e.g. https://huggingface.co/unsloth/GLM-5.2-GGUF/resolve/main/UD-Q2_K_XL/GLM-5.2-UD-Q2_K_XL
    nshards = int(sys.argv[2])
    out = sys.argv[3]
    all_shards = {}
    for i in range(1, nshards + 1):
        url = f"{base}-{i:05d}-of-{nshards:05d}.gguf"
        h = parse_header(url, want_kv=(i == 1))
        all_shards[i] = {"url": url, "data_start": h["data_start"], "tensors": h["tensors"], "kv": h.get("kv", {})}
        print(f"shard {i}: {h['n_tensors']} tensors", file=sys.stderr)
    json.dump(all_shards, open(out, "w"))
    print(f"wrote {out}", file=sys.stderr)