diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index d262039..3d6877a 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -218,8 +218,8 @@ with st.sidebar.expander("Parallelism", expanded=True): st.session_state["cp_placement"] = _default_suggestion.cp_placement st.session_state["_pl_active_scope"] = "full" st.rerun() - # Three latency-optimal apply buttons — pick a scope, load the - # latency-min Pareto config into the sidebar sliders. + # Three memory-optimal apply buttons — pick a scope, load the + # smallest-fit Pareto config into the sidebar sliders. from tests.analytical_visualization.auto_explore import ( run_auto_explore as _run_auto_explore_sb, ) @@ -235,14 +235,14 @@ with st.sidebar.expander("Parallelism", expanded=True): _sb_clicks[_sb_scope] = st.button( _sb_scope_meta[_sb_scope][2], width='stretch', key=f"_sb_apply_{_sb_scope}", - help=f"Apply latency-optimal parallelism for " + help=f"Apply memory-optimal (smallest-fit) parallelism for " f"{_sb_scope_meta[_sb_scope][2]} scope.", ) for _sb_scope, _sb_click in _sb_clicks.items(): if not _sb_click: continue _sb_ia, _sb_ff, _sb_lbl = _sb_scope_meta[_sb_scope] - with st.spinner(f"Finding latency-optimal parallelism ({_sb_lbl})..."): + with st.spinner(f"Finding memory-optimal parallelism ({_sb_lbl})..."): _sb_res = _run_auto_explore_sb( model, _default_machine, s_kv=s_kv, mode=mode, include_attention=_sb_ia, include_ffn=_sb_ff, b=b_batch, @@ -553,17 +553,18 @@ tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([ # ── TAB 1: Physical layout ────────────────────────────────────── with tab_layout: - # Quick shortcut: three buttons that pick the latency-optimal Pareto - # config for a chosen scope and load it into the sidebar. Layout - # re-renders immediately on the next Streamlit run with the new config. + # Quick shortcut: three buttons that pick the memory-optimal + # (smallest-fit) Pareto config for a chosen scope and load it into + # the sidebar. Layout re-renders immediately on the next Streamlit + # run with the new config. from tests.analytical_visualization.auto_explore import ( run_auto_explore as _run_auto_explore_layout, ) st.markdown( - "**Quick auto-suggest layout** — pick a scope; the latency-optimal " - "Pareto config for that scope is loaded into the sidebar and this " - "layout redraws around it." + "**Quick auto-suggest layout** — pick a scope; the memory-optimal " + "(smallest-fit) Pareto config for that scope is loaded into the " + "sidebar and this layout redraws around it." ) _pl1, _pl2, _pl3, _pl4 = st.columns([1, 1, 1, 1]) with _pl1: @@ -591,7 +592,7 @@ with tab_layout: if not _pl_clicked: continue _pl_ia, _pl_ff, _pl_lbl = _pl_scope_meta[_pl_scope] - with st.spinner(f"Finding latency-optimal parallelism ({_pl_lbl})..."): + with st.spinner(f"Finding memory-optimal parallelism ({_pl_lbl})..."): _pl_res = _run_auto_explore_layout( model, machine, s_kv=s_kv, mode=mode, include_attention=_pl_ia, include_ffn=_pl_ff, b=b_batch, @@ -629,7 +630,7 @@ with tab_layout: # stage tables (attn/ffn/full). st.session_state["_pl_active_scope"] = _pl_scope st.success( - f"Loaded latency-optimal {_pl_lbl} config: CP={_pl_best.cp} " + f"Loaded memory-optimal {_pl_lbl} config: CP={_pl_best.cp} " f"TP={_pl_best.tp} PP={_pl_best.pp} DP={_pl_best.dp} → " f"{_pl_best.latency_ms:.2f} ms, {_pl_best.pes_used} PEs. " "Redrawing layout..." diff --git a/tests/analytical_visualization/stage_latencies.py b/tests/analytical_visualization/stage_latencies.py index 7c06f3c..7b11fe8 100644 --- a/tests/analytical_visualization/stage_latencies.py +++ b/tests/analytical_visualization/stage_latencies.py @@ -59,12 +59,13 @@ def stage_rmsnorm(cfg: FullConfig) -> StageCost: vis, bnd = _visible(cmp_s, mem_s, 0) return StageCost( name="S1 RMSNorm", - formula=f"bytes = {T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM", + formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM", compute_s=cmp_s, memory_s=mem_s, comm_s=0.0, bound=bnd, visible_s=vis, flops=flops, mem_bytes=bytes_, - flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}", - mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B", + flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}", + mem_formula=f"B*T_q*d*b + T_q*d*b (weight) " + f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B", ) @@ -88,13 +89,16 @@ def stage_wq(cfg: FullConfig) -> StageCost: vis, bnd = _visible(cmp_s, mem_s, 0) return StageCost( name="S2 W_Q GEMM", - formula=f"FLOPs = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; " - f"weight = {weight_B/1e6:.1f} MB", + formula=f"FLOPs = 2*B*T_q*d*(H_q/TP*d_h) " + f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; " + f"weight = {weight_B/1e6:.1f} MB (shared across batch)", compute_s=cmp_s, memory_s=mem_s, comm_s=0.0, bound=bnd, visible_s=vis, flops=int(flops), mem_bytes=int(weight_B), - flops_formula=f"2*T_q*d*(H_q/TP*d_h) = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}", - mem_formula=f"d*(H_q/TP*d_h)*b = {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB", + flops_formula=f"2*B*T_q*d*(H_q/TP*d_h) " + f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}", + mem_formula=f"d*(H_q/TP*d_h)*b (weight, B-invariant) " + f"= {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB", ) @@ -113,12 +117,15 @@ def stage_wkv(cfg: FullConfig) -> StageCost: vis, bnd = _visible(cmp_s, mem_s, 0) return StageCost( name="S3 W_K + W_V GEMM", - formula=f"FLOPs per proj = 2*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2", + formula=f"FLOPs per proj = 2*B*T_q*d*(H_kv/TP*d_h) " + f"= 2*{B}*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2", compute_s=cmp_s, memory_s=mem_s, comm_s=0.0, bound=bnd, visible_s=vis, flops=int(flops), mem_bytes=int(weight_B), - flops_formula=f"2*(2*T_q*d*(H_kv/TP*d_h)) = 2*(2*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}", - mem_formula=f"2*(d*(H_kv/TP*d_h)*b) = 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB", + flops_formula=f"2*(2*B*T_q*d*(H_kv/TP*d_h)) " + f"= 2*(2*{B}*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}", + mem_formula=f"2*(d*(H_kv/TP*d_h)*b) (weights, B-invariant) " + f"= 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB", ) @@ -132,12 +139,14 @@ def stage_kv_append(cfg: FullConfig) -> StageCost: mem_s = bytes_ / cfg.machine.bw_hbm return StageCost( name="S4 KV cache append", - formula=f"bytes = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B", + formula=f"bytes = 2*B*T_q*(H_kv/TP)*d_h*b " + f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B", compute_s=0, memory_s=mem_s, comm_s=0, bound="memory", visible_s=mem_s, flops=0, mem_bytes=bytes_, flops_formula="0 (no matmul, just cache write)", - mem_formula=f"2*T_q*(H_kv/TP)*d_h*b = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B", + mem_formula=f"2*B*T_q*(H_kv/TP)*d_h*b " + f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B", ) @@ -172,19 +181,23 @@ def stage_qkT(cfg: FullConfig) -> StageCost: S_local = cfg.topo.s_local dh = cfg.model.d_head hq_per_pe = cfg.h_q_per_pe - flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe + B = max(1, cfg.topo.b) + flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe total_flops = flops_per_hop * passes _hop_word = "pass" if passes == 1 else "hops" return StageCost( name=f"S5 Q.K^T (x{passes} {_hop_word})", - formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop", + formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = " + f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = " + f"{flops_per_hop:.2g} FLOPs/hop", compute_s=total, memory_s=0, comm_s=0, bound="compute", visible_s=total, hop_multiplier=passes, flops=int(total_flops), mem_bytes=0, flops_formula=( - f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = " - f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}" + f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = " + f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = " + f"{total_flops:.3g}" ), mem_formula="0 (scores accumulated on-chip)", ) @@ -205,15 +218,16 @@ def stage_softmax(cfg: FullConfig) -> StageCost: _hop_word = "pass" if passes == 1 else "hops" return StageCost( name=f"S6 softmax (x{passes} {_hop_word})", - formula=f"elts/hop = {hq_per_pe}*{T_q}*{S_local} = {elems:.2g}", + formula=f"elts/hop = B*(H_q/TP)*T_q*S_local " + f"= {B}*{hq_per_pe}*{T_q}*{S_local} = {elems:.2g}", compute_s=0, memory_s=total, comm_s=0, bound="memory", visible_s=total, hop_multiplier=passes, flops=0, mem_bytes=int(total_bytes), flops_formula="~O(elts) (negligible)", mem_formula=( - f"{passes}*2*b*(H_q/TP)*T_q*S_local = " - f"{passes}*2*{b}*{hq_per_pe}*{T_q}*{S_local} = " + f"{passes}*2*b*B*(H_q/TP)*T_q*S_local = " + f"{passes}*2*{b}*{B}*{hq_per_pe}*{T_q}*{S_local} = " f"{total_bytes/1e6:.2f} MB" ), ) @@ -227,19 +241,23 @@ def stage_pv(cfg: FullConfig) -> StageCost: S_local = cfg.topo.s_local dh = cfg.model.d_head hq_per_pe = cfg.h_q_per_pe - flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe + B = max(1, cfg.topo.b) + flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe total_flops = flops_per_hop * passes _hop_word = "pass" if passes == 1 else "hops" return StageCost( name=f"S7 P.V (x{passes} {_hop_word})", - formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop", + formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = " + f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = " + f"{flops_per_hop:.2g} FLOPs/hop", compute_s=total, memory_s=0, comm_s=0, bound="compute", visible_s=total, hop_multiplier=passes, flops=int(total_flops), mem_bytes=0, flops_formula=( - f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = " - f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}" + f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = " + f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = " + f"{total_flops:.3g}" ), mem_formula="0 (accumulated on-chip)", ) @@ -302,13 +320,15 @@ def stage_merge(cfg: FullConfig) -> StageCost: return StageCost( name=f"S8 online-softmax merge (x{max(0, cfg.topo.cp-1)}){name_suffix}", - formula=f"~6*{T_q}*{hq_per_pe}*{dh}*(C-1) = {flops:.2g} FLOPs", + formula=f"~6*B*T_q*(H_q/TP)*d_h*(C-1) " + f"= 6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} " + f"= {flops:.2g} FLOPs", compute_s=cmp_s, memory_s=0, comm_s=comm_s, bound=bound, visible_s=visible_s, flops=int(flops), mem_bytes=0, comm_bytes=comm_bytes, flops_formula=( - f"6*T_q*(H_q/TP)*d_h*(CP-1) = " - f"6*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}" + f"6*B*T_q*(H_q/TP)*d_h*(CP-1) = " + f"6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}" ), mem_formula="0 (in-register)", comm_formula=comm_formula or "0", @@ -324,11 +344,13 @@ def stage_normalize(cfg: FullConfig) -> StageCost: cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util) return StageCost( name="S9 normalize O/l", - formula=f"{T_q}*{hq_per_pe}*{dh} = {flops} divisions", + formula=f"B*T_q*(H_q/TP)*d_h = {B}*{T_q}*{hq_per_pe}*{dh} " + f"= {flops} divisions", compute_s=cmp_s, memory_s=0, comm_s=0, bound="trivial", visible_s=cmp_s, flops=int(flops), mem_bytes=0, - flops_formula=f"T_q*(H_q/TP)*d_h = {T_q}*{hq_per_pe}*{dh} = {flops}", + flops_formula=f"B*T_q*(H_q/TP)*d_h " + f"= {B}*{T_q}*{hq_per_pe}*{dh} = {flops}", mem_formula="0", ) @@ -346,13 +368,16 @@ def stage_wo(cfg: FullConfig) -> StageCost: vis, bnd = _visible(cmp_s, mem_s, 0) return StageCost( name="S10 W_O GEMM", - formula=f"FLOPs = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; " - f"weight = {weight_B/1e6:.1f} MB", + formula=f"FLOPs = 2*B*T_q*(H_q/TP*d_h)*d " + f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; " + f"weight = {weight_B/1e6:.1f} MB (shared across batch)", compute_s=cmp_s, memory_s=mem_s, comm_s=0.0, bound=bnd, visible_s=vis, flops=int(flops), mem_bytes=int(weight_B), - flops_formula=f"2*T_q*(H_q/TP*d_h)*d = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}", - mem_formula=f"(H_q/TP*d_h)*d*b = {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB", + flops_formula=f"2*B*T_q*(H_q/TP*d_h)*d " + f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}", + mem_formula=f"(H_q/TP*d_h)*d*b (weight, B-invariant) " + f"= {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB", ) @@ -404,7 +429,7 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost: return StageCost( name=f"C1 O/m/l all-reduce ({tier}, x{cp} ranks)", formula=(f"decode: gather partial (O,m,l) once at end; " - f"M = T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; " + f"M = B*T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; " f"AR: 2*(CP-1)/CP*M / BW + 2*(CP-1)*alpha"), compute_s=0, memory_s=0, comm_s=ar_time, bound="comm", visible_s=ar_time, @@ -432,16 +457,16 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost: M_KV = B * ((2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)) variant_desc = "Q+O/m/l ring" formula_bytes = ( - f"2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b " - f"= 2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b} " + f"B*(2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b) " + f"= {B}*(2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b}) " f"= {M_KV} B/hop" ) else: M_KV = 2 * B * S_local * hkv_per_pe * dh * b variant_desc = "K/V ring" formula_bytes = ( - f"2*S_local*(H_kv/TP)*d_h*b " - f"= 2*{S_local}*{hkv_per_pe}*{dh}*{b} " + f"2*B*S_local*(H_kv/TP)*d_h*b " + f"= 2*{B}*{S_local}*{hkv_per_pe}*{dh}*{b} " f"= {M_KV/1e6:.3f} MB/hop" ) @@ -527,7 +552,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost: total_comm_bytes = int(2 * (tp - 1) / tp * bytes_) return StageCost( name=f"C2 TP AllReduce W_O (TP={tp} ranks, {scope})", - formula=f"2*(TP-1)/TP * {T_q}*{d}*{b} B / BW + 2(TP-1)*alpha [{scope}]", + formula=f"2*(TP-1)/TP * B*T_q*d*b B / BW + 2(TP-1)*alpha " + f"[B={B}, {scope}]", compute_s=0, memory_s=0, comm_s=comm_time, bound="comm", visible_s=comm_time, flops=0, mem_bytes=0, comm_bytes=total_comm_bytes, @@ -541,7 +567,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost: f"every rank ends up with the full hidden vector for the next\n" f"stage (RMSNorm -> FFN). Fires ONCE per layer over {scope} links.\n" f"---\n" - f"2*(TP-1)/TP * T_q*d*b = 2*({tp}-1)/{tp} * {T_q}*{d}*{b} " + f"2*(TP-1)/TP * B*T_q*d*b " + f"= 2*({tp}-1)/{tp} * {B}*{T_q}*{d}*{b} " f"= {total_comm_bytes/1e6:.2f} MB over {scope} at " f"{bw/1e9:.0f} GB/s" ), @@ -576,7 +603,9 @@ def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost: total_comm_bytes = int(2 * (split - 1) / split * bytes_per_hop * cfg.topo.cp) return StageCost( name=f"C3 Score AllReduce ({split}-way, xCP hops)", - formula=f"2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} / BW + latency", + formula=f"2*({split}-1)/{split} * B*(H_q/TP)*T_q*S_local*b / BW " + f"= 2*({split}-1)/{split} * {B}*{hq_per_pe}*{T_q}*{S_local}*{b} " + f"/ BW + latency", compute_s=0, memory_s=0, comm_s=total, bound="comm", visible_s=total, flops=0, mem_bytes=0, comm_bytes=total_comm_bytes, @@ -647,12 +676,13 @@ def stage_ffn_rmsnorm(cfg: FullConfig) -> StageCost: mem_s = bytes_ / cfg.machine.bw_hbm return StageCost( name="F1 RMSNorm (pre-FFN)", - formula=f"{T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM", + formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM", compute_s=0, memory_s=mem_s, comm_s=0, bound="memory", visible_s=mem_s, flops=flops, mem_bytes=bytes_, - flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}", - mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B", + flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}", + mem_formula=f"B*T_q*d*b + T_q*d*b (weight) " + f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B", ) @@ -668,12 +698,16 @@ def _ffn_gemm(cfg: FullConfig, name: str, ffn_per_pe: int) -> StageCost: vis, bnd = _visible(cmp_s, mem_s, 0) return StageCost( name=name, - formula=f"FLOPs = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB", + formula=f"FLOPs = 2*B*T_q*d*(ffn/div) " + f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; " + f"weight = {weight_B/1e6:.1f} MB (shared across batch)", compute_s=cmp_s, memory_s=mem_s, comm_s=0, bound=bnd, visible_s=vis, flops=int(flops), mem_bytes=int(weight_B), - flops_formula=f"2*T_q*d*(ffn/div) = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}", - mem_formula=f"d*(ffn/div)*b = {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB", + flops_formula=f"2*B*T_q*d*(ffn/div) " + f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}", + mem_formula=f"d*(ffn/div)*b (weight, B-invariant) " + f"= {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB", ) @@ -698,12 +732,14 @@ def stage_ffn_swiglu(cfg: FullConfig) -> StageCost: mem_s = bytes_ / cfg.machine.bw_hbm return StageCost( name="F4 SwiGLU act", - formula=f"3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM", + formula=f"3*B*T_q*(ffn/div)*b " + f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM", compute_s=0, memory_s=mem_s, comm_s=0, bound="memory", visible_s=mem_s, flops=flops, mem_bytes=bytes_, - flops_formula=f"~3*T_q*(ffn/div) = 3*{T_q}*{ffn_per_pe} = {flops}", - mem_formula=f"3*T_q*(ffn/div)*b = 3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B", + flops_formula=f"~3*B*T_q*(ffn/div) = 3*{B}*{T_q}*{ffn_per_pe} = {flops}", + mem_formula=f"3*B*T_q*(ffn/div)*b " + f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B", ) @@ -720,12 +756,16 @@ def stage_ffn_down(cfg: FullConfig) -> StageCost: vis, bnd = _visible(cmp_s, mem_s, 0) return StageCost( name="F5 W_down GEMM", - formula=f"FLOPs = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB", + formula=f"FLOPs = 2*B*T_q*(ffn/div)*d " + f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; " + f"weight = {weight_B/1e6:.1f} MB (shared across batch)", compute_s=cmp_s, memory_s=mem_s, comm_s=0, bound=bnd, visible_s=vis, flops=int(flops), mem_bytes=int(weight_B), - flops_formula=f"2*T_q*(ffn/div)*d = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}", - mem_formula=f"(ffn/div)*d*b = {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB", + flops_formula=f"2*B*T_q*(ffn/div)*d " + f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}", + mem_formula=f"(ffn/div)*d*b (weight, B-invariant) " + f"= {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB", ) @@ -758,7 +798,8 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost: total_comm_bytes = int(2 * (divisor - 1) / divisor * bytes_) return StageCost( name=f"CF1 FFN AllReduce (scope={scope}, x{divisor})", - formula=f"2*({divisor}-1)/{divisor} * {T_q}*{d}*{b} B / BW + latency", + formula=f"2*({divisor}-1)/{divisor} * B*T_q*d*b B / BW + latency " + f"[B={B}]", compute_s=0, memory_s=0, comm_s=comm_time, bound="comm", visible_s=comm_time, flops=0, mem_bytes=0, comm_bytes=total_comm_bytes, @@ -771,8 +812,8 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost: f"rank has the full FFN output for the next layer. Larger scope\n" f"= less FFN weight memory per PE but bigger AR bill.\n" f"---\n" - f"2*(div-1)/div * T_q*d*b = 2*({divisor}-1)/{divisor} * " - f"{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at " + f"2*(div-1)/div * B*T_q*d*b = 2*({divisor}-1)/{divisor} * " + f"{B}*{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at " f"{bw/1e9:.0f} GB/s (scope={scope})" ), )