论文
随笔
Tags
Slides
⤴
Search
## Pushing the Limits of Low-Bit Optimizers with a Focus on EMA Dynamics
<h3 id="background">Background</h3> <ul> <li>模型大小飞速增加 vs. 硬件价格居高不下</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312203012.png" alt="Image" style="max-width: 65%; height: auto; margin: 0 auto;"> </div> <ul> <li>解决方案: <ul> <li>MoE, LoRA; ZeRO, FSDP;</li> <li>Network Quantization; <span style="color: red;">Lightweight Optimizers</span></li> </ul> </li> </ul>
<h3 id="background-1">Background</h3> <ul> <li> <p>Optimizer States (2x model size):</p> $$ m_{t+1} \leftarrow \beta_1 \cdot m_t + (1 - \beta_1) \cdot g, \\ v_{t+1} \leftarrow \beta_2 \cdot v_t + (1 - \beta_2) \cdot g^2. $$</li> <li> <p>DeepSeek-v3 训练框架: $g \overset{\text{BF16}}{\rightarrow} m, v \overset{\text{FP32}}{\rightarrow} \theta$</p> </li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312204230.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="quantization-and-dequantization">Quantization and Dequantization</h3> <ul> <li> <p>Quantization:</p> $$ q = Q(x) := \mathop{\text{argmin}} \limits_{k=0}^{2^b - 1} \big|\frac{x}{\textcolor{red}{\Delta}} - \textcolor{red}{\iota_k} \big|. $$</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312205652.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div> <ul> <li> <p>Dequantization:</p> $$ \tilde{x} = Q^{\dagger}(q) := \iota_{q} \cdot \Delta. $$<p>
<h3 id="stateful-optimizers-in-ultra-low-bits"><u>S</u>tateful <u>O</u>ptimizers in Ultra <u>LO</u>w Bits</h3> <ul> <li>Low-Bitwidth EMA update:</li> </ul> $$ \begin{array}{rl} \text{Dequantization: } & \tilde{x}_t = Q^{\dagger}(q_t) = \iota_{q_t} \cdot \Delta_t, \\ \text{EMA update: } & \hat{x}_{t+1} \leftarrow \beta \cdot \tilde{x}_t + (1 - \beta) \cdot z_{t + 1}, \\ \text{Quantization: } & q_{t+1} = Q(\hat{x}_{t+1}). \end{array} $$<div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">Dettmers T., et al. 8-bit Optimizers via Block-wise Quantization. ICLR, 2022.</p> <p style="margin: 2px 0;">Li B., et al. Memory Efficient Optimizers with 4-bit States. NeurIPS, 2023.</p> </div>
<h3 id="quantization-for-unsigned-ema-update">Quantization for Unsigned EMA Update</h3> <ul> <li><em>Signal Swamping</em> (<u>large-to-small number addition</u>)</li> </ul> $$ \text{EMA update: } \hat{x}_{t+1} \leftarrow \beta \cdot \tilde{x}_t + \underbrace{\textcolor{red}{(1 - \beta) \cdot z_{t + 1}}}_{\text{very small for } \beta \rightarrow 1}. $$<div style="text-align: center; margin-top: 50px; margin-bottom: -80px; padding: 0"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312211840.png" alt="Image" style="max-width: 70%; height: auto;margin: 0 auto;"> </div> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">Higham N. J. The Accuracy of Floating Point Summation. SIAM Journal on Scientific Computing. 1993.</p> </div>
<h3 id="signal-swamping">Signal Swamping</h3> <ul> <li>总结</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312212039.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="case-study">Case Study</h3> <div class="slide-cols"> <div class="slide-col-6"> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312212821.png" alt="Image" style="max-width: 90%; height: auto;margin: 0 auto;"> </div> </div> <div class="slide-col-4"> <ul> <li> <p>一定<span style="color: red">条件</span>下:</p> <ul> <li>Linear 下全部不更新</li> <li>DE 下部分更新</li> </ul> </li> <li> <p>实际上 $\beta \ge 0.9$ 为<span style="color: red">相当常见的 setting</span></p> </li> </ul> </div> </div>
<h3 id="case-study-1">Case Study</h3> <div class="slide-cols"> <div class="slide-col-4"> <ul> <li>$X \in \mathbb{R}^{1000}$</li> <li>$Z \sim \mathcal{U}[0, 1]$</li> </ul> <p><span style="color: red;"><em>X</em></span> Fixed $\Delta$</p> <p><span style="color: red;"><em>X</em></span> $z \le \Delta$</p> <ul> <li>理论收敛至: 0.5</li> </ul> </div> <div class="slide-col-6"> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312213810.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div> </div> </div>
<h3 id="stochastic-rounding">Stochastic Rounding</h3> <ul> <li> <p>假设 $\iota_{k-1} \le x / \Delta \le \iota_k$:</p> $$ Q_{sr}(x) := \left \{ \begin{array}{ll} k-1 & w.p. \quad \frac{\iota_k - x / \Delta}{ \iota_k - \iota_{k-1}}, \\ k & w.p. \quad \frac{x / \Delta - \iota_{k-1}}{ \iota_k - \iota_{k-1}}. \end{array} \right . $$</li> <li> <p>High variance:</p> </li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250313112908.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="logarithmic-quantization">Logarithmic Quantization</h3> $$ \begin{array}{ll} Q(x) &=\text{Clip}(\lfloor \log_{\alpha} \frac{x}{\Delta} + \xi \rceil; 0, 2^b - 1) \\ &\approx \mathop{\text{argmin}} \limits_{k=0}^{2^b - 1} \big|\frac{x}{\Delta} \cdot \alpha^\xi - \iota_k \big|, \end{array} $$<ul> <li>3-bit quantization anchors:</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250313113440.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="logarithmic-quantization-1">Logarithmic Quantization</h3> <ul> <li>2-bit quantization illustration</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250313113535.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="logarithmic-quantization-2">Logarithmic Quantization</h3> <ul> <li> <p>Easy to implement</p> </li> <li> <p>State Decay Alignment</p> </li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250313115306.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="quantization-for-signed-ema-update">Quantization for Signed EMA Update</h3> <p><strong>X</strong> <span style="color: gray">Singal Swamping</span></p> <p><strong>✓</strong> <strong>Sign representation</strong></p> <p><strong>✓</strong> <strong>Descent direction</strong></p> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250314115701.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="theoretical-analysis">Theoretical Analysis</h3> <div class="slide-cols"> <div class="slide-col-6"> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250314115959.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div> </div> <div class="slide-col-4"> <p>$\rightarrow$ <span style="color: red">Bits $\downarrow$</span> or <span style="color: red">$\beta \uparrow$ </span></p> <p>$\rightarrow$ Quantization errors <span style="color: red">$\uparrow$</span></p> <p>$\rightarrow$ gradient variance <span style="color: red"> $\uparrow$ </span></p> <p>$\rightarrow$ <span style="color: red"> bad </span> convergence</p> </div> </div>
<h3 id="momentum-adjustment">Momentum Adjustment</h3> <ul> <li><strong>方差控制:</strong> 选择 $\beta'$ 满足:</li> </ul> $$ \underbrace{\frac{\textcolor{gray}{\beta'}}{1 - \textcolor{gray}{\beta'}} r_{\text{median}}(b')}_{\textcolor{gray}{\text{undetermined}}} \le \underbrace{\frac{\beta}{1 - \beta} r_{\text{median}}(b)}_{\textcolor{green}{\text{valid setup}}}. $$<ul> <li><strong>查表:</strong> (<u>灰色区域代表了经验可行的参数推荐</u>)</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250314121510.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div>
<h3 id="experiments">Experiments</h3> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250417145811.png" alt="Image" style="max-width: 85%; height: auto;margin: 0 auto;"> </div>
<h3 id="loss">Loss</h3> <ul> <li>损失正常收敛</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250319170139.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div>
<h3 id="quantile">Quantile $x_p$</h3> <ul> <li>基本上 $p \in [0.05, 0.3]$ 都有不错的性能</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250319170604.png" alt="Image" style="max-width: 55%; height: auto;margin: 0 auto;"> </div>
<h3 id="beta-block-size">Beta, Block size</h3> <ul> <li>损失正常收敛</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250407200935.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div>
<h3 id="2nd-state-distribution">2nd State Distribution</h3> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250407201114.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="adabelief">AdaBelief</h3> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250407202155.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
Thanks!