| |
| """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 |
|
|
| |
| |
| HYBRID_EXPERT_BUDGET_GB = 310.0 - 38.0 - 43.5 - 5.4 |
|
|
| |
| HOT_FRACTION_TARGET = 0.08 |
| MIN_HOT, MAX_HOT = 8, 96 |
| MAX_COLD = 176 |
|
|
|
|
| def main(): |
| 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_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) |
|
|
| |
| order = np.dstack(np.unravel_index( |
| np.argsort(mass, axis=None)[::-1], mass.shape))[0] |
| hot = [set() for _ in range(n_layers)] |
| |
| 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 |
|
|
| |
| base_all = total_experts * COST_BASE |
| extra_hot = n_hot * (COST_HOT - COST_BASE) |
| budget_bytes = HYBRID_EXPERT_BUDGET_GB * 1e3 |
| need_savings = base_all + extra_hot - budget_bytes |
| 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)) |
| |
| 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() |
|
|
