File size: 4,692 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 | #!/usr/bin/env python3
"""Per-expert precision assignment from routing stats, budget-constrained.
Tiers per expert within each hybrid layer (whole-NVFP4 layers and MTP are
untouched):
hot = NVFP4 21.23 MB/expert (4.5 bpw)
base = AQLM w13 1x16, w2 2x16 12.58 MB/expert (2 / 4 bpw)
cold = AQLM w13 1x16, w2 1x16 9.44 MB/expert (2 / 2 bpw)
Greedy global assignment: experts ranked by routing mass; the hottest get
NVFP4, the coldest fund them by dropping w2 to one book, subject to the
total expert-byte budget. Per-layer floor/cap keeps every layer sane.
Output: /data/glm52-expert-assignment.json
{layer: {"hot": [ids], "cold": [ids]}} (rest = base)
"""
import json
import os
import sys
import numpy as np
STATS = sys.argv[1] if len(sys.argv) > 1 else "/data/glm52-expert-stats.npz"
OUT = "/data/glm52-expert-assignment.json"
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MB = 1e6
COST_HOT, COST_BASE, COST_COLD = 21.234, 12.583, 9.437 # MB / expert
# weight budget: total <= 292 GB
# non-expert 38.0 GB, whole-NVFP4 layers (3,4,5,8,74-77) 43.5 GB, MTP 5.4 GB
HYBRID_EXPERT_BUDGET_GB = 310.0 - 38.0 - 43.5 - 5.4
# knobs
HOT_FRACTION_TARGET = 0.08 # ~20 experts/layer on average
MIN_HOT, MAX_HOT = 8, 96 # per-layer bounds
MAX_COLD = 176 # never cold more than this many per layer
def main():
d = np.load(STATS)
counts = d["counts"].astype(np.float64) # [L, 256]
layer_ids = d["layer_ids"].tolist()
# exclude whole-NVFP4 layers (incl. MTP) from per-expert assignment
plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))
whole_nvfp4 = set(plan["nvfp4_layers"])
keep_rows = [i for i, li in enumerate(layer_ids) if li not in whole_nvfp4]
counts = counts[keep_rows]
layer_ids = [layer_ids[i] for i in keep_rows]
n_layers, n_exp = counts.shape
mass = counts / counts.sum(axis=1, keepdims=True).clip(min=1)
total_experts = n_layers * n_exp
n_hot_target = int(total_experts * HOT_FRACTION_TARGET)
# global ranking for hot with per-layer floor/cap
order = np.dstack(np.unravel_index(
np.argsort(mass, axis=None)[::-1], mass.shape))[0] # (layer, expert)
hot = [set() for _ in range(n_layers)]
# floors first
for li in range(n_layers):
top = np.argsort(mass[li])[::-1][:MIN_HOT]
hot[li].update(top.tolist())
n_hot = sum(len(h) for h in hot)
for li, ei in order:
if n_hot >= n_hot_target:
break
if len(hot[li]) >= MAX_HOT or ei in hot[li]:
continue
hot[li].add(int(ei))
n_hot += 1
# bytes so far with everyone else at base
base_all = total_experts * COST_BASE
extra_hot = n_hot * (COST_HOT - COST_BASE)
budget_bytes = HYBRID_EXPERT_BUDGET_GB * 1e3 # MB
need_savings = base_all + extra_hot - budget_bytes # MB to claw back
per_cold_saving = COST_BASE - COST_COLD
cold = [set() for _ in range(n_layers)]
if need_savings > 0:
n_cold_needed = int(np.ceil(need_savings / per_cold_saving))
# coldest experts globally (ascending mass), respecting caps
order_asc = np.dstack(np.unravel_index(
np.argsort(mass, axis=None), mass.shape))[0]
n_cold = 0
for li, ei in order_asc:
if n_cold >= n_cold_needed:
break
ei = int(ei)
if ei in hot[li] or len(cold[li]) >= MAX_COLD:
continue
cold[li].add(ei)
n_cold += 1
total = (n_hot * COST_HOT
+ sum(len(c) for c in cold) * COST_COLD
+ (total_experts - n_hot - sum(len(c) for c in cold)) * COST_BASE)
out = {}
hot_mass, cold_mass = [], []
for li in range(n_layers):
gid = layer_ids[li]
out[str(gid)] = {
"hot": sorted(hot[li]),
"cold": sorted(cold[li]),
}
hot_mass.append(mass[li, sorted(hot[li])].sum())
if cold[li]:
cold_mass.append(mass[li, sorted(cold[li])].sum())
json.dump(out, open(OUT, "w"))
n_cold_t = sum(len(c) for c in cold)
print(f"hot experts: {n_hot} ({n_hot/total_experts:.1%}), "
f"mean mass covered {np.mean(hot_mass):.1%}")
print(f"cold experts: {n_cold_t} ({n_cold_t/total_experts:.1%}), "
f"mean mass affected {np.mean(cold_mass) if cold_mass else 0:.2%}")
print(f"hybrid expert bytes: {total/1e3:.1f} GB "
f"(budget {HYBRID_EXPERT_BUDGET_GB:.1f} GB)")
print(f"projected total weights: {total/1e3 + 38.0 + 43.5 + 5.4:.1f} GB (ceiling 310)")
print("saved", OUT)
if __name__ == "__main__":
main()
|