Kaihao Ma

dblp:249/5421 · DBLP profile ↗
← Back
8ranked-venue papers
4as first author
8since 2021 · last 2026
0009-0001-3176-7816ORCID · corroborated

Domains — the database's venue-derived domains; a paper can count in several

Systems, architecture and hardware · 6 · 2 first-author · 6 since 2021Databases, data management, data science and information retrieval · 2 · 2 first-author · 2 since 2021Artificial intelligence and machine learning · 1 · 1 first-author · 1 since 2021
YearPublicationVenuePosition
2026 KBest: Efficient Vector Search on Kunpeng CPU
Kaihao Ma, Oleg Senkevich, Daihao Xue, Dmitriy Malyshev, Yangming Lv, Shihai Xiao, Xiao Yan 0002, Alexander Radionov, Weidi Zeng, Yuanzhan Gao, Zhiyu Zou, Xin Yao 0008, Yaoyao Fu, Gongyi Wang, Gong Zhang 0001, Fei Yi, Yingfan Liu
KDD (1)1
2025 Adaptive Parallel Training for Graph Neural Networks
abstract
There are several strategies to parallelize graph neural network (GNN) training over multiple GPUs. We observe that there is no consistent winner (i.e., with the shortest running time), and the optimal strategy depends on the graph dataset, GNN model, training algorithm, and hardware configurations. As such, we design the APT system to automatically select efficient parallelization strategies for GNN training tasks. To this end, we analyze the trade-offs of the strategies and design simple yet effective cost models to compare their execution time and facilitate strategy selection. Moreover, we also propose a general abstraction of the strategies, which allows to implement a unified execution engine that can be configured to run different strategies. Our experiments show that APT usually chooses the optimal or a close to optimal strategy, and the training time can be reduced by over 2x compared with always using a single strategy. APT is open-source at https://github.com/kaihaoma/APT.
Kaihao Ma, Renjie Liu 0001, Xiao Yan 0002, Zhenkun Cai, Xiang Song 0003, James Cheng
PPoPP1
2024 PPS: Fair and efficient black-box scheduling for multi-tenant GPU clusters
Kaihao Ma, Zhenkun Cai, Xiao Yan 0002, Zhi Liu 0002, Yihui Feng, Chao Li 0009, Wei Lin 0022, James Cheng
Parallel Comput.1
2023 FEC: Efficient Deep Recommendation Model Training with Flexible Embedding Communication
abstract
Embedding-based deep recommendation models (EDRMs), which contain small dense models and large embedding tables, are widely used in industry. Embedding communication constitutes the main cost for the distributed training of EDRMs, and thus we propose two strategies to improve its efficiency, i.e.,embedding tiering andpre-fetching. In particular, embedding tiering uses AllReduce to communicate popular embeddings that are accessed frequently. This is counter-intuitive as embeddings belong to the sparse embedding tables, but reasonable because the access pattern of popular embeddings resembles dense models. Pre-fetching starts communication early for embeddings that receive no updates such that they are removed from the critical path of training. We implement embedding tiering and pre-fetching in a system called FEC and compare it with the state-of-the-art systems on real datasets. The results show that FEC consistently outperforms the existing methods on all datasets, and its speed can be up to 6.65x and 2.42x in terms of embedding communication time and training throughput compared with the best performing baseline.
Kaihao Ma, Xiao Yan 0002, Zhenkun Cai, Yidi Wu 0001, James Cheng
Proc. ACM Manag. Data1
2022 TensorOpt: Exploring the Tradeoffs in Distributed DNN Training With Auto-Parallelism
abstract
Effective parallelization strategies are crucial for the performance of distributed deep neural network (DNN) training. Recently, several methods have been proposed to search parallelization strategies but they all optimize a single objective (e.g., execution time, memory consumption) and produce only one strategy. We proposeFrontier Tracking(FT), an efficient algorithm that findsa set of Pareto-optimal parallelization strategiesto explore the best trade-off among different objectives. FT can minimize the memory consumption when the number of devices is limited and fully utilize additional resources to reduce the execution time. Based onFT, we develop a user-friendly system, calledTensorOpt, which allows users to run their distributed DNN training jobs without caring the details about searching and coding parallelization strategies. Experimental results show that TensorOpt is more flexible in adapting to resource availability compared with existing frameworks.
Zhenkun Cai, Xiao Yan 0002, Kaihao Ma, Yidi Wu 0001, James Cheng, Teng Su, Fan Yu 0004
IEEE Trans. Parallel Distributed Syst.3
2022 Elastic Deep Learning in Multi-Tenant GPU Clusters
abstract
We study how to support elasticity, that is, the ability to dynamically adjust the parallelism (i.e., the number of GPUs), for deep neural network (DNN) training in a GPU cluster. Elasticity can benefit multi-tenant GPU cluster management in many ways, for example, achieving various scheduling objectives (e.g., job throughput, job completion time, GPU efficiency) according to cluster load variations, utilizing transient idle resources, and supporting performance profiling, job migration, and straggler mitigation. We propose EDL, which enables elastic deep learning with a simple API and can be easily integrated with existing deep learning frameworks such as TensorFlow and PyTorch. EDL also incorporates techniques that are necessary to reduce the overhead of parallelism adjustments, such as stop-free scaling and dynamic data pipeline. We demonstrate with experiments that EDL can indeed bring significant benefits to the above-listed applications in GPU cluster management.
Yidi Wu 0001, Kaihao Ma, Xiao Yan 0002, Zhi Liu 0002, Zhenkun Cai, James Cheng, Fan Yu 0004
IEEE Trans. Parallel Distributed Syst.2
2021 DGCL: an efficient communication library for distributed GNN training
abstract
Graph neural networks (GNNs) have gained increasing popularity in many areas such as e-commerce, social networks and bio-informatics. Distributed GNN training is essential for handling large graphs and reducing the execution time. However, for distributed GNN training, a peer-to-peer communication strategy suffers from high communication overheads. Also, different GPUs require different remote vertex embeddings, which leads to an irregular communication pattern and renders existing communication planning solutions unsuitable. We propose the distributed graph communication library (DGCL) for efficient GNN training on multiple GPUs. At the heart of DGCL is a communication planning algorithm tailored for GNN training, which jointly considers fully utilizing fast links, fusing communication, avoiding contention and balancing loads on different links. DGCL can be easily adopted to extend existing single-GPU GNN systems to distributed training. We conducted extensive experiments on different datasets and network configurations to compare DGCL with alternative communication schemes. In our experiments, DGCL reduces the communication time of the peer-to-peer communication by 77.5% on average and the training time for an epoch by up to 47%.
Zhenkun Cai, Xiao Yan 0002, Yidi Wu 0001, Kaihao Ma, James Cheng, Fan Yu 0004
EuroSys4
2021 Seastar: vertex-centric programming for graph neural networks
abstract
Graph neural networks (GNNs) have achieved breakthrough performance in graph analytics such as node classification, link prediction and graph clustering. Many GNN training frameworks have been developed, but they are usually designed as a set of manually written, GNN-specific operators plugged into existing deep learning systems, which incurs high memory consumption, poor data locality, and large semantic gap between algorithm design and implementation. This paper proposes the Seastar system, which presents a vertex-centric programming model for GNN training on GPU and provides idiomatic python constructs to enable easy development of novel homogeneous and heterogeneous GNN models. We also propose novel optimizations to produce highly efficient fused GPU kernels for forward and backward passes in GNN training. Compared with the state-of-the art GNN systems, DGL and PyG, Seastar achieves better usability, up to 2 and 8 times less memory consumption, and 14 and 3 times faster execution, respectively.
Yidi Wu 0001, Kaihao Ma, Zhenkun Cai, Tatiana Jin, Boyang Li 0016, Chengguang Zheng, James Cheng, Fan Yu 0004
EuroSys2