The substantial memory demands of pre-training and fine-tuning large language models (LLMs) require memory-efficient optimization algorithms. One promising approach is layer-wise optimization, which treats each transformer block as a single layer and optimizes it sequentially, while freezing the other layers to save optimizer states and activations. Although effective, these methods ignore the varying importance of the modules within each layer, leading to suboptimal performance. Moreover, layer-wise sampling provides only limited memory savings, as at least one full layer must remain active during optimization. To overcome these limitations, we propose Module-wise Importance SAmpling (MISA), a novel method that divides each layer into smaller modules and assigns importance scores to each module. MISA uses a weighted random sampling mechanism to activate modules, provably reducing gradient variance compared to layer-wise sampling. Additionally, we establish an \(\mathcal{O}(1/\sqrt{K})\) convergence rate under non-convex and stochastic conditions, where $K$ is the total number of block updates, and provide a detailed memory analysis showcasing MISA's superiority over existing baseline methods. Experiments on diverse learning tasks validate the effectiveness of MISA. Source code is available at https://github.com/pkumelon/MISA.

翻译：预训练与微调大语言模型（LLMs）所需的大量内存资源，催生了对内存高效优化算法的需求。层间优化是一种前景广阔的方法，它将每个Transformer块视为独立层进行顺序优化，同时冻结其他层以节省优化器状态和激活值内存。尽管该方法行之有效，但其忽略了层内各模块的重要性差异，导致性能未能达到最优。此外，层间采样仅能提供有限的内存节省，因为优化过程中至少需要保持一个完整层的激活状态。为突破这些限制，我们提出模块重要性采样（MISA）这一创新方法，将每个层划分为更小的模块并为每个模块分配重要性分数。MISA采用加权随机采样机制激活模块，理论上相比层间采样能有效降低梯度方差。此外，我们在非凸随机条件下建立了\(\mathcal{O}(1/\sqrt{K})\)收敛率（其中$K$为块更新总次数），并通过详细的内存分析证明了MISA相对于现有基线方法的优越性。在多类学习任务上的实验验证了MISA的有效性。源代码发布于https://github.com/pkumelon/MISA。