Scaling Up On-Device LLMs via Active-Weight Swapping Between DRAM and Flash

Intensive Reading

01Author Info#

• About Ting Cao - Dr. Ting Cao: A Professor at the Institute of AI Industry Research (AIR), Tsinghua University.

02Background#

03Challenges#

• 挑战一：如何准确地识别出下一步计算到底需要哪些“活跃权重” 。如果识别错误，会降低模型的准确度。
• 挑战二：如何能足够早地预测出需要的活跃权重，从而将缓慢的闪存加载过程与当前的计算过程并行处理，以隐藏延迟。
• 现有的一些方法依赖 ReLU 激活函数来预测稀疏性，但这不适用于 Llama 等为追求高精度而未使用 ReLU 的现代 LLM.

04Insights#

利用了 Top-K 的稀疏性，实现了在非 ReLu 上的权重值预测和预取

论文提出了两个核心观察：

• Similarities in Cross-Layer Activations

  • The input activations of the attention and MLP blocks in LLMs exhibit high cross-layer similarity due to residuals to the input activations.
  • 由于激活值相似度很高，所以用当前层最重要的 K 个激活通道去预测下一层最重要的 K 个激活通道，准确度也很高

    

• Contextual Hot Active Weights During Decoding

  • Contextual active weights exhibit high temporal locality across inference iterations during decoding.
  • 在一个具体的对话或任务中（上下文层面），“热点权重”的重复使用率，远高于在所有通用任务中（任务层面）的平均重复使用率
  • 所以根据上下文的激活频率设计缓存会更有效（缓存命中率会更高）

05Approaches#

5.1Cross-Layer Active Weight Preloading#

当计算第 N 层时，ActiveFlow 利用第 N 层的激活值来预测并提前加载第 N+1 层到第 N+k 层（一个“层组”）所需要的活跃权重

和 LLM in a flash 的观察一样，IO 每次读的块越大，IO 吞吐量越高，所以 ActiveFlow 为了让预加载更高效，它还对存储在闪存中的权重布局进行了重排，将多个层中相同 channel 的权重存储在一起。这样，一次I/O操作就可以读取多个层的同一部分权重，显著增大了数据块大小，从而最大化闪存的读取带宽。

5.2Active Weight Swapping Pipeline#

The pipeline consists of four main operations:

• Computing (C) – Performs the required computations.
• Top-k (T) - Extracts the top-k mask from activations to determine the indices of the activated weight channels.
• On-demand loading (L) – Loads weights for the current layer group.
• Preloading (PL) – Preloads weights for the next layer group.

空间上，它将完整的、经过重排的模型权重存放在大容量但慢速的闪存中，只将计算必需的活跃权重和热点权重放在高速但容量有限的DRAM中。

时间上，它通过“跨层预加载”技术，让负责I/O的“加载线程”和负责计算的“计算线程”协同工作，实现了“计算当前层组，同时预加载下一层组”的高效循环，从而最大程度地隐藏了从闪存读取数据所带来的延迟。

4.1 部分的数学建模跳过

ActiveFlow 通过上下文的激活频率设计了缓存

5.3Self-Distillation for Top-K Sparse LLM#

Self-distillation uses the full model’s output distribution as a soft target to preserve essential distributional details and weight correlations.

对于反向传播时遇到的挑战，ActiveFlow 使用了两个技术：

• 梯度直通估计 (Gradient Straight-Through Estimation, STE)：为了解决稀疏化导致的大量梯度为零、模型难以训练的问题，该方法在反向传播时绕过了稀疏掩码（mask）的梯度清零效应，确保模型参数能得到有效更新。

• 组合损失函数: 使用了KL散度（Kullback-Leibler Divergence）和交叉熵（Cross-Entropy）组合的损失函数，以更好地对齐稀疏模型和全量模型的输出。

06Evaluation#

07Thoughts#

7.1When Reading#

这篇文章每张图都画得很好，干净简洁，一看就懂，值得学习！

再温故一下 contextual sparsity 的定义

再温故一下 hot active weights 中 hot 和 active 的含义

contextual sparsity 强调 active: small, input-dependent sets of attention heads and MLP parameters that lead to (approximately) the same output as the full model for an input

• ☐ contextual sparsity 和 power law distribution 强调的不同，以及对应的优化策略方向？

08Related Works#

• ToRead: Training-Free Activation Sparsity in Large Language Models.