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

News

Selected Publications

Ziheng's work has been cited more than 3000 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. 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.

  3. DietCode: Automatic Optimization for Dynamic Tensor Programs
    Bojian Zheng*, Ziheng Jiang*, Cody Hao Yu, Haichen Shen, Joshua Fromm, Yizhi Liu, Yida Wang, Luis Ceze, Tianqi Chen, Gennady Pekhimenko
    the Conference on Machine Learning and Systems (MLSys), 2022.

  4. 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%)

  5. A Hardware–Software Blueprint for Flexible Deep Learning Specialization
    Thierry Moreau, Tianqi Chen, Luis Vega, Jared Roesch, Eddie Yan, Lianmin Zheng, Josh Fromm, Ziheng Jiang, Luis Ceze, Carlos Guestrin, Arvind Krishnamurthy
    IEEE/ACM International Symposium on Microarchitecture (IEEE MICRO), 2019.
    Best Paper

  6. 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.

  7. 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

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