Deploying PyTorch Models To the Browser and Beyond With Transformers.js - Joshua Lochner

In this talk, Joshua, a machine learning engineer at Hugging Face, discusses his journey and the creation of Transformers.js, a JavaScript library that enables the running of state-of-the-art pre-trained models directly in web browsers without requiring a server. He highlights the project's growth, achieving 4.4 million monthly downloads, and makes a case for the advantages of in-browser AI, including increased security, lower costs, and scalability. Joshua outlines the process of exporting PyTorch models to ONNX and optimizing them for browser-based inference. He introduces WebGPU, which enhances performance for GPU computing in the browser, along with other emerging technologies. Throughout the talk, he presents numerous demos showcasing the capabilities of Transformers.js, emphasizing its user-friendly API and the exciting potential of web AI applications.

So, let's turn the clock back 3 years to 2023. Picture this. You've just graduated from university, your GPU poor, and watching YouTube in your bedroom. You've just seen a super interesting video, and you're scrolling through the comments to see what others have to say. Unfortunately, it's full of bots and scammers. Some quite easy to identify, but others are sophisticated, spanning multiple comments, manufacturing social proof in order

to deceive. You have an idea. You've done some machine learning projects in the past. Nothing big, but for some reason, you have enough confidence or naivity to do something that the thousands of YouTube engineers and goo and Google scientists couldn't. You're going to fix YouTube's spam problem. You're also in luck because there are some great tutorials for text classification uh fine-tuning on a platform called Hugging Face.

Weird name, but you go along with it. After a bit of work, you fine-tune your own text classification model and it works. You're able to correctly classify comments. Time to deploy. It should be able to interact with web pages. You want others to be able to use it. Uh infrastructure should be scalable. Uh it should be affordable to run. So you think a browser extension seems like

your best bet. But remember, you're GPU poor. So you can't let the millions of users uh that would make inference calls. So seems like we're stuck and doomed to fail. Well, that's unfortunate. Uh looks like the project will fade into obscurity just like your many others. Unless you could run the model locally in the browser. Now yes, that was me in the story. And although I never

did end up finishing that side project, I think the journey that it would ultimately lead me on has been pretty crazy to say the least. and is the reason I'm here today presenting at the PyTorch conference. So I guess things worked out for the better. Let's dive in. So hi, my name is Joshua and in 2023 I joined HuggingFace as a machine learning engineer to continue a

little project I was working on at the time. Transformers.js. Transformers.js is a JavaScript library that allows you to run state-of-the-art pre-trained models directly in the browser. No server required. It's designed to be functionally equivalent to the PyTorch uh to the Python transformers library, meaning you can run the same pre-trained models using a very similar API and in the most performant way possible. Since the project began in

2023, TransformersJS has seen some amazing growth. Just last month, we hit a record of 4.4 million monthly downloads on npm. We also doubled the size of the core team from well one to two. Uh shout out Niko. Anyway, my goal today is to introduce you to web AI and make a compelling argument for why you would want to run models in the browser. We'll then explore how

to export PyTorch models to Onyx and to demonstrate how to optimize these exported models. Next, we'll show how to deploy and run them using Transformers.js s to create performant, interactive, and visually stunning web AI experiences, leveraging emerging web technologies like web GPU for crossplatform inference. In fact, every demo I'm going to show today is something you can try out yourself. No installation required. Simply visit a website

and run the AI models locally. That's the power of web AI. And speaking of which, let's talk let's talk about machine learning on the web. So what are the benefits of in browser inference? So the first one of course uh security and privacy. If you're handling sensitive documents, maybe considering you know uploading images of your faces of of your face, maybe something that's not you're you're not

interested in that. So we'd rather run things locally. Everything stays on device. It also allows you to build real-time applications because there's no server and a client to server latency. Um you're also not sending large files. So imagine you are doing a video stream. Uh you're not being you're not sending every frame over the wire and coming back with a request with a response to your request.

Uh that could introduce quite a lot of latency. Next, lower costs both for the user as well as the developer. So the user you just run on your device. So everything whatever the electricity costs of your device are, that's what you're paying. Um and from the for the developer uh those who want to distribute their applications or models without having to pay for servers or inference API

uh costs everything is distributed on client uh devices. So you're not paying for inference. And lastly, it's highly scalable. So because you're distributing on the web, uh I guess visiting a website is seemingly the easiest way to show your uh to show your your your projects or uh for users to interact with it. Remember the model gets loaded into the browser, downloaded into the browser and then

cached by the browser. So if the user re, you know, goes back to the website, it'll all be ready and cached. Um and that means that any of your users can just try out your model. uh and it's pretty scalable. I guess that's the power of the web. Uh we also need to consider a few things when optimizing for in browser inference. First uh quantization uh which

is a technique to reduce the computational uh and memory costs of running inference when by representing the weights using lower precision uh like 8 bit integers or uh 32-bit uh instead of the original uh 32-bit or 16- bit floats. Uh next we should be taking advantage of the various browser APIs we have at our disposal like webgu and webn which both of which I'll talk about uh

later and we should be doing our best to optimize our model exports as much as possible. So how does it work? First of all we're converting our model to onyx using uh you can either use one of the many pre-trained uh deployed ones on the hing phase hub. We have around 3,000 transformers.js compatible uh checkpoints or you can convert your own uh in ways I'll show soon.

Um and then you write your JavaScript code often just a few lines of code. Uh we try and make this as simple as possible using the pipeline API uh which I'll also show a bit more a bit later. Behind the scenes the model is executed via Onyx runtime using the selected u execution provider to run on either your CPU, GPU or NPU. So let's take a look

at how you would export PyTorch models to Onyx. Onyx or open neural network exchange is a standard which defines a common set of operators and a common file format for representing machine learning models. Uh PyTorch has fantastic support for making model exports seamless and ina enabling interoperability and compatibility with a wide uh variety of frameworks and deployment runtimes. It aims to solve the issue of framework uh

dependency by providing a middle ground for model portability making it easier to move models across uh tools and and hardware. So one of the tools I like to use to uh showcase I guess onyx models is Netron. Um it's basically really simple. You drag your model your Onyx model in and you're able to see the graph as well as the input output signatures. So here's an example

of small LM, one of the models we uh work on at Hugging Face. Um and you can clearly see the inputs, the outputs, um as well as the graph format, which makes debugging a lot easier. Um yeah, pretty cool. Uh and so a typical workflow usually involves training, export, and inference. First, we build and train the model with PyTorch. Uh then we export to onyx uh which

is a standardized framework agnostic format uh by mapping operations uh and saving weights for inference and finally we execute the onyx model uh with kernels optimized for the target hardware uh CPU GPU and so on using these so-called execution providers. Uh so Dynamo is PyTorch's modern approach to exporting replacing the legacy torch script which used tracing that often had issues with data dependent or conditional behavior. At

a high level, uh, Dynamo is PI is P is a Python level just in time compiler designed to, uh, make unmodified PyTorch programs faster and for our purposes makes onyx conversion uh, extremely easy. So the pipeline creates your model uh, as torch.onx.exported exported program normalizes it into functional operators, records uh shape constraints and then translates that graph into Onyx using Onyx script uh based translation logic uh

producing a torch. Onyx program that you can further optimize, serialize and validate. Uh finally, this uh approach makes it really easy to integrate custom or fused ops into your exported graph and also enables the use of more recent uh onyx operator sets. So let's start with the hello world of PyTorch to Onyx conversion. The way you write your modeling code is quite important. As a general rule

of thumb, keeping your keeping everything in PyTorch land uh and invo and avoid using Python specific features like type casts or for loops, while loops, those kinds of things. Uh if you find yourself needing to write uh and resort to these features, then I'd recommend you rethinking your approach, trying to vectorize as much as possible. Uh here's the exported model uh displayed in Netron. Uh as you

can see, pretty simple. Now let's consider a bit more of a complex example. Uh using a model that's defined in hugging face transformers. Uh let's take a B embedding model. Uh we can prepare inputs using the corresponding tokenizer and then run torch.x.export with the following parameters. Now this kind of works but has one major flaw. The shapes are static. So when we try run the model again

but with a different sequence length or batch size, we would get an error saying that it expected the same shape it was exported with. Luckily, it's a pretty easy fix. The easiest way is with dimension hints like dim. Dim do.static or dim.dynamic, but I prefer to be explicit in defining the names and shapes directly. It also shows up quite nicely in the in Netron. uh it shows

like the the the dimension names as well as the input names are uh much better. So now models often follow a similar input output signature. For example, image classification models have pixel values as the input and then the scores attached to every label um as the output. So at HuggingFace we maintain a separate library called Optimum which defines these mappings and export configurations for you. uh meaning

you can export hugging face transformers uh using a single command from your CLI. Now that we know the basics of onyx conversion, let's see how we can optimize it further. Dynamo also uh supports registering custom operators enabling runtime and hardware specific optimizations. So let's take a look at how this is done in Onyx runtime. For example, uh the language used there is uh contrabop, which means it's

basically built into the runtime, but it's a non-standard uh operator designed to support, let's say, new experimental or highly optimized uh operations that are not yet part of the official Onyx spec. Uh they allow superior performance using uh through operator fusion uh while remaining directly available in the core uh Onyx runtime codebase. uh Onyx runtime contrabops are defined under the com.microsoft namespace. So they're pretty easy to

identify when you're looking through Netron or inspecting the graph yourself. So let's consider a typical attention module with the implementation and Onyx export code uh shown here. You can register a custom OP in three simple steps. So step one, register the op. So Dynamo uh traces this STPA node as a single opaque node instead of decomposing it into elementary operations like you know the matmo softmax etc.

Uh step two we subclass or patch the model to call the custom op instead and then step three is to build the translation table which tells the Dynamo exporter how to convert our custom op into onyx nodes. The onyx script library makes this really easy to do as well. Then when you want to export, it's pretty much identical to what we saw before. Uh but you just

have to specify the custom translation table. And when viewed in Netron, the differences between exports uh becomes immediately obvious. The default approach decomposes the basic implementation into many separate operations while the custom version collapses it all down into a single multi head attention node which is part of the com.microsoft namespace. And when benchmarked using the web GPU execution provider, we can really see how much of a

difference this makes. Uh in this example, we're seeing between three to 4x speed up uh across different uh sequence lengths and uh head and number of heads. Um now moving on to a real world example. Let's see how this optimization can take place in the BERT embedding model um that we saw a little bit earlier. On the left, I've colorcoded the different parts of the attention calculation.

And on the right, pretty much all of that is is absorbed into the single multiattension node, making our graph significantly more um uh simple. Now, when benchmarking this model, we can see clear performance improvements, especially at large batch sizes. Onyx runtime uh which uh also provides a really simple API for contisation uh which can be used to reduce the size of the models and of course that

is extremely important in browser contexts because you don't want your users downloading massive amounts of gigabytes of uh model weights. This produces a new onyx model that replaces all matmo ops with matmo nbit ops and gather nodes with gather block quantiz nodes. Um reducing the size uh quite significantly depending on the number of bits you choose obviously that that decides how much smaller the model will be.

Um I also wanted to mention two other libraries that I use on a daily basis. The first being Onyx Slim, which is a toolkit for optimizing Onyx models by rewriting the graph in a more efficient way to reduce node counts by dduplicating weights, uh let's say rewiring the graph in certain places and performing other uh graph uh transformations. So in this real world example of a couple

months ago when we exported a zeroot voice cloning model called chatterbox uh we needed to unroll many uh loops leading to an exceptionally large graph uh thousands and thousands of nodes and Onyx slim was able to rewrite the graph using significantly fewer nodes um in this example for replacing over 23,500 constant nodes replacing with them with initializers um and many many more uh optimizations uh I believe

it cut down the model size by a good 40 50%. Which is pretty crazy. Um, and although this is less needed uh now with uh torch dynamo, it's always a good idea to test if your graphs can be optimized further. Now, Onyx runtime genai is also a pretty fantastic library. Uh it focuses specifically on LM architectures. uh and their model builder API uh allows you to have

finer control over the exports by essentially refining and um or redefining the architectures in onyx operations directly. So now let's talk about transformers.js. Uh we hope that we try um try to make our API as similar as possible to the Python transformers library. of course taking into into account the weird quirks of JavaScript like many asynchronous operations. So it's pretty much three things. You import the library,

you create the pipeline depending on the task that you're trying to achieve and then you run the model on your input. So in this in this case we're doing sentiment analysis and we get an an output score a high output score with a label positive. Um you can also specify custom models. So in this case we want to perform the background removal task. uh and we specify

the ben 2 model uh to to use. Um similarly we pass an input image and then we run the background removal removes the background. Uh we also have the pipeline API that has many loading and runtime uh parameters. So for example when you're loading the model you can specify the device and the quantization as well as when you are running the model it'll be like let's say

max number of tokens sampling parameters that kind of stuff. We also support um advanced usage of model and processor classes. So for example here we're running SAM in the browser and the if you want to for example let's say premputee embeddings so you can do multiple decodes later uh you can just use it via the raw uh model and processor classes. And we also have some pretty

exciting news. Uh in the past I think it was 10 days ago um we released the biggest update of transformers.js yet uh transformers.js JS version 4. It took over a year of development since we started in around March of 2025, and we're finally in a stable enough position to release it to the public. The biggest change was undoubtedly the adoption of a new WebGPU runtime completely rewritten

in C++. Uh we've worked closely with Onyx runtime team to thoroughly test this across over 200 supported architectures uh in And in addition to better operator support for better performance, accuracy, and coverage, the new runtime also allows the same Transformers.js code to be run in a wide variety of JavaScript environments, including browsers, server, runtimes, and desktop applications. So you can now run web GPU accelerated models directly

in node, bun, or dino. But what is webgu? Well, WebGPU is a modern web graphics and computation API that allows high performance GPU uh access across uh your web browsers and other JavaScript environments. It's the successor to WebGL and provides a much more advanced much more advanced capabilities for both graphics processing as well as general purpose GPU computing. Although to be honest, the name web GPU is

a bit of a misnomer because it makes you think the technology is locked into the browser, but this couldn't be further from the truth. uh implementations like Dawn, which is the one that Chrome uses behind the scenes, um uh allows you to use the same technology in a truly cross crossplatform way. So, it can even be run from languages like Python. Um and in TransformersJS, it's as

simple as specifying a single line of code to change the device. Another new uh web API you should be aware of is webn which is which allows you to enable um hardware accelerated neural network inference directly in the browser using a highle interface. Um to use it in transformersjs all you need to do once again is to set the device to webn npu webn cp or webn

gpu. Uh things are still quite experimental especially in the browser APIs themselves. Um, so you'd need to enable it via a feature flag. But I think this is quite an interesting thing to um to be aware of. And now let's consider and show some things that are possible. Uh first of all, starting with uh this uh 1.25 uh billion parameter model from liquid AI. Uh we're able

to achieve quite significant token uh tokens per second. So around uh I think it was around a 200 tokens per second on my M4 Max. um everything running locally in the browser of course. Um and yeah, pretty good. I guess this model also does thinking mode. So you can uh achieve some pretty high level tasks even using such a small model. Um it's even possible to run

significantly larger models. So here's actually GPTOSS 20B. Now you might ask yourself why you would want to run a 20 billion parameter model in the browser, but you know why not? Um so in this case we use a custom op um with 4bit quantized weights and FP16 activations. Um the file size is around 10 GB. Uh so getting that loaded into your browser and cached is a

bit of a mission than the first time but um you can see pretty good 40 tokens a second. Not bad. Um but of course this is not limited to text modalities. So here's a vision language model which uh trans which describes what it's seeing on um on as a webcam or video input. Um, the model is based on, I guess, their their LFM2 architecture. So, it's actually

extremely quick. Um, what I've actually had to do here is slow down the trans slow down the captioning by around 200 milliseconds because it was kind of too fast and I couldn't keep I couldn't read what it was saying. Uh, it's still pretty fast. I don't know if anyone can read that fast, but uh pretty good to know that this is possible in the browser. Uh we

also uh collaborated recently with Google to uh support days zero uh to have day zero support of uh their new gem of 4 model uh in transformers.js. So of course it was I think it was the day after 2 launch. So I tested out there. It supports um audio input. So I I asked it a question first and then you can also take images, screenshots of the

video and then uh it is able to see what tell you what it's seeing. Um, next up we also have uh Cohes uh transcribe uh automatic speech recognition model which runs blazingly fast uh as well. Um, we also actually recently added support for voxal real time enabling real time streamed speech recognition with sub 500 millisecond uh latencies. So there's just me a demo of me speaking. Um,

and of course we also want to support much smaller models. It's the ones we actually try to promote and show that you can actually really do amazing things with much smaller models. So in this case it was like a a 20 million parameter image classification model uh to turn into a game. So if you're familiar with Google's uh quick draw game, it's basically you draw on the

canvas and you hope the neural network can predict what you're doing. In this case, the predictions are, you know, running extremely quickly as every time you draw as well as it can even run on your phone. Um, and taking this a bit of a step further, we actually built a game using Google's recently released function gemma. So in this case, you describe where you want to place

the object on uh on the grid, let's say using natural language. And in this case, yeah, you you go to a target and um keep going across the levels. Um, another fun way of visualizing data is using uh embeddings, of course. But I guess what this shows is you're combining the power of machine learning with the expressiveness and interactivity of web browser uh demonstrations. So in this

case using 3JS as um a way to visualize documents in a 3D space of course clustered together based by distance and you can do semantic uh search based on what you're looking at. So I searched for weather and all the the weather points uh showed up and everything's scored I guess. Um next up we also support various image uh and computer vision tasks. So in this case

uh image segmentation uh or object detection. So you can uh and we recently collaborated with uh GeoAI uh startup they um allow you to select boxes on like you know Google maps or anything and to annotate what you're what you're looking at. And then of course you can do other things like depth estimation. This really adds another level of immersion to your web applications. Um, someone actually

worked on like a Wikipedia article uh renderer that kind of displayed all the images in a bit more 3D depth, which is I thought pretty cool. Um, another one way of using uh transformers.js is for educational tooling. So in this case, we're learning how vision transformers work by inspecting uh the attention heat maps across all the layers and attention Um, this is a way to show how

you can interact with web pages. So, in this case, we're doing like summarization. You can imagine this interacting as a a browser extension or something. So, let's say you're asking to summarize this this paragraph. It'll summarize it for you. And you can even interact with the web page in like a chat format, of course. So, you know, who are the crew members? And then it'll pop up

with with the answer. Uh, you can also interact with various browser APIs. So in this case we're doing uh tool calling to ask um and obviously I request permission of the user to access their location as well as the the date and time as well as like other things like random number generation I guess. Um and you're also able to create interesting uh interfaces. So u similar

to the semantic galaxy demo we showed earlier. This is a way of like scrolling through your documents or your books or anything um and visualizing it uh in a new and interesting way. And then the final last demo is uh document parsing with Dockling. Uh this is of course a perfect example of privacy focused um uh demos where you upload your your sensitive documents and you want

to you know convert it to a more computer readable format um and everything runs in the browser. The model's like really really small. I think it's like 300 million parameters or something and it's extremely extremely uh good. So yeah, I uh think there are basically endless possibilities for you to explore create interesting applications. Let's say you're a model developer and you want to showcase your model in

a new and interesting way. Um I think browser and web AI is um is a really interesting way to to achieve that and I hope you build it with transformers.js. So thanks so