analyze: add MoE router gating/share diagnostic plots
New "model" family in the gallery: router_gating (mean soft gate weight vs. pre-step energy, showing the router's soft decision boundaries) and router_share_by_pdg/router_share_by_process (stacked top-1 dispatch share by species / true physics process). Needs a live checkpoint's Router, so it's a documented exception to the rest of the package's polars/numpy-only contract; gracefully degrades to a placeholder for non-MoE checkpoints. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
@@ -137,12 +137,89 @@ def _render_bar(r: Reduced, params: dict):
|
||||
return fig
|
||||
|
||||
|
||||
def _render_router_gating(r: Reduced, params: dict):
|
||||
n_experts = r.payload["n_experts"]
|
||||
log_x = r.payload.get("log_x", False)
|
||||
fig, axes = ps.new_figure(
|
||||
"slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False
|
||||
)
|
||||
flat = axes.ravel()
|
||||
for ax, key in zip(flat, ("rollout", "reference")):
|
||||
side = r.payload.get(key, {})
|
||||
centers = np.asarray(side.get("centers", []))
|
||||
means = np.asarray(side.get("means", []))
|
||||
if len(centers) and means.size:
|
||||
cum = np.zeros(len(centers))
|
||||
for i in range(n_experts):
|
||||
ax.fill_between(
|
||||
centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}"
|
||||
)
|
||||
cum = cum + means[:, i]
|
||||
if log_x:
|
||||
ax.set_xscale("log")
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_title(_SERIES_LABELS[key], fontsize=8)
|
||||
ax.set_xlabel(r.xlabel)
|
||||
flat[0].set_ylabel("mean gate weight")
|
||||
ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
|
||||
return fig
|
||||
|
||||
|
||||
def _render_router_share(r: Reduced, params: dict):
|
||||
categories = r.payload["categories"]
|
||||
n_experts = r.payload["n_experts"]
|
||||
x = np.arange(len(categories))
|
||||
present = [k for k in ("rollout", "reference") if k in r.payload]
|
||||
fig, axes = ps.new_figure(
|
||||
"slide-16x9",
|
||||
title=r.title,
|
||||
params=params,
|
||||
nrows=1,
|
||||
ncols=len(present),
|
||||
squeeze=False,
|
||||
)
|
||||
flat = axes.ravel()
|
||||
for ax, key in zip(flat, present):
|
||||
side = r.payload[key]
|
||||
shares = np.array([side[c] for c in categories]) # (n_cat, n_experts)
|
||||
bottom = np.zeros(len(categories))
|
||||
for i in range(n_experts):
|
||||
ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
|
||||
bottom += shares[:, i]
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(categories, rotation=45, ha="right")
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_title(_SERIES_LABELS[key], fontsize=8)
|
||||
flat[0].set_ylabel("share of rows dispatched to expert")
|
||||
ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
|
||||
return fig
|
||||
|
||||
|
||||
def _render_unavailable(r: Reduced, params: dict):
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
ax.axis("off")
|
||||
ax.text(
|
||||
0.5,
|
||||
0.5,
|
||||
r.payload.get("note", "not available"),
|
||||
ha="center",
|
||||
va="center",
|
||||
wrap=True,
|
||||
fontsize=10,
|
||||
transform=ax.transAxes,
|
||||
)
|
||||
return fig
|
||||
|
||||
|
||||
_RENDERERS = {
|
||||
"overlay_hist": _render_overlay,
|
||||
"single_hist": _render_single,
|
||||
"grouped_hist": _render_grouped,
|
||||
"profile": _render_profile,
|
||||
"bar": _render_bar,
|
||||
"router_gating": _render_router_gating,
|
||||
"router_share": _render_router_share,
|
||||
"unavailable": _render_unavailable,
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user