arXiv CSopen access

Poseidon: DAG-Guided Parallelism Search for LLM Pre-Training on Heterogeneous Clusters

arXiv:2609.06086v1 [cs.DC] 5 Sep 2026

Poseidon: DAG-Guided Parallelism Search for LLM Pre-Training on Heterogeneous Clusters Xiaosong Chen∗ University of Macau [email protected]

Shaoheng Nie∗ Fudan University [email protected]

Zhongmin Zhao University of Macau [email protected]

Zizhao Mo University of Macau [email protected]

Jiapeng Chen Fudan University [email protected]

Huanle Xu† University of Macau [email protected]

Zeren Li† Independent researcher [email protected]

Weiwei Sun† Fudan University [email protected]

ChengZhong Xu University of Macau [email protected]

∗ Equal contribution.

† Corresponding authors.

memory can vary significantly within the same accelerator class, leading to diverse training performance. For example, limited memory constrains batch sizes on older accelerators, while differences in computational power create significant gaps in training time for identical workloads. To optimize cluster resource utilization, pre-training LLMs on heterogeneous accelerator devices has become increasingly important. Several studies have explored this problem [15, 43, 49], focusing on optimizing parallelization configurations to minimize training time on such resources. These configurations typically combine pipeline, tensor, and data parallelism to improve computational efficiency, reduce memory usage, and lower communication overhead. However, accelerator heterogeneity introduces significant complexity in identifying optimal parallelization configurations. This is because configuration searching is exposed to a new dimension, where the solution space grows exponentially with the number of accelerator types involved. In witnessing this dilemma, existing works propose heuristic methods to balance searching time and training throughput. Although they reduce the searching time to some extent, we identify three key obstacles that can lead to low training efficiency. First, the cost models utilized by existing works to capture the training time are inaccurate. Since they cannot precisely characterize the overlap among concurrently proceeding operations, we demonstrate up to 1.32× deviation in our experimental testbed. Second, the heuristic pruning strategies, which serve as the key to reducing searching time, risk neglecting the optimal configuration. This is because the rigid strategies proposed cannot flexibly adapt to various cluster environments, where inter-accelerator connections and computational power gaps can be highly diverse. In this sense, we witness up to 1.91× degradation when evaluating the pruning strategies for baseline systems. Third, existing modeling techniques present low generalizability, due to the tight coupling with a specific training mode, i.e., 1F1B [28], where each pipeline stage alternates between executing one

Abstract With the rapid advancement of accelerator technologies, pre-training large language models (LLMs) on heterogeneous accelerator clusters has become increasingly crucial for maximizing hardware utilization. Existing systems, however, suffer from inaccurate training time modeling, which undermines the parallelization optimizations built upon it. Moreover, for current approaches, the vast configuration search space makes exhaustive exploration infeasible, forcing a trade-off between search time and training efficiency. To overcome these limitations, we introduce Poseidon, an efficient and scalable LLM training framework designed with heterogeneity awareness. Its core is an explicit training time model based on a directed acyclic graph. Building on this graph, Poseidon employs two efficient, theoretically grounded strategies: stage-level pruning via early stopping with partial estimation, and layer-to-stage mapping exploiting a ridge-like distribution pattern. These strategies reduce the search space without sacrificing optimal training efficiency. Experiments on heterogeneous clusters show that Poseidon improves training throughput by up to 2.76× over state-of-the-art systems.

1

Introduction

LLMs [1, 22, 41] have demonstrated unprecedented capabilities, advancing applications such as code completion [9, 42], story writing [52], and conversational agents [54]. However, the scaling law drives the development of increasingly larger foundational models to fully exploit their potential, requiring substantial computational power. For instance, training LLaMA-3 requires 3.8 × 1025 FLOPs, as reported by Meta [8]. This has led to widespread deployment of accelerators such as GPUs [26] and NPUs [5] to support these resource-intensive workloads. Meanwhile, accelerators evolve through multiple generations, resulting in notable inter-generation heterogeneity. Computational power and 1

forward pass and one backward pass of micro-batches after pipeline warm-up. As a result, popular training variants, including Eager 1F1B [56] and Interleaved 1F1B [30], cannot be well supported by this design. In this paper, we propose Poseidon, a novel LLM pretraining system designed to address the aforementioned limitations. Poseidon efficiently explores optimal parallel configurations across heterogeneous accelerator resources while maintaining low search overhead. To achieve this, it incorporates an explicit directed acyclic graph (DAG)-based training time estimation model that captures stage- and batch-level execution dependencies. This graph-based representation models the data execution order within each stage and the dependencies across micro-batches for arbitrary parallel training configurations, using distinct node and edge types. It enables training time estimation over heterogeneous resources by computing the longest path. Furthermore, the graph accounts for communication-computation overlap in heterogeneous environments, significantly enhancing the accuracy of training time predictions. Leveraging this estimation framework, Poseidon integrates two optimization strategies to efficiently navigate the vast configuration search space. First, it employs stage-level pruning with early stopping, guided by a custom lower-bound function. This mechanism terminates the evaluation of suboptimal configurations once partial stage analysis reveals inefficiencies, avoiding exhaustive search. Second, Poseidon exploits a ridge-like distribution pattern in layer-to-stage mapping to prune suboptimal regions of the search space. This approach, backed by theoretical guarantees, substantially accelerates search while guaranteeing optimal solutions. Together, these strategies enable Poseidon to achieve unprecedented efficiency in identifying optimal parallel configurations for LLM training over heterogeneous clusters. In summary, we have made the following contributions: Ž We introduce a novel DAG-based cost model for LLM training, capturing various factors and therefore enabling precise training time estimation. This model also offers high generalizability to multiple training variants. Ž We propose two pruning strategies for configuration exploration that are theoretically guaranteed to identify the optimal parallel configuration. Ž We evaluate Poseidon via extensive experiments on heterogeneous clusters. Results show that it consistently outperforms baselines, achieving up to 2.76× higher throughput.

2

Background and Motivation

2.1

LLM Training under Accelerator Heterogeneity

Micro-batches

Stage 1

2

1

1

Stage 2

Estimate

𝑡1

Forward

1 1

2

2

𝑡2

DP all-reduce DP all-reduce

2

𝑚𝑎𝑥(𝑡1 , 𝑡2 )

Backward

DP all-reduce

Communication

Figure 1. An illustrative example to demonstrate why the cost model adopted by existing works cannot accurately characterize the training time. Here, 𝑡 1 and 𝑡 2 denote the total computation and communication costs for each microbatch in stages 1 and 2, respectively, as defined in Eq. (1). (PP) [13, 28, 29], tensor parallelism (TP) [30, 33], and data parallelism (DP) [6, 20], dividing computation at layer, operator, and data levels, respectively. Meanwhile, modern clusters increasingly contain accelerator generations with different compute power, memory capacities, and interconnect bandwidths [51]. Such heterogeneity complicates LLM parallelization because synchronous operations, such as parameter updates, are bottlenecked by the slowest device [3, 44]. Recent work therefore seeks to balance computation across devices while reducing communication overhead [43, 49]. 2.2

Limitations of Existing Pre-training Systems

Several heterogeneity-aware LLM pre-training systems have been proposed to fully exploit cluster resources [19, 37, 43, 49, 55]. To improve LLM training efficiency, they coordinate model layer allocation across accelerator devices, balance computational workload within each stage, and minimize communication via network-efficient parallelization. However, these systems still exhibit several limitations that introduce substantial impact on training efficiency: 𝑳 1 : The adopted models cannot characterize training time accurately. Current heterogeneity-aware LLM training systems [19, 37, 43, 49, 55] are commonly built on top of an approximated training time estimation model: 𝑆 ∑︁ 𝑡𝑠 + (𝐵 − 1) · max 𝑡 𝑗 + 𝐷𝑃all, (1) 𝑠=1

1≤ 𝑗 ≤𝑆

where 𝑆 is the stage count, 𝐵 is the number of micro-batches, and 𝐷𝑃 all represents the DP all-reduce overhead. Each 𝑡𝑠 includes the forward, backward, and inter-stage communication costs per micro-batch at stage 𝑠. However, this formulation results in significant deviation from the ground-truth training time, which adversely impacts training efficiency. To illustrate this issue, we construct an example in Fig. 1, highlighting the discrepancy between the estimated training time in Eq. (1) and the actual training time. Specifically, in this example, involving two micro-batches pipelined across two stages, the estimation model fails to account for the

Training an LLM involves processing a batch of data through a sequence of layers during the forward pass (FP) and backward pass (BP) to compute gradients. As models scale, distributed training partitions parameters and data across devices to alleviate memory bottlenecks and accelerate computation [11]. Common strategies include pipeline parallelism 2

1 0

1

Real 1.23×

2

1.32×

3

1.0 0.0

Metis

Brute-Force 1.9×

1.2×1.0× 1.2×1.0× 1

2

Settings

1.0× 3

(a) Degradation in Metis.

Settings Figure 2. The cost model adopted by existing training systems consistently exhibits a significant training time gap.

HexiScale Brute-Force 1.5× 1.4× 1.0× 1.0×

1.2×1.0× 1.0 0.0

1

2

Settings

3

(b) Degradation in HexiScale.

Figure 3. The gap between training time under configurations found by existing systems’ search strategies and under their optimal configurations obtained via brute-force.

overlap between communication and computation, substantially overestimating training time. As shown in Fig. 2, such models consistently produce inaccurate training time estimates across various cluster configurations, with deviations of up to 1.32×. Consequently, these estimation errors can distort the ranking of candidate configurations, misleading the search process and ultimately reducing training efficiency. 𝑳 2 : The adopted pruning strategies yield suboptimal training efficiency. Exploring vast configuration spaces is computationally expensive, so existing approaches often balance training and search efficiency. However, such compromises risk overlooking the optimal configuration, ultimately degrading training performance. For instance, the pruning strategy in Metis [43] imposes strong assumptions on the low variance of per-stage device counts, and on the precedence of data parallelism over tensor parallelism during search. However, the optimal configuration varies with the training hyperparameters, such as batch size and the inter-accelerator connection in clusters, making these rigid strategies potentially inefficient. Fig. 3 shows the notable time gap for training one iteration over LLaMA-3 models [8, 41] in the cluster consisting of 16 Atlas A2-2 and 16 Atlas A2-3 NPUs (detailed NPU parameters are provided in § 6.1). The results compare the training time using configurations derived from existing pruning strategies with that of the ideal configuration found via exhaustive search. Specifically, Metis and HexiScale [49] exhibited up to 1.91× and 1.47× degradation, respectively, from improper pruning. 𝑳 3 : Existing modeling techniques lack generalizability. Since existing works employ coarse-grained analytical approximations, their models cannot accommodate diverse pipeline schedules. For example, Eq. (1) fails to accurately model training time under the Interleaved 1F1B paradigm [30], which minimizes pipeline bubbles through enhanced computation overlap and thus reduces training time. This same limitation applies to other 1F1B variants like Eager 1F1B [56] and Seq1F1B [36], as their similar bubble-reduction mechanisms likewise elude coarse-grained modeling. 2.3

2.0

Training Time (Normalized)

Estimated 1.29×

Training Time (Normalized)

Training Time (Normalized)

2

Fortunately, we found that modeling LLM training as a graph offers a valuable property that significantly enhances modeling performance [12, 14, 16, 21, 27, 39, 45–47]. Specifically, across any training paradigm, total LLM training time is determined by the maximum execution time across all micro-batches. Therefore, accurate training time estimation can be achieved by analyzing the computation dependencies across micro-batches and the computation time within each stage. Graph-based modeling techniques are particularly effective in capturing various types of dependencies inherent in LLM training. For instance, computation on the (𝑘 + 1)-th stage can only begin after the activations from the 𝑘-th stage are delivered, and operations on a given device must wait for preceding computations to complete. By capturing these dependencies, the model can be constructed in a more finegrained and precise manner, leading to improved accuracy in performance prediction and optimization. While graph-based modeling is highly effective in capturing complex dependencies, finding an optimal training configuration remains computationally expensive. The primary challenge lies in the vast space of feasible configurations, which renders exhaustive search impractical. To address this, a highly efficient search method is essential—one that can rapidly identify high-quality configurations without traversing numerous candidates, posing a significant challenge in heterogeneous environments.

3

Poseidon Overview

In light of the opportunities and challenges above, we introduce Poseidon, a scalable LLM training system for heterogeneous accelerator resources. Specifically, Poseidon leverages an efficient training time estimation model and search strategy to enhance training efficiency by identifying more effective parallelization configurations while maintaining minimal search overhead. It also generalizes across different pipeline schedules, including 1F1B, Eager 1F1B, and Interleaved 1F1B.

Graph-based Modeling: Opportunities and Challenges

3.1

Based on the preceding analysis of training time evaluation, effective modeling should provide accurate time estimates and capture the characteristics of different training paradigms.

Key Design Ideas

In this subsection, we highlight the key ideas underlying Poseidon to explain its performance advantages. 𝑰 1 : Explicit DAG construction for training time modeling. To achieve precise training time modeling, Poseidon 3

Number of Layers

Device-1 1F1B 11 10 9 12345678 Stage ID

12 11 10 9

Device-2 Eager-1F1B

12345678 Stage ID

10 9 8 7

Device-3 Interleaved-1F1B

Profiling

Enum DP, PP, TP

Memory Profiling

Devices of Stages Time Profiling

Searching

Layer-level Pruning

Layers of Stages

Simulator DAG-based Cost Model

123456789 Stage ID

Time Cost

Training Config

Figure 4. Optimal parallelization configurations consistently exhibit a ridge-like pattern, i.e., the number of layers allocated to the same device type first increases and then decreases.

Stage 1

Stage 2 Stage 3

Stage-level Pruning

Figure 5. The system overview of Poseidon. without compromising solution quality. For instance, once the number of layers allocated to a particular device type starts to decrease, configurations with a higher layer count than the previous stage can be safely pruned, substantially reducing search time.

utilizes a DAG-based technique to capture the complex dependencies and represent the execution times among operations. Specifically, this approach characterizes the execution at both the stage- and batch-level, allowing us to simultaneously incorporate the data execution order on each stage and the data dependencies within each micro-batch. In this DAG representation, each node corresponds to either an FP or BP operation for a specific micro-batch on a given stage, with the node weight indicating the computation overhead. Edges in different directions represent two types of execution dependencies: inter-stage activation transfers and the execution of different batches within the same stage. The edge weights denote the latency incurred by these dependencies. Consequently, the training time for each iteration can be represented as the longest path through the graph. This approach inherently captures the overlap between communication and computation. 𝑰 2 : Stage-level pruning via early stopping with partial estimation. To address the challenge of achieving high search efficiency, Poseidon introduces a stage-level pruning strategy to significantly reduce the search overhead. This approach employs an early exit mechanism that halts exploration of inefficient configurations when stages are partially allocated. The mechanism is driven by a lower-bound function we design, which estimates whether the training time of a partially determined configuration is likely to exceed the best result observed so far. If this condition is met, further exploration of that configuration path can be safely skipped, substantially saving search time without compromising solution quality. 𝑰 3 : Layer-level pruning via leveraging ridge-like distribution. To further reduce the search overhead in identifying optimal configurations, Poseidon incorporates an efficient layer-level pruning approach that maps model layers across stages. This leverages a key observation: in the optimal parallelization configuration, the number of layers assigned to all stages of the same device type consistently follows a ridge-like pattern, where the count first increases and then decreases, as shown in Fig. 4. Exploiting this property, Poseidon can skip the evaluation of certain configurations

3.2

System Architecture

The Poseidon system comprises three functional modules: profiling, searching, and simulator. The representation of these modules and their interaction are depicted in Fig. 5. Poseidon first incorporates a profiling module in the system to conduct computation of model layers and collect the corresponding execution time information. With the help of this module, we can obtain the computation time, communication bandwidth, and memory footprint—including activation and gradient storage—for a single layer under different TP degrees across heterogeneous accelerator devices. Additionally, the memory profiling provides critical constraints on the maximum number of layers that can be assigned to each stage. This profiling data underpins the DAG-based cost model, enabling accurate evaluation of candidate parallelization configurations. Given the vast search space introduced by accelerator heterogeneity, Poseidon’s search module explores optimal training parallelization configurations by applying two complementary pruning strategies. It first uses stage-level pruning to early-terminate branches whose projected training time is clearly suboptimal, and then exploits the insight that, within each accelerator type, optimal layer allocations across stages follow a ridge-like pattern to perform layer-level pruning. Together, these techniques swiftly eliminate low-efficiency candidates and dramatically accelerate the search process. Moreover, the search process relies on feedback from a simulator module, which uses a DAG to model the computation and communication dependencies of each parallelization configuration. By applying a dynamic programming technique over the DAG, the simulator module consistently captures training time information from the path with the longest time overhead for the provided parallelization configurations. 4

S𝑡𝑎𝑟𝑡: 𝑆

Pipeline-3 𝐹1,1

𝐹1,2

𝐹1,3

𝐹1,1𝐹2,1

𝐹1,2𝐹2,2

𝐹1,3𝐹2,3

𝐹1,4 𝐵1,1 Pipeline-2 𝐹1,4𝐵2,1 Pipeline-1

𝐵1,1𝐹2,4

𝐵1,2

𝐵1,3

𝐵1,4

𝐵1,2𝐵2,2

𝐵1,3𝐵2,3

𝐵1,4𝐵2,4

Stage-1

𝐹1,1𝐹2,1 𝐹 𝐹 𝐹 𝐵 𝐵 𝐵 𝐵 𝐹3,1 1,2𝐹2,2𝐹3,2 1,3𝐹2,3𝐵3,1 1,4𝐵2,1𝐹3,3 1,1𝐹2,4𝐵3,2 1,2𝐵2,2𝐹3,4 1,3𝐵2,3𝐵3,3 1,4𝐵2,4𝐵3,4

Stage-2

𝐵 𝐵 𝐹2,1𝐹3,1 𝐹 𝐹 𝐵 𝐵 𝐹4,1 2,2𝐹3,2𝐵4,1 2,3𝐵3,1𝐹4,2 2,1𝐹3,3𝐵4,2 𝐹2,4𝐵3,2𝐹4,3 2,2𝐹3,4𝐵4,3 2,3𝐵3,3𝐹4,4 2,4𝐵3,4𝐵4,4

Stage-3

𝐹3,1𝐹4,1

𝐹3,2𝐵4,1

𝐵3,1𝐹4,2

𝐹3,3𝐵4,2

𝐵3,2𝐹4,3

𝐹3,4𝐵4,3

𝐵3,3𝐹4,4

𝐵3,4𝐵4,4

Stage-4

𝐹4,1

𝐵4,1

𝐹4,2

𝐵4,2

𝐹4,3

𝐵4,3

𝐹4,4

𝐵4,4

𝐸𝑛𝑑: 𝐸

Figure 6. Graph-based pipeline cost model (3 pipelines, 4 stages per pipeline, 4 micro-batches per pipeline, 1F1B schedule) with computational graph annotations: 𝐹𝑠,𝑚 /𝐵𝑠,𝑚 denote forward/backward of micro-batch 𝑚 at stage 𝑠, 𝑅𝑠 indicates DP all-reduce; blue edges represent inter-stage dependencies and green edges represent intra-stage dependencies.

4

Micro-batch computation nodes. For each micro-batch at every stage, two nodes are generated: a forward pass node 𝐹𝑠,𝑚 and a backward pass node 𝐵𝑠,𝑚 , where 𝑠 ∈ [1, 𝑁 ] represents the stage index and 𝑚 is the micro-batch index. This results in 2𝑁 𝑀 computation nodes for each pipeline. DP all-reduce nodes. Each stage includes a synchronization node 𝑅𝑠 for gradient aggregation, resulting in a total of 𝑁 communication nodes. Control nodes. A virtual start node 𝑆 and end node 𝐸 define the global execution boundaries.

Design Details of Poseidon

This section presents Poseidon’s design, starting with a discussion of the DAG-based modeling technique in § 4.1. We then introduce the stage-level and layer-level pruning strategies in § 4.2.1 and § 4.2.2. Finally, we propose an efficient search algorithm in § 4.2.3, which leverages both strategies to accelerate the search process without compromising training efficiency. 4.1

DAG-based Cost Model Design

Before presenting our DAG-based training time model, we first contextualize it within prior work. DAG abstractions have been widely used to model DNN and LLM training pipelines for purposes such as compilation optimization [27, 46], execution scheduling [12, 14, 16, 39], communication optimization [45–47], and formal reasoning [21]. These works convincingly demonstrate the expressive power of DAGs. In Poseidon, we leverage this expressive power for a fundamentally different purpose. Our contribution is not the DAG abstraction itself, but its novel application as a quantitative, microbatch-level cost model for guiding optimal parallelism search in LLM training. While prior work uses DAGs to represent the system for the purpose of optimization or reasoning, we use the DAG to simulate and estimate training time. We formulate our training time modeling technique as follows: given a parallelization configuration 𝐶 with 𝐷 pipelines, 𝑁 stages per pipeline, and 𝑀 micro-batches, the model yields a precise training time estimation 𝑇 . The modeling process involves three steps: node construction, edge construction, and weight assignment. Once 𝐶 is represented by this DAG, we compute 𝑇 —defined as the critical path length—using dynamic programming. A representative example illustrating this DAG-based cost model is provided in Fig. 6.

4.1.2 Edge construction. We leverage directed edges to capture dependencies among nodes in the DAG. Specifically, we define the dependencies with the following types of edges: Inter-stage edges. For each micro-batch 𝑚, the forward computation flow is represented by 𝐹𝑠,𝑚 → 𝐹𝑠+1,𝑚 , while the backward propagation flow is represented by 𝐵𝑠,𝑚 → 𝐵𝑠 −1,𝑚 . This requires 2(𝑁 − 1)𝑀 edges for each pipeline. Intra-stage edges. The pipeline schedule, such as 1F1B, dictates micro-batch execution order within a stage. For 𝑀 micro-batches, each stage needs 2𝑀 − 1 edges to represent this order, totaling 𝑁 (2𝑀 − 1) edges per pipeline. All-reduce edges. Each 𝑅𝑠 depends on 𝐷 pipeline instances completing their final backward pass 𝐵𝑠,𝑀 . This requires 𝑁 𝐷 edges, where the edges are represented as 𝐵𝑠,𝑀 → 𝑅𝑠 . Auxiliary edges. The start node 𝑆 connects to all firststage forward nodes 𝐹 1,𝑚 via 𝐷 edges, and all 𝑅𝑠 nodes converge to the end node 𝐸 via 𝑁 edges, yielding 𝐷 + 𝑁 edges. 4.1.3 Weight assignment according to profiling data. Poseidon assigns weights to nodes and edges in the graph based on detailed profiling data. Most LLMs consist of repeatedly stacked homogeneous layers, such as LLaMA [41] and Qwen3-32B [50], so profiling a single representative layer is sufficient. Some models adopt mixed layer designs; for example, DeepSeek-V3 [7] contains both dense Transformer layers and MoE Transformer layers. In such cases, Poseidon profiles one representative layer for each distinct layer type,

4.1.1 Node construction. In the DAG-based cost model, nodes represent diverse operations. In particular, we model the training process with the following node types based on event type: 5

which introduces only marginal overhead due to the small number of distinct layer types. To support the cost model, Poseidon profiles three types of performance metrics under different accelerator types and TP degrees: (1) the forward and backward execution times of a single layer, (2) the activation transfer time, and (3) the DP all-reduce synchronization time of a single layer. These per-layer profiling results are then aggregated to estimate the stage-level execution and communication costs for weight assignment as follows: Computation nodes 𝑭𝒔,𝒎 , 𝑩 𝒔,𝒎 . The forward and backward pass times are assigned to these nodes. Each value equals the profiled per-layer computation time multiplied by the number of layers in stage 𝑠. Inter-stage edges. The weights of these edges represent activation transfer time, obtained from profiling under the corresponding device type and TP degree. Intra-stage edges. These edges have zero weight and represent only the batch execution order within a stage. DP all-reduce nodes 𝑹 𝒔 . Gradient synchronization time is assigned to these nodes, which equals the per-layer allreduce time multiplied by the number of layers in stage 𝑠. Control nodes and auxiliary edges. These are assigned zero weight, indicating no time cost.

To achieve this, Poseidon introduces a lower-bound function 𝑔(𝑆𝑥 ), used to check if the current search is worth continuing. Here, 𝑆𝑥 denotes the temporal parallelization configuration with only the number of model layers in stages 1 to 𝑥 determined. The lower-bound function (𝑥 −1is defined ) as: 𝑥 −1 ∑︁ ∑︁ 𝑔(𝑆𝑥 ) = 𝑓𝑖 + 𝑀 · (𝑓𝑥 + 𝑏𝑥 ) + max 𝑏 𝑗 + 𝑑𝑖 , (2) 𝑖=1

𝑗=𝑖

where 𝑓𝑖 , 𝑏𝑖 , and 𝑑𝑖 denote the forward, backward, and DP all-reduce time for stage 𝑖, and 𝑀 is the number of microbatches. This formulation provides an explicit lower bound of training time given the determined layer allocation under 𝑆𝑥 , where adjusting the allocation of remaining layers can only yield a longer training time. If the temporal configuration 𝑆𝑥 yields a training time longer than 𝑇min , i.e., the best training time observed so far, Poseidon stops further exploration based on this configuration. Otherwise, Poseidon continues exploring configurations extended from 𝑆𝑥 and updates 𝑇min only when a complete configuration with a shorter estimated training time is found. Due to the exponential growth of the search space with respect to the number of stages traversed, the early stopping strategy effectively mitigates the computational overhead of exhaustive parallel training search. We formalize this observation in the subsequent theorem, with a full proof by contradiction detailed in the Appendix A. Theorem 1. If 𝑔(𝑆𝑥 ) > 𝑇min , all extensions of 𝑆𝑥 can be safely pruned without loss of optimality.

4.1.4 Modeling various pipeline schedules. Pipeline schedules such as Eager 1F1B and Interleaved 1F1B have different execution dependency structures. The flexible DAGbased model accommodates these variations by adjusting node counts and dependency edges without changing the overall framework. We use Interleaved 1F1B to illustrate the corresponding modifications to node and edge construction. Node construction. In the Interleaved 1F1B schedule with a virtual pipeline size of 𝑉 , there are 𝑉 𝑁 virtual stages. Thus, the number of computation nodes increases to 2𝑉 𝑁 𝑀. Edge construction. As the number of nodes increases under the Interleaved 1F1B schedule, the number of interstage edges also scales accordingly, reaching 2(𝑁𝑉 − 1)𝑀, and the number of intra-stage edges increases to 𝑁 (2𝑉 𝑀 −1). 4.2

1≤𝑖 ≤𝑥

4.2.2 Layer-level pruning. In addition to stage-level pruning, Poseidon further introduces layer-level pruning to improve search efficiency. These two strategies operate at different levels and can be effectively integrated with each other. As shown in Fig. 4, the optimal parallelization configuration for LLM training consistently exhibits a ridge-like layer allocation across devices of the same type, i.e., the number of layers assigned to stages within a specific accelerator type first increases and then decreases. This ridge-like pattern reveals a key property guiding our pruning strategy. However, one may wonder if this property will always hold, especially as the search space grows exponentially with cluster size and number of accelerator types. Therefore, a theoretical guarantee is essential for search efficiency based on this property. Accordingly, we propose a rigorous theorem that works for any heterogeneous environment at an arbitrary scale as follows: Theorem 2. Given 𝑁 pipeline stages, there exists an optimal layer allocation scheme 𝑙 1, 𝑙 2, . . . , 𝑙 𝑁 , where 𝑙𝑖 denotes the number of layers assigned to stage 𝑖, such that for any set of device groups of the same type assigned to 𝑘 stages 𝑠 1 < 𝑠 2 < · · · < 𝑠𝑘 , the corresponding subsequence 𝑙𝑠1 , 𝑙𝑠2 , . . . , 𝑙𝑠𝑘 forms a ridge-like distribution. This theoretical result provides a principled foundation for search and enables the pruning of suboptimal configurations. An intuitive explanation of Theorem 2 is as follows:

Search Policy for Optimal Parallel Training

Due to the expansive search space from heterogeneous accelerator clusters, effective pruning techniques are essential to identify the most efficient training configuration. 4.2.1 Stage-level pruning. Instead of applying heuristics to trade off between training efficiency and search overhead, we first propose a stage-level pruning strategy to maintain an optimally efficient configuration with reduced search time. Specifically, it enables early exit in branches deemed to be inefficient, even if only partial stages of them are determined, to avoid unnecessary configuration exploration. This capability is critical as today’s LLMs consist of dozens of layers, which inevitably expands the search space across stages. 6

earlier stages need to cache more forward activations [28], resulting in higher memory pressure on devices of the same type. To alleviate this, fewer layers are allocated to early stages, leading to gradually more layers toward intermediate stages. Conversely, latter stages wait longer for their first micro-batch; assigning more layers here prolongs training time and creates a bottleneck. Thus, the number of layers should decrease toward the tail of the model. This dual pressure—balancing memory constraints in early stages and latency bottlenecks in later ones—creates a ridge-like layer allocation across stages of the same device type. We defer the full proof to the Appendix. Furthermore, this theorem reveals a counterintuitive finding: merely balancing computation time across stages increases training time. Consequently, existing approaches focused on minimizing stage time variance are inherently less efficient.

search process, the search and simulator modules are implemented in C++ over 3,000 lines of code. The profiling module is implemented in Python and also surpasses 3,000 lines. Simulator design. The simulator module provides precise training time estimates for a given parallelization configuration, with a user-specified batch size and sequence length. It evaluates the training time using profiling data and the DAG-based cost model. Specifically, if the PP degree remains unchanged, the DAG does not need to be rebuilt; only the edge and node weights are updated. Training on NPUs. To fully unleash the potential of NPUs, Poseidon performs profiling and executes LLM training on top of MindSpeed-LLM [2], a high-performance training framework built on Megatron-LM [30] and optimized for NPUs. Consequently, Poseidon provides full-stack NPU support for LLM training. Automatic training workflow. Poseidon employs a usertransparent integration mechanism to interface with LLM training frameworks. It injects search and profiling logic into user training requests transparently. The optimal configuration is applied automatically, enabling end-to-end training.

4.2.3 Search policy. Building upon the properties described above, Poseidon performs an efficient search over parallelization configurations. Since TP requires frequent communication, Poseidon restricts each TP group to a single node to avoid inter-node bandwidth bottlenecks [33, 34, 43], whereas DP and PP are not subject to this restriction. To mitigate straggler effects, where faster pipelines are blocked by peers during gradient aggregation, Poseidon enforces identical configurations across all pipelines. Before the main search, Poseidon randomly samples 500 feasible configurations and estimates their training times using the DAG-based cost model. The best sampled result initializes 𝑇min , providing an effective upper bound for subsequent pruning. The impact of the number of warm-up samples is evaluated in § 6.4.2. Poseidon then enumerates feasible 3D parallelism configurations and device-to-stage assignments while recursively determining the layer allocation of each stage. For every complete configuration, the simulator estimates its training time and updates 𝑇min when a better configuration is found. During this process, Poseidon applies stage- and layer-level pruning cooperatively.At stage 𝑥, if the lower bound 𝑔(𝑆𝑥 ) of a partial configuration exceeds 𝑇min , the entire search branch is discarded. Meanwhile, the layer-level strategy explores only allocations satisfying both memory constraints and the theoretically derived ridge-like pattern. The two strategies are complementary: stage-level pruning is particularly effective for configurations with many pipeline stages, whereas layer-level pruning substantially reduces the search space when many candidate layer allocations exist for each stage.

5

6

Evaluation

We evaluate Poseidon on heterogeneous clusters across five aspects: training throughput, search efficiency, cost model accuracy, pruning effectiveness, and system scalability. In this section, training throughput is defined as the reciprocal of per-iteration training time; therefore, higher throughput indicates shorter training time. 6.1

Experiment Setup

Cluster environments. Experiments are conducted on both NPU clusters and GPU clusters: (1) NPU clusters. Three NPU clusters are used, equipped with Atlas A2-2, A2-3, and A2-4 devices, respectively. Each cluster consists of homogeneous nodes with eight NPUs per node. The intra-node, intra-cluster inter-node, and cross-cluster bandwidths are 1400 Gbps, 200 Gbps, and 75 Gbps, respectively. The peak performance of the A2-2, A2-3, and A2-4 is 376 TFLOPs, 313 TFLOPs, and 280 TFLOPs, respectively. The A2-2 and A2-3 provide 64 GB of memory, while the A2-4 provides 32 GB. (2) GPU clusters. Three GPU clusters are used, equipped with A100, RTX-4090, and RTX-3090 GPUs, respectively. Each cluster consists of homogeneous nodes with eight GPUs per node. The intra-node, intra-cluster inter-node, and crosscluster bandwidths are 256 Gbps, 100 Gbps, and 75 Gbps, respectively. The peak performance of the A100, RTX-4090, and RTX-3090 is 312 TFLOPs, 330 TFLOPs, and 71 TFLOPs, with memory capacities of 80 GB, 24 GB, and 24 GB, respectively. Baselines. We compare Poseidon with state-of-the-art heterogeneity-aware pre-training systems, including HexiScale [49] and Metis [43].

Poseidon Implementation

The Poseidon framework consists of three core components: a search module, a simulator module, and a profiling module. To meet stringent performance requirements during the 7

0.0

Setting 1

0.0 Poseidon

Setting 2 Metis

0.0 HexiScale

Searching Time (s)

Normalized Thpt.

LLaMA-3(8B) LLaMA-2(13B) LLaMA(30B) 1.15× 1.05× 1.00× 1.94× 1.00× 0.89× 1.0 0.86× 1.6 1.0 1.00× 1.15× 0.5 0.5 0.8 Setting 3

16

LLaMA-3(8B) 15.54

8 6.73 0

3.80 Setting 1

LLaMA-2(13B) 55.34 1000 50 49.36 42.99 25

500

0

0

Poseidon

Setting 2 Metis

LLaMA(30B) 1003.05

88.42

HexiScale

89.87

Setting 3

Figure 8. The search time of all systems across various model scales and heterogeneous cluster settings.

Figure 7. The training throughput of Poseidon compared to baseline systems, across different models and cluster settings.

6.2.1 Throughput across model and cluster scales. For a fair comparison, all systems adopt the 1F1B schedule, the only one supported by the baselines. As shown in Fig. 7, Poseidon achieves the highest training throughput across all settings, outperforming Metis by 1.05–1.94× and HexiScale by 1.21–1.69×, with both maximum gains observed on LLaMA-30B. Metis’s layer allocation suffers from two issues: (1) heuristics that prioritize computational balance across stages, and (2) iterative error-correction for OOM failures. In 1F1B scheduling, its greedy allocation often triggers OOMs because earlier stages require higher activation memory. As a result, the subsequent heuristic reallocation lacks global optimality and degrades the quality of final allocations. Similarly, HexiScale relies on a two-phase graph-partitioning heuristic without optimality guarantees. In contrast, Poseidon exhaustively explores the search space under memory constraints, leveraging an accurate DAG-based cost model and provable pruning strategies to consistently find optimal configurations.

7 hours of search time given a per-search cost of 0.025 ms on our testbed. The estimated search time further increases to 4.81 days in Setting 2, with each search taking 0.0408 ms. When scaled to Setting 3, the estimated total search time exceeds 5,000 years—each search taking 0.189 ms. Although searching a single configuration is on the order of milliseconds, the enormous size of the configuration space renders brute-force search impractical. These results highlight the need for efficient search algorithms that reduce exploration time while preserving solution quality, as discussed in § 2.3. Fig. 8 shows the search time taken by each method for the experiments in Fig. 7. Although Poseidon does not achieve the shortest search time in all settings, it effectively balances search efficiency and configuration quality. For smaller models (LLaMA-8B/13B), differences in search time across methods are marginal, with search times remaining within the same order of magnitude. Considering the overall training cost, this overhead is negligible. For instance, assuming a training dataset of 10 billion tokens under Setting 1, the configuration selected by Poseidon completes training in approximately 21.2 hours, which is 2.8 hours faster than that produced by Metis. Such throughput gains become increasingly significant as the size of the datasets scales up. Moreover, since brute-force search is feasible under Settings 1 and 2, we employ brute-force search in these cases. The results show that Poseidon identifies exactly the same optimal parallelism configurations as brute-force search, achieving identical throughput. On larger models (LLaMA-30B), Poseidon reduces search time by up to 11.34× compared to Metis, whose depth-first search becomes less efficient as the search space grows. Conversely, Poseidon’s pruning mechanism aggressively eliminates suboptimal configurations early, reducing search overhead. While HexiScale also achieves rapid search via a twophase graph-partitioning heuristic, it compromises configuration quality, as shown in Fig. 7.

6.2.2 Analysis of search overhead. All methods require a one-time profiling step to collect data for the cost model. In our evaluation, this step takes up to 2 hours and is reused across all methods. Thus, the profiling time is excluded from the reported search overhead. The brute-force search over the full configuration space is prohibitively time-consuming. For example, under Setting 1, there are approximately 1.01 × 109 feasible configurations according to dynamic programming analysis, leading to at least

6.2.3 Case study of configuration results. Table 1 presents the parallelism strategies discovered by all methods under Setting 2. Poseidon identifies a ridge-like configuration through fine-grained device and layer allocation. Metis produces similar DP and TP settings but retains a coarse pipeline layout due to its limited search space, which hinders exploitation of inter-device heterogeneity. HexiScale limits configurability by grouping similar devices using a graph-partitioning heuristic, narrowing opportunities for better configurations.

Experimental settings. We evaluate Poseidon under a diverse set of models and heterogeneous configurations: –Setting 1: LLaMA-3(8B) [8] on one 8×Atlas A2-2 node and one 8×Atlas A2-3 node. –Setting 2: LLaMA-2(13B) [42] on two 8×Atlas A2-2 nodes and two 8×Atlas A2-3 nodes. –Setting 3: LLaMA(30B) [41] on two 8×Atlas A2-2, two 8×Atlas A2-3, and two 8×Atlas A2-4 nodes. –Setting 4: Mixtral(8×7B) [17], an MoE model, on one 8×A22 node, two 8×A2-3 nodes, and four 8×A2-4 nodes. –Setting 5: LLaMA-2(13B) on one 8×A100 node, two 8×RTX4090 nodes, and two 8×RTX-3090 nodes. –Setting 6: LLaMA-3(70B) on six 8×Atlas A2-2, six 8×Atlas A2-3, and twenty-four 8×Atlas A2-4 nodes. 6.2

End-to-end Training Performance

8

Table 1. Configurations discovered by Poseidon, Metis, and HexiScale Poseidon

Metis

HexiScale

PP

16 (2, 1, A3, 2)×1, (2, 1, A3, 3)×2, (2, 1, A3, 2)×1, (2, 1, A2, 3)×4, (2, 1, A3, 2)×2, (2, 1, A3, 1)×2, (2, 1, A2, 3)×4

14

10

(2, 1, A2, 3)×8, (2, 2, A3, 3)×2, (2, 1, A3, 3)×1, (2, 1, A3, 2)×1, (2, 1, A3, 3)×1, (2, 1, A3, 2)×1

(1, 4, A3, 5)×3, (1, 2, A3, 2)×2, (1, 4, A2, 5)×1, (1, 4, A2, 6)×2, (1, 2, A2, 2)×2

Estimated

1 1.02× 0

1

1.02×

2 Setting

Real

0.98×

3

2.5 2.0 1.5

1F1B Eager-1F1B Interleaved-1F1B

1.0 0.5 32 64 128

256

Global Batch Size

512

Figure 9. Accuracy of the Figure 10. Training throughDAG-based model in estimat- put of Poseidon under different pipeline schedules. ing training time. 6.3

Poseidon

3.50 3.25 3.000

10

20

30

40

Searching Time (s)

50

60

25 20 15 10 5 0

Sampling Counts

Figure 11. Breakdown anal- Figure 12. Impact of samysis of Poseidon’s pruning pling frequency on search efstrategies. fectiveness. – Effectiveness of individual pruning strategies. Both stagelevel and layer-level pruning accelerate convergence across all batch sizes, demonstrating that each strategy effectively filters out poor configurations. – Synergistic gains from combined pruning. When combined, Poseidon achieves the fastest reduction in training time, highlighting the synergy: layer-level pruning guides the search toward high-quality regions, while stage-level pruning tightens bounds and prunes suboptimal branches.

Evaluation of the DAG-Based Cost Model

6.3.1 Training time estimation in end-to-end experiments. We compare Poseidon’s DAG-based estimation with actual end-to-end training times. For actual iteration time, we run 30 iterations, discard the first five to avoid warm-up, and average the remaining 25. As shown in Fig. 9, Poseidon maintains low estimation errors (≤2%) across Settings 1–3, indicating DAG-based cost model accurately captures runtime behaviors, including computation–communication overlap and dependencies across heterogeneous devices. We further evaluate its fidelity on LLaMA-3 (70B) under 20 distinct heterogeneous deployments, using two nodes of 8×Atlas A2-2, four nodes of 8×Atlas A2-3, and five nodes of 8×Atlas A2-4. For each deployment, 300 configurations are sampled (100 per pipeline schedule: 1F1B, Interleaved 1F1B, Eager 1F1B). These configurations cover diverse DP/PP/TP degrees (DP=1–4, PP=1–16, TP=1–8) to reflect realistic intraand inter-node parallelism. Across all configurations, the simulator achieves an average prediction accuracy of 98%.

6.4.2 Impact of warm-up sampling count. Poseidon adopts a warm-up sampling phase before stage-level pruning, in order to establish a high-quality initial baseline. A better baseline enables stricter pruning, which reduces simulator overhead. We tested warm-up counts from 0 to 20,000 and found a trade-off: more samples improve pruning but increase upfront cost. Without sampling, weak initial configurations lead to 23.2× more simulator calls compared to the optimal warm-up sampling count, as shown in Fig. 12. Moderate sampling (e.g., 500 samples) quickly sets a competitive baseline, enabling efficient pruning. Thus, Poseidon defaults to 500 warm-up samples to balance cost and effectiveness.

6.3.2 Generalizability to pipeline schedules. We evaluate Poseidon under different pipeline schedules using Setting 2. As shown in Fig. 10, varying the pipeline schedule alone can change throughput by up to 1.3× with other factors held constant. This highlights the importance of pipeline schedule as an additional optimization dimension—Poseidon fully exploits this flexibility, whereas existing baselines cannot. 6.4

No pruning Only layer-level Only stage-level

3.75

0 10 20 50 10 0 20 0 50 0 1K 2K 5K 10 K 20 K

Notation: (𝑎, 𝑏, 𝑐, 𝑑 ) × 𝑛 denotes 𝑛 consecutive stages with identical configurations. 4.00

Training Time (s)

2

Normalized Thpt.

Training Time (Normalized)

(DP, TP, NPU, Layer)

Total Calls (Normalized)

Methods

6.5

Evaluation in General Settings

6.5.1 Training MoE models. The Mixture-of-Experts (MoE) architecture is increasingly popular for scaling LLMs, as adopted in Mixtral [17], DeepSeek-V3 [7], and Qwen3 [50]. Although MoE models introduce multiple experts per layer, the profiling patterns remain similar to non-MoE models. Thus, the proposed training paradigm can be seamlessly applied to MoE architectures without modification. We evaluated Poseidon on the Mixtral(8×7B) model under Setting 4. As shown in Fig. 13, Poseidon achieves 1.06× and 1.61× higher training throughput than Metis and HexiScale, respectively, while reducing search time by 52.07× and 3.15×.

Effectiveness of the Search Algorithm

6.4.1 Ablation studies for pruning strategies. Poseidon’s pruning strategies are evaluated on LLaMA-30B under four scenarios: stage-level only, layer-level only, full pruning, and no pruning. Fig. 11 shows the evolution of the best-found training time, yielding two insights: 9

1.00× 0.66×

0.5 0.0

300

Metis

Poseidon

656.1

600

Normalized Thpt.

1.06×

Searching Time (s)

Normalized Thpt.

1.0

39.7

0 12.6

HexiScale

1.6 0.8

1.74× 1.00× 0.63×

0.0 Poseidon

Searching Time (s)

Normalized Thpt.

Figure 13. Training throughput and search time of Poseidon compared to baseline systems on a MoE model.

Metis

1

Poseidon

1.51×

Metis

1.58×

HexiScale 1.50×

1.00×1.00×

1.00×1.09×

1.00×0.97×

96

192

288

Number of NPUs(A2-2:A2-3:A2-4 = 1:1:4)

Table 2. The search time of all systems on heterogeneous NPU clusters with varying scales

150 49.2 0

2

Figure 15. Training throughput of Poseidon and baselines on heterogeneous NPU clusters at different scales.

330.7

300

LLaMA-3(70B)

52.5

HexiScale

Figure 14. The training throughput and searching time of Poseidon compared to baselines on heterogeneous GPU clusters. 6.5.2 LLM training across GPU Clusters. Since the searching module and runtime system are loosely coupled, Poseidon’s search module can be easily integrated with other runtime systems designed for different types of accelerators. To verify Poseidon’s scalability in GPU environments, we integrated its search module into the HexiScale runtime and evaluated it under Setting 5. As shown in Fig. 14, Poseidon achieves 1.74× and 2.76× higher training throughput than Metis and HexiScale, respectively, while reducing search time by up to 6.72×. Notably, the GPU clusters exhibit greater compute and memory heterogeneity than the NPU clusters. The results show that Poseidon still delivers substantial throughput gains in such highly heterogeneous environments. More generally, parallelism configuration search is independent of accelerator type, relying only on generic hardware profiles such as compute capability, memory capacity, and communication bandwidth. At this level of abstraction, there is no fundamental difference between GPUs and NPUs. 6.5.3 Large-scale Simulation Experiments. To further evaluate Poseidon at larger scales, we simulate its and baselines’ performance under Setting 6 using LLaMA-3 (70B). All simulations use our simulator driven by the DAG-based cost model, emulating heterogeneous clusters with 96, 192, and 288 NPUs. Superior scalability with increasing cluster size. As shown in Fig. 15, Poseidon consistently delivers the highest training throughput across all simulated scales in terms of performance. Specifically, it achieves a 1.50× to 1.58× improvement over Metis and a 1.45× to 1.55× improvement over HexiScale. Efficient search at scale. As shown in Table 2, Poseidon substantially reduces search overhead, achieving 37.8×–53.0× speedups over Metis and 4.5×–5.3× over HexiScale. These results confirm that Poseidon’s theory-driven pruning strategies scale effectively, maintaining high efficiency despite exponential growth in device counts and search complexity. 10

7

Number of NPUs

96

192

288

Poseidon Metis HexiScale

26.0s 1378.7s 116.8s

47.8s 2382.6s 251.7s

75.8s 2862.9s 394.7s

Related Work

Automating parallelism on heterogeneous devices. Several recent studies have focused on automatically tuning parallelism for heterogeneous environments [43, 49]. Earlier work includes Piper [40], which first attempted to solve the problem of optimal partitioning via a two-level dynamic programming approach, and AMP [19], which defines a valid space of model-parallel strategies and underpins the cost models used by both Metis and HexiScale. Alpa [55] introduces a hierarchical search space and compiler-driven passes to derive efficient execution plans across parallelism levels—Metis itself builds on Alpa’s infrastructure. Galvatron [25] employs a decision-tree heuristic to prune the search space, followed by dynamic programming to identify optimal parallelism configurations. While recent work [48] adopts a similar stage-level pruning strategy, it overlooks pipeline parallelism. Other contemporary approaches, such as Sailor [34] and H2 [38], leverage dynamic programming combined with heuristic pruning or depth-first search with heuristics, respectively, to prune the search space. However, the cost models that support their pruning mechanisms rely on static formulas similar to those used in Metis. In contrast, Poseidon is the first system to provide theoretical guarantees of optimal parallel configurations on heterogeneous clusters for LLM training. DAG-based modeling for training systems. DAGs are a fundamental abstraction for representing dependency structures in deep neural network (DNN) and LLM training. Prior work has leveraged DAGs for compilation optimization [27, 46], execution scheduling [12, 14, 16, 39], communication optimization [45–47], and formal reasoning [21]. While these works demonstrate the expressive power of DAG abstractions, their objectives differ fundamentally from Poseidon. Notably, FlexFlow [16] is conceptually closest in its use of DAG modeling to Poseidon, it predicts parallel strategy performance via DAG simulation, but it focuses on operatorlevel parallel simulation for DNNs. Poseidon, by contrast,

performs microbatch-level simulation tailored to LLM parallelism and explicitly supports large-scale parallel configuration search. Pipeline schedule. Early pipeline parallelism schemes suffered significant device idle time ("pipeline bubbles"), limiting hardware utilization. While GPipe [13] introduced micro-batch pipelining to mitigate bubbles, it increases peak memory demands. Subsequent innovations like PipeDream’s 1F1B schedule [28] reduced bubbles through interleaved forward/backward passes, inspiring variants including Eager 1F1B [56], Interleaved 1F1B [30], and Seq1F1B [36]. Zerobubble [32] achieved full stall elimination via backward-pass decomposition—later refined in [31] to reduce activation memory without reintroducing bubbles. Most recently, Zorse [10] introduces interleaved GPipe-style pipeline scheduling with offloading parameters and activations to CPU memory. Mario [24] enables near-zero-cost activation checkpointing by overlapping recomputation with residual bubbles, enhancing both memory and compute efficiency. Poseidon’s flexible modeling readily integrates all such variants. Training frameworks and systems. Numerous systems have been developed to support efficient large-scale model training [4, 35, 53]. For example, Megatron-LM [33] pioneered model parallelism for multi-billion-parameter models; Whale [15] optimizes communication and memory efficiency on heterogeneous GPU clusters; and MegaScale [18] enables LLM training across over 10,000 GPUs by addressing scheduling, failure recovery, and network bottlenecks. These frameworks provide robust infrastructure for parallel training. To fully unleash the power of NPUs in our cluster, we built our system on top of MindSpeed-LLM, which offers similar functionalities with the above systems.

8

[2] Ascend. 2025. MindSpeed-LLM. https://github.com/Ascend/ MindSpeed-LLM. Accessed: 2025-04-18. [3] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems 33 (2020), 1877–1901. https://dl.acm.org/doi/10.5555/3495724.3495883 [4] Zihan Chang, Sheng Xiao, Shuibing He, Siling Yang, Zhe Pan, and Dong Li. 2024. Frenzy: A Memory-Aware Serverless LLM Training System for Heterogeneous GPU Clusters. arXiv:2412.14479 [cs.DC] https://arxiv.org/abs/2412.14479 [5] Yiran Chen, Yuan Xie, Linghao Song, Fan Chen, and Tianqi Tang. 2020. A Survey of Accelerator Architectures for Deep Neural Networks. Engineering 6, 3 (2020), 264–274. https://doi.org/10.1016/j.eng.2020. 01.007 [6] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc’aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. 2012. Large scale distributed deep networks. Advances in neural information processing systems (2012). https://dl.acm.org/doi/ 10.5555/2999134.2999271 [7] DeepSeek-AI, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, and et al. Chong Ruan. 2025. DeepSeek-V3 Technical Report. arXiv:2412.19437 [cs.CL] https://arxiv.org/abs/2412.19437 [8] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. 2024. The llama 3 herd of models. arXiv (2024). https://arxiv.org/abs/2407.21783 [9] Daya Guo, Qihao Zhu, Dejian Yang, Zhenda Xie, Kai Dong, Wentao Zhang, Guanting Chen, Xiao Bi, Yu Wu, YK Li, et al. 2024. DeepSeekCoder: When the Large Language Model Meets Programming–The Rise of Code Intelligence. arXiv (2024). https://arxiv.org/abs/2401.14196 [10] Runsheng Benson Guo, Utkarsh Anand, Khuzaima Daudjee, and Rathijit Sen. 2025. Zorse: Optimizing LLM Training Efficiency on Heterogeneous GPU Clusters. https://arxiv.org/abs/2507.10392 [11] Julia Gusak, Daria Cherniuk, Alena Shilova, Alexandr Katrutsa, Daniel Bershatsky, Xunyi Zhao, Lionel Eyraud-Dubois, Oleh Shliazhko, Denis Dimitrov, Ivan V Oseledets, et al. 2022. Survey on Efficient Training of Large Neural Networks.. In IJCAI. 5494–5501. https://doi.org/10. 24963/ijcai.2022/769 [12] Samuel Hsia, Alicia Golden, Bilge Acun, Newsha Ardalani, Zachary DeVito, Gu-Yeon Wei, David Brooks, and Carole-Jean Wu. 2024. MADmax beyond single-node: Enabling large machine learning model acceleration on distributed systems. In 2024 ACM/IEEE 51st Annual International Symposium on Computer Architecture (ISCA). IEEE, 818–833. https://doi.org/10.1109/ISCA59077.2024.00064 [13] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. 2019. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural informatSion processing systems 32 (2019). https://dl.acm.org/doi/10.5555/3454287.3454297 [14] Ho Young Jhoo, Chung-Kil Hur, and Nuno P Lopes. 2025. Pfeife: Automatic pipeline parallelism for pytorch. In Forty-second International Conference on Machine Learning. https://dl.acm.org/doi/10.5555/ 3780338.3781408 [15] Xianyan Jia, Le Jiang, Ang Wang, Wencong Xiao, Ziji Shi, Jie Zhang, Xinyuan Li, Langshi Chen, Yong Li, Zhen Zheng, et al. 2022. Whale: Efficient giant model training over heterogeneous GPUs. In 2022 USENIX Annual Technical Conference (USENIX ATC 22). 673–688. https://www.usenix.org/conference/atc22/presentation/jia-xianyan [16] Zhihao Jia, Matei Zaharia, and Alex Aiken. 2019. Beyond data and model parallelism for deep neural networks. Proceedings of Machine Learning and Systems 1 (2019), 1–13. https://proceedings.mlsys.org/paper_files/paper/2019/file/

Conclusion

We present Poseidon, an automated parallelization system for LLM training on heterogeneous accelerator clusters. At its core, Poseidon designs a delicate DAG-based cost model to capture critical factors such as dependencies and overlap, enabling accurate training time estimation via criticalpath analysis. Built upon this model, Poseidon incorporates two theoretically grounded searching strategies, i.e., stageand layer-level pruning, to reduce the search space without compromising training efficiency. Moreover, Poseidon is empowered with strong generalizability to enhance efficiency for diverse paradigms like Eager 1F1B and Interleaved 1F1B. These innovations collectively improve Poseidon’s efficiency, unlocking the scalable, efficient, and cost-effective LLM training in real-world heterogeneous environments.

References [1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. 2023. Gpt-4 technical report. arXiv (2023). https://arxiv.org/abs/2303.08774 11

large-scale language model training on gpu clusters using megatronlm. In Proceedings of the international conference for high performance computing, networking, storage and analysis. 1–15. https: //doi.org/10.1145/3458817.3476209 [31] Penghui Qi, Xinyi Wan, Nyamdavaa Amar, and Min Lin. 2024. Pipeline parallelism with controllable memory. arXiv preprint arXiv:2405.15362 (2024). https://arxiv.org/abs/2405.15362 [32] Penghui Qi, Xinyi Wan, Guangxing Huang, and Min Lin. 2024. Zero bubble (almost) pipeline parallelism. In The Twelfth International Conference on Learning Representations. https://iclr.cc/virtual/2024/poster/ 17595 [33] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-lm: Training multibillion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053 (2019). https://arxiv.org/abs/1909.08053 [34] Foteini Strati, Zhendong Zhang, George Manos, Ixeia Sánchez Périz, Qinghao Hu, Tiancheng Chen, Berk Buzcu, Song Han, Pamela Delgado, and Ana Klimovic. 2025. Sailor: Automating Distributed Training over Dynamic, Heterogeneous, and Geo-distributed Clusters. In Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles (Lotte Hotel World, Seoul, Republic of Korea) (SOSP ’25). Association for Computing Machinery, New York, NY, USA, 204–220. doi:10.1145/ 3731569.3764839 [35] Suhas Jayaram Subramanya, Daiyaan Arfeen, Shouxu Lin, Aurick Qiao, Zhihao Jia, and Gregory R Ganger. 2023. Sia: Heterogeneityaware, goodput-optimized ML-cluster scheduling. In Proceedings of the 29th Symposium on Operating Systems Principles. ACM, 1–17. https: //doi.org/10.1145/3600006.3613175 [36] Ao Sun, Weilin Zhao, Xu Han, Cheng Yang, Xinrong Zhang, Zhiyuan Liu, Chuan Shi, and Maosong Sun. 2024. Seq1f1b: Efficient sequencelevel pipeline parallelism for large language model training. arXiv preprint arXiv:2406.03488 (2024). https://arxiv.org/abs/2406.03488 [37] Zhenbo Sun, Huanqi Cao, Yuanwei Wang, Guanyu Feng, Shengqi Chen, Haojie Wang, and Wenguang Chen. 2024. Adapipe: Optimizing pipeline parallelism with adaptive recomputation and partitioning. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3. 86–100. https://doi.org/10.1145/3620666.3651359 [38] Ding Tang, Jiecheng Zhou, Jiakai Hu, Shengwei Li, Huihuang Zheng, Zhilin Pei, Hui Wang, and Xingcheng Zhang. 2025. H2:Towards Efficient Large-Scale LLM Training on Hyper-Heterogeneous Cluster over 1,000 Chips. arXiv:2505.17548 [cs.DC] https://arxiv.org/abs/2505.17548 [39] Zhenheng Tang, Xueze Kang, Yiming Yin, Xinglin Pan, Yuxin Wang, Xin He, Qiang Wang, Rongfei Zeng, Kaiyong Zhao, Shaohuai Shi, et al. 2024. FusionLLM: a decentralized LLM training system on geo-distributed GPUs with adaptive compression. arXiv preprint arXiv:2410.12707 (2024). https://arxiv.org/abs/2410.12707 [40] Jakub M Tarnawski, Deepak Narayanan, and Amar Phanishayee. 2021. Piper: Multidimensional planner for dnn parallelization. Advances in Neural Information Processing Systems 34 (2021), 24829–24840. https: //dl.acm.org/doi/10.5555/3540261.3542163 [41] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, MarieAnne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. 2023. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 (2023). https://arxiv.org/abs/2302.13971 [42] Hugo Touvron, Louis Martin, et al. 2023. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 (2023). https://arxiv.org/abs/2307.09288 [43] Taegeon Um, Byungsoo Oh, Minyoung Kang, Woo-Yeon Lee, Goeun Kim, Dongseob Kim, Youngtaek Kim, Mohd Muzzammil, and Myeongjae Jeon. 2024. Metis: Fast Automatic Distributed Training on Heterogeneous {GPUs}. In 2024 USENIX Annual Technical Conference (USENIX ATC 24). 563–578. https://dl.acm.org/doi/10.5555/3691992.3692027

b422680f3db0986ddd7f8f126baaf0fa-Paper.pdf [17] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. 2024. Mixtral of experts. arXiv (2024). https://arxiv.org/abs/2401.04088 [18] Ziheng Jiang, Haibin Lin, Yinmin Zhong, Qi Huang, Yangrui Chen, Zhi Zhang, Yanghua Peng, Xiang Li, Cong Xie, Shibiao Nong, et al. 2024. {MegaScale}: Scaling large language model training to more than 10,000 {GPUs}. In 21st USENIX Symposium on Networked Systems Design and Implementation (NSDI 24). 745–760. https://www.usenix. org/conference/nsdi24/presentation/jiang-ziheng [19] Dacheng Li, Hongyi Wang, Eric Xing, and Hao Zhang. 2022. Amp: Automatically finding model parallel strategies with heterogeneity awareness. Advances in Neural Information Processing Systems 35 (2022), 6630–6639. https://dl.acm.org/doi/10.5555/3600270.3600750 [20] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. 2020. Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprint arXiv:2006.15704 (2020). https://arxiv. org/abs/2006.15704 [21] Antian Liang, Zhigang Zhao, Kai Zhang, Xuri Shi, Chuantao Li, Chunxiao Wang, Zhenying He, Yinan Jing, and X Sean Wang. 2025. HAPT: Heterogeneity-Aware Automated Parallel Training on Heterogeneous Clusters. arXiv (2025). https://arxiv.org/abs/2509.24859 [22] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. 2024. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437 (2024). https://arxiv.org/abs/2412.19437 [23] Guodong Liu, Youshan Miao, Zhiqi Lin, Xiaoxiang Shi, Saeed Maleki, Fan Yang, Yungang Bao, and Sa Wang. 2024. Aceso: Efficient parallel dnn training through iterative bottleneck alleviation. In Proceedings of the Nineteenth European Conference on Computer Systems. 163–181. https://doi.org/10.1145/3627703.3629554 [24] Weijian Liu, Mingzhen Li, Guangming Tan, and Weile Jia. 2025. Mario: Near Zero-cost Activation Checkpointing in Pipeline Parallelism. In Proceedings of the 30th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming. 197–211. https://doi.org/10.1145/ 3710848.3710878 [25] Xupeng Miao, Yujie Wang, Youhe Jiang, Chunan Shi, Xiaonan Nie, Hailin Zhang, and Bin Cui. 2022. Galvatron: Efficient transformer training over multiple gpus using automatic parallelism. Proceedingsof the VLDB Endowment (2022). https://doi.org/10.14778/3570690.3570697 [26] Sparsh Mittal and Shraiysh Vaishay. 2019. A survey of techniques for optimizing deep learning on GPUs. J. Syst. Archit. 99 (2019). https: //doi.org/10.1016/j.sysarc.2019.101635 [27] Philipp Moritz, Robert Nishihara, Stephanie Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael I Jordan, et al. 2018. Ray: A distributed framework for emerging {AI} applications. In 13th USENIX symposium on operating systems design and implementation (OSDI 18). 561–577. https: //www.usenix.org/conference/osdi18/presentation/moritz [28] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. 2019. PipeDream: Generalized pipeline parallelism for DNN training. In Proceedings of the 27th ACM symposium on operating systems principles. 1–15. https://doi.org/10.1145/3341301.3359646 [29] Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. 2021. Memory-efficient pipeline-parallel dnn training. In International Conference on Machine Learning. PMLR, 7937–7947. https://proceedings.mlr.press/v139/narayanan21a.html [30] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. 2021. Efficient 12

[44] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processing systems 30 (2017). https://dl.acm.org/doi/10.5555/3295222.3295349 [45] Xinchen Wan, Xinyu Yang, Kaiqiang Xu, Xudong Liao, Yilun Jin, Yijun Sun, Zhenghang Ren, Han Tian, and Kai Chen. 2025. Coflow Scheduling for LLM Training. In Proceedings of the ACM SIGCOMM 2025 Conference. 1232–1234. https://doi.org/10.1145/3718958.3750467 [46] Shibo Wang, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, et al. 2022. Overlap communication with dependent computation via decomposition in large deep learning models. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1. 93–106. https://doi.org/10.1145/3567955.3567959 [47] William Won, Taekyung Heo, Saeed Rashidi, Srinivas Sridharan, Sudarshan Srinivasan, and Tushar Krishna. 2023. Astra-sim2. 0: Modeling hierarchical networks and disaggregated systems for largemodel training at scale. In 2023 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS). IEEE, 283–294. https://doi.org/10.1109/ISPASS57527.2023.00035 [48] Ruilong WU, Xinjiao Li, Yisu Wang, Xinyu Chen, and Dirk Kutscher. 2025. Rethinking Dynamic Networks and Heterogeneous Computing with Automatic Parallelization. In Proceedings of the 9th Asia-Pacific Workshop on Networking (APNET ’25). Association for Computing Machinery, New York, NY, USA, 164–171. doi:10.1145/3735358.3735382 [49] Ran Yan, Youhe Jiang, Xiaonan Nie, Fangcheng Fu, Bin Cui, and Binhang Yuan. 2026. HexiScale: Accommodating Large Language Model Training over Heterogeneous Environment. Proceedings of Machine Learning and Systems (2026). https://arxiv.org/abs/2409.01143 [50] An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, and others. 2025. Qwen3 Technical Report. arXiv:2505.09388 [cs.CL] https://arxiv.org/abs/2505.09388 [51] Zhisheng Ye, Wei Gao, Qinghao Hu, Peng Sun, Xiaolin Wang, Yingwei Luo, Tianwei Zhang, and Yonggang Wen. 2024. Deep learning workload scheduling in gpu datacenters: A survey. Comput. Surveys 56, 6 (2024), 1–38. https://doi.org/10.1145/3638757 [52] Ann Yuan, Andy Coenen, Emily Reif, and Daphne Ippolito. 2022. Wordcraft: story writing with large language models. In Proceedings of the 27th International Conference on Intelligent User Interfaces. 841–852. https://doi.org/10.1145/3490099.3511105 [53] Zihao Zeng, Chubo Liu, Xin He, Juan Hu, Yong Jiang, Fei Huang, Kenli Li, and Wei Yang Bryan Lim. 2025. AutoHete: An Automatic and Efficient Heterogeneous Training System for LLMs. arXiv:2503.01890 [cs.LG] https://arxiv.org/abs/2503.01890 [54] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Tianle Li, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zhuohan Li, Zi Lin, Eric P Xing, et al. 2023. Lmsys-chat-1m: A large-scale real-world llm conversation dataset. arXiv preprint arXiv:2309.11998 (2023). https://arxiv.org/abs/ 2309.11998 [55] Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P Xing, et al. 2022. Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22). 559– 578. https://www.usenix.org/conference/osdi22/presentation/zhenglianmin [56] Yonghao Zhuang, Lianmin Zheng, Zhuohan Li, Eric Xing, Qirong Ho, Joseph Gonzalez, Ion Stoica, Hao Zhang, and Hexu Zhao. 2023. On optimizing the communication of model parallelism. Proceedings of Machine Learning and Systems 5 (2023), 526– 540. https://proceedings.mlsys.org/paper_files/paper/2023/file/ a42cbafcabb6dc7ce77bfe2e80f5c772-Paper-mlsys2023.pdf 13

Appendix

B.1

In this appendix, we provide a detailed theoretical analysis of the stage-level pruning guarantee (Theorem 1) and the layer-level pruning guarantee (Theorem 2).

We consider a scenario of training LLMs on a pipeline with heterogeneous device groups. We introduce the following definitions and assumptions to formalize this process. Definition 1. Let 𝐿 denote the total number of layers in the model, 𝑁 the number of pipeline stages, and 𝑀 the number of micro-batches. Definition 2. Let 𝑙𝑖 denote the number of layers assigned to stage 𝑖. Let 𝑓𝑖 and 𝑏𝑖 denote the per-layer forward and backward computation time at stage 𝑖, respectively. Then, let 𝐹𝑖 and 𝐵𝑖 denote the total forward and backward computation time of a micro-batch at stage 𝑖, respectively. Note that 𝑓𝑖 and 𝑏𝑖 are determined by the device group assigned to stage 𝑖. Additionally, we have:

A

Stage-level Pruning Guarantee

Recall that stage-level pruning operates on a partial configuration 𝑆𝑥 during the recursive search. The lower-bound function 𝑔(𝑆𝑥 ), defined in Eq. (2), characterizes the minimum possible training time that any complete configuration extending 𝑆𝑥 can achieve. Therefore, if this lower bound already exceeds the best observed time 𝑇min , continuing to explore the remaining stages cannot lead to a better configuration. The following proof formalizes this intuition and shows that the pruning rule in Theorem 1 preserves optimality.

Preliminaries

𝐹𝑖 = 𝑙𝑖 · 𝑓𝑖 , 𝐵𝑖 = 𝑙𝑖 · 𝑏 𝑖 .

The proof of Theorem 1. We proceed by contradiction. Suppose there exists a full configuration 𝑆 ∗ extending 𝑆𝑥 such that its corresponding per-iteration training time 𝑇 (𝑆 ∗ ) < 𝑇min . Since 𝑆 ∗ extends 𝑆𝑥 , the layer allocations for stages 1 through 𝑥 remain unchanged, and only the assignments for stages 𝑥 + 1 onward may vary. By the definition of 𝑔(𝑆𝑥 ), which provides a valid lower bound on the total training time for any such extension, we must have:

Definition 3. For any two stages 𝑖 and 𝑗, let 𝑡𝑖,𝑗 denote the communication time required to transfer activations from stage 𝑖 to stage 𝑗. Definition 4. Let 𝐹𝑖,𝑗 and 𝐵𝑖,𝑗 denote the forward and backward nodes, respectively, corresponding to the 𝑖-th pipeline stage and the 𝑗-th micro-batch in the computation DAG. Without loss of generality, we consider a single pipeline in the DAG for the purpose of this proof. Therefore, all DAGs referred to in the following analysis implicitly assume a single pipeline structure. Definition 5. In a DAG, each node has at most two outgoing edges. One edge represents the dependency within the same micro-batch across adjacent pipeline stages, referred to as the inter-stage edge. The other edge represents the execution dependency within the same stage, referred to as the intra-stage edge. As shown in Figure 16, intra-stage edges are directed from left to right, while inter-stage edges of forward nodes are directed downward and those of backward nodes are directed upward. Definition 6. In a DAG, the critical path is the longest weighted path, defined as the path with the maximum sum of weights over all its nodes and edges. In addition, we introduce four assumptions to support the subsequent theoretical analyses. Assumption 1. For any 𝑖 (1 ≤ 𝑖 ≤ 𝑁 ), 𝑏𝑖 = 𝑟 · 𝑓𝑖 , where 𝑟 > 1 is the ratio of backward to forward computation time of a single layer. Backward propagation is typically more time-consuming than forward propagation due to gradient computation [32]. Assumption 2. For any communication time 𝑡𝑖,𝑗 (1 ≤ 𝑖, 𝑗 ≤ 𝑁 ) and any stage 𝑘 (1 ≤ 𝑘 ≤ 𝑁 ), it holds that

𝑇 (𝑆 ∗ ) ≥ 𝑔(𝑆𝑥 ). However, we are assuming both 𝑇 (𝑆 ∗ ) < 𝑇min and 𝑔(𝑆𝑥 ) > 𝑇min , which leads to a contradiction: 𝑇 (𝑆 ∗ ) ≥ 𝑔(𝑆𝑥 ) > 𝑇min > 𝑇 (𝑆 ∗ ). Therefore, no extension of 𝑆𝑥 can yield a per-iteration training time smaller than 𝑇min , and all such configurations can be safely pruned. □

B

(3)

Layer-level Pruning Guarantee

The layer-level pruning guarantee is built upon the ridgelike property stated in Theorem 2. Its proof is more involved but can be summarized in three steps. First, we identify a canonical structure of the critical path. Second, based on this characterization, we show that for stages assigned to device groups of the same type, any layer allocation can be transformed into a non-increasing ordering without increasing the critical-path cost, when temporarily ignoring the impact of device memory constraints. Finally, we combine this property with the fact that, under the 1F1B pipeline paradigm, activation storage requirements decrease monotonically with increasing stage indices [15, 23, 28]. Together, these arguments establish the optimality of the ridge-like layer distribution used in Neptune’s layer-level pruning.

𝑟 −1 𝐹𝑘 . 2 Assumption 3. It holds that 𝑀 ≥ 2𝑁 . 𝑡𝑖,𝑗 ≤

14

Figure 16. Warmup (blue), steady (orange), and ending (green) phases in the DAG when 𝑁 = 4 and 𝑀 = 8.

Figure 17. A schematic representation of the path 𝐾 (𝑥, 𝑦, 𝑤) In practice, the number of micro-batches is typically much larger than the number of pipeline stages [23], making this assumption reasonable. Assumption 4. Let 𝐹 max = max1≤𝑖 ≤𝑁 𝐹𝑖 and 𝐹 min = min1≤𝑖 ≤𝑁 𝐹𝑖 . We assume that 𝐹 min ≥ 𝛼 · 𝐹 max,
2 𝑟 where 𝛼 = max 3 , 1+𝑟 . In this proof, we focus on the Eager-1F1B pipeline scheduling strategy; other scheduling variants can be analyzed similarly. Under Eager-1F1B, for the 𝑖-th stage (1 ≤ 𝑖 ≤ 𝑁 ), the execution consists of three phases: (i) a warmup phase with 2(𝑁 − 𝑖) + 1 forward micro-batches; (ii) a steady phase consisting of 𝑀 −2(𝑁 −𝑖) −1 groups of alternating backwardand-forward micro-batches; (iii) a cooldown (ending) phase with 2(𝑁 − 𝑖) + 1 backward micro-batches. We refer to these three parts as the warmup phase, steady phase, and ending phase, respectively. Figure 16 presents an example. B.2

Critical Path Characterization

In this part, we identify a canonical structure of the critical path in the DAG and present its formal definition as follows. Definition 7. In a DAG, let 𝐾 (𝑥, 𝑦, 𝑤) denote a path constructed as follows: • start from 𝐹 1,1 ; • move downward until reaching stage 𝑤; • move right for 2(𝑁 − 𝑥) steps; • move downward until reaching stage 𝑥; 15

• move right by one step into the steady phase; • move upward until reaching stage 𝑦; • move right until entering the ending phase; • move upward until reaching stage 𝑤; • move right until reaching 𝐵 𝑤,𝑀 ; • finally, move upward until reaching 𝐵 1,𝑀 , where 1 ≤ 𝑤 ≤ 𝑦 ≤ 𝑥 ≤ 𝑁 , with 𝑦 = arg max𝑖 ∈ [𝑦,𝑥 ] 𝐹𝑖 and 𝑤 = arg max𝑖 ∈ [1,𝑥 ] 𝐹𝑖 . Figure 17 provides an example of 𝐾 (𝑥, 𝑦, 𝑤). Next, we prove the existence of the critical path 𝐾 (𝑥, 𝑦, 𝑤). Before that, we introduce several definitions and lemmas. Definition 8. The following terms describe canonical path structures in a DAG, which will be used in the subsequent lemmas. • Segment is a path consisting of consecutive steps in a single direction, which can be upward, downward, or rightward. • Plateau is a segment consisting of at least 2 consecutive rightward steps in the steady phase. • Peak consists of three consecutive segments in the steady phase: first, ℎ consecutive upward steps; second, 𝑗 consecutive rightward steps; and third, ℎ consecutive downward steps, where ℎ ≥ 0, 𝑗 ≥ 1, and ℎ is referred to as the height of the peak. The peak-top is the segment formed by the 𝑗 rightward steps. • Valley consists of three consecutive segments in the steady phase: first, 𝑑 consecutive downward steps; second, 𝑗 consecutive rightward steps; and third, 𝑑 consecutive upward

Figure 18. The step transfer operation during steady phase, transfering 2 rightward steps from one stage to another stage.

Figure 19. An example of Peak Flattening. Replacing a peak(blue) with a rightward segment(red). steps, where 𝑑 ≥ 0, 𝑗 ≥ 1, and 𝑑 is referred to as the depth of the valley. The valley-bottom is the segment formed by the 𝑗 rightward steps. Lemma 1. If there exists a critical path containing plateaus in the steady phase, then there exists a critical path such that all its plateaus in the steady phase are contained in a single stage, denoted as the 𝑥 ∗ -th stage. Moreover, let 𝑆 denote the set of stages visited by the path in the steady phase. Then it holds that

local cost change satisfies 𝐹𝑥 + 𝐵𝑥 = (1 + 𝑟 )𝐹𝑥 ≤ (1 + 𝑟 )𝐹 𝑦 = 𝐹 𝑦 + 𝐵 𝑦 ,

(4)

where the inequality follows from 𝐹𝑥 ≤ 𝐹 𝑦 . Therefore, the modified path is no worse than the original critical path. In the case 𝐹𝑥 < 𝐹 𝑦 , the operation strictly improves the path cost, contradicting the optimality of the original critical path. Hence, we must have 𝐹𝑥 = 𝐹 𝑦 . By repeatedly applying this relocation argument, all plateaus can be merged into a single stage 𝑥 ∗ such that 𝐹𝑥 ∗ = max𝑖 ∈𝑆 𝐹𝑖 , where 𝑆 is the set of stages visited in the steady phase. This completes the proof. □

𝐹𝑥 ∗ = max 𝐹𝑖 . 𝑖 ∈𝑆

Proof. We prove the lemma via a plateau relocation (exchange) argument. Consider a critical path in which the steady phase contains plateaus located in at least two different stages. Let these stages be 𝑥 and 𝑦 with 1 ≤ 𝑥 < 𝑦 ≤ 𝑁 , and assume without loss of generality that 𝐹𝑥 ≤ 𝐹 𝑦 . Let 𝜎𝑥 denote a plateau in stage 𝑥 consisting of 𝑘 ≥ 2 consecutive rightward steps. We construct a new path by transferring two rightward steps from 𝜎𝑥 to a plateau in stage 𝑦. Figure 18 illustrates this operation. After this operation, the modified segment in stage 𝑥 becomes a shorter plateau 𝜎𝑥′ with (𝑘 − 2) steps; if 𝑘 = 2, it disappears. Next, we show that this operation does not decrease the total path cost. Since each rightward step in a plateau contributes one forward and one backward computation, the

Lemma 2. If a critical path contains peaks or valleys in the steady phase, then there exists a critical path whose peak-tops and valley-bottoms are all located at stages with the maximum forward computation time among the stages traversed in the steady phase. That is, for any such stage 𝑥, 𝐹𝑥 = max 𝐹𝑖 , 𝑖 ∈𝑆

where 𝑆 denotes the set of stages traversed by this path in the steady phase. Proof. We prove the lemma by contradiction. We first consider the case of a peak; the case of a valley can be proved symmetrically by filling the valley-bottom. Suppose there exists a critical path whose steady phase contains a peak with its peak-top located at stage 𝑥. Let 𝑆 denote the set of stages traversed by this path in the steady 16

(a)

(b)

(c)

(d)

Figure 20. Examples of P (𝑥, 𝑦, 𝑧). Since 𝑟 > 1 and 𝐹 𝑦 > 𝐹𝑥 , we obtain

phase. Assume, for contradiction, that there exists a stage 𝑦 ∈ 𝑆 such that 𝐹 𝑦 > 𝐹𝑥 . (5)

𝐶 flat > 𝐶 peak .

Thus, replacing the original peak subpath with the flat segment yields a strictly longer path, contradicting the assumption that the original path is critical. Therefore, the peak-top of a critical path cannot be located at a stage whose forward computation time is smaller than that of another stage traversed in the steady phase. That is, the peak-top must lie in a stage 𝑥 satisfying

Without loss of generality, consider the case 𝑦 > 𝑥. Let ℎ = 𝑦 − 𝑥 be the height between stages 𝑥 and 𝑦, and let the peak-top contain 𝑐 ≥ 1 rightward steps. Then the peak subpath first moves upward from stage 𝑦 to stage 𝑥, then moves right along the peak-top, and finally moves downward from stage 𝑥 back to stage 𝑦. We replace this peak subpath with a flat rightward segment at stage 𝑦, consisting of 𝑐 + 4ℎ rightward steps. This operation is referred to as cutting the peak-top. Figure 19 illustrates this replacement. We now compare the costs of the two subpaths. The cost of the original peak subpath can be written as 𝐶 peak =

𝑦−1 ∑︁ 𝑖=𝑥

𝐹𝑥 = max 𝐹𝑖 . 𝑖 ∈𝑆

By Assumption 1, we have 𝐵𝑖 = 𝑟 𝐹𝑖 . Since 𝐹 𝑦 > 𝐹𝑥 and 𝐹 𝑦 is no smaller than the stages involved in this replacement, we have 𝐹𝑖 ≤ 𝐹 𝑦 for 𝑖 ∈ [𝑥, 𝑦]. Moreover, by Assumption 2, (7)

𝐹𝑥 ∗ = max 𝐹𝑖 .

Therefore, 𝐶 peak ≤

𝑖 ∈𝑆

𝑦−1 ∑︁ 𝑖=𝑥

𝑐 −1 (1 + 𝑟 )𝐹 𝑦 + (𝑟 − 1)𝐹 𝑦 + (1 + 𝑟 )𝐹𝑥 2

This completes the proof.

(12) □

Theorem 3. Consider a critical path that enters the steady phase at stage 𝑥 and exits the steady phase at stage 𝑧. There exists a critical path P with the same entry and exit stages, and a stage 𝑦, such that the following properties hold: • In the steady phase, P consists of three segments: 1. starting from the entry at stage 𝑥, it moves continuously upward or downward to stage 𝑦; 2. it then moves continuously rightward along stage 𝑦;

(8)

𝑐 −1 (1 + 𝑟 )𝐹𝑥 . 2 After cutting the peak-top, the replacement subpath at stage 𝑦 has cost     𝑐 −1 𝑐 −1 𝐶 flat = + 2ℎ (𝐹 𝑦 +𝐵 𝑦 ) = + 2ℎ (1 +𝑟 )𝐹 𝑦 . (9) 2 2

(11)

The same argument applies to a valley by reversing the vertical directions, which corresponds to filling the valleybottom. Hence, every peak-top and valley-bottom can be placed at a stage with the maximum forward computation time among stages in 𝑆. If multiple stages attain the same maximum value, the above replacement can be applied without decreasing the path cost, allowing all peak-tops and valley-bottoms to be moved to a single such stage 𝑥 ∗ . Therefore, there exists a critical path in which all peak-tops and valley-bottoms lie in the same stage 𝑥 ∗ , with

𝑐 −1 (𝐹𝑥 + 𝐵𝑥 ). (6) 𝐹𝑖 + 𝐵𝑖 + 𝑡𝑖,𝑖+1 + 𝑡𝑖+1,𝑖 + 2

𝑡𝑖,𝑖+1 + 𝑡𝑖+1,𝑖 ≤ (𝑟 − 1)𝐹 𝑦 .

(10)

= 2𝑟ℎ𝐹 𝑦 +

17

3. from stage 𝑦, it moves continuously upward or downward to the exit at stage 𝑧. • Stage 𝑦 has the maximum forward computation time among all stages traversed by P in the steady phase, i.e., 𝐹𝑦 = max 𝐹𝑖 . min(𝑥,𝑦,𝑧 ) ≤𝑖 ≤max(𝑥,𝑦,𝑧 )

Figure 20 presents four examples of the critical-path structures characterized in Theorem 3.

Figure 21. Relocating rightward steps from one stage to another stage(from blue steps to red steps). The cost change depends on the backward computation time difference of the two stages

Proof. By Lemma 1, there exists a critical path in which all plateaus in the steady phase are located in a single stage that attains the maximum forward computation time among the stages traversed in the steady phase. By Lemma 2, all peak-tops and valley-bottoms can also be placed in such a stage without decreasing the path cost. Therefore, there exists a critical path whose all rightward movements in the steady phase are concentrated in a single stage, denoted as stage 𝑦, where 𝐹𝑦 =

max

min(𝑥,𝑦,𝑧 ) ≤𝑖 ≤max(𝑥,𝑦,𝑧 )

𝐹𝑖 .

(13)

After entering the steady phase at stage 𝑥, the path must first move vertically to stage 𝑦, then move rightward along stage 𝑦, and finally move vertically from stage 𝑦 to the exit stage 𝑧. Hence, the steady-phase portion of P consists of the three segments stated above. This completes the proof. □

Figure 22. Delaying the position of the upward step, the cost change depends on the difference between a Backward time(blue) in stage 𝑤 and a Forward+Backward time(red) in stage 𝑦.

Definition 9. Let P (𝑥, 𝑦, 𝑧) denote the critical path characterized in Theorem 3. After characterizing the critical path in the steady phase, we next analyze its structure in the warmup and ending phases. Lemma 3. Let P be a critical path that enters the ending phase at stage 𝑧. The subpath of P in the ending phase, denoted by Pend , can be adjusted such that it consists of three segments: • starting from stage 𝑧, move continuously upward to stage 𝑤, where 𝐹 𝑤 = max 𝐹𝑖 ;

Lemma 4. Consider a critical path whose ending-phase subpath moves rightward at stage 𝑤. Suppose that, after the last rightward segment in the steady phase, the path moves continuously upward from stage 𝑦 to stage 𝑤. Then 𝐵𝑤 ≥ 𝐹𝑦 + 𝐵𝑦 . Proof. We prove this lemma by contradiction. Suppose that 𝐵 𝑤 < 𝐹 𝑦 + 𝐵 𝑦 . Then one backward computation at stage 𝑤 in the ending phase can be replaced by one forward computation and one backward computation at stage 𝑦 in the steady phase, as illustrated in Figure 22. This replacement preserves the dependency constraints of the DAG and strictly increases the path cost. This contradicts the assumption that the original path is critical. Therefore, we must have 𝐵𝑤 ≥ 𝐹𝑦 + 𝐵𝑦 . □

1≤𝑖 ≤𝑧

• move continuously rightward along stage 𝑤; • from stage 𝑤, move continuously upward to stage 1. Proof. This lemma can also be proved by relocating rightward steps. If Pend contains rightward steps in stages other than 𝑤, these steps can be replaced by rightward steps in stage 𝑤 without decreasing the cost of P, as illustrated in Figure 21. By repeatedly applying this relocation operation until all rightward steps lie in stage 𝑤, we obtain the desired critical path. □

Theorem 4. There exists a critical path P (𝑥, 𝑦, 𝑧), as defined in Definition 9, such that 𝑦 ≤𝑥

and 𝑦 ≤ 𝑧.

Moreover, stage 𝑦 satisfies 𝐹𝑦 ≥

Moreover, lemma 3 has a symmetric lemma for warmup phase. That is, for any critical path P, all of the rightward steps in warmup phase can be relocated to a stage 𝑚, where 𝐹𝑚 = max 𝐹𝑖 and 𝑥 is the enter stage of P. It can be proved

max

𝐹𝑖 .

𝑖 ∈ [𝑦,max(𝑥,𝑧 ) ]

Proof. We first prove that there exists a critical path satisfying 𝑦 ≤ 𝑧. Suppose, for contradiction, that a critical path P (𝑥, 𝑦, 𝑧) satisfies 𝑦 > 𝑧.

1≤𝑖 ≤𝑥

similarly.

18

Figure 23. Comparing the difference between the path P (𝑥, 𝑦, 𝑧) and the path going directly from stage 𝑦 to ending phase, where 𝑧 > 𝑦. Since 𝑦 > 𝑧, after the last rightward segment in the steady phase, the path must move continuously upward from stage 𝑦 to stage 𝑧. After entering the ending phase at stage 𝑧, by Lemma 3, the ending-phase subpath can be adjusted such that its rightward steps are located at a stage 𝑤, where 𝐹 𝑤 = max 𝐹𝑖 . 1≤𝑖 ≤𝑧

Therefore, the case 𝑦 > 𝑥 is also impossible, and we have Combining the two results, there exists a critical path P (𝑥, 𝑦, 𝑧) such that 𝑦 ≤ 𝑥 and 𝑦 ≤ 𝑧. Moreover, by Theorem 3, stage 𝑦 has the maximum forward computation time among the stages traversed by the steady-phase subpath. Since 𝑦 ≤ 𝑥 and 𝑦 ≤ 𝑧, this gives

(14)

Thus, the path continues moving upward from stage 𝑧 to stage 𝑤 before taking rightward steps in the ending phase. By Lemma 4, we have 𝐵𝑤 ≥ 𝐹𝑦 + 𝐵𝑦 .

𝐹𝑦 ≥ This completes the proof.

(15)

𝐹𝑖 .

(19) □

Furthermore, we show that there exists a critical path P (𝑥, 𝑦, 𝑧) with 𝑧 = 𝑦, as illustrated in Figure 20d. Theorem 5. There exists a critical path P (𝑥, 𝑦, 𝑧) such that 𝑧 = 𝑦 ≤ 𝑥. Moreover, it satisfies 𝐹 𝑦 ≥ max 𝐹𝑖 . 𝑖 ∈ [𝑦,𝑥 ]

Proof. We prove the theorem by contradiction. By Theorem 4, there exists a critical path P (𝑥, 𝑦, 𝑧) such that 𝑦 ≤ 𝑥 and 𝑦 ≤ 𝑧. Suppose, for contradiction, that no such critical path satisfies 𝑧 = 𝑦. Then every critical path of this form must have 𝑧 > 𝑦. Consider such a critical path P (𝑥, 𝑦, 𝑧) with 𝑧 > 𝑦. By Theorem 4, we have

(16)

Hence, this replacement strictly increases the path cost while preserving the DAG dependencies, as illustrated in Figure 22. This contradicts the assumption that P (𝑥, 𝑦, 𝑧) is a critical path. Therefore, the case 𝑦 > 𝑧 is impossible, and there exists a critical path satisfying 𝑦 ≤ 𝑧.

max

𝑖 ∈ [𝑦,max(𝑥,𝑧 ) ]

Therefore, replacing one forward-backward pair at stage 𝑦 in the steady phase by one backward computation at stage 𝑤 does not decrease the path cost. We can repeatedly apply this replacement until the path can no longer exit the steady phase at stage 𝑧. Next, we further replace one forward-backward pair at stage 𝑦 by one forward computation and one backward computation at stage 𝑤. Since 𝐹 𝑤 > 0, we have 𝐹𝑤 + 𝐵𝑤 > 𝐹𝑦 + 𝐵𝑦 .

(18)

𝑦 ≤ 𝑥.

𝐹 𝑦 ≥ max 𝐹𝑖 . 𝑖 ∈ [𝑦,𝑧 ]

(20)

By Lemma 3, the ending-phase subpath can be adjusted such that its rightward steps are located at a stage 𝑤, where

(17)

It remains to prove that there exists a critical path satisfying 𝑦 ≤ 𝑥. This part follows by a symmetric argument on the warmup phase. Specifically, if 𝑦 > 𝑥, then before reaching the rightward segment at stage 𝑦 in the steady phase, the path must move continuously downward from stage 𝑥 to stage 𝑦. By applying the analogous replacement argument in the warmup phase, one can construct a path with strictly larger cost, again contradicting the criticality of P (𝑥, 𝑦, 𝑧).

𝐹 𝑤 = max 𝐹𝑖 . 1≤𝑖 ≤𝑧

(21)

We choose 𝑤 as the smallest stage satisfying the above equality. Since stage 𝑦 is also no smaller than all stages in [𝑦, 𝑧] in terms of forward computation time, we have 𝑤 ≤ 𝑦 and 𝐹𝑤 ≥ 𝐹𝑦 . Now consider the portion of P (𝑥, 𝑦, 𝑧) from the last rightward segment in the steady phase to the rightward segment 19

Figure 24. The example of comparing the difference between the path 𝐾 (𝑥, 𝑢, 𝑤)(blue) and the path 𝐾 (𝑣, 𝑣, 𝑤)(red) when 𝑁 = 4, 𝑢 = 2, 𝑣 = 4, 𝑤 = 1, 𝑥 = 3. in the ending phase. Since 𝑧 > 𝑦, this portion first moves downward from stage 𝑦 to stage 𝑧, enters the ending phase, and then moves upward from stage 𝑧 to stage 𝑤. We replace this subpath with one that exits the steady phase directly from stage 𝑦, as shown in Figure 23. This replacement changes the exit stage of the steady phase from 𝑧 to 𝑦. The gain of this replacement is 𝑧 ∑︁

2𝐵 𝑤 − 𝐹𝑖 + 𝐵𝑖 + 𝑡𝑖 −1,𝑖 + 𝑡𝑖,𝑖 −1

.

B.3

Ridge-like Allocation Guarantee

Based on the canonical critical-path structure identified above, we next establish the ridge-like allocation guarantee for layerlevel pruning. Definition 10. Let 𝐶 (P) denote the cost of a path P in the DAG. Theorem 7. For any pair of stages 𝑢 and 𝑣 satisfying 1 ≤ 𝑢 < 𝑣 ≤ 𝑁 , if 𝐹𝑢 ≤ 𝐹 𝑣 , then there exist 𝑥, 𝑦, and 𝑤, with 𝑦 ≠ 𝑢, such that 𝐾 (𝑥, 𝑦, 𝑤) is a critical path.

(22)

𝑖=𝑦+1

Proof. Consider any pair of stages 𝑢 and 𝑣 such that 1 ≤ 𝑢 < 𝑣 ≤ 𝑁 and 𝐹𝑢 ≤ 𝐹 𝑣 . Suppose that 𝐾 (𝑥, 𝑢, 𝑤) is a critical path. It suffices to show that

Since 𝐹 𝑤 ≥ 𝐹 𝑦 ≥ 𝐹𝑖 for all 𝑖 ∈ [𝑦, 𝑧], and by Assumption 2, the communication cost is bounded by the computation cost. Therefore, each term in the summation is non-negative, and the replacement does not decrease the total path cost. Thus, we obtain another critical path whose steady-phase exit stage becomes 𝑧 = 𝑦, contradicting the assumption that no critical path of this form satisfies 𝑧 = 𝑦. Therefore, there exists a critical path P (𝑥, 𝑦, 𝑧) such that 𝑧 = 𝑦 ≤ 𝑥. Finally, since 𝑧 = 𝑦, the maximality property in Theorem 4 reduces to 𝐹 𝑦 ≥ max 𝐹𝑖 . 𝑖 ∈ [𝑦,𝑥 ]

𝐶 (𝐾 (𝑥, 𝑢, 𝑤)) ≤ 𝐶 (𝐾 (𝑣, 𝑣, 𝑤)) .

(24)

Indeed, if Eq. (24) holds, then 𝐾 (𝑣, 𝑣, 𝑤) has a cost no smaller than that of the critical path 𝐾 (𝑥, 𝑢, 𝑤), and hence 𝐾 (𝑣, 𝑣, 𝑤) is also a critical path. Since its middle index is 𝑣 ≠ 𝑢, this gives the desired critical path. Figure 24 illustrates this comparison. We prove Eq. (24) through the following steps. 1. From Definition 7, the cost of 𝐾 (𝑥, 𝑦, 𝑤) can be written as ! 𝑥 −1 ∑︁
𝐶 (𝐾 (𝑥, 𝑦, 𝑤)) = 𝐹𝑖 + 𝐵𝑖 + 𝑡𝑖,𝑖+1 + 𝑡𝑖+1,𝑖 + 𝐹𝑥 + 𝐵𝑥

(23)

This completes the proof.

𝑖=1

□

|

Therefore, we obtain the following theorem. Theorem 6. There exist 𝑥, 𝑦, and 𝑤 such that 𝐾 (𝑥, 𝑦, 𝑤) is a critical path.

{z

downward and upward steps

} (25)

+ 2(𝑁 − 𝑥)𝐹 𝑤 + 2(𝑁 − 𝑦)𝐵 𝑤 | {z } rightward steps in stage 𝑤

+ (𝑀 − 2(𝑁 − 𝑦) − 1) (𝐹 𝑦 + 𝐵 𝑦 ) . | {z }

Proof. By Theorem 5 and Lemma 3, there exists a critical path whose steady-phase and ending-phase structures satisfy the properties stated in Definition 7. Therefore, this critical path can be represented as 𝐾 (𝑥, 𝑦, 𝑤) for some 𝑥, 𝑦, and 𝑤. □

rightward steps in stage 𝑦

2. To prove Eq. (24), we first use the following equivalent form: 𝐶 (𝐾 (𝑥, 𝑢, 𝑤)) ≤ 𝐶 (𝐾 (𝑣, 𝑣, 𝑤)) (26) ⇐⇒ 𝐶 (𝐾 (𝑣, 𝑣, 𝑤)) − 𝐶 (𝐾 (𝑥, 𝑢, 𝑤)) ≥ 0.

Figure 17 illustrates the structure of 𝐾 (𝑥, 𝑦, 𝑤) defined in Definition 7. Theorem 6 reduces the critical path to a concise canonical form, which facilitates the subsequent analysis. 20

Substituting Eq. (25) and 𝐵𝑖 = 𝑟 𝐹𝑖 for 1 ≤ 𝑖 ≤ 𝑁 into Eq. (26), we obtain the following equivalent condition: ! 𝑣 ∑︁
0≤ (1 + 𝑟 )𝐹𝑖 + 𝑡𝑖 −1,𝑖 + 𝑡𝑖,𝑖 −1 − 2(𝑣 − 𝑥)𝐹 𝑤 𝑖=𝑥+1

+ (𝑀 − 2(𝑁 − 𝑢) − 1) (1 + 𝑟 ) (𝐹 𝑣 − 𝐹𝑢 )

Lemma 5. Given the layer sequence 𝑙𝑠1 , 𝑙𝑠2 , · · · , 𝑙𝑠𝑘 , if device memory constraints are temporarily ignored and the sequence is not monotonically non-increasing, then the following adjustment procedure can transform it into a monotonically nonincreasing sequence without increasing the cost of the critical path in the DAG. Adjustment procedure. 1. Identify an index 𝑖 (1 ≤ 𝑖 < 𝑘) such that 𝑙𝑠𝑖 < 𝑙𝑠𝑖+1 . If no such index exists, terminate. 2. Reallocate one layer from stage 𝑠𝑖+1 to stage 𝑠𝑖 : 𝑙𝑠𝑖 ← 𝑙𝑠𝑖 + 1, 𝑙𝑠𝑖+1 ← 𝑙𝑠𝑖+1 − 1.

(27)

+ 2(𝑣 − 𝑢) ((1 + 𝑟 )𝐹 𝑣 − 𝑟 𝐹 𝑤 ) . Therefore, proving Eq. (27) is equivalent to proving Eq. (24). 3. Next, consider the following sufficient condition: 0 ≤ 𝛼 (𝑣 − 𝑥) (1 + 𝑟 )𝐹 𝑤 − 2(𝑣 − 𝑥)𝐹 𝑤 + 2(𝑣 − 𝑢) ((1 + 𝑟 )𝐹 𝑣 − 𝑟 𝐹 𝑤 ) .

(28)

3. Repeat from Step 1.

We have Eq. (28) =⇒ Eq. (27), because the right-hand side of Eq. (28) is no larger than the right-hand side of Eq. (27). This follows from the following facts: • 𝐹 𝑣 − 𝐹𝑢 ≥ 0; • 𝐹𝑖 ≥ 𝛼𝐹 𝑤 for all 𝑖 ∈ [𝑥 + 1, 𝑣]; • 𝑡𝑖,𝑖 −1 + 𝑡𝑖 −1,𝑖 ≥ 0 for all 𝑖 ∈ [𝑥 + 1, 𝑣]. 4. By rearranging the terms in Eq. (28), moving the terms involving 𝐹 𝑤 to the left-hand side and the terms involving 𝐹 𝑣 to the right-hand side, and considering the sign of the coefficient, Eq. (28) can be equivalently written as   2(𝑣 − 𝑢) (1 + 𝑟 ) 𝐹𝑤 ≤ 𝐹 𝑣 . (29) 2(𝑣 − 𝑥) + 2𝑟 (𝑣 − 𝑢) − 𝛼 (1 + 𝑟 ) (𝑣 − 𝑥)

Proof. We show that the layer reallocation in Step 2 does not increase the cost of the critical path. There are two cases. • If 𝑙𝑠𝑖 ≤ 𝑙𝑠𝑖+1 still holds after the adjustment, then by Theorem 7, the case 𝑦 = 𝑠𝑖 can be excluded from the analysis of the critical path 𝐾 (𝑥, 𝑦, 𝑤). Therefore, the cost of the critical path does not increase. • If 𝑙𝑠𝑖 > 𝑙𝑠𝑖+1 holds after the adjustment, then the initial difference between the two stages must be exactly one. In this case, by direct calculation, 𝐶 (𝐾 (𝑥, 𝑠𝑖 , 𝑤)) in the DAG after the adjustment is no greater than 𝐶 (𝐾 (𝑥, 𝑠𝑖+1, 𝑤)) in the DAG before the adjustment. Hence, the cost of the critical path still does not increase. Therefore, each adjustment step preserves or reduces the cost of the critical path. Repeating this procedure eventually transforms the sequence into a monotonically nonincreasing sequence, which proves the lemma. □

Thus, Eq. (29) is equivalent to Eq. (28). 5. Since 𝑢 ≤ 𝑥 < 𝑣, it can be verified that when   2 𝑟 , 𝛼 = max , 3 1+𝑟

Lemma 6. Let 𝑙𝑖𝑚 denote the memory-constrained layer limit of stage 𝑖 (1 ≤ 𝑖 ≤ 𝑁 ). Then the sequence 𝑙𝑠𝑚1 , 𝑙𝑠𝑚2 , · · · , 𝑙𝑠𝑚𝑘 is monotonically non-decreasing.

the following inequality holds: 1 2(𝑣 − 𝑢) (1 + 𝑟 ) ≤ . 𝛼 2(𝑣 − 𝑥) + 2𝑟 (𝑣 − 𝑢) − 𝛼 (1 + 𝑟 ) (𝑣 − 𝑥) Moreover, by Assumption 4, we have 1 𝐹𝑤 ≤ 𝐹𝑣 . 𝛼 Therefore, Eq. (30) implies Eq. (29). 6. Finally, Assumption 4 gives   2 𝑟 𝛼 = max , . 3 1+𝑟

(30)

Proof. Under the 1F1B pipeline paradigm, activation storage requirements decrease monotonically as the stage index increases. Since the considered stages are assigned to device groups of the same type and thus have identical memory capacities, their memory-constrained layer limits form a monotonically non-decreasing sequence. □ Finally, we prove Theorem 2.

(31)

The proof of Theorem 2. Lemma 5 shows that, when device memory constraints are ignored, there exists an optimal layer allocation whose layer sequence over device groups of the same type is monotonically non-increasing. Lemma 6 further shows that the corresponding memory-constrained layer limits form a monotonically non-decreasing sequence. Combining these two monotonicity properties, the feasible optimal allocation must first follow the non-increasing tendency induced by the critical-path cost, while also respecting the non-decreasing upper bounds imposed by memory constraints. Therefore, the optimal layer distribution over device groups of the same type exhibits a unimodal structure. This completes the proof. □

Combining the above implications, we have (31) =⇒ (30) =⇒ (29) ⇐⇒ (28) =⇒ (27) ⇐⇒ (24). Therefore, Eq. (24) holds, and the theorem follows.

□

Theorem 7 shows that when there exist two stages 𝑢 and 𝑣 with 𝐹𝑢 ≤ 𝐹 𝑣 , the case 𝑦 = 𝑢 can be safely excluded from the analysis of the critical path 𝐾 (𝑥, 𝑦, 𝑤). Before proving Theorem 2, we present two supporting lemmas. For any layer allocation scheme 𝑙 1, 𝑙 2, · · · , 𝑙 𝑁 , we focus on the stages 𝑠 1, 𝑠 2, · · · , 𝑠𝑘 that are assigned to device groups of the same type, and consider their corresponding layer sequence 𝑙𝑠1 , 𝑙𝑠2 , · · · , 𝑙𝑠𝑘 .

21