Hetu v2: A General and Scalable Deep Learning System with Hierarchical and Heterogeneous Single Program Multiple Data Annotations

TL;DR

Hetu v2 introduces Hierarchical and Heterogeneous SPMD (HSPMD) to extend the traditional SPMD model with asymmetric sharding and hierarchical communication, enabling dynamic strategy switching in distributed DL. It provides two-tier sharding annotations ($HDim$, $HSize$, $DG$ Union, $DS$ Union), a hierarchical communication resolution pipeline, progressive graph specialization to instantiate device-specific execution, and a fused batched-send-receive (BSR) switching mechanism for runtime adaptation. The authors show Hetu achieves comparable or superior performance to specialized systems across heterogeneous clusters, elastic training, and mixed-length data, while preserving a single-program, declarative model. This work offers a flexible, scalable framework for modern large-scale training on diverse hardware and data layouts.

Abstract

The Single Program Multiple Data (SPMD) paradigm provides a unified abstraction to annotate various parallel dimensions in distributed deep learning (DL) training. With SPMD, users can write training programs from the viewpoint of a single device, and the system will automatically deduce the tensor sharding and communication patterns. However, with the recent development in large-scale DL models, distributed training exhibits spatial and temporal workload heterogeneity, arising from both device disparities (e.g., mixed hardware, failures) and data variations (e.g., uneven sequence lengths). Such heterogeneity violates SPMD's assumption of uniform workload partitioning, which restricts its ability to express and optimize heterogeneous parallel strategies effectively. To address this, we propose HSPMD within the Hetu v2 system to achieve general and scalable DL training. HSPMD extends SPMD's annotations to support asymmetric sharding and composes standard communication primitives for hierarchical communication, all while retaining the simplicity of a single-device declarative programming model. Leveraging HSPMD, Hetu handles spatial heterogeneity through progressive graph specialization, enabling device-specific execution logic, and addresses temporal heterogeneity via dynamic graph switching. Evaluations on heterogeneous clusters, elastic training, and mixed-length data scenarios show that HSPMD matches or outperforms specialized systems, providing a flexible and efficient solution for modern large-scale model training.

Figures (19)

• Figure 1: Three cases use heterogeneous parallel strategies to mitigate workload imbalance. "DP, CP, TP, PP" represent data, context, tensor and pipeline parallelism, respectively.
• Figure 2: From SPMD to HSPMD. Left: SPMD annotations and their deduction through the computation graph when employing data parallelism (DP) and tensor parallelism (TP) across GPUs 0-4. These semantics can only express uniform workload partitioning and symmetric communication across devices. Right: An example using HSPMD to express heterogeneous parallel strategies: tensor parallelism (TP) between GPU 0,3 and GPU 5,6; pipeline parallelism (PP) between GPU1 and GPU5,6; context parallelism (CP) between GPU2,4; and heterogeneous data parallelism (DP) across these groups. "RS, BSR" represents reduce-scatter and batched-send-receive, respectively.
• Figure 3: HSPMD's sharding annotations (aligned with Figure \ref{['fig:SPMD_vs_HSPMD']}). The tensor is first sharded along the HDim dimension into HSize sharding subgroup. Each sharding subgroup then applies a corresponding DS from the DS Union. The devices in each sharding subgroup are given by the corresponding DG within the DG Union.
• Figure 4: HSPMD's communication resolution. Orange denotes bottom-tier communication, executed independently within sharding subgroups; blue indicates top-tier communication, involving interactions between subgroups. "SR": send-receive; "AR, RS, AG": all-reduce, reduce-scatter, all-gather; "BSR": batched-send-receive; "$\times$": unsupported; "BC": bottom-tier communication.
• Figure 5: Bottom-tier collective communication. The transformation from Partial ($d=-2$) to Duplicate ($d=-1$) triggers an all-reduce (AR), while the transformation from Partial ($d=-2$) to Split ($d\ge0$) or from Split ($d\ge0$) to Duplicate ($d=-1$) triggers reduce-scatter (RS) and all-gather (AG), respectively..