展开查看完整源码minimal_moe.py · 225 行
"""Train a 32-parameter top-1 MoE diagnostic toy, not a language model.
CPU/NumPy only. Saves data, weights, traces, and an inspectable HTML heatmap.
Discrete choices/capacity masks are locally constant for differentiation.
"""
import argparse
import json
import math
import platform
from pathlib import Path
import numpy as np
def softmax(z):
m = z.max(axis=1, keepdims=True)
e = np.exp(z - m)
return e / e.sum(axis=1, keepdims=True), (m[:, 0] + np.log(e.sum(axis=1)))
def dispatch(p, factor=1.25):
n, experts = p.shape
capacity = math.ceil(factor * n / experts)
selected = p.argmax(axis=1) # ties go to lowest index, intentionally explicit
onehot = np.eye(experts)[selected]
positions = (onehot.cumsum(axis=0) * onehot).sum(axis=1)
accepted = positions <= capacity
pre = onehot.sum(axis=0).astype(int)
post = (onehot * accepted[:, None]).sum(axis=0).astype(int)
return selected, accepted, pre, post, capacity
def forward(w, v, x, y, alpha=0.0, beta=0.0, factor=1.25):
n, experts = len(x), w.shape[1]
p, logz = softmax(x @ w)
selected, accepted, pre, post, capacity = dispatch(p, factor)
f = pre / n
probability = p.mean(axis=0)
gate = p[np.arange(n), selected]
expert_out = np.zeros(n)
for e in range(experts):
idx = (selected == e) & accepted
expert_out[idx] = x[idx] @ v[e]
pred = 0.25 * x[:, 0] + gate * expert_out
error = pred - y
task = np.mean(error ** 2)
bal = experts * np.dot(f, probability) # BEFORE capacity, f stop-gradient
zloss = np.mean(logz ** 2)
gp = np.broadcast_to(alpha * experts * f / n, p.shape).copy()
gp[np.arange(n), selected] += 2 * error * expert_out / n
gz = p * (gp - (gp * p).sum(axis=1, keepdims=True))
gz += beta * 2 * logz[:, None] * p / n
gw = x.T @ gz
gv = np.zeros_like(v)
for e in range(experts):
idx = (selected == e) & accepted
gv[e] = x[idx].T @ (2 * error[idx] * gate[idx] / n)
metrics = dict(mse=float(task), balance=float(bal), z_loss=float(zloss),
pre=pre.tolist(), post=post.tolist(), capacity=capacity,
drop=float(1 - accepted.mean()), max_share=float(f.max()),
load_cv=float(f.std() / f.mean()), dead=int((pre == 0).sum()),
entropy=float(-(p * np.log(np.maximum(p, 1e-300))).sum(axis=1).mean()))
trace = dict(probabilities=p.tolist(), selected=selected.tolist(),
accepted=accepted.tolist(), predictions=pred.tolist())
return task + alpha * bal + beta * zloss, (gw, gv), metrics, trace
def make_data():
rng = np.random.default_rng(7000)
result = {}
for split, n in [('train', 256), ('dev', 128), ('test', 256)]:
x = np.column_stack([rng.normal(size=(n, 3)), np.ones(n)])
group = (x[:, 0] > 0).astype(int) + 2 * (x[:, 1] > 0).astype(int)
y = 0.25*x[:, 0] + np.array([1., -1., 2., -2.])[group]*x[:, 2]
y += 0.4*np.array([1., -1., -1., 1.])[group]
result[split] = dict(x=x.tolist(), y=y.tolist(), group=group.tolist())
return result
def evaluate(w, v, data, factor=1.25):
x, y = np.array(data['x']), np.array(data['y'])
metrics, traces = [], []
for start in range(0, len(x), 64):
_, _, m, t = forward(w, v, x[start:start+64], y[start:start+64], factor=factor)
metrics.append(m)
traces.append(t)
avg = {k: float(np.mean([m[k] for m in metrics])) for k in
['mse', 'balance', 'z_loss', 'drop', 'max_share', 'load_cv', 'dead', 'entropy']}
avg['pre'] = np.sum([m['pre'] for m in metrics], axis=0).tolist()
avg['post'] = np.sum([m['post'] for m in metrics], axis=0).tolist()
avg['batches'] = metrics
return avg, traces
def checks():
p, _ = softmax(np.zeros((64, 4)))
s, a, pre, post, cap = dispatch(p)
bal = 4 * np.dot(pre/64, p.mean(axis=0))
assert bal == 1 and pre.tolist() == [64, 0, 0, 0]
assert cap == 20 and post.tolist() == [20, 0, 0, 0] and (~a).sum() == 44
# High z-loss can coexist with IDENTICAL probabilities/routes after a shift.
p2, logz2 = softmax(np.full((64, 4), 100.0))
assert np.array_equal(p, p2)
# Counterexample to measuring balance only after capacity clipping.
indices = np.repeat(np.arange(4), [28, 12, 12, 12])
q = np.eye(4)[indices] * 0.96 + 0.01
_, _, before, after, _ = dispatch(q, 0.5)
assert before.tolist() == [28, 12, 12, 12] and after.tolist() == [8, 8, 8, 8]
# Dropped-to-zero argmax would mislabel overflow as expert zero.
original = np.full(64, 2)
mask = np.eye(4)[original]
mask[20:] = 0
assert np.bincount(mask.argmax(1), minlength=4).tolist() == [44, 0, 20, 0]
rng = np.random.default_rng(170)
x = np.column_stack([rng.normal(size=(16, 3)), np.ones(16)])
y = rng.normal(size=16)
w, v = rng.normal(size=(4, 4)), rng.normal(size=(4, 4))
args = (x, y, 0.1, 0.001)
_, grads, _, base_trace = forward(w, v, *args)
maximum = 0.0
for param, grad in zip([w, v], grads):
for index in np.ndindex(param.shape):
old, eps = param[index], 1e-6
param[index] = old + eps
plus, _, _, tp = forward(w, v, *args)
param[index] = old - eps
minus, _, _, tm = forward(w, v, *args)
param[index] = old
assert tp['selected'] == tm['selected'] == base_trace['selected']
assert tp['accepted'] == tm['accepted'] == base_trace['accepted']
maximum = max(maximum, abs((plus-minus)/(2*eps) - grad[index]))
assert maximum < 1e-7
# Zero expert gradient on an entirely unused expert; residual on dropped rows.
w = np.zeros((4, 4)); w[3, 0] = 8
_, (_, gv), _, tr = forward(w, v, x, y)
assert np.all(gv[1:] == 0)
dropped = ~np.array(tr['accepted'])
assert np.allclose(np.array(tr['predictions'])[dropped], 0.25*x[dropped, 0])
return dict(gradient_max_abs_error=maximum, uniform_balance=bal,
uniform_pre=pre.tolist(), uniform_post=post.tolist(), uniform_drop=44/64,
shifted_z_loss=float(np.mean(logz2**2)), post_clip_counterexample=[before.tolist(), after.tolist()],
checks='finite difference, tie, capacity, overflow index, shift, dead gradient, residual: passed')
def train(data, seed, name, alpha, beta, out):
rng = np.random.default_rng(seed)
w = rng.normal(0, 0.03, (4, 4)); w[3, 0] += 2
v = rng.normal(0, 0.1, (4, 4))
moments = [np.zeros_like(w), np.zeros_like(v)]
variances = [np.zeros_like(w), np.zeros_like(v)]
x, y = np.array(data['train']['x']), np.array(data['train']['y'])
history = []
for step in range(601):
if step % 100 == 0:
dev, _ = evaluate(w, v, data['dev'])
train_metrics, _ = evaluate(w, v, data['train'])
history.append(dict(step=step, dev=dev, train=train_metrics))
if step == 600:
break
idx = rng.choice(len(x), 64, replace=False)
_, grads, _, _ = forward(w, v, x[idx], y[idx], alpha, beta)
for i, (param, grad) in enumerate(zip([w, v], grads)):
moments[i] = 0.9*moments[i] + 0.1*grad
variances[i] = 0.999*variances[i] + 0.001*grad**2
mhat = moments[i] / (1 - 0.9**(step+1))
vhat = variances[i] / (1 - 0.999**(step+1))
param -= 0.03 * mhat / (np.sqrt(vhat) + 1e-8)
test, traces = evaluate(w, v, data['test'])
capacity_sweep = {str(c): evaluate(w, v, data['test'], c)[0] for c in [0.5, 1.0, 1.25, 2.0, 4.0]}
key = f'{name}_{seed}'
result = dict(name=name, seed=seed, alpha=alpha, beta=beta, steps=600,
history=history, test=test, capacity_sweep=capacity_sweep, traces=traces)
(out/f'{key}_weights.json').write_text(json.dumps(dict(w=w.tolist(), v=v.tolist()), indent=2))
(out/f'{key}.json').write_text(json.dumps(result, indent=2))
return result
def html_report(results, out):
# A numeric HTML table, not an image and not a substitute for the imagegen figures.
sections = []
for r in results:
rows = []
t = r['traces'][0]
for j, probs in enumerate(t['probabilities']):
cells = ''.join(f'<td style="background:rgba(0,135,170,{p:.4f})">{p:.4f}</td>' for p in probs)
rows.append(f'<tr><th>{j}</th>{cells}<td>{t["selected"][j]}</td><td>{t["accepted"][j]}</td></tr>')
sections.append(f'<details><summary>{r["name"]}, seed {r["seed"]}: MSE {r["test"]["mse"]:.4f}, drop {r["test"]["drop"]:.3f}</summary>'
'<table><thead><tr><th>Token</th><th>E0</th><th>E1</th><th>E2</th><th>E3</th><th>Choice</th><th>Accepted</th></tr></thead><tbody>'
+ ''.join(rows) + '</tbody></table></details>')
html = '<!doctype html><html lang="en"><meta charset="utf-8"><title>MoE actual routing traces</title>'
html += '<style>body{font:16px system-ui;max-width:960px;margin:40px auto;padding:16px;color:#122c3d}table{border-collapse:collapse;margin:20px 0}td,th{padding:5px 14px;border:1px solid #ccd}summary{padding:16px;cursor:pointer}details{border-bottom:1px solid #ccd}</style>'
html += '<h1>MoE actual routing traces</h1><p>First fixed test batch: 64 synthetic tokens, four experts. All runs use capacity 20 per expert. Cell opacity equals probability (0 to 1); numeric values remain visible. Choice is BEFORE capacity. False means expert branch skipped, residual retained. These are toy results, not paper benchmarks. Expand any run.</p>'
(out/'routing_heatmap.html').write_text(html + ''.join(sections) + '</html>')
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--output', type=Path, default=Path(__file__).parent/'results')
parser.add_argument('--check-only', action='store_true')
args = parser.parse_args()
check = checks()
print(json.dumps(check, indent=2))
if args.check_only:
return
out = args.output; out.mkdir(parents=True, exist_ok=True)
data = make_data()
(out/'data.json').write_text(json.dumps(data, indent=2))
print('CPU', platform.python_version(), 'NumPy', np.__version__, 'X=[64,4], logits=[64,4], experts=[4,4], parameters=32')
results = [train(data, seed, name, a, b, out) for seed in [70,71,72]
for name, a, b in [('task', 0, 0), ('balance', 0.1, 0), ('balance_z', 0.1, 0.001)]]
summary = {}
for name in ['task', 'balance', 'balance_z']:
subset = [r for r in results if r['name'] == name]
summary[name] = {k: dict(mean=float(np.mean([r['test'][k] for r in subset])),
sd=float(np.std([r['test'][k] for r in subset], ddof=1)))
for k in ['mse','drop','max_share','load_cv','balance','z_loss','entropy']}
print(name, json.dumps(summary[name]))
(out/'summary.json').write_text(json.dumps(dict(checks=check, results=summary,
environment=dict(python=platform.python_version(), numpy=np.__version__, device='CPU', dtype='float64'),
protocol='Fixed600 steps; Adam0.03; batch64; capacity1.25; same data/initial weights/batch stream within seed; no dev/test selection'), indent=2))
html_report(results, out)
if __name__ == '__main__':
main()
