Training LLMs with MXFP4 Permalink
Albert Tseng, Tao Yu, Youngsuk Park.
CV
Education
- Ph.D in Computer Science, Cornell University, June 2024
- Thesis: Towards Expressive and Robust Learning With Hyperbolic Geometry
- Committee: Christopher De Sa (Chair), Kilian Weinberger, Noah Stephens-Davidowitz
- B.S. in Mathematics and Applied Mathematics (ZhiYuan Honors), Shanghai Jiao Tong University, June 2019
Industrial Experience
- Applied Scientist, Algorithm @ AWS Neuron Science Amazon Web Services, June. 2024 - present
- Developed efficient ML training and inference algorithms for foundation models on cutting-edge hardware, focusing on low-precision training, as well as quantization and sparsity for accelerated inference
- Research Intern, AWS AI Research Amazon Web Services, Aug. 2023 - March 2024
- Developed low-precision strategies for LLMs to mitigate the trade-off between memory and throughput without accuracy degradations
- Research Intern, ML Privacy Team Apple, 2020, 2021, 2022, 2023 summer
- Researched and designed various algorithms to enhance model performance, privacy guarantees, and system latency in federated learning settings
- Research Intern, Visual Computing Microsoft Research Asia, 2019 summer
- Studied and proposed efficient spatial attention mechanisms in graph attention networks
- Research Intern, Dept. of Computer Science Cornell University, July. 2018 - Dec. 2018
- Researched efficient representation learning (e.g. graph networks, hyperbolic space representations)
- Proposed algorithms for effectively detecting white-box adversarial examples to assess model robustness.
Service
PC/Reviewer: NeurIPS, ICML, AISTATS, ICLR, KDD, SDM
Publications
Stochastic Rounding for LLM Training: Theory and Practice Permalink
Kaan Ozkara, Tao Yu, Youngsuk Park.
Collage: Light-Weight Low-Precision Strategy for LLM Training Permalink
Tao Yu, Gaurav Gupta, Karthick Gopalswamy, Amith R Mamidala, et al.
Momentum Approximation in Asynchronous Private Federated Learning Permalink
Tao Yu, Congzheng Song, Jianyu Wang, Mona Chitnis.
Shadow Cones: Unveiling Partial Orders in Hyperbolic Space Permalink
Tao Yu*, Toni J.B. Liu*, Albert Tseng, Christopher De Sa.
Coneheads: Hierarchy Aware Attention Permalink
Albert Tseng, Tao Yu, Toni J.B. Liu, Christopher De Sa.
Random Laplacian Features For Learning with Hyperbolic Space Permalink
Tao Yu, Christopher De Sa.
Understanding Hyperdimensional Computing for Parallel Single-Pass Learning Permalink
Tao Yu*, Yichi Zhang*, Zhiru Zhang, Christopher De Sa.
MCTensor: A High-Precision Deep Learning Library with Multi-Component Floating-Point Permalink
Tao Yu, Wentao Guo, Jianan Canal Li, Tiancheng Yuan, Christopher De Sa.
Representing Hyperbolic Space Accurately using Multi-Component Floats Permalink
Tao Yu, Christopher De Sa.
Salvaging Federated Learning by Local Adaptation Permalink
Tao Yu, Eugene Bagdasaryan, Vitaly Shmatikov.
Numerically Accurate Hyperbolic Embeddings Using Tiling-Based Models Permalink
Tao Yu, Christopher De Sa.
A New Defense Against Adversarial Images: Turning a Weakness into a Strength Permalink
Tao Yu*, Shengyuan Hu*, Chuan Guo, Weilun Chao, Kilian Q. Weinberger.
Simplifying Graph Convolutional Networks Permalink
Felix Wu*, Tianyi Zhang*, Amauri Holanda de Souza Jr.*, Christopher Fifty, Tao Yu, Kilian Q. Weinberger.
Curvature-based Comparison of Two Neural Networks Permalink
Tao Yu, Huan long, John Hopcroft.
Talks
Invited Talk on Revisiting Low Precision Strategies for LLM Training
Talk at SANDS Seminar Series, S3S 2024, KAUST
Invited Talk on Efficient Asynchronous Private Federated Learning
Workshop talk at Workshop on Privacy Preserving Machine Learning, PPML 2024, Online
Invited Talk on Understanding Hyperdimensional Computing for Parallel Single-pass Learning
Talk at VSAONLINE 2022, Online
Conference Proceeding Talk on Numerically Accurate Hyperbolic Embeddings Using Tiling-Based Models
Talk at NeurIPS 2019 Spotlight, Vancouver, BC, Canada