Lightning Talk: Accelerating On-Device ML Inference With ExecuTorch and Arm SME2 - Jason Zhu, Arm

This talk presents the collaborative work of Tyler Mullen and Jason Zoo from ARM, focusing on enhancing edge inference using Exeutorch and ARM's new scalable matrix extension technology, SME2. The discussion highlights the importance of efficient compute platforms for developers and how ARM is leveraging SME2 for on-device inference. Jason emphasizes advancements in computer vision, particularly through the squeeze sam model, which optimizes performance for mobile devices while maintaining a compact architecture. The session includes a detailed analysis of end-to-end latency reductions and highlights the significance of identifying bottlenecks at an operator level to enhance performance further. Additionally, the speakers discuss the collaborative efforts in optimizing model architectures to ensure smoother execution paths.

Good afternoon everyone. Thank you for joining us today. My name is Tyler Mullen and I'll be uh speaking here with my colleague Jason Zoo on some of the work that we've been doing at ARM to kind of marry Exeutorch and our technology 2 to enable uh inference on the edge with uh devices. So I think Mgan did a great job of explaining why we care about on

edge inference and I won't go through a lot of that again as well as some of the pieces that come into play when you're trying to enable that but I do think it's worth maybe a quick primer on ARM and the edge itself and 2. So at ARM we're very focused on making sure that we have the most power efficient and most high performance compute platforms that

we can provide to developers. And one of the ways that we're doing that recently is with our scalable matrix extension 2 or SME2. And what this is is this is a math acceleration piece that is within the CPU itself. It allows developers to take advantage of this acceleration without leaving the CPU. And one of the ways that you can do that now is through chips like the

ones shown on the screen here that have been available for roughly six months in the ecosystem. And as we start to marry Executive TTorch with this technology, SME2, we've learned a number of things along the way that we'd like to share with you today. So I will turn it over to Jason who's going to walk through some of the details of those. >> Thank you, Tyler. Um,

so I think Morgan did a good job on the voice user case on device. Uh, actually I'm going to talk a slightly another one is the vision. Uh, we know we use some of smartphone for your image, your video processing, capturing. So actually it's already running over there right if you're looking for the computer vision um field is actually the pace has been dramatically uh evolved just

think about the history one decade ago that is image net alex net is pretty much actually triggered the modern computer vision um research on this field but that days the task is really simple image classification right simple task then until 2023 I think meta released the model called segment anything model that's really bring the a new type of interactive experience particularly for segmentation. So for instance in

here you can see the middle panel that's actually the Instagram uh cutout feature if you play use that um and then since then every actually there's like a new version of the same model come up um given this actually let's say this v vision workload is more interactive one really let's say critical in requirement is come up is actually have have to have a low latency is

actually have much stronger low latency compared even maybe before the image classification problem because this is a real time user engagement experience. So let's take this Instagram cutout feature um as a focus case study in here. So the model behind that feature called squeeze sam and squeam I want to highlight to point is is actually the original or the general sam model architecture is actually still too

big for ondevice deployment. So it's only suitable for cloud-based deployment. So squeam is actually uh I like this uh model is is really let's optimize for um for the uh mobile and is still maintain a hybrid convolutional and the transform architecture which means architecture is really not like outdated is actually has a lot of representative learnings you also apply to online for instance. Second is even this

model the squeeze sim the researcher did a very great job to reduce the model size from I think the SAM has like 600 million parameters to down to like a 10 million parameters so it's much smaller however the computer workload is actually still very high estimate about 26 billion matrix multiply communication cycles to handle one image 512 x 512 so this case which means like the the

optim optimized kernel the optimized run runtime framework like XTOch is still very critical to really make this model to be practical on device. So before we dive into the result for the squid 7 I just want to quickly refresh our memory about the full path from the model definition to reach the hardware execution. Right at high level is actually have to go five layers. You have the

PyTorch to create your model. Then you have XTOch. You have XMP pack at a CPU back end. Then you reach Clyde AI which is a kernel library for ARM. And then you reach the hardware the instruction sets over there. For instance, SME2. So for model developers, you don't have to write any model code specific for SME2. The runtime stack will handle for you. However, one thing I

mention to here I highlight in here is when you run your model the inference of course you get the inference number but don't stop in here because the the operator for your model each operator go through these five layers things can happens in any of the stage. So it's critical to ensure to check your operation actually has been on the fastest path as you expected. So let

me show you some concrete example. So this is the main result for the squeeze sam model. Uh I show you both the end to end latency and also the breakdown view of where the time spent at operate level. So let's take FP16 data format for this model example. So you can see for the end to end latency has reduced from over one second to 300 milliseconds. This

is based on a single thread um CPU setting. So it's equal to more like 3.9 speed up quite big. Then if we look at where actually the time get reduced if you will immediately see this blue box which is actually corresponded to the convolutional operations is actually reduced drastically from about 900 milliseconds to 100 milliseconds. And then you also if you cross FP16 to int 8 you

will see actually withme on you actually the fp16 can reach the similar latency with in 8 that's another actually more complex problems or topic about quantization because quantization has overhead about quantize and recquantise during the computer side. So another thing I actually want to highlight in here is right now if you look at those SME on columns the blue box is not necessary the bottleneck anymore that

was the matrix compute is actually the bottleneck right now is on the orange that's is a data movement and if you look deep that count like 40% of the the runtime right now and almost like 100 milliseconds spent on data transports. This data transport has actually no value to the model. It does not impact output. It's really just overhead. So by look at this operate operator level

timing breaking down view and then we want to trace back okay what really caused this transport in this case. So we did root cause analysis together actually with um meta xorch team uh including Morgan and a few other colleagues over there. So we found actually trace back to the model graph. We see actually this particular squeeze model has a repeated block that contain a common 2D then

activation function then normalization. So this particular order of the stack is actually give us a two problems. Firstly the the activation in the middle actually will break the possible to con to fuse the convert and the batch romization and the secondly is now right now actually batch roam operator will be sitting in PyTorch XTOch and had not be delegated to XMP and the two framework have a

different layout requirements for the 4D tensor and particularly XMP require channel last NHWC and the Python torch usually um support or use uh NCHW. So this the layout difference between the two framework actually cause continuously transports or repeated transports. So a quick fix uh actually implemented also proposed by the um mad xorch team is to delegate the batch normalization. So by simply doing that then basically you

you make all the operations within XMP pack they don't have this transport immediately we able to reduce 70% of the transports overhead of course for long term also I think it's good to really study also working with the community for model developers when you develop your model design architecture would be good actually also have awareness of what actually type of architecture is execution friendly right I did

it the the more execution friendly will be convolution first follow up batch normalization and then activation function. In this way you really ensure the continuity of the layout and that you can eliminate the overhead the source level. So this actually the pattern ofme accelerate the matrix compute and then a new bottleneck could be rised if you're looking at operator level is actually can be cross the model

families I have been show you the squeeze sam but also I also show you language model and the the result in here is actually from Google's gem 4 model which actually launched last week we were also collaborate over there to do model optimization so You can see that the big swins on matrix compute accelerate about 4 to 9x and of course there are some operations will be

remain and those are what could be could be the new areas actually we can optimize further through um the framework level. So the key takeaway for this is really uh two. One is CPU withme is actually could be a good way to actually accelerate the matrix compute and then make some uh ondevice model become possible which probably not able to before and the second is when you

have the end latency don't stop there look deeper at operate level and at operate level you might discover something actually you can optimize further so I just want to thank um actually quite a colleague for ARM but also especially for the meta xtorch team to work together on this from the real world user case and also the implementation for optimization. We do have a pytorch blog post

on this um if you want to find more details and also we have a code repo to show actually how to use x torch to do operate level um um profiling and also identify bottlenecks for optimization. Thank you.