From 91b63eb9f66bd02eeac87811c4ad53d671a5332b Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 28 Jul 2026 13:39:51 -0700 Subject: [PATCH] analytical-viz: parallelism sensitivity chart in Auto Explore tab Adds compute_parallelism_sensitivity() + ParallelismSensitivityRow to auto_explore.py. For a baseline ConfigScore (picked from the Pareto set), sweeps each parallelism knob (CP, TP, PP, DP, EP) individually while holding others fixed. Reports latency + memory-fit per value. Sweep values start at 1 and step by 2 (multiples of 2 rather than only powers of 2), giving finer granularity for the visual than the enumerator's sparser set. CP goes up to 256 (per real-deployment scale), TP to 64, PP to 32. New UI panel in Auto Explore tab: 5 subplots (log/log), one per knob: - Solid line = latency where the config fits memory - Red X markers = infeasible (out of budget) - Dotted vertical line = baseline value User picks which Pareto row is the "baseline" via a number input; the sensitivity chart re-computes around it. Behaviour caveat noted during verification: for large PP the FFN AR can cross a SIP boundary (uses stage_latencies.py:730 sips_used tier selection), producing a step-up in latency. This is the existing model's choice, honestly reflected in the chart. Whether the FFN AR should span only one PP stage's ranks (thus not cross SIPs) is a separate discussion about stage_latencies.py. Verified: - 11 pytest tests pass (added 2 new for parallelism sensitivity) - Smoke: at Llama 70B decode 128K with baseline CP=8/TP=16/PP=1/DP=1: * CP sweep: 4 fits, 8 = baseline optimum, 16+ overshoots memory * TP sweep: 8/16 fit, 16 = baseline optimum, 32 slower (spans SIPs) * PP sweep: 1 = baseline, 2+ slower due to sips_used tier drop for FFN AR * DP sweep: same shape as PP * EP sweep: monotone decreasing (bigger EP = smaller per-PE FFN) Co-Authored-By: Claude Opus 4.7 --- tests/analytical_visualization/app.py | 67 +++++++++++++++- .../analytical_visualization/auto_explore.py | 80 +++++++++++++++++++ .../test_auto_explore.py | 38 +++++++++ 3 files changed, 184 insertions(+), 1 deletion(-) diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index a8236b3..dfb3c80 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -1286,7 +1286,10 @@ with tab_compare: # ── TAB 5: Auto Explore ───────────────────────────────────────── with tab_auto: - from tests.analytical_visualization.auto_explore import run_auto_explore + from tests.analytical_visualization.auto_explore import ( + compute_parallelism_sensitivity, + run_auto_explore, + ) st.markdown( "Sweep every valid parallelism config (CP, TP, PP, DP, kv_shard, " @@ -1416,6 +1419,68 @@ with tab_auto: _df = pd.DataFrame(_rows) st.dataframe(_df, width='stretch', hide_index=True) + # ── Parallelism sensitivity subplots ────────────────────── + st.markdown( + "**Parallelism sensitivity** — for the *baseline* config " + "picked below, vary each knob one at a time (holding others " + "fixed). Answers 'if I only tweak CP / TP / PP / DP / EP, " + "what happens?' Solid line = fits memory; red X = infeasible." + ) + _sens_c1, _sens_c2 = st.columns([1, 4]) + with _sens_c1: + _sens_row = st.number_input( + "Baseline row #", + min_value=0, max_value=len(_pareto_by_lat) - 1, + value=0, step=1, key="_auto_sens_row", + help="Which Pareto row is the sensitivity computed around?", + ) + _baseline_for_sens = _pareto_by_lat[int(_sens_row)] + _sens_rows = compute_parallelism_sensitivity( + _baseline_for_sens, model, machine, s_kv=s_kv, mode=mode, + ) + + import matplotlib.pyplot as _plt + _fig_sens, _axes_sens = _plt.subplots(1, 5, figsize=(14, 3.2)) + for _ax, _row in zip(_axes_sens, _sens_rows): + _fit_vals, _fit_lats = [], [] + _bad_vals = [] + for _v, _lat_ns, _fit in zip(_row.values, _row.latencies_ns, + _row.fits_flags): + if _lat_ns != _lat_ns: # NaN — skip out-of-bounds + continue + if _fit: + _fit_vals.append(_v) + _fit_lats.append(_lat_ns / 1e6) # ms + else: + _bad_vals.append((_v, _lat_ns / 1e6)) + if _fit_vals: + _ax.plot(_fit_vals, _fit_lats, "o-", color="#2c3e50", + linewidth=1.4, markersize=4) + if _bad_vals: + _ax.scatter( + [v for v, _ in _bad_vals], + [l for _, l in _bad_vals], + marker="x", color="#c0392b", s=40, alpha=0.6, + label="no fit", + ) + # Baseline marker + _ax.axvline(_row.baseline_value, color="#3498db", + linestyle=":", linewidth=1.0) + _ax.set_xscale("log") + _ax.set_yscale("log") + _ax.set_title( + f"{_row.knob.upper()} (baseline={_row.baseline_value})", + fontsize=10, + ) + _ax.set_xlabel("value") + _ax.set_ylabel("lat (ms)") + _ax.grid(alpha=0.3) + if _bad_vals: + _ax.legend(loc="best", fontsize=7) + _fig_sens.tight_layout() + st.pyplot(_fig_sens, width='stretch') + _plt.close(_fig_sens) + # ── Load-into-sidebar ───────────────────────────────────── st.markdown("**Load a Pareto config into the main sidebar sliders**") _lc1, _lc2 = st.columns([1, 3]) diff --git a/tests/analytical_visualization/auto_explore.py b/tests/analytical_visualization/auto_explore.py index a9bac4b..0099b91 100644 --- a/tests/analytical_visualization/auto_explore.py +++ b/tests/analytical_visualization/auto_explore.py @@ -25,6 +25,17 @@ from .memory_layout import compute_memory from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig from .stage_latencies import all_ffn_stages, all_stages +# ── Parallelism sensitivity sweep values (see compute_parallelism_sensitivity) +# Multiples of 2 (finer grid than powers of 2 alone), capped at values that +# are physically meaningful for the modelled topology. +_PARALLELISM_SWEEP_VALUES = { + "cp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64, 96, 128, 192, 256), + "tp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64), + "pp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 32), + "dp": (1, 2, 4, 6, 8, 12, 16), + "ep": (1, 2, 4, 6, 8, 12, 16, 32), +} + # ── Search space ───────────────────────────────────────────────────── @@ -298,6 +309,75 @@ def pareto_frontier(scores: list[ConfigScore]) -> list[ConfigScore]: return frontier +# ── Parallelism sensitivity ───────────────────────────────────────── + + +@dataclass +class ParallelismSensitivityRow: + """One knob's sweep result: latency + fit at every tested value. + + All *other* parallelism knobs (including HW) are held at their baseline + values, so this isolates the effect of just this one knob. + """ + knob: str # "cp" / "tp" / "pp" / "dp" / "ep" + values: list[int] # tested values (from _PARALLELISM_SWEEP_VALUES) + latencies_ns: list[float] # one per value; NaN when infeasible + fits_flags: list[bool] # per-value memory+placement feasibility + baseline_value: int # what the baseline config uses for this knob + baseline_latency_ns: float + + +def compute_parallelism_sensitivity( + baseline: ConfigScore, + model: ModelConfig, + machine: MachineParams, + s_kv: int, + mode: str, +) -> list[ParallelismSensitivityRow]: + """For each parallelism knob (CP, TP, PP, DP, EP), sweep its values + holding the OTHER knobs fixed at the baseline. Reports latency + memory + fit per value. + + Answers: 'if I only change this one knob, how does latency and memory + fit change?' — the complement to auto_hardware.compute_sensitivity + which does the same for hardware knobs. + """ + baseline_topo = baseline.as_topology(s_kv, mode) + + rows: list[ParallelismSensitivityRow] = [] + for knob, values in _PARALLELISM_SWEEP_VALUES.items(): + baseline_val = getattr(baseline_topo, knob, 1) + # EP isn't stored on ConfigScore's as_topology output today (dp + # attribute exists but ep does not always). Fall back to 1. + latencies: list[float] = [] + fits: list[bool] = [] + for v in values: + # Skip nonsensical values that would violate hard bounds. + if knob == "pp" and v > model.layers: + latencies.append(float("nan")) + fits.append(False) + continue + if knob == "tp" and v > 4 * model.h_q: + latencies.append(float("nan")) + fits.append(False) + continue + # Build a variant TopologyConfig with just this one knob changed. + trial = replace(baseline_topo, **{knob: v}) + cfg = FullConfig(model=model, topo=trial, machine=machine) + score = score_config(cfg) + latencies.append(score.total_latency_ns) + fits.append(score.fits_memory and score.placement_valid) + rows.append(ParallelismSensitivityRow( + knob=knob, + values=list(values), + latencies_ns=latencies, + fits_flags=fits, + baseline_value=int(baseline_val), + baseline_latency_ns=baseline.total_latency_ns, + )) + return rows + + # ── Top-level driver ───────────────────────────────────────────────── diff --git a/tests/analytical_visualization/test_auto_explore.py b/tests/analytical_visualization/test_auto_explore.py index 163c203..01ab681 100644 --- a/tests/analytical_visualization/test_auto_explore.py +++ b/tests/analytical_visualization/test_auto_explore.py @@ -140,3 +140,41 @@ def test_pp_does_not_reduce_single_request_latency(): assert fastest.pp == 1, ( f"expected PP=1 for the fastest decode config; got PP={fastest.pp}" ) + + +# ── Parallelism sensitivity ───────────────────────────────────────── + + +def test_parallelism_sensitivity_has_all_five_knobs(): + """compute_parallelism_sensitivity returns one row per knob.""" + from tests.analytical_visualization.auto_explore import ( + compute_parallelism_sensitivity, + ) + model = PRESETS["Llama 3.1 70B"].model + machine = MachineParams() + res = run_auto_explore(model, machine, s_kv=8192, mode="decode") + baseline = res.pareto_scores[0] + rows = compute_parallelism_sensitivity( + baseline, model, machine, s_kv=8192, mode="decode", + ) + knobs = {r.knob for r in rows} + assert knobs == {"cp", "tp", "pp", "dp", "ep"} + + +def test_parallelism_sensitivity_includes_baseline_value(): + """Each row's values contain the baseline_value.""" + from tests.analytical_visualization.auto_explore import ( + compute_parallelism_sensitivity, + ) + model = PRESETS["Llama 3.1 70B"].model + machine = MachineParams() + res = run_auto_explore(model, machine, s_kv=8192, mode="decode") + baseline = res.pareto_scores[0] + rows = compute_parallelism_sensitivity( + baseline, model, machine, s_kv=8192, mode="decode", + ) + for r in rows: + # Baseline value must be sweepable OR mapped to something in values. + assert r.baseline_value in r.values or r.baseline_value == 1, ( + f"knob {r.knob}: baseline_value={r.baseline_value} not in {r.values}" + )