Memory-efficient optimization methods have recently gained increasing attention for scaling full-parameter training of large language models under the GPU-memory bottleneck. Existing approaches either lack clear convergence guarantees, or only achieve the standard O(ε−4) iteration complexity in the nonconvex settings. We propose Omni-Masked Gradient Descent (OMGD), an optimization method based on mask traversal for memory efficient training, and provide a nonconvex convergence analysis that establishes a strictly improved iteration complexity of O~(ε−3) for finding an ε-approximate stationary point. Empirically, OMGD is a lightweight, plug-and-play approach that integrates seamlessly into most mainstream optimizers, yielding consistent improvements over competitive baselines in both fine-tuning and pre-training tasks.
Figures & tables
Algorithm
Iter. complexity
Nonconvex
Convex
SGD ( Ghadimi and Lan, 2013 )
O(ϵ−4)
O(ϵ−2)
RR-SGD ( Nguyen et al., 2021 )
O~(ϵ−3)
O~(ϵ−1)
RRM-SGD ( Qiu and Milzarek, 2024 )
O~(ϵ−3)
-
GoLore ( He et al., 2024 )
O(ϵ−4)
-
LISA ( Pan et al., 2024 )
-
O(ϵ−2)
Table 1: Comparison of methods by iteration complexity.
Figure 1: Illustration of the epochwise mask application in OMGD with d=8,M=4,N=4 . ① Generated masks satisfy the condition given in ( 3 ). ② Outer loop processes the M masks sequentially, corresponding to M consecutive epochs within one cycle. ③ Inner loop performs a full dataset pass for each mask, computing the masked gradient defined in ( 4 ) to update the model parameters.
z(i)
i th data point
zt
Sampled data point at step t
S(j)
j th mask
St
Sampled mask at step t
×
Cartesian product
[M]
{1,…,M}
⊙
Hadamard product
⌈⋅⌉
Ceiling function
Table 2: Notation list.
Figure 2: Squared L2 norm of overall error, decay term, data-reshuffle term, and compression-error term.
Algorithm
CoLA
STS-B
MRPC
RTE
SST2
MNLI
QNLI
QQP
Avg
AdamW (Full params)
64.16
90.81
92.07
80.51
94.84
87.97
92.93
89.12
86.55
GoLore ( He et al., 2024 )
62.62
90.49
91.95
78.70
94.72
87.33
92.35
87.83
85.75
SIFT ( Song et al., 2023 )
62.39
90.28
92.73
77.98
95.18
87.40
92.59
88.72
85.91
LISA ( Pan et al., 2024 )
61.76
90.19
92.25
78.34
94.50
87.54
92.68
88.77
85.75
LISA-scale
61.51
90.20
91.91
76.17
94.27
87.55
92.71
88.81
85.39
LISA-wor-no-scale
62.35
90.45
92.36
78.34
94.84
87.55
92.59
88.73
85.90
Table 3: Fine-tuning results on GLUE benchmark using pre-trained RoBERTa-Base.
Figure 6
Figure 7: Training loss of fine-tuning RoBERTa-Base on CoLA.
Table 6: Fine-tune RoBERTa-Base on the CoLA dataset, evaluated by Matthews Correlation Coefficient (MCC). Ablations on sampling period K and sampling layers γ .
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
CoLA
STS-B
MRPC
RTE
SST2
MNLI
QNLI
QQP
Batch Size
32
16
16
32
16
16
16
64
Learning Rate ( ×10−5 )
2.5
2.0
3.5
1.0
1.0
1.0
1.0
1.0
Max Seq. Len.
512
512
512
512
512
512
512
512
Appendix
Table 7: Hyperparameters used in fine-tuning pre-trained RoBERTa-Base model on the GLUE benchmark.
Method
Model
Gradients
Optimizer
Others
Total
Full params
12.55
12.55
25.10
14.66
64.86
GaLore/GoLore
12.55
12.55
1.73
4.40
31.23
LISA/LISA-wor
12.55
1.24
2.48
3.29
19.56
Appendix
Table 8: Memory breakdown (GB) for different training methods when pre-training LLaMA-7B on C4 dataset.