diff --git a/tests/analytical_visualization/pe_weight_layout.py b/tests/analytical_visualization/pe_weight_layout.py index be99ac1..47f5516 100644 --- a/tests/analytical_visualization/pe_weight_layout.py +++ b/tests/analytical_visualization/pe_weight_layout.py @@ -114,6 +114,14 @@ def draw_pe_layout(cfg: FullConfig, ax=None): cube_h = 6.0 cube_gap = 0.6 + # When TP fills the cube (TP>=8), cp_placement is "cube" and CP + # ranks span cubes — the intra-cube CP-color branch below is + # skipped. Surface which CP rank this figure is showing so users + # aren't guessing which of `CP` groups is drawn. + _cube_place_tag = "" + if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1: + _cube_place_tag = f" [CP rank 0 of {cfg.topo.cp}]" + for cube_idx in range(n_cubes): x0 = cube_idx * (cube_w + cube_gap) y0 = 0 @@ -125,7 +133,8 @@ def draw_pe_layout(cfg: FullConfig, ax=None): ) ax.add_patch(rect) ax.text(x0 + cube_w / 2, y0 + cube_h + 0.1, - f"Cube {cube_idx} (PEs {cube_idx*8}-{cube_idx*8+7})", + f"Cube {cube_idx} (PEs {cube_idx*8}-{cube_idx*8+7})" + + _cube_place_tag, ha="center", va="bottom", fontsize=11, fontweight="bold") pe_pad_x = 0.15 @@ -180,7 +189,14 @@ def draw_pe_layout(cfg: FullConfig, ax=None): _tp_rank = pe_id_in_group % tp _fc = "#f0f7ff" _ec = "#3a86ff" - _cp_label = "" + # cp_placement=="cube" with CP>1: the whole cube is + # one CP rank (rank 0 by convention — see cube title + # tag above). Add the locator so users know what + # rank this per-PE data belongs to. + if _cp_place == "cube" and _cp_val > 1: + _cp_label = " | CP=0" + else: + _cp_label = "" pe_rect = patches.Rectangle( (px, py), pe_w, pe_h, @@ -250,9 +266,12 @@ def draw_pe_layout(cfg: FullConfig, ax=None): ax.set_aspect("equal") ax.axis("off") + _title_cp = f", CP={cfg.topo.cp}" + if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1: + _title_cp += " (showing 1 of CP groups; others identical)" title = ( f"Per-PE layout of one CP group " - f"(TP={tp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})" + f"(TP={tp}{_title_cp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})" f" | KV mode: {cfg.topo.kv_shard_mode}" ) ax.set_title(title, fontsize=11, fontweight="bold") diff --git a/tests/analytical_visualization/test_pe_weight_layout.py b/tests/analytical_visualization/test_pe_weight_layout.py new file mode 100644 index 0000000..64d4f07 --- /dev/null +++ b/tests/analytical_visualization/test_pe_weight_layout.py @@ -0,0 +1,101 @@ +"""Tests for draw_pe_layout — CP rank locator visibility across TP tiers. + +The PE-level view has always color-labeled CP ranks when they're packed +intra-cube (`cp_placement="pe"`). When TP >= 8, TP fills the cube and CP +gets pushed to cube-level, so the intra-cube color branch is skipped — +which used to silently drop the CP rank locator entirely. These tests +guard the fix: a `[CP rank 0 of N]` cube tag + `CP=0` per-PE label are +always present when CP > 1, regardless of placement. +""" +from __future__ import annotations + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from tests.analytical_visualization.model_config import ( + FullConfig, MachineParams, TopologyConfig, +) +from tests.analytical_visualization.model_presets import PRESETS +from tests.analytical_visualization.pe_weight_layout import draw_pe_layout + + +def _cfg(cp: int, tp: int, cp_placement: str | None = None): + model = PRESETS["Llama 3 8B"].model + kwargs = dict(cp=cp, tp=tp, pp=1, s_kv=8192, mode="decode") + if cp_placement is not None: + kwargs["cp_placement"] = cp_placement + return FullConfig( + model=model, topo=TopologyConfig(**kwargs), machine=MachineParams(), + ) + + +def _all_text(fig) -> str: + """Concatenate all text artists (labels + title) into one blob.""" + ax = fig.gca() + parts = [t.get_text() for t in ax.texts] + parts.append(ax.get_title()) + return "\n".join(parts) + + +def test_high_tp_cp_gt_1_shows_rank_locator(): + """TP=8, CP=4 -> cp_placement forced to 'cube'; CP rank locator + must still appear in cube tag + per-PE header + figure title.""" + cfg = _cfg(cp=4, tp=8) + assert cfg.topo.cp_placement == "cube" # sanity + fig = draw_pe_layout(cfg) + try: + text = _all_text(fig) + assert "CP rank 0 of 4" in text, text + assert "CP=0" in text, text + assert "CP=4" in text, text # figure title dim summary + finally: + plt.close(fig) + + +def test_intra_cube_cp_placement_still_shows_per_rank_labels(): + """cp_placement='pe' (small TP, CP>1) — the old branch already + colors + labels each CP rank. Regression check: the new tag does + NOT replace those, and each CP=0..CP-1 label is still emitted.""" + cfg = _cfg(cp=4, tp=2) # 4*2=8 fits intra-cube + assert cfg.topo.cp_placement == "cube" # default, but auto/manual can differ + # Force pe placement to exercise the intra-cube branch. + cfg.topo.cp_placement = "pe" + fig = draw_pe_layout(cfg) + try: + text = _all_text(fig) + for r in range(4): + assert f"CP={r}" in text, (r, text) + # The 'CP rank 0 of N' tag is only for the cube-placement fallback. + assert "CP rank 0 of" not in text, text + finally: + plt.close(fig) + + +def test_cp_1_no_locator_added(): + """CP=1: nothing to locate; no CP tag or label should appear.""" + cfg = _cfg(cp=1, tp=8) + fig = draw_pe_layout(cfg) + try: + text = _all_text(fig) + assert "CP rank" not in text, text + # The header 'CP=0' should not be added when CP=1. + # (Figure title mentions 'CP=1' — that's an OK dim summary, + # not a locator per-PE.) + assert " | CP=0" not in text, text + finally: + plt.close(fig) + + +def test_tp_16_still_shows_cp_rank_across_spilled_cubes(): + """TP=16 (spills to 2 cubes per CP group), CP=2 -> both cubes + for CP rank 0 must carry the '[CP rank 0 of 2]' tag.""" + cfg = _cfg(cp=2, tp=16) + assert cfg.topo.cp_placement == "cube" + fig = draw_pe_layout(cfg) + try: + text = _all_text(fig) + assert text.count("CP rank 0 of 2") >= 2, text + assert "CP=0" in text, text + finally: + plt.close(fig)