File size: 4,326 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 | #!/usr/bin/env python3
"""Generate download manifests (url, byte range, output path) for the hybrid build.
- NVFP4 repo: parse each shard's safetensors header remotely; select all
non-routed-expert tensors plus routed experts of plan.nvfp4_layers;
coalesce adjacent ranges.
- GGUF repos: use the tensor maps produced by gguf_remote.py to select
blk.N.ffn_{gate,up,down}_exps.weight for the planned layers.
"""
import json
import re
import struct
import sys
import urllib.request
from concurrent.futures import ThreadPoolExecutor
ROOT = "/home/coder/git/glm52"
DL = "/tmp/glm52-dl"
plan = json.load(open(f"{ROOT}/hybrid_plan.json"))
NVFP4_LAYERS = set(plan["nvfp4_layers"])
GGUF_LAYERS = {**{l: "ud_q3" for l in plan["ud_q3_layers"]},
**{l: "ud_q2" for l in plan["ud_q2_layers"]}}
COALESCE_GAP = 4 << 20
def fetch(url, start=None, length=None):
req = urllib.request.Request(url)
if start is not None:
req.add_header("Range", f"bytes={start}-{start+length-1}")
for attempt in range(5):
try:
with urllib.request.urlopen(req, timeout=60) as r:
return r.read()
except Exception as e:
if attempt == 4:
raise
return None
def needed_nvfp4(name):
m = re.match(r"model\.layers\.(\d+)\.mlp\.experts\.", name)
if m:
return int(m.group(1)) in NVFP4_LAYERS
return True
def st_header(url):
n = struct.unpack("<Q", fetch(url, 0, 8))[0]
hdr = json.loads(fetch(url, 8, n))
return n, hdr
def nvfp4_manifest():
idx = json.load(open("/tmp/glm52-recon/nvfp4_model.safetensors.index.json"))
shards = sorted(set(idx["weight_map"].values()))
entries, headers = [], {}
def one(shard):
url = f'{plan["sources"]["nvfp4"]}/{shard}'
hlen, hdr = st_header(url)
data_base = 8 + hlen
ranges = []
for name, info in hdr.items():
if name == "__metadata__":
continue
if needed_nvfp4(name):
b, e = info["data_offsets"]
ranges.append((b, e))
ranges.sort()
merged = []
for b, e in ranges:
if merged and b - merged[-1][1] <= COALESCE_GAP:
merged[-1][1] = max(merged[-1][1], e)
else:
merged.append([b, e])
ents = [{"url": url,
"start": data_base + b,
"length": e - b,
"out": f"{DL}/nvfp4/regions/{shard}/{b}.bin",
"rel_start": b} for b, e in merged]
return shard, hlen, hdr, ents
with ThreadPoolExecutor(16) as ex:
for shard, hlen, hdr, ents in ex.map(one, shards):
headers[shard] = {"header_len": hlen, "header": hdr}
entries.extend(ents)
print(f"{shard}: {len(ents)} regions, "
f"{sum(e['length'] for e in ents)/1e9:.2f} GB", file=sys.stderr)
json.dump(headers, open(f"{DL}/nvfp4/headers.json", "w"))
return entries
def gguf_manifest(tag):
m = json.load(open(f"/tmp/glm52-recon/{'q3kxl' if tag=='ud_q3' else 'q2kxl'}_map.json"))
layers = {l for l, t in GGUF_LAYERS.items() if t == tag}
entries = []
tinfo = {}
for sh in m.values():
for t in sh["tensors"]:
mm = re.match(r"blk\.(\d+)\.(ffn_(?:gate|up|down)_exps)\.weight", t["name"])
if mm and int(mm.group(1)) in layers:
out = f"{DL}/gguf/{tag}/{t['name']}.{t['type']}.bin"
entries.append({"url": sh["url"], "start": t["abs_off"],
"length": t["nbytes"], "out": out})
tinfo[t["name"]] = {"type": t["type"], "dims": t["dims"],
"nbytes": t["nbytes"], "file": out}
json.dump(tinfo, open(f"{DL}/gguf/{tag}_tensors.json", "w"))
return entries
def main():
import os
for d in ["nvfp4/regions", "gguf/ud_q3", "gguf/ud_q2"]:
os.makedirs(f"{DL}/{d}", exist_ok=True)
ents = []
ents += gguf_manifest("ud_q3")
ents += gguf_manifest("ud_q2")
ents += nvfp4_manifest()
total = sum(e["length"] for e in ents)
json.dump(ents, open(f"{DL}/manifest.json", "w"))
print(f"manifest: {len(ents)} entries, {total/1e9:.1f} GB total", file=sys.stderr)
if __name__ == "__main__":
main()
|