Add flow layout: cursor placement, left rail, declared band heights
Figures were positioned by hand-written offsets. Every gap was a magic number tuned against the content that happened to be there, so a label that grew two characters landed on the next tensor, and two stages started from two different x shared no rail. Both failures compile cleanly. Replace it with a cursor. Objects placed with an empty coordinate argument reserve their own width -- including a stack's offset sheets and a bracket's overhang -- and gaps are declared once (\stgutter, \strowgap, \stblockgap). The gap belongs to the object that follows it and the first object in a band gets none, so every band starts flush on a shared rail and gap=0pt states that two shards tile exactly. \stlink makes a connector's label a flow object, which is what removes the label-wider-than-its-arrow failure entirely. \stcol is a vertical sub-flow for a split along the contracted axis. \strow declares its height, so an object that does not fit -- or a column that does not add up to what it declared, and is therefore drawn off-center -- becomes a package warning, which build.sh fails on. Absolute placement is unchanged: passing a coordinate takes the original code path, and \sttrack folds a hand-placed node back into the cursor. All three golden examples and the new tests/flow.tex are converted and build clean.
This commit is contained in:
dela
@@ -1,6 +1,8 @@
|
||||
% 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.
|
||||
% rank held across every stage, a collective node as a real operation, and a
|
||||
% split along the contracted axis drawn as a column sub-flow.
|
||||
% Layout is entirely by cursor: no absolute coordinate appears below.
|
||||
% ../scripts/build.sh tp-ffn-allreduce.tex
|
||||
\documentclass[border=10pt]{standalone}
|
||||
\usepackage[cjk]{supertensor}
|
||||
@@ -14,74 +16,71 @@
|
||||
\stdim{bt}{6} % B*T rows
|
||||
\stdim{d}{4} % model width
|
||||
\stdim{dffl}{4} % d_ff / p (per-rank hidden width)
|
||||
\stdim{dff}{8} % d_ff = p * (d_ff/p); the height of the stacked W_2 column
|
||||
|
||||
\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}
|
||||
\ststage{SA}{列切 $\mathbf W_1$：无通信}
|
||||
\strow{rowA}{bt}
|
||||
\stface[role=act, bracket=true]{X}{}{bt}{d}
|
||||
\stglyph{mA}{$\times$}
|
||||
% Two shards, tiled exactly: gap=0pt makes them adjacent by construction,
|
||||
% so neither stretching nor an eyeballed offset can creep in.
|
||||
\stface[role=r1]{W1a}{}{d}{dffl}
|
||||
\stface[role=r2, gap=0pt]{W1b}{}{d}{dffl}
|
||||
\stglyph{eA}{$=$}
|
||||
\stface[role=r1]{Ha}{}{bt}{dffl}
|
||||
\stface[role=r2, gap=0pt]{Hb}{}{bt}{dffl}
|
||||
\strowend
|
||||
\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}
|
||||
\ststage{SB}{行切 $\mathbf W_2$：一次 All-Reduce}
|
||||
\strow{rowB}{dff} % the stacked W_2 column is the tallest object here
|
||||
\stface[role=r1]{Ga}{}{bt}{dffl}
|
||||
\stface[role=r2, gap=0pt]{Gb}{}{bt}{dffl}
|
||||
\stglyph{mB}{$\times$}
|
||||
% W_2 is split along the CONTRACTED axis: the shards stack vertically and
|
||||
% together have exactly the height of H's width. Splitting reverses concat,
|
||||
% which is what gap=0pt inside the column states.
|
||||
\stcol{W2}{dff}
|
||||
\stface[role=r1]{W2a}{}{dffl}{d}
|
||||
\stface[role=r2, gap=0pt]{W2b}{}{dffl}{d}
|
||||
\stcolend
|
||||
\stglyph{eB}{$=$}
|
||||
\stface[role=r1]{Pa}{}{bt}{d}
|
||||
\stglyph{plus}{$+$}
|
||||
\stface[role=r2]{Pb}{}{bt}{d}
|
||||
\stlink{lc}{}
|
||||
\stcomm{ar}{All-Reduce}
|
||||
\stlink{ly}{}
|
||||
\stface[role=act, bracket=true]{Y}{}{bt}{d}
|
||||
\strowend
|
||||
\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{W2}{$\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
|
||||
|
||||
% ================================================================= formula ==
|
||||
% Placed last so it is centered on the figure that was actually drawn.
|
||||
\sttopformula{F}{$\begin{gathered}
|
||||
\mathbf{XW}_1=[\,\mathbf{XW}_1^{(1)}\mid
|
||||
\mathbf{XW}_1^{(2)}\,]=[\,\mathbf H^{(1)}\mid\mathbf H^{(2)}\,]\\[1.2mm]
|
||||
[\,\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)}
|
||||
\end{gathered}$}
|
||||
|
||||
% ============================================================== meaning box ==
|
||||
\node[inner sep=0pt, fit=(F)(rowA)(rowB)(Y-shape)(Ga-shape)] (all) {};
|
||||
\stbbox{all}
|
||||
\stmeaningbox{mb}{16.4cm}{all}
|
||||
{$BT$ 展平后的 token 数，$d$ 模型宽度，$d_{\mathrm{ff}}$ 前馈中间宽度，
|
||||
$p$ TP 并行度（图中 $p=2$）}
|
||||
|
||||
Reference in New Issue
Block a user