论文
随笔
Tags
Slides
⤴
Search
## Pushing the Limits of Low-Bit Optimizers with a Focus on EMA Dynamics
<h3 id="background">Background</h3> <p>↗️ 模型大小飞速增加 vs. 硬件价格居高不下</p> <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> <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. $$<ul> <li>DeepSeek-v3 训练框架: $g \overset{\text{BF16}}{\rightarrow} m, v \overset{\text{FP32}}{\rightarrow} \theta$</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}{y_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) := y_{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> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617110019.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </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 as } \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/20250722100933.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> <p>💡 总结</p> <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"> <p> </p> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250722101123.png" alt="Image" style="max-width: 90%; height: auto;margin: 0 auto;"> </div> </div> <div class="slide-col-4"> <p> </p> <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"> <p> </p> <ul> <li> <p>随机信号:</p> <ul> <li>$X \in \mathbb{R}^{1000}$</li> <li>$Z \sim \mathcal{U}[0, 1]$</li> </ul> </li> <li> <p>Relaxed 条件:</p> </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/20250717131423.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>假设 $y_{k-1} \le x / \Delta \le y_k$:</p> $$ Q_{sr}(x) := \left \{ \begin{array}{ll} k-1 & w.p. \quad \frac{y_k - x / \Delta}{ y_k - y_{k-1}}, \\ k & w.p. \quad \frac{x / \Delta - y_{k-1}}{ y_k - y_{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/20250722101414.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 - y_k \big|, \end{array} $$<ul> <li>3-bit quantization levels:</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> <p>✅ Easy to implement</p> <p>✅ State decay alignment</p> <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>❎ <span style="color: gray">Singal Swamping</span></p> <p>☑️ <strong>Sign representation</strong></p> <p>☑️ <strong>Descent direction</strong></p> <p>💡 总结:</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"> <p> </p> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250722101739.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div> </div> <div class="slide-col-4"> <p> </p> <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"> worse </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/20250617111401.png" alt="Image" style="max-width: 100%; height: auto;margin: 0 auto;"> </div> <p>😒 <strong>传统方法:</strong> $\underset{\text{Training from scratch}}{\xrightarrow{\text{Ultra-Low-Bit}}}$ degeneration/collapse</p> <p>😊 <strong>SOLO:</strong> Robust to bits/tasks/models</p>
<h3 id="experiments-giant-models">Experiments (Giant Models)</h3> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617112643.png" alt="Image" style="max-width: 95%; 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="state-changes">State Changes</h3> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617112843.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="generalizability-of-solo">Generalizability of SOLO</h3> <ul> <li>AdaBelief</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617113113.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div> <ul> <li>Larger-scale models:</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617113029.png" alt="Image" style="max-width: 100%; height: auto;margin: 0 auto;"> </div>
Thanks!