浏览器跑大语言模型-WEBLLM

前言

陈天奇团队新发布了一个可将大型语言模型和基于 LLM 的聊天机器人引入 Web 浏览器的项目 —— Web LLM。一切都在浏览器内运行，无需服务器支持，并使用 WebGPU 加速。GitHub 地址：web-llm。

体验一下

根据官网上的操作说明，体验一下。

由于依赖于浏览器的 WebGPU 能力，需要安装最新的 Chrome 113 版本，在这之前的版本是不支持的。因此需要下载 Chrome Canary。

1.下载完成开发版浏览器之后启动浏览器；如果是苹果芯片的 mac，用以下命令启动浏览器：

/Applications/Google\ Chrome.app/Contents/MacOS/Google\ Chrome --enable-dawn-features=allow_unsafe_apis,disable_robustness

2.打开对应的 demo 地址：chat-demo。

3.选择模型，输入问题，点击“send”。此时，应用就会获取模型参数，并缓存在本地，这个下载的过程耗时比较长，大概需要几分钟，不过只有第一次需要，这个过程很有可能会失败，多试几次。之后的问答相对快一些（但也需要几十秒，在我 intel 芯片的 mac 上运行并不是很快。）。

web-llm 包

web-llm 是一个JavaScript 包，将聊天语言模型直接运行在具有硬件加速的 Web 浏览器上，依靠 WebGPU 加速，无需服务端支持。

直接引用

// 1.引入模块
import * as webllm from "@mlc-ai/web-llm";




async function main() {

	// 2.创建聊天模块
	const chat = new webllm.ChatModule();
  // 3.显示模块初始化进度
  chat.setInitProgressCallback((report: webllm.InitProgressReport) => {
    // setLabel("init-label", report.text);
  });

	// 4.加载模型
  await chat.reload("vicuna-v1-7b-q4f32_0");

  // 5.生成回答回调函数
  const generateProgressCallback = (_step: number, message: string) => {
    setLabel("generate-label", message);
  };


  // 6.输入问题，返回答案
  const prompt0 = "What is the capital of Canada?";
  // setLabel("prompt-label", prompt0);
  const reply0 = await chat.generate(prompt0, generateProgressCallback);
  console.log(reply0);

  // 7.输入问题，返回答案
  const prompt1 = "Can you write a poem about it?";
  // setLabel("prompt-label", prompt1);
  const reply1 = await chat.generate(prompt1, generateProgressCallback);
  console.log(reply1);

  console.log(await chat.runtimeStatsText());
}
main();

使用 Web Worker

web-llm 支持在 web worker 中运行，从而可以避免阻塞 UI 的渲染。

先在 worker 中创建 ChatModule，并使该模块绑定到 handler 上。

// worker.ts
import { ChatWorkerHandler, ChatModule } from "@mlc-ai/web-llm";



// Hookup a chat module to a worker handler
const chat = new ChatModule();
const handler = new ChatWorkerHandler(chat);
self.onmessage = (msg: MessageEvent) => {
  handler.onmessage(msg);
};

在主逻辑中，创建一个 ChatWorkerClient 对象，其余逻辑和之前保持一致。

// main.ts
import * as webllm from "@mlc-ai/web-llm";




async function main() {

  // Use a chat worker client instead of ChatModule here
  const chat = new webllm.ChatWorkerClient(new Worker(
    new URL('./worker.ts', import.meta.url),
    {type: 'module'}
  ));
  // everything else remains the same
}

自定义模型和参数

WebLLM 包中有两个关键参数变量，可以用来修改模型和参数：

• model_url：模型的参数及元数据
• model_lib：模型的 wasm 库，可以加速模型的计算。

async main() {
  const myLlamaUrl = "/url/to/my/llama";
  const appConfig = {
  "model_list": [
    {
      "model_url": myLlamaUrl,
      "local_id": "MyLlama-3b-v1-q4f32_0"
    }
  ],
  "model_lib_map": {
    "llama-v1-3b-q4f32_0": "/url/to/myllama3b.wasm",
  };
 
  const chatOpts = {
    "repetition_penalty": 1.01
  };

  const chat = new ChatModule();


  await chat.reload("MyLlama-3b-v1-q4f32_0", chatOpts, appConfig);
}

由于 web-llm 是 MLC LLM 项目的一部分，因此模型的创建依赖于 MLC LLM 的工作流。至于如何通过 MLC LLM 编译出上述的 modelUrl 和 modelLib，目前相关文档还没有给出，预计在未来几周内发布。

工作原理

机器学习编译

其核心关键技术是机器学习编译 (machine learning compilation, MLC) 。那么什么是机器学习编译呢？此处的编译与传统编译的含义有些区别，这个过程不一定涉及代码生成，只是过程类似，具体指将机器学习算法从开发阶段，通过变换和优化算法，使其变成部署状态。

• 开发形式 是指我们在开发机器学习模型时使用的形式。典型的开发形式包括用 PyTorch、TensorFlow 或 JAX 等通用框架编写的模型描述，以及与之相关的权重。
• 部署形式 是指执行机器学习应用程序所需的形式。它通常涉及机器学习模型的每个步骤的支撑代码、管理资源（例如内存）的控制器，以及与应用程序开发环境的接口（例如用于 android 应用程序的 java API）。

通常有多种方法可以在特定环境中实现模型执行。上图展示了一个例子。 值得注意的是，有两个区别： 首先，第一个 linear 层和 relu 计算被折叠成一个 linear_relu 函数；另外，这有一个特定的 linear_relu 的详细实现。当然，现实世界的用例，linear_relu 可以通过各种代码优化技术来实现。机器学习编译的过程就是将上图左侧的内容转换为右侧的过程。在不同的场景中，这个过程可以是手动完成的，也可以使用一些自动转换工具，或两者兼而有之。

一句话总结：开发形式偏向理论，部署形式偏向实际实现。

通过机器学习编译，可以自动优化模型的计算图、内存分配和运算操作，从而提高模型的执行效率和性能。

TVM (Tensor Virtual Machine）是机器学习编译的一个重要组成部分，它提供了一个完整的深度学习编译器和运行时系统，可用于将机器学习模型编译为高效的代码，并在不同的硬件平台上进行部署。

Apache TVM is an open source machine learning compiler framework for CPUs, GPUs, and machine learning accelerators. It aims to enable machine learning engineers to optimize and run computations efficiently on any hardware backend.

而 TVM 正是陈天奇团队的主要研究方向，他还有一门课程专门讲机器学习编译。

对比 TensorFlow.js/ONNX.js

TensorFlow.js 和 ONNX.js 一直在尝试把机器学习应用在浏览器中，但和原生相比，web 端仍然存在一些明显的性能差距。导致差距的原因在于 TensorFlow.js 和 ONNX.js 借助 WebGL 的 Vertex Shader(顶点着色器) 和 Fragment Shader(片元着色器)发挥 GPU 的计算能力。而 WebGPU 新增了 Compute Shaders(计算着色器) ，从而更高性能地利用 GPU 实现计算。

在实现原理上，web-llm 和 TensorFlow.js、ONNX.js 也有很大的差别。传统的方式需要通过手写着色器来实现各个算子，这一部分的实现细节，我之前分享过，可以参见《WebGL 为什么可以被用来做计算》。

而 Web LLM，利用 TVM 将高级框架的模型，例如 Tensorflow，Keras，Pytorch，MXNET 和 ONNX，自动生成低级代码，在本案例中，以 SPIR-V 格式表示计算着色器。生成的低级代码可以被打包成一个部署模块。为了充分发布 WebGPU 的潜力，TVM 针对 WASM 和 WebGPU 做了专门的编译优化。

先在 TVM 中创建了一个语言模型的 IRModule，以此来减少了计算量和内存使用。IRModule 是一个中间表示模块，它包含了深度学习模型的计算图、参数和优化策略等信息。

TVM 的 IRModule 中的每个函数都可以被进一步转换并生成可运行的代码，这些代码可以被普遍部署在任何最小 TVM 运行时支持的环境中（JavaScript 就是其中之一）。

其次，TensorIR 是生成优化程序的关键技术，通过结合专家知识和自动调度程序快速转换 TensorIR 程序，来提供高效的解决方案。使用张量程序抽象的主要目的是表示循环和相关的硬件加速选择，如多线程、特殊硬件指令的使用和内存访问。

经过 TVM 的编译，将模型转换成：

• Host code: 主程序，计算启动参数并启动设备。
• SPIRV shaders: 着色器程序，用于模型计算。是一种开放标准的着色器中间表示，最终转换为可以在GPU上执行的二进制代码。
• tvm web runtime: 运行时。可以运行 host code 和 SPIRV shaders。

将上面三部分打包到一个 wasm 文件，但此时还有两个问题需要解决：

1. 如何访问系统的服务和资源？
2. 如何和 WebGPU 驱动进行交互？

问题 1 的标准解法是 WASI，但目前在web上没有成熟的 WASI，只能通过 emscripten 生成 WASI-like 库来和系统交互。

问题 2 的解法是将 WebGPU runtime 打到 tvmjs runtime 中，在 wasm 中通过回调的方式触发 GPU 代码。

其中，tvmjs runtime 这个库用来和 wasm 进行交互，其产生步骤：

• 进入到 web 目录：github.com/mlc-ai/rela…
• 执行打包编译命令：npm run bundle
• 生成 tvmjs library

性能对比

最后看下官方提供的性能评测。以模型 MobileNet 为例，Chrome WebGPU 下的性能和 Native Metal （MacOS）相差并不是很大。

未来规划

• 通过 WASI 使得 WebGPU 的 native API 成为标准，从而 WASM 可以直接调用 WebGPU。
• TVM 社区在积极推动 Rust-based runtime。
• …

名词解释

LLM 大语言模型：LLM（Large Language Model）大语言模型是指一种基于深度学习技术的自然语言处理模型，能够处理大量的文本数据。它可以预测给定上下文中的下一个单词或字符的概率，从而实现文本生成、机器翻译、语音识别、问答系统等任务。LLM模型通常采用循环神经网络（RNN）或变换器（Transformer）等架构，通过在大规模语料库上训练模型来学习语言的结构和规律。目前，一些知名的LLM模型包括GPT-3、BERT、XLNet等。

Hugging Face：是一家人工智能软件公司，致力于为自然语言处理领域开发各种工具和平台。该公司最著名的产品是Transformers，这是一个面向自然语言处理的机器学习平台，可以轻松构建、训练和部署各种深度学习模型。Hugging Face还提供了一系列开源的自然语言处理工具和库，如PyTorch-Transformers、Tokenizers、Datasets等，这些工具广泛应用于文本分类、情感分析、机器翻译、对话系统等领域。除此之外，Hugging Face还开发了一款名为Hugging Face Hub的平台，可以方便地共享和发现各种自然语言处理模型和数据集。

Apache TVM（Apache Tensor Virtual Machine） ：是一个开源的深度学习编译器和运行时系统，可以将深度学习模型优化、编译和部署到各种硬件平台上。它支持多种深度学习框架，如TensorFlow、PyTorch、MXNet等，并且可以自动优化模型的计算图、内存分配和运算操作，从而提高模型的执行效率和性能。Apache TVM同时也支持多种硬件平台，包括CPU、GPU、FPGA、ASIC等，可以在不同的硬件平台上部署相同的深度学习模型。此外，Apache TVM还提供了一系列 API 和工具，方便用户进行模型调试、性能分析和部署。Apache TVM 已经成为深度学习编译器和运行时领域的重要开源项目之一，并且被广泛应用于工业界和学术界的深度学习应用中。

IRModule：在TVM中，IRModule（Intermediate Representation Module）是一个中间表示模块，用于组织、管理和优化深度学习模型的中间表示。IRModule 是 TVM 中的一个重要概念，它包含了深度学习模型的计算图、参数和优化策略等信息，可以通过TVM IR语言来描述和表示。IRModule可以在不同的模块之间进行组合和复用，从而实现模型的灵活组合和优化。

具体来说，IRModule主要包括以下几个部分：

1.全局函数定义（GlobalFunctionDef）：描述了模型的计算图结构和参数信息。

2.全局变量定义（GlobalVarDef）：描述了模型的中间结果和状态信息。

3.优化策略（Passes）：描述了模型的优化策略和优化阶段。

通过IRModule，可以进行模型的优化、编译和部署，从而实现对深度学习模型的高效执行和推理。同时，IRModule也是TVM中实现模型组合和复用的重要手段，可以将多个模型组合成一个复合模型，从而实现更加复杂的深度学习任务。

张量程序抽象（Tensor Program Abstraction，TPA） ：是一个用于描述和优化深度学习模型的高层次抽象。TPA将深度学习模型抽象为一组张量计算操作，以及这些操作之间的依赖关系和数据流向。通过对TPA的分析和优化，可以实现对深度学习模型的高效执行和推理。

具体来说，TPA包括以下几个方面：

1.张量计算操作：TPA将深度学习模型中的各种计算操作（如卷积、池化、全连接等）抽象为张量计算操作，以及相关的数据输入和输出张量。这些计算操作可以用数学表达式或图形模型来表示和描述。

2.计算图：TPA将计算操作和数据张量组合成一个计算图，以描述它们之间的依赖关系和数据流向。计算图可以用有向无环图（DAG）来表示和优化。

3.优化策略：TPA还包括一些优化策略，如张量融合、内存分配、计算图剪枝等，用于提高深度学习模型的执行效率和性能。

通过TPA的抽象和优化，可以将深度学习模型转化为高效的张量计算操作序列，从而实现对模型的快速执行和推理。TPA在深度学习编译器和运行时系统中发挥着重要的作用，如Apache TVM等。

SPIRV（可编程着色器中间表示） ：是一种开放标准的着色器中间表示，用于将高级语言编写的图形着色器（如GLSL和HLSL）转换为可以在GPU上执行的二进制代码。SPIRV可以在不同的GPU厂商之间实现可移植性，并且支持调试和优化。SPIRV还可以支持新兴的GPU编程语言，如Vulkan和OpenCL。

WASI（WebAssembly System Interface） ：是一种标准接口，用于在WebAssembly（Wasm）运行时中访问操作系统服务和资源。它为Wasm提供了一种可移植的方式，以访问文件系统、网络、随机数生成器和其他系统级资源。WASI 的目标是在不牺牲安全性、可移植性和性能的情况下，将Wasm推向更广泛的场景，例如云、物联网和边缘计算。

参考

浏览器就能跑大模型了，陈天奇团队发布 WebLLM，无需服务器支持

simple-chat

What is ML Compilation

WebGPU powered machine learning in the browser with Apache TVM

Compiling Machine Learning to WASM and WebGPU with Apache TVM

简单了解一下 SPIR-V

tvm web