analytical-viz: add Auto Explore Streamlit tab
New tab consumes auto_explore.run_auto_explore(). UI:
- Header showing current model + workload + per-PE HBM budget
- Run sweep button (spinner while ~28k configs run in ~5-10s)
- 4 metric cards: enumerated, feasible, pareto count, best latency
- 2-panel scatter:
* latency vs PEs (feasible in grey, Pareto coloured by efficiency,
dashed line connects Pareto in PE order to show the trade-off curve)
* HBM utilization vs latency for Pareto, coloured by PE count
- Sortable Pareto table
- "Load into sidebar" widget: pick a row, sets the sidebar
session_state keys (cp, tp, pp, dp, tp_placement, cp_placement,
cp_ring_variant, kv_mode, ffn_scope_label) and st.rerun()s so the
user can flip to another tab and see the full breakdown.
Session-state caches the last sweep result under _auto_explore_result;
if the model/workload/HBM change without re-running, a warning suggests
refreshing.
Ffn_scope_label mapping is dynamic (contains substituted divisors), so
the Load button reconstructs the exact label using the target
cp/tp/dp values before assigning to session_state["ffn_scope_label"].
Verified:
- app.py parses cleanly
- All 9 pytest tests in test_auto_explore.py still pass
- Smoke: matplotlib + pandas + auto_explore imports round-trip
Next: verification pass across Qwen 3 8B and Mixtral 8x7B; polish.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -419,9 +419,9 @@ if _warnings:
|
|||||||
|
|
||||||
|
|
||||||
# ── Tabs ─────────────────────────────────────────────────────────
|
# ── Tabs ─────────────────────────────────────────────────────────
|
||||||
tab_layout, tab_memory, tab_stages, tab_compare = st.tabs([
|
tab_layout, tab_memory, tab_stages, tab_compare, tab_auto = st.tabs([
|
||||||
"Physical layout", "Memory breakdown", "Per-stage latency",
|
"Physical layout", "Memory breakdown", "Per-stage latency",
|
||||||
"Save & compare",
|
"Save & compare", "Auto Explore",
|
||||||
])
|
])
|
||||||
|
|
||||||
|
|
||||||
@@ -1282,3 +1282,181 @@ with tab_compare:
|
|||||||
"'-' means the stage is absent for that config (e.g. C1 in "
|
"'-' means the stage is absent for that config (e.g. C1 in "
|
||||||
"decode, since the O/m/l all-reduce is folded into S8)."
|
"decode, since the O/m/l all-reduce is folded into S8)."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── TAB 5: Auto Explore ─────────────────────────────────────────
|
||||||
|
with tab_auto:
|
||||||
|
from tests.analytical_visualization.auto_explore import run_auto_explore
|
||||||
|
|
||||||
|
st.markdown(
|
||||||
|
"Sweep every valid parallelism config (CP, TP, PP, DP, kv_shard, "
|
||||||
|
"ffn_scope, tp_placement, cp_placement, cp_ring) for the currently-"
|
||||||
|
"selected model + workload and extract the 3D Pareto frontier on "
|
||||||
|
"(latency ↓, PEs ↓, efficiency ↑). Throughput = 1/latency for a "
|
||||||
|
"single request so it collapses with latency and is shown for info only."
|
||||||
|
)
|
||||||
|
|
||||||
|
ax_l, ax_r = st.columns([3, 1])
|
||||||
|
with ax_l:
|
||||||
|
st.markdown(
|
||||||
|
f"**Context:** {model.name} | S_kv=**{s_kv:,}** | mode=**{mode}** | "
|
||||||
|
f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**"
|
||||||
|
)
|
||||||
|
with ax_r:
|
||||||
|
_run_sweep = st.button("Run sweep", type="primary", width='stretch')
|
||||||
|
|
||||||
|
if _run_sweep or st.session_state.get("_auto_explore_result") is not None:
|
||||||
|
if _run_sweep:
|
||||||
|
with st.spinner("Sweeping ~28k configs..."):
|
||||||
|
_res = run_auto_explore(model, machine, s_kv=s_kv, mode=mode)
|
||||||
|
st.session_state["_auto_explore_result"] = _res
|
||||||
|
st.session_state["_auto_explore_ctx"] = (
|
||||||
|
model.name, s_kv, mode, machine.pe_hbm_gb,
|
||||||
|
)
|
||||||
|
_res = st.session_state["_auto_explore_result"]
|
||||||
|
|
||||||
|
# Warn if the cached result is stale relative to the current config.
|
||||||
|
_cached_ctx = st.session_state.get("_auto_explore_ctx")
|
||||||
|
_cur_ctx = (model.name, s_kv, mode, machine.pe_hbm_gb)
|
||||||
|
if _cached_ctx != _cur_ctx:
|
||||||
|
st.warning(
|
||||||
|
"Sweep result is from a different model/workload. "
|
||||||
|
"Click **Run sweep** to refresh."
|
||||||
|
)
|
||||||
|
|
||||||
|
_m1, _m2, _m3, _m4 = st.columns(4)
|
||||||
|
_m1.metric("Enumerated", f"{_res.total_enumerated:,}")
|
||||||
|
_m2.metric("Feasible", f"{_res.total_feasible:,}")
|
||||||
|
_m3.metric("Pareto configs", len(_res.pareto_scores))
|
||||||
|
if _res.pareto_scores:
|
||||||
|
_best = min(_res.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||||
|
_m4.metric("Best latency", f"{_best.latency_ms:.2f} ms")
|
||||||
|
|
||||||
|
if not _res.pareto_scores:
|
||||||
|
st.error(
|
||||||
|
"No feasible configs — every parallelism choice exceeds "
|
||||||
|
"per-PE HBM. Try raising `pe_hbm_gb` in the sidebar or "
|
||||||
|
"reducing `s_kv`."
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# ── Pareto scatter: latency vs PEs, coloured by efficiency ─
|
||||||
|
import matplotlib.pyplot as _plt
|
||||||
|
_fig, _axs = _plt.subplots(1, 2, figsize=(10, 3.5))
|
||||||
|
|
||||||
|
# All feasible = grey points; Pareto = coloured by efficiency.
|
||||||
|
_feas = [s for s in _res.all_scores
|
||||||
|
if s.fits_memory and s.placement_valid]
|
||||||
|
_axs[0].scatter(
|
||||||
|
[s.pes_used for s in _feas],
|
||||||
|
[s.latency_ms for s in _feas],
|
||||||
|
c="lightgrey", s=8, alpha=0.4, label="feasible",
|
||||||
|
)
|
||||||
|
_pareto_sorted = sorted(_res.pareto_scores,
|
||||||
|
key=lambda s: s.pes_used)
|
||||||
|
_sc = _axs[0].scatter(
|
||||||
|
[s.pes_used for s in _pareto_sorted],
|
||||||
|
[s.latency_ms for s in _pareto_sorted],
|
||||||
|
c=[s.efficiency_score for s in _pareto_sorted],
|
||||||
|
cmap="viridis", s=60, edgecolors="black", linewidths=0.8,
|
||||||
|
label="Pareto",
|
||||||
|
)
|
||||||
|
# Connect Pareto in PE order to show the trade-off curve.
|
||||||
|
_axs[0].plot(
|
||||||
|
[s.pes_used for s in _pareto_sorted],
|
||||||
|
[s.latency_ms for s in _pareto_sorted],
|
||||||
|
"k--", linewidth=0.8, alpha=0.5,
|
||||||
|
)
|
||||||
|
_axs[0].set_xlabel("PEs used"); _axs[0].set_ylabel("Latency (ms)")
|
||||||
|
_axs[0].set_xscale("log", base=2)
|
||||||
|
_axs[0].set_yscale("log")
|
||||||
|
_axs[0].set_title("Pareto: latency vs PEs")
|
||||||
|
_axs[0].grid(alpha=0.3)
|
||||||
|
_axs[0].legend(loc="upper right", fontsize=8)
|
||||||
|
_fig.colorbar(_sc, ax=_axs[0], label="efficiency")
|
||||||
|
|
||||||
|
# HBM utilization vs latency for Pareto — where does the config land?
|
||||||
|
_pareto_by_lat = sorted(_res.pareto_scores,
|
||||||
|
key=lambda s: s.total_latency_ns)
|
||||||
|
_axs[1].scatter(
|
||||||
|
[s.hbm_utilization * 100 for s in _pareto_by_lat],
|
||||||
|
[s.latency_ms for s in _pareto_by_lat],
|
||||||
|
c=[s.pes_used for s in _pareto_by_lat],
|
||||||
|
cmap="plasma", s=60, edgecolors="black", linewidths=0.8,
|
||||||
|
)
|
||||||
|
_axs[1].set_xlabel("HBM utilization (%)")
|
||||||
|
_axs[1].set_ylabel("Latency (ms)")
|
||||||
|
_axs[1].set_yscale("log")
|
||||||
|
_axs[1].set_title("HBM usage vs latency (Pareto)")
|
||||||
|
_axs[1].grid(alpha=0.3)
|
||||||
|
_sc2 = _axs[1].collections[0]
|
||||||
|
_fig.colorbar(_sc2, ax=_axs[1], label="PEs")
|
||||||
|
|
||||||
|
_fig.tight_layout()
|
||||||
|
st.pyplot(_fig, width='stretch')
|
||||||
|
_plt.close(_fig)
|
||||||
|
|
||||||
|
# ── Pareto table + Load button ─────────────────────────────
|
||||||
|
st.markdown("**Pareto configurations** (sorted by latency)")
|
||||||
|
_rows = []
|
||||||
|
for i, s in enumerate(_pareto_by_lat):
|
||||||
|
_rows.append({
|
||||||
|
"#": i,
|
||||||
|
"CP": s.cp, "TP": s.tp, "PP": s.pp, "DP": s.dp,
|
||||||
|
"kv": s.kv_shard_mode,
|
||||||
|
"ffn": s.ffn_shard_scope,
|
||||||
|
"tp_place": s.tp_placement,
|
||||||
|
"cp_place": s.cp_placement,
|
||||||
|
"cp_ring": s.cp_ring_variant,
|
||||||
|
"lat (ms)": round(s.latency_ms, 3),
|
||||||
|
"eff": round(s.efficiency_score, 4),
|
||||||
|
"PEs": s.pes_used,
|
||||||
|
"SIPs": s.sips_used,
|
||||||
|
"HBM %": round(s.hbm_utilization * 100, 1),
|
||||||
|
})
|
||||||
|
_df = pd.DataFrame(_rows)
|
||||||
|
st.dataframe(_df, width='stretch', hide_index=True)
|
||||||
|
|
||||||
|
# ── Load-into-sidebar ─────────────────────────────────────
|
||||||
|
st.markdown("**Load a Pareto config into the main sidebar sliders**")
|
||||||
|
_lc1, _lc2 = st.columns([1, 3])
|
||||||
|
with _lc1:
|
||||||
|
_pick = st.number_input(
|
||||||
|
"Row #", min_value=0, max_value=len(_pareto_by_lat) - 1,
|
||||||
|
value=0, step=1, key="_auto_pick_row",
|
||||||
|
)
|
||||||
|
with _lc2:
|
||||||
|
if st.button("Load into sidebar", type="secondary"):
|
||||||
|
_s = _pareto_by_lat[int(_pick)]
|
||||||
|
# Set the sidebar's session_state keys directly. The
|
||||||
|
# sidebar reads these on the next rerun.
|
||||||
|
st.session_state["cp"] = _s.cp
|
||||||
|
st.session_state["tp"] = _s.tp
|
||||||
|
st.session_state["pp"] = _s.pp
|
||||||
|
st.session_state["dp"] = _s.dp
|
||||||
|
st.session_state["tp_placement"] = _s.tp_placement
|
||||||
|
st.session_state["cp_placement"] = _s.cp_placement
|
||||||
|
st.session_state["cp_ring_variant"] = _s.cp_ring_variant
|
||||||
|
st.session_state["kv_mode"] = _s.kv_shard_mode
|
||||||
|
# ffn_scope label is dynamic: "TP+CP (div=…)" etc.
|
||||||
|
# Reconstruct with the target cp/tp/dp values.
|
||||||
|
_tp, _cp, _dp = _s.tp, _s.cp, _s.dp
|
||||||
|
_ffn_label_map = {
|
||||||
|
"TP": f"TP only (div={max(1, _tp)})",
|
||||||
|
"TP+CP": f"TP*CP (div={max(1, _tp * _cp)})",
|
||||||
|
"TP+CP+DP": f"TP*CP*DP (div={max(1, _tp * _cp * _dp)})",
|
||||||
|
}
|
||||||
|
st.session_state["ffn_scope_label"] = \
|
||||||
|
_ffn_label_map[_s.ffn_shard_scope]
|
||||||
|
st.success(
|
||||||
|
f"Loaded row {_pick}: CP={_s.cp} TP={_s.tp} PP={_s.pp} "
|
||||||
|
f"DP={_s.dp} → latency {_s.latency_ms:.2f} ms, "
|
||||||
|
f"{_s.pes_used} PEs. Flip to another tab to see the "
|
||||||
|
f"full breakdown."
|
||||||
|
)
|
||||||
|
st.rerun()
|
||||||
|
else:
|
||||||
|
st.info(
|
||||||
|
"Click **Run sweep** to enumerate all valid parallelism "
|
||||||
|
"configurations and rank them on the Pareto frontier. "
|
||||||
|
"Takes ~5–10 s for the full search."
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user