PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xiaosong Chen, Shaoheng Nie, Zhongmin Zhao, Zizhao Mo, Jiapeng Chen, Huanle Xu, Zeren Li, Weiwei Sun, ChengZhong Xu

arXiv:2609.06086v1cs.DC

TL;DR

Poseidon addresses inaccurate training-time models and costly parallelization search for LLM pre-training on heterogeneous accelerators. It combines a DAG-based estimator with theoretically grounded pruning, and experiments report up to 2.76× higher throughput than baselines.

  • Problem

    Prior approximate models miss communication-computation overlap, and the large configuration space makes efficient parallelization search challenging.

    Source: §2.2 · p. 2, §1 · p. 2

  • Method

    Poseidon estimates training time with a dependency-aware DAG and uses stage-level early stopping and layer-to-stage pruning to search configurations.

    Source: §1 · p. 2

  • Results

    Up to 2.76× higher throughput than baselines on heterogeneous clusters.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The pruning strategies are theoretically guaranteed to identify optimal parallel configurations, while the cost model generalizes to multiple training variants.

    Source: §1 · p. 2

    The motivating limitation is that prior approximate models fail to account for communication-computation overlap, leading to inaccurate training-time estimates.

    Source: §2.2 · p. 2

Abstract

from arXiv · show

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\times$ over state-of-the-art systems.

1 Introduction

Poseidon targets inaccurate training-time estimates and costly parallelization search for LLM pre-training on heterogeneous accelerators. Its DAG model and two pruning strategies support efficient configuration search, and experiments report higher throughput than baselines.

1 Introduction

  • Existing cost models miss communication-computation overlap, while rigid pruning can miss optimal configurations and coarse models fail to generalize across training variants.The search space also grows exponentially with the number of accelerator types, and existing modeling is tied to 1F1B.Source: §1 · pp. 1–2
  • Poseidon models training time with a DAG that captures stage- and batch-level dependencies and estimates time by computing the longest path.The graph also accounts for communication-computation overlap in heterogeneous environments.Source: §1 · p. 2
  • Stage-level early stopping and ridge-pattern layer-to-stage mapping prune the search space with theoretical guarantees of optimal configurations.The stage-level strategy uses partial estimation, while layer-to-stage mapping exploits a ridge-like distribution pattern.Source: §1 · p. 2
  • Up to 2.76× higher throughput than baselines on heterogeneous clusters.The experiments report that Poseidon consistently outperforms baselines.Source: §1 · p. 2

2 Background and Motivation

Accelerator differences complicate LLM parallelization, while prior systems face inaccurate time models, risky pruning, and limited support for pipeline variants. Graph-based modeling can capture execution dependencies, but finding good configurations remains computationally expensive.

2.1 LLM Training under Accelerator Heterogeneity

  • Compute power, memory, and interconnect bandwidth vary across accelerator generations, complicating workload balance and parallelization.LLM training commonly partitions computation across layers, operators, and data using pipeline, tensor, and data parallelism.Source: §2.1 · p. 2

2.2 Limitations of Existing Pre-training Systems

  • Up to 1.32× deviations occur because approximate models miss communication-computation overlap, distorting candidate rankings and search.Figure 2 concerns estimation gaps across cluster configurations.Source: §2.2 · pp. 2–3, Figure 2 · p. 3
  • Coarse-grained models tied to 1F1B do not accurately model Interleaved 1F1B, Eager 1F1B, or Seq1F1B schedules.These variants use bubble-reduction mechanisms that such models do not capture.Source: §2.2 · p. 3
  • Up to 1.91× degradation for Metis and 1.47× for HexiScale results from improper pruning compared with brute-force optimal configurations.The cited comparison uses LLaMA-3 on a cluster of 16 Atlas A2-2 and 16 Atlas A2-3 NPUs.Source: Figure 3 · p. 3, §2.2 · p. 3

2.3 Graph-based Modeling: Opportunities and Challenges

  • Graph-based modeling can represent dependencies across stages and micro-batches, but exhaustive configuration search is impractical in the vast feasible space.Training time is determined by the maximum execution time across micro-batches, motivating dependency-aware estimation.Source: §2.3 · p. 3

3 Poseidon Overview

Poseidon combines a DAG-based training-time model with theoretically grounded pruning to search heterogeneous accelerator configurations efficiently. Its profiling, search, and simulator modules support this process, which also generalizes across pipeline schedules.

3 Poseidon Overview

  • Poseidon supports 1F1B, Eager 1F1B, and Interleaved 1F1B while searching heterogeneous accelerator configurations.Its search combines stage-level and layer-level pruning to reduce search overhead without compromising training efficiency.Source: §3 · p. 3

3.1 Key Design Ideas

  • The DAG represents stage- and batch-level operations and dependencies, with iteration time estimated as the graph’s longest path.This captures overlap between communication and computation.Source: §3.1 · p. 4
  • Stage-level pruning stops exploring a partially allocated configuration when a lower bound indicates it is inefficient, reducing search overhead without compromising solution quality.The lower bound compares the partial configuration’s projected training time with the best result observed so far.Source: §3.1 · p. 4
  • Within each device type, optimal layer allocations rise and then fall across stages, allowing Poseidon to prune allocations that violate this ridge-like pattern.Once allocations begin decreasing, configurations assigning more layers than the preceding stage can be safely pruned.Source: §3.1 · p. 4, Figure 4 · p. 4

3.2 System Architecture

  • Profiling supplies per-layer execution and memory information; the search module applies complementary pruning, and the simulator estimates configuration time from DAG dependencies.The simulator uses dynamic programming to find the longest-time path.Source: §3.2 · p. 4

4 Design Details of Poseidon

Poseidon models training time with a profiled DAG whose critical path estimates execution time, then searches configurations using stage- and layer-level pruning. The search uses a safe partial-configuration lower bound and a theoretically supported ridge-like allocation pattern.

4.1 DAG-based Cost Model Design

  • Poseidon uses a DAG as a quantitative, microbatch-level cost model, computing training time as the critical path length.The model captures execution dependencies across stages and micro-batches.Source: §4.1 · p. 5
  • Each pipeline’s DAG includes forward and backward computation nodes, stage-level all-reduce nodes, and directed edges for execution dependencies.Edges represent inter-stage forward/backward flow, intra-stage schedule order, synchronization, and auxiliary boundaries.Source: §4.1.1 · p. 5, §4.1.2 · p. 5

4.1.3 Weight assignment according to profiling data.

  • Profiling assigns computation and all-reduce costs to nodes and activation-transfer costs to inter-stage edges; intra-stage and control edges have zero weight.Poseidon profiles representative layer types and execution, transfer, and synchronization times under device types and TP degrees.Source: §4.1.3 · pp. 5–6

4.1 DAG-based Cost Model Design

  • The DAG accommodates schedules such as Eager 1F1B and Interleaved 1F1B by adjusting node counts and dependency edges.For Interleaved 1F1B, a virtual pipeline size V increases the computation-node count to 2VNM.Source: §4.1.4 · p. 6

4.2 Search Policy for Optimal Parallel Training

  • Stage-level pruning stops a search branch when the lower bound for its partial stage allocation exceeds the best observed time, without losing optimality.The lower bound uses determined stages’ forward, backward, and DP all-reduce times and the number of micro-batches.Source: §4.2.1 · p. 6
  • For any heterogeneous environment, an optimal allocation exists whose same-device-type stage subsequences follow a ridge-like pattern.The pattern supports pruning; the paper attributes it to early-stage memory pressure and late-stage latency bottlenecks.Source: §4.2.2 · pp. 6–7
  • Merely balancing computation time across stages can increase training time, making stage-time-variance minimization less efficient.This is the paper’s stated implication of the layer-allocation theorem.Source: §4.2.2 · p. 7
  • Poseidon samples 500 feasible configurations to initialize the best-time bound, then searches feasible parallelism and device assignments using both pruning strategies.TP groups are restricted to one node, while DP and PP are unrestricted; all pipelines use identical configurations to mitigate stragglers.Source: §4.2.3 · p. 7

5 Poseidon Implementation

Poseidon is implemented with profiling, search, and simulator modules, with an automated workflow that applies selected configurations to training requests. It also supports LLM training on NPUs through MindSpeed-LLM.

5 Poseidon Implementation

  • Poseidon has profiling, search, and simulator modules; the search and simulator are implemented in C++, and profiling is implemented in Python.The search and simulator exceed 3,000 lines of code each, as does the profiling module.Source: §5 · p. 7
  • The simulator estimates training time from profiling data and the DAG-based cost model, updating node and edge weights without rebuilding the DAG when PP degree is unchanged.Source: §5 · p. 7
  • Poseidon transparently injects search and profiling into training requests and applies the chosen configuration automatically; its NPU support uses MindSpeed-LLM.MindSpeed-LLM is built on Megatron-LM and optimized for NPUs.Source: §5 · p. 7

6 Evaluation

Poseidon is evaluated across heterogeneous accelerator clusters, model scales, and training settings. Its reported results cover throughput, search efficiency, estimation accuracy, pruning, and scalability.

6.1 Experiment Setup

  • Experiments span NPU and GPU clusters, comparing Poseidon with Metis and HexiScale across LLM, MoE, and heterogeneous configurations.The evaluation includes multiple model sizes and cluster deployments.Source: §6.1 · pp. 7–8

6.2 End-to-end Training Performance

  • 1.05–1.94× over Metis and 1.21–1.69× over HexiScale: Poseidon achieves the highest throughput across the evaluated model and cluster settings.Both maximum gains occur on LLaMA-30B; the comparison uses the 1F1B schedule supported by all systems.Source: Figure 7 · p. 8, §6.2.1 · p. 8
  • Poseidon matches brute-force configurations and throughput in Settings 1–2; on LLaMA-30B, it cuts search time by up to 11.34× versus Metis.Profiling time is excluded from reported search overhead and reused across methods.Source: Figure 8 · p. 8, §6.2.2 · p. 8

6.3.1 Training time estimation in end-to-end experi-

  • Estimation errors stay at or below 2% across Settings 1–3; sampled LLaMA-3 70B configurations average 98% prediction accuracy, while pipeline schedule alone changes throughput by up to 1.3×.The schedule comparison holds other factors constant.Source: Figure 9 · p. 9, §6.3.1 · p. 9, Figure 10 · p. 9

6.4 Effectiveness of the Search Algorithm

  • Combined stage- and layer-level pruning converges fastest; without warm-up sampling, simulator calls rise 23.2× versus the optimal count, motivating Poseidon’s 500-sample default.The two pruning strategies guide search toward high-quality regions and tighten bounds to prune suboptimal branches.Source: Figure 11 · p. 9, §6.4.1 · p. 9, Figure 12 · p. 9

6.5 Evaluation in General Settings

  • On Mixtral(8×7B), Poseidon achieves 1.06× and 1.61× higher throughput than Metis and HexiScale, while reducing their search times by 52.07× and 3.15×.The evaluation tests applying the training approach to a MoE model.Source: §6.5.1 · p. 9, Figure 13 · p. 10
  • On GPU clusters, Poseidon achieves 1.74× and 2.76× higher throughput than Metis and HexiScale, respectively, while reducing search time by up to 6.72×.The GPU clusters exhibit greater compute and memory heterogeneity than the NPU clusters.Source: Figure 14 · p. 10, §6.5.2 · p. 10
  • Across 96–288 NPUs, throughput improves 1.50×–1.58× over Metis and 1.45×–1.55× over HexiScale; search speedups are 37.8×–53.0× and 4.5×–5.3×, respectively.These large-scale experiments simulate LLaMA-3 (70B) deployments.Source: Figure 15 · p. 10, §6.5.3 · p. 10, Table 2 · p. 10

7 Related Work

Related work covers automatic parallelism search, DAG-based training models, pipeline schedules, and training frameworks. Poseidon differs through microbatch-level DAG simulation and support for large-scale LLM configuration search.

Automating parallelism on heterogeneous devices

  • Prior heterogeneous-device parallelism systems use dynamic programming, heuristics, graph partitioning, or depth-first search; Poseidon provides theoretical guarantees for optimal LLM configurations on heterogeneous clusters.The discussion names Piper, AMP, Alpa, Galvatron, Sailor, and H2 [38] among prior approaches.Source: §7 · p. 10

DAG-based modeling for training systems

  • Unlike FlexFlow ’s operator-level DNN simulation, Poseidon models LLM parallelism at the microbatch level and supports large-scale configuration search.The passage distinguishes their objectives and simulation granularity.Source: §7 · pp. 10–11

Pipeline schedule

  • Pipeline scheduling work addresses bubbles, memory, and recomputation; Poseidon’s flexible modeling integrates the discussed schedule variants.The passage discusses GPipe, PipeDream, Eager 1F1B, Interleaved 1F1B, Seq1F1B [36], zero-bubble, Zorse [10], and Mario.Source: §7 · p. 11

Training frameworks and systems

  • Existing training systems provide infrastructure for large-scale training; Poseidon’s NPU system is built on MindSpeed-LLM.Examples include Megatron-LM [33], Whale, and MegaScale.Source: §7 · p. 11

8 Conclusion

Poseidon combines DAG-based estimation with stage- and layer-level pruning to search heterogeneous accelerator clusters for efficient LLM training configurations.

8 Conclusion

  • Poseidon uses critical-path analysis of a DAG cost model to capture dependencies and overlap, then applies theoretically grounded stage- and layer-level pruning to reduce search without compromising training efficiency.The conclusion says the system generalizes to paradigms including Eager 1F1B and Interleaved 1F1B.Source: §8 · p. 11

Appendix

The appendix analyzes the theoretical guarantees for both pruning strategies.

Appendix

  • The appendix provides detailed theoretical analyses of the guarantees for stage-level and layer-level pruning.These are identified as Theorem 1 and Theorem 2, respectively.Source: Appendix · p. 14

A Stage-level Pruning Guarantee

Poseidon’s pruning guarantees rest on a lower bound for partial configurations and a ridge-like structure in layer allocation. The arguments show when search branches can be discarded and how same-type stages can be ordered without increasing critical-path cost.

A Stage-level Pruning Guarantee

  • A partial configuration can be safely pruned when its valid lower bound exceeds the best observed training time, because no extension can improve on that time.The guarantee applies to the remaining stage assignments in the recursive search.Source: A Stage-level Pruning Guarantee · p. 14

B Layer-level Pruning Guarantee

  • For stages assigned to device groups of the same type, layer allocations can be transformed into non-increasing order without increasing critical-path cost when memory constraints are ignored.The argument combines this ordering with monotonically non-decreasing memory-constrained layer limits under 1F1B to establish the ridge-like allocation guarantee.Source: B Layer-level Pruning Guarantee · p. 14

B.1 Preliminaries

The preliminaries specify the pipeline variables and execution dependencies used in the DAG analysis. They also state assumptions for the theoretical proofs and describe the phases of Eager-1F1B scheduling.

B.1 Preliminaries

  • The setup defines L as the layer count, N as the pipeline-stage count, M as the micro-batch count, and per-stage layer and forward/backward computation times.It distinguishes per-layer times f_i and b_i from per-micro-batch times F_i and B_i.Source: §B.1 · p. 14
  • DAG edges represent inter-stage and intra-stage execution dependencies, and the critical path is the longest weighted path through the graph.Forward-node inter-stage edges point downward, backward-node edges upward, and intra-stage edges point rightward.Source: §B.1 · p. 14
  • The theoretical analysis assumes backward per-layer time b_i equals r times forward time f_i, with r greater than 1.The text introduces four assumptions for the subsequent theoretical analysis.Source: §B.1 · p. 14
  • Under Eager-1F1B, each stage has warmup, steady, and ending phases, with warmup and ending lengths dependent on its index and N.The caption illustrates the three phases for N = 4 and M = 8.Source: §B.1 · p. 15, Figure 16 · p. 15

B.2 Critical Path Characterization

The critical-path analysis reduces steady-phase paths to a canonical form by relocating plateaus, peaks, and valleys without reducing path cost. It then constrains the warmup and ending portions, yielding a canonical critical path K(x,y,w).

B.2 Critical Path Characterization

  • Plateaus in the steady phase can be relocated to a single stage with maximum forward computation time among the stages visited, without reducing path cost.The proof repeatedly relocates plateau steps until they are merged at one such stage.Source: §B.2 · p. 16
  • Peak-tops and valley-bottoms can likewise be moved to stages with maximum forward computation time among those traversed, without reducing path cost.When multiple stages share the maximum, the extrema can be moved to one common stage.Source: §B.2 · p. 17
  • In the steady phase, a critical path can concentrate all rightward steps at one stage y, reached vertically from entry stage x and left vertically for exit stage z.Stage y has the maximum forward computation time among stages traversed in that steady-phase path.Source: §B.2 · pp. 17–18
  • The ending-phase structure can be adjusted so its rightward steps lie at one stage, with the path moving vertically to that stage and then upward to stage 1.A replacement argument supports relocating the rightward steps without decreasing path cost.Source: §B.2 · p. 18
  • The ending-phase rightward steps of a critical path can be relocated to a single stage w without decreasing path cost.The analogous warmup argument relocates all warmup rightward steps to a stage m with maximum forward computation time.Source: §B.2 · p. 18
  • Theorem 4 constrains the steady-phase middle stage y to have maximum forward computation time over the interval from y to the greater of x and z.Theorem 5 further establishes a critical path with z = y ≤ x.Source: §B.2 · pp. 18, 20
  • Theorem 6 establishes that a critical path can be represented in the canonical form K(x,y,w).The form follows from the characterized steady-phase path and ending-phase structure.Source: §B.2 · p. 20

B.3 Ridge-like Allocation Guarantee

The ridge-like allocation guarantee combines a critical-path property with layer reallocation and memory limits. For stages of the same device type, it supports a unimodal optimal layer distribution.

B.3 Ridge-like Allocation Guarantee

  • When F_u ≤ F_v for u < v, a critical path K(x,y,w) exists with y different from u, excluding that middle-stage case from analysis.This is the critical-path property used in the layer-allocation argument.Source: §B.3 · pp. 20–21
  • Ignoring device-memory constraints, layer counts across stages of the same device type can be made non-increasing without increasing critical-path cost.The adjustment transfers layers from a later stage with fewer layers to the preceding stage until the sequence is non-increasing.Source: §B.3 · p. 21
  • Combining non-increasing cost-driven layer counts with non-decreasing memory-constrained limits yields a unimodal optimal layer distribution across same-type stages.The memory-limit ordering follows from monotonically decreasing activation storage requirements with increasing stage indices under 1F1B.Source: §B.3 · p. 21

Related papers

Loading 2609.06086v1…