mukesh
5e92b89821
The Per-PE layout diagram used to color-code + label CP ranks only
when cp_placement=='pe' (CP packed intra-cube). At TP>=8 the TP dim
fills the cube, cp_placement gets pushed to 'cube', the intra-cube CP
branch is skipped, and the CP rank tag silently disappeared from
cubes, per-PE headers, and the figure title. Users at TP>=8 had no
way to tell which of the CP groups the diagram was drawing.
Fix: when cp_placement=='cube' and CP>1, add
- cube header suffix: '[CP rank 0 of {CP}]'
- per-PE header suffix: ' | CP=0'
- figure title: 'CP={CP} (showing 1 of CP groups; others identical)'
Behavior at cp_placement=='pe' is unchanged — the existing per-CP-rank
palette + 'CP={rank}' label still fire. CP=1 adds no locators at all.
test_pe_weight_layout covers four cases: TP=8/CP=4 (cube placement,
locator appears), TP=2/CP=4 forced to pe placement (regression:
per-rank labels still there, cube tag not added), CP=1 (no locators),
TP=16/CP=2 (both spilled cubes carry the CP-rank-0 tag).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
280 lines
11 KiB
Python
280 lines
11 KiB
Python
"""Draw a PE-level view of one CP group (one TP group of cubes).
|
||
|
||
Shows how weights + KV cache are distributed across the PEs of ONE CP
|
||
rank. Attention weights are sharded by head (H_q, H_kv split across TP);
|
||
FFN is sharded by dim (ffn_dim / TP / EP); KV cache is sharded by CP
|
||
(sequence) × TP (head).
|
||
|
||
If TP > 8 the group spans multiple cubes — layout draws all of them.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import matplotlib.pyplot as plt
|
||
import matplotlib.patches as patches
|
||
|
||
from .model_config import FullConfig
|
||
|
||
|
||
PE_COLS = 4
|
||
PE_ROWS = 2
|
||
PES_PER_CUBE = PE_COLS * PE_ROWS
|
||
|
||
|
||
def layers_per_stage_val(cfg: FullConfig) -> int:
|
||
return (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
|
||
|
||
|
||
def _q_heads_for_pe(pe_id_in_group: int, tp: int, h_q: int) -> list[int]:
|
||
"""Return the list of Q-head indices assigned to this PE."""
|
||
heads_per_pe = h_q / tp
|
||
start = int(round(pe_id_in_group * heads_per_pe))
|
||
end = int(round((pe_id_in_group + 1) * heads_per_pe))
|
||
return list(range(start, end))
|
||
|
||
|
||
def _kv_heads_for_pe(pe_id_in_group: int, tp: int, h_kv: int,
|
||
kv_shard_mode: str) -> tuple[list[int], str]:
|
||
"""KV heads assigned to this PE and a note about replication/split."""
|
||
if h_kv >= tp:
|
||
heads_per_pe = h_kv // tp
|
||
start = pe_id_in_group * heads_per_pe
|
||
end = start + heads_per_pe
|
||
return list(range(start, end)), ""
|
||
# TP > H_kv
|
||
if kv_shard_mode == "replicate":
|
||
rep_factor = tp // h_kv
|
||
head_idx = pe_id_in_group // rep_factor
|
||
return [head_idx], f"replicated x{rep_factor}"
|
||
# split mode
|
||
split_factor = tp // h_kv
|
||
head_idx = pe_id_in_group // split_factor
|
||
part = pe_id_in_group % split_factor
|
||
return [head_idx], f"head-split ({part+1}/{split_factor})"
|
||
|
||
|
||
def _per_pe_bytes(cfg: FullConfig, pe_id_in_group: int) -> dict:
|
||
"""Return per-tensor bytes for one PE within a TP group."""
|
||
m = cfg.model
|
||
tp = cfg.topo.tp
|
||
pp = cfg.topo.pp
|
||
ep = max(1, cfg.topo.ep)
|
||
b = m.bytes_per_elem
|
||
layers_per_stage = (m.layers + pp - 1) // pp
|
||
|
||
hq_per_pe = m.h_q / tp
|
||
if cfg.topo.kv_shard_mode == "replicate":
|
||
hkv_per_pe = max(1.0, m.h_kv / tp)
|
||
else:
|
||
hkv_per_pe = m.h_kv / tp
|
||
ffn_div = cfg.ffn_shard_divisor * ep
|
||
ffn_per_pe = m.ffn_dim // ffn_div
|
||
|
||
per_layer = {
|
||
"W_Q": int(m.hidden * hq_per_pe * m.d_head * b),
|
||
"W_K": int(m.hidden * hkv_per_pe * m.d_head * b),
|
||
"W_V": int(m.hidden * hkv_per_pe * m.d_head * b),
|
||
"W_O": int(hq_per_pe * m.d_head * m.hidden * b),
|
||
"W_gate": int(m.hidden * ffn_per_pe * b),
|
||
"W_up": int(m.hidden * ffn_per_pe * b),
|
||
"W_down": int(ffn_per_pe * m.hidden * b),
|
||
}
|
||
kv_per_layer = int(2 * cfg.topo.s_local * hkv_per_pe * m.d_head * b)
|
||
|
||
all_layers = {k: v * layers_per_stage for k, v in per_layer.items()}
|
||
all_layers["KV cache"] = kv_per_layer * layers_per_stage
|
||
# transient estimate
|
||
if cfg.topo.mode == "decode":
|
||
all_layers["Transient"] = int(4 * (m.hidden + hq_per_pe * m.d_head) * b)
|
||
else:
|
||
TILE = 1024
|
||
all_layers["Transient"] = int(2 * hq_per_pe * cfg.topo.T_q * TILE * b
|
||
+ cfg.topo.T_q * m.hidden * b)
|
||
return all_layers
|
||
|
||
|
||
def draw_pe_layout(cfg: FullConfig, ax=None):
|
||
"""Draw one CP group's PEs with per-tensor weight + KV breakdown.
|
||
|
||
Each PE shows: PE id, Q heads, KV heads, W_Q / W_K / W_V / W_O in
|
||
MB, FFN (gate+up+down) in MB, KV cache in GB, total in GB.
|
||
|
||
TP=8 -> 1 cube. TP=16 -> 2 cubes. TP=32 -> 4 cubes.
|
||
Under placement, n_cubes reflects cubes-per-group (may include CP spill).
|
||
"""
|
||
tp = cfg.topo.tp
|
||
# cubes per one cube-level group (= intra-cube dims / PEs-per-cube, rounded up).
|
||
n_cubes = max(1, (cfg.topo.intra_cube_dims + PES_PER_CUBE - 1) // PES_PER_CUBE)
|
||
|
||
if ax is None:
|
||
fig, ax = plt.subplots(figsize=(7.5 * n_cubes, 8.5))
|
||
else:
|
||
fig = ax.figure
|
||
|
||
cube_w = 7.0
|
||
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
|
||
|
||
rect = patches.FancyBboxPatch(
|
||
(x0, y0), cube_w, cube_h,
|
||
boxstyle="round,pad=0.05",
|
||
facecolor="#f8f9fa", edgecolor="#212529", linewidth=1.5,
|
||
)
|
||
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})"
|
||
+ _cube_place_tag,
|
||
ha="center", va="bottom", fontsize=11, fontweight="bold")
|
||
|
||
pe_pad_x = 0.15
|
||
pe_pad_y = 0.2
|
||
pe_gap = 0.10
|
||
inner_w = cube_w - 2 * pe_pad_x
|
||
inner_h = cube_h - 2 * pe_pad_y
|
||
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
|
||
pe_id_in_group = cube_idx * PES_PER_CUBE + pe_local
|
||
if pe_id_in_group >= cfg.topo.intra_cube_dims:
|
||
continue
|
||
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_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,
|
||
facecolor=_fc, edgecolor=_ec, linewidth=1.2,
|
||
)
|
||
ax.add_patch(pe_rect)
|
||
|
||
# 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(
|
||
_tp_rank, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
|
||
)
|
||
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
|
||
trans_mb = bytes_["Transient"] / 1e6
|
||
total_gb = sum(bytes_.values()) / 1e9
|
||
|
||
q_str = (f"{q_heads[0]}-{q_heads[-1]}"
|
||
if len(q_heads) > 1 else str(q_heads[0])
|
||
if q_heads else "-")
|
||
kv_str = str(kv_heads[0]) if kv_heads else "-"
|
||
if kv_note:
|
||
kv_str += f"({kv_note})"
|
||
|
||
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} 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,
|
||
ha="left", va="top",
|
||
fontsize=7.0, fontweight="bold",
|
||
family="monospace", color="#0d47a1")
|
||
|
||
# Per-tensor bytes with explicit division divisors.
|
||
# Divisor is always the PRODUCT of the parallel dims listed
|
||
# (TP*CP means TP times CP, i.e. multiplication).
|
||
tp_d = cfg.topo.tp
|
||
ffn_div = cfg.ffn_shard_divisor
|
||
ffn_scope = cfg.topo.ffn_shard_scope.replace("+", "*")
|
||
kv_div = cfg.topo.cp * cfg.topo.tp
|
||
L = layers_per_stage_val(cfg)
|
||
detail = (
|
||
f"weights (all {L} layers):\n"
|
||
f" W_Q /TP={tp_d} : {bytes_['W_Q']/1e6:7.1f}MB\n"
|
||
f" W_K /TP={tp_d} : {bytes_['W_K']/1e6:7.1f}MB\n"
|
||
f" W_V /TP={tp_d} : {bytes_['W_V']/1e6:7.1f}MB\n"
|
||
f" W_O /TP={tp_d} : {bytes_['W_O']/1e6:7.1f}MB\n"
|
||
f" FFN /({ffn_scope})={ffn_div}: {ffn_mb:7.1f}MB\n"
|
||
f"runtime:\n"
|
||
f" KV /(CP*TP)={kv_div}: {kv_gb*1000:7.1f}MB\n"
|
||
f" trans : {trans_mb:7.1f}MB\n"
|
||
f"W total: {weights_gb*1000:7.1f}MB\n"
|
||
f"TOTAL : {total_gb*1000:7.1f}MB"
|
||
)
|
||
ax.text(px + pe_w / 2, py + 0.06, detail,
|
||
ha="center", va="bottom",
|
||
fontsize=5.8, family="monospace", color="#1a1a1a")
|
||
|
||
total_w = n_cubes * cube_w + (n_cubes - 1) * cube_gap
|
||
ax.set_xlim(-0.2, total_w + 0.2)
|
||
ax.set_ylim(-0.3, cube_h + 0.6)
|
||
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}{_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")
|
||
|
||
return fig
|