AI/ML frameworks for cloud TPUs

Key Takeaways

[music] If you're already working in AI [music] or learning about it, you've probably dealt with a few frameworks. PyTorch, Caris, Jax, Triton, VLLM. But what are they all for? How are they different? And when should you use one over the other? We'll break them down and look at some of the trade-offs, give them a definition, and talk about which tests they're good at. In other words, we'll do a little if this then that comparison. First, let's cover model training frameworks. We all like shortcuts, right? Think of a framework as a toolkit. A shortcut from zero to production code. You could build everything from scratch if you're some kind of super genius. Or you could lean on these tools to pre-build part of the puzzle for you. Whether that's doing complex math, data handling, or optimization. For many, the perfect starting point is Caris. Caris is an easy to use interface or API for building models. It's like a clean dashboard that sits on top of a powerful engine. You can build a powerful neural network in just a few lines of caris code. And a huge advantage today is that caris is multibackend. That means you can write your caris code once and run it using jacks or pietorrch as the underlying execution engine. That kind of flexibility is incredible for defining your models. Underneath caris or for more granular control, we have pietorch and jax. PyTorch, an open-source framework originally created by Meta, is loved by researchers for its Pythonic feel and flexibility. Jax from Google is a high performance numerical computing library, excellent for research and large-scale model development due to its speed and automatic differentiation capabilities. These are the engines that power your model definitions. It's important to remember that while Caris, PyTorch, and Jax define models, they're also fully capable of training models directly. Researchers often leverage these frameworks for their full access to parameters, allowing for highly customized and experimental training workflows, allowing them full control over the process. But now that we've trained the model, what's next? Once you've defined and trained your model, you need to put it to work for end users. making predictions or generating outputs based upon new unseen data. While training is about learning, inference is about making it quick and easy for them to apply that knowledge, often in real time and at scale, allowing for them to serve predictions and generate text for whoever is using the model. To do that, you need dedicated inference frameworks and servers. These tools allow your models to handle anywhere from 10 to 10 million requests quickly and reliably in a production environment. Let's look at some key players on the inference stage. For GPUs, Nvidia's Triton inference server is a highly performant open-source inference server. It supports models from various training frameworks like TensorFlow, PyTorch and ON&NX runtime offering features like dynamic batching and concurrent model execution for maximum efficiency. TGI from hugging face is another popular solution specifically designed for deploying LLM for fast inference providing features like continuous batching and quantitization. It works great if you're solely working with textonly models deployed through hugging faces model garden. For more than just texton models, consider vlm. VLM is a wellknown as a fast and affordable library for LLM serving using techniques like page detention to manage memory efficiently during text generation. And when it comes to scaling up your inferencing workload to handle more incoming traffic, perhaps you might want to consider LLMD. LLMD is a powerful open-source framework designed for high throughput LLM serving, offering advanced features to optimize performance and cost for large-scale deployments. While many developers focus on inference, understanding fine-tuning is fundamental. This involves taking a pre-trained model, one that has already learned from a general data set, and further training it on a smaller specific data set to adapt it for a particular task or domain. It's way more efficient than training a model from scratch. For efficient fine-tuning of large models, especially LLM, a powerful approach is parameter efficient fine-tuning or PFT. Libraries like hugging faces PFT provide methods like Laura. Low rank adaptation that allow for you to adapt massive models to a new task with less compute and less data. But keep in mind that HuggingFace only offers tools for PyTorch. So to quickly recap, when you're looking to define your models, caris stands out for its user-friendly API, offering a highle abstraction that can run on powerful engines like PyTorch and Jax. For more granular control or cutting edge research, PyTorch provides a Pythonic and flexible environment, while Jax excels in high performance numerical computing and automatic differentiation, making them ideal choices for the underlying model definition and training. For inference, which is your primary interaction, tools like Triton Inference Server offer multi-framework support and GPU acceleration for maximum efficiency. TGI is a strong contender specifically for fast LLM deployment with features like continuous batching. And VLM provides an affordable and fast solution for LM serving through techniques like page attention. And for training and fine-tuning, you'll leverage the full capabilities of TensorFlow, PyTorch, and Jax. For efficient adaptation of large models, particularly LLM's hugging faces PFT provides methods like Laura, allowing you to fine-tune with significantly less compute and data, though it currently focuses on PyTorch. Although, don't forget there's no substitute for writing the code and seeing [music] these powerful frameworks in action for yourself. To learn more, check out the link in the description below. Heat. Heat. >> [music]