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