mukesh
4ded67a443
Replace the three memory-optimal Pareto-search buttons in the Parallelism sidebar with three scope-aware memory-min autosuggest buttons. Each button sizes for a different per-PE memory footprint: - Attn -> attention weights + KV cache (no FFN weights) - FFN/MoE -> FFN weights only (no attention, no KV) - Attn+FFN -> everything (traditional autosuggest) memory_layout: per_pe_weight_bytes and compute_memory take include_attention / include_ffn flags; KV cache is zeroed when attention scope is excluded. autosuggest: _score_candidate and auto_suggest forward the flags into compute_memory, so the smallest-fit search now respects the chosen scope. app.py: single "Apply memory-min" caption + 3 column buttons. Each button runs auto_suggest with its scope filter, snaps the sliders to the picked (CP, TP, PP), and sets the Physical Layout scope filter to match. Removes the standalone Apply memory-min button (was duplicating the Attn+FFN case) and the Pareto-search import. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
249 lines
9.2 KiB
Python
249 lines
9.2 KiB
Python
"""Per-PE memory footprint: weights + KV cache + activations + slack.
|
||
|
||
All formulas are per-PE, per-PP-stage. If PP > 1, each PE holds only
|
||
its stage's layers (layers/PP).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
|
||
from .model_config import FullConfig
|
||
|
||
|
||
@dataclass
|
||
class MemoryBreakdown:
|
||
weights_bytes: int
|
||
kv_cache_bytes: int
|
||
transient_bytes: int
|
||
budget_bytes: int
|
||
|
||
@property
|
||
def used_bytes(self) -> int:
|
||
return self.weights_bytes + self.kv_cache_bytes + self.transient_bytes
|
||
|
||
@property
|
||
def slack_bytes(self) -> int:
|
||
return max(0, self.budget_bytes - self.used_bytes)
|
||
|
||
@property
|
||
def over_budget(self) -> bool:
|
||
return self.used_bytes > self.budget_bytes
|
||
|
||
|
||
def per_pe_weight_bytes(cfg: FullConfig,
|
||
include_attention: bool = True,
|
||
include_ffn: bool = True) -> int:
|
||
"""Attention + FFN weights per PE, bf16. Divided by TP, PP, and EP.
|
||
|
||
EP divides the FFN (experts) across ranks; attention weights are
|
||
unaffected. When kv_shard_mode='replicate' and TP > H_kv, each KV
|
||
head is replicated (per-PE W_K/W_V size doesn't drop below one head).
|
||
|
||
Scope flags let callers count only attention or only FFN weights —
|
||
useful for "smallest deployment for just this block" sizing. Both
|
||
True (default) matches the traditional full-model per-PE weight
|
||
footprint.
|
||
"""
|
||
m = cfg.model
|
||
tp = cfg.topo.tp
|
||
pp = cfg.topo.pp
|
||
ep = max(1, cfg.topo.ep)
|
||
|
||
hq_per_pe = cfg.h_q_per_pe
|
||
if cfg.topo.kv_shard_mode == "replicate":
|
||
# Each KV head held fully; replicated across ranks if TP > H_kv.
|
||
hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
|
||
else: # "split"
|
||
# Head-dim split: fractional head allowed (< 1 head bytes when TP > H_kv).
|
||
hkv_per_pe_bytes = m.h_kv / tp
|
||
|
||
per_layer_attn = 0
|
||
if include_attention:
|
||
per_layer_attn = (
|
||
m.hidden * hq_per_pe * m.d_head + # W_Q
|
||
m.hidden * hkv_per_pe_bytes * m.d_head + # W_K
|
||
m.hidden * hkv_per_pe_bytes * m.d_head + # W_V
|
||
hq_per_pe * m.d_head * m.hidden # W_O
|
||
)
|
||
|
||
per_layer_ffn = 0
|
||
if include_ffn:
|
||
# FFN divisor: TP (default) or TP*CP or TP*CP*DP if user opts in.
|
||
# EP further divides (MoE experts).
|
||
ffn_div = cfg.ffn_shard_divisor * ep
|
||
per_layer_ffn = 3 * m.hidden * (m.ffn_dim // ffn_div)
|
||
|
||
per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
|
||
layers_per_stage = (m.layers + pp - 1) // pp
|
||
return int(per_layer * layers_per_stage)
|
||
|
||
|
||
def per_pe_kv_cache_bytes(cfg: FullConfig) -> int:
|
||
"""K + V per PE, across all layers this stage holds.
|
||
|
||
- CP shards the sequence dim -> S_local = S_kv/CP tokens per PE.
|
||
- TP splits KV heads across ranks. If kv_shard_mode='replicate' and
|
||
TP > H_kv, each KV head is duplicated across TP/H_kv ranks so
|
||
per-PE storage doesn't shrink below 1 head.
|
||
- PP: only layers_per_stage layers stored per PE.
|
||
- Batch (B): each concurrent request keeps its own KV cache slice.
|
||
"""
|
||
m = cfg.model
|
||
pp = cfg.topo.pp
|
||
tp = cfg.topo.tp
|
||
B = max(1, cfg.topo.b)
|
||
if cfg.topo.kv_shard_mode == "replicate":
|
||
hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
|
||
else:
|
||
hkv_per_pe_bytes = m.h_kv / tp
|
||
layers_per_stage = (m.layers + pp - 1) // pp
|
||
per_layer = 2 * cfg.topo.s_local * hkv_per_pe_bytes * m.d_head * m.bytes_per_elem
|
||
return int(per_layer * layers_per_stage * B)
|
||
|
||
|
||
def per_pe_transient_bytes(cfg: FullConfig) -> int:
|
||
"""Rough peak transient (activations, GEMM outputs) per PE."""
|
||
m = cfg.model
|
||
T_q = cfg.topo.T_q
|
||
hq_per_pe = cfg.h_q_per_pe
|
||
|
||
if cfg.topo.mode == "decode":
|
||
return 4 * (m.hidden + m.head_dim_total_q // cfg.topo.tp) * m.bytes_per_elem
|
||
else:
|
||
TILE = 1024
|
||
tile_score = hq_per_pe * T_q * TILE * m.bytes_per_elem
|
||
return int(2 * tile_score + T_q * m.hidden * m.bytes_per_elem)
|
||
|
||
|
||
def compute_memory(cfg: FullConfig,
|
||
include_attention: bool = True,
|
||
include_ffn: bool = True) -> MemoryBreakdown:
|
||
"""Per-PE memory breakdown. Scope flags let callers count only the
|
||
attention or only the FFN block (KV cache goes with the attention
|
||
block; transient activation buffer is small either way and left in).
|
||
"""
|
||
kv_bytes = per_pe_kv_cache_bytes(cfg) if include_attention else 0
|
||
return MemoryBreakdown(
|
||
weights_bytes=per_pe_weight_bytes(
|
||
cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||
),
|
||
kv_cache_bytes=kv_bytes,
|
||
transient_bytes=per_pe_transient_bytes(cfg),
|
||
budget_bytes=cfg.machine.pe_budget_bytes,
|
||
)
|
||
|
||
|
||
def _one_layer_row(name: str,
|
||
global_shape: tuple[int, int],
|
||
per_pe_shape: tuple[int, int],
|
||
bytes_per_elem: int) -> dict:
|
||
"""Per-tensor row for ONE layer (both global and per-PE shard)."""
|
||
p_global = global_shape[0] * global_shape[1]
|
||
p_per_pe = per_pe_shape[0] * per_pe_shape[1]
|
||
return {
|
||
"Tensor": name,
|
||
"Global shape": f"({global_shape[0]}, {global_shape[1]})",
|
||
"Params/layer": f"{p_global/1e6:.2f} M",
|
||
"Bytes/layer (global)": f"{p_global * bytes_per_elem / 1e6:.2f} MB",
|
||
"Per-PE shape": f"({per_pe_shape[0]}, {per_pe_shape[1]})",
|
||
"Bytes/layer (per PE)": f"{p_per_pe * bytes_per_elem / 1e6:.2f} MB",
|
||
"_p_global": p_global,
|
||
"_p_per_pe": p_per_pe,
|
||
}
|
||
|
||
|
||
def attention_weight_rows(cfg: FullConfig) -> list[dict]:
|
||
"""Per-tensor rows for attention weights (one layer)."""
|
||
m = cfg.model
|
||
hq_per_pe = cfg.h_q_per_pe
|
||
hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
|
||
return [
|
||
_one_layer_row("W_Q", (m.hidden, m.h_q * m.d_head),
|
||
(m.hidden, hq_per_pe * m.d_head),
|
||
m.bytes_per_elem),
|
||
_one_layer_row("W_K", (m.hidden, m.h_kv * m.d_head),
|
||
(m.hidden, hkv_per_pe * m.d_head),
|
||
m.bytes_per_elem),
|
||
_one_layer_row("W_V", (m.hidden, m.h_kv * m.d_head),
|
||
(m.hidden, hkv_per_pe * m.d_head),
|
||
m.bytes_per_elem),
|
||
_one_layer_row("W_O", (m.h_q * m.d_head, m.hidden),
|
||
(hq_per_pe * m.d_head, m.hidden),
|
||
m.bytes_per_elem),
|
||
]
|
||
|
||
|
||
def ffn_weight_rows(cfg: FullConfig) -> list[dict]:
|
||
"""Per-tensor rows for FFN weights (one layer, activated for MoE)."""
|
||
m = cfg.model
|
||
ffn_per_pe = m.ffn_dim // cfg.topo.tp
|
||
return [
|
||
_one_layer_row("W_gate", (m.hidden, m.ffn_dim),
|
||
(m.hidden, ffn_per_pe), m.bytes_per_elem),
|
||
_one_layer_row("W_up", (m.hidden, m.ffn_dim),
|
||
(m.hidden, ffn_per_pe), m.bytes_per_elem),
|
||
_one_layer_row("W_down", (m.ffn_dim, m.hidden),
|
||
(ffn_per_pe, m.hidden), m.bytes_per_elem),
|
||
]
|
||
|
||
|
||
def kv_cache_rows(cfg: FullConfig) -> list[dict]:
|
||
"""Per-tensor row for K and V cache (one layer)."""
|
||
m = cfg.model
|
||
hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
|
||
b = m.bytes_per_elem
|
||
global_k = cfg.topo.s_kv * m.h_kv * m.d_head
|
||
per_pe_k = cfg.topo.s_local * hkv_per_pe * m.d_head
|
||
return [
|
||
{
|
||
"Tensor": "K cache",
|
||
"Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
|
||
"Params/layer": f"{global_k/1e6:.2f} M",
|
||
"Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
|
||
"Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
|
||
"Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
|
||
"_p_global": global_k,
|
||
"_p_per_pe": per_pe_k,
|
||
},
|
||
{
|
||
"Tensor": "V cache",
|
||
"Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
|
||
"Params/layer": f"{global_k/1e6:.2f} M",
|
||
"Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
|
||
"Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
|
||
"Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
|
||
"_p_global": global_k,
|
||
"_p_per_pe": per_pe_k,
|
||
},
|
||
]
|
||
|
||
|
||
def sum_bytes_all_layers(rows: list[dict], cfg: FullConfig,
|
||
per_pe: bool = False) -> int:
|
||
"""Sum bytes across all rows × N layers (or layers/PP for per_pe)."""
|
||
b = cfg.model.bytes_per_elem
|
||
layers = cfg.model.layers
|
||
if per_pe:
|
||
layers = (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
|
||
key = "_p_per_pe" if per_pe else "_p_global"
|
||
return sum(r[key] for r in rows) * b * layers
|
||
|
||
|
||
def total_weight_bytes_full_model(cfg: FullConfig) -> int:
|
||
"""Total bytes of ALL weights (attention + FFN) for the full unsharded model."""
|
||
m = cfg.model
|
||
attn_per_layer = (
|
||
m.hidden * m.h_q * m.d_head + # W_Q
|
||
m.hidden * m.h_kv * m.d_head * 2 + # W_K, W_V
|
||
m.h_q * m.d_head * m.hidden # W_O
|
||
)
|
||
ffn_per_layer = 3 * m.hidden * m.ffn_dim
|
||
return (attn_per_layer + ffn_per_layer) * m.bytes_per_elem * m.layers
|
||
|
||
|
||
def total_kv_bytes_full_model(cfg: FullConfig) -> int:
|
||
"""Total bytes of KV cache for full unsharded model at cfg.topo.s_kv."""
|
||
m = cfg.model
|
||
per_layer = 2 * cfg.topo.s_kv * m.h_kv * m.d_head * m.bytes_per_elem
|
||
return per_layer * m.layers
|