mukesh
fc1dcdde24
Three additions to the tensor sharding view when batch B > 1:
1. Shape label appends "B={B}": e.g. (128, S_kv=131072, B=8).
2. Note under the KV cache: "batch B={B} => {B}x KV bytes per PE".
3. Visual "stack" — up to 5 dashed offset rectangles drawn behind
the KV cache to hint at the batch stacking dimension. Capped at 5
so a big B doesn't overwhelm the diagram.
4. Title also gets B={B} between CP and FFN scope.
Attention/FFN weight tensors are NOT stacked — weights are shared
across the batch (correct: only activations + KV scale with B).
At B=1, all four additions are no-ops so the diagram looks unchanged
from before.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
361 lines
14 KiB
Python
361 lines
14 KiB
Python
"""Pictorial view of how each weight/KV tensor is sharded across ranks.
|
|
|
|
Every tensor is drawn as a rectangle labelled with its (rows, cols)
|
|
shape. Grid lines show shard boundaries; one shard (this PE's share)
|
|
is highlighted so it's obvious which dim the split is along.
|
|
|
|
Sharding pattern:
|
|
- W_Q (hidden, H_q*d_head) : split on **output cols** across TP
|
|
- W_K (hidden, H_kv*d_head): split on **output cols** across TP
|
|
- W_V (hidden, H_kv*d_head): split on **output cols** across TP
|
|
- W_O (H_q*d_head, hidden) : split on **input rows** across TP
|
|
- W_gate/up (hidden, ffn_dim): split on **output cols** across FFN scope
|
|
- W_down (ffn_dim, hidden) : split on **input rows** across FFN scope
|
|
- K/V cache (S_kv, H_kv*d_head): split on **rows (S axis)** by CP AND
|
|
**cols (head axis)** by TP -> 2D shard grid
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.patches as patches
|
|
|
|
from .model_config import FullConfig
|
|
|
|
|
|
_UNSHARDED_FACE = "#e9ecef"
|
|
_UNSHARDED_EDGE = "#adb5bd"
|
|
_MY_SHARD_FACE = "#3a86ff"
|
|
_MY_SHARD_EDGE = "#0057b3"
|
|
_TP_LINE_COLOR = "#d90429" # TP split lines
|
|
_CP_LINE_COLOR = "#0f9d58" # CP split lines
|
|
_FFN_LINE_COLOR = "#e37400" # FFN-scope split lines
|
|
|
|
|
|
def _draw_tensor(ax, x, y, w, h, name, shape_str,
|
|
row_splits: int = 1, col_splits: int = 1,
|
|
my_row_shard: int = 0, my_col_shard: int = 0,
|
|
row_line_color: str = _TP_LINE_COLOR,
|
|
col_line_color: str = _TP_LINE_COLOR,
|
|
row_label: str = "", col_label: str = "",
|
|
note: str = "",
|
|
cell_annot=None):
|
|
"""Draw a tensor with sharding grid.
|
|
|
|
cell_annot: optional callable (row, col) -> str returning the physical
|
|
label to render inside shard cell (r, c). If it returns "", cell is not
|
|
annotated.
|
|
"""
|
|
"""Draw a tensor as a rectangle with grid lines showing shards."""
|
|
# Base rectangle (unsharded background)
|
|
rect = patches.Rectangle(
|
|
(x, y), w, h,
|
|
facecolor=_UNSHARDED_FACE, edgecolor=_UNSHARDED_EDGE, linewidth=1.0,
|
|
)
|
|
ax.add_patch(rect)
|
|
|
|
# Highlight this PE's shard
|
|
shard_w = w / max(1, col_splits)
|
|
shard_h = h / max(1, row_splits)
|
|
sx = x + my_col_shard * shard_w
|
|
sy = y + (row_splits - 1 - my_row_shard) * shard_h
|
|
highlight = patches.Rectangle(
|
|
(sx, sy), shard_w, shard_h,
|
|
facecolor=_MY_SHARD_FACE, edgecolor=_MY_SHARD_EDGE, linewidth=1.5,
|
|
alpha=0.8,
|
|
)
|
|
ax.add_patch(highlight)
|
|
|
|
# Horizontal grid lines (row splits)
|
|
for r in range(1, row_splits):
|
|
yy = y + r * shard_h
|
|
ax.plot([x, x + w], [yy, yy],
|
|
color=row_line_color, linewidth=1.2, linestyle="--")
|
|
# Vertical grid lines (col splits)
|
|
for c in range(1, col_splits):
|
|
xx = x + c * shard_w
|
|
ax.plot([xx, xx], [y, y + h],
|
|
color=col_line_color, linewidth=1.2, linestyle="--")
|
|
|
|
# Per-cell physical annotations (e.g. "cube3 PE5") if provided.
|
|
if cell_annot is not None:
|
|
for r in range(row_splits):
|
|
for c in range(col_splits):
|
|
text = cell_annot(r, c)
|
|
if not text:
|
|
continue
|
|
cx = x + c * shard_w + shard_w / 2
|
|
cy = y + (row_splits - 1 - r) * shard_h + shard_h / 2
|
|
ax.text(cx, cy, text,
|
|
ha="center", va="center",
|
|
fontsize=6.0, color="#0d1b2a",
|
|
alpha=0.85)
|
|
|
|
# Stacked labels ABOVE rectangle: note (italic, top), shape, name (bold, closest)
|
|
if note:
|
|
ax.text(x + w / 2, y + h + 0.55, note,
|
|
ha="center", va="bottom", fontsize=6.5,
|
|
color="#555", style="italic")
|
|
ax.text(x + w / 2, y + h + 0.30, shape_str,
|
|
ha="center", va="bottom", fontsize=7, color="#666")
|
|
ax.text(x + w / 2, y + h + 0.08, name,
|
|
ha="center", va="bottom",
|
|
fontsize=10, fontweight="bold")
|
|
|
|
# Split labels BELOW rectangle (stacked if both present)
|
|
if row_label and col_label:
|
|
ax.text(x + w / 2, y - 0.10, row_label,
|
|
ha="center", va="top", fontsize=6.5,
|
|
color=row_line_color, fontweight="bold")
|
|
ax.text(x + w / 2, y - 0.35, col_label,
|
|
ha="center", va="top", fontsize=6.5,
|
|
color=col_line_color, fontweight="bold")
|
|
elif row_label:
|
|
ax.text(x + w / 2, y - 0.10, row_label,
|
|
ha="center", va="top", fontsize=6.5,
|
|
color=row_line_color, fontweight="bold")
|
|
elif col_label:
|
|
ax.text(x + w / 2, y - 0.10, col_label,
|
|
ha="center", va="top", fontsize=6.5,
|
|
color=col_line_color, fontweight="bold")
|
|
|
|
|
|
def _physical_label(cfg: FullConfig, dim: str, rank: int) -> str:
|
|
"""Where does rank `rank` of dim (`tp`|`cp`|`ffn`) live physically?
|
|
Returns a short "cubeX PEy" style label based on the placement config.
|
|
"""
|
|
topo = cfg.topo
|
|
if dim == "tp":
|
|
if topo.tp_placement == "pe":
|
|
# TP fills PEs within one cube (spilling if tp > 8)
|
|
pe_in_cube = rank % topo.pes_per_cube_hw
|
|
cube = rank // topo.pes_per_cube_hw
|
|
return f"c{cube}p{pe_in_cube}"
|
|
else: # tp on cube
|
|
return f"cube{rank}"
|
|
if dim == "cp":
|
|
if topo.cp_placement == "pe":
|
|
pe_in_cube = rank % topo.pes_per_cube_hw
|
|
cube = rank // topo.pes_per_cube_hw
|
|
return f"c{cube}p{pe_in_cube}"
|
|
else: # cp on cube
|
|
return f"cube{rank}"
|
|
return ""
|
|
|
|
|
|
def _kv_cell_annot(cfg: FullConfig):
|
|
"""Return an annotation function for KV cache cells: (cp_r, tp_r) -> label."""
|
|
def _ann(r, c):
|
|
cp_lbl = _physical_label(cfg, "cp", r)
|
|
tp_lbl = _physical_label(cfg, "tp", c)
|
|
# Two lines if both are non-trivial; otherwise a single line
|
|
if cp_lbl and tp_lbl:
|
|
return f"{cp_lbl}\n{tp_lbl}"
|
|
return cp_lbl or tp_lbl
|
|
return _ann
|
|
|
|
|
|
def _tp_cell_annot(cfg: FullConfig):
|
|
"""Col-parallel tensors (W_Q/W_K/W_V/W_gate/W_up): each col = one TP rank."""
|
|
def _ann(r, c):
|
|
return _physical_label(cfg, "tp", c)
|
|
return _ann
|
|
|
|
|
|
def _tp_row_cell_annot(cfg: FullConfig):
|
|
"""Row-parallel tensors split by TP (W_O): each row = one TP rank."""
|
|
def _ann(r, c):
|
|
return _physical_label(cfg, "tp", r)
|
|
return _ann
|
|
|
|
|
|
def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
|
|
show_physical: bool = False):
|
|
"""Draw all tensors with their sharding overlays for PE `my_pe`.
|
|
|
|
show_physical: annotate each shard cell with the owning cube/PE.
|
|
"""
|
|
m = cfg.model
|
|
tp = cfg.topo.tp
|
|
cp = cfg.topo.cp
|
|
|
|
# PE's assigned shard index for each dim
|
|
q_shard = my_pe % tp # W_Q col shard
|
|
kv_shard_col = min(my_pe % tp, m.h_kv - 1) if m.h_kv >= tp else 0
|
|
ffn_shard = my_pe % cfg.ffn_shard_divisor
|
|
kv_row_shard = 0 # per-CP; showing CP=0 view for clarity
|
|
|
|
# Enlarged mode gives more room for physical annotations.
|
|
_fig_w, _fig_h = (22, 13) if show_physical else (15, 9)
|
|
_cell_w = 3.4 if show_physical else 2.6
|
|
_cell_h = 3.0 if show_physical else 2.2
|
|
_gap_x = 0.8 if show_physical else 0.55
|
|
_gap_y = 2.0 if show_physical else 1.7
|
|
|
|
if ax is None:
|
|
fig, ax = plt.subplots(figsize=(_fig_w, _fig_h))
|
|
else:
|
|
fig = ax.figure
|
|
|
|
# Grid layout: 2 rows x 4 cols of tensor rects
|
|
# Attention row: W_Q, W_K, W_V, W_O
|
|
# FFN + KV row: W_gate, W_up, W_down, KV cache
|
|
cell_w = _cell_w
|
|
cell_h = _cell_h
|
|
gap_x = _gap_x
|
|
gap_y = _gap_y
|
|
|
|
_kv_ann = _kv_cell_annot(cfg) if show_physical else None
|
|
_tp_col_ann = _tp_cell_annot(cfg) if show_physical else None
|
|
_tp_row_ann = _tp_row_cell_annot(cfg) if show_physical else None
|
|
|
|
positions = [
|
|
(0, 1), # W_Q at row 0
|
|
(1, 1), # W_K
|
|
(2, 1), # W_V
|
|
(3, 1), # W_O
|
|
(0, 0), # W_gate at row 1
|
|
(1, 0), # W_up
|
|
(2, 0), # W_down
|
|
(3, 0), # KV cache
|
|
]
|
|
|
|
# ── attention tensors ────────────────────────
|
|
tp_txt = f"TP={tp}"
|
|
ffn_txt = f"{cfg.topo.ffn_shard_scope}={cfg.ffn_shard_divisor}"
|
|
cp_txt = f"CP={cp}"
|
|
|
|
def px(col): return col * (cell_w + gap_x)
|
|
def py(row): return row * (cell_h + gap_y)
|
|
|
|
# W_Q: (hidden, H_q*d_head), split on COL by TP
|
|
_draw_tensor(ax, px(0), py(1), cell_w, cell_h,
|
|
"W_Q", f"({m.hidden}, {m.h_q * m.d_head})",
|
|
row_splits=1, col_splits=tp,
|
|
my_col_shard=q_shard,
|
|
col_line_color=_TP_LINE_COLOR,
|
|
col_label=f"cols split by {tp_txt}",
|
|
note="col-parallel (output dim)",
|
|
cell_annot=_tp_col_ann)
|
|
|
|
# W_K, W_V: (hidden, H_kv*d_head)
|
|
for (name, col_idx) in [("W_K", 1), ("W_V", 2)]:
|
|
splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
|
|
_draw_tensor(ax, px(col_idx), py(1), cell_w, cell_h,
|
|
name, f"({m.hidden}, {m.h_kv * m.d_head})",
|
|
row_splits=1, col_splits=splits,
|
|
my_col_shard=min(q_shard, splits - 1),
|
|
col_line_color=_TP_LINE_COLOR,
|
|
col_label=f"cols split by {tp_txt}",
|
|
note="col-parallel (heads)",
|
|
cell_annot=_tp_col_ann)
|
|
|
|
# W_O: (H_q*d_head, hidden), split on ROW by TP
|
|
_draw_tensor(ax, px(3), py(1), cell_w, cell_h,
|
|
"W_O", f"({m.h_q * m.d_head}, {m.hidden})",
|
|
row_splits=tp, col_splits=1,
|
|
my_row_shard=q_shard,
|
|
row_line_color=_TP_LINE_COLOR,
|
|
row_label=f"rows split by {tp_txt}",
|
|
note="row-parallel (input dim)",
|
|
cell_annot=_tp_row_ann)
|
|
|
|
# ── FFN tensors ──────────────────────────────
|
|
# W_gate, W_up: (hidden, ffn_dim), split on COL by FFN scope
|
|
for (name, col_idx) in [("W_gate", 0), ("W_up", 1)]:
|
|
_draw_tensor(ax, px(col_idx), py(0), cell_w, cell_h,
|
|
name, f"({m.hidden}, {m.ffn_dim})",
|
|
row_splits=1, col_splits=cfg.ffn_shard_divisor,
|
|
my_col_shard=ffn_shard,
|
|
col_line_color=_FFN_LINE_COLOR,
|
|
col_label=f"cols split by {ffn_txt}",
|
|
note="col-parallel")
|
|
|
|
# W_down: (ffn_dim, hidden), split on ROW by FFN scope
|
|
_draw_tensor(ax, px(2), py(0), cell_w, cell_h,
|
|
"W_down", f"({m.ffn_dim}, {m.hidden})",
|
|
row_splits=cfg.ffn_shard_divisor, col_splits=1,
|
|
my_row_shard=ffn_shard,
|
|
row_line_color=_FFN_LINE_COLOR,
|
|
row_label=f"rows split by {ffn_txt}",
|
|
note="row-parallel")
|
|
|
|
# ── KV cache: (B, S_kv, H_kv*d_head), 2D shard on S_kv x heads ─
|
|
# Batch (B) is a stacking dimension — each concurrent request holds
|
|
# its own KV slice per PE, so total KV per PE = B * (per-slice size).
|
|
kv_row_splits = cp
|
|
kv_col_splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
|
|
_B = max(1, cfg.topo.b)
|
|
_batch_suffix = f", B={_B}" if _B > 1 else ""
|
|
_batch_note = (f"; batch B={_B} => {_B}x KV bytes per PE"
|
|
if _B > 1 else "")
|
|
# Draw B-1 offset "shadow" rectangles behind the KV cache to give
|
|
# a visual hint of the batch stacking dimension. Cap at 5 shadows
|
|
# so a large B doesn't overwhelm the diagram.
|
|
if _B > 1:
|
|
_kv_shadow_x = px(3)
|
|
_kv_shadow_y = py(0)
|
|
_n_shadows = min(_B - 1, 5)
|
|
_offset_step = 0.15
|
|
for _si in range(_n_shadows, 0, -1):
|
|
_dx = _si * _offset_step
|
|
_dy = _si * _offset_step
|
|
_shadow = patches.Rectangle(
|
|
(_kv_shadow_x + _dx, _kv_shadow_y + _dy),
|
|
cell_w, cell_h,
|
|
facecolor="#f5f5f5",
|
|
edgecolor="#adb5bd",
|
|
linewidth=0.8, linestyle="--", alpha=0.55,
|
|
)
|
|
ax.add_patch(_shadow)
|
|
_draw_tensor(ax, px(3), py(0), cell_w, cell_h,
|
|
"KV cache (K and V)",
|
|
f"({m.h_kv * m.d_head}, S_kv={cfg.topo.s_kv:,}"
|
|
f"{_batch_suffix})",
|
|
row_splits=kv_row_splits,
|
|
col_splits=kv_col_splits if show_physical else 1,
|
|
my_row_shard=kv_row_shard,
|
|
row_line_color=_CP_LINE_COLOR,
|
|
row_label=f"rows: S split by {cp_txt}",
|
|
col_label=f"cols: heads split by {tp_txt}",
|
|
col_line_color=_TP_LINE_COLOR,
|
|
note=f"2D shard: seq axis (CP) x head axis (TP)"
|
|
f"{_batch_note}",
|
|
cell_annot=_kv_ann)
|
|
# In non-physical mode, KV was drawn with col_splits=1; overlay TP col
|
|
# splits as dotted lines to hint at the 2D nature. In physical mode
|
|
# the _draw_tensor call already drew proper vertical dashes.
|
|
if not show_physical:
|
|
kv_x = px(3)
|
|
kv_y = py(0)
|
|
shard_col_w = cell_w / max(1, kv_col_splits)
|
|
for c in range(1, kv_col_splits):
|
|
xx = kv_x + c * shard_col_w
|
|
ax.plot([xx, xx], [kv_y, kv_y + cell_h],
|
|
color=_TP_LINE_COLOR, linewidth=1.2, linestyle=":")
|
|
|
|
ax.set_xlim(-0.3, 4 * (cell_w + gap_x) - gap_x + 0.3)
|
|
ax.set_ylim(-0.9, 2 * (cell_h + gap_y) - gap_y + 0.9)
|
|
ax.set_aspect("auto")
|
|
ax.axis("off")
|
|
|
|
_B_title = max(1, cfg.topo.b)
|
|
_b_title_suffix = f", B={_B_title}" if _B_title > 1 else ""
|
|
ax.set_title(
|
|
f"Tensor sharding view (PE {my_pe}) | TP={tp}, CP={cp}"
|
|
f"{_b_title_suffix}, "
|
|
f"FFN scope={cfg.topo.ffn_shard_scope} | "
|
|
f"blue = this PE's shard",
|
|
fontsize=11, fontweight="bold",
|
|
)
|
|
# Legend for line colors
|
|
ax.plot([], [], color=_TP_LINE_COLOR, linewidth=1.2, linestyle="--",
|
|
label=f"TP split ({tp} ranks)")
|
|
ax.plot([], [], color=_CP_LINE_COLOR, linewidth=1.2, linestyle="--",
|
|
label=f"CP split ({cp} ranks, sequence dim)")
|
|
ax.plot([], [], color=_FFN_LINE_COLOR, linewidth=1.2, linestyle="--",
|
|
label=f"FFN scope split ({cfg.ffn_shard_divisor} ranks)")
|
|
ax.legend(loc="lower center", bbox_to_anchor=(0.5, -0.05),
|
|
ncol=3, fontsize=8, frameon=False)
|
|
|
|
return fig
|