| |
| """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() |
|
|
