![Jason Mayes](https://web.dev/images/authors/jasonmayes.jpg) Jason Mayes [X](https://twitter.com/jason_mayes) [GitHub](https://github.com/jasonmayes) [LinkedIn](https://www.linkedin.com/in/webai) ![François Beaufort](https://web.dev/images/authors/beaufortfrancois.jpg) François Beaufort [GitHub](https://github.com/beaufortfrancois)

<br />

Great news! You've built a cool [Web AI application](https://goo.gle/made-with-tfjs)
that runs machine learning models directly on a user's device. It runs entirely
on the client-side web browser, without relying on the cloud. This on-device
design enhances user privacy, boosts performance, and reduces costs
significantly.

However, there's a hurdle. Your
[TensorFlow.js](https://www.tensorflow.org/js) model can operate on
both CPUs (WebAssembly) and more powerful GPUs (through
[WebGL](https://developer.mozilla.org/docs/Web/API/WebGL_API) and
[WebGPU](https://developer.chrome.com/blog/webgpu-release.md.txt)). The question is:
**how can you consistently automate browser testing with the selected hardware?**

Maintaining consistency is crucial for comparing machine learning model
performance over time as you iterate and improve them, prior to deployment for
real-world users to use on their device.

Setting up a consistent testing environment with GPUs can be harder than
expected. In this blog post, we'll share the problems we faced and how we solved
them, so you can improve your application's performance.

This isn't just for Web AI developers! If you're working on web gaming or
graphics, this post is valuable for you, too.

## What's in our automation toolbox

Here is what we are using:

- **Environment** : A Linux-based Google Colab [notebook](https://colab.google/notebooks/) connected to an NVIDIA T4 or V100 GPU. You can use other cloud platforms, such as Google Cloud (GCP), if preferred.
- **Browser** : Chrome supports [WebGPU](https://developer.chrome.com/docs/web-platform/webgpu), a powerful [successor to WebGL](https://developer.chrome.com/blog/from-webgl-to-webgpu), that brings the advancements of modern GPU APIs to the web.
- **Automation** : [Puppeteer](https://developer.chrome.com/docs/puppeteer) is a Node.js library that lets you control browsers programmatically with JavaScript. With Puppeteer, we can automate Chrome in headless mode, which means the browser runs without a visible interface, on a server. We are using the improved [new headless mode](https://developer.chrome.com/docs/chromium/new-headless), not the [legacy](https://developer.chrome.com/blog/headless-chrome) form.

> [!TIP]
> **Tip** : A [step-by-step
> guide](https://developer.chrome.com/docs/web-platform/webgpu/colab-headless) is available in our documentation, so you can try out the complete setup yourself.

## Verify the environment

The best way to check whether hardware acceleration is turned on in Chrome is to
type `chrome://gpu` into the address bar. You can
programmatically [perform the equivalent with Puppeteer](https://github.com/jasonmayes/headless-chrome-nvidia-t4-gpu-support/blob/main/examples/puppeteer/jPuppet.js)
with `console.log` or save the full report as PDF to check manually:

```javascript
/* Incomplete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters: Expands later
const browser = await puppeteer.launch({
  headless: 'new',
  args:  ['--no-sandbox']
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({ path: './gpu.pdf' });

await browser.close();
```

Open `chrome://gpu` and you should have the following results:

| Graphics feature status ||
|---|---|
| OpenGL: | Disabled |
| Vulkan: | Disabled |
| WebGL: | Software only, hardware acceleration unavailable. |
| WebGL2: | Software only, hardware acceleration unavailable. |
| WebGPU: | Disabled |
| **Problems detected.** WebGPU has been disabled via blocklist or the command line. ||

Not a great start. It's fairly clear that hardware detection was failing.
**WebGL, WebGL2, and WebGPU are essentially disabled or software only** . We
aren't alone in this problem - there are numerous discussions online of people
in a similar situation, including on the official Chrome support channels
([1](https://bugs.chromium.org/p/chromium/issues/detail?id=1459930)),
([2](https://bugs.chromium.org/p/chromium/issues/detail?id=765284)).

## Enable WebGPU and WebGL support

By default, Headless Chrome
[disables GPU](https://bugs.chromium.org/p/chromium/issues/detail?id=1416283).
To enable it on Linux, apply all of the following flags when launching Headless
Chrome:

- **`--no-sandbox`** flag disables [Chrome's security sandbox](https://developer.chrome.com/docs/puppeteer/troubleshooting#set_up_a_chrome_linux_sandbox), which isolates the browser process from the rest of the system. Running Chrome as root without this sandbox is not supported.
- **`--headless=new`** flag runs Chrome with the new and improved [headless mode](https://developer.chrome.com/docs/chromium/new-headless), without any visible UI.
- **`--use-angle=vulkan`** flag tells Chrome to use the [Vulkan backend](https://chromium.googlesource.com/angle/angle/+/HEAD/src/libANGLE/renderer/vulkan/README.md) for [ANGLE](https://chromium.googlesource.com/angle/angle/), which translates OpenGL ES 2/3 calls to Vulkan API calls.
- **`--enable-features=Vulkan`** flag enables Vulkan graphics backend for compositing and rasterization in Chrome.
- **`--disable-vulkan-surface`** flag disables the `VK_KHR_surface` vulkan instance extension. Instead of using a swapchain, [Bit blit](https://en.wikipedia.org/wiki/Bit_blit) is used for the present render result on screen.
- **`--enable-unsafe-webgpu`** flag enables the experimental WebGPU API in Chrome on Linux and disables the adapters blocklist.

Now we combine all the changes we have made so far. Here is the complete script.

```javascript
/* Complete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters
const browser = await puppeteer.launch({
  headless: 'new',
  args: [
    '--no-sandbox',
    '--headless=new',
    '--use-angle=vulkan',
    '--enable-features=Vulkan',
    '--disable-vulkan-surface',
    '--enable-unsafe-webgpu',
  ]
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({path: './gpu.pdf'});

await browser.close();
```

Run the script again. No WebGPU problems are detected and the value changes from
disabled to software only.

| Graphics feature status ||
|---|---|
| OpenGL: | Disabled |
| Vulkan: | Disabled |
| WebGL: | Software only, hardware acceleration unavailable. |
| WebGL2: | Software only, hardware acceleration unavailable. |
| WebGPU: | Software only, hardware acceleration unavailable. |

However, hardware acceleration is still unavailable, the NVIDIA T4 GPU isn't
detected.

## Install the correct GPU drivers

We investigated more closely the output of `chrome://gpu`, with some GPU experts
on the Chrome team. We found issues with the [default drivers installed on the
Linux Colab](https://github.com/googlecolab/colabtools/issues/3556#issuecomment-1499397023)
instance, causing issues with Vulkan, leading to Chrome unable to detect the
NVIDIA T4 GPU at the `GL_RENDERER` level as shown in the following output. This
causes problems with Headless Chrome.

| Driver information ||
|---|---|
| GL_RENDERER | ANGLE (Google, Vulkan 1.3.0 (SwiftShader Device (Subzero) (0x0000C0DE)), SwiftShader driver-5.0.0) |
[The default output doesn't detect NVIDIA T4 GPU.]

Installing the correct drivers that were compatible therefore fixes the issue.

| Driver information ||
|---|---|
| GL_RENDERER | ANGLE (NVIDIA Corporation, Tesla T4/PCIe/SSE2, OpenGL ES 3.2 NVIDIA 525.105.17) |
[Updated output after drivers are installed.]

To install the correct drivers, run the following commands during setup. The
last two lines help you to log the outputs of what NVIDIA drivers detects along
with `vulkaninfo`.

```console
apt-get install -y vulkan-tools libnvidia-gl-525

// Verify the NVIDIA drivers detects along with vulkaninfo
nvidia-smi
vulkaninfo --summary
```

Now run the script again and we get the following result. 🎉

| Graphics feature status ||
|---|---|
| OpenGL: | Enabled |
| Vulkan: | Enabled |
| WebGL: | Hardware accelerated but at reduced performance. |
| WebGL2: | Hardware accelerated but at reduced performance. |
| WebGPU: | Hardware accelerated but at reduced performance. |

By using the correct drivers and flags when running Chrome, we now have WebGPU
and WebGL support using the shiny, new headless mode.

## Behind the scenes: Our team's investigation

After much research, we didn't find working methods for the environment we
needed to execute in Google Colab, although there were some
[hopeful posts](https://mirzabilal.com/how-to-enable-hardware-acceleration-on-chrome-chromium-puppeteer-on-aws-in-headless-mode)
that worked in other environments, which was promising. Ultimately, we weren't
able to replicate their success in the Colab NVIDIA T4 environment, as we had 2
key issues:

1. Some combinations of flags allow detection of the GPU, but don't allow you to actually use the GPU.
2. Examples of working solutions by third parties used the old Chrome headless version, which at some point will be deprecated in favor of the [new version](https://developer.chrome.com/docs/chromium/new-headless). We needed a solution that worked with the new Headless Chrome to be better future proofed.

We confirmed the under utilization of the GPU by running an
[example TensorFlow.js web page for image recognition](https://tensorflowjs-fashion-mnist-classifier.glitch.me/),
whereby we trained a model to recognize clothing samples (sort of like a "hello
world" of machine learning).

On a regular machine, 50 training cycles (known as epochs) should run in less
than 1 second each. Calling Headless Chrome in its default state, we could log
the JavaScript console output to the Node.js server-side command line to see how
fast these training cycles were actually taking.

As expected, each training epoch took much longer than expected (several
seconds), which suggests Chrome has fallen back to plain old JS CPU execution
instead of utilizing the GPU:
![The training epochs move at a slower cadence.](https://developer.chrome.com/blog/image/initial_epoch_timing.gif) **Figure 1**: Real-time capture showing how long each training epoch took to execute (seconds).

After fixing the drivers and using the right combination of flags for Headless
Chrome, rerunning the TensorFlow.js training example results in much faster
training epochs.
![There's an increase in speed for epochs..](https://developer.chrome.com/blog/image/epoch_timing_with_headless.gif) **Figure 2**: Real-time capture showing the speed up of epochs.

## Summary

[Web AI has grown exponentially](https://www.youtube.com/watch?v=r7hOoCY6uGo)
since its creation in 2017. With browser technologies such as WebGPU, WebGL, and
[WebAssembly](https://webassembly.org/), a machine learning model's
mathematical operations can be further accelerated on the client side.

As of 2023 TensorFlow.js and MediaPipe Web crossed over 1 billion downloads of
models and libraries---a historic milestone and a sign of how web
developers and engineers are shifting to embrace [AI in their next generation
web apps to make some truly incredible solutions](https://goo.gle/made-with-tfjs).

With great success in usage comes great responsibility. At this level of usage
in production systems, the need arises for testing client-side, browser-based AI
models in a true browser environment, while also being scalable, automatable,
and within a known standardized hardware setup.
When we say "true browser environment," we mean that this is a real Chrome
browser with hardware support, as opposed to a software emulation.

By harnessing the combined power of the new Headless Chrome and Puppeteer, you
can confidently test such workloads in a standardized and replicable
environment, ensuring consistent and reliable results.

## Wrap up

A [step-by-step guide](https://developer.chrome.com/docs/web-platform/webgpu/colab-headless.md.txt) is available in
our documentation, so you can try out the complete setup yourself.

If you found this useful, drop a shout out over on
[LinkedIn](https://www.linkedin.com/in/WebAI),
[X (formerly Twitter)](https://twitter.com/jason_mayes), or whatever
social network you use using hashtag **#WebAI**. It would be great to hear any
feedback you have so we know to write more stuff like this in the future.

[Add a star on the GitHub repo](https://github.com/jasonmayes/headless-chrome-nvidia-t4-gpu-support)
to receive any future updates.

## Acknowledgements

A huge thank you to everyone on the Chrome team who helped debug the driver and
WebGPU issues we faced in this solution, with a special thanks to
[Jecelyn Yeen](https://jec.fish/) and
[Alexandra White](https://heyawhite.com) for helping to wordsmith
this blog post. Thanks to Yuly Novikov, Andrey Kosyakov, and
[Alex Rudenko](https://mastodon.online/@orkon) who were instrumental
in creating the final, working solution.