SPB Git

spb/wp33_uqo

Public MIT
7.3 KB · 100 lines latex
1\section{Mathematical Formulation}2\label{sec:formulation}34\subsection{Objects}5Let $\E$ be the set of (layer, expert) pairs, $|\E| = LE$, each of size $b$ bytes (3.375~MiB for OLMoE at 4 bits; equal6within an architecture). Let $\Nset = \{0,1,\dots,N\}$ be the nodes, node $0$ being the coordinator that runs attention, the7router and the KV cache, and let tier $\mathcal{T}_n = \{\mathrm{ram}, \mathrm{ssd}\}$ be the storage tiers of node $n$.8Tokens are steps $t = 1,2,\dots$; the router of layer $l$ at step $t$ selects $S_{t,l} \subset \E_l$, $|S_{t,l}| = K$.9Decision variables are binary and indexed by step:10\begin{align}11  x_{e,n,t} &\in \{0,1\} \quad \text{expert $e$ is resident in the memory of node $n$ at step $t$,} \\12  z_{e,n,t} &\in \{0,1\} \quad \text{expert $e$ is \emph{executed} on node $n$ at step $t$,} \\13  p_{e,n,t} &\in \{0,1\} \quad \text{a prefetch of $e$ into node $n$ is issued during step $t$.}14\end{align}15Every expert exists on the local NVMe of each node that holds a copy of the checkpoint (the SSD tier is unbounded for our16purposes and costs only time), so residency refers to memory.1718\subsection{Constraints}19\begin{align}20  \sum_{e \in \E} b\, x_{e,n,t} &\le M_n \qquad \forall n, t && \text{(memory budget of node $n$)} \label{eq:mem}\\21  \sum_{n \in \Nset} z_{e,n,t} &= \ind[e \in S_{t,l(e)}] \qquad \forall e, t && \text{(every selected expert executed exactly once)} \label{eq:exec}\\22  z_{e,n,t} &\le x_{e,n,t} + m_{e,n,t} \qquad \forall e, n, t && \text{(execution needs residency or a demand miss $m$)} \label{eq:res}\\23  x_{e,n,t+1} &\le x_{e,n,t} + p_{e,n,t} + m_{e,n,t} \qquad \forall e, n, t && \text{(residency only grows by prefetch or miss)} \label{eq:dyn}24\end{align}25where $m_{e,n,t} \in \{0,1\}$ is the demand-miss indicator, itself a consequence of the router and of $x$. Constraint26\eqref{eq:exec} encodes the correctness invariant: the set of executed experts is the router's set, whatever the placement.2728\subsection{Costs}29The per-step latency decomposes along the layer sequence, because layer $l$ cannot start before layer $l-1$'s output exists:30\begin{equation}31  \Lambda_t \;=\; \sum_{l=1}^{L} \Big[ \underbrace{\tau^{\mathrm{attn}}_{0}}_{\text{attention, node 0}}32  + \max_{n \in \Nset}\big( \underbrace{c_n \cdot \textstyle\sum_{e \in S_{t,l}} z_{e,n,t}}_{\text{expert compute on $n$}}33  + \underbrace{\rho_{0n}\,\ind[\exists e\!:\! z_{e,n,t}=1, n\neq 0]}_{\text{activation round trip}} \big)34  + \underbrace{\sigma_{t,l}}_{\text{stall}} \Big],35  \label{eq:latency}36\end{equation}37with $c_n$ the measured per-expert compute time of node $n$, $\rho_{0n}$ the measured round-trip time of a38$d$-dimensional activation and its result between node $0$ and $n$ (on our LAN, $0.5$--$1.0$~ms), and the stall39\begin{equation}40  \sigma_{t,l} \;=\; \max\Big(0,\; \max_{e \in S_{t,l},\, m_{e,n,t} = 1} \big( \lambda_{\mathrm{tier}(e)} - \eta_{e,t} \big) \Big),41  \label{eq:stall}42\end{equation}43where $\lambda_{\mathrm{tier}}$ is the fetch latency of the tier the expert comes from ($\lambda_{\mathrm{pc}} = 0.21$~ms page44cache, $\lambda_{\mathrm{ssd}} = 0.43$~ms cold NVMe, $\lambda_{\mathrm{lan}} \approx 28$~ms for 3.4~MB over 0.945~Gbit/s, all45measured on our hardware) and $\eta_{e,t}$ is the lead time between the moment a prefetch of $e$ was issued and the moment46$e$ is needed ($\eta = 0$ for a demand miss). Energy per step is $\Phi_t = \sum_n P_n \cdot \Lambda^{(n)}_t$ with $P_n$ the47package power of node $n$ measured by \code{powermetrics} during the run, and communication cost $\Gamma_t = \sum_{n \ne 0}48\beta_{0n}\,(\text{activation bytes} + \text{weight bytes moved})$.4950\subsection{Problem}51\begin{equation}52  \min_{x,z,p}\; \Ex\Big[\sum_{t} \Lambda_t + \lambda\,\Phi_t + \mu\,\Gamma_t\Big] \quad \text{s.t. \eqref{eq:mem}--\eqref{eq:dyn}},53  \label{eq:problem}54\end{equation}55where the expectation is over the router's sequence $\{S_{t,l}\}$, which is revealed online: at step $t$, $S_{t,l}$ becomes56known only after layer $l-1$'s output is computed. The problem is a stochastic online integer program; its offline relaxation57with known $S$ is the paging problem with multiple caches and non-uniform fetch costs, for which Belady's furthest-in-future58rule is optimal per cache and which our simulator's \emph{oracle} predictor approximates.5960\subsection{When can a prefetch pay?}61\label{sec:prefetch-condition}62Consider a predictor issuing, when layer $l-d$'s routing is known at step $t$, a set $\hat S_{t,l}$ of $m$ experts for layer63$l$, with recall $r = \Ex|\hat S \cap S|/K$ and precision $q = \Ex|\hat S \cap S|/m$. With per-layer decode time $\tau_\ell$64the lead time is $\eta = d\,\tau_\ell$. Per layer, the expected stall without prefetch is65$\sigma_0 = \lceil K(1-h)/w \rceil\, \lambda$ for hit rate $h$ and $w$ parallel loaders, and with prefetch66\begin{equation}67  \sigma_d \;\approx\; \lceil K(1-h)(1-r')/w \rceil\,\lambda \;+\; \lceil K(1-h)\, r'/w \rceil\,\max(0, \lambda - d\,\tau_\ell),68  \label{eq:prefetch-gain}69\end{equation}70where $r'$ is the recall among the experts that would have missed. The prefetch also adds $m(1-q)$ useless fetches whose71bandwidth and install cost $\kappa$ each is paid on the critical path when the loader is shared with demand misses.72A prefetch can only reduce latency if73\begin{equation}74  r'\,\min(\lambda, d\,\tau_\ell) \;>\; m(1-q)\,\kappa\;/\;K(1-h),75  \label{eq:condition}76\end{equation}77and it can \emph{remove} the stall only if $d\,\tau_\ell \ge \lambda$. On the laptop, $\tau_\ell = 0.128$~ms (489~\tps{} over7816 layers), so a page-cache fetch ($\lambda = 0.21$~ms) needs $d \ge 2$ and a cold-NVMe fetch ($\lambda = 0.43$~ms) needs79$d \ge 4$; a LAN fetch ($\lambda \approx 28$~ms) can never be hidden within a token on this model, which is why weights are not80moved over the LAN during decode in \titan{} and remote experts are instead executed where they reside. The one-layer online81experiment of Section~\ref{sec:results-prefetch} violated \eqref{eq:condition} ($d = 1$, $q \approx 0.5$) and lost throughput, as82predicted; the depth-$d$ experiments test the condition directly.8384\begin{proposition}[Demand loading is $\mathcal{O}(1)$-competitive per layer only when the table exceeds the per-step85working set]86\label{prop:capacity}87With per-layer tables of capacity $C \ge K$, LRU eviction, and no prefetch, the number of demand misses at a decode step is at88most $K$ and at least $K - |S_{t,l} \cap S_{t-1,l}| $ when $C = K$; for $C \ge |\bigcup_{s \le \Delta} S_{t-s,l}|$ the misses at89step $t$ are exactly the experts not used in the last $\Delta$ steps at that layer.90\end{proposition}91\begin{proof}92With $C = K$ the table after step $t-1$ holds exactly $S_{t-1,l}$, so every expert of $S_{t,l} \setminus S_{t-1,l}$ misses and93no other does; misses are bounded by $K$ because $|S_{t,l}| = K$. For the second claim, LRU keeps the most recently used94$C$ distinct experts; if $C$ is at least the number of distinct experts used in the last $\Delta$ steps, all of them are95resident, so a miss at step $t$ is precisely an expert unused during those $\Delta$ steps.96\end{proof}97The proposition is elementary but useful: it says that the hit rate of LRU at capacity $C$ is the fraction of a step's98experts that recur within the $C$-recent window of that layer, a quantity we measure directly from traces99(Section~\ref{sec:results-locality}) and which, for OLMoE, is 0.57 at $C = 16$ and 0.80 at $C = 32$.100
Files · d3deca489711afcaa521d6c5c41bc20dea077b28
Menu

Theme