Mejoras de WebAssembly y WebGPU para una IA web más rápida (parte 2)

Este documento es la continuación de las mejoras de WebAssembly y WebGPU para la IA web más rápida, parte 1. Te recomendamos que leas esta publicación o mires la charla en IO 24 antes de continuar.

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

WebGPU

WebGPU brinda a las aplicaciones web acceso al hardware de la GPU del cliente para que realicen un procesamiento eficiente y altamente paralelo. Desde el lanzamiento de WebGPU en Chrome, vimos demostraciones increíbles de inteligencia artificial (IA) y aprendizaje automático (AA) en la Web.

Por ejemplo, Web Stable Diffusion demostró que era posible usar IA para generar imágenes a partir de texto, directamente en el navegador. A principios de este año, el propio equipo de Mediapipe de Google publicó compatibilidad experimental para la inferencia de modelos de lenguaje grandes.

En la siguiente animación, se muestra Gemma, el modelo de lenguaje grande (LLM) de código abierto de Google, que se ejecuta completamente en el dispositivo en Chrome, en tiempo real.

En la siguiente demostración de Hugging Face del modelo de Segment Any de Meta, se producen máscaras de objetos de alta calidad en el cliente por completo.

Estos son solo algunos de los asombrosos proyectos que muestran la potencia de WebGPU para IA y AA. WebGPU permite que estos modelos y otros se ejecuten significativamente más rápido de lo que podrían en la CPU.

La comparativa de WebGPU de Hugging Face para la incorporación de texto de Hugging Face demuestra grandes velocidades en comparación con una implementación de CPU del mismo modelo. En una laptop Apple M1 Max, la WebGPU era 30 veces más rápida. Otros informaron que WebGPU acelera las comparativas más de 120 veces.

Mejora las funciones de WebGPU para IA y AA

WebGPU es excelente para modelos de IA y AA, que pueden tener miles de millones de parámetros gracias a la compatibilidad con sombreadores de cómputos. Los sombreadores de cómputos se ejecutan en la GPU y ayudan a ejecutar operaciones de array paralelas en grandes volúmenes de datos.

Entre las numerosas mejoras a WebGPU del año pasado, seguimos agregando más capacidades para mejorar el rendimiento de la IA y el AA en la Web. Recientemente, lanzamos dos funciones nuevas: los productos de punto flotante de 16 bits y los productos de punto entero empaquetados.

Número de punto flotante de 16 bits

Recuerda que las cargas de trabajo de AA no requieren precisión. shader-f16 es una función que permite el uso del tipo f16 en el lenguaje de sombreado WebGPU. Este tipo de punto flotante usa 16 bits en lugar de los 32 bits habituales. f16 tiene un rango más pequeño y es menos preciso, pero para muchos modelos de AA, esto es suficiente.

Esta función aumenta la eficiencia de las siguientes maneras:

  • Memoria reducida: Los tensores con elementos f16 ocupan la mitad del espacio, lo que reduce a la mitad el uso de memoria. Los cálculos de la GPU suelen tener un cuello de botella en el ancho de banda de la memoria, por lo que la mitad de la memoria a menudo puede significar que los sombreadores se ejecuten el doble de rápido. Técnicamente, no necesitas f16 para ahorrar en ancho de banda de memoria. Es posible almacenar los datos en un formato de baja precisión y, luego, expandirlos a f32 por completo en el sombreador para realizar el procesamiento. Sin embargo, la GPU consume potencia de procesamiento adicional para empaquetar y desempaquetar los datos.

  • Conversión de datos reducida: f16 usa menos procesamiento, ya que minimiza la conversión de datos. Los datos de baja precisión se pueden almacenar y usar directamente sin conversión.

  • Mayor paralelismo: Las GPU modernas pueden ajustar más valores simultáneamente en las unidades de ejecución de la GPU, lo que le permite realizar una mayor cantidad de procesamientos paralelos. Por ejemplo, una GPU que admite hasta 5 billones de operaciones de punto flotante f32 por segundo podría admitir 10 billones de operaciones de punto flotante f16 por segundo.

Captura de pantalla de las comparativas de WebGPU para la incorporación de texto
Con shader-f16, la comparativa de WebGPU de Hugging Face para incorporación de texto la ejecuta 3 veces más rápido que f32 en una laptop Apple M1 Max.

WebLLM es un proyecto que puede ejecutar varios modelos grandes de lenguaje. Usa Apache TVM, un framework de compilación de aprendizaje automático de código abierto.

Le pedí a WebLLM que planificara un viaje a París con el modelo de ocho mil millones de parámetros de Llama 3. Los resultados muestran que durante la fase de autocompletado del modelo, f16 es 2.1 veces más rápido que f32. Durante la fase de decodificación, es más de 1.3 veces más rápido.

En primer lugar, las aplicaciones deben confirmar que el adaptador de GPU sea compatible con f16 y, si está disponible, habilitarlo explícitamente cuando soliciten un dispositivo GPU. Si f16 no es compatible, no puedes solicitarlo en el array 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);
}

Luego, en tus sombreadores WebGPU, debes habilitar f16 de manera explícita en la parte superior. Después de eso, podrás usarlo dentro del sombreador como cualquier otro tipo de datos de número de punto flotante.

// 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];
  ...
}

Productos escalares de números enteros empaquetados

Muchos modelos aún funcionan bien con solo 8 bits de precisión (la mitad de f16). Esto es popular entre los LLM y los modelos de imagen para la segmentación y el reconocimiento de objetos. Dicho esto, la calidad de salida de los modelos se degrada con menos precisión, por lo que la cuantización de 8 bits no es adecuada para todas las aplicaciones.

Pocas GPU admiten valores de 8 bits de forma nativa. Aquí es donde entran en juego los productos de punto entero empaquetados. Enviamos la DP4a en Chrome 123.

Las GPU modernas tienen instrucciones especiales para tomar dos números enteros de 32 bits, interpretarlos como 4 de 8 bits empaquetados de forma consecutiva y calcular el producto escalar entre sus componentes.

Esto es particularmente útil para la IA y el aprendizaje automático porque los kernels de multiplicación de matrices están compuestos por muchos productos escalares.

Por ejemplo, multipliquemos una matriz de 4 x 8 por un vector de 8 x 1. El cálculo implica tomar productos de 4 puntos para calcular cada uno de los valores en el vector de salida. A, B, C y D.

Diagrama de ejemplo de multiplicación de vector matricial

El proceso para calcular cada uno de estos resultados es el mismo. veremos los pasos involucrados en el procesamiento de una de ellas. Antes de cualquier cálculo, primero debemos convertir los números enteros de 8 bits a un tipo con el que podamos realizar operaciones aritméticas, como f16. Luego, ejecutamos una multiplicación por elementos y, por último, sumamos todos los productos juntos. En total, para toda la multiplicación de vectores de matrices y matrices, realizamos 40 conversiones de número entero a flotantes para descomprimir los datos, 32 multiplicaciones de número de punto flotante y 28 adiciones de número de punto flotante.

Para matrices más grandes con más operaciones, los productos de punto entero empaquetados pueden ayudar a reducir la cantidad de trabajo.

Para cada una de las salidas en el vector de resultados, realizamos dos operaciones de producto de punto empaquetados con el lenguaje sombreado de WebGPU integrado dot4U8Packed y, luego, sumamos los resultados. En total, para la multiplicación completa de vectores matriciales, no realizamos ninguna conversión de datos. Ejecutamos 8 productos de punto empaquetados y 4 adiciones de números enteros.

Diagrama de ejemplo de multiplicación de vectores de matrices y números enteros empaquetados

Probamos productos de punto enteros empaquetados con datos de 8 bits en una variedad de GPU de consumidores. En comparación con el punto flotante de 16 bits, podemos ver que 8 bits es entre 1.6 y 2.8 veces más rápido. Cuando también usamos productos de punto entero empaquetados, el rendimiento es aún mejor. Es entre 1.7 y 2.9 veces más rápido.

Captura de pantalla de la aceleración de la multiplicación de vectores matriciales: f16 frente a u8
Gráfico 1: Aceleración del vector de matriz, comparando f16 con U8 y U8 con dot4U8Packed.

Verifica la compatibilidad del navegador con la propiedad wgslLanguageFeatures. Si la GPU no admite productos de puntos empaquetados de forma nativa, el navegador aplica polyfills su propia implementación.

// main.js

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

La siguiente diferencia de fragmento de código (diferencia) en la que se destacan los cambios necesarios para usar productos de números enteros empaquetados en un sombreador de WebGPU.

Antes: Sombreador de WebGPU que acumula productos escalares parciales en la variable `suma`. Al final del bucle, "sum" contiene el producto escalar completo entre un vector y una fila de la 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);
  }
}

After: sombreador WebGPU escrito para usar productos de punto enteros empaquetados. La principal diferencia es que, en lugar de cargar 4 valores flotantes del vector y la matriz, el sombreador carga un único número entero de 32 bits. Este número entero de 32 bits contiene los datos de cuatro valores enteros de 8 bits. Luego, llamamos a dot4U8Packed para calcular el producto escalar de los dos 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);
  }
}

Tanto los productos de punto flotante de 16 bits como los de punto entero empaquetado son las funciones incluidas en Chrome que aceleran la IA y el AA. El punto flotante de 16 bits está disponible si el hardware lo admite, y Chrome implementa productos de punto enteros empaquetados en todos los dispositivos.

Puedes usar estas funciones en la versión estable de Chrome hoy mismo para obtener un mejor rendimiento.

Componentes propuestos

De cara al futuro, estamos investigando dos funciones más: los subgrupos y la multiplicación de matrices cooperativas.

La función de subgrupos permite que el paralelismo de nivel SIMD se comunique o realice operaciones matemáticas colectivas, como una suma de más de 16 números. Esto permite compartir datos entre subprocesos de forma eficiente. Las APIs modernas de GPU admiten subgrupos, con nombres variables y formas ligeramente diferentes.

Resumimos el conjunto común en una propuesta que llevamos al grupo de estandarización de WebGPU. Además, creamos prototipos de subgrupos en Chrome detrás de una marca experimental y presentamos los resultados iniciales en el debate. El problema principal es cómo garantizar un comportamiento portátil.

La multiplicación de matrices cooperativa es una adición más reciente a las GPU. Una multiplicación de matrices grande se puede desglosar en múltiples multiplicaciones de matrices más pequeñas. La multiplicación de matrices cooperativas realiza multiplicaciones en estos bloques más pequeños de tamaño fijo en un solo paso lógico. En ese paso, un grupo de subprocesos cooperan de manera eficiente para calcular el resultado.

Encuestamos la compatibilidad con las APIs de GPU subyacentes y planeamos presentar una propuesta al grupo de estandarización de WebGPU. Al igual que con los subgrupos, esperamos que gran parte del debate se centre en la portabilidad.

Para evaluar el rendimiento de las operaciones de subgrupos, en una aplicación real, integramos la asistencia experimental para subgrupos en MediaPipe y la probamos con el prototipo de Chrome para las operaciones de subgrupos.

Usamos subgrupos en los kernels de GPU de la fase de autocompletado del modelo grande de lenguaje, por lo que solo estoy informando la aceleración de la fase de autocompletado. En una GPU Intel, vemos que los subgrupos tienen un rendimiento dos veces y medio más rápido que el modelo de referencia. Sin embargo, estas mejoras no son coherentes en diferentes GPU.

Captura de pantalla de aceleración de subgrupos en la inferencia MediaPipe LLM
Gráfico 2. Los subgrupos permiten que el autocompletado se ejecute 2.5 veces más rápido en la GPU Intel Tiger Lake GT2, con compatibilidad experimental en Chrome y Mediapipe.

En el siguiente gráfico, se muestran los resultados de la aplicación de subgrupos para optimizar una matriz multiplicando las microcomparativas en varias GPU de consumo. La multiplicación de matrices es una de las operaciones más pesadas en los modelos grandes de lenguaje. Los datos muestran que en muchas de las GPU, los subgrupos aumentan la velocidad dos, cinco o incluso trece veces la línea base. Sin embargo, ten en cuenta que en la primera GPU, los subgrupos no son mucho mejores.

Captura de pantalla de Aceleración de subgrupos para la multiplicación de matrices
Gráfico 3. Aplicar subgrupos para la multiplicación de matrices puede aumentar aún más el rendimiento.

La optimización de la GPU es difícil

En última instancia, la mejor manera de optimizar tu GPU depende de lo que ofrezca el cliente. Usar funciones nuevas y sofisticadas de GPU no siempre da los resultados que esperas, ya que pueden implicar muchos factores complejos. Es posible que la mejor estrategia de optimización en una GPU no sea la mejor en otra.

Quieres minimizar el ancho de banda de la memoria mientras usas por completo los subprocesos de procesamiento de la GPU.

Los patrones de acceso a la memoria también pueden ser muy importantes. Las GPU suelen tener un rendimiento mucho mejor cuando los subprocesos de procesamiento acceden a la memoria en un patrón óptimo para el hardware. Importante: Debes esperar características de rendimiento diferentes en distintos hardware de GPU. Es posible que debas ejecutar diferentes optimizaciones según la GPU.

En el siguiente gráfico, tomamos el mismo algoritmo de multiplicación de matrices, pero agregamos otra dimensión para demostrar con más detalle el impacto de varias estrategias de optimización y la complejidad y la variación entre las diferentes GPU. Hemos introducido una nueva técnica aquí: llamaremos "Swizzle". Swizzle optimiza los patrones de acceso a la memoria para que sean más óptimos para el hardware.

Puedes ver que el swizzle de memoria tiene un impacto significativo; a veces tienen incluso más impacto que los subgrupos. En GPU 6, Swizzle proporciona una aceleración de 12 veces, mientras que los subgrupos proporcionan una aceleración de 13 veces. Combinados, tienen una velocidad increíble de 26 veces. Para otras GPU, a veces, Swizzle y subgrupos combinados tienen un mejor rendimiento que uno solo. Y en otras GPU, el uso exclusivo de swizzle funciona mejor.

Captura de pantalla de la aceleración de las estrategias de multiplicación de matrices
Gráfico 4.

El ajuste y la optimización de los algoritmos de GPU para que funcionen bien en cada pieza de hardware puede requerir mucha experiencia. Sin embargo, por suerte, hay una gran cantidad de trabajo talentoso que se lleva a cabo en frameworks de bibliotecas de nivel superior, como Mediapipe, Transformers.js, Apache TVM y ONNX Runtime Web, entre otros.

Las bibliotecas y los frameworks están bien preparados para manejar la complejidad de administrar diversas arquitecturas de GPU y generar código específico de la plataforma que se ejecutará bien en el cliente.

Conclusiones

El equipo de Chrome sigue ayudando a evolucionar los estándares de WebAssembly y WebGPU con el objetivo de mejorar la plataforma web para las cargas de trabajo de aprendizaje automático. Invertimos en primitivas de procesamiento más rápidas, una mejor interoperabilidad entre los estándares de la Web y nos aseguramos de que los modelos grandes y pequeños puedan ejecutarse de manera eficiente en todos los dispositivos.

Nuestro objetivo es maximizar las capacidades de la plataforma y, al mismo tiempo, conservar lo mejor de la Web: su alcance, usabilidad y portabilidad. Y no estamos haciendo esto solos. Estamos trabajando en colaboración con otros proveedores de navegadores de W3C y con muchos socios de desarrollo.

Esperamos que recuerdes lo siguiente cuando trabajes con WebAssembly y WebGPU:

  • La inferencia de IA ya está disponible en la Web y en todos los dispositivos. Esto ofrece la ventaja de ejecutar aplicaciones en los dispositivos cliente, como un costo de servidor reducido, una latencia baja y una mayor privacidad.
  • Aunque muchas de las funciones mencionadas son relevantes principalmente para los autores del framework, tus aplicaciones pueden beneficiarse sin mucha sobrecarga.
  • Los estándares de la Web son fluidos y evolucionando, por lo que siempre estamos buscando comentarios. Comparte el tuyo para WebAssembly y WebGPU.

Agradecimientos

Queremos agradecer al equipo de gráficos web de Intel, que fue fundamental para impulsar la WebGPU f16 y las funciones del producto empaquetadas con valores enteros. Queremos agradecer a los demás miembros de los grupos de trabajo de WebAssembly y WebGPU de W3C, incluidos los otros proveedores de navegadores.

Gracias a los equipos de IA y AA de Google y de la comunidad de código abierto por ser socios increíbles. Y, por supuesto, todos nuestros compañeros de equipo que hacen todo esto posible.