PaperlayerView exact revision on arXiv

Source-linked AI summary

HAPMoE: Heterogeneity-Aware Automatic Parallelism Planning for Mixture-of-Experts Models Training

Mengyuan Fan, Peizhuang Cong, Zixiao Huang, Si Xu, Tong Qiao, Yanghao Li, Jing Yang, Tong Yang, Quanlu Zhang, Yu Wang

arXiv:2609.39350v1cs.DCcs.LG

TL;DR

Automatic parallelization faces a gap because existing approaches generally target MoE architectures or heterogeneous clusters, rather than both together. HAPMoE profiles a target cluster, builds an MoE-aware cost model, and searches six parallel dimensions. It reports up to 3.2× higher end-to-end throughput across representative mixed-hardware clusters and up to 78% gains from non-uniform partitioning.

  • Problem

    Existing automatic-parallelization approaches typically target MoE architectures or heterogeneous clusters, leaving their combination insufficiently addressed despite MoE routing and bandwidth asymmetry complicating planning.

    Source: §1 · p. 1

  • Method

    HAPMoE profiles compute kernels and collectives, builds an MoE-aware cost model, and searches DP, PP, TP, CP, EP, and TPE to produce Megatron-LM plans.

    Source: §1 · p. 2

  • Results

    Across representative mixed-hardware clusters, HAPMoE improves end-to-end training throughput by up to 3.2×; non-uniform partitioning yields 4%–78% gains depending on cluster composition.

    Source: §1 · p. 2

  • Takeaways & Limitations

    HAPMoE provides a practical framework for scaling Transformer models on heterogeneous accelerator clusters, with search overhead under one minute.

    Source: §6 · p. 9

    HAPMoE targets Megatron-LM-style Transformer and MoE training, plans from warm-up profiles before training, and does not yet support online scheduling or dynamic reconfiguration.

    Source: Limitations · p. 10

Abstract

from arXiv · show

As model sizes continue to scale, distributed training has become inevitable. Automatic parallelization techniques can derive efficient training parallelism strategies at low cost while achieving superior performance. The difficulty of this problem is jointly determined by the complexity of the model and the underlying compute cluster. Meanwhile, mixture-of-experts (MoE) models are increasingly emerging as the dominant architecture and the rapid evolution of accelerator hardware has made cluster heterogeneity commonplace, posing substantial challenges to automatic parallelization. However, existing approaches typically target either MoE architectures or heterogeneous clusters, failing to generalize to scenarios where both challenges coexist. To this end, we present HAPMoE, a heterogeneity-aware automatic parallelism planner for MoE training. HAPMoE builds a lightweight MoE-aware cost model and efficiently searches a six-dimensional parallel space, producing parallel plans directly deployable on Megatron-LM. Experiments show that HAPMoE improves end-to-end training throughput by up to 3.2$\times$ over baselines across heterogeneous clusters. Its non-uniform pipeline partitioning yields an additional up to 78% gains, and its pruning-enhanced dynamic programming algorithm completes the search within 1 minute, demonstrating high efficiency and practical value in complex hardware environments.

1 Introduction

HAPMoE addresses the gap between MoE-aware planning and heterogeneous-cluster planning with a system that models MoE behavior and searches a six-dimensional parallel space. It produces deployable Megatron-LM plans and reports throughput gains alongside a sub-minute search.

1 Introduction

  • Existing systems typically address MoE or hardware heterogeneity separately, although MoE routing, communication, and load imbalance complicate planning on heterogeneous clusters.Bandwidth asymmetry makes MoE training especially sensitive to hardware differences.Source: §1 · p. 1
  • HAPMoE profiles compute kernels and collectives, builds an MoE-aware cost model, and searches DP, PP, TP, CP, EP, and TPE to generate Megatron-LM plans.Plans include parallel configuration, stage partitioning, device mapping, and recomputation policy.Source: §1 · p. 2
  • Across representative mixed-hardware clusters, HAPMoE improves end-to-end training throughput by up to 3.2×.Ablations also show that it consistently outperforms widely used baselines on dense models and in non-heterogeneous settings.Source: §1 · p. 2
  • 4%–78% throughput gains, depending on cluster composition, come from non-uniform pipeline/data-parallel partitioning.The design jointly optimizes stage-wise parallelism knobs and layer allocation rather than enforcing uniform settings.Source: §1 · p. 2
  • The pruning-enhanced dynamic-programming search reduces the effective search space and completes within <1 minute in all evaluated settings.The stated search-space reduction enables strategy search with sub-minute overhead.Source: §1 · p. 2

2 Background and Related Work

The background reviews parallelism strategies and MoE training constraints, then surveys automatic-parallelism, heterogeneous-cluster, and MoE-acceleration systems. It highlights the need to account jointly for routing, communication, memory, and hardware differences.

2.1 Preliminary

  • DP, PP, TP, and CP divide training data, layers, operators, and sequence dimensions, while MoE adds EP and TPE for distributing or sharding experts.These strategies trade memory and compute capacity against communication, pipeline bubbles, or configuration complexity.Source: §2.1 · p. 2
  • Sparse activation keeps active compute near that of smaller dense models, but routing drift, expert-load variation, memory pressure, and All-to-All communication complicate MoE training.The text calls for resource-aware parallelism that jointly accounts for routing, communication, memory, and hardware heterogeneity.Source: §2.1 · p. 2

2.2 Related Studies

  • Automatic-parallelism systems use search, cost models, and scheduling, but many target dense models or homogeneous settings, limiting applicability to heterogeneous MoE training.Examples include Alpa and FlexFlow [Zheng et al., 2022; Jia et al., 2019] and LLM-oriented Galvatron and Merak [Miao et al., 2022; Lai et al., 2023].Source: §2.2 · pp. 2–3
  • Heterogeneous-cluster and MoE systems improve placement, communication, routing, or load balancing, yet do not fully integrate heterogeneous hardware modeling with global multi-dimensional parallelism search.The review also notes brittleness under extreme network variability and sparse routing.Source: §2.2 · p. 3

3 HAPMoE Design

HAPMoE profiles heterogeneous hardware, models MoE costs, and searches for a feasible parallel plan that minimizes iteration latency. Its design combines stage-level cost modeling with a pruned search over non-uniform pipeline partitions and device assignments.

3 HAPMoE Design

  • HAPMoE searches the 6D parallel configuration, pipeline partition, device assignment, and recomputation policy using a pruning-enhanced pipeline.The profile–model–search workflow produces a plan directly deployable in Megatron-LM.Source: §3 · p. 3

3.1 Problem Formulation

  • The formulation minimizes per-iteration latency for an MoE Transformer while respecting per-stage memory constraints on a heterogeneous cluster.Inputs include model and training parameters, hardware inventory, memory budgets, and profiled bandwidth and collective characteristics.Source: §3.1 · p. 3

3.2 MoE-Aware Performance Modeling

  • The MoE-aware model estimates profiled layer compute, routing and dispatch/combine communication, and stage memory costs, accounting for load imbalance and feasible recomputation.It uses lookup tables for most terms and models the exposed dispatch/combine communication separately from overlap reflected in profiled native collective times.Source: §3.2 · p. 4

3.3 Heterogeneity-Aware Device Modeling

  • Heterogeneity-aware device modeling uses warm-up profiles of compute, communication primitives, and memory to score candidate plans under the 1F1B schedule.The steady-state latency term is determined by the bottleneck pipeline stage.Source: §3.3 · p. 5

3.4 Search Algorithm

  • The search allows stage layer counts to vary so workloads can match heterogeneous device capacities, evaluating candidates under latency and memory-feasibility constraints.Its dynamic-programming state tracks assigned layers and remaining device-group availability; pruning uses load balance, memory feasibility, and heterogeneous communication criteria.Source: §3.4 · p. 5

4 Megatron-LM Framework Integration

HAPMoE integrates with Megatron-LM through a script-in, script-out workflow. It extracts runtime configuration, profiles a few warm-up iterations, and emits a launcher for the selected plan.

4 Megatron-LM Framework Integration

  • HAPMoE accepts an unmodified Megatron-LM training script and emits a runnable launcher without model-code refactoring or manual distributed-knob tuning.The launcher includes the parallel configuration, non-uniform layer partitions, stage-to-device mapping, and recomputation directives.Source: §4 · p. 6
  • It parses runtime arguments for model shape, MoE settings, batch sizes, and optimizer/runtime options to supply metadata for profiling, feasibility checks, and plan materialization.The extracted metadata serves as the interface between Megatron-LM and HAPMoE.Source: §4 · p. 6
  • A few warm-up iterations instrument iteration timing, router statistics, and memory footprints, producing cached tables for the planner.Instrumentation uses a few stable training-loop hooks and is removable for normal training.Source: §4 · p. 6

5 Evaluation

HAPMoE is evaluated across dense and MoE models on homogeneous and heterogeneous clusters, with results covering throughput, MFU, latency, search efficiency, and non-uniform partitioning. Its heterogeneous MoE results show the strongest overall performance, while ablations highlight the importance of non-uniform PP/DP.

5.1 Experimental setup

  • Experiments use Mixtral-S/L and LLaMA-2 7B/13B, H800, MI300X, and 910B accelerators, and baselines tailored to each model and cluster setting.Training uses micro-batch size 1, maximum sequence length 4096, and a six-dimensional search space.Source: §5.1 · p. 6

5.2 Results of Dense Model on Heter. Clusters

  • On a cluster with 2×8 A100 and 2×8 910B accelerators, HAPMoE achieves higher throughput and MFU than MI, Alpa, and Metis for both LLaMA-2 sizes.Its gains increase with model scale.Source: Figure 2 · p. 7, §5.2 · p. 7

5.3 Results of MoE Model on Homo. Clusters

  • Across homogeneous clusters and model scales, HAPMoE consistently outperforms in throughput and MFU, with gains becoming more pronounced as model scale increases.Its routing-aware search and cost model account for routing imbalance and dispatch costs.Source: Figure 3 · p. 7, §5.3 · p. 7

5.4 Results of MoE Model on Heter. Clusters

  • HAPMoE achieves the best average performance, improving geometric-mean throughput over MI by 1.67× on Mixtral-S and 1.78× on Mixtral-L.Iteration latency falls to 0.56× and 0.58×, respectively.Source: Table 2 · p. 8, §5.4 · p. 8

5.5 Search Time and Accuracy Analysis

  • HAPMoE search remains below 60 s across model and cluster scales, while latency and VRAM estimation errors remain low.The authors report reliable prediction under complex configurations.Source: Figure 5 · p. 8, §5.5 · p. 8

5.6 Disabling Non-uniform PP/DP

  • Disabling non-uniform PP/DP raises heterogeneous-cluster latency to 2.3–3.4× and reduces throughput to 45–62%.Separately, four-node profiles predict 16-node latency within about 10% error across homogeneous and heterogeneous settings.Source: Table 3 · p. 8, §5.6 · pp. 8–9

6 Conclusion

HAPMoE combines lightweight profiling, an MoE-specific cost model, and six-dimensional parallelism search to address MoE sparsity and hardware diversity. Experiments report throughput improvements up to 3.2×, pipeline gains up to 78%, and search overhead under one minute.

6 Conclusion

  • HAPMoE addresses MoE sparsity and hardware diversity through lightweight profiling, an MoE-specific cost model, and efficient six-dimensional parallelism search.It generates hybrid strategies, including non-uniform pipeline and expert parallelism, while balancing compute, communication, and memory.Source: §6 · p. 9
  • Experiments on dense LLMs and MoE models show up to 3.2× throughput improvement and up to 78% pipeline gain over SOTA baselines.Search overhead is under one minute.Source: §6 · p. 9
  • HAPMoE provides a practical, robust solution for scaling large Transformer models on heterogeneous accelerator clusters.Source: §6 · p. 9

Limitations

HAPMoE targets Megatron-LM-style Transformer and MoE training, planning from warm-up profiles before training starts. Its current scope excludes online scheduling and dynamic reconfiguration.

Limitations

  • HAPMoE targets Megatron-LM-style Transformer and MoE training, building plans from warm-up profiles before training starts.Source: Limitations · p. 10
  • Larger production clusters may introduce stronger network contention, failures, and topology effects than those covered in the experiments.Source: Limitations · p. 10
  • HAPMoE does not yet support online scheduling or dynamic reconfiguration as routing, workloads, or device availability change during training.The authors leave larger-scale deployment studies and online adaptive scheduling to future work.Source: Limitations · p. 10

Ethical considerations

This work studies training-system efficiency using existing model architectures and infrastructure measurements. It introduces no new datasets, human-subject data, or user-facing model capabilities.

Ethical considerations

A Model Details

The evaluation uses two Mixtral-style MoE configurations, Mixtral-S and Mixtral-L, with BF16 mixed-precision training and a ZeRO-1 distributed optimizer.

A Model Details

  • Mixtral-S (M1) is a 24-layer MoE Transformer with 8 experts, top-k = 2 routing, and an all-gather dispatcher.It has hidden size 4096, FFN dimension 14336, 32 attention heads, and global batch size 256.Source: A Model Details · p. 13
  • Mixtral-L (M2) is a 48-layer, approximately 122B-parameter MoE model with 8 experts per layer and top-k = 2 routing.It uses an all-to-all dispatcher, hidden size 6144, FFN dimension 16384, and global batch size 512.Source: A Model Details · p. 13
  • Both models use BF16 mixed-precision training with a ZeRO-1 distributed optimizer; Mixtral-L on the smallest clusters offloads FP32 optimizer states to host memory.Source: A Model Details · p. 13

B 1F1B Pipeline Schedule Illustration

The 1F1B schedule model separates warm-up and cooldown from steady state, whose throughput is set by the bottleneck pipeline stage. The appendix gives the complete HS-DP search routine.

B 1F1B Pipeline Schedule Illustration

  • In the 1F1B schedule, steady-state throughput is determined by the bottleneck pipeline stage after warm-up and before cooldown.The schedule illustration supports the iteration-level latency model; the bottleneck is defined using stage forward and backward times.Source: B 1F1B Pipeline Schedule Illustration · p. 13
  • HS-DP considers device and end-layer candidates, selects feasible recomputation under stage memory limits, evaluates stage time and communication, then updates its dynamic program.The search reconstructs the optimal path by backtrace.Source: B 1F1B Pipeline Schedule Illustration · p. 13

C HS-DP algorithm pseudo-code

  • The appendix’s complete HS-DP routine covers stage construction, feasibility filtering, and dynamic-programming transitions.The main text presents the modeling and pruning principles.Source: C HS-DP algorithm pseudo-code · p. 13

D Detailed Results of Profile Scaling

Profiling on 4 nodes and reusing those results for 16-node searches produced latency predictions close to measurements for two models in homogeneous and heterogeneous clusters.

D Detailed Results of Profile Scaling

  • Within approximately ±10% in homogeneous clusters and at most about 14% in heterogeneous clusters, 4-node profiles predicted 16-node latency for both models.The largest reported deviation was −13.92% for M1 on a scaled H800/910B cluster.Source: D Detailed Results of Profile Scaling · pp. 14–15
  • The evaluation profiled two models on a 4-node cluster and reused those results to guide parallel-strategy searches on a 16-node cluster.It tested homogeneous and heterogeneous cluster configurations.Source: D Detailed Results of Profile Scaling · p. 14
  • For M1 on a scaled H800/910B cluster, predicted latency was 1150 ms versus 1336 ms measured, with −13.92% error.The table compares latency predictions from 4-node profiles with measurements at larger scale.Source: D Detailed Results of Profile Scaling · p. 15
Loading 2609.39350v1…