supertensor: shape-aware tensor figure toolkit
Extracted from the tensor-formula-viz skill and rebuilt around the idea that the geometry rules should be enforced by construction rather than restated as prose an agent has to remember. - assets/supertensor.sty: faces, stacks, index faces, shared caption lanes, meaning box, signature. Macros take a declared axis and a declared role, so equal shapes get equal edges, a x a is square, a transpose swaps the face, and contracted axes share an edge length -- without any manual alignment. - scripts/preflight.sh: decide the TikZ/CJK path before drawing. - scripts/build.sh: compile and fail on silent corruption (missing CJK glyphs, overfull boxes, undeclared roles), then export pdf/svg/png/thumb. - scripts/test.sh: build every figure as a regression test for the package. - examples/: three golden figures (TP-FFN, causal MHA, MoE top-k gather) plus an anti-pattern gallery of figures that compile cleanly and still lie. - SKILL.md + references/: lean entry point, details loaded on demand.
This commit is contained in:
dela
@@ -0,0 +1,84 @@
|
||||
% Anti-pattern gallery -- four ways to draw a figure that compiles cleanly and
|
||||
% still teaches the reader something false. Left of each pair is wrong.
|
||||
% ../scripts/build.sh antipatterns.tex
|
||||
\documentclass[border=10pt]{standalone}
|
||||
\usepackage[cjk]{supertensor}
|
||||
|
||||
\stsetrole{k}{stOrange}
|
||||
\stsetrole{r1}{stTeal}
|
||||
\stsetrole{r2}{stCoral}
|
||||
\stsetrole{idx}{stViolet}
|
||||
\stsetrole{act}{stTeal}
|
||||
|
||||
\stdim{T}{5}
|
||||
\stdim{dh}{3}
|
||||
\stdim{d}{6}
|
||||
\stdim{k}{2}
|
||||
|
||||
\newcommand{\bad}[1]{{\color{stCoral}$\times$}\;#1}
|
||||
\newcommand{\good}[1]{{\color{stTeal}$\checkmark$}\;#1}
|
||||
|
||||
\begin{document}
|
||||
\begin{tikzpicture}
|
||||
|
||||
% ------------------------------------------------------------- pair 1 + 2 ---
|
||||
\node[st stage, anchor=north west] (S1) at (0,0) {(1) 转置只改了标签};
|
||||
\coordinate (p1) at ($(S1.west)+(1.0,-1.7)$);
|
||||
|
||||
% WRONG: same face as K, relabelled. The reader cannot see the contracted axis.
|
||||
\stface[role=k, bracket=true]{A1}{(p1)}{T}{dh}
|
||||
\stface[role=k, bracket=true]{A2}{($(A1.east)+(2.2,0)$)}{dh}{T}
|
||||
|
||||
\node[st stage, anchor=north west] (S2) at ($(S1.west)+(7.6,0)$) {(2) 分片没有铺满母体};
|
||||
\coordinate (p2) at ($(S2.west)+(1.0,-1.7)$);
|
||||
|
||||
% WRONG: a gap between the shards. The parent's width is now a lie.
|
||||
\stface[role=r1]{B1a}{(p2)}{T}{dh}
|
||||
\stface[role=r2]{B1b}{($(B1a.east)+(2.6*\stunit,0)$)}{T}{dh}
|
||||
\stface[role=r1]{B2a}{($(B1b.east)+(2.4,0)$)}{T}{dh}
|
||||
\stface[role=r2]{B2b}{($(B2a.east)+(1.5*\stunit,0)$)}{T}{dh}
|
||||
|
||||
\node[inner sep=0pt, fit=(A1)(A2)(B1a)(B2b)] (row1) {};
|
||||
\stlane{row1}
|
||||
\stcaption{A1}{\bad{$\mathbf K^{\top}$}}{面仍是 $T\times d_h$}
|
||||
\stcaption{A2}{\good{$\mathbf K^{\top}$}}{面已换成 $d_h\times T$}
|
||||
\stcaption{B1a}{\bad{$[\mathbf W^{(1)}\mid\mathbf W^{(2)}]$}}{中间凭空多出空隙}
|
||||
\stcaption{B2a}{\good{$[\mathbf W^{(1)}\mid\mathbf W^{(2)}]$}}{两片正好铺满}
|
||||
\stnolane
|
||||
|
||||
% ------------------------------------------------------------- pair 3 + 4 ---
|
||||
\coordinate (y2) at ($(A1-shape.south)+(0,-10mm)$);
|
||||
\node[st stage, anchor=north west] (S3) at (S1.west |- y2) {(3) 索引画成了热力图};
|
||||
\coordinate (p3) at ($(S3.west)+(1.0,-1.6)$);
|
||||
|
||||
% WRONG: a lightness ramp on an index invites `expert 3 > expert 0'.
|
||||
\stface[role=idx]{C1}{(p3)}{T}{k}
|
||||
\stindexface[role=idx]{C2}{($(C1.east)+(2.4,0)$)}{T}{k}{0,1, 1,2, 2,3, 3,0, 0,2}
|
||||
|
||||
\node[st stage, anchor=north west] (S4) at (S2.west |- y2) {(4) 全图一个色阶};
|
||||
\coordinate (p4) at ($(S4.west)+(1.0,-1.6)$);
|
||||
|
||||
% WRONG: one pale level everywhere -- nothing is legible at thumbnail size.
|
||||
\stface[role=act, pattern=solid, level=1]{D1}{(p4)}{T}{d}
|
||||
\stface[role=act]{D2}{($(D1.east)+(2.6,0)$)}{T}{d}
|
||||
|
||||
\node[inner sep=0pt, fit=(C1)(C2)(D1)(D2)] (row2) {};
|
||||
\stlane{row2}
|
||||
\stcaption{C1}{\bad{$\mathcal I$}}{深浅暗示大小可比}
|
||||
\stcaption{C2}{\good{$\mathcal I$}}{离散符号，无色阶}
|
||||
\stcaption{D1}{\bad{$\mathbf X$}}{只有一档淡色}
|
||||
\stcaption{D2}{\good{$\mathbf X$}}{三档亮度，非周期}
|
||||
\stnolane
|
||||
|
||||
% ============================================================== meaning box ==
|
||||
\node[inner sep=0pt, fit=(S1)(row1)(row2)(D2-shape)(B2a-shape)] (all) {};
|
||||
\stmeaningbox{mb}{16.2cm}{all}
|
||||
{}
|
||||
{四张“错图”都能干净编译。编译器只检查 \TeX{} 的语法，不检查图讲的事情对不对，
|
||||
所以交付前必须按 \texttt{references/checklist.md} 做一次人眼审图}
|
||||
{(1) 转置要物理交换面宽高；(2) 分片必须精确铺满母体，省略要画省略号；
|
||||
(3) 索引、掩码、数值是三套画法；(4) 亮度分档是对比度的来源，不是饱和度}
|
||||
\stsignature{反面示例：编译通过但讲错的四种画法}{mb}
|
||||
|
||||
\end{tikzpicture}
|
||||
\end{document}
|
||||
@@ -0,0 +1,105 @@
|
||||
% Golden example 2 -- causal multi-head attention.
|
||||
% Shows: leading axes as stack depth, a transpose that physically swaps the
|
||||
% face, equal edge length on the contracted axis, and a Boolean mask drawn in
|
||||
% a different grammar from the scores it gates.
|
||||
% ../scripts/build.sh mha-causal.tex
|
||||
\documentclass[border=10pt]{standalone}
|
||||
\usepackage[cjk]{supertensor}
|
||||
|
||||
\stsetrole{q}{stTeal}
|
||||
\stsetrole{k}{stOrange}
|
||||
\stsetrole{v}{stViolet}
|
||||
\stsetrole{s}{stCoral}
|
||||
\stsetrole{w}{stGray}
|
||||
|
||||
\stdim{T}{6}
|
||||
\stdim{dh}{3}
|
||||
\stdim{d}{9} % d = h * d_h, h = 3
|
||||
|
||||
\begin{document}
|
||||
\begin{tikzpicture}
|
||||
|
||||
% ================================================================= formula ==
|
||||
\node (F) at (0,0) {\stformula{$\displaystyle
|
||||
\mathbf A^{(i)}=\mathrm{softmax}\!\left(
|
||||
\frac{\mathbf Q^{(i)}\mathbf K^{(i)\top}}{\sqrt{d_h}}+\mathbf M\right),\qquad
|
||||
\mathbf O^{(i)}=\mathbf A^{(i)}\mathbf V^{(i)}$}};
|
||||
|
||||
% ============================================================ stage A row ===
|
||||
\node[st stage, below=7mm of F] (SA) {每头打分：沿 $d_h$ 收缩};
|
||||
\coordinate (a) at ($(SA)+(-3.9,-1.9)$);
|
||||
|
||||
\ststack[role=q, bracket=true]{Q}{(a)}{T}{dh}{3}
|
||||
\node[st op, right=6mm of Q] (mA) {$\times$};
|
||||
% K^T: the face is physically swapped, not relabelled. Its height equals Q's
|
||||
% width -- that is the contracted axis d_h, drawn at one edge length.
|
||||
\ststack[role=k, bracket=true]{KT}{($(mA)+(1.9,0)$)}{dh}{T}{3}
|
||||
\node[st op, right=6mm of KT] (eA) {$=$};
|
||||
\ststack[role=s]{S}{($(eA)+(2.0,0)$)}{T}{T}{3}
|
||||
|
||||
\node[inner sep=0pt, fit=(Q)(KT)(S)] (rowA) {};
|
||||
\stlane{rowA}
|
||||
\stcaption{Q}{$\mathbf Q^{(i)}$}{$h\times T\times d_h$}
|
||||
\stcaption{KT}{$\mathbf K^{(i)\top}$}{$h\times d_h\times T$}
|
||||
\stcaption{S}{$\mathbf S^{(i)}$}{$h\times T\times T$}
|
||||
\stnolane
|
||||
|
||||
% ============================================================ stage B row ===
|
||||
\node[st stage, below=9mm of Q-shape.south west, anchor=north west] (SB)
|
||||
{因果掩码与加权求和};
|
||||
\coordinate (b) at ($(SB)+(0.6,-2.0)$);
|
||||
|
||||
% The mask is a Boolean support, not a magnitude: one flat level, exact
|
||||
% triangle, no stack -- it is shared by every head.
|
||||
\stface[role=w, pattern=data,
|
||||
data={300000,330000,333000,333300,333330,333333}]{M}{(b)}{T}{T}
|
||||
\ststack[role=s, pattern=causal]{A}{($(M.east)+(3.75,0)$)}{T}{T}{3}
|
||||
\starrowlabel{M.east}{A.west}{softmax}
|
||||
\node[st op, right=6mm of A] (mB) {$\times$};
|
||||
\ststack[role=v, bracket=true]{V}{($(mB)+(1.4,0)$)}{T}{dh}{3}
|
||||
\node[st op, right=6mm of V] (eB) {$=$};
|
||||
\ststack[role=v]{O}{($(eB)+(1.4,0)$)}{T}{dh}{3}
|
||||
|
||||
\node[inner sep=0pt, fit=(M)(A)(V)(O)] (rowB) {};
|
||||
\stlane{rowB}
|
||||
\stcaption{M}{$\mathbf M$}{$T\times T$}
|
||||
\stcaption{A}{$\mathbf A^{(i)}$}{$h\times T\times T$}
|
||||
\stcaption{V}{$\mathbf V^{(i)}$}{$h\times T\times d_h$}
|
||||
\stcaption{O}{$\mathbf O^{(i)}$}{$h\times T\times d_h$}
|
||||
\stnolane
|
||||
|
||||
% ============================================================ stage C row ===
|
||||
\node[st stage, below=9mm of M-shape.south west, anchor=north west] (SC)
|
||||
{沿 $d_h$ 拼接后投影};
|
||||
\coordinate (c) at ($(SC)+(1.2,-2.0)$);
|
||||
|
||||
% Concatenation reverses the split: three h-shards of width d_h tile a face of
|
||||
% width d exactly.
|
||||
\stface[role=v]{C1}{(c)}{T}{dh}
|
||||
\stface[role=v]{C2}{($(C1.east)+(1.5*\stunit,0)$)}{T}{dh}
|
||||
\stface[role=v]{C3}{($(C2.east)+(1.5*\stunit,0)$)}{T}{dh}
|
||||
\node[st op, right=6mm of C3] (mC) {$\times$};
|
||||
\stface[role=w]{WO}{($(mC)+(2.5,0)$)}{d}{d}
|
||||
\node[st op, right=6mm of WO] (eC) {$=$};
|
||||
\stface[role=v, bracket=true]{Y}{($(eC)+(2.5,0)$)}{T}{d}
|
||||
|
||||
\node[inner sep=0pt, fit=(C1)(WO)(Y)] (rowC) {};
|
||||
\stlane{rowC}
|
||||
\stcaption{C2}{$[\,\mathbf O^{(1)}\mid\mathbf O^{(2)}\mid\mathbf O^{(3)}\,]$}{$T\times d$}
|
||||
\stcaption{WO}{$\mathbf W_O$}{$d\times d$}
|
||||
\stcaption{Y}{$\mathbf Y$}{$T\times d$}
|
||||
\stnolane
|
||||
|
||||
% ============================================================== meaning box ==
|
||||
\node[inner sep=0pt, fit=(F)(rowA)(rowB)(rowC)(Y-shape)(C2-shape)] (all) {};
|
||||
\stmeaningbox{mb}{16.8cm}{all}
|
||||
{$T$ 序列长度，$d_h$ 单头宽度，$h$ 头数（图中 $h=3$，即堆叠的三张面），
|
||||
$d=h\,d_h$；批轴 $B$ 省略}
|
||||
{$\mathbf S,\mathbf A$ 是分数与概率（行和为 $1$）；$\mathbf M\in\{0,-\infty\}^{T\times T}$
|
||||
是布尔支撑而非数值，被所有头共享，故只画一张、不堆叠；紫色一族标记 $\mathbf V\rightarrow\mathbf O\rightarrow\mathbf Y$ 同一数据流}
|
||||
{$(T\times d_h)(d_h\times T)\rightarrow(T\times T)$：$\mathbf K^{\top}$ 的面高即收缩维 $d_h$；
|
||||
拼接是切分的逆运算，$3$ 个 $d_h$ 恰好铺满 $d$}
|
||||
\stsignature{因果多头注意力（掩码 + 拼接投影）}{mb}
|
||||
|
||||
\end{tikzpicture}
|
||||
\end{document}
|
||||
@@ -0,0 +1,114 @@
|
||||
% Golden example 3 -- MoE top-k routing and gather.
|
||||
% Shows the three cell semantics side by side in one figure: a SCORE face
|
||||
% (graded lightness), an INDEX face (discrete symbols, no ramp), and a BOOLEAN
|
||||
% support face (one flat level) -- plus a gather whose output heights are data
|
||||
% dependent and must sum back to T*k.
|
||||
% ../scripts/build.sh moe-topk-gather.tex
|
||||
\documentclass[border=10pt]{standalone}
|
||||
\usepackage[cjk]{supertensor}
|
||||
|
||||
\stsetrole{act}{stTeal} % X and the per-expert buffers: same data, regrouped
|
||||
\stsetrole{wg}{stViolet} % learned router weight
|
||||
\stsetrole{s}{stCoral} % scores
|
||||
\stsetrole{idx}{stOrange} % indices / token ids
|
||||
\stsetrole{m}{stGray} % boolean support
|
||||
|
||||
\stdim{T}{6}
|
||||
\stdim{d}{4}
|
||||
\stdim{E}{4}
|
||||
\stdim{k}{2}
|
||||
\stdim{ne}{3} % capacity per expert in this instance: n_e = 3
|
||||
|
||||
\begin{document}
|
||||
\begin{tikzpicture}
|
||||
|
||||
% ================================================================= formula ==
|
||||
\node (F) at (0,0) {\stformula{$\displaystyle
|
||||
\mathbf G=\mathrm{softmax}(\mathbf{XW}_g),\quad
|
||||
\mathcal I_t=\operatorname*{top-}k_{e}\,\mathbf G_{t,e},\quad
|
||||
\mathbf D_{t,e}=\mathbf 1[e\in\mathcal I_t],\quad
|
||||
\mathbf X^{(e)}=\mathrm{gather}(\mathbf X,\mathbf D_{:,e})$}};
|
||||
|
||||
% ============================================================ stage A row ===
|
||||
\node[st stage, below=7mm of F] (SA) {打分：token 对专家};
|
||||
\coordinate (a) at ($(SA)+(-3.6,-1.8)$);
|
||||
|
||||
\stface[role=act, bracket=true]{X}{(a)}{T}{d}
|
||||
\node[st op, right=6mm of X] (mA) {$\times$};
|
||||
\stface[role=wg]{Wg}{($(mA)+(1.6,0)$)}{d}{E}
|
||||
\node[st op, right=6mm of Wg] (eA) {$=$};
|
||||
\stface[role=s]{G}{($(eA)+(1.6,0)$)}{T}{E}
|
||||
|
||||
\node[inner sep=0pt, fit=(X)(Wg)(G)] (rowA) {};
|
||||
\stlane{rowA}
|
||||
\stcaption{X}{$\mathbf X$}{$T\times d$}
|
||||
\stcaption{Wg}{$\mathbf W_g$}{$d\times E$}
|
||||
\stcaption{G}{$\mathbf G$}{$T\times E$}
|
||||
\stnolane
|
||||
|
||||
% ============================================================ stage B row ===
|
||||
\node[st stage, below=9mm of X-shape.south west, anchor=north west] (SB)
|
||||
{取前 $k$：连续分数 $\rightarrow$ 离散选择};
|
||||
\coordinate (b) at ($(SB.west)+(0.6,-1.7)$);
|
||||
|
||||
% Indices are drawn as symbols, not as magnitudes: expert 3 is not "bigger"
|
||||
% than expert 0, so the index face gets no lightness ramp.
|
||||
\stindexface[role=idx]{I}{(b)}{T}{k}{0,1, 1,2, 2,3, 3,0, 0,2, 1,3}
|
||||
% The same routing decision as a boolean support: one flat level, exactly k
|
||||
% cells per row, and every unselected cell left unfilled.
|
||||
\stface[role=m, pattern=data, level=3,
|
||||
data={3300,0330,0033,3003,3030,0303}]{D}{($(I.east)+(3.1,0)$)}{T}{E}
|
||||
\starrowlabel{I.east}{D.west}{one-hot}
|
||||
|
||||
\node[st comm, right=9mm of D] (gz) {Gather};
|
||||
\starrow{D.east}{gz.west}
|
||||
|
||||
\node[inner sep=0pt, fit=(I)(D)(gz)] (rowB) {};
|
||||
\stlane{rowB}
|
||||
\stcaption{I}{$\mathcal I$}{$T\times k$}
|
||||
\stcaption{D}{$\mathbf D$}{$T\times E$}
|
||||
\stnolane
|
||||
|
||||
% ============================================================ stage C row ===
|
||||
% Every stage heading starts on the same left rail; only the vertical position
|
||||
% follows the previous row.
|
||||
\coordinate (cy) at ($(I-shape.south)+(0,-9mm)$);
|
||||
\node[st stage, anchor=north west] (SC) at (SB.west |- cy)
|
||||
{按专家聚合：每个缓冲区的高度是数据决定的};
|
||||
\coordinate (c) at ($(SC.west)+(0.5,-1.9)$);
|
||||
|
||||
% Each buffer keeps X's width d -- gather regroups rows, it never reshapes the
|
||||
% feature axis. The heights are n_e, and they must sum to T*k.
|
||||
\stindexface[role=idx, border=false]{t0}{(c)}{ne}{1}{1,4,5}
|
||||
\stface[role=act]{B0}{($(t0.east)+(2*\stunit,0)$)}{ne}{d}
|
||||
\stindexface[role=idx, border=false]{t1}{($(B0.east)+(1.1,0)$)}{ne}{1}{1,2,6}
|
||||
\stface[role=act]{B1}{($(t1.east)+(2*\stunit,0)$)}{ne}{d}
|
||||
\stindexface[role=idx, border=false]{t2}{($(B1.east)+(1.1,0)$)}{ne}{1}{2,3,5}
|
||||
\stface[role=act]{B2}{($(t2.east)+(2*\stunit,0)$)}{ne}{d}
|
||||
\stindexface[role=idx, border=false]{t3}{($(B2.east)+(1.1,0)$)}{ne}{1}{3,4,6}
|
||||
\stface[role=act]{B3}{($(t3.east)+(2*\stunit,0)$)}{ne}{d}
|
||||
|
||||
\stcaptiontop{t0}{\stshapefont{token}}
|
||||
|
||||
\node[inner sep=0pt, fit=(t0)(B3)] (rowC) {};
|
||||
\stlane{rowC}
|
||||
\stcaption{B0}{$\mathbf X^{(1)}$}{$n_1\times d$}
|
||||
\stcaption{B1}{$\mathbf X^{(2)}$}{$n_2\times d$}
|
||||
\stcaption{B2}{$\mathbf X^{(3)}$}{$n_3\times d$}
|
||||
\stcaption{B3}{$\mathbf X^{(4)}$}{$n_4\times d$}
|
||||
\stnolane
|
||||
|
||||
% ============================================================== meaning box ==
|
||||
\node[inner sep=0pt, fit=(F)(rowA)(rowB)(rowC)(B3-shape)(t0-top)] (all) {};
|
||||
\stmeaningbox{mb}{16.6cm}{all}
|
||||
{$T$ token 数，$d$ 模型宽度，$E$ 专家数（图中 $E=4$），$k$ 每 token 选中的专家数
|
||||
（图中 $k=2$），$n_e$ 落到第 $e$ 个专家的 token 数}
|
||||
{$\mathbf G$ 是分数，深浅可比大小；$\mathcal I$ 是索引，格内是符号不是数值，
|
||||
故不用深浅；$\mathbf D$ 是布尔支撑，只有一档灰、每行恰好 $k$ 格；
|
||||
$\mathbf X^{(e)}$ 与 $\mathbf X$ 同色，因为它是同一批数据换了分组}
|
||||
{top-$k$ 把连续分数截成离散选择，这一步不可微；gather 只重排行、不动特征轴，
|
||||
故每个缓冲区仍是 $d$ 宽；$\sum_e n_e=Tk$，图中 $4\times3=6\times2$}
|
||||
\stsignature{MoE 路由：top-$k$ 选择与按专家 gather}{mb}
|
||||
|
||||
\end{tikzpicture}
|
||||
\end{document}
|
||||
@@ -0,0 +1,95 @@
|
||||
% Golden example 1 -- tensor parallel FFN, column-then-row sharding + AllReduce.
|
||||
% Shows: partition geometry (shards tile the parent exactly), one hue per TP
|
||||
% rank held across every stage, a collective node as a real operation.
|
||||
% ../scripts/build.sh tp-ffn-allreduce.tex
|
||||
\documentclass[border=10pt]{standalone}
|
||||
\usepackage[cjk]{supertensor}
|
||||
|
||||
% --- role ledger: one hue per TP rank, held from W through H to P ----------
|
||||
\stsetrole{act}{stViolet} % activations that every rank sees
|
||||
\stsetrole{r1}{stTeal} % rank 1
|
||||
\stsetrole{r2}{stOrange} % rank 2
|
||||
|
||||
% --- geometry ledger: one physical edge per symbolic axis ------------------
|
||||
\stdim{bt}{6} % B*T rows
|
||||
\stdim{d}{4} % model width
|
||||
\stdim{dffl}{4} % d_ff / p (per-rank hidden width)
|
||||
|
||||
\begin{document}
|
||||
\begin{tikzpicture}
|
||||
|
||||
% ================================================================= formula ==
|
||||
\node (F) at (0,0) {\stformula{$\mathbf{XW}_1=[\,\mathbf{XW}_1^{(1)}\mid
|
||||
\mathbf{XW}_1^{(2)}\,]=[\,\mathbf H^{(1)}\mid\mathbf H^{(2)}\,]$}};
|
||||
\node[below=1.2mm of F] (F2) {\stformula{$\displaystyle
|
||||
[\,\mathbf H^{(1)}\mid\mathbf H^{(2)}\,]
|
||||
\begin{bmatrix}\mathbf W_2^{(1)}\\[-1pt]\mathbf W_2^{(2)}\end{bmatrix}
|
||||
=\sum_{r}\mathbf H^{(r)}\mathbf W_2^{(r)}=\sum_r\mathbf P^{(r)}$}};
|
||||
|
||||
% ============================================================ stage A row ===
|
||||
\node[st stage, below=7mm of F2] (SA) {列切 $\mathbf W_1$：无通信};
|
||||
\coordinate (a) at ($(SA)+(-5.6,-1.8)$);
|
||||
|
||||
\stface[role=act, bracket=true]{X}{(a)}{bt}{d}
|
||||
\node[st op, right=5mm of X] (mA) {$\times$};
|
||||
% Two shards, tiled exactly: adjacent faces, no stretching, no gap.
|
||||
\stface[role=r1]{W1a}{($(mA)+(1.5,0)$)}{d}{dffl}
|
||||
\stface[role=r2]{W1b}{($(W1a.east)+(2*\stunit,0)$)}{d}{dffl}
|
||||
\node[st op, right=5mm of W1b] (eA) {$=$};
|
||||
\stface[role=r1]{Ha}{($(eA)+(1.5,0)$)}{bt}{dffl}
|
||||
\stface[role=r2]{Hb}{($(Ha.east)+(2*\stunit,0)$)}{bt}{dffl}
|
||||
|
||||
\node[inner sep=0pt, fit=(X)(W1a)(Ha)(Hb)] (rowA) {};
|
||||
\stlane{rowA}
|
||||
\stcaption{X}{$\mathbf X$}{$BT\times d$}
|
||||
\stcaption{W1a}{$\mathbf W_1^{(1)}$}{$d\times d_{\mathrm{ff}}/p$}
|
||||
\stcaption{W1b}{$\mathbf W_1^{(2)}$}{$d\times d_{\mathrm{ff}}/p$}
|
||||
\stcaption{Ha}{$\mathbf H^{(1)}$}{$BT\times d_{\mathrm{ff}}/p$}
|
||||
\stcaption{Hb}{$\mathbf H^{(2)}$}{$BT\times d_{\mathrm{ff}}/p$}
|
||||
\stnolane
|
||||
|
||||
% ============================================================ stage B row ===
|
||||
\node[st stage, below=9mm of X-shape.south west, anchor=north west] (SB)
|
||||
{行切 $\mathbf W_2$：一次 All-Reduce};
|
||||
\coordinate (b) at ($(SB)+(-0.4,-1.9)$);
|
||||
|
||||
\stface[role=r1]{Ga}{(b)}{bt}{dffl}
|
||||
\stface[role=r2]{Gb}{($(Ga.east)+(2*\stunit,0)$)}{bt}{dffl}
|
||||
\node[st op, right=5mm of Gb] (mB) {$\times$};
|
||||
% W_2 is split along the CONTRACTED axis: the two shards stack vertically and
|
||||
% together have exactly the height of H's width. Splitting reverses concat.
|
||||
\stface[role=r1]{W2a}{($(mB)+(1.35,0.46)$)}{dffl}{d}
|
||||
\stface[role=r2]{W2b}{($(W2a.south)+(0,-2*\stunit)$)}{dffl}{d}
|
||||
\node[st op, right=5mm of W2a.east |- W2a.south] (eB) {$=$};
|
||||
\stface[role=r1]{Pa}{($(eB)+(1.3,0)$)}{bt}{d}
|
||||
\node[st op, right=4mm of Pa] (plus) {$+$};
|
||||
\stface[role=r2]{Pb}{($(plus)+(1.3,0)$)}{bt}{d}
|
||||
|
||||
\node[st comm, right=9mm of Pb] (ar) {All-Reduce};
|
||||
\stface[role=act, bracket=true]{Y}{($(ar)+(1.9,0)$)}{bt}{d}
|
||||
\starrow{Pb.east}{ar.west}
|
||||
\starrow{ar.east}{Y.west}
|
||||
|
||||
\node[inner sep=0pt, fit=(Ga)(W2a)(W2b)(Pa)(Pb)(Y)] (rowB) {};
|
||||
\stlane{rowB}
|
||||
\stcaption{Ga}{$\mathbf G^{(1)}$}{$BT\times d_{\mathrm{ff}}/p$}
|
||||
\stcaption{Gb}{$\mathbf G^{(2)}$}{$BT\times d_{\mathrm{ff}}/p$}
|
||||
\stcaption{W2b}{$\mathbf W_2^{(r)}$}{$d_{\mathrm{ff}}/p\times d$}
|
||||
\stcaption{Pa}{$\mathbf P^{(1)}$}{$BT\times d$}
|
||||
\stcaption{Pb}{$\mathbf P^{(2)}$}{$BT\times d$}
|
||||
\stcaption{Y}{$\mathbf Y$}{$BT\times d$}
|
||||
\stnolane
|
||||
|
||||
% ============================================================== meaning box ==
|
||||
\node[inner sep=0pt, fit=(F)(rowA)(rowB)(Y-shape)(Ga-shape)] (all) {};
|
||||
\stmeaningbox{mb}{16.4cm}{all}
|
||||
{$BT$ 展平后的 token 数，$d$ 模型宽度，$d_{\mathrm{ff}}$ 前馈中间宽度，
|
||||
$p$ TP 并行度（图中 $p=2$）}
|
||||
{$\mathbf X,\mathbf Y$ 每个 rank 完整持有；$\mathbf W_1^{(r)},\mathbf W_2^{(r)},
|
||||
\mathbf H^{(r)},\mathbf P^{(r)}$ 仅本 rank 持有，$\mathbf P^{(r)}$ 是部分和而非最终输出}
|
||||
{$\mathbf G^{(r)}=\mathrm{GeLU}(\mathbf H^{(r)})$ 逐元素、无跨 rank 依赖；列切 $\mathbf W_1$ 使 $\mathbf H$ 沿 $d_{\mathrm{ff}}$ 切分；行切 $\mathbf W_2$ 沿收缩维切分，
|
||||
故 $\mathbf Y=\sum_r\mathbf P^{(r)}$ 需一次 All-Reduce，前向每层仅此一次通信}
|
||||
\stsignature{TP-FFN（GeLU + All-Reduce）}{mb}
|
||||
|
||||
\end{tikzpicture}
|
||||
\end{document}
|
||||
Reference in New Issue
Block a user