From ed28eea1568b0797c868c7b4c6186d99793c8aee Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 28 Jul 2026 21:59:18 -0700 Subject: [PATCH] analytical-viz: color-code CP groups in PE layout when cp_placement=pe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When cp_placement=pe packs multiple CP ranks intra-cube, the PE-level layout now: 1. Colors each PE's background by its CP rank (8-color palette wraps for cp > 8). Cp_placement=cube keeps the historical light-blue styling (only one CP rank per cube tile). 2. Adds "TP=N | CP=M" to the per-PE header so users can read the (cp_rank, tp_rank) pair without inferring from position. 3. Fixes a head-assignment bug: _q_heads_for_pe / _kv_heads_for_pe / _per_pe_bytes used to be called with the raw PE-in-group index, which treated every PE as a distinct TP rank. Under cp_placement=pe this gave wrong Q/KV head lists for PEs beyond the first TP group (PE 2..7 with tp=2, cp=4 would ask for head slots 32..127 in a 32-head model). Now called with tp_rank = pe_id % tp, so all CP ranks in the cube share the correct head split. Verified via a headless matplotlib smoke test with CP=4/TP=2 on Qwen 3 8B: 8 PE patches + 1 cube patch → 5 distinct facecolors (cube + 4 CP ranks), no errors. Co-Authored-By: Claude Opus 4.7 --- .../pe_weight_layout.py | 49 +++++++++++++++++-- 1 file changed, 44 insertions(+), 5 deletions(-) diff --git a/tests/analytical_visualization/pe_weight_layout.py b/tests/analytical_visualization/pe_weight_layout.py index f3d1d16..be99ac1 100644 --- a/tests/analytical_visualization/pe_weight_layout.py +++ b/tests/analytical_visualization/pe_weight_layout.py @@ -136,6 +136,26 @@ def draw_pe_layout(cfg: FullConfig, ax=None): pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS + # When cp_placement=pe, CP ranks are packed intra-cube along with + # TP ranks. Color-code by CP rank so it's obvious which PEs belong + # to which sequence-shard group. Palette wraps beyond 8 CP ranks. + _cp_place = cfg.topo.cp_placement + _cp_val = cfg.topo.cp + _cp_palette = [ + "#f0f7ff", # baseline blue (also the fallback for cp_placement=cube) + "#e8f5e9", # light green + "#fff3e0", # light orange + "#f3e5f5", # light purple + "#ffebee", # light red + "#e0f7fa", # light cyan + "#fce4ec", # light pink + "#f9fbe7", # light lime + ] + _cp_edge_palette = [ + "#3a86ff", "#2e7d32", "#ef6c00", "#7b1fa2", + "#c62828", "#0097a7", "#c2185b", "#9e9d24", + ] + for pr in range(PE_ROWS): for pc in range(PE_COLS): pe_local = pr * PE_COLS + pc @@ -145,17 +165,36 @@ def draw_pe_layout(cfg: FullConfig, ax=None): px = x0 + pe_pad_x + pc * (pe_w + pe_gap) py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap) + # Which (cp_rank, tp_rank) does this PE hold? + # When cp_placement=pe: cp_rank = pe_id // tp, tp_rank = pe_id % tp. + # When cp_placement=cube: every PE in this cube is one CP rank + # (drawn on its own tile), tp_rank = pe_id % tp. + if _cp_place == "pe" and _cp_val > 1: + _cp_rank = pe_id_in_group // tp + _tp_rank = pe_id_in_group % tp + _fc = _cp_palette[_cp_rank % len(_cp_palette)] + _ec = _cp_edge_palette[_cp_rank % len(_cp_edge_palette)] + _cp_label = f" | CP={_cp_rank}" + else: + _cp_rank = None + _tp_rank = pe_id_in_group % tp + _fc = "#f0f7ff" + _ec = "#3a86ff" + _cp_label = "" + pe_rect = patches.Rectangle( (px, py), pe_w, pe_h, - facecolor="#f0f7ff", edgecolor="#3a86ff", linewidth=1.0, + facecolor=_fc, edgecolor=_ec, linewidth=1.2, ) ax.add_patch(pe_rect) - q_heads = _q_heads_for_pe(pe_id_in_group, tp, cfg.model.h_q) + # Head assignment uses TP rank (not raw pe_id) so that under + # cp_placement=pe, all CP ranks share the same head split. + q_heads = _q_heads_for_pe(_tp_rank, tp, cfg.model.h_q) kv_heads, kv_note = _kv_heads_for_pe( - pe_id_in_group, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode, + _tp_rank, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode, ) - bytes_ = _per_pe_bytes(cfg, pe_id_in_group) + bytes_ = _per_pe_bytes(cfg, _tp_rank) weights_gb = sum(v for k, v in bytes_.items() if k not in ("KV cache", "Transient")) / 1e9 kv_gb = bytes_["KV cache"] / 1e9 @@ -172,7 +211,7 @@ def draw_pe_layout(cfg: FullConfig, ax=None): ffn_mb = (bytes_["W_gate"] + bytes_["W_up"] + bytes_["W_down"]) / 1e6 # Header (PE id and heads) bigger, then per-tensor breakdown - header = (f"PE {pe_id_in_group}\n" + header = (f"PE {pe_id_in_group} TP={_tp_rank}{_cp_label}\n" f"Q heads: {q_str}\n" f"KV head: {kv_str}") ax.text(px + 0.08, py + pe_h - 0.05, header,