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