| |
| """Turn a dense npz into an N-expert npz whose experts are all IDENTICAL. |
| |
| This is the control that isolates the mixture's cartridge overhead from its |
| arithmetic. A cartridge built from it streams a different bank depending on |
| the routed token, walks the extra region boundaries, and does one router |
| lookup per token - but every expert holds the same weights, so it must |
| generate exactly the tokens the dense cartridge generates. Any difference in |
| cycles is therefore structure, and any difference in output is a bug. |
| """ |
| import argparse |
| import sys |
|
|
| import numpy as np |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("src") |
| ap.add_argument("dst") |
| ap.add_argument("--nexp", type=int, default=4) |
| ap.add_argument("--route", default="bal", |
| help="routing kind from train/route.py, or 'spread' for " |
| "tok %% N (which maximises how often the routed bank " |
| "changes between consecutive tokens)") |
| a = ap.parse_args() |
|
|
| z = dict(np.load(a.src)) |
| if "_moe" in z: |
| raise SystemExit("%s is already a mixture" % a.src) |
| n = a.nexp |
| for k in list(z): |
| if k.endswith("_W1") or k.endswith("_W2"): |
| for e in range(n): |
| z["%s_e%d" % (k, e)] = z[k].copy() |
| del z[k] |
| if a.route == "spread": |
| rt = [t % n for t in range(64)] |
| else: |
| sys.path.insert(0, "train") |
| import route as R |
| fit = np.load("data/fit_bpe64.npy") |
| rt = R.build(a.route, n, fit, seed=1) |
| z["_route"] = np.array(rt, dtype=np.int16) |
| z["_moe"] = np.array([n, 1], dtype=np.int16) |
| np.savez(a.dst, **z) |
| print("wrote %s: %d identical experts, route %s" % (a.dst, n, a.route)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
