Turbine os testes do modelo de IA da Web: WebGPU, WebGL e Headless Chrome

François Beaufort
François Beaufort

Temos uma ótima notícia! Você criou um aplicativo de IA da Web legal que executa modelos de machine learning diretamente no dispositivo de um usuário. Ele é executado inteiramente no navegador da Web do lado do cliente, sem depender da nuvem. Esse design no dispositivo melhora a privacidade do usuário, melhora a performance e reduz custos significativamente.

No entanto, há um obstáculo. O modelo do TensorFlow.js pode operar em ambas as CPUs (WebAssembly) e as GPUs mais potentes (por meio de WebGL e WebGPU). A pergunta é: como automatizar de maneira consistente os testes do navegador com o hardware selecionado?

Manter a consistência é crucial para comparar o desempenho do modelo de machine learning ao longo do tempo à medida que você itera e melhora os modelos, antes da implantação para que usuários do mundo real usem nos dispositivos deles.

Configurar um ambiente de teste consistente com GPUs pode ser mais difícil do que o esperado. Nesta postagem do blog, vamos compartilhar os problemas que enfrentamos e como os resolvemos para que você possa melhorar o desempenho do seu aplicativo.

Isso não é apenas para desenvolvedores de IA da Web! Se você estiver trabalhando com jogos na Web ou gráficos, esta postagem também é útil para você.

Recursos de automação

Veja o que estamos usando:

  • Ambiente: um notebook do Google Colab baseado em Linux conectado a uma GPU NVIDIA T4 ou V100. Se preferir, use outras plataformas de nuvem, como o Google Cloud (GCP).
  • Navegador: o Chrome é compatível com a WebGPU, uma poderosa sucessora da WebGL que leva os avanços das APIs modernas de GPU para a Web.
  • Automação: o Puppeteer é uma biblioteca do Node.js que permite controlar os navegadores de maneira programática com JavaScript. Com o Puppeteer, é possível automatizar o Chrome no modo headless, o que significa que o navegador será executado em um servidor sem uma interface visível. Estamos usando o novo modo headless aprimorado, não o formulário legado.

Verificar o ambiente

A melhor maneira de verificar se a aceleração de hardware está ativada no Chrome é digitando chrome://gpu na barra de endereço. É possível executar o equivalente com o Puppeteer (em inglês) de maneira programática com console.log ou salvar o relatório completo como PDF para verificação manual:

/* 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();

Abra chrome://gpu. Você verá estes resultados:

Status do recurso gráfico
OpenGL: Desativada
Vulkan: Desativada
WebGL: Somente software, aceleração de hardware indisponível.
WebGL2: Somente software, aceleração de hardware indisponível.
WebGPU: Desativada

Problemas detectados.
A WebGPU foi desativada por meio da lista de bloqueio ou da linha de comando.

Não foi um ótimo começo. Está claro que a detecção do hardware estava falhando. WebGL, WebGL2 e WebGPU estão essencialmente desativados ou somente de software. Não estamos sozinhos nesse problema. Existem inúmeras discussões on-line de pessoas em uma situação semelhante, inclusive nos canais de suporte oficiais do Chrome (1, 2).

Ativar o suporte à WebGPU e ao WebGL

Por padrão, o Headless Chrome desativa a GPU. Para ativá-lo no Linux, aplique todas as flags abaixo ao iniciar o Headless Chrome:

  • A sinalização --no-sandbox desativa o sandbox de segurança do Chrome, que isola o processo do navegador do restante do sistema. Não é possível executar o Chrome como raiz sem esse sandbox.
  • A sinalização --headless=new executa o Chrome com o novo e aprimorado modo headless, sem nenhuma interface visível.
  • A flag --use-angle=vulkan instrui o Chrome a usar o back-end do Vulkan para o ANGLE, que converte chamadas do OpenGL ES 2/3 em chamadas de API do Vulkan.
  • A sinalização --enable-features=Vulkan ativa o back-end gráfico do Vulkan para composição e varredura no Chrome.
  • A sinalização --disable-vulkan-surface desativa a extensão da instância do vulkan VK_KHR_surface. Em vez de uma cadeia de troca, a Bit blit é usada para o resultado de renderização presente na tela.
  • A sinalização --enable-unsafe-webgpu ativa a API WebGPU experimental no Chrome no Linux e desativa a lista de bloqueio de adaptadores.

Agora, combinamos todas as mudanças feitas até o momento. Veja o script completo.

/* 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();

Execute o script novamente. Nenhum problema de WebGPU é detectado, e o valor muda de desativado para somente software.

Status do recurso gráfico
OpenGL: Desativada
Vulkan: Desativada
WebGL: Somente software, aceleração de hardware indisponível.
WebGL2: Somente software, aceleração de hardware indisponível.
WebGPU: Somente software, aceleração de hardware indisponível.

No entanto, a aceleração de hardware ainda está indisponível, e a GPU NVIDIA T4 não é detectada.

Instalar os drivers de GPU corretos

Investigamos mais de perto a saída de chrome://gpu com alguns especialistas em GPU da equipe do Chrome. Encontramos problemas com os drivers padrão instalados na instância do Linux Colab, causando o Vulkan, fazendo com que o Chrome não consiga detectar a GPU NVIDIA T4 no nível GL_RENDERER, conforme mostrado no resultado a seguir. Isso causa problemas com o Headless Chrome.

A saída padrão não detecta a GPU NVIDIA T4.
Informações do motorista
GL_RENDERER ANGLE (Google, Vulkan 1.3.0 (SwiftShader Device (Subzero) (0x0000C0DE)), driver SwiftShader -5.0.0)

A instalação dos drivers corretos que eram compatíveis, portanto, corrige o problema.

Resposta atualizada após a instalação dos drivers.
Informações do motorista
GL_RENDERER ANGLE (NVIDIA Corporation, Tesla T4/PCIe/SSE2, OpenGL ES 3.2 NVIDIA 525.105.17)

Para instalar os drivers corretos, execute os comandos a seguir durante a instalação. As duas últimas linhas ajudam a registrar as saídas dos drivers da NVIDIA detectados com vulkaninfo.

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

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

Execute o script novamente e o resultado a seguir será exibido. 🎉

Status do recurso gráfico
OpenGL: Ativado
Vulkan: Ativado
WebGL: Aceleração de hardware, mas com desempenho reduzido.
WebGL2: Aceleração de hardware, mas com desempenho reduzido.
WebGPU: Aceleração de hardware, mas com desempenho reduzido.

Com o uso dos drivers e sinalizações corretos ao executar o Chrome, agora oferecemos suporte à WebGPU e ao WebGL usando o novo modo headless.

Nos bastidores: investigação da nossa equipe

Após muita pesquisa, não encontramos métodos de trabalho para o ambiente que precisávamos executar no Google Colab, embora houvesse algumas postagens esperançasas que funcionassem em outros ambientes, o que era promissor. No final das contas, não conseguimos replicar o sucesso deles no ambiente Colab NVIDIA T4, já que tivemos dois problemas principais:

  1. Algumas combinações de sinalizações permitem a detecção da GPU, mas não permitem o uso real dela.
  2. Exemplos de soluções de trabalho de terceiros usavam a antiga versão headless do Chrome, que em algum momento será descontinuada e substituída pela nova versão. Precisávamos de uma solução que funcionasse com o novo Headless Chrome para ser mais preparada para o futuro.

Confirmamos a baixa utilização da GPU executando uma página da Web de exemplo do TensorFlow.js para reconhecimento de imagens, em que treinamos um modelo para reconhecer amostras de roupas (como um "Hello World" do machine learning).

Em uma máquina regular, cada 50 ciclos de treinamento, conhecidos como períodos, precisam ser executados em menos de 1 segundo. Ao chamar o Headless Chrome no estado padrão, poderíamos registrar a saída do Console JavaScript na linha de comando do lado do servidor do Node.js para ver a velocidade real desses ciclos de treinamento.

Como esperado, cada período de treinamento levou muito mais tempo do que o esperado (vários segundos), o que sugere que o Chrome voltou à execução da CPU do JS simples em vez de usar a GPU:

Os períodos de treinamento são mais lentos.
Figura 1: captura em tempo real mostrando quanto tempo cada período de treinamento levou para ser executado (segundos).

Depois de corrigir os drivers e usar a combinação certa de sinalizações para o Headless Chrome, executar novamente o exemplo de treinamento do TensorFlow.js resulta em períodos de treinamento muito mais rápidos.

Há um aumento na velocidade para períodos.
Figura 2: captura em tempo real mostrando a aceleração dos períodos.

Resumo

A IA da Web cresceu exponencialmente desde que foi criada, em 2017. Com tecnologias de navegadores como WebGPU, WebGL e WebAssembly, as operações matemáticas de um modelo de machine learning podem ser ainda mais aceleradas no lado do cliente.

Em 2023, o TensorFlow.js e o MediaPipe Web ultrapassaram mais de 1 bilhão de downloads de modelos e bibliotecas, um marco histórico e um sinal de como os desenvolvedores e engenheiros da Web estão migrando para adotar a IA nos apps da Web de última geração para criar soluções realmente incríveis.

O uso bem-sucedido traz grandes responsabilidades. Nesse nível de uso em sistemas de produção, surge a necessidade de testar modelos de IA baseados em navegador no lado do cliente em um ambiente de navegador real, além de serem escalonáveis, automatizáveis e dentro de uma configuração de hardware padronizada conhecida.

Ao aproveitar a capacidade combinada dos novos Headless Chrome e Puppeteer, é possível testar essas cargas de trabalho com confiança em um ambiente padronizado e replicável, garantindo resultados consistentes e confiáveis.

Conclusão

Um guia passo a passo está disponível em nossa documentação para que você possa testar a configuração completa por conta própria.

Se você achou isso útil, envie uma mensagem no LinkedIn, X (antigo Twitter) ou em qualquer rede social que você use usando a hashtag #WebAI. Seria ótimo receber seu feedback para que possamos escrever mais coisas como esta no futuro.

Adicione uma estrela ao repositório do GitHub para receber atualizações futuras.

Agradecimentos

Agradecemos imensamente a todos da equipe do Chrome que ajudaram a depurar os problemas do driver e da WebGPU que enfrentamos nessa solução, em um agradecimento especial a Jecelyn Yeen e Alexandra White por ajudarem com o Wordsmith nesta postagem do blog. Graças a Yuly Novikov, Andrey Kosyakov e Alex Rudenko, que foram fundamentais na criação da solução final de trabalho.