File size: 2,531 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 | #!/usr/bin/env python3
"""Two-tier per-expert assignment: hot = NVFP4 (4.5bpw), cold = full 2-bpw
AQLM (w13 1x16, w2 1x16). No base tier. Hot count is set byte-exactly from
the expert budget; hot experts chosen by global routing mass with per-layer
floor/cap.
Usage: solve_assignment_2tier.py <stats.npz> [expert_budget_gb]
Output: /data/glm52-expert-assignment.json {layer: {hot: [...], cold: [...]}}
"""
import json
import os
import sys
import numpy as np
STATS = sys.argv[1]
EXPERT_BUDGET_GB = float(sys.argv[2]) if len(sys.argv) > 2 else 250.0
OUT = "/data/glm52-expert-assignment.json"
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
COST_HOT, COST_COLD = 21.234, 9.437 # MB/expert
MIN_HOT, MAX_HOT = 16, 176
d = np.load(STATS)
counts = d["counts"].astype(np.float64)
layer_ids = d["layer_ids"].tolist()
plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))
whole = set(plan["nvfp4_layers"]) - {78} # none expected now
keep = [i for i, li in enumerate(layer_ids) if li not in whole]
counts = counts[keep]
layer_ids = [layer_ids[i] for i in keep]
n_layers, n_exp = counts.shape
mass = counts / counts.sum(axis=1, keepdims=True).clip(min=1)
total = n_layers * n_exp
budget_mb = EXPERT_BUDGET_GB * 1e3
n_hot = int((budget_mb - total * COST_COLD) / (COST_HOT - COST_COLD))
print(f"budget {EXPERT_BUDGET_GB} GB -> {n_hot} hot experts "
f"({n_hot/total:.1%})")
hot = [set() for _ in range(n_layers)]
for li in range(n_layers):
hot[li].update(np.argsort(mass[li])[::-1][:MIN_HOT].tolist())
placed = sum(len(h) for h in hot)
order = np.dstack(np.unravel_index(
np.argsort(mass, axis=None)[::-1], mass.shape))[0]
for li, ei in order:
if placed >= n_hot:
break
if len(hot[li]) >= MAX_HOT or int(ei) in hot[li]:
continue
hot[li].add(int(ei))
placed += 1
out = {}
hot_mass = []
for li in range(n_layers):
h = sorted(int(x) for x in hot[li])
out[str(layer_ids[li])] = {
"hot": h,
"cold": [e for e in range(n_exp) if e not in hot[li]],
}
hot_mass.append(mass[li, h].sum())
json.dump(out, open(OUT, "w"))
tot_mb = placed * COST_HOT + (total - placed) * COST_COLD
print(f"hot: {placed} ({placed/total:.1%}) covering mean "
f"{np.mean(hot_mass):.1%} of routing mass "
f"(min {np.min(hot_mass):.1%} max {np.max(hot_mass):.1%})")
print(f"expert bytes: {tot_mb/1e3:.1f} GB | hot/layer: "
f"min {min(len(h) for h in hot)} max {max(len(h) for h in hot)} "
f"mean {placed/n_layers:.0f}")
print("saved", OUT)
|