ZeroLock: Concurrent Memory-Efficient LLM Training via Modular Update Decoupling

ZeroLock: Concurrent Memory-Efficient LLM Training via Modular Update Decoupling

arXiv:2608.07974v1 [cs.LG] 8 Aug 2026

Wentao Dai∗ , Xuanran Li∗ , Yuxiang Zhang, Ming Tang, Chao Huang

Time Abstract—Large language model (LLM) fine-tuning at the Time B1 F1 F2 B2 F3 B3 F1 B1 F2 B2 F3 B3 S2 edge adapts the model to scenario-specific data while preserving F1 B1 F2 B2 F3 B3 F1 B1 F2 B2 F3 B3 S2 F1 F2 B1 F3 B2 B3 F1 F2 B1 F3 B2 B3 privacy. Although existing studies proposed pipeline parallelism S1 F1 F2 B1 F3 B2 B3 F1 F2 B1 F3 B2 B3 to address the limited memory and computing resources of edge S1 B1 B2 B3 F1 F2 F3 B1 B2 B3 S0 F1 F2 F3 B11 B2 B3 F1 F2 F3 B12 B2 B3 devices, they commonly rely on backpropagation (BP) training, S0 F1 F2 F3 Iteration Iteration Time Iteration 1 Iteration 2 which has a fundamental limitation of update locking and could (a) S2 F1 B1 F2 B2 F3 B3 F1 B1 F2 B2 F3Time B3 experience severe throughput and memory bottlenecks. In this S2 F1 F2 B1 B2 F2 F3 B2 B3 F3 F1 B3 B1 F1 F2 B1 B2 F2 F3 B2 B3 F3 B3 work, we propose a BP-free algorithm, called ZeroLock, that S1 F1 B1 decouples the model updates into independent chunk updates by S1 F3 F1 B3 B1 F1 B1 F2 F3 B2 B3 B1 B2 F2 F3 B2 B3 F3 B3 F1 F2 S0 F1 B1 F2 B2 local objective construction. It breaks the update locking of BP S0 F1 B1 Iteration F2 B2 1F3 B3 F1 B1 F2Iteration B2 F32 B3 and hence can improve throughput at the algorithm level and Iteration 1 Iteration 2 lower memory usage by reducing activation storage. To the best (b) Time Time Time Time of our knowledge, we provide the first theoretical framework A3 A3 A3 A3 A2 A2 for such local objective construction-based approaches under S2 S2 A1 A2 S2 A2 S2 A1 A1 A1 A3 A3 A3 A3 general model chunk division by mapping local objectives to the A2 A2 S1 S1 A2 S1 A2 S1 A1 A1 A1 A1 A3 A3 global objective. We prove that ZeroLock has a convergence rate A3 A3 √ A2 A2 S0 S0 S0 A2 S0 A2 A1 A1 A1 of Õ(1/ T ), which differs from BP only by polylogarithmic A1 Iteration 1 Iteration 1 1 Iteration Iteration 1 factors. We design a system for ZeroLock and build real-world (c) (d) prototypes, incorporating techniques such as early forwarding and failure recovery for efficient and robust implementation. Fig. 1. Throughput with (a) BP baseline and (b) ZeroLock; memory with (c) BP Experiments on the prototype show that compared to BP-based baseline and (d) ZeroLock. In this figure, we use the typical 1F1B framework baselines, ZeroLock reduces the memory by 26.5% and improves as the BP baseline, because many existing studies (e.g., [5]) are built upon it. S0, S1, and S2 are stages, each corresponding to the updates of a model chunk throughput by 4.9%. Index Terms—Edge intelligence, collaborative edge training and assigned to a device for execution. Fn, Bn, and An denote the forward pass, backward pass, and activations of the n-th microbathch, repetitively. The system, pipeline parallelism, backpropogation-free training. stripes in (c) and (d) indicate the duration that the corresponding activations

I. I NTRODUCTION

are kept in memory. For example, at time slot 3 of S0, A1-A3 are kept in memory in (c), while only A2 is kept in memory in (d).

Edge intelligence deploys large language models (LLMs) at the network edge, achieving privacy-preserving and providing However, devices at the network edge usually have limited real-time inference services. In practice, it is important to finecomputation and memory resources, so it is difficult for a tune LLMs at the edge for two reasons. First, data distributions device or a GPU to fine-tune the entire LLM. To address this, are usually scenario- and individual-specific, making it imporexisting studies have proposed pipeline parallelism approaches tant to adapt the models to maintain inference accuracy. Second, [3]–[5]. The main idea is to partition the model vertically into data privacy concern and possibly frequent model adaptation chunks of consecutive layers. The update of each model chunk requirements make it difficult to upload the data to the cloud corresponds to a stage and is assigned to a different device, for fine-tuning. There are many such examples requiring both and these devices collaborate to update the chunks in a pipeline model adaptation and privacy preserving. For example, test-time fashion. GPipe [3] is a typical approach and processes microtraining [1] adapts LLMs to maintain user-specific memory for batches sequentially. 1F1B [4] interleaves forward pass and personalized service provision. Electroencephalogram signalbackward update to improve pipeline utilization. Building upon based tasks [2] (e.g., emotion recognition, sleep staging) require 1F1B, PipeDream [4], [6] enables asynchronous pipeline trainmodel adaptation for individuals, as those signals demonstrate ing via weight stashing. Recent studies [7]–[10] reduce bubbles strong individual-dependent patterns. in pipeline through placement search, virtual or bidirectional stages, and finer backward decomposition. Other studies [11]– Wentao Dai, Yuxiang Zhang, and Ming Tang are with the Department of Computer Science and Engineering, Southern University of Science [15] focused on pipeline planning considering preemptions, and Technology, China. Xuanran Li is with the Department of Mathstragglers, device heterogeneity, or communication efficiency. ematics, Southern University of Science and Technology, China. Chao Huang is with the School of Computing, Montclair State University, New Confidant [5] implements 1F1B on smartphones. Jersey, USA. Email: {12311217,12312110,12410823}@mail.sustech.edu.cn, Most of these pipeline parallelism works (e.g., [3]–[16]) tangm3@sustech.edu.cn, huangch@montclair.edu. (Corresponding Author: rely on backpropagation (BP) training and focus system-level Ming Tang) * Equal Contribution scheduling optimization. BP training has a fundamental

limitation of update locking [17]–[19]. That is, the model update contains a forward pass followed by a backward pass, so the update of upstream layers needs to wait for the forward and backward computation of downstream layers. As a result, the approaches based on BP could have crucial limitations:

Q2 Does the modular decoupling of local objective construction theoretically harm the model convergence? Q3 How to design and construct real-world prototype systems (for both multi-GPU server and Android phone scenarios) to achieve efficient and robust implementation?

Throughput Bottleneck: As shown in Fig. 1 (a), due to update locking, bubbles arise because an upstream stage (e.g., S0) needs to wait for the backward pass of all downstream stages (e.g., S1, S2) to accomplish its own backward pass. Although existing works (e.g., [4], [6]) have proposed approaches to eliminate the bubbles, the updates of stages remain coupled due to the BP training, and hence the fundamental issue has not been resolved. • Memory Bottleneck: As shown in Fig. 1 (c), due to update locking, an upstream stage (e.g., S0) needs to keep its activations until all its downstream stages and itself (e.g., S0, S1, S2) accomplish their backward passes. This leads to significant memory waste for upstream stages.

Answering Q1 requires concrete design to incorporate LLMspecific characteristics (e.g., token sequence, low-rank adaptation for efficient fine-tuning) into local objective construction framework. Answering Q2 is very challenging. This is because the modular update decoupling breaks global objective into local objectives of model chunks, while recent studies lack frameworks to quantifying how optimizing the local objectives improves the global objective. Answering Q3 is non-trivial due to the real-world engineering requirements of light-weight, high-throughput, and robust implementation. Existing works tried to address one or two of Q1–Q3. PPLL [26] places gradient-isolated vision blocks on different GPUs and transfers their features for block-wise updates. FluidPipe [27] focuses on LLM and adds an auxiliary task head to the first part of a two-stage model, avoiding the local updates on the first part from waiting for the second part’s gradient. SCPL [28] decouples BP via per-segment supervised contrastive losses for synchronized multi-GPU model parallelism. However, these works [26]–[28] considered hidden state alignment for Q1, failing to characterize task-specific information in local objective construction, and they did not address Q2. For Q3, although they propose high-level pipeline logic, their designs remain at a conceptual level, lacking critical implementation details such as continuous cross-batch pipelining, stage I/O queues, and explicit RPC primitives. Meanwhile, they did not provide deployable prototype for real-world execution. Although recent work LoPT [29] addressed Q2, its analysis limits to two chunks and fails to provide a general analytical framework for connecting local objective to global objective. We address those challenges and answer Q1–Q3. Our main contributions are summarized as follows:

•

The pervasive use of BP training and the resulting throughput and memory bottlenecks motivate the main question: Main Question. How can we design a system that fundamentally overcome the throughput and memory bottlenecks at the algorithm level by breaking the update locking in BP? BP-free training was proposed to break the update locking in BP. There are two categories. (I) Backward Gradient Estimation. For example, direct feedback alignment [20], [21] bypasses the chain rule by propagating target errors directly to all layers. Zero-order optimization [22], [23] computes perturbation-based loss differences during the forward pass to estimate gradients. However, this category usually leads to either significant model accuracy degradation for complex tasks or significant computing overhead. (II) Objective Reconstruction. For example, NoProp [24] reframes neural modules as independent denoising units that map noisy target embeddings back to clean targets. Predictive coding [25] alternates between minimizing chunkwise prediction errors and updating weights using the local errors. Local objective construction (e.g., depth-progressive monotonic learning [18]) equips each layer with an independent local objective, allowing purely local gradient computation. Within category (II), local objective construction decouples the modular updates across stages and has comparable model accuracy to BP approaches, so it is a promising candidate to address the Main Question. Specifically, since each chunk is updated independently using its local loss, its update does not need to wait for the forward and backward passes of its downstream chunks, so bubbles due to waiting can be removed; meanwhile, each chunk needs to keep its activations only until its own updates, reducing memory usage. In this work, based on local objective construction-based BPfree approaches, we aim to propose a pipeline parallelism LLM training system to fundamentally overcome the throughput and memory bottlenecks. However, this design is not straightforward and needs to address the following questions: Q1 How to fine-tune LLM with local objective construction?

ZeroLock Algorithm: To answer Q1, we propose a local objective construction-based BP-free algorithm, called ZeroLock, for LLM fine-tuning, which incorporates lowrank adaptation (LoRA) [30] and LLM-compatible readout head and loss function. It breaks the update coupling across model chunks, and hence reduces the pipeline bubbles and activation storage at the algorithm level. • Theoretical Analysis: To answer Q2, we establish theoretical equivalences to represent the local optimum of model chunks in a form of global objective. Then, the local updates of model chunks can be equivalently represented as global updates, with which the global convergence can be derived. To the best of our knowledge, this is the first analytical framework for the local objective construction algorithms under general model chunk division, which provides a systematic approach for analyzing global convergence given decoupled local updates. We √ prove that ZeroLock has a convergence rate of Õ(1/ T ), which differs from BP by only polylogarithmic factors.

•

Real-World System Design: To answer Q3, based on In (2), Wlm is the head weight that maps the normalized hidden ZeroLock, we design systems and build prototypes for representation Norm(hk ) to vocabulary logits z k ∈ RS×V , both multi-GPU server and Android phone scenarios. This with V being the vocabulary size. Vector z k,i is the i-th row of system can achieve concurrent throughput (see Fig. 1 (b)) z k ∈ RS×V . With softmax(·), pk,i ∈ R1×V is the predictive and lower memory usage by reducing activation storage distribution of next token over the vocabulary space. (see Fig. 1 (d)). Techniques such as early forwarding and Local Objective: With the output of the readout head, the failure recovery are proposed to ensure light-weight, high- local loss of each chunk k consists of a task-dependent term throughput, and robust implementation. To the best of our LTask (·) and a consistency term Lk→k−1 Consis (·): knowledge, we build the first prototype on Android system Lk (pk ) = αLTask (pk ) + (1 − α)Lk→k−1 (3) Consis (pk ), for local objective construction-based BP-free algorithm. • Evaluation on Real-World Prototype: Experiments on where α ∈ (0, 1] is a weight, and pk = (pk,i , i ∈ Ω) with Ω multi-GPU server show that when compared with BP- as the token sequence removing prompt or padding positions. based baselines, ZeroLock system reduces the memory (i) Task-dependent term LTask is a local version of the global usage by 26.5% and improves the throughput by 4.9%. The objective. It aligns the output of the readout head of each chunk fine-tuning of TinyLlama on Android phones experiences k to the global ground-truth target: a peak PSS of less than 4000 MiB, a battery temperature
1 X of around 37◦ C, and a wall-clock time of 1644.1 seconds, LTask (pk ) = − Dψ pk,i , py , (4) |Ω| indicating that the implementation is practically feasible. i∈Ω Our code is available at https://anonymous.4open.science/ where py ∈ R1×V is the one-hot target. Let P denote the space r/unlock_trainer-105B. of p. Dψ (·) is the Bregman divergence induced by a strictly The rest of this paper is organized as follows. Sections II and III convex and differentiable potential function ψ : P → R, i.e., present ZeroLock algorithm and its system design, respectively. D (u, v) = ψ(u) − ψ(v) − ⟨∇ψ(v), u − v⟩. KL divergence is ψ Experiments are in Section IV. Section V concludes this work. a specific example of Bregman divergence by selecting negative entropy as the potential function. Importantly, since minimizing II. Z ERO L OCK A LGORITHM AND A NALYSIS KL divergence is equivalent to minimizing cross-entropy, (4) A. ZeroLock Algorithm can be replaced with cross-entropy loss in practice. (ii) Consistency term aligns the outputs of chunk k to those Consider an LLM with an embedding operation E(x) and a set of transformer layers L. We introduce LoRA [30] to reduce of chunk k − 1, ensuring coherent outputs across chunks. the trainable parameters in fine-tuning. Let wl ∈ Rd1 ×d2 denote
1 X Lk→k−1 Dψ pk,i , sg(pk−1,i ) , (5) the base weights of the l-th transformer layer, which is frozen Consis (pk ) = |Ω| d3 ×d2 d1 ×d3 i∈Ω during fine-tuning. Let Al ∈ R and Bl ∈ R denote the trainable low-rank decomposition matrices in LoRA. Then, where sg(·) is stop-gradient operator with sg(p) = p and the parameter updates of transformer layer l ∈ L is represented ∇sg(p) = 0. Similarly, KL divergence can be used in practice. by ∆wl = Bl Al , so the weights after fine-tuning are wl + ∆wl . Fine-Tuning Process: This process contains T iterations. To achieve modular update decoupling in LLM fine-tuning, In each iteration, the low rank matrices Al and Bl of layer we partition the entire model into model chunks and introduce l ∈ Lk are fine-tuned using stochastic gradient descents: local objectives to these chunks for independent updates.1 ∂fk ∂fk Chunk Partitioning: Consider an LLM split into an embedAl ← Al − η , Bl ← Bl − η , (6) ∂Al ∂Bl ding operation E(·) and K chunks, where chunk k contains one or multiple consecutive transformer layers, denoted by set Lk . where η is the learning rate. Let Wk = (wl | l ∈ Lk ) and ∆Wk = (∆wl | l ∈ Lk ) denote Remark 1. According to (6), layers within the same chunk are the frozen base weights and the trainable LoRA parameters in updated using BP, while the updates of layers from different chunk k, respectively, leading to a mapping fk (·; Wk + ∆Wk ) chunks are decoupled and can be executed in parallel. This from the chunk input to output. Let x be the LLM input. Define removes the parallelism bubbles at the algorithm level (see Fig. h0 ∈ RS×d and hk ∈ RS×d as the token embedding and the 1 (b)) and avoids the need for keeping the activations during output of chunk k, respectively, with S as the sequence length: the downstream’s updates (see Fig. 1 (d)). h0 = E(x), hk = fk (hk−1 ; Wk +∆Wk ), k = 1, · · · , K. (1) B. Theoretical Analysis: Chunk-Wise Performance Readout Head: To enable local updates, we introduce a With chunk-wise performance analysis, we aim to provide frozen local readout head after chunk k = 1, · · · , K: insights into how the global objective value changes across model chunks. We focus on one token in the sequence and ⊤ z k = Norm(hk )Wlm , pk,i = softmax(z k,i ). (2) omit the token subscript. Recall that p is the readout head output of a chunk and is a distribution over vocabulary space. 1 The main idea is inspired by the local objective construction approach Let L : P → R denote the global objective given p ∈ P, i.e., in [18]. Different from [18] on classification tasks, ZeroLock handles LLM fine-tuning by incorporating LoRA and LLM-compatible head and loss design. L(p) ≜ LTask (p), which is the loss function characterizing •

how much p deviates from global ground-truth target. Let p∗k denote the optimal pk of chunk k that minimizes the local loss Lk (·), i.e., p∗k ≜ arg minp∈P Lk (p): Lemma 1 (Local Optimum). The local optimum p∗k can be equivalently represented by   1 ∗ pk = arg min ⟨∇L(pk−1 ), p−pk−1 ⟩+ Dψ (p, pk−1 ) . (7) p∈P α Proof. Recall the definition of the Bregman divergence, Dψ (u, v) = ψ(u) − ψ(v) − ⟨∇ψ(v), u − v⟩.

(8)

Substituting (u = p, v = pk−1 ) and (u = p, v = py ) into (8) respectively and rearranging the former, we have ψ(p) = ψ(pk−1 )+⟨∇ψ(pk−1 ),p − pk−1 ⟩+Dψ (p, pk−1 ), Dψ (p, py ) = ψ(p) − ψ(py ) − ⟨∇ψ(py ), p − py ⟩. Thus, Dψ (p, py )=ψ(pk−1 ) − ψ(py ) − ⟨∇ψ(py ), pk−1 − py ⟩ − ⟨∇ψ(py ), p − py ⟩ + ⟨∇ψ(py ), pk−1 − py ⟩ + ⟨∇ψ(pk−1 ), p − pk−1 ⟩ + Dψ (p, pk−1 ) = Dψ (pk−1 , py ) + ⟨∇L(pk−1 ), p − pk−1 ⟩ (9)

where (a) holds as the definition of Dψ (pk−1 , py ) and L(pi−1 ). Substituting (9) into the local loss function Lk (p) leads to Lk (p) = αDψ (pk−1 , py ) + α⟨∇L(pk−1 ), p − pk−1 ⟩ + Dψ (p, pk−1 ). Given pk−1 , term αDψ (pk−1 , py ) is a constant, so it can be omitted in optimization. Proof completes. Lemma 1 represents the local optimum p⋆ in a form of the global objective L(·), connecting local and global objectives. Meanwhile, it demonstrates a balance between moving toward the steepest descent direction of global objective L(p) and penalizing deviations from the previous distribution pk−1 . As in many existing convergence analysis on BP or BP-free (e.g., [29], [31]), we assume the smoothness of global objective. Assumption 1 (Global Objective Smoothness). The global objective L(·) is β-smooth relative to ψ in Bregman geometry: L(u) ≤ L(v)+⟨∇L(v), u−v⟩+βDψ (u, v), u, v ∈ P. (10) Then, we derive the chunk-wise performance of ZeroLock. Proposition 1 (Chunk-Wise Performance). Suppose pk−1 , p⋆k , and pk lie in a compact set of the relative interior of the simplex, and ψ(·) is locally strongly convex. Define δk ≜ Dψ (pk , p∗k ). Under Assumption 1 and α < β1 ,  K  K X 1 1X L(pK ) ≤ L(p0 )− −β Dψ (pk , pk−1 )+ δk . (11) α α k=1

k=1

Proof. Substituting u = pk and v = pk−1 in Assumption 1, L(pk ) ≤ ⟨∇L(pk−1 ), pk − p∗k ⟩ + ⟨∇L(pk−1 ), p∗k − pk−1 ⟩ |

{z

(i)

}

|

{z

(ii)

α⟨∇L(pk−1 ), pk − p⋆k ⟩ (a)

= ⟨∇ψ(pk−1 ) − ∇ψ(p⋆k ), pk − p⋆k ⟩

(b)

=δk + Dψ (p⋆k , pk−1 ) − Dψ (pk , pk−1 ).

(13)

Here, (a) holds due to the KKT condition and 1⊤ (pk − p⋆k ) = 0, where the latter holds since pk − p⋆k lies in the simplex tangent space. (b) holds based on Bregman three-point identity. According to the definition of p∗k , term (ii) satisfies ⟨g k , p∗k − pk−1 ⟩ + α1 Dψ (p∗k , pk−1 ) ≤ 0. Substituting (i) and (ii) into (12) and summing over k = 1, . . . , K complete the proof. Proposition 1 shows that the global loss L(pK ) of the final chunk tends to decrease PK as the number of chunks increases, with a bounded error k=1 δk /α resulting from the fine-tuning suboptimality of model chunks due to finite parameters.2 C. Theoretical Analysis: Algorithm Convergence

(a)

+ Dψ (p, pk−1 ),

Since p⋆k lies in the relative interior of the probability sim1 ⋆ plex, the KKT
condition of (7) gives ∇L(pk−1 )+ α ∇ψ(pk )− ∇ψ(pk−1 ) + λk 1 = 0. Accordingly, term (i) satisfies

}

+ L(pk−1 ) + βDψ (pk , pk−1 ). (12)

We provide the first analytical framework for local objective construction-based BP-free approach under general chunk division. As mentioned earlier, the major challenge comes from characterizing global convergence based on the local updates of each model chunk. To overcome this, we equivalently map local updates to global updates (Lemma 3); based on√this, we prove that the convergence rate √ of ZeroLock is Õ(1/ T ) (Theorem 1), differing from O(1/ T ) of BP by a polylogarithmic factor.3 For iteration t, let ω t = (ω tk , k = 0, 1, · · · , K) denote the parameters of chunks to be updated. We introduce this notation to generalize the analysis for various scenarios (e.g., LLM, CNN). Following the notation in Section II-A, ω t0 is the fixed parameters of operator E(·), and ω tk ≜ ∆Wkt + Wk for chunk k. Let pt0 , pt1 , . . . , ptK denote the readout head output of the associated model chunks. For analytical simplicity, we set the Bregman divergence in (4) and (5) as KL divergence. Then, the local loss is equivalent to Ltk (p; α) = αDKL (p ∥ py ) + (1 − α)DKL (p ∥ ptk−1 ). (14) Mapping from Local to Global Updates: First, we show the equivalence on local objective as follows. Lemma 2 (Objective Equivalence). Minimizing Ltk (p; α) is
∗,t equivalent to minimizing DKL p ∥ pk (α) , where
1−α py (l)α ptk−1 (l) ∗,t pk (l; α) = Pm , (15)
′ α pt ′ 1−α l′ =1 py (l ) k−1 (l ) with l denoting the l-th element of the associated vector and m denoting the length of the vector. 2 Although a relevant analysis on chunk-wise performance is presented in [18], our Proposition 1 generalizes that in [18] (i) from KL divergence to general Bregman divergence and (ii) by explicitly deriving the specific form of the error δk (rather than simply assuming an error). 3 f (T ) = Õ(g(T )) means f (T ) = O(g(T ) logc T ) for some constant c. The extra log T is a slow-varying factor that becomes negligible for large T .

Proof. Local loss Ltk (p; α) can be equivalently represented as Ltk (p; α) =

m X

p(l) log

l=1

Define Zkt (α) =

Pm

α l=1 py (l)

p(l)

1−α . py (l)α (ptk (l))

(16)

1−α ptk−1 (l) . Then,

t Ltk (p; α) = DKL (p ∥ p∗,t k (α)) − log Zk (α).

(17)

Note that for any chunk k, vectors py and ptk−1 are given, so Zkt (α) is a constant. Hence, this lemma is proved. ∗,t Note that p∗,t k (α) ≜ (pk (l; α), l = 1, · · · , m) is essentially the optimal solution that minimizes Ltk (p; α),
with proof omitted here. Thus, minimizing DKL p ∥ p∗,t k (α) is equivalent to minimizing the difference between p and the optimal solution. Then, we can represent global updates in ZeroLock as a form of local updates. We define ptk = fˆk (ptk−1 ; ω tk ) as the mapping from the readout head output ptk−1 of chunk k to ptk , given parameter ω tk of chunk k. Although this mapping cannot be directly obtained using the readout head operation (as it is not a one-to-one correspondence), it can be approximated given the statistics of ptk across chunks. Also, this mapping is introduced for theoretical analysis, and it does not need to be obtained in practice. Define the Jacobian matrix J tθk ,k ≜ ∂ fˆk (ptk−1 ; ω tk )/∂ω tk . The stop-gradient operation induces the following Jacobian matrix for all parameters: J t = blkdiag(J tω1 ,1 , J tω2 ,2 , · · · , J tωK ,K ),

(18)

where blkdiag(·) denotes block-diagonal matrix with all elements outside the diagonal blocks being zero. Define etk (α) = log ptk − log p∗,t k (α), which is the gap between the recent ptk and the optimal p∗,t k (α) in logarithmic form. Based on Lemma 2 and ptk = fˆk (ptk−1 ; ω tk ), minimizing Ltk (p; α) for all chunks over dataset D is equivalent to finding the parameter ω = (ω k , k = 0, 1, . . . , K) that minimizes "K #   X ∗,t t Rt (ω; α) = E(x,y)∼D DKL fˆk (p ; ω k )∥p (α) . (19) k−1

k

k=1

Lemma 3 (Global Update Equivalence). Based on Lemma 2, the update rule of ZeroLock can be represented as follows:4 ω t+1 ← ω t − ηt ∇ω Rt (ω t ; α),

(20)

where ∇ω Rt (ω; α) = (J t )⊤ et , et ≜ (etk , k = 0, 1, · · · , K). Proof. According to (17) in Lemma 2, given output ptk , ∇ptk Ltk (ptk ; α) = log ptk − log p∗,t k (α) + 1.

(21)

We apply the chain rule. Let ℓtk (ω tk ) ≜ Ltk (ptk ), and we omit α for presentation simplicity. For the u-th coordinate of ω k , m

m

l=1

l=1

∂ℓtk (ω tk ) X ∂Ltk (ptk ) ∂ptk (l) (a) X t ∂ptk (l) . (22) = = ek (l) ∂ω tk,u ∂ptk (l) ∂ω tk,u ∂ω k,u 4 The algorithm in [18] differs from ZeroLock by maintaining an encoder and having its output c(x) as the input for each chunk. To generalize it, we can append a column of J tωc ,k for k = 1, · · · , K before the columns in J t , where J tωc ,k is the Jacobian matrix of fkt (·) with respect to the parameters ω c of the encoder. Under this modification, Lemma 3 still holds.

Pm t Pm ∂pt (l) ∂ Equality (a) holds because l=1 ∂ωk t = ∂ω t l=1 pk (l) = u u Pm t 0, since l=1 pk (l) = 1, and (21). Finally, stacking all coordinates and the gradients of all chunks yields this lemma. PI Convergence: Define ∥a∥2D ≜ Ex∼DX [ i=1 ∥ai (x)∥22 ], PI ⟨a, b⟩D ≜ Ex∼DX [ i=1 ⟨ai (x), bi (x)⟩]. Here, I is the size of a and b, x is the input of the model, and DX is the distribution of x in dataset D. Let p(ω) = (pk (ω) ≜ fˆk (pk−1 ; ω k ), k = 0, · · · , K) denote the readout head output given ω. We rewrite Rt (ω; α) = R(ω; α, p̃t ) where pt = p(ω t ) for presentation simplicity. The following Assumptions 2–4 are commonly considered in existing analysis (e.g., [29], [31]). Assumption 5 is reasonable given the bounded gradient in Assumption 2 and the probability simplex space of p̃t . Assumption 6 is reasonable under simple random sampling and given the bounded gradient. Assumption 2 (Bounded Gradient). The gradient is upper bounded, i.e., E[∥∇ω Rt (ω; α)∥2 ] ≤ M (α). Assumption 3 (Lipschitz Continuity). Readout output p(ω) is Lipschitz continuous with respect to ω, i.e., ∥p(ω)−p(ω ′ )∥D ≤ ρ∥ω − ω ′ ∥, for all ω and ω ′ in parameter space. Assumption 4 (Smoothness). For any ω in parameter space, Rt (ω; α) is βR (α)-smooth, i.e., Rt (ω ′ ; α) ≤ Rt (ω; α) + ⟨∇Rt (ω; α), ω ′ − ω⟩ + βR2(α) ∥ω ′ − ω∥2 . Assumption 5 (Stability). The first-order objective variation is controlled alongthe distribution update, i.e.,  quadratically ∇p̃t R(ω t ; α, p̃t ), p̃t+1 − p̃t D + ≤ Cfo (α)∥p̃t+1 − p̃t ∥2D . Assumption 6 (Sampling). The mini-batch sampling is unbiased and has a bounded variance, i.e., ||∇ω R̂t (ω; α) − ∇ω Rt (ω; α)||2 ≤ σ 2 (α)/B, where ∇ω R̂t (ω; α) is the gradient under sampled mini-batch and B is the mini-batch size.   Define drift δtref (α) ≜ Rt+1 (ω t+1 ; α) − Rt (ω t+1 ; α) + . This drift exists as distribution ptk−1 from chunk k −1 changes, and characterizes the change of Rt (·) across iterations. Lemma 4 (Bounded Drift). The drift of Rt (·) is PTcumulative −1 upper bounded, i.e., DT (α) ≜ t=0 E[δtref (α)] ≤ (Cfo (α) + PT −1 βR (α)/2)ρ2 M (α) t=0 ηt2 . Proof. According to the definition of δtref (α), by substituting ω ′ = ω t+1 and ω = ω t in Assumption 4, δtref (α) ≤ (Cfo (α)+ βR (α) t+1 − p̃t ∥2D . Based on Lipschitz continuity of p̃t ≜ 2 )∥p̃ t p̃(ω ) in Assumption 3, ∥p̃t+1 − p̃t ∥2D ≤ ρ2 ∥ω t+1 − ω t ∥2 . Let Cref ≜ (Cfo (α) + βR (α)/2)ρ2 . Based on Lemma 3 and Assumption 2, E[δtref ] ≤ Cref (α)E[∥ω t+1 − ω t ∥2 ] = Cref (α)ηt2 E[∥∇ω Rt (ω; α)∥2 ] ≤ Cref (α)M (α)ηt2 . Summing over t = 0, . . . , T − 1 completes the proof. √ Theorem 1 (Convergence). Let ηt = η0 / t + γ ≤ 1/βR (α), where γ > 0. Based on Lemma 4 and Assumptions 4–6, T −1

1 X 2(R0 − Ropt + DT (α)) √ E∥∇ω Rt (ω; α)∥2 ≤ T t=0 T η0 / T + γ PT −1   βR (α)σ(α)2 t=0 ηt2 /B 1 √ + = Õ √ , (23) T η0 / T + γ T

Coordinator

micro-batch updates across stages; failure recovery for tracking the state of committed requests and assigning replay requests Forward Update window Commited state Registration Micro-batch schedular Replay request Heartbeat when failure detected. In the scheduler, update window defines Hidden state a contiguous sequence of micro-batches that executors are Executor 0 Executor 1 Readout head (Stage 0) (Stage 1) allowed to and must handle before advancing its process Membership Membership and serving as the unit of checkpoint; micro-batch scheduler Local Execution monitor Execution monitor backward indicates the sequence of micro-batches in the update window Failure recovery Failure recovery update and will assist with the micro-batch replay after failure detected. Stage Worker Stage Worker input Hidden ... state CheckCheckCheckEach executor contains both control plane and data plane. Buffer Buffer Buffer point point point Input The control plane contains three components: membership for reporting its registration and hearbeat; execution monitor for Fig. 2. System overview with one coordinator and multiple executors, each monitoring its execution status; failure recovery for execution corresponding to a stage (i.e., the update of a model chunk). replay. The data plane at an executor is responsible for handing a stage, i.e., the update of a model chunk. It retrieves either where R0 and Ropt are the initial and optimal R(·). the model input or the hidden state of the upstream stage from Proof. For presentation simplicity, we omit notation α in Rt (·). its buffer, performs forward pass to determine its hidden state, By Assumption 4 and considering update with sampled mini- transfers the hidden state to the buffer of the downstream stage, and performs a backward update based on its local objective. batch, i.e., ω t+1 ← ω t − η t ∇R̂t (ω t ), we have We propose the following techniques to achieve lightweight, t+1 t t t high-throughput, and robust implementation. (I) Early ForRt (ω ) ≤ Rt (ω ) − η ⟨∇Rt (ω ), ∇R̂t (ω)⟩ warding. The forwarding of the hidden state is designed to βR (α)ηt2 + ∥∇R̂t (ω t )∥2 . (24) perform ahead of the local updates, avoiding unnecessary inter2 stage waiting. (II) Independent Execution and Checkpoint. Taking conditional expectation and using Assumption 6 and Each executor maintains only its own trainable parameters, ηt ≤ 1/βR (α), we have optimizer state, and checkpoint, enabling recovery at the stage 2 2 granularity, and independently constructs its local signal and ηt βR (α)ηt σ(α) E[Rt (ω t+1 )] ≤ E[Rt (ω t )]− ∥∇Rt (ω t )∥2 + . perform local updates. (III) State-Only Inter-Stage Exchange. 2 2B (25) Executors exchange only the hidden state (but not the gradients Based on the definition of δtref (α), Rt (ω t+1 ) ≥ Rt+1 (ω t+1 )− or optimizer states) from other stages. (IV) Buffer-Assisted δtref (α). By substituting Rt (ω t+1 ) into (25) and rearranging, State Exchange. Each executor maintains a bounded hidden state buffer, storing the hidden states of its upstream stage ηt E∥∇Rt (ω t )∥2 ≤ E[Rt (ω t )] − E[Rt+1 (ω t+1 )] for iterations, enabling replay after failure. Overall, (I)-(III) 2 contribute to lightweight and high-throughput; (II) and (IV) 2 2 βR (α)ηt σ(α) + + E[δtref (α)]. (26) contribute to failure recovery capability. 2B Summing from t = 0 to T − 1 and based on Lemma 4 and B. System Design Details the definition of Ropt , the proof is complete. We first present the execution interfaces. Then, we introduce Theorem 1 √ shows that ZeroLock algorithm has a convergence inter-stage runtime, in-stage execution, and failure recovery. rate of √ Õ(1/ T ), which differs from the convergence rate 1) Stage-Execution-Related Interface: The following defines O(1/ T ) of conventional BP by only polylogarithmic factor. the unified interfaces related to stage execution at the executors. Meanwhile, the study in this section provides the first analytical Input Interface. For stage 0 executor, it inputs either token framework for the local objective construction-based approach IDs or pre-computed initial hidden representation (constructed under general model chunk division. via a local input embedding layer). For each of the subsequent III. S YSTEM D ESIGN stages, it inputs the hidden states of the upstream stage from the We propose a pipeline training system for ZeroLock. It is hidden state buffer, filled by its upstream executor. Meanwhile, applicable to both multi-GPU server and multi-device scenarios. all stages input the attention mask, position IDs, and labels We first present the overview and detail the system design. Then, associated with a specific micro-batch to facilitate local training. Forward-and-Output Interface. The executor performs we discuss the specific design for parallelism at mobile devices. forward pass using its chunk to produce its hidden state, which A. System Overview is then detached from the computation graph and written to As shown in Fig. 2, this system contains one coordinator and the outbound buffer for sending to the downstream stage. multiple executors (i.e., GPUs or devices), and is separated into Local Update Interface. The executor utilizes the input control plane and data plane. The coordinator falls in the control hidden states and labels to construct a local loss and performs plane and has three components: membership for tracking a backward pass and optimizer step to update its trainable active executors; scheduler for managing the scheduling of parameters. This interface supports two ways to construct loss:

Data Control plane plane

Membership

Scheduler

Failure Recovery

Stage Worker

(a) Static Model Head. It reuses the pre-trained model’s final normalization layer and frozen model head as the readout head, and the token-level cross entropy and task-specific labels are used to construct local loss, as in Section II-A. This design does not induce new trainable parameters, while it assumes the intermediate hidden states can be interpretable by the final head, which may not hold for complex generative tasks. (b) Readout Adapter. To bridge the gap between the intermediate hidden states and the frozen model head’s readout space, executor can insert a stage-wise readout adapter (i.e., a small residual MLP initialized approximate the identity mapping). During training, this adapter aligns the intermediate hidden space to be more amenable to the frozen model head. Besides (a) and (b), this interface also accommodates other lightweight head, e.g., bottleneck projection head or restricted vocabulary head, to reduce computation and memory overhead. 2) Inter-Stage Runtime and In-Stage Execution: The interstage runtime (acting as a scheduling engine) organizes multiple micro-batches into an update window at the coordinator. These micro-batches enter the stage pipeline sequentially. For each micro-batch, the stage worker performs in-stage execution:

Proactive Checkpointing: Each executor periodically captures local checkpoints, including trainable parameters, optimizer states, and progress metadata. • Reactive Recovery: Upon detecting a failure, the runtime restores the failed stage from its latest checkpoint and replays necessary update windows to catch up to the pipeline frontier. Successful stages retain their execution.

•

C. Mobile Execution Backend

For deployment at mobile devices, we use the ExecuTorch framework, Meta’s official PyTorch-native edge runtime. It performs ahead-of-time (AOT) compilation to produce a static computational representation, serialized as a .pte program, which enables deterministic execution on resourceconstrained end devices without a Python interpreter. However, ExecuTorch bundles forward and backward computations into a single Python ExecuTorch (PTE) method, where the forward hidden state is exposed only after the local backward gradient computation. This prohibit the early-forwarding opportunity. To address this, we insert a lightweight pipeline marker operator within the PTE method—specifically after the detached hidden state output but before the parameter-gradient subgraph. S1 Input: Retrieve hidden states of the upstream stage from This operator performs no tensor computation; instead, it the hidden state buffer; load the attention mask, position signals the runtime to capture the hidden outputs and suspend IDs, and labels related to the micro-batch. the PTE method while preserving its execution state (stack, S2 Forward-and-Output: Perform forward pass to produce its tensors, optimizer buffers). We extend the ExecuTorch training hidden state; output it to the outbound buffer. runtime with a two-phase Android interface. The first phase S3 Local Training: Backpropagate the local loss and accumu- executes up to the marker and returns the hidden state, which late gradients for the current update window. The optimizer is immediately transferred to the downstream device. The step is applied after the final micro-batch in that window. second phase resumes the same method to complete the local Note that downstream execution can start immediately after backward pass, after which an on-device AdamW optimizer S2 without waiting for the upstream’s local backward in S3, updates the stage’s LoRA parameters. In this case, the transfer to the downstream device overlaps with the local backward leading to a higher degree of parallelism. computation, making the resulting stage program support local Across stages, inter-stage runtime is responsible for managtraining and early-forwarding without incurring cross-stage ing the hidden state flow. The management is in the form of gradient Remote Procedure Call (RPC) overhead. an entry, which contains the hidden state tensor of a microbatch and the related metadata for sending, receiving, and IV. E XPERIMENTS status. The runtime tracks each entry’s transition from the We build prototype systems using both multi-GPU/CPU outbound buffer of upstream stage to the hidden state buffer of downstream stage, allowing the coordinator to accurately server and Android devices respectively. The evaluation on control the buffer occupancy of each stage without tracking multi-GPU/CPU server shows (E1) the method comparison on each tensor. Meanwhile, we introduce an in-flight depth to memory and throughput and (E2) failure recovery. The evaluarestrict the number of entries between each pair of stages, tion on Android devices shows (E3) on-device performance. preventing upstream stages from overwhelming downstream stages. Furthermore, on GPU, preposted receive is allowed, A. Experimental Settings such that downstream stage can submit receive_entry request Model and data. We use TinyLlama as the pre-trained model. in advance. Each entry binds to a readiness event. Compute Its Transformer layers are partitioned into three consecutive stream waits only for this specific event, avoiding unrelated chunks, each assigned to one NVIDIA L40 GPU or one Andriod communication operations from blocking the compute stream. phone. We apply LoRA with rank 4 and scaling factor 16. On 3) Failure Recovery: The failure recovery capability comes default, experiments are conducted on a fixed 10,000-example from two facts: each stage checkpoints its state independently; subset of AG News; all sequences are padded or truncated the hidden state buffer introduces resilience to failure. Conse- to 128 tokens. Let b denote the number of physical batches quently, after a failure, only the failed stages need to roll back, processed by one microbatch call; m denotes the number of which retrieves hidden state from the buffer and replaying the such microbatch calls before an optimizer update; thus, batch requests. The brief failure recovery mechanism is as follows: B = bm. We repeat the experiment with three random seeds.

50

10

0

Throughput (req/s)

BP (Single GPU) GPipe 1F1B PipeDream ZeroLock (α = 1) ZeroLock (α = 0.5)

Max per-stage memory

75

NLL

Accuracy (%)

B = 32 4 ZeroLock GPipe

1F1B PipeDream

2 0

800

1200

0

Optimizer updates

400

800

1200

1/32

Optimizer updates

Fig. 4. Peak GPU memory usage and the memory usage of each component under the peak (b = 8 and m = 4).

1/32

Max per-stage memory

15

B = 128

(a)

5

8/16

4/8

8/4

(b)

10

4/32

2/16

Batch geometry b/m

16/8

B = 128 200 100 0

32/4

4/32

Batch geometry b/m

8/16

16/8

32/4

Batch geometry b/m

(c)

(d)

Fig. 5. (a) and (c), maximum stage memory usage under B = 32 and B = 128, respectively; (b) and (d), throughput under B = 32 and B = 128, respectively. When b/m = 8/4, when compared with PipeDream, ZeroLock reduces the memory by 26.5% and improves throughput by 4.9%.

250 200 150 2

3

4

Pipeline stages / GPUs

Methods. We compare ZeroLock with three approaches. (i) GPipe [3] completes all forward passes in a logical batch before the backward pass. (ii) 1F1B [4] interleaves forward and backward within a logical batch and flushes the pipeline before the shared optimizer update. (iii) PipeDream [4] maintains 1F1B execution across update windows via weight stashing (i.e., storing multiple copies), ensuring forward and backward passes use consistent parameter versions.5

0

8/4

0

Throughput (req/s)

17.1 0.0 94.0 21.0 0.0 84.5 21.0 0.0 85.5 5.1 0.0 22.0 13.0 0.0 41.8 21.0 0.0 66.8 4.3 0.8 22.0 12.5 1.4 42.9 16.8 1.4 65.1 4.3 0.0 22.0 4.3 0.0 19.0 4.3 0.0 20.0 . k e n d er a he cv w io at Oth l pe l + st ctivat loss/b put/re ht cac a t / A To ig In de Fwd We Mo

4/8

100

Batch geometry b/m

Fig. 3. (a) Accuracy and (b) negative-likelihood loss.

GPipe S2 4179.1 799.4 2625.3 643.3 GPipe S1 3034.5 590.1 2247.1 91.8 GPipe S0 3128.5 715.1 2135.1 171.8 1F1B S2 1935.6 799.4 656.3 452.9 1F1B S1 1839.7 590.1 1112.5 82.2 1F1B S0 2546.0 715.1 1595.8 147.3 PipeDream S2 1935.6 800.2 656.3 452.1 PipeDream S1 1840.9 590.1 1112.5 81.5 PipeDream S0 2541.4 715.8 1595.8 146.6 ZeroLock S2 1868.3 799.4 628.3 414.4 ZeroLock S1 1699.8 715.1 547.3 414.1 ZeroLock S0 1825.8 840.1 547.3 414.1

2/16

Throughput (req/s)

400

Throughput (req/s, log)

0

B = 32 200

(a)

10

10

2

1

ZeroLock GPipe

Local

1F1B PipeDream

Wi-Fi Mobile Constr.

Sender-side link profile

(b)

Fig. 6. Throughput under diverse scenarios: (a) number of stages, each executed by a GPU; (b) diverse sender-side links.

end-to-end backpropagation, which need to retain activations for backward gradient computation. Meanwhile, given m = 4, GPipe, 1F1B, PipeDream, and ZeroLock experience out-ofmemory when b = 10, 20, 20, 28, respectively, showing that the system with ZeroLock can afford a larger batch size. B. E1: Memory and Throughput Figs. 5 (a) and (c) show that ZeroLock is more beneficial Despite the modular update decoupling, ZeroLock has a when the physical batch to microbatch ratio (i.e., b/m) is comparable accuracy and negative log-likelyhood (NLL) to larger. Specifically, more physical batches (in a microbatch call) baseline approaches (see Fig. 3). Build upon this, we show its indicates larger forward graphs retained by baseline approaches, leading to more obvious benefits for modular decoupling. memory reduction and throughput improvement as follows. 2) Throughput: Figs. 5 (b) and (d) and Fig. 6 show the 1) Memory: As shown in Fig. 4, when compared with GPipe, 1F1B, and PipeDream, ZeroLock reduces the mean per- throughput comparison. The default setting contains three stage peak memory by 47.8%, 14.7%, and 14.6%, respectively, stages, with b = 8 and m = 4. First, in Fig. 5 (b), under and by 55.3%, 26.6%, and 26.5% for the maximum stage the default setting, when compared with GPipe, 1F1B, and usage, respectively. Most of the reduction comes from the PipeDream, ZeroLock improves the throughput by 55.8%, elimination of activations, where the reduction are 75.4%, 62.8%, and 4.9%, respectively. Second, in Figs. 5 (b) and 40.8%, and 48.8% on per-stage average, respectively. This (d) and Fig. 6 (a), as the increase of the physical batches reduction results from the modular update decoupling of in a microbatch call and stages, ZeroLock and PipeDream ZeroLock; in comparison, the other approaches implement demonstrate more significant increase in throughput, showing the better potentials for reducing bubbles in each call and ex5 Although existing works [26]–[29] considered modular decoupling in ploiting additional GPUs, respectively. Finally, Fig. 6 (b) shows pipeline parallelism, [28] is not for LLM fine-tuning; [26], [27] do not provide the throughput under simulated sender-side communication open-sourced code; the open-sourced code of [29] does not implement systemlevel pipeline parallelism (e.g., it places all chunks of a certain batch in one link: Wi-Fi, 2 ms/1000 Mbps (latency/bandwidth); mobile, 10 GPU), so comparing it with those deployed in multiple GPUs may not be fair. ms/200 Mbps; constrained, 30 ms/50 Mbps. ZeroLock is more

CUDA kernel

TABLE II O N - DEVICE EVALUATION WITH A NDRIOD PHONES .

Idle

S0

S0

S1

S1

Stage/device

PSS

Temp.

S2

S2

S0/NX809J 11.43/13.56 12.44/17.56 38.28/46.90 3423.5 S1/Lenovo L71091 7.56/8.32 0/0.02 15.61/27.51 3408.2 S2/Pixel 10 Pro XL 9.53/11.40 0/5.79 10.30/18.99 3824.3

37.0 36.0 36.7

500

1k

0

S0

S0

S1

S1

S2

500

1k

500

1k

S2

0

500

1k

0

Forward Local backward Optimizer

Queue

Other local Queued wait

12.2 11.4

S0 active S0 wait

Time in one-second interval (ms)

Fig. 7. CUDA activity: (a) GPipe; (b) 1F1B; (c) PipeDream; (d) ZeroLock.

S1 active S1 wait

0.2

S2 active S2 wait

TABLE I

10.0

0.9

0

FAILURE RECOVERY.

8.0

Recovery (ms)

Transfer (MiB)

Synchronous 1F1B ZeroLock

2381.2 ± 66.6 2013.1 ± 6.4

192 96

Total

S0 S1 S2

15 10 5 0

5

10

Mean stage time (s)

Method

Stage-local loss

0

PTE

(a)

0

50

100

Local optimizer step

(b)

Fig. 8. (a) Latency breakdown; (b) Local loss convergence.

D. E3: On-Device Evaluation with Andriod Phones For implementation on Android devices, we export three beneficial when the link is slower. This is because baselines TinyLlama training PTEs containing transformer layers [0,6], need to send both forward hidden state and backward hidden [7,13], and [14,21] and deploy them on NX809J, Lenovo gradients, while ZeroLock transfers only forward hidden state. L71091, and Pixel 10 Pro XL, respectively. We use 128 data Fig. 7 visualizes the CUDA activity traces under b = 8, samples for training with sequence length 128, b = 1, LoRA m = 4, 512 data samples, and 16 optimizer updates. The plots rank 8, scale 16, and learning rate 10−4 . Each stage commits show a one-second steady-state interval, during which all three GPUs exhibit compute activity; this interval selection was fixed 128 optimizer steps and writes nine checkpoints. Table II shows the median/95th percentile PTE, queuing, and prior to inspecting method-specific behavior to avoid bias. The total duration (in seconds) for a single batch request with the Kernel and device copy timestamps are measured externally, latency break down shown in Fig. 8 (a), the application peak without injecting device-side synchronization. As shown in the PSS in MiB, and battery temperature in ◦ C. The total wall-clock figure, ZeroLock and PipeDream demonstrate more concurrent stage activities and less GPU idle time. Furthermore, when time of the entire fine-tuning is 1644.1 s with a throughput of compared with PipeDream, ZeroLock shows a more balanced 0.0779 records/s. Fig. 8 (b) shows the loss convergence. These activity load across stages and a higher active fraction, with results indicate the implementation is practically feasible. an improvement of 13.7%. Note that this active fraction serves V. C ONCLUSION as a diagnostic indicator of pipeline scheduling rather than a proxy for throughput, as idle periods also encompass host In this work, we propose a ZeroLock algorithm, which dispatch and Gloo communication overhead. achieves modular update decoupling by local objective construction, for pipeline parallelism in LLM fine-tuning. It effectively C. E2: Failure Recovery removes pipeline bubbles by alleviating the inter-stage waiting We evaluate the outage recovery capability of ZeroLock. It and reduces memory usage by mitigating unnecessary activation is compared with 1F1B. Specifically, 1F1B recovers via global storage at the algorithm level. We provide the first analytical rollback and full replay. In contrast, ZeroLock enables local framework for local objective construction-based approaches recovery by reusing buffered hidden states and retaining com- under general chunk division √ and prove that ZeroLock has a mitted previous updates, eliminating redundant recomputation. convergence rate of Õ(1/ T ), comparable to BP training. We consider three stages, and set b = 1 and m = 8. The Meanwhile, we design a system for deploying ZeroLock, failure setting is as follows: W0–W3, four prelude windows; incorporating techniques such as early forwarding and failure W4–W7, four failure windows at Stage 1; W8–W11, four recovery to improve system throughout and robustness. We resumed windows. The worker processes and GPU contexts build real-world prototype and show that when compared with remain alive during the failure. Recovery latency is defined BP-based baselines, ZeroLock reduces the memory by 26.5% as the during between when Stage 1 resumes and when Stage and improves throughput by 4.9%. For future direction, it is 2 commits W7. In Table I, ZeroLock reduces the recovery meaningful to further incorporate operator-level optimization latency by 368 ms and the transfer traffic from 192 to 96 MiB. to further improve throughput and reduce memory usage.

VI. U SE OF AI D ISCLOSURE ChatGPT was used to assist with code development and experimental setup. Specifically, it was used to configure and troubleshoot the environments required for benchmark reproduction, run and migrate benchmark implementations across different environments, generate Kotlin code for deploying and executing PTE models on mobile devices, and generate batch experiment scripts using torchrun. The AI-generated code and configuration instructions were reviewed, adapted where necessary, tested, and validated by the author. The author takes full responsibility for the final implementation, experimental results, and their interpretation. R EFERENCES [1] A. Tandon, K. Dalal, X. Li, D. Koceja, M. Rød, S. Buchanan, X. Wang, J. Leskovec, S. Koyejo, T. Hashimoto et al., “End-to-end test-time training for long context,” arXiv preprint arXiv:2512.23675, 2025. [2] A. M. Alghamdi, M. U. Ashraf, A. A. Bahaddad, K. A. Almarhabi, W. A. Al Shehri, and A. Daraz, “Cross-subject eeg signals-based emotion recognition using contrastive learning,” Scientific Reports, vol. 15, no. 1, p. 28295, 2025. [3] Y. Huang, Y. Cheng, A. Bapna, O. Firat, M. X. Chen, D. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, and Z. Chen, “GPipe: Efficient training of giant neural networks using pipeline parallelism,” in Proc. NeurIPS, Vancouver, BC, Canada, Dec. 2019. [4] D. Narayanan, A. Harlap, A. Phanishayee, V. Seshadri, N. R. Devanur, G. R. Ganger, P. B. Gibbons, and M. Zaharia, “PipeDream: Generalized pipeline parallelism for DNN training,” in Proc. ACM SOSP, Huntsville, ON, Canada, Oct. 2019. [5] Y. Chen, Y. Yan, S. Ge, Y. Qin, Y. Zheng, Q. Yang, S. He, Z. Shi, J. Chen, and Y. Shu, “Confidant: Customizing transformer-based llms via collaborative training on mobile devices,” in Proc. ACM MobiCom, 2025. [6] D. Narayanan, A. Phanishayee, K. Shi, X. Chen, and M. Zaharia, “Memory-efficient pipeline-parallel DNN training,” in Proc. ICML, Virtual Conf., Jul. 2021. [7] S. Fan, Y. Rong, C. Meng, Z. Cao, S. Wang, Z. Zheng, C. Wu, G. Long, J. Yang, L. Xia, L. Diao, X. Liu, and W. Lin, “DAPPLE: A pipelined data parallel approach for training large models,” in Proc. ACM PPoPP, Virtual Conf., Feb.–Mar. 2021. [8] D. Narayanan, M. Shoeybi, J. Casper, P. LeGresley, M. Patwary, V. A. Korthikanti, D. Vainbrand, P. Kashinkunti, J. Bernauer, B. Catanzaro, A. Phanishayee, and M. Zaharia, “Efficient large-scale language model training on GPU clusters using Megatron-LM,” in Proc. ACM/IEEE SC, St. Louis, MO, USA, Nov. 2021. [9] S. Li and T. Hoefler, “Chimera: Efficiently training large-scale neural networks with bidirectional pipelines,” in Proc. ACM/IEEE SC, St. Louis, MO, USA, Nov. 2021. [10] P. Qi, X. Wan, G. Huang, and M. Lin, “Zero bubble (almost) pipeline parallelism,” in Proc. ICLR, Vienna, Austria, May 2024. [11] X. Lian, S. A. Jacobs, L. Kurilenko, M. Tanaka, S. Bekman, O. Ruwase, and M. Zhang, “Universal checkpointing: A flexible and efficient distributed checkpointing system for large-scale DNN training with reconfigurable parallelism,” in Proc. USENIX ATC, Boston, MA, USA, Jul. 2025. [12] S. Gandhi and C. Kozyrakis, “Sparse checkpointing for fast and reliable MoE training,” in Proc. USENIX NSDI, Renton, WA, USA, May 2026. [13] T. Wu, L. Cao, H. Lu, X. Jiang, Y. Yu, S. Yang, G. Yang, J. Wang, L. Qu, L. Zhang, and W. Wang, “Attack of the bubbles: Straggler-resilient pipeline parallelism for large model training,” in Proc. USENIX NSDI, Renton, WA, USA, May 2026. [14] J. Chen, X. Deng, Z. Xiong, S. Guo, X. Qiu, P. Wang, and D. Niyato, “CollaPipe: Adaptive segment-optimized pipeline parallelism for collaborative LLM training in heterogeneous edge networks,” arXiv:2509.19855, Sep. 2025. [15] M. Ryabinin, T. Dettmers, M. Diskin, and A. Borzunov, “SWARM parallelism: Training large models can be surprisingly communicationefficient,” in Proc. ICML, Honolulu, HI, USA, Jul. 2023.

[16] A. Borzunov, D. Baranchuk, T. Dettmers, M. Ryabinin, Y. Belkada, A. Chumachenko, P. Samygin, and C. Raffel, “Petals: Collaborative inference and fine-tuning of large models,” in Proc. ACL Syst. Demonstrations, Toronto, ON, Canada, Jul. 2023. [17] Z. Huo, B. Gu, qian Yang, and H. Huang, “Decoupled parallel backpropagation with convergence guarantee,” in Proc. ICML, Jul. 2018. [18] C. Ye, R. Ye, Y. Zhang, and M. Tang, “Depth-progressive monotonic learning without global backpropagation,” in Proc. ICML, 2026. [19] R. Ye, C. Ye, C. Huang, M. Tang, and Y. Liu, “Beyondbackpropagation training: Methods, applications, and perspectives,” TechRxiv, vol. 2026, no. 0103, 2026. [Online]. Available: https: //www.techrxiv.org/doi/abs/10.36227/techrxiv.176740426.63642005/v1 [20] A. Nøkland, “Direct feedback alignment provides learning in deep neural networks,” Proc. NeurIPS, 2016. [21] M. Akrout, C. Wilson, P. Humphreys, T. Lillicrap, and D. B. Tweed, “Deep learning without weight transport,” in Proc. NeurIPS, vol. 32, 2019. [22] S. Malladi, T. Gao, E. Nichani, A. Damian, J. D. Lee, D. Chen, and S. Arora, “Fine-tuning language models with just forward passes,” in Proc. NeurIPS, 2023. [23] S. Wang, Z. Chen, and M. Tang, “Curvzo: Adaptive curvature-guided sparse zeroth-order optimization for efficient llm fine-tuning,” in Proc. ICML, 2026. [24] Q. Li, Y. W. Teh, and R. Pascanu, “Noprop: Training neural networks without back-propagation or forward-propagation,” in Conference on Lifelong Learning Agents. PMLR, 2026, pp. 525–544. [25] B. Millidge, A. Tschantz, and C. L. Buckley, “Predictive coding approximates backprop along arbitrary computation graphs,” Neural Computation, vol. 34, no. 6, pp. 1329–1368, 2022. [26] X. Guo, C. Xu, G. Guo, F. Zhu, C. Cai, P. Wang, X. Wei, J. Su, and J. Gao, “Faster multi-GPU training with PPLL: A pipeline parallelism framework leveraging local learning,” arXiv:2411.12780, Nov. 2024. [27] M. Aljahdali and M. Canini, “Idle no more: Boosting distributed pipeline training via FluidPipe,” in Proc. IEEE ICDCSW, Glasgow, U.K., Jul. 2025. [28] M.-Y. Ho, C.-K. Wang, Y.-T. Lin, and H.-H. Chen, “Scpl: Enhancing neural network training throughput with decoupled local losses and model parallelism,” ACM Trans. Manage. Inf. Syst., vol. 17, no. 2, Jun. 2026. [29] H. Shi, T. Han, P. Wang, Z. Wang, X. Yang, and J. Su, “Rethinking local learning: A cheaper and faster recipe for LLM post-training,” arXiv:2605.04913, May 2026. [30] E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen, “LoRA: Low-rank adaptation of large language models,” in Proc. ICLR, Virtual Conf., Apr. 2022. [31] P. Han, C. Huang, G. Tian, M. Tang, and X. Liu, “Convergence analysis of split federated learning on heterogeneous data,” Advances in Neural Information Processing Systems, vol. 37, pp. 103 476–103 544, 2024.

Related documents

• AIM-ZO: Activation-Informed Subspace Maintenance for Zeroth-Order LLM Fine-Tuning#1129851
• SubZero+: Efficient Zeroth-Order LLM Fine-Tuning via Large Learning Rates#616813
• SubZero+: Efficient Zeroth-Order LLM Fine-Tuning via Large Learning Rates#616313
• TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems#6521
• Fast and Memory-Efficient Wavelet Convolutions via I/O-Aware Reformulation#444696