WebLLM: Inferencia de LLMs en el Navegador con Aceleración WebGPU

La inferencia de modelos de lenguaje grandes (LLMs) tradicionalmente ha requerido infraestructura de servidor significativa debido a sus requisitos computacionales y de memoria. Esto introduce latencia de red, costos operativos y preocupaciones de privacidad, ya que los datos deben ser transmitidos a un servidor remoto. WebLLM aborda este problema fundamental al trasladar la inferencia de LLMs directamente al navegador del usuario, aprovechando la capacidad de WebGPU para la aceleración de hardware.

Esta aproximación democratiza el acceso a capacidades avanzadas de IA, permitiendo que aplicaciones web ofrezcan experiencias de usuario ricas y sensibles al contexto sin depender de servicios en la nube para cada interacción. La capacidad de ejecutar LLMs localmente transforma el paradigma de diseño de aplicaciones de IA, abriendo nuevas posibilidades para asistentes personales, herramientas de productividad y experiencias interactivas donde la privacidad y la baja latencia son críticas.

Arquitectura del Sistema

WebLLM se construye sobre una arquitectura modular que integra varias tecnologías web y de compilación de ML. El núcleo del sistema es el MLCEngine, una interfaz que abstrae la carga y ejecución del modelo. Los modelos se distribuyen en formato MLC, que incluye los pesos del modelo y una librería WebAssembly (.wasm) que contiene los ejecutables optimizados para la computación. Esta librería .wasm es generada por TVMjs, un runtime de Apache TVM compilado para WebAssembly, que permite la ejecución eficiente de operaciones tensoriales en el navegador.

La aceleración de hardware se logra mediante WebGPU, una API web que expone capacidades de GPU modernas a JavaScript, permitiendo la ejecución de cómputos paralelos intensivos. WebLLM gestiona la descarga y el almacenamiento en caché de los artefactos del modelo (configuración, WASM, tokenizer) utilizando APIs de almacenamiento del navegador como Cache API, IndexedDB u Origin Private File System (OPFS). Para mejorar la reactividad de la UI, WebLLM soporta la ejecución del MLCEngine en Web Workers o Service Workers, delegando las operaciones computacionalmente intensivas a hilos separados. La compatibilidad con la API de OpenAI se logra mediante una capa de abstracción que traduce las llamadas de la API a las operaciones internas del MLCEngine, incluyendo funcionalidades como streaming, JSON-mode y, en desarrollo, function-calling. La integridad de los artefactos del modelo puede ser verificada mediante Subresource Integrity (SRI) hashes para garantizar la autenticidad y prevenir la manipulación.

import { CreateMLCEngine } from "@mlc-ai/web-llm";

const initProgressCallback = (initProgress) => {
  console.log(initProgress);
};

const selectedModel = "Llama-3.1-8B-Instruct-q4f32_1-MLC";
const engine = await CreateMLCEngine(
  selectedModel,
  { initProgressCallback: initProgressCallback },
);

const messages = [
  { role: "system", content: "You are a helpful AI assistant." },
  { role: "user", content: "Hello!" },
];

const chunks = await engine.chat.completions.create({
  messages,
  temperature: 1,
  stream: true,
  stream_options: { include_usage: true },
});

let reply = "";
for await (const chunk of chunks) {
  reply += chunk.choices[0]?.delta.content || "";
  console.log(reply);
  if (chunk.usage) {
    console.log(chunk.usage);
  }
}

import { CreateWebWorkerMLCEngine } from "@mlc-ai/web-llm";

async function main() {
  const engine = await CreateWebWorkerMLCEngine(
    new Worker(new URL("./worker.ts", import.meta.url), {
      type: "module",
    }),
    selectedModel,
    { initProgressCallback },
  );
  // ... el resto de la lógica de chat
}

Fundamentos Teóricos

La capacidad de ejecutar modelos complejos como los LLMs en entornos restringidos como el navegador web se basa en décadas de investigación en compiladores, optimización de hardware y sistemas distribuidos. El uso de WebAssembly como objetivo de compilación para TVMjs se alinea con el concepto de una máquina virtual portable y segura, similar a los principios de la JVM o CLR, pero optimizada para la web. Esto permite la ejecución de código de alto rendimiento escrito en lenguajes como C++ o Rust en el navegador, un concepto explorado en papers sobre la virtualización de runtimes.

La optimización de operaciones tensoriales en GPUs, fundamental para WebGPU, se remonta a la investigación en computación paralela y arquitecturas SIMD. Algoritmos como la multiplicación de matrices (GEMM) y las convoluciones, que son la base de las redes neuronales, han sido extensamente estudiados y optimizados en la literatura académica. La gestión de la memoria y el caching en el navegador, aunque a un nivel de aplicación, se relaciona con principios de jerarquía de memoria y políticas de reemplazo de caché, conceptos fundamentales en la arquitectura de computadoras y sistemas operativos, como se describe en trabajos clásicos sobre sistemas de archivos y bases de datos.