# Trillion Parameter AI Serving Infrastructure for Scientific Discovery: A Survey and Vision

arXiv: https://arxiv.org/abs/2402.03480 · HTML: https://mlsanity.com/paper/2402.03480/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 65 references that are in the map.
- First posted on arXiv 5 Feb 2024.

## Builds on (most cited first, 30 of 34)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Graph Neural Networks: A Review of Methods and Applications](https://mlsanity.com/paper/1812.08434/) (2018)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [LaMDA: Language Models for Dialog Applications](https://mlsanity.com/paper/2201.08239/) (2022)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Galactica: A Large Language Model for Science](https://mlsanity.com/paper/2211.09085/) (2022)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Eliciting Latent Predictions from Transformers with the Tuned Lens](https://mlsanity.com/paper/2303.08112/) (2023)
- [Dissecting Recall of Factual Associations in Auto-Regressive Language Models](https://mlsanity.com/paper/2304.14767/) (2023)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [Progressive Neural Architecture Search](https://mlsanity.com/paper/1712.00559/) (2017)
- [DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale](https://mlsanity.com/paper/2207.00032/) (2022)
- [Caffe: Convolutional Architecture for Fast Feature Embedding](https://mlsanity.com/paper/1408.5093/) (2014)
- [A Comprehensive Survey of Neural Architecture Search: Challenges and Solutions](https://mlsanity.com/paper/2006.02903/) (2020)
- [graph2vec: Learning Distributed Representations of Graphs](https://mlsanity.com/paper/1707.05005/) (2017)
- [Clipper: A Low-Latency Online Prediction Serving System](https://mlsanity.com/paper/1612.03079/) (2016)
- [Neural Architecture Optimization](https://mlsanity.com/paper/1808.07233/) (2018)
- [PanGu-Sigma: Towards Trillion Parameter Language Model with Sparse Heterogeneous Computing](https://mlsanity.com/paper/2303.10845/) (2023)
- [word2vec, node2vec, graph2vec, X2vec: Towards a Theory of Vector Embeddings of Structured Data](https://mlsanity.com/paper/2003.12590/) (2020)
- [funcX: A Federated Function Serving Fabric for Science](https://mlsanity.com/paper/2005.04215/) (2020)
- [MOFormer: Self-Supervised Transformer model for Metal-Organic Framework Property Prediction](https://mlsanity.com/paper/2210.14188/) (2022)
- [Attention Lens: A Tool for Mechanistically Interpreting the Attention Head Information Retrieval Mechanism](https://mlsanity.com/paper/2310.16270/) (2023)
- [On Neural Architecture Search for Resource-Constrained Hardware Platforms](https://mlsanity.com/paper/1911.00105/) (2019)

## Built on by (most cited first, 5 of 5)

- [A Primer on the Inner Workings of Transformer-based Language Models](https://mlsanity.com/paper/2405.00208/) (2024)
- [On the Foundations of Earth and Climate Foundation Models](https://mlsanity.com/paper/2405.04285/) (2024)
- [MLPerf Power: Benchmarking the Energy Efficiency of Machine Learning Systems from Microwatts to Megawatts for Sustainable AI](https://mlsanity.com/paper/2410.12032/) (2024)
- [Large Physics Models: Towards a collaborative approach with Large Language Models and Foundation Models](https://mlsanity.com/paper/2501.05382/) (2025)
- [FrugalPrompt: Reducing Contextual Overhead in Large Language Models via Token Attribution](https://mlsanity.com/paper/2510.16439/) (2025)

## Abstract (arXiv)

Deep learning methods are transforming research, enabling new techniques, and ultimately leading to new discoveries. As the demand for more capable AI models continues to grow, we are now entering an era of Trillion Parameter Models (TPM), or models with more than a trillion parameters -- such as Huawei's PanGu-Sigma. We describe a vision for the ecosystem of TPM users and providers that caters to the specific needs of the scientific community. We then outline the significant technical challenges and open problems in system design for serving TPMs to enable scientific research and discovery. Specifically, we describe the requirements of a comprehensive software stack and interfaces to support the diverse and flexible requirements of researchers.