AsyncHZP: Hierarchical ZeRO Parallelism with Asynchronous Scheduling for Scalable LLM Training

TL;DR

AsyncHZP introduces asynchronous hierarchical ZeRO parallelism to address memory and communication bottlenecks in large-scale LLM training. It adaptively reshards parameters, gradients, and optimizer states across multiple ZeRO dimensions ($Z_1$, $Z_2$, $Z_3$) and employs a multi-stream scheduling workflow to prefetch parameters and overlap gradient reduction with computation. Empirical results on Dense and MoE models show consistent performance gains over traditional ND parallelism, with robust stability and improved scalability across long sequences and massive clusters. The approach is designed to be integration-friendly with existing ND techniques, offering a practical, scalable alternative for state-of-the-art large-scale training without heavy tuning.

Abstract

The training efficiency and scalability of language models on massive clusters currently remain a critical bottleneck. Mainstream approaches like ND parallelism are often cumbersome and complex, while flexible alternatives such as the Zero Redundancy Optimizer (ZeRO) are frequently hampered by communication overhead. In this paper, we propose Asynchronous Hierarchical Zero Parallelism (AsyncHZP), a novel asynchronous variant of ZeRO designed to achieve superior performance while maintaining simplicity and memory efficiency. Unlike traditional ZeRO, which employs over-fine-grained sharding that can lead to inefficient communication, AsyncHZP adaptively reshards parameters, gradients, and optimizer states across different replica groups. This strategy optimizes device memory utilization and significantly reduces communication overhead. In addition, we also design a multi-stream asynchronous scheduling method that executes parameter all-gather and gradient reduce-scatter operations in dedicated background threads, effectively overlapping communication with computation while incurring negligible memory fragmentation. Empirical evaluations on both Dense and Mixture-of-Experts (MoE) models confirm that AsyncHZP maintains robust stability at scale. It consistently outperforms classic ND parallelism, achieving state-of-the-art performance without complex strategic tuning, thereby simplifying the path to efficient large-scale training.

Figures (5)

• Figure 1: Illustration of Asynchronous Hierarchical Zero Parallelism. Unlike traditional Zero implementation that reshards parameters, gradients and optimizer states across the whole Data Parallelism dimension, AsyncHZP could adaptively reshards these with some additional dimension, such as ZeRO1 Parallelism and ZeRO3 Parallelism. And we could asynchronously pre-launch the all-gather of parameters according to the available memory to achieve less blocking for the layer-wise computation and hungrily consume the reduce-scatter of gradients to reduce memory peak.
• Figure 2: Illustration of Forwarding-backwarding for Hierarchical Zero Parallelism with Asynchronous Scheduling.
• Figure 3: Multi-stream asynchronous scheduling for non-blocking communications.
• Figure 4: Comparison between tensor parallelism(TP) and AsyncHZP when coupled with context parallelism(CP).
• Figure 5: AsyncHZP with Pipeline Parallelism with shared communication operations.