This presentation introduces "VM," a program or framework (distinct from "UMD" and "LMD" for large-scale disaggregated inference, though related) developed by Google in collaboration with the Berkeley group and industry partners. The primary goal of VM is to significantly simplify the process of serving machine learning models on Google's Tensor Processing Units (TPUs). The speaker highlights that VM already supports TPUs, leveraging JAX kernels within its PyTorch code paths, with future plans for deeper JAX integration. Addressing the prevalent challenge users face in deploying open-source models on TPUs, the project aims to make VM extensible to support TPUs, JAX, and custom applications, all while maintaining an open-source commitment. Driven by the increasing dominance of inference in modern ML workflows, Google's strategy is to embrace the ecosystem, providing seamless deployment for PyTorch models on TPUs. The long-term vision encompasses "Day 0" immediate PyTorch model support with solid performance, evolving to "Day X" where popular models are further optimized with JAX for superior performance on Google's hardware. Future plans include enhancing VM with JAX support for a new "Ironwood" AI chip, expanding model coverage, and offering advanced performance tuning capabilities.
VM on TPU: Simplifying ML Model Serving
0:05 Introduction to VM: Fe, a Google software engineer, introduces "VM" (a model serving framework, not a traditional virtual machine) and clarifies its relation to "UMD" and "LMD" (large-scale disaggregated inference, based on VM).
0:40 Existing JAX & TPU Integration: VM already supports Google's TPUs since early 2024, leveraging JAX kernels to boost performance in its PyTorch code paths. There are plans for even more direct JAX integration in the future.
1:18 Core Problem & Goal: The project aims to make it "as easy as possible" for users to serve their machine learning models on TPUs, addressing the common challenge of deploying open-source models.
1:42 Collaborative Development: VM is a collaborative effort involving various Google teams, the Berkeley group (original developers of VM), and external partners such as Red Hat, IBM, and numerous universities.
1:59 Extensibility and Open Source: A key objective is to make VM extensible to support TPUs, JAX, and custom features, with an emphasis on open-source development, which is crucial for TPU compatibility.
2:33 Growth in Inference: The increasing sophistication of ML models means that most use cases, including aspects of training and agentic systems, are driven by inference, indicating a significant growth opportunity.
3:02 Google's Ecosystem Approach: Google aims to embrace the ML ecosystem by providing effortless ways for PyTorch model users to serve their models on TPUs using Google's stack, moving beyond simple reference implementations.
3:33 Day Zero to Day X Vision: The vision includes "Day 0" support for PyTorch models (from Hugging Face, Model Garden) with good performance, evolving to "Day X" where popular models are further optimized with JAX versions for superior performance on Google chips.
4:11 Future Strategic Focus: The team foresees an "age of inference" where usability and efficiency will be critical for the success of the ML ecosystem.
4:26 Upcoming Developments & Call to Action: Future plans include launching a new AI chip ("Ironwood"), adding comprehensive JAX support to VM, increasing model coverage, offering performance tuning, and committing to open-source all developments. The speaker invites interested parties to collaborate.