Lightning Talk: Training Embedding Model Resiliently for Multimodal M... Huamin Chen & Haichen Zhang

This talk covers the importance and functionalities of a semantic router in model inference for applications like chatbots and agentic workloads. The speaker discusses how a semantic router intelligently selects the most efficient machine learning models based on criteria such as cost, latency, and accuracy. Different routing methods are explored, including auto modes offered by model providers that optimize computational resources. The session highlights the use of embedding models for classification, emphasizing their role in maintaining efficiency and scalability in various tasks. Furthermore, the speaker elaborates on the training of multimodal and text embedding models, presenting the advantages these methods bring to the AI ecosystem, especially in regard to minimizing token usage. Finally, he introduces the AMD ROCm framework and its capabilities in optimizing PyTorch performance for model training and inference.

First of the four uh problem statements, why do we need a semantic router for our model inference? Uh for multiple reasons. Well, you are using the models for either chatbots or agentic workloads, there's a number of issues over there. Uh the cost, how much you pay for the token, and what's the latency you are going to experience, and most importantly, how accurate does the model answer your

questions or organic workloads. So, based on this uh three criteria, you may have a variety of model choices based on specific workloads. And that's come into the inception called semantic router. Right. Thank you for joining. You are not late. You just somehow I started early. Uh I'm trying to spice up this atmosphere to give up people some ideas what we are doing. But I guess it's almost

time. Okay. Um so, let us just uh have a quick rundown real quick. You know, many of the um frontier labs, including OpenAI, using certain type of routers to route their queries. So, if you do not pay attention, if you are just run for itself, and you do not specifically select any models, chances are the uh the model providers may provide the best cost-efficient models uh for

your workloads or for your prompt. Um if you are using, for example, um the Sorry. There's auto mode, right? So, the auto mode will specifically look at what you are asking, and then pick the best models um you know, for your answer, for your code generation. And if you are using the open router, for example, there's also an auto mode, meaning that's uh if you specify the

model, you just specify auto in the model field, and then from there it's going to the cost-efficient models that will you know most likely give you the best choice without using a lot of computations in background. All right. So, this is just basically what the semantic router is trying to do. So, we are using this concept of a mixed mixture of model in contrast with a mixture

of experts in a single model. So, in the back end, we can have a variety of models. You know, they could be text models, it could be multimodal models, it could be big models, it could be small models small models. And based on the prompt, we identify what is the best inference engine a um question the best way. And then we make the decisions at the routing

level routing level. Um so, what we do that is that we create the ecosystem under the VM projects. So, this is a part of the VM ecosystems. The idea is that we are using the you know, if you're looking at our organizations, there are two models. One is called a VM uh it's called one is called a router, the other one is called a semantic router. And

semantic router is the is the project that we are just talking about uh today. So, what's is different from the other projects is that we did not just look at the you know, the back end statistics. a queue length and load balancer is not our goal. Our goal is to look at the contents itself and make a intelligent decision which is the best models can most likely

satisfy your needs without using a lot of tokens. So, that's the basic idea. Uh we have a layered architecture as the we do the classification. So, there's two ways if you're working in the AI field, there's a classifier using like a birds or just a generative models to predict what is going to use or you can use the embedded models, you know, just like the way you

are doing the rack. Embedding models give you the clustering opportunities based on which you pick up the cluster and then do the routing on that front. What we are using we started from class classifiers. We have the birds you know classifiers fine tune birds classifiers but increasingly we are using embedding models for efficiency as well as the scalability. In this embedding models all you need to do

is just to augment the knowledge base. You do not need to have additional training and this knowledge base we are just tell the clustering how they are formed and which centroid they can do tilt so they can make a decisions on that clustering might ID. Basically cost efficient training free mechanisms to do the classification. Um so the talk today is about how we build the classifiers cost

efficiently and augmented you know uh extensively and use on the energy systems. Next. Um so we have both text embeddings as well as multimodal embeddings. For text embeddings the challenging we are using because remember this is router. Router has special requirements. You have to be efficient and you have to be accurate. Efficient means that latency in the classification is models has to be as fast as possible

in milliseconds or tens of milliseconds range and accuracy means that you are have to support multilingual English, French and so on. And we have to be very you know generate generous enough to support both small, short and long prompts right so that's the that's we are doing. We searched all these existing embedding models and we could not find any of these models can do it. You know

for efficiency at least you are needs to support flash attention and the only flash attention in binding models is modern birds. But modern birds does not support multilingual. So there's a variant called a multi modal And then we do not find it support long context. So we have to augment it, you know, using fine tuning to extend it to 32K. So we are able to classify the

whole contents, not just the 8K that you would typically see. So we are not excluding any of the opportunities to use that inference service we use. So that's what we do for modern birds extensions. remember earlier I said we are not using classification anymore. We are using the invited models. Modern birds is still a classifier. It's a birds model. So we have to fine tune it into

a modern into a invited And so the way we are fine tuning it and it's based on the BGM, which is a open source invited models. We using the same training mechanisms to achieve the accuracy as as good as, you know, some of the frontier invited models. Right? We also to the research called a 2D MC, meaning that if you're using invited models, you have two ways

to scale. One is the dimension scale down and up. You can scale down the dimensions in invited models. It started with a 784, but you can scale down all the way to 64 dimensions. The other one is the number of layers. With 32K 32 layers, but we can always also scale down to fewer layers for efficiency. Right? So that's is the text invited models. The multi modal

invited models is along those ideas. So we are using some of the existing greatest encoders, different model encoders to build up two variants. One is for data center use, 1.1 billion. It's not done yet. It's still work in progress. The other one is the mobile violence, it's only 200 million 200 parameters. We are using this model training this imbalance for multi-model classifications and they work wonderfully. Um

with the details how we train the models I will talk from AMD how to using the AMD PyTorch ecosystems for model training. Uh thank you for me and thank you for joining our session. Uh as mentioned earlier, all of this embedding model uh is uh practiced on AMD GPU. And uh next I'm going to talk about how to run PyTorch well on AMD data center GPU series.

Uh this is my agenda today. Uh as a very beginning I'm going to introduce how to and then I'm going to introduce some new functionality and uh features on AMD ROCm and final part I'm going to talk about deep dive about AMD training fundamentals. Okay, uh without further ado, let's get started. Uh the ROCm open software world. As you know, ROCm is something similar as CUDA and

once you have AMD GPU, you have to install ROCm to support the training and inference as well. And uh we can provide some different layers. This is our key. Building on the drivers and runtimes is our compiler and the tools and also we can provide some providers for the user and sitting on top of that is our high performance kernel libraries like the ROCm plus and CK

kernels like this. And AMD has sponsored the PyTorch community since the very beginning. And this is our installation ways. You can use the Docker image or PP wild installation. And many of you may be questioned is as CUDA PyTorch. Actually, they are actually same with zero code changes And about the new functionalities and performance update, I'm going to talk about the first one is low precision and

format enablement. We can support the MXM data type on the latest advanced AMD GPU. And we can provide lots of my poor compiler and backends for the user AOT inductor with CK backend. And we can also support some gym operators like the STPA means. And we can also support some things like the memory and distributed optimizations like the symmetric memory enablement mentioned here. And this is our

Torch Inductor performance dashboard. We use this to monitor the Torch Inductor performance for user to we can to see how this Torch Inductor features run well on the GPU. And this is our advanced tuning strategies by using the tunable ops. We can use this and auto tuning to boost the current to accelerate the training as well. And this is the performance in terms of the preview latency

and decoding throughput. Here we use the Llama 3 models to give you example. About the Torch Titan, we have released a specific Docker to support the Torch Titan. You can just Docker pull and run And we all can also provide the Torch compile for the user. For example, you know, Torch compile support multiple modes like reduce overhead and max auto And I will give you example of

how the Torch compile boost the model performance when running on AMD GPU. As you can see here, we compared all the other modes with baseline, the eager mode. And this is our unit test coverage. Because I think actually practice PyTorch in your product or workload process, you want to make sure the quality is good and we actually we have sponsored lots of cluster for the PyTorch team

to run the CI/CD test and we will update it in in this test coverage And then you can check this online. Okay, that's all for my talk today. Thank you for joining our session. Thank you.