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 <noreply@anthropic.com>
This commit is contained in:
@@ -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 ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user