File size: 3,092 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 | #!/usr/bin/env python3
"""Manifest of NVFP4 per-expert tensors for the 2-tier hot sets.
Fetches the HF repo's safetensors index + shard headers via range requests
and emits (url, offset, length, out) entries for every tensor of every hot
expert in the assignment, excluding layers whose NVFP4 bytes are already in
the local checkpoint's compacted arrays (LOCAL_LAYERS).
"""
import json
import os
import struct
import sys
import urllib.request
from concurrent.futures import ThreadPoolExecutor
BASE = "https://huggingface.co/lukealonso/GLM-5.2-NVFP4/resolve/main"
DL = "/tmp/glm52-hot-dl"
ASSIGN = "/data/glm52-expert-assignment.json"
LOCAL_LAYERS = {3, 4, 5, 8, 74, 75, 76, 77}
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+0}-{start+length-1}")
for attempt in range(6):
try:
with urllib.request.urlopen(req, timeout=90) as r:
return r.read()
except Exception:
if attempt == 5:
raise
return None
def main():
os.makedirs(f"{DL}/regions", exist_ok=True)
assign = {int(k): v for k, v in json.load(open(ASSIGN)).items()}
want = set()
for li, a in assign.items():
if li in LOCAL_LAYERS:
continue
for e in a["hot"]:
want.add((li, e))
print(f"{len(want)} hot experts to download", file=sys.stderr)
idx = json.loads(fetch(f"{BASE}/model.safetensors.index.json"))
wm = idx["weight_map"]
import re
by_shard: dict[str, list[str]] = {}
pat = re.compile(r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.")
for name, shard in wm.items():
m = pat.match(name)
if m and (int(m.group(1)), int(m.group(2))) in want:
by_shard.setdefault(shard, []).append(name)
def one(shard):
url = f"{BASE}/{shard}"
hlen = struct.unpack("<Q", fetch(url, 0, 8))[0]
hdr = json.loads(fetch(url, 8, hlen))
base = 8 + hlen
ranges = sorted(
hdr[n]["data_offsets"] for n in by_shard[shard] if n in hdr
)
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": base + b, "length": e - b,
"out": f"{DL}/regions/{shard}/{b}.bin", "rel_start": b}
for b, e in merged]
return shard, hlen, hdr, ents
entries, headers = [], {}
with ThreadPoolExecutor(16) as ex:
for shard, hlen, hdr, ents in ex.map(one, sorted(by_shard)):
headers[shard] = {"header_len": hlen, "header": hdr}
entries.extend(ents)
json.dump(headers, open(f"{DL}/headers.json", "w"))
json.dump(entries, open(f"{DL}/manifest.json", "w"))
total = sum(e["length"] for e in entries)
print(f"{len(entries)} regions, {total/1e9:.1f} GB", file=sys.stderr)
if __name__ == "__main__":
main()
|