Melhorias com o WebAssembly e a WebGPU para uma IA da Web mais rápida, parte 2

Este documento é uma continuação das melhorias do WebAssembly e da WebGPU para uma IA da Web mais rápida, parte 1. Recomendamos que você leia esta postagem ou assista à palestra no IO 24 antes de continuar.

Austin Eng
Austin Eng
Deepti Gandluri
Deepti Gandluri
François Beaufort
François Beaufort

WebGPU

A WebGPU fornece aos aplicativos da Web acesso ao hardware da GPU do cliente para realizar cálculos eficientes e altamente paralelos. Desde o lançamento da WebGPU no Chrome, temos visto demonstrações incríveis de inteligência artificial (IA) e machine learning (ML) na Web.

Por exemplo, o Web Stable Diffusion (em inglês) demonstrou que era possível usar IA para gerar imagens a partir de textos diretamente no navegador. No início deste ano, a equipe do Mediapipe do Google publicou o suporte experimental à inferência de modelos de linguagem grandes.

A animação a seguir mostra o Gemma, o modelo de linguagem grande (LLM) de código aberto do Google, executado inteiramente no dispositivo no Chrome, em tempo real.

A seguinte demonstração do Hugging Face do modelo Segment Anywhere da Meta produz máscaras de objetos de alta qualidade inteiramente no cliente.

Esses são apenas alguns dos projetos incríveis que mostram o poder da WebGPU para IA e ML. A WebGPU permite que esses e outros modelos sejam executados de maneira significativamente mais rápida do que poderiam na CPU.

O comparativo de mercado da WebGPU do Hugging Face para incorporação de texto demonstra enormes velocidades em comparação com uma implementação de CPU do mesmo modelo. Em um laptop Apple M1 Max, a WebGPU era 30 vezes mais rápida. Outras pessoas relataram que a WebGPU acelera o comparativo de mercado em mais de 120 vezes.

Como melhorar recursos da WebGPU para IA e ML

A WebGPU é ótima para modelos de IA e ML, que podem ter bilhões de parâmetros, graças ao suporte a sombreadores de computação. Sombreadores de computação são executados na GPU e ajudam a executar operações de matriz paralela em grandes volumes de dados.

Entre as várias melhorias da WebGPU no último ano, continuamos adicionando mais recursos para aprimorar o desempenho de ML e IA na Web. Recentemente, lançamos dois novos recursos: produtos de ponto flutuante de 16 bits e produtos inteiros de pontos compactados.

Ponto flutuante de 16 bits

Lembre-se, cargas de trabalho de ML não exigem precisão. O shader-f16 é um recurso que permite o uso do tipo f16 na linguagem de sombreamento da WebGPU. Esse tipo de ponto flutuante ocupa 16 bits, em vez dos 32 bits normais. A f16 tem um intervalo menor e é menos preciso, mas isso é suficiente para muitos modelos de ML.

Esse recurso aumenta a eficiência de algumas maneiras:

  • Memória reduzida: os tensores com elementos f16 ocupam metade do espaço, o que reduz o uso da memória pela metade. Os cálculos da GPU geralmente têm gargalos na largura de banda da memória, portanto, metade da memória pode fazer com que os sombreadores sejam executados duas vezes mais rápido. Tecnicamente, você não precisa de f16 para economizar na largura de banda de memória. É possível armazenar os dados em um formato de baixa precisão e, em seguida, expandi-los para f32 completo no sombreador para fins de computação. No entanto, a GPU gasta poder de computação extra para compactar e descompactar os dados.

  • Redução da conversão de dados: f16 usa menos computação minimizando a conversão de dados. Os dados de baixa precisão podem ser armazenados e usados diretamente, sem conversão.

  • Maior paralelismo: as GPUs modernas são capazes de encaixar mais valores simultaneamente nas unidades de execução da GPU, permitindo um número maior de cálculos paralelos. Por exemplo, uma GPU que oferece suporte a até 5 trilhões de operações de ponto flutuante f32 por segundo pode aceitar 10 trilhões de operações de ponto flutuante f16 por segundo.

.
Captura de tela do comparativo da WebGPU para embedding de texto
Com a shader-f16, o comparativo de mercado WebGPU da Hugging Face para incorporação de texto executa o comparativo três vezes mais rápido do que o f32 no laptop Apple M1 Max.

O WebLLM é um projeto que pode executar vários modelos de linguagem grandes. Ele usa o Apache TVM, um framework de compilador de machine learning de código aberto.

Pedi ao WebLLM para planejar uma viagem a Paris usando o modelo de parâmetros de 8 bilhões do Llama 3. Os resultados mostram que, durante a fase de preenchimento automático do modelo, f16 é 2,1 vezes mais rápido do que f32. Durante a fase de decodificação, ele é 1,3 vez mais rápido.

Primeiro, os aplicativos precisam confirmar se o adaptador da GPU oferece suporte a f16 e, se disponível, ativá-lo explicitamente ao solicitar um dispositivo GPU. Se f16 não for compatível, não será possível solicitá-lo na matriz requiredFeatures.

// main.js

const adapter = await navigator.gpu.requestAdapter();
const supportsF16 = adapter.features.has('shader-f16');
if (supportsF16) {
  // Use f16.
  const device = await adapter.requestDevice({
    requiredFeatures: ['shader-f16'],
  });
  initApp(device);
}

Em seguida, nos sombreadores da WebGPU, você precisa ativar explicitamente a f16 na parte superior. Depois disso, você pode usá-lo no shader como qualquer outro tipo de dado flutuante.

// my-shader.wgsl

enable f16;

struct Data {
  values : array<vec4<f16>>
}
@group(0) @binding(0) var<storage, read> data : Data;
@compute @workgroup_size(64) fn main(@builtin(global_invocation_id) gid : vec3u) {
  let value : vec4<f16> = data.values[gid.x];
  ...
}

Produtos inteiros e escalados embalados

Muitos modelos ainda funcionam bem com apenas 8 bits de precisão (metade de f16). Isso é muito utilizado entre LLMs e modelos de imagem para segmentação e reconhecimento de objetos. Dito isso, a qualidade da saída dos modelos diminui com menos precisão, então a quantização de 8 bits não é adequada para todos os aplicativos.

Relativamente poucas GPUs oferecem suporte nativo a valores de 8 bits. É aqui que entram os produtos de ponto inteiro compactados. Nós enviamos o DP4a no Chrome 123.

As GPUs modernas têm instruções especiais para pegar dois números inteiros de 32 bits, interpretá-los como quatro números inteiros de 8 bits agrupados consecutivamente e calcular o produto escalar entre os componentes deles.

Isso é particularmente útil para IA e machine learning, porque os kernels de multiplicação de matrizes são compostos de muitos produtos pontuais.

Por exemplo, vamos multiplicar uma matriz 4 x 8 por um vetor 8 x 1. A computação disso envolve usar 4 produtos escalares para calcular cada um dos valores no vetor de saída. A, B, C e D.

Diagrama de exemplo de multiplicação de vetor de matriz

O processo para computar cada uma dessas saídas é o mesmo. vamos analisar as etapas envolvidas na computação de uma delas. Antes de qualquer cálculo, primeiro precisamos converter os dados inteiros de 8 bits em um tipo com o qual podemos realizar aritmética, como f16. Em seguida, fazemos uma multiplicação por elemento e, por fim, adicionamos todos os produtos. No total, para toda a multiplicação de vetor de matriz, realizamos 40 conversões de número inteiro em flutuação para descompactar os dados, 32 multiplicações flutuantes e 28 adições de pontos flutuantes.

Para matrizes maiores com mais operações, produtos de ponto inteiro compactados podem ajudar a reduzir a quantidade de trabalho.

Para cada uma das saídas no vetor de resultado, executamos duas operações de produto escalar usando o dot4U8Packed integrado da WebGPU Shading Language e, em seguida, adicionamos os resultados. No total, para toda a multiplicação de vetores de matriz, não realizamos nenhuma conversão de dados. Executamos 8 produtos escalar compactados e 4 adições de números inteiros.

Diagrama do exemplo de multiplicação de vetor de matriz de número inteiro empacotado

Testamos produtos inteiros e empacotados com dados de 8 bits em várias GPUs de consumo. Em comparação com o ponto flutuante de 16 bits, podemos observar que os 8 bits são de 1,6 a 2,8 vezes mais rápidos. Quando também usamos produtos inteiros e pontilhados compactados, o desempenho é ainda melhor. É de 1,7 a 2,9 vezes mais rápido.

Captura de tela da aceleração de multiplicação de vetores de matriz: f16 x u8
Gráfico 1: aceleração do vetor de matriz, comparando f16 a U8 e U8 com dot4U8Packed.

Verifique a compatibilidade do navegador com a propriedade wgslLanguageFeatures. Se a GPU não oferecer suporte nativo a produtos empacotados, o navegador realizará o polyfill da própria implementação.

// main.js

if (navigator.gpu.wgslLanguageFeatures.has('packed_4x8_integer_dot_product')) {
  // Use dot4U8Packed, dot4I8Packed builtin
  // functions in the shaders.
}

A seguinte diferença de snippet de código (diferença) destacando as mudanças necessárias para usar produtos de números inteiros compactados em um sombreador da WebGPU.

Antes: um sombreador da WebGPU que acumula produtos escalares parciais na variável "sum". No final da repetição, `sum` contém o produto escalar completo entre um vetor e uma linha da matriz de entrada.

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f16;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : vec4<f16> = vector.values[i];
    let v2 : vec4<f16> = matrix.values[start + i];
    sum += dot(v1, v2);
  }
}

Depois: um sombreador da WebGPU criado para usar produtos de ponto de números inteiros compactados. A principal diferença é que, em vez de carregar quatro valores flutuantes do vetor e da matriz, esse shader carrega um único número inteiro de 32 bits. Esse número inteiro de 32 bits contém os dados de quatro valores inteiros de 8 bits. Em seguida, chamamos dot4U8Packed para calcular o produto escalar dos dois valores.

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f32;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : u32 = vector.values[i];
    let v2 : u32 = matrix.values[start + i];
    sum += dot4U8Packed(v1, v2);
  }
}

Os produtos de ponto flutuante de 16 bits e inteiros empacotados são os recursos incluídos no Chrome que aceleram a IA e o ML. O ponto flutuante de 16 bits fica disponível quando é compatível com hardware, e o Chrome implementa produtos de pontos inteiros compactados em todos os dispositivos.

Você pode usar esses recursos no Chrome Stable hoje mesmo para ter um desempenho melhor.

Recursos propostos

No futuro, estamos investigando dois outros atributos: subgrupos e multiplicação de matrizes cooperativas.

O recurso de subgrupos permite que o paralelismo no nível de SIMD se comunique ou realize operações matemáticas coletivas, como uma soma de mais de 16 números. Isso permite um compartilhamento eficiente de dados entre linhas de execução. Os subgrupos são compatíveis com APIs modernas de GPUs, com nomes variados e em formas ligeiramente diferentes.

Criamos uma proposta com esse conjunto comum para o grupo de padronização da WebGPU. Além disso, criamos um protótipo de subgrupos no Chrome para criar uma sinalização experimental e colocamos nossos resultados iniciais na discussão. A principal questão é como garantir o comportamento portátil.

A multiplicação de matrizes cooperativas é uma adição mais recente às GPUs. Uma multiplicação de matrizes grande pode ser dividida em várias multiplicações de matrizes menores. A multiplicação de matrizes cooperativas realiza multiplicações nesses blocos menores de tamanho fixo em uma única etapa lógica. Nessa etapa, um grupo de linhas de execução colabora de forma eficiente para calcular o resultado.

Pesquisamos o suporte das APIs de GPU subjacentes e planejamos apresentar uma proposta para o grupo de padronização da WebGPU. Assim como com os subgrupos, esperamos que grande parte da discussão seja centrada na portabilidade.

Para avaliar o desempenho das operações de subgrupos, em um aplicativo real, integramos o suporte experimental a subgrupos ao MediaPipe e o testamos com o protótipo do Chrome para operações de subgrupos.

Usamos subgrupos nos kernels da GPU da fase de pré-preenchimento do modelo de linguagem grande. Portanto, estou relatando apenas a aceleração para a fase de pré-preenchimento. Em uma GPU Intel, vemos que os subgrupos têm desempenho duas vezes e meia mais rápido do que a linha de base. No entanto, essas melhorias não são consistentes em GPUs diferentes.

Captura de tela da aceleração de subgrupos na inferência do MediaPipe LLM
Gráfico 2. Os subgrupos fazem o preenchimento automático ser executado 2,5 vezes mais rápido na GPU Intel Tiger Lake GT2, com suporte experimental no Chrome e no Mediapipe.

O próximo gráfico mostra os resultados da aplicação de subgrupos para otimizar uma matriz multiplicada por microcomparação em várias GPUs de consumo. A multiplicação de matrizes é uma das operações mais pesadas em modelos de linguagem grandes. Os dados mostram que, em muitas das GPUs, os subgrupos aumentam a velocidade de 2, 5 e até 13 vezes a linha de base. No entanto, observe que na primeira GPU, os subgrupos não são muito melhores.

Captura de tela da aceleração de subgrupo para multiplicação de matrizes
Gráfico 3. A aplicação de subgrupos para multiplicação de matrizes pode aumentar ainda mais o desempenho.

A otimização da GPU é difícil

Em última análise, a melhor forma de otimizar sua GPU depende da GPU que o cliente oferece. Nem sempre o uso de recursos avançados da GPU compensa o esperado, porque pode haver muitos fatores complexos envolvidos. A melhor estratégia de otimização em uma GPU pode não ser a melhor em outra GPU.

Você quer minimizar a largura de banda da memória enquanto usa totalmente as linhas de execução de computação da GPU.

Os padrões de acesso à memória também podem ser muito importantes. As GPUs tendem a ter um desempenho muito melhor quando as linhas de execução de computação acessam a memória em um padrão ideal para o hardware. Importante: as características de desempenho variam de acordo com o hardware da GPU. Pode ser necessário executar otimizações diferentes dependendo da GPU.

No gráfico a seguir, usamos o mesmo algoritmo de multiplicação de matrizes, mas adicionamos outra dimensão para demonstrar melhor o impacto de várias estratégias de otimização e a complexidade e variância em diferentes GPUs. Introduzimos uma nova técnica aqui, que chamaremos de "Swizzle". O swizzle otimiza os padrões de acesso à memória para que sejam mais adequados ao hardware.

O swizzle de memória tem um impacto significativo. às vezes isso é ainda mais impactante do que subgrupos. Na GPU 6, o swizzle proporciona 12x mais velocidade, enquanto os subgrupos proporcionam 13x mais. Combinadas, elas têm uma aceleração incrível de 26x. Para outras GPUs, às vezes o swizzle e os subgrupos combinados têm um desempenho melhor do que os dois sozinhos. E em outras GPUs, só usar o swizzle tem o melhor desempenho.

Captura de tela da aceleração para estratégias de multiplicação de matrizes
Gráfico 4.

Ajustar e otimizar os algoritmos da GPU para funcionar bem em qualquer hardware pode exigir muito conhecimento. Felizmente, há uma quantidade enorme de trabalhos talentosos sendo direcionados a frameworks de bibliotecas de nível superior, como Mediapipe, Transformers.js, Apache TVM, ONNX Runtime Web e muito mais.

As bibliotecas e os frameworks estão bem posicionados para lidar com a complexidade de gerenciar várias arquiteturas de GPU e gerar códigos específicos da plataforma que vão funcionar bem no cliente.

Aprendizados

A equipe do Chrome continua a desenvolver os padrões WebAssembly e WebGPU com o objetivo de melhorar a plataforma Web para cargas de trabalho de aprendizado de máquina. Estamos investindo em primitivos de computação mais rápidos, melhor interoperabilidade entre os padrões da Web e para garantir que modelos grandes e pequenos possam ser executados de maneira eficiente em vários dispositivos.

Nosso objetivo é maximizar os recursos da plataforma e, ao mesmo tempo, manter o melhor da Web: seu alcance, usabilidade e portabilidade. E não estamos fazendo isso sozinhos. Estamos trabalhando em colaboração com outros fornecedores de navegadores do W3C e com vários parceiros de desenvolvimento.

Esperamos que você se lembre do seguinte ao trabalhar com o WebAssembly e a WebGPU:

  • A inferência de IA já está disponível na Web em todos os dispositivos. Isso traz as vantagens da execução em dispositivos clientes, como menor custo de servidor, baixa latência e maior privacidade.
  • Embora muitos recursos discutidos sejam relevantes principalmente para os autores do framework, seus aplicativos podem se beneficiar sem muita sobrecarga.
  • Os padrões da Web são fluidos e estão em evolução, e estamos sempre em busca de feedback. Compartilhe seu arquivo com o WebAssembly e a WebGPU.

Agradecimentos

Gostaríamos de agradecer à equipe de gráficos da Web da Intel, que foi fundamental na condução da WebGPU f16 e nos recursos de produtos inteiros de pontos. Gostaríamos de agradecer aos outros membros dos grupos de trabalho WebAssembly e WebGPU do W3C, incluindo outros fornecedores de navegadores.

Agradecemos às equipes de IA e ML do Google e da comunidade de código aberto por serem parceiros incríveis. E, claro, a todos os nossos colegas de equipe, que possibilitam tudo isso.