Scalable Zeroth-Order Fine-Tuning for Extremely Large Language Models with Limited GPU Memory
Authors: Liangyu Wang
Jie Ren
Hang Xu
Junxiao Wang
Huanyi Xie
David E. Keyes
Di Wang
Venue: COLM 2025 (2025)
Keywords: Zeroth-Order Optimization
LLMs
Fine-Tuning
TL;DR: ZO2 is a memory-efficient framework that enables zeroth-order fine-tuning of large language models like OPT-175B on a single 18GB GPU.

Abstract: Fine-tuning large pre-trained LLMs generally demands extensive GPU memory. Traditional first-order optimizers like SGD encounter substantial difficulties due to increased memory requirements from storing activations and gradients during both the forward and backward phases as the model size expands. Alternatively, zeroth-order (ZO) techniques can compute gradients using just forward operations, eliminating the need to store activations. Furthermore, by leveraging CPU capabilities, it's feasible to enhance both the memory and processing power available to a single GPU.
We propose a novel framework, ZO2 (Zeroth-Order Offloading), for efficient zeroth-order fine-tuning of LLMs with only limited GPU memory. Our framework dynamically shifts model parameters between the CPU and GPU as required, optimizing computation flow and maximizing GPU usage by minimizing downtime. This integration of parameter adjustments with ZO's double forward operations reduces unnecessary data movement, enhancing the fine-tuning efficacy. Additionally, our framework supports an innovative low-bit precision approach in AMP (Automatic Mixed Precision) mode to streamline data exchanges between the CPU and GPU.
Employing this approach allows us to fine-tune extraordinarily large models, such as the OPT-175B with 175 billion parameters, on a mere 18GB GPU. Moreover, our framework achieves these results with almost no additional time overhead and absolutely no accuracy loss compared to standard zeroth-order methods. ZO2's code has been open-sourced in https://github.com/liangyuwang/zo2.

Published as a conference paper at COLM 2025
Scalable Zeroth-Order Fine-Tuning for Extremely Large Lan- guage Models with Limited GPU Memory
Liangyu Wang', Jie Ren', Hang Xu', Junxiao Wang', Huanyi Xie David E. Keyes', Di Wang
King Abdullah University of Science and Technology (KAUST) 2Guangzhou University
ekasy. edw.sstttedu..... edu.
Abstract
Fine-tuning large pre-trained LLMs generally demands extensive GPU memory. Traditional first-order optimizers like SGD encounter substantial difficulties due to increased memory requirements from storing activations and gradients during both the forward and backward phases as the model size expands. Alternatively, zeroth-order (ZO) techniques can compute gradients using just forward operations, eliminating the need to store ac- tivations. Furthermore, by leveraging CPU capabilities, it's feasible enhance both the memory and processing power a available to a single GPU. We propose novel framework, ZO2 (Zeroth-Order Offloading), for effi- cient zeroth-order fine-tuning of LLMs with only limited GPU memory. Our framework dynamically shifts model parameters between the CPÚ and GPU required, optimízing computation flow and maximizing GPU usage by minimizing downtime. This integration of parameter adjustments with ZO's double forward operations reduces unnecessary data move- ment, enhancing the fine-tuning efficacy. Additionally, our framework supports innovative low-bit precision approach in AMP (Automatic Mixed Precision) mode to streamline data exchanges between the CPU and GPU. Employing this approach allows large models, such as the OPT-175B with 175 billion parameters, on a mere 18GB GPU. Moreover, our framework achieves these results with almost no additional time overhead and absolutely no accuracy loss compared standard zeroth-order methods. ZO2's code has been open-sourced in http:://gitluub.oo///innggg//
fine-tune extraordinarily
1 Introduction
As Large Language Models (LLMs) grow in scale-reaching hundreds of billions of pa- rameters (like OPT-175B (Zhang et al., 2022), Llama 3.1 405B (Dubey et al., 2024)) and their applications diversify (Pei et al., 2024; Huang et al., 2025; Pei et al., 2025), managing GPU memory resources effectively becomes crucial. Efficient GPU memory management is crucial not only because it directly influences model performance and training speed, but also because GPU memory is both expensive and limited in quantity. However, this creates significant challenge in handling ever-larger models within the physical constraints of current hardware technologies. CPU offloading has become a crucial technique for over- coming the challenge. It involves transferring computations and data from the GPU to the CPU, specifically targeting data or parameters that are less frequently accessed ("inactive"). Specifically, it leverages the typically larger and more cost-effective CPU memory (DDR SDRAM) compared to the more expensive and less abundant GPU memory (HBM). By offloading these inactive tensors of the neural network, CPU offloading effectively alleviates the memory and computational pressures on GPUs. While CPU offloading has been com- monly applied in inference to manage memory-intensive tasks like KV cache offloading (Ge

---

Published as a conference paper at COLM 2025
et al., 2023; Sheng et al., 2023) and Mixture of Experts (MoE) offloading (Eliseev & Mazur, 2023; Xue et al., 2024), its application in training, especially fine-tuning (Ouyang et al., 2022; Hu et al., 2022; Shao et al., 2024; Zheng et al., 2025; Wang et al., 2025), remains less explored.
Recently, some works (Rajbhandari et al., 2020; Ren et al., 2021) have tried to introduce CPU offloading into LLM training. However, they are typically constrained by the capabilities of first-order optimizers such as SGD and Adaptive Moment Estimation (AdamW) (Loshchilov & Hutter, 2017), and limited GPU memory, restricting large-scale model scalability on single GPU systems. In detail, using first-order optimizers introduces two major inefficiencies in CPU offloading (Section 4.1): (1) Multiple communication operations: During the training of LLMs, parameters are used not only for computing the loss during the forward pass but also for gradient computation in the backward pass. This necessitates offloading the same data (parameter) twice-one for each pass (see Figure 2a for an illustration). Such redundancy not only doubles the communication volume between the CPU and GPU but also introduces significant latency and inefficiency due to repetitive data transfers. (2) Huge data transfer volume per communication operation: Furthermore, both parameters and activations (hidden states) are required in the backward pass to complete gradient computations. This means that parameters and activation values must be offloaded during each forward pass and re-uploaded to the GPU for the backward pass. The result is significant increase in the volume of data transferred, which severely impacts training throughput and efficiency.
On the other hand, compared to first-order optimization methods, zeroth-order (ZO) meth- ods offer a novel approach to fine-tuning LLMs (Zhang et al., 2024; Malladi et al., 2023; Gautam et al., 2024). These methods utilize dual-forward passes to estimate parameter gradients and subsequently update parameters, as illustrated in Figure 2b. This approach eliminates the traditional reliance on backward passes, thereby streamlining the training process by significantly reducing the number of computational steps required.
Based on the above observations, we conjecture that ZO's architecture is particularly well- suited for CPU offloading strategies. Intuitively, by eliminating backward passes and the need to store activation values, it can significantly reduce GPU memory demands through efficient parameter offloading. However, despite these advantages, 2o training via CPU offloading introduces new challenges, particularly in the realm of CPU-to-GPU communication.
One challenge lies in ensuring the alignment of random perturbations used in duwl-forward passes across different transformer blocks. In contrast to the monolithic forward pass in MeZO, our block-wise execution requires careful management of the random number generator (RNG) state to preserve gradient correctness. Without this mechanism, ZO training may suffer from accuracy mismatch.
Another key challenge is designing an efficient communictioo-compuaatioo scheduling mechanism. First-order offloading systems (Ren et al., 2021; Rajbhandari et al., 2021) and inference pipelines (Sheng et al., 2023) require intricate and often brittle scheduler logic to pipeline forward/backwart steps with parameter movement. In contrast, we find that the dual-forward nature of ZO optimization naturally doubles computation time, enabling the communication be fully hidden behind computation in most cases. This allows our dynamic scheduler to be extremely simple yet highly effective, requiring only three CUDA streams with minimal control logic.
To tackle these challenges, we introduce ZO2 (Zeroth-Order Offloading), a novel framework specifically designed for zeroth-order fine-tuning of LLMSs under tight memory constraints. Our design (1) preserves the mathematical integrity of MeZO by introducing a robust RNG state manager, and (2) achieves high throughput via a lightweight dynamic scheduler that overlaps parameter uploads, durl-forward computations, and offloads. In addition to the scheduler and RNG state management, ZO2 incorporates several complementary system optimizations to further enhance throughput and memory efficiency. These include a reusable memory strategy that eliminates repeated CUDA allocations, a reordering of parameter update operations minimize data transfers, and AMP- MoP-mode-awwan parameter
2

---

Published as a conference paper at COLM 2025
compression that reduces communication overhead during mixed-precision training. Our contributions can be summarized as follows:
Innovative use of CPU-offloading for ZO methods: We determine that ZO is inherently more suitable for CPU offloading than first-order optimizers, due to its forward-only execution and minimal activation dependencies. By combining ZO with CPU-offloading, ZO2 enables efficient training of models like OPT-175B on single 18GB GPU.
Completely lossless accuracy, high-throughput core design: (1) We introduce an RNG state manager ensure accuracy alignment between perturbation and param- eter update phases during block-wise ZÖ training. (2) We propose a lightweight dynamic scheduler that overlaps parameter uploading, dual-forward computation, and offloading using only three CUDA streams. Compared to other complex sched- uler logic design, 7O2 achieves simpler and more robust execution with minimal scheduling complexity.
Complementary system-level optimizations: ZO2 integrates additional key im- provements maximize throughput and minimize overhead: (1) Reusable memory allocation to avoid costly CUDA malloc/ free during block transfers (Section D.2). (2) Efficient parameter update reordering to reduce redundant data transfers by fusing updates into the forward phase (Section D.3). (3) AMP-mode-aware com- pression that reduces CPU-GPU communication volume via low-bit representations like FP16 and FP8 (Section D.4).
Empirical validation at scale: ZO2 is able to fine-tune OPT-175B using just 18GB of GPU memory, with no additional runtime or accuracy degradation. Our abla- tion study confirms the contribution of each optimization component to overall throughput.
Open-source implementation with a clean and extensible API. We release the full implementation of ZO2 with lightweight and modular codebase. Our API (Appendix E) is intentionally designed be simple and minimal, enabling users to easily plug ZO2 into existing Py Torch training workflows.
2 Related Work
Zeroth-Order (ZO) Optimization. ZO optimization offers gradient-free alternative to first-order (FO) optimization by approximating gradients through function value-based estimates. These estimates theoretically require only two forward passes but are believed to be prohibitively slow for optimizing large models. Despite this limitation, ZO methods have been utilized in deep learning to generate adversarial examples or adjust input embeddings (Sun et al., 2022a;b), though they have not been widely adopted for direct optimization of large-scale models (Liu et al., 2020). Several acceleration techniques have been proposed to address the scaling challenges of ZO optimization and some of them have been used for LLM fine-tuning. These include using historical data to improve gradient estimators (Cheng et al., 2021), exploiting gradient structures (Singhal et al., 2023) or sparsity to reduce the dependence of ZO methods on the size the problem (Chen et al., 2024; Cai et al., 2022; 2021), and reusing intermediate features (Chen et al., 2024) and random perturbation vectors (Malladi et al., 2023) during the optimization process. These advancements suggest that ZO optimization could increasingly be applied to more complex and large-scale ML problems. While previous ZO optimization efforts have primarily targeted algorithmic improvements for GPU memory efficiency, our approach extends these optimizations to the system level, enabling more robust memory management and enhanced performance for large-scale machine learning applications. More related work can be found in Appendix B.
3 Preliminaries on Zeroth-Order Optimization
Zeroth-order (ZO) optimization provides a gradient-fres alternative to first-order (FO) methods, estimating gradients using only function evaluations. In this paper, focus
3

---

Published as a conference paper at COLM 2025
the Randomized Gradient Estimator (RGE) (Nesterov & Spokoiny, 2017), which forms the foundation MeZO (Malladi et al., 2023) and our work.
RGE approximates the gradient VL(0) by evaluating the loss function along randomly sampled direction NO, 1). The projected gradient is defined as:
L(e Ez) I( EZ)
g
E ER
(1)
2
and the parameter update is given by:
0 - 0 17-8.2
(2)
where is small smoothing parameter and 17 is the learning rate.
Importantly, ZO methods avoid the need for backward computation and activation storage by relying solely duwl-forward passes. This forward-only nature, combined with the fact that is scalar, makes ZO highly memory-efficient and well-suited for offloading- -aaenng training. For full derivation and the complete MeZO algorithm, please refer to Appendix C.
4 Framework Overview
Dual forward (Block i, Epoch
Input
Set RNG state: S
Upload
Input
Weight W-i--
Projected Gradient j-1
Update
Block
Embedding
Weight wi,
Block
Dual forward
Set RNG state: Sy
O
Offoad
Perturb Parameters :0 = +
Dual forward
Block
+
forward(8, output
output,
Block n
Repeat
Block
Set RNG state: S,
Dual forward
Perturb Parameters 2ez
LM head
Block
forward output
output;
Get RNG state: S
Projected Gradient
Perturb Parameters :
Time
CPU Upload
Compute
Stream
Figure 1: Workflow of the ZO2 framework for fine-tuning LLMs.
4.1 Insight
Our framework, ZO2, exclusively designed for zeroth-order optimizer, incorporates a CPU offloading strategy optimized for this specific approach. This design leverages the CPU for storage and the GPU for computation, streamlining parameter offloading in manner uniquely suited to the operational characteristics of zeroth-order methods:
(1) Äs illustrated in Figure ((a), the first-order optimizer typically requires both forward and backward passes with opposing workflow directions, and each pass necessitates pa- rameter availability. Consequently, when parameters offloaded to the CPU are required for computation on the GPU, communication is necessary twice (forward and backward passes). In contrast, the zeroth-order optimizer necessitates only two forward passes with the same workflow direction (Figure 2(b)), allowing parameters to be reused with just one communication cycle. This modification effectively halves the communication frequency.
4

---

Published as a conference paper at COLM 2025
W;
W; freorr
W; frroor
W3 Ynwood
hoo
ax-
dIF; Iartwar
dW, boskec
dif techerd
Wy Rewert
fower
W2
Wy lnwad
(a) Model using first-order optimizer with (b) Model using zero-order optimizer with only dual-forward passes workflow
forward-bacc -bakwad passes workflow
Figure 2: Motivation. Using linear operator as an example. (a) First-Order Optimizer: Requires forward and backward passes, storing intermediate activations X2) and incurring repeated CPU-GPU parameter transfers. (b) Zeroth-Order Optimizer: Uses dual-forward passes with perturbed weights W,) to estimate gradients. No activation storage is needed, and each parameter is transferred only once, reducing memory and communication overhead.
(2) As discussed in Section 5.2, it is essential to overlap communication and compu- tation tasks. Normally, communication between the CPU and GPU consumes signifi- cantly more time than GPU computations. The zeroth-order optimizer, with its dual- forward passes in the same direction, maintains the duration of communication while increasing computation time. This setup effectively reduces the overall communication overhead by half and reduce the complexity of the scheduler.
(3) The zeroth-order optimizer does not involve a backward pass, eliminating the need to store and offload activations. This feature significantly reduces memory requirements.
(4) The gradient calculation the zeroth-order optimizer is achieved by multiplying a projected gradient with Gaussian distribution. The projected gradient is single value (Equation 1), and the Gaussian distribution can be generated using seed. Consequently, the actual gradient can be computed in-place during parameter updates, eliminating the need for dedicated storage space for real gradients.
4.2 ZO2 Framework Overview
In this part, we first provide an overview and brief introduction to our ZO2 (Zeroth-Order Offloading) framework. To better illustrate our idea, we first describe the computation workflow of the original ZO optimization procedure for LLM fine-tuning, Consider the architecture simple decoder-only LLM,, which typically comprises an embedding layer, N transformer blocks, and a language model (LM) head layer. Our offloading strategy involves offloading all transformer blocks to the CPU, while retaining the remaining compo- nents on the GPU. This approach is similarly applicable to more complex LLM architectures like OPT (Zhang et al., 2022). Initially, input data is loaded from the disk into the CPU and subsequently transferred to the GPU. Within the GPU, each ding layer, transformer blocks, and the language model (LM) head-executes dual-forward computations sequentially to estimate the projected gradient and update parameters.
including the embed-
The above approach divides the entire model's dual-forward process, as outlined in the original ZO workflow (Algorithm 3), into discrete block-level operations. However, this could potentially introduce an accuracy mismatch. This challenge arises because the core of the ZO method relies on the uniformity of Gaussian random vectors applied during both the perturbation and the parameter update phases, which must be consistently synchronized across all computations (Algorithm 3). To address this, we propose random eendom number generator (RNG) state manager (Section 5.1) that meticulously aligns the random vectors. This management ensures that the random perturbations and the subsequent parameter updates across different transformer blocks maintain identical stochastic characteristics.
From the efficiency perspective, naive implementation with deep learning frameworks like PyTorch (Paszke et al., 2019) typically manage both communication (via interconnections, e.g., PCIe) and computation tasks with a single CUDA stream, leading to significant ineffi- ciencies. Specifically, for ZO optimization, the i-th transformer block is uploaded from the tasks using its CUDA and Tensor Cores, and the CPU memory is used for parameter storage), undergoes dual-forward
CPU to the GPU (the GPU is designated for computation-intent
5

---

Published as a conference paper at COLM 2025
computation, and then is offloaded back to the CPU. The i 1-th block must wait for the offloading of the i-th block to finish before its uploading, leading to idle CUDA and Tensor Cores during communication while the interconnection remains idle during computation. See Figure 5 for an illustration. Our ZO2 framework implementation achieves the strategic utilization of CPU and GPU resources (Section 5.2). This approach involves dynamically offloading model parameters to the CPU and uploading them back to the GPU as needed for computation. Specifically, for the transformer model structure, each transformer block is individually uploaded for processing and subsequently offloaded post-computation, thus balancing communication and computation across blocks. As illustrated in Figure while the i-th transformer block is being computed, the i + 1-th block is pre-uploaded, and the i- 1-th block offloaded simultaneously. This strategic overlapping ensures continuous and efficient computation, reducing idle times and maximizing GPU utilization. In the uploading phase ZO2, transformer blocks are transferred into a reusable memory space on the GPU, eliminating the extra time typically required for CUDA memory allocation (Section D.2). Moreover, parameter updates are ingeniously fused with the dual-forward passes minimize redundant data transfers, thereby enhancing the overall efficiency of the model training process (Section D.3).
Algorithm 1 ZO2 Computation with RNG State Manager
Require:
Transformer blocks (W;), with number of transformer blocks N, embedding parameters Embheddigg, and LM head LMhead, module parameter module forward function forward, loss function L, training iterations T, perturbation step size e, data batch B, learning rate //t, random seed s, random state buffer rsb, last iteration's random state /rs, function Updatel' arameters aarater'sr
and Perturbl'a Parameterss from Algorithm 3.
1: Initialize
2: for 1,,,,,, do
3: Set random seed s and sample batch D Get random state GetRngState(s), and push rs into rsh.
4:
5:
j > 1 then
6:
Update last iteration's random state Irs PopLeft(rsb)
7:
else
8:
Irs None
9:
end
out rs, Irs Dualfoorwad(mmmedddiinneeess.s. B,
11:
for t= to N do
out rs, Irs DualForwaard(,cc,7.77ssrrr out end for
13:
ouf out_, rs, lrs DualTorvardl.MMheddd..s.. out ou!-)
15:
4 = L(outt f_ -- L(out))
g (e, (2c)
17:
end for
18: function DUALFOORARRDG.CCSSTT.III.. inn.::/..............ii...iiiii inputt,,
19:
if 0 then
20:
SetRngState(s, Irs)
21:
0, <- UpdateParameters((.g)
22:
end if
23:
SetRngState(s, rs), 0 Pr
24:
out < forward (0: (0;input+)
25:
SetRngState(s, rs), 0 t PerturbParamettess(. -2c)
out_ + (f; SetRngState(s, rs), 0 Perturbaarammtees///
27:
rs GetRngState(s)
29:
return out out ITs
30:
end function
Our ZO2 framework further integrates a novel low-bit precision technique that efficiently manages data transfers between the CPU and GPU in the AMP mode (see Figure 4 for

---

Published as a conference paper at COLM 2025
illustration). This technique is aligned with AMP protocols by ensuring that high- bit precision is maintained for parameter updates, while low-bit precision data is used for computation the GPU (Section D.4). This dual-precision approach significantly reduces the communication overhead, optimizing memory usage without compromising computational accuracy.
In the following section, we will provide implementation challenges and details of our framework.
Design and Implementation Details
5.1 ZO2 with RNG State Manager
The core principle of 'O algorithms is the uniform application of Gaussian random vector for each parameter during both the perturbation and update phases. MeZO (Algorithm 3) accomplishes this by resetting the seed at each iteration to control the state the random number generator (RNG), ensuring consistent execution. However, unlike MeZO, ZO2 disaggregates the model's dual-forward process across different model blocks (Figure 1, Algorithm 1), which could potentially lead to discrepancies in the RNG states between perturbation and parameter updates.
To maintain the precision of the MeZO workflow within the ZO2 framework, we metic- ulously record the RNG states (rng_state) during each module's dual-forward operation (Algorithm 1 Line 18-30). Specifically, rng_state saved prior executing any parameter perturbations and before the module forward pass. This ensures that outputs are consis- tently generated across iterations. Additionally, given that parameters are updated using the gradient projected from the last iteration (refer to Section D.3), we preserve the last random state (last_rsfate) to accurately replicate the Gaussian perturbations that were applied during the perturbation process (Algorithm Line 4-9).
This precise synchronization of rng rng_stale and last_rslate across different model blocks in ZO2 guarantees that each parameter update adheres to the same stochastic path as established in MeZO, Consequently, this methodological rigor ensures the preservation of exact accuracy throughout the workflow, facilitating reliable and reproducible outcomes in line with the original algorithmic design.
5.2
Dynamic Scheduler Design for Efficient Overlap
Algorithm 2 ZO2 Dynamic Scheduler
Require: Transformer blocks { ( AN WW
with number of transformer blocks N, embedding parameters Embedding, and LM head LMhead.
1: Initialize a dynamic scheduler S(.} to control dual-forward computation C(-), upload- ing u(-), and offloading O(.) operations.
2: Asynchronously launch S S(C((mbedding) u(W)))
3: for 1 N do
4:
Synchronously wait until U(W,) finished.
5: if 1 then
6:
Asynchronously launch S(C(W,), u(W;
7:
else
8:
Synchronously wait until C(W-1) finished.
9:
Asynchronously launch S(O(W, 1), (W,, end if
10:
11: end for
Synchronously wait until C(WN-1) and U(Wx) finished.
13: Asynchronously launch S(O(W,-)), (-
Synchronously wait until C(WN) finished.
15: Asynchronously launch S(O(Wx), C(MMhedd)
7

---

Published as a conference paper at COLM 2025
While Section 5.1 ensures the correctness and precision alignment of zeroth-order updates by managing random number generator (RNG) states within the DualForward routine (Algo- rithm 1), this section focuses on improving efficiency by overlapping CPU-GPU communica- tion with computation. Specifically, we aim to accelerate the core task in Algorithm 1-the DualForward computation for each transformer block-by overlapping it with parameter uploads and offloads.
Figure 5 illustrates the evolution from naive sequential execution (Figure 5a), where all communication and computation are serialized, to an overlapped execution model (Figure 5b), which forms the basis of our dynamic scheduler. In this optimized setting, computation of block i (ooolodwwwrrd))) uploadingg of block i+1, and offloading of block i- are launched concurrently on three dedicated CUDA streams.
A key enabler of this simplified scheduling design is the dual-forward nature of ZO opti- mization, which inherently doubles the computation time while leaving the communication time per block unchanged. This high cumpute-to-communiic to achieve full overlap, effectively hiding communication delays. This behavior stands in sharp contrast with traditional first-order offloading training methods (Ren et al., 2021; Rajbhandari et al., 2021), where scheduling must delicately interleave forward, backward, and optimizer operations. Similarly, inference frameworks like (Sheng et al., 2023) require complex runtime stream orchestration. In contrast, 7O2 achieves efficient scheduling with minimal control logic.
ratio makes it much easier
To ensure correctness during asynchronous execution, we enforce two dependency rules: (1) Intra-block: offload of block waits for completion of compute, which waits for upload. (2) Inter-block: each upload/ upload/compute /offload task waits for the previous block's corre- sponding task. Our actual execution pattern follows the sequence O(W;-1) + C(W)) + U(W;+1), as formalized in Algorithm 2. This pattern ensures safety without requiring global synchronization, thus avoiding pipeline "bubbles".
We also retain the embedding layer and LM head on the GPU throughout. This not only avoids unnecessary transfers but also enables uploading of the first transformer block while the embedding layer is being computed-furte design section can be seen in Appendix D.I.
improving overlap efficiency. The full scheduler
Experiment
We evaluate ZO2 using the OPT model family (Zhang al., 2022), ranging from 1.3B to 175B parameters (Table 4), to assess scalability across different model sizes. Our baseline is MeZO (Malladi et al., 2023), the most memory- and throughput-effccient zeroth-order optimizer to date. ZO2 builds upon MeZO, aiming to further reduce GPU memory usage without sacrificing throughput or accuracy. More experimental setups and additional experiments can be found in Appendix F, G.
Moreexpprimenta
6.1 Main Results
Table 1: Main results of ZO2 performance various model configurations and both FP32 and FP16 modes. Instances of '-' in the table indicate scenarios where the corresponding method failed to execute due to memory constraints. The values in parentheses (x) represent the ratio of each measurement compared to the baseline MeZO (first column) configuration.
Model
McZO
GPU Memory Usage (MB)
(IT)22
MczO (FPI6
(IP16)
(F132)
Throughput (lokens/sec)
202 (1732)
(FT16)
(IP16)
OPT-1.38
8898
5098/x0.57)
581 6814(x0.65)
3750(x0.42)
1998
1955(x0.97)
66291x3.32)
6448(x3.23)
OPT-227B
14514
59300x0.41)
9054(x0.62)
4142(x0.29)
1104
1086(x0.95)
4229(x3.83)
4220kx3.2)
OPT-6.7B
32930
8420(x0.26)
16686(x0.55 50)
4992(x0 15)
492
4850x0.99
2349kx4.77)
2270(x4.61)
OPT-138
58762
10736(xt.18)
296900(00511
6180(x0.11)
266
239(x0 97)
1326165.87)
1251(x5.54)
OPT-300
15981
63896
8856
122
641
514
OPT-66B
22295
12071
273
OPT-1758
34015
15039
14
37
8

---

Published as a conference paper at COLM 2025
The performance results of our experiments are presented in Table 1, where we compare the GPU memory usage and throughput of the MeZO and ZO2 frameworks, employing both FP32 and FP16 data formats. The results demonstrate consistent advantage of ZO2 in terms of GPU memory utilization across all model sizes, highlighting significant efficiency improvements, especially in large-scale models like OPT-175B. This efficiency is attributed to 2O2's s design, which strategically utilizes GPU memory to temporarily store only a limited number of transformer blocks for computation rather than the entire model. Notably, the memory savings become more pronounced as the model size increases. For smaller models, the GPÚ memory savings are less pronounced due to the significant proportion of memory allocated for input data, which diminishes the relative impact of the memory optimization.
In terms of throughput, ZO2 maintains a performance comparable to MeZ0 in most tested scenarios without any additional time overhead. The instances where ZO2 exhibits de- crease throughput, such as with the OPT-1.3B model in FP32 format, can be primarily attributed to the dynamics of computation and communication. In these cases, the com- putation of each transformer block's dual-forward passes completes quicker than their corresponding communication tasks, leading to idle times the dynamic scheduler (dis- cussed in Section 5.2) synchronizes and waits for these communication tasks to conclude. It is important to note that our results do not show consistent pattern where either smaller or larger models benefit more significantly from the computation-come pnpuuaaiooncomunicattionn overlap, indicating that the effectiveness of this overlap does not linearly correlate with model size.
6.2 Results on Accuracy Alignment
Table 2:
Main results of ZO2 precision on OPT-13B
Method
SST-2 (%)
RTE (%)
(%)
BoolQ (%)
WSC (%)
WIC (%)
MultiRC (%)
MeZO
91.4
66.1
67.9
67.6
63.5
61.1
60.1
702
91.4
66.1
67.9
67.6
63.5
61.1
60.1
We evaluate ZO2's accuracy alignment on the OPT-13B model across seven NLP benchmarks: SST-2, RTE, CB, BoolQ, WSC, WIC, and MultiRC (Socher et al., 2013; Dagan et al., 2005; De Marneffe al., 2019; Clark et al., 2019; Levesque et al., 2012; Pilehvar & Camacho- Collados, 2018; Khashabi et al., 2018). These tasks span sentiment analysis, entailment, coreference, QA, and word sense disambiguation, covering diverse aspects of language understanding.
As shown in Table 2, ZO2 achieves identical precision rates to the baseline MeZO approach across all evaluated benchmarks. This parity in performance is significant as it not only validates the effectiveness of our RNG manager but also highlights ZO2's capability to maintain model precision while reducing GPU memory usage. Since ZO2 achieves the same accuracy as MeZO-which has already been extensively compared against FO methods in its original paper-ww omitted redundant FO comparisons to keep the scope focused.
6.3
Ablation Study
Table 3: Throughput (tokens/sec) results to validate proposed features.
ZO2
ZO2
ZO2 (no efficient update)
Model
MeZO
ZO2
(no scheduler overlap)
(no reusable memory)
OPT-1.3B
1998
1109 (x0.56)
735 (x0.37)
1567 (x0.78)
1955 (x0.97)
OPT-2.7B
1104
573 (x0.52)
422 (x0.38)
849 (x0.77)
1086 (x0.98)
OPT-6.7B
492
225 (x0.46)
184 (x0.37)
373 (x0.76)
485 (x0.98)
OPT-13B
266
105 (x0.39)
103 (x0.39)
198 (x0.74)
259 (x0.97)
OPT-30B
35
46
81
122
OPT-66B
22
15
36
40
OPT-175B
8
5
13
14
In order to discern the individual contributions of key features within the ZO2 framework to its overall performance, an ablation study was conducted focusing on three critical

---

Published as a conference paper at COLM 2025
components: the dynamic scheduler (Sec. 5.2), reusable memory (Sec. D.2), and efficient parameter updating (Sec. D.3). This study focused on throughput, as the three features were designed to improve it without affecting ZO2's built-in memory efficiency. Since CPU offloading already ensures low memory usage, an ablation on memory was unnecessary: Given the tightly integrated nature of our system, traditional ablation methodologies that add one feature at a time to baseline are impractical. Instead, we adopted a reverse ablation approach where each feature was individually disabled. This allowed us observe the decrement in throughput relative to the fully operational framework, thereby highlighting the significance of each component.
The results, presented in Table 3, provide clear illustration of how the absence of each feature impacts the system's 's throughput: (1) Horizontal Comparison. Across all models, the removal of reusable memory results in the most substantial decrease in throughput, followed by the dynamic scheduler, and finally, the efficient parameter updating. This order of impact suggests that while all three features are pivotal, the overhead introduced by CUDA malloc operations, which are eliminated by reusable memory, significantly outweighs the communication delays between the CPU and GPU, managed by the dynamic scheduler and efficient parameter updating. For instance, when reusable memory is not employed, the throughput drops to 37% of the fully optimized framework for the OPT-6.7B model, highlighting its critical role in enhancing performance. (2) Vertical Compan Comparison. As the model size increases, the relative importance of the dynamic scheduler and efficient parameter updating grows more pronounced. This trend is observable from the throughput: for larger models like OPTT-6.7B, the reduction in throughput when the scheduler and efficient update features are disabled is relatively larger than in small models. This indicates that as models become larger, the complexities and overheads associated with managing and optimizing communications between CPU and GPU become more critical to maintaining performance.
7 Conclusion
In this paper, we presented ZO2, an efficient framework that enables the training of ex- tremely large language models, such as the OPT-175B, with 18GB GPU memory-a capa- bility previously unattainable with traditional methods. By effectively integrating CPU offloading, random number generator manager, high-performance dynamic scheduler, effi- cient memory management, efficient parameter updating, and AMP support, our framework reduces GPÚ memory demands while maintaining high throughput without additional time costs. These innovations not only lower the bar for teams with limited hardware resources and advance the democratization of large models, but also open new avenues for advancing AI technology more efficiently.
Acknowledgment
Di Wang and Liangyu Wang are partially supported by KAUST through awards BAS/1/ 11660-0000,,, URF/1/550S-01-1 and by the KAUST Center of Excellence for Genera- tive Al under award number 5940.
References
HanQin Cai, Yuchen Lou, Daniel McKenzie, and Wotao Yin. A zeroth-order block coordinate descent algorithm for huge-scale black-box optimization. In International Conference on Machine Learning, pp. 1193-1203. PMLR, 2021.
HanQin Cai, Daniel McKenzie, Wotao Yin, and Zhenliang Zhang. Zeroth-order regularized optimization (zoro): Approximately sparse gradients and adaptive sampling. SIAM Journal on Optimization, 32(2):687-714, 2022.
Aochuan Chen, Yimeng Zhang, Jinghan Jia, James Diffenderfer, Konstantinos Parasyris, Jiancheng Liu, Yihua Zhang, Zheng Zhang, Bhavya Kailkhura, and Sijia Liu. Deepzero:
10

---

Published as a conference paper at COLM 2025
Scaling up zeroth-order optimization for deep model training. In The Twelfth International Conference on Learning Representations, 2024.
Shuyu Cheng, Guoqiang Wu, and Jun Zhu. On the convergence of pproo-guided zeroth- order optimization algorithms. Advances in Neural Information Processing Systems, 34: 14620-14631, 2021.
Christopher Clark, Kenton Lee, Ming-We Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. aiv preprint ar xXiv:1905.10044. 2019.
Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entail- ment challenge. In Machine learning challenges workshop, pp. 177-190. Springer, 2005,
Marie-Catherine De Marneffe, Mandy Simons, and Judith Tonhauser. The commitmentbank: Investigating projection in naturally occurring discourse. In proceedings of Sinn und Bedentung, volume 23, pp. 107-124, 2019.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama herd models. Xrviv preprint ar Xiv:24407 21783,, 2024.
John C Duchi, Michael I Jordan, Martin J Wainwright, and Andre Wibisono. Optimal rates for zero-order convex optimization: The power of two function evaluations. IEEE Transactions on Information Theory, 6188)22788-8806, 2015.
Artyom Eliseev and Denis Mazur. Fast inference mixture-of-experts language models with offloading. arXiv preprint ar XrI::22312.17238, 2023.
Tanmay Gautam, Youngsuk Park, Hao Zhou, Parameswaran Raman, and Wooseok Ha. Variance-reduced zeroth-order methods for fine-tuning language models. arXio preprint arXiv:2404.08080, 2024.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for 1lms. arXiv preprint ar arXie:2310.01801 2023.
Georgi Gerganov. llama.cpp, 2023. URL http:://gggttbbbbommmeeeer/////a cpp.
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, al. Lora: Low-rank adaptation of large language models. ICLR, 1,2):3, 2022.
Haochen Huang, Jiahuan Pei, Mohammad Aliannejadi, Xin Sun, Moonisa Ahsan, Pablo Cesar, Chuang Yu, Zhaochun Ren, and Junxiao Wang. Fine-grained vision-language modeling for multimodal training assistants in augmented reality. arXio preprint mmXi::2507.05555, 2025.
Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Looking beyond the surface: A challenge set for reading comprehension over multiple sentences. In Proceedings of the 2018 Conference of the North American Chapter of the Associa- fion for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), pp. 252-262, 2018
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium Operating Systems Principles, pp. 611-626, 2023.
Hector Levesque, Emnest Davis, and Leora Morgenstemn The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning, 2012.
11

---

Published as a conference paper at COLM 2025
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. Pytorch distributed: Experi- ences on accelerating data parallel training. ar preprint ar xxXo::2006.1500 2020.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp- 3747-5766 2020.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ar preprint arXio:1711.05101, 2017.
Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems, 11153338-530755 2023.
Yurii Nesterov and Vladimir Spokoiny, Random gradient-free minimization of convex functions. Foundations of Computational Mathematics, 17(2)5527-566, 2017.
Long Ouyang, Jeffrey Wu, Xu liang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An im- perative style, high-peeffrrmnce deep learning library. Advances in neural information processing systems, 32, 2019.
Jiahuan Pei, Irene Viola, Haochen Huang, Junxiao Wang, Moonisa Ahsan, Fanghua Ye, Jiang Yiming, Yao Sai, Di Wang, Zhumin Chen, et al. Autonomous workflow for multimodal fine-grained training assistants towards mixed reality. In Findings of the Association for Computational Linguistics ACL 2024, pp. 4051-4066, 2024.
Jiahuan Pei, Fanghua Ye, Xin Sun, Wentao Deng, Koen Hindriks, and Junxiao Wang. Con- versational education at scale: A multi-lim agent workflow for procedural learning and pedagogic quality assessment. Xiv preprint Xrx:::2507 xxX::2507700528, 2025.
Mohammad Taher Pilehvar and Jose Camacho-Collados. Wic: the word-in-context dataset for evaluating context-sensitive meaning representations. ar Xio preprint ar Xixx:1808.09121, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero- infinity: Breaking the gpu memory wall for extreme scale deep learning, 2021. URL http:://arxiv..og/bbb//111 07857.
Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. (Zero-offload]: Democratizing (illiinn- scale) model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551-564, 2021.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathe- matical reasoning in open language models. ar preprint arXi::2402.0000 03300, 2024.
12

---

Published as a conference paper at COLM 2025
Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, lon Stoica, and Ce Zhang. Flexgen: High-throughput generative inference of large language models with single gpu. In International Conference Machine Learning, pp- 31094-31116 PMLR, 2023.
Utkarsh Singhal, Brian Cheung, Kartik Chandra, Jonathan Ragan-Kelley, Joshua Tenen- baum, Tomaso A Poggio, and Stella X Yu. How guess a gradient. arXiv preprint (tXiio:3122.4479/ 2023.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. Recursive deep models for semantic com compoittionality over sentiment treebank. In Proceedings of the 2013 conference empirical methods in natural language processing, Pp. 1631-1642, 2013.
Tianxiang Sun, Zhengfu He, Hong Qian, Yunhua Zhou, Xuan-Jing Huang, and Xipeng Qiu. Bbtv2: Towards gradient-free future with large language models. In Proceedings the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 3916-3930, 2022a.
Tianxiang Sun, Yunfan Shao, Hong Qian, Xuanjing Huang, and Xipeng Qiu. Black-box tuning for
-asooee------sevvicee In International Conference on Machine Learning, pp. 20841-20855. PMLR, 2022b.
Liangyu Wang, Huanyi Xie, Xinhai Wang, Tianjin Huang, Mengdi Li, and Di Wang. Infinite sampling: Efficient and stable grouped rl training for large language models. arXiv preprint ar arii::2506.22500 2025.
Leyang Xue, Yao Fu, Zhan Lu, Luo Mai, and Mahesh Marina. Moe-infinity: Activation-aware expert offloading for efficient moe serving. arXiv preprint #rXic:2401.13611 2024.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arxio preprint arXiv:2505.09388, 2025.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv=2205001068, 2022.
Yihua Zhang, Pingzhi Li, Junyuan Hong, Jiaxiang Li, Yimeng Zhang, Wenqing Zheng, Pin- Yu Chen, Jason D Lee, Wotao Yin, Mingyi Hong, et al. Revisiting zeroth-order optimization for memory-efficient 1lm fine-tuning: benchmark. ar preprint arXiv:2402. 11592, 2024.
Chujie Zheng, Shixuan Liu, Mingze Li, Xiong-Hui Chen, Bowen Yu, Chang Gao, Kai Dang, Yuqiong Liu, Rui Men, An Yang, et al. Group sequence policy optimization. arviv preprint ar *Xit:2507.1807 2025.
13

---

Published as a conference paper at COLM 2025
A
More Figures, Algorithms, and Tables
AdamW
60
57
MeZO
ZO2
Memory Usage (GB)
40
29
20
188
10
X
X
OPT-6.7B
OPT-13B
OPT-30B
OPT-175B
Figure Single GPU memory usage comparison for training LLMs across different optimiz- ers (AdamW, SGD, MeZO, and ZO2 (Zeroth-Order Offload)) and model sizes (OPT-6.7B, OPT-13B, OPT-30B, OPT-175B). The 'X' indicates that training was not feasible due to excessive memory demand.
Computatin
CPU
Inpal
Black
Black 2
-xX
Black n
Commaniotinn
Dat & Uppart
Upiad
Uptaa vira oxtrad grakiai
Projected Gradiced
Upaa
Block
Embedding
Bleck I
Bleck 2
Black
LM lead
Ohipot
Endeoling:
Block
Bliik
Bllck
LM
lhe tat
Othhed with kvr bit
CRU
Black
Bleck
Figure 4: Workflow of the ZO2 framework (AMP mode) for fine-tuning LLMs.
(a) ZO2 without overlap.
(b) 702 with overlap.
Figure 5: Sequential Task Execution in the Naive ZO2 Framework Depicting Non- overlapping Furl-Forward Passes and Associated Inefficiencies
14

---

Published as a conference paper at COLM 2025
Algorithm MeZO (Malladi et al., 2023)
Require: Model parameters
R loss function R" + R, training iterations T,
perturbation step size E, data batch B, learning rate
1: for == L,...., do
2:
Set random seed and sample batch B C D
3:
f- PerturbParrmmeters(B.ee
4:
L <- L(0; B)
5:
Reset RNG with seed
6:
PerturbParameters/ -2c)
7:
L_ - L(e; B)
8:
Reset RNG with seed
9:
+- PerturblParameters/ (,,
10:
8 8- (
11:
Reset RNG with seed
+ Updatel'arameters: (0,
13:
end for
14:
function UPDATEPARAMETERSÓ.
15:
for each 0; 0, where 8 E R4 do
16:
2; NO,11) E IR4:
17:
8 -
1/1 -2
18: end for
19: return
20:
end function
function PERTURBPARAMETEESS..
22:
for each 0; 0, where 0; E R4 do R4
23:
Zi (0,1))
24:
0 k 0 + EZ;
end for
26:
return
end function
Table 4: OPT model family configs in experiments.
Model Size
Layers
Heads
Dimension
Sequence Length
1.3B
24
32
2048
2.7B
32
32
2560
6.7B
32
32
4096
13B
40
5120
2048
30B
48
56
7168
66B
72
9216
175B
96
96
12288
15

---

Published as a conference paper at COLM 2025
B More Related Work
CPU Offloading for LLMs. With recent advancements in LLMs, several approaches have emerged offload data to CPU memory, mitigating GPU memory limitations. One such method is vLLM (Kwon et al., 2023), which utilizes PagedAAtentionn
to dynamically manage the key-value (KV) cache at a granular block level. Portions of the KV cache can be temporar- ily swapped out of GPU memory to accommodate new requests. Llama.cpp (Gerganov, 2023) addresses oversized LLMs inference by using static layer partitioning. It stores certain contiguous layers in CPU memory while keeping others in GPU memory. During computa- tion, the CPU handles the layers in its memory, followed by the GPU computing its assigned layers. FlexGen (Sheng et al., 2023), GPU-centric inter-layer pipeline LLMs inference method, seeks to improve throughput by pinning some model weights in GPU memory for each layer. During inference, it overlaps GPU processing of the current layer with data loading for the next. DeepSpeed (Rajbhandari et al., 2020) introduces a technique to offload the first-order optimizer state to the CPU, significantly reducing GPU memory requirements during training. Zero-offload (Ren et al., 2021) extends the DeepSpeed approach by not only offloading data to the CPU but also engaging the CPU in computational tasks. Despite these ad advancements, the predominant focus of previous research has been on optimizing LLM inference or first-order optimization through strategic CPU-GPU data transfers. Our work, in contrast, introduces a novel approach by implementing CPU offloading specifically for zeroth-order optimization and line-tuning of LLMs.
C Full Preliminaries on ZO and ZO-SGD
ZO optimization offers gradient-free alternative to first-order (FO) optimization by ap- proximating gradients through function value-based estimates. There are different ZO optimizers for estimating the gradient. To better illustrate our framework, in this paper, we focus on the randomized gradient estimator (RGE) proposed by (Nesterov Spokoiny, 2017), which approximates the FO gradient using finite differences of function values along randomly chosen direction vectors and has been used widely in the ZO optimization literature. Our idea can be applied to other ZO optimizers.
Given loss function L(-) and model with parameter O E Rd, the RGE employed by MeZO (Malladi al., 2023), referred to as ((),, is to approximate (())) and is expressed using central difference:
L = gz Rd
(3)
L(e L(0+ L(e ez)
g
E R
(4)
2
where is random direction vector drawn from the standard Gaussian distribution (0 I), and e > 0 is a small perturbation step size, also known as the smoothing parameter. g represents the projected gradient computed using the model's dual-forward passes. Notably, g R is just scalar value and requires minimal memory space. The rationale behind RGE stems from the concept of the directional derivative (Duchi et al., 2015). As approaches 0, the directional derivative provides us an unbiased gradient estimator of f(x).. Thus, the RGE fx(x) can be interpreted as an approximation of the FO gradient ff(x) using the directional derivative (Zhang et al., 2024). Zeroth-order stochastic gradient descent (ZÖ-SGD) follows a similar algorithmic framework to its first-order counterpart, SGD, but updates the parameters with the estimated gradient V fx) via zeroth order (function value) information for the descent direction. The parameters update of ZO-SGD is defined by:
0 = 0-
where 7 is the learning rate. is important to note that the variable in Equation 2 should be identical to the z in Equation 1.
Specifically, the entire MeZO workflow is shown by Algorithm 3. The process initializes with parameters 0 and iterates over a predetermined number of steps T. Each iteration samples a batch B from dataset D and employs a perturbation strategy. Parameters @ are
16

---

Published as a conference paper at COLM 2025
first perturbed positively to compute loss L,, followed by negative perturbation of 2e to compute k_.. Parameters are then reset to their original state for gradient estimation. The gradient is approximated by the difference (4+ - / (2e), representing the directional derivative along perturbed parameters. This projected gradient, combined with random Gaussian noise 2, updates each parameter 8, optimizing the loss function.
D More ZO2 Implementation Details
D.1 Specific Dynamic Scheduler Design for Efficient Overlap
Figure 5a offers schematic representation of the sequential, non-overlapping task execution within the basic ZO2 framework, particularly highlighting the inefficiencies of dual-forward input passes without task overlap. Initially, input data is transferred from the CPU to the GPU,
beginning with input processing through the embedding layer. Following this, each trans- former block-ranging from Block to Block n-undergoes a distinct cycle: upload to the GPU, execution of dual-forward computations, and subsequent offload back to the CPU upon completion of computations.
This linear processing sequence reveals a critical inefficiency: the GPU must idle while awaiting each block's 's offload back to the CPU, delaying the upload and processing of the subsequent block. This causes significant downtime for the GPU during offloads and for the CPU during uploads, as each must wait for the other's task completion before proceeding. The depicted lack of overlap between computation (green arrows) and communication (blue arrows) tasks pinpoints a crucial area for enhancement. Implementing an overlapped or asynchronous task management strategy, like Figure 5b, could markedly improve system efficiency and throughput, potentially reducing training times and optimizing the use of both CPU and GPU resources.
To overlap the data loading and computation process, we propose a dynamic scheduler, utilizing the asynchronous execution on different CUDA streams. Specifically, our scheduler includes three CUDA streams (Figure 1), which are utilized to control the i-th transformer block's computation, the i 1-th block's uploading, and the i 1-th block's offloading can occur concurrently. This design minimizes data transfer conflicts and maximizes GPU utilization by keeping computational and communication channels active.
However, designing this dynamic scheduler presents challenges when communication tasks outlast computation tasks, leading to potential errors. For example, if the upload of the i-th block is incomplete when its computation begins, this can lead to errors, as the GPU computes with an incomplete set of parameters. Similarly, if the computation of the i-th block is still ongoing when its offloading begins, it can also result in errors because the computation is disrupted by the removal of necessary data. An intuitive solution is to perform a global synchronization directly after initiating U(W++1), C(Wi), and O(W;1) asynchronously. However, this approach is likely to introduce delays, or "bubbles," " due to the varying execution speeds of the three tasks. To mitigate these issues, we establish (1) The offloading task of the i-th block is contingent upon the completion of its computation task, which, in turn, depends on the completion of its uploading task. (2) Each task must wait for the preceding task of the same type to complete; hence, the i-th task cannot commence until the i 1-st task has concluded. To efficiently manage task execution, we use an asynchronous task launch sequence: O(W,-1), C(W,), and U(W;11). Here, O(W;1) synchronously awaits the completion of both C(W; and O(W; 2). Similarly, C(W;) will not start until both U(W,) and C(W, 1) have finished. Lastly, U(W++1) must wait for the completion of UWW)), ensuring orderly and error-free task progression.
two critical dependency relationships:
A significant advantage of our framework is that the dual-forward mechanism doubles the computation time, while the communication time per block remains unchanged. This enhancement substantially increases the likelihood of complete overlap between communi- cation and computation tasks, especially since communication between the CPU and GPU is generally slower than computation on the GPU. Our following evaluations (Section 6) show that with zO's unique dual-forward passes, which extend computation times compared
17

---

Published as a conference paper at COLM 2025
with the single forward pass, communication delays are no longer the primary bottleneck in most scenarios.
Moreover, special attention needs to be given to the embedding parameters and the LM head, as they represent the beginning and end of the model, respectively. By consistently maintaining both the embedding and LM head on the GPU, circumvent the overhead linked to frequent transfers. For the embedding layer, simultaneous uploading of input data and embedding parameters could compete for interconnection bandwidth. Moreover, keeping the embedding layer on the GPU enables the pre-uploading of the first transformer block, effectively overlapping with the computations of the embedding layer. Meanwhile, continuously keeping the LM head on the GPU removes delays associated with its offload- ing-since no subsequent block computations overlap with this offloading- and facilitates weight sharing with the embedding layer, as noted in some conditions (Radford et al., 2019), thus consolidating related computations and enhancing operational efficiency. The detailed scheduler design to apply ZO2 on LLMs is shown in Algorithm 2.
D.2 Efficient Memory Management via Reusable One Block Space GPU
We adopt the strategy by (Ren et al., 2021) to optimize memory management by pre- allocating a reusable transformer block of memory on the GPU, eliminating the overhead of repeated memory allocations and deallocations during data transfers between the CPU and GPU. This memory is dynamically reassigned each transformer block in sequence, speed- ing up data transfers and stabilizing GPU memory usage, thereby enhancing computational efficiency.
We also adopt the strategy outlined by Li et al. (2020), leveraging communication buckets to enhance the throughput of block communications. Specifically, we concatenate parameter fragments within blocks into contiguous memory buckets, thus improving communication efficiency.
D.3 Efficient Parameter Update Strategy
Pay Duo Arrwod
We-ty Dual- bosaandd
Deal fenward
Du bowed
Desifoneer
Desifoovrt
caree
ap-s
Wast
Projectod
(a) Model parameter updates without the effi- cient strategy.
(b) Model parameter updates with the efficient strategy.
Figure 6: Comparison of model parameters updates without/with efficient strategy. (a) illustrates the process where, at the j-th iteration, the model computes the projected gradient 8i using the dual-forward method and subsequently updates the model parameters. (b) demonstrates that at the j-th iteration, the model first updates the parameters using the previously saved projected gradient 8j-1, and then performs the dual-forward pass to compute the new projected gradient 8i-
In the ZO2 framework, the parameter update strategy is meticulously designed to precede the dual-forward computations of each transformer block. Traditionally, each transformer block is subjected to two distinct data transfer phases (Figure 6a, two green dotted boxes for each block): one for the dual-forward computations and another for applying gradient updates. This requirement stems from the fact that the (approximated) gradients are obtained only after completing the dual-forward computations for all blocks. For the first- order methods, the same offloading strategy requires parameters to be uploaded for the computation phase, offloaded upon completion, and then re-uploaded and offloaded again for the parameter update phase, given that the gradient dimensions match those of the parameters. This iterative process effectively doubles the communication load and extends the duration of training.
However, compared to the interdependence of dual-forward calculations across blocks, the parameter update process remains independent for each block, allowing us to reorder the operations. Once blocks are updated with the last iteration's gradients, only single upload
18

---

Published as a conference paper at COLM 2025
and offload cycle is necessary for each block. This streamlined approach is only feasible in the ZO framework because, unlike first-order methods where the gradient dimensions are identical to those of the parameters, the projected gradient from ZO is only scaler and can be persistently stored on the GPU. By implementing preemptive parameter updates, the framework significantly curtails the number of data transfers required per iteration (Figure 6b, one dotted box for each block). This adjustment not only halves the usage of interconnection bandwidth but also enhances the efficiency of the training process, thereby streamlining operations and reducing overhead.
D.4 ZO2 in AMP Mode
Figure 4 illustrates the workflow of the ZO2 framework under AMP mode, which employs reduced precision formats to accelerate the training of LLMs. AMP leverages formats such as Tensor Float Point 32 (TF32), which provides higher computational throughput compared to Float Point 32 (FP32). AMP represents a compromise between FP32 and FP16, retaining the data storage format of 32 bits while offering computational speeds comparable to FP16. This adaptation allows for faster computation while maintaining the same communication speed as purely FP32 workflow, which can complicate the overlap. Therefore, specialized framework must be designed to manage these unique challenges effectively. This acceleration is critical for enhancing training efficiency but introduces challenges in maintaining effective computation-comit computationcoommnniattonn overlap, the data transfer still utilizes the FP32 format.
To address this, the ZO2 framework incorporates a compression mechanism where pa- rameters are compressed to low-bit formats during offloading from GPU to CPU. This compression significantly reduces the data volume, enabling quicker transfers and mitigat- ing bandwidth limitations. The current compression settings include bfloat16 and float16, which reduce the data size by 50%, and more aggressive reductions like float8, which compress to 25% of the original size.
Upon uploading these compressed parameters back to the GPU, they are decompressed and restored to FP32 for high-precision parameter updates. Subsequent computations, particularly the dual-forward passes, are then performed using the TF32 format to exploit the computational speed.
E Codebase
tarch
from torch.optia Lnport adams
froe ant Inport 200mfigg 20_if_init
Free mowecfreeet suport ORTFonCacssii
a Hodel and aptimize indt
Redi ltit
20 con.o ossonfig/ecctooeeeeee "scoooo""" 2004trae, officarinng "opu". working destem "cacaso", Srile-s)
mode] OHIForCussINN http://.orrrppptttteeeeeeeeeeedd pwwtrstteedd "acebook/Sstti..0 fotprroostteeetooooookkottita.. toceea"
with nujr__inittan_cmmig
Snport
probratrad) facebookkstt-2..99
(ereeeee Imste-5)
Traintng
Trxising
in wwrpleee_trainnng..o
Traaa
8 Tradn
training Input.Idk, 1dh, training labela aet cratntig data tatch mode). traint)
tradrdng Seput, 1ds training Jabuls Jabia #sdc).00, tradn()
get traindng date batch
spetat www. grodt)
Loxs irdel tou-training, input Ids, sawlectrainingg Jsbels)
cutput modeltinput ids training irput iit, Latala training. latela] wopptt. ancherl))
gyclatzee.stss()
extpor engellingst (ddseva)__iiirtt_t__oo Ixtelsseval_aablll
soil Iepur_tds, Jatels
hatet
model, evall)
rinzl graising updats wini, opt. (pspps/s/////
with tarchuno_grott)
outpee- #stel/fapat
(dt-rat Ado-rssl trout, Inteleweval, (abbel)
(a) PyTorc first-order optimizer training.
(b) 7O2 training.
Figure 7: Comparison of model training API with PyTorch first-order optimizer and with ZO2.
19

---

Published as a conference paper at COLM 2025
ZO2 is encapsulated in approximately 5,500 lines of Python code, designed for ease of use, paralleling standard PyTorch training paradigms. The framework facilitates seamless integration and customization, enabling both researchers and practitioners to adapt it to diverse requirements efficiently. Figure 6 provides an illustrative example that demonstrates the API's similarity to conventional PyTorch usage, underscoring the user-friendly nature of ZO2.
F Full Experimental Setup
The experimental evaluation of our framework was conducted using the PyTorch deep learning library, integrated with NVIDIA CUDA streams to optimize parallel computation tasks. We selected the Open Pre-trained Transformer (OPT) (Zhang et al., 2022) model family (Table 4) as the subject of our experiments due to its open-source availability, widespread adoption in the research community, and diverse range of model sizes, ranging from 1.3 billion to 175 billion parameters, which allows for comprehensive assessment of our framework's performance across varying scales of model complexity.
In our evaluation, MeZO (memory-efficient zerothorder optimizer) (Malladi et al., 2023) serves as the baseline method, as it is the most memory-throw senoryythrouuhhut efficient ZO method cur- rently. Our framework builds upon MeZO, reducing GPU memory usage while maintaining throughput and precision. All performance, including measurements of GPU memory usage and throughput, were conducted using the Stanford Sentiment Treebank (SST-2) Socher et al. (2013). The tests were performed on system equipped with an NVIDIA A100 GPU with 80GB of memory and an AMD Milan CPU, operating under Python version 3.11, PyTorch 2.4.0, and CUDA 12.1. Our experiments further employed hyperparameters such as a learning rate of 1 10-7, 7, batch size of 1, 100 training steps, and sequence length of 2048, to rigorously test the framework under controlled conditions.
G More Experiments
G.1 Evaluation of AMP Mode
Table 5: Throughput (tokens/sec) results to validate AMP Mode. AMP auto-cast with FP16
(top) and BF16 (below).
702
702
702
702
Model
(non-compress)
(FP16)
(BF16)
(FP8)
OPT-1.3B
4827
4770 (x0.988)
4760 (x0.986)
4802 (x0.995)
OPT-2.7B
2811
2974 (x1.058)
2974 (x1.058)
2997 (x1.066)
OPT-6.7B
1271
1641 (x1.291)
1641 (x1.291)
1662 (x1.308)
OРT-13B
561
930 (x1.658)
930 (x1.658)
951 (x1.695)
OРT-30B
286
416 (x1.455)
416 (x1.455)
425 (x1.486)
OPT-66B
127
(x1.512)
192 (x1.512)
198 (x1.559)
OPT-175B
43
65 (x1.512)
65 (x1.512)
68 (x1.584)
OPT-1.3B
4565
4430 (x0.970)
4430 (x0.970)
4463 (x0.978)
OFT-2.7B
2778
2816 (x1.014)
2816 (x1.014)
2818 (x1.014)
OPT-6.7B
1273
1594 (x1 x1.252))
1594 (x1.252)
1612 (x1.266)
OPT-13B
678
910 (x1.342)
910 (x1.342)
924 (x1.363)
OРT-30B
285
407 (xl (114488)
407 (x1.428)
415 (x1.456)
OPT-66B
127
188 (x1.480)
188 (x1.480)
194 (x1.528)
OPT-175B
43
64 (x1.488)
64 (x1.488)
67 (x1.565)
The efficiency of the AMP mode shown Table 5, where we evaluate the throughput using two AMP auto-cast computational data formats: FP16 and BF16. Additionally, we investigate the impact of various compression formats (FP16, BF16, and FP8) on communi- cation and compulation performance detailed in Section D.4.
20

---

Published as a conference paper at COLM 2025
Across all models tested, clear trend emerges: lower-bit compression formats consistently yield higher throughput. Notably, there no significant difference in throughput between the 16-bit formats, FP16 and BF16, suggesting that the compression efficiency rather than the specific format type is the crucial factor in enhancing communication speed.
In most scenarios (specifically for the OPT models greater than 2.7B), employing low-bit compression results in superior throughput, underscoring the benefits of reducing data transfer volumes. However, exceptions are observed, such with the OPT-1.3B model, where non-compressed data slightly outperforms the compressed formats. This outcome is attributed to the system being computation-boune rather than communication-boutt such contexts, the additional computational demands imposed by the compression process do not sufficiently offset the benefits of reduced data transfer times, thereby introducing an overhead that detracts from the overall system efficiency.
In
G.2 Analysis of More Experimental Settings
Table 6: Different batch-size analysis.
Memory Usage (MB)
Throughput (tokens/sed
Model
Batch-size
MeZO
ZO2
Me7O
702
OPT-1.3B
8898
5098 (x0.57)
1998
1955 (x0.97)
OPT-2.7B
14514
5930 (x0.41)
1104
1086 (x0.98)
OPT-6.7B
32930
8420 (x0.26)
492
485 (x0.98)
OРT-13B
58762
10736 (x0.18)
266
(x0.97)
OPT-1.3B
11733
8481 (x0.72)
2126
2105 (x0.99)
OPT-2.7B
18373
9312 (x0.50)
1204
1193 (x0.99)
OPT-6.7B
34943
9683 (x0.27)
561
556 (x0.99)
OPT-13B
61919
12315 (x0.19)
297
294 (x0.99)
OPT-1.3B
16067
12887 (x0.80)
2309
2296 (x0.99)
OPT-2.7B
23731
14211 (x0.59)
1268
1262 (x0.99)
OPT-6.7B
43683
17233 (x0.39)
589
587 (x0.99)
OPT-13B
72159
19591 (x0.27)
312
311 (x0.99)
OPT-1.3B
27529
22775 (x0.82)
2358
2348 (x0.99)
OPT-2.7B
37113
23639 (x0.63)
1297
1294 (x0.9)
OPT-6.7B
57541
26207 (x0.45)
604
602 (x0.99)
OPT-13B
31399
320
Differential Batch-size and Sequence Length Analysis. This analysis explores the impact of varying batch sizes and sequence lengths on the performance of the ZO2 compared to the MeZO baseline. Tables 6 and 7 present the memory usage and throughput metrics for different configurations of the OPT models, ranging from 1.3B to 13B parameters. Table 6 shows the results for different batch-sizes. As batch size increases, there is consistent trend where ZO2 maintains throughput equivalency with MeZO across all model sizes, despite significant reductions in memory usage. Even higher batch sizes, ZO2 demonstrates robust performance, showing almost no decrease in throughput relative to its MeZO counterpart. For example, in the OPT-1.3B model a batch size of 8, the throughput remains constant at 2348 tokens/sec /sec, maintaining operational efficiency irrespective of the increased computational load.
Table 7 illustrates the impact of sequence length throughput. Similar to the batch- size analysis, increasing the sequence length does not compromise the throughput of ZO2, maintaining parity with the MeZO model across varying lengths. Notably, even at a sequence length of 8192 for the OPT-1.3B model, ZO2 sustains a throughput 1279 seec effectively handling larger input sizes without drop in performance.
The analyses confirm that ZO2 effectively manages larger batch sizes and sequence lengths without sacrificing throughput. This resilience is crucial for practical deployments where varying input sizes and batch configurations are common, underscoring the scalability and robustness of the 702 approach in diverse operational environments.
21

---

Published as a conference paper at COLM 2025
Table 7: Different sequence length analysis.
Memory
Usage (MB)
Throughput (tokens/sec)
Model
Length
MeZO
ZO2
MeZ0
Z02
OPT-1.3B
8665
4471 (x0.51)
1901
1830 (x0.96)
OPT-2.7B
1024
14465
5389 (x0.37)
1051
1013 (x0.96)
OPT-6.7B
31379
8405 (x0.26)
439
426 (x0.97)
OPT-13B
56183
10717 (x0.19)
244
243 (x0.99)
OPT-1.3B
8898
5098 (x0.57)
1998
1955 (x0.97)
OPT-2.7B
2048
14514
5930 (x0.41)
1104
1086 (x0.98)
OPT-6.7B
32930
8420 (x0.26)
492
485 (x0.98)
OPT-13B
58762
10736 (x0.18)
266
259 (x0.97)
OPT-1.3B
11379
7581 (x0.67)
1707
1692 (x0.99)
OPT-2.7B
4096
19655
9023 (x0.45)
1008
1001 (x0.99)
OPT-6.7B
38047
11763 (x0.30)
509
505 (x0.99)
OPT-13B
64519
14915 (x0.23)
275
272 (x0.99)
OPT-1.3B
32499
20507 (x0.63)
1282
1279 (x0.99)
OPT-2.7B
8192
38127
21351 (x0.55)
786
784 (x0.99)
OPT-6.7B
54495
24115 (x0.44)
446
445 (x0.99)
OPT-13B
30355
246
G.3 Evaluation on Qwen3 Models
To evaluate the generalizability of ZO2 beyond the OPT family, we conduct experiments on Qwen3 (Yang et al., 2025). Compared to MeZO, ZO2 achieves significant GPU memory reduction with only slight drop in throughput. This confirms that our framework remains highly efficient, even when applied to different LLM architectures. This further supports the adaptability of ZO2 across diverse model families.
Table 8: Results of ZO2 performance for various Qwen3 models and FP16 modes. The values parentheses (x) represent the ratio of each measurement compared the baseline MeZO (first column) configuration.
GPU Memory Usage (MB)
Throughput (tokens/sed sec) MeZO (FP16)
Model
MeZO (FP16)
ZO2 (FP16)
ZO2 (FP16)
Qwen3-1.7B
9329
8649(x0.92)
8320
6873(x0.82)
Qwen3-4B
14357
9287(x0.64)
3955
3600(x0.91)
Qwen3-8B
22995
10403(x0.45)
2391
1827(x0.76)
Qwen3-14B
36877
13033(x0.35)
1421
1114(x0.78)
Qwen3-32B
71195
13691(x0.19)
668
500(x0.74)
22


Reviews:
[{"review_id": "rJhTxuWcV8", "reviewer": "Reviewer_pm1a", "summary": "This paper presents a memory-efficient language model fine-tuning system by combining an existing zeroth-order optimization system with CPU offloading techniques. The authors first identified that zeroth-order optimization, due to its forward-pass-only nature, incurs less communication overhead when using CPU offloading. Then the authors introduced a series of technique: (1) the algorithm loads and offloads parameters in blocks, and the two forward passes and the parameter update from last iteration are done at the same time to avoid multiple communications; (2) the authors introduced a RNG (random number generator) manager to handle doing parameter update and forward passes at the same time; (3) the algorithm overlaps communication and computations; (4) reusable memory allocation to avoid malloc/free overhead; (5) support of low precision. \n\nThe experiment results are very strong: the algorithm performs the same as MeZO (a zeroth-order implementation) and saves up to 10x memory; there is very little speed degradation even given the heavy communication. The authors also present ablations to demonstrate the importance of each of the techniques above.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The proposed ZO2 method is both innovative and very useful. \n\n(1) The authors identified that zeroth-order optimization is more suitable for CPU offloading than first-order methods, as all the operations can be packed into one scan of the parameters and no activations need to be saved. This saves at least twice the amount of communication. The use of RNG manager such that the parameter update and forward passes can be done at the same time is also very cool.\n\n(2) The authors put in a lot of engineering effort (scheduler, support low precision, reusable memory) to enable the high performance of the algorithm.\n\n(3) The memory reduction is significant with almost no speed degradation, thanks to the above design. Wherever people use ZO and need further memory reduction, they should refer to this method.", "reasons_to_reject": "The following is not a reason for rejection (and the authors do not need to add these for rebuttal), but just thoughts on the ZO method and experiments:\n\n(1) Do we need more memory reduction? Considering the performance gap (and the fact that ZO still needs many more steps) between ZO and first-order optimization methods, I can imagine one of the main use cases of ZO is on-device update. Many of these mobile devices nowadays have unified memory (such as the M chips from apple), and hence \"offloading\" no longer makes a difference. It is of course out of the scope of this paper, but I think it would be great if the authors provide some discussions on this topic.\n\n(2) How effective is ZO? Again, this is out of this paper's scope (which assumes ZO is effective and useful). I also understand that this paper follows the original MeZO paper's experiment setup. But it would be cool to see more experiments on other models (Llama for example) and more tasks (instruction tuning for example). Again, the authors do not need to conduct these for rebuttal and this is merely a suggestion.", "questions_to_authors": "Please see \"reasons to reject\".", "extra_scores": {}}, {"review_id": "egPzZOHriz", "reviewer": "Reviewer_Ka4C", "summary": "This paper introduces ZO2 with source code available. ZO2 implements an efficient ZO-offloading training framework on PyTorch with the following contributions:\n(1) asynchronous operations with overlapping: we use 3 CUDA streams that are responsible for upload, compute, and offload. Each CUDA stream can have overlap, and their synchronization order is determined by the correctness of ZO (e.g. offload stream synchronizes with the compute stream, etc.)\n(2) reusable memory space mapping \n(3) RNG state manager and \"fused\" ZO updates with forward.\n\nThe evaluation mostly focus on the memory efficiency: MeZO on OPT-30B will cause CUDA OOM but ZO2 won't. ZO2 can effectively optimize OPT-175B with offloading without CUDA OOM.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The asynchronous computation and overlapping design in uploading/compute/offloading is unique to ZO and discussed sufficiently in the paper. This is a non-trivial system contribution. Figure 1, 4 & 5 is clear about this design.", "reasons_to_reject": "My main concern is the lack of technical novelty on the system side. In my view, the ZO2's core contribution/technical challenge that worth  of discussion is only on the clever design of 3 CUDA stream and figuring out the right dependency and do overlapping (Figure 1, 4, `zo2/optimizer/mezo_sgd/zo2.py inner_zo_forward call`). The other techniques such as AMP during upload/offload, RNG manager, etc. is not worthy of discussion (although they are needed for correctness or efficiency purpose, the technique is either well-known or not hard to develop). \n\nNeed clarification / modification:\n1.  \"parameter updates are ingeniously fused with the dual forward passes\" The code is implemented purely in PyTorch and in eager mode, the updated parameter will be read and written back to DRAM before the forward pass proceeds. **The use of the word \"fused\" is not accurate as it implies that we are performing *in-memory/online* update just before matmul or other operation in SRAM**. I would suggest to replace it with \"combine/merge the call\" unless the framework is actually doing fused *online* update before matmul/other ops in SRAM\n\n2. \"Reusable memory allocation to avoid costly CUDA malloc/free during block transfers\" According to `zo2/optimizer/mezo_sgd/zo2.py, line 439-520`, we are just allocating a large contiguous tensor for upload and offload. This mainly saves the overhead time not necessarily the actual malloc/free time. **The \"reusable\" term is misleading as we are just allocating the large block right before upload/offload**. The same memory space won't be reused across layer and there is no custom memory allocator that pre-allocate a large memory pool and we reuse the pre-allocated space. *It is likely that we here only save the PyTorch overhead time instead of any actual NVidia CUDA malloc time.* \n\nMinor but needed for more precise discussion on the system benefits:\n\nAs ZO2's main contribution is on the efficiency side with several system designs, a low-level trace (PyTorch profiler trace on CUDA stream is acceptable, ideally with detailed nsight trace) is definitely needed in the appendix or in the ablation study.", "questions_to_authors": "I don't have any other questions besides the \"need clarification/modification\" section in the weakness section.", "extra_scores": {}}, {"review_id": "myPDnpe9IJ", "reviewer": "Reviewer_SBMa", "summary": "The paper proposes to shift some computation from GPU to CPU with LLM's zeroth-order fine-tuning. Specifically, it will load the transformer block one by one from CPU memory to GPU memory, and several tricks such as pipelining have been proposed to increase the throughput. Experimental results show that the proposed framework can handle larger models' fine-tuning while maintaining the same performance.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper studies an important problem in reducing the GPU memory requirement for the large language model fine-tuning.\n2. The experimental results demonstrate the effectiveness of the proposed framework.", "reasons_to_reject": "1. The paper's organization and presentation need to be improved. Figure 1 is not helpful in showing the design of the proposed framework, and the figure that the paper refers to several times (Figure 5) is in the appendix. Two of the proposed tricks mentioned in the ablation study are in the appendix. Overall, I believe the paper can be much improved by prioritizing the content and details.  \n2. The proposed framework is only tested in the OPT model families. It would be better to show other models, although I believe the proposed framework should work in other model families. \n3. I am not very familiar with the MLsys community, so I am not sure about the significance of the novelty in the paper. To me, the novelty is OK but not very significant.", "questions_to_authors": "1. Could you explain why the proposed framework works very well in FP16 but not well in FP32? It seems to me the proposed framework should be the same with different precisions.", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces ZO2, a system enabling memory-efficient zeroth-order (ZO) fine-tuning of large language models—up to OPT-175B—on a single 18GB GPU. By leveraging the dual-forward structure of ZO optimization, ZO2 achieves efficient CPU-GPU offloading with minimal overhead. Key contributions include a 3-stream CUDA scheduler, RNG state management, and low-bit precision support.

Strengths:
- Addresses a highly practical challenge in LLM fine-tuning.
- Demonstrates strong empirical results with up to 10× memory savings and no accuracy loss.
- Well-integrated system design, with meaningful insights and thoughtful engineering.

Concerns:
- reviewer Ka4C noted that individual components (e.g., RNG manager, memory buffers) are not novel in isolation, though the integration and insight are recognized as valuable.

Two reviewers strongly supported acceptance (ratings of 8 and 6), emphasizing the method’s novelty and practical impact. A third reviewer, initially more critical (score 5), acknowledged the effectiveness of the system (especially the overlapping design) and raised their score after engaging with the authors' clarifications and data. Although concerns about the originality of some system components remain, the consensus supports that the work constitutes a valuable contribution.

Conclusion:
With general support from reviewers and effective rebuttals, this paper makes a timely and practical contribution. I recommend acceptance.
Decision: Accept