On the Opportunities and Risks of Foundation Models

Computer systems are one of the largest bottlenecks to developing foundation models. Foundation models are frequently too large to fit in the main memory of a single accelerator (e.g., GPU) and require an immense amount of computation to train (e.g., >1000 petaFLOP/s-days for GPT-3 (Brown et al., 2020)). Additionally, these models will likely get larger over time: for instance, the compute and memory requirements of state-of-the-art language models have grown by three orders of magnitude in the last three years, and are projected to continue growing far faster than hardware capabilities (Figure 19). Once trained, these large models are expensive to perform inference with and difficult to debug, monitor, and maintain in production applications. We believe that further advances in the performance and usability of foundation models will require careful co-design across algorithms, models, software, and hardware systems, as well as new interfaces for programming and deploying ML applications. In this section, we discuss the key computer systems challenges in developing and productionizing large-scale foundation models.

Improving performance through co-design

Today, training large-scale foundation models (Brown et al., 2020; Rae et al., 2021; NVIDIA and Microsoft, 2021; Wang and Komatsuzaki, 2021) can often require custom software systems such as Megatron, DeepSpeed, or Mesh Transformer JAX (Shoeybi et al., 2019; Rasley et al., 2020; Wang, 2021), built on top of standard frameworks like PyTorch, TensorFlow, and JAX (Paszke et al., 2019; Abadi et al., 2016; Bradbury et al., 2018). These software systems rely on a number of innovations across the stack to train models efficiently at scale: new parallelization dimensions such as pipeline parallelism (Huang et al., 2019; Narayanan et al., 2019) that limit communication while keeping devices busy, state-sharding optimizers to reduce memory usage (Rajbhandari et al., 2020), just-in-time (JIT) compilers to optimize the computation graph (PyTorch, 2021), and optimized libraries like cuDNN and NCCL (NVIDIA, 2021). Megatron and DeepSpeed are efficient to a particular scale; for example, Megatron can extract up to 52% of the theoretical peak throughput of modern hardware with approximately 3000 GPUs on a model with a trillion parameters (Narayanan et al., 2021b). However, scaling to larger models with more GPUs still is challenging, since existing parallelization strategies break down at larger GPU counts. Data parallelism is limited by the batch size (Li et al., 2020e), pipeline parallelism by the number of layers in the model (Huang et al., 2019; Narayanan et al., 2019), and tensor model parallelism by the number of GPUs in a single server (Shoeybi et al., 2019).

While we will continue to realize performance gains from new hardware, growth in the resource requirements of large models far outstrips generational hardware improvements (Brown et al., 2020). To facilitate the next major leap in model capacity and to democratize the advances in model quality, it will be increasingly critical to co-design training algorithms, models, software, and hardware, because many of the avenues to dramatically increase performance alter the semantics of the training computation. For example, executing operations in lower precision (such as fp16) can help increase throughput on modern hardware (e.g., the V100 and A100 GPUs have dedicated tensor core units for lower-precision matrix multiplication), but also affect the numerics of the optimization procedure (Micikevicius et al., 2017). Similarly, exploiting weight sparsity can significantly improve training and inference times (Elsen et al., 2020; Gale et al., 2020) by only performing mathematical operations on the non-zeros in the model, but requires different training algorithms (Jayakumar et al., 2021; Evci et al., 2020; Dettmers and Zettlemoyer, 2019). Other examples of co-design include model architectures that map more efficiently to hardware (So et al., 2019; Child et al., 2019; Wang et al., 2020c; Lee-Thorp et al., 2021; Kitaev et al., 2020; Beltagy et al., 2020; Tay et al., 2020; Ren et al., 2021), efficient optimizers (Anil et al., 2020; Shazeer and Stern, 2018), novel tokenization alternatives (Xue et al., 2021; Tay et al., 2021), specially architected hardware training platforms (Jouppi et al., 2017; Mudigere et al., 2021; Selene, 2021), and distributed parallelization strategies with relaxed weight update semantics (Narayanan et al., 2019; Narayanan et al., 2021a).

Case study: efficient knowledge representation

As a concrete case study of successful co-design, retrieval-based models such as REALM, RAG, ColBERT-QA, and RETRO (Guu et al., 2020; Lewis et al., 2020b; Khattab et al., 2020; Borgeaud et al., 2021) take a different approach to model design than simply increasing the number of model parameters. Instead of trying to accumulate implicit knowledge from ever-larger datasets directly into a DNN model with billions of parameters (like GPT-3), retrieval-based models store knowledge outside the model parameters in the form of text passages, capturing knowledge within the passages with dense vector representations. These models then use scalable top-k search mechanisms to extract knowledge pertinent to each input, while keeping the DNN model itself small (§4.1.4: modeling-memory). This design improves computational efficiency as well as maintainability of the model in production: for example, developers can update the knowledge of the model just by replacing a text passage, without needing to retrain a large DNN.

Retrieval-based models have achieved promising initial results by leveraging several new cross-functional ideas, including backpropagating the loss through the retriever during training (Guu et al., 2020) (which requires approximating the gradient through a knowledge store consisting of millions of passages) and modeling fine-grained interactions between queries and passages (Khattab and Zaharia, 2020; Khattab et al., 2020) (which requires decomposing the computation into vector-level nearest-neighbor search operations). These techniques allow retrieval-based models to be accurate and efficient, but demand functionality not readily supported by popular ML frameworks and nearest-neighbor indexes such as FAISS (Johnson et al., 2019).

Automated optimization

Another important challenge in systems is to automate the application of optimizations that straddle algorithms, models, software, and hardware. While many optimizations and parallelization strategies are complementary, identifying the most effective combination of optimizations is challenging since the joint search space grows combinatorially and optimizations interact in non-trivial ways (Narayanan et al., 2021b). Foundation models heighten the need for automated optimization as manual experimentation is extremely expensive and time-consuming at the scale of thousands of GPUs.

Recent work in this area has focused on systems targeting semantics-preserving optimizations. In particular, systems have been proposed to automatically discover mathematically-equivalent graph substitutions (Jia et al., 2019a; Wang et al., 2021c), facilitate the distributed execution of operator graphs through both high-level APIs and low-level compilers (Rasley et al., 2020; Mandeep Baines, 2021; Bradbury et al., 2018; Shazeer et al., 2018; Lepikhin et al., 2020), and automate the selection of hybrid distribution strategies (Jia et al., 2019b; Santhanam et al., 2021). These systems have helped deploy many foundation models in industry (Fedus et al., 2021; M2M-100, 2020; Turing-NLG, 2020).

Unfortunately, automated optimization becomes much harder when composing semantics-altering optimizations (§4.5.1: systems-co-design), as it is often unclear how to jointly model the statistical impacts of these techniques (e.g., how many training iterations are needed to reach a specific accuracy?). We will therefore need new software tools, libraries, and compilers to automatically identify compositions of optimizations that target comprehensive metrics like time-to-accuracy (Coleman et al., 2017; Mattson et al., 2020). Building such tools will require tight collaboration between systems and machine learning experts.

Execution and programming models

The unique multi-task nature of foundation models provides an opportunity to amortize training and inference costs over many applications. In particular, paradigms such as adaptation mean more sharing across model instances (Raffel, 2021). For example, two models prefix-tuned (Li and Liang, 2021) from the same pretrained model can share the same model “stem,” reducing the storage footprint (the shared stem only needs to be stored once), while also making it possible for execution to be shared and batched across the prefix-tuned models (Shen et al., 2019; Narayanan et al., 2018). Consequently, the specific adaptation mechanism used informs system optimization (§4.3: adaptation).

It is an open question as to what programming interface should be used to specify that various adapted models are derived from the same pretrained model (e.g., models Y and Z are derived from the same pretrained model X), or that various components of two models share parameters (e.g., two models A and B share the same stem up to layer i). Ludwig (Molino et al., 2019) and PyTorch’s Module offer easy ways to compose functionality within a model, but no system today supports cross-model dependencies. Giving users the opportunity to provide annotations will allow training and inference systems to optimize and orchestrate computation more efficiently; without such annotations, systems will not have visibility into what computation and parameters can be shared across model instances. A model’s “adaptation history” (what models is this particular model adapted from) can also be used for debugging: an adapted model’s errors on particular types of inputs could originate from the pretrained model, pointing to issues in the pretraining process versus adaptation process. Frameworks like PyTorch, as well as software libraries for training foundation models such as HuggingFace Transformers (Wolf et al., 2020), do not allow for fine-grained lineage information across entire model instances to be specified.

Building and maintaining a cluster of thousands of accelerators also requires tremendous effort. New training paradigms like Learning@Home (Ryabinin and Gusev, 2020; Diskin et al., 2021) explore leveraging volunteer compute over the internet to train foundation models collaboratively. Such fundamentally new execution models can decrease the cost of training for any one entity, but require collaboration across a number of different areas like security (to ensure that a malicious volunteer cannot significantly alter the training process), distributed systems (to deal with fault tolerance issues as volunteers drop), and crowdsourcing.

Productionization of foundation models

As the community continues to push the capabilities of foundation models, realizing their potential will require addressing the challenges associated with deploying these resource-intensive models in production. These challenges include performing model inference with tight latency targets, and ensuring that models and data are monitored in an automated way.

For applications with strict cost and latency constraints, model compression techniques like distillation (Hinton et al., 2015; Li et al., 2020d; Sanh et al., 2019), quantization (Polino et al., 2018; Gholami et al., 2021; Zhou et al., 2018), pruning (LeCun et al., 1990; Gordon et al., 2020; McCarley et al., 2019; Wang et al., 2019c; Sajjad et al., 2020), and sparsity (Gale et al., 2020; Elsen et al., 2020) could aid deployment by transforming larger models to obtain desired inference-time properties. These techniques were originally intended for smaller models (e.g., BERT-L) in low-memory environments (e.g., mobile phones), but are now necessary to handle the extreme scale of modern foundation models in datacenter deployments. Parallelization techniques like tensor model parallelism (Shoeybi et al., 2019), traditionally used for training, might also be useful to reduce inference latency, and also provide additional memory capacity across GPUs to fit the model’s parameters.

In addition to these practical constraints, increases in the size and complexity of foundation models and the datasets used to train them pose new challenges to model and dataset lifecycle management. Since models with a large number of parameters are hard to manually inspect by humans, we need better systems for automated dataset curation (§4.6: data) and model quality assurance. Techniques like behavioral testing (Ribeiro et al., 2020) and model assertions (Kang et al., 2020) facilitate easier model maintenance in production by providing analogs to unit tests, runtime monitoring (in the form of test-time assertions), and continuous model improvement (as new inputs come in) for models deployed in end applications. These tools can help address issues of fairness and bias (§5.1: fairness), and reduce model mispredictions.