avatar

Ziheng Jiang

Senior Staff Research Scientist
ByteDance
ziheng.misc@gmail.com


About Me

Ziheng works on large language model (LLM) systems at Bytedance, focusing on scaling and optimizing LLM training and inference. He was a Ph.D. student advised by Luis Ceze and Tianqi Chen in the Paul G. Allen School of Computer Science & Engineering at the University of Washington. He received his Bachelor’s degree from Fudan University, where he was a member of Fudan NLP Lab, working with Xipeng Qiu, and Zheng Zhang.

He is/was heavily involved in projects: MegaScale, Apache TVM, Apache MXNet

Research Interests

Selected Publications

Ziheng's work has been cited more than 4000 times. Here are some selected papers in the field of Machine Learning Systems. See a full-list of publication at Google Scholar.

  1. MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs
    Ziheng Jiang*, Haibin Lin*, Yinmin Zhong*, Qi Huang, Yangrui Chen, Zhi Zhang, Yanghua Peng, Xiang Li, Cong Xie, Shibiao Nong, Yulu Jia, Sun He, Hongmin Chen, Zhihao Bai, Qi Hou, Shipeng Yan, Ding Zhou, Yiyao Sheng, Zhuo Jiang, Haohan Xu, Haoran Wei, Zhang Zhang, Pengfei Nie, Leqi Zou, Sida Zhao, Liang Xiang, Zherui Liu, Zhe Li, Xiaoying Jia, Jianxi Ye, Xin Jin, Xin Liu
    21st USENIX Symposium on Networked Systems Design and Implementation (NSDI), 2024.

  2. MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism
    Ruidong Zhu*, Ziheng Jiang*, Chao Jin*, Peng Wu, Cesar A. Stuardo, Dongyang Wang, Xinlei Zhang, Huaping Zhou, Haoran Wei, Yang Cheng, Jianzhe Xiao, Xinyi Zhang, Lingjun Liu, Haibin Lin, Li-Wen Chang, Jianxi Ye, Xiao Yu, Xuanzhe Liu, Xin Jin, Xin Liu
    the ACM Special Interest Group on Data Communication (SIGCOMM), 2025.

  3. Understanding Stragglers in Large Model Training Using What-if Analysis
    Jinkun Liu, Ziheng Jiang, Zuquan Song, Sida Zhao, Menghan Yu, Zhanghan Wang, Chenyuan Wang, Zuocheng Shi, Xiang Shi, Wei Jia, Zherui Liu, Shuguang Wang, Haibin Lin, Xin Liu, Aurojit Panda, Jinyang Li
    USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2025

  4. Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts
    Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Chengquan Jiang, Qi Hou, Weihao Cui, Size Zheng, Li-Wen Chang, Quan Chen, Xin Liu
    the Conference on Machine Learning and Systems (MLSys), 2025.

  5. TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives
    Size Zheng, Jin Fang, Xuegui Zheng, Qi Hou, Wenlei Bao, Ningxin Zheng, Ziheng Jiang, Dongyang Wang, Jianxi Ye, Haibin Lin, Li-Wen Chang, Xin Liu
    the Conference on Machine Learning and Systems (MLSys), 2025.

  6. FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion
    Li-Wen Chang, Wenlei Bao, Qi Hou, Chengquan Jiang, Ningxin Zheng, Yinmin Zhong, Xuanrun Zhang, Zuquan Song, Ziheng Jiang, Haibin Lin, Xin Jin, Xin Liu
    Preprint, 2024.

  7. Characterizing Structural Regularities of Labeled Data in Overparameterized Models
    Ziheng Jiang*, Chiyuan Zhang*, Kunal Talwar, Michael C. Mozer
    International Conference on Machine Learning (ICML), 2021.
    Long Oral (3%)

  8. Learning to optimize tensor programs
    Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin, Arvind Krishnamurthy
    the Conference on Neural Information Processing Systems (NIPS), 2018.

  9. TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
    Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, Arvind Krishnamurthy
    USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2018

  10. Efficient Deep Learning Inference on Edge Devices
    Ziheng Jiang, Tianqi Chen, Mu Li
    the Conference on Machine Learning and Systems (MLSys), 2018.