这篇文章介绍了一个基于 Rust 和 WebAssembly (Wasm) 的解决方案,用于在异构边缘计算设备上快速和便携地进行 Llama2 模型的推理。与 Python 相比,这种 Rust+Wasm 应用程序的体积仅为 Python 的 1/100,速度提升 100 倍,并且可以在全硬件加速环境中安全运行,不需要更改二进制代码。文章基于 Georgi Gerganov 创建的 llama.cpp 项目,将原始的 C++ 程序适配到 Wasm 上。安装过程包括安装 WasmEdge 和 GGML 插件,下载预构建的 Wasm 应用和模型,然后使用 WasmEdge 运行 Wasm 推理应用,并传递 GGUF 格式的模型文件。此外,文章还提供了多个命令行选项,用于配置与模型的交互方式。
原文链接:https://www.secondstate.io/articles/fast-llm-inference/
curl -sSf https://raw.githubusercontent.com/WasmEdge/WasmEdge/master/utils/install.sh | bash -s -- --plugin wasi_nn-ggml
步骤2. 下载预构建的 Wasm 应用和模型
curl -LO https://github.com/second-state/llama-utils/raw/main/chat/llama-chat.wasm
你还应该下载一个 GGUF 格式的 llama2 模型。下面的例子下载了一个调整为5位权重的 llama2 7B聊天模型。
curl -LO https://huggingface.co/wasmedge/llama2/resolve/main/llama-2-7b-chat-q5_k_m.gguf
步骤3. 运行使用 WasmEdge 运行 wasm 推理应用程序,同时加载 GGUF 模型。现在,你可以输入问题与模型进行聊天了。
wasmedge --dir .:. --nn-preload default:GGML:AUTO:llama-2-7b-chat-q5_k_m.gguf llama-chat.wasm
配置模型行为你可以使用命令行选项配置与模型的交互方式。
Options:
-m, --model-alias <ALIAS>
Model alias [default: default]
-c, --ctx-size <CTX_SIZE>
Size of the prompt context [default: 4096]
-n, --n-predict <N_PRDICT>
Number of tokens to predict [default: 1024]
-g, --n-gpu-layers <N_GPU_LAYERS>
Number of layers to run on the GPU [default: 100]
-b, --batch-size <BATCH_SIZE>
Batch size for prompt processing [default: 4096]
-r, --reverse-prompt <REVERSE_PROMPT>
Halt generation at PROMPT, return control.
-s, --system-prompt <SYSTEM_PROMPT>
System prompt message string [default: "[Default system message
for the prompt template]"]
-p, --prompt-template <TEMPLATE>
Prompt template. [default: llama-2-chat] [possible values: llama-2-chat, codellama-instruct, mistral-instruct-v0.1, mistrallite, openchat, belle-llama-2-chat, vicuna-chat, chatml]
--log-prompts
Print prompt strings to stdout
--log-stat
Print statistics to stdout
--log-all
Print all log information to stdout
--stream-stdout
Print the output to stdout in the streaming way
-h, --help
Print help
例如,以下命令指定了 2048 个 token 的上下文长度和每次响应的最大 512 个 token 。它还告诉 WasmEdge 以流式方式将模型响应逐个 token 返回到stdout。该程序在低端 M2 MacBook 上每秒生成大约 25 个 token。
wasmedge --dir .:. --nn-preload default:GGML:AUTO:llama-2-7b-chat-q5_k_m.gguf \
llama-chat.wasm -c 2048 -n 512 --log-stat --stream-stdout
[USER]:
Who is the "father of the atomic bomb"?(谁是“原子弹之父”?)
---------------- [LOG: STATISTICS] -----------------
llama_new_context_with_model: n_ctx = 2048
llama_new_context_with_model: freq_base = 10000.0
llama_new_context_with_model: freq_scale = 1
llama_new_context_with_model: kv self size = 1024.00 MB
llama_new_context_with_model: compute buffer total size = 630.14 MB
llama_new_context_with_model: max tensor size = 102.54 MB
[2023-11-10 17:52:12.768] [info] [WASI-NN] GGML backend: llama_system_info: AVX = 0 | AVX2 = 0 | AVX512 = 0 | AVX512_VBMI = 0 | AVX512_VNNI = 0 | FMA = 0 | NEON = 1 | ARM_FMA = 1 | F16C = 0 | FP16_VA = 1 | WASM_SIMD = 0 | BLAS = 0 | SSE3 = 0 | SSSE3 = 0 | VSX = 0 |
The "father of the atomic bomb" is a term commonly associated with physicist J. Robert Oppenheimer. Oppenheimer was the director of the Manhattan Project, the secret research and development project that produced the atomic bomb during World War II. He is widely recognized as the leading figure in the development of the atomic bomb and is often referred to as the "father of the atomic bomb."
llama_print_timings: load time = 15643.70 ms
llama_print_timings: sample time = 2.60 ms / 83 runs ( 0.03 ms per token, 31886.29 tokens per second)
llama_print_timings: prompt eval time = 7836.72 ms / 54 tokens ( 145.12 ms per token, 6.89 tokens per second)
llama_print_timings: eval time = 3198.24 ms / 82 runs ( 39.00 ms per token, 25.64 tokens per second)
llama_print_timings: total time = 18852.93 ms
----------------------------------------------------
下面是它在 Nvidia A10G 机器上以每秒 50 个 token 的速度运行示例。
wasmedge --dir .:. --nn-preload default:GGML:AUTO:llama-2-7b-chat-q5_k_m.gguf \
llama-chat.wasm -c 2048 -n 512 --log-stat
[USER]:
Who is the "father of the atomic bomb"? (谁是“原子弹之父”?)
---------------- [LOG: STATISTICS] -----------------
llm_load_tensors: using CUDA for GPU acceleration
llm_load_tensors: mem required = 86.04 MB
llm_load_tensors: offloading 32 repeating layers to GPU
llm_load_tensors: offloading non-repeating layers to GPU
llm_load_tensors: offloaded 35/35 layers to GPU
llm_load_tensors: VRAM used: 4474.93 MB
..................................................................................................
llama_new_context_with_model: n_ctx = 2048
llama_new_context_with_model: freq_base = 10000.0
llama_new_context_with_model: freq_scale = 1
llama_kv_cache_init: offloading v cache to GPU
llama_kv_cache_init: offloading k cache to GPU
llama_kv_cache_init: VRAM kv self = 1024.00 MB
llama_new_context_with_model: kv self size = 1024.00 MB
llama_new_context_with_model: compute buffer total size = 630.14 MB
llama_new_context_with_model: VRAM scratch buffer: 624.02 MB
llama_new_context_with_model: total VRAM used: 6122.95 MB (model: 4474.93 MB, context: 1648.02 MB)
[2023-11-11 00:02:22.402] [info] [WASI-NN] GGML backend: llama_system_info: AVX = 1 | AVX2 = 1 | AVX512 = 0 | AVX512_VBMI = 0 | AVX512_VNNI = 0 | FMA = 1 | NEON = 0 | ARM_FMA = 0 | F16C = 1 | FP16_VA = 0 | WASM_SIMD = 0 | BLAS = 1 | SSE3 = 1 | SSSE3 = 1 | VSX = 0 |
llama_print_timings: load time = 2601.44 ms
llama_print_timings: sample time = 2.63 ms / 84 runs ( 0.03 ms per token, 31987.81 tokens per second)
llama_print_timings: prompt eval time = 203.90 ms / 54 tokens ( 3.78 ms per token, 264.84 tokens per second)
llama_print_timings: eval time = 1641.84 ms / 83 runs ( 19.78 ms per token, 50.55 tokens per second)
llama_print_timings: total time = 4254.95 ms
----------------------------------------------------
[ASSISTANT]:
The "father of the atomic bomb" is a term commonly associated with physicist J. Robert Oppenheimer. Oppenheimer was the director of the Manhattan Project, the secret research and development project that produced the first atomic bomb during World War II. He is widely recognized as the leading figure in the development of the atomic bomb and is often referred to as the "father of the atomic bomb."
(“原子弹之父”这个词通常与物理学家 J. Robert Oppenheimer 联系在一起。奥本海默是曼哈顿计划的主管,曼哈顿计划是二战期间研发出第一颗原子弹的秘密研发项目。他被广泛认为是原子弹研发的领导人物,常被称为“原子弹之父”。)
LLM 代理和应用我们利用 Rust 和 WasmEdge 技术,构建了一个与 OpenAI 兼容的 API 服务。它允许你使用任何兼容 OpenAI 的开发工具,如 flows.network,来创建 LLM 代理和应用。注:原始文档均为英文,括号内的中文为翻译内容。
边缘上的 Llama,图片由 Midjourney 生成
为什么不使用 Python?LLM(大型语言模型)如 llama2 通常在 Python(例如 PyTorch、Tensorflow 和 JAX)中进行训练。在 AI 计算的推理应用中,占比约 95%,Python 并不适合。LLM 工具链中常用的 Python 包相互冲突
Chris Lattner(参与 LLVM、Tensorflow 和 Swift 语言开发的著名软件工程师)在 This Week in Startup 播客上进行了精彩的采访。他讨论了 Python 在模型训练中的优势,但并不适合推理。 Rust+Wasm 的优势Rust+Wasm 技术栈构建了统一的云计算基础设施。它涵盖了从设备、边缘云到本地服务器和公有云的全方位服务。它在 AI 推理应用中是 Python 技术栈的有效替代方案。此外,埃隆·马斯克曾评价 Rust 为 AGI(通用人工智能)的理想语言。fn main() {
let args: Vec<String> = env::args().collect();
let model_name: &str = &args[1];
let graph =
wasi_nn::GraphBuilder::new(wasi_nn::GraphEncoding::Ggml, wasi_nn::ExecutionTarget::AUTO)
.build_from_cache(model_name)
.unwrap();
let mut context = graph.init_execution_context().unwrap();
let system_prompt = String::from("<<SYS>>You are a helpful, respectful and honest assistant. Always answer as short as possible, while being safe. <</SYS>>");
let mut saved_prompt = String::new();
loop {
println!("Question:");
let input = read_input();
if saved_prompt == "" {
saved_prompt = format!("[INST] {} {} [/INST]", system_prompt, input.trim());
} else {
saved_prompt = format!("{} [INST] {} [/INST]", saved_prompt, input.trim());
}
// 将用户问题字符串预处理成张量格式后,设置为模型输入张量,以供模型进行下游推理计算。
let tensor_data = saved_prompt.as_bytes().to_vec();
context
.set_input(0, wasi_nn::TensorType::U8, &[1], &tensor_data)
.unwrap();
// 执行推理计算
context.compute().unwrap();
// 获取输出
let mut output_buffer = vec![0u8; 1000];
let output_size = context.get_output(0, &mut output_buffer).unwrap();
let output = String::from_utf8_lossy(&output_buffer[..output_size]).to_string();
println!("Answer:\n{}", output.trim());
saved_prompt = format!("{} {} ", saved_prompt, output.trim());
}
}
要自行构建应用,只需安装 Rust 编译器及其 wasm32-wasi 编译目标。
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
rustup target add wasm32-wasi
然后,检出源项目,并运行cargo命令从 Rust 源项目构建 Wasm 文件。
在云端或边缘运行获取 Wasm 字节码文件后,你便可以在任何支持 WasmEdge 运行时的设备上进行部署。你只需要安装带有 GGML 插件的 WasmEdge。我们目前提供的 GGML 插件支持包括通用 Linux 和 Ubuntu Linux 在内的多种操作系统,适用于 x86 和 ARM CPU、Nvidia GPU,以及 Apple M1/M2/M3。基于 llama.cpp,WasmEdge GGML 插件将自动利用设备上的任何硬件加速来运行 llama2模 型。例如,如果你的设备有 Nvidia GPU,安装程序将自动安装优化了 CUDA 的 GGML 插件版本。对于 Mac 设备,我们专门为 Mac OS 构建了 GGML 插件,它利用 Metal API 在 M1/M2/M3 内置的神经处理引擎上执行推理工作负载。Linux CPU 构建的GGML 插件使用 OpenBLAS 库来自动检测和利用现代 CPU 上的高级计算特性,如 AVX 和 SIMD。通过使用 Rust 和 Wasm 技术,我们可以实现 AI 模型在异构硬件和平台上的可移植性,同时又不损失性能。 下一步计划虽然 WasmEdge GGML 工具目前已经投入使用,许多云原生客户都在使用,但它仍处于发展的初期阶段。如果你有兴趣贡献于开源项目,并对影响未来的大型语言模型(LLM)推理基础设施的发展方向产生影响,并为开源项目共享力量。# 克隆源项目
git clone https://github.com/second-state/llama-utils
cd llama-utils/chat/
# 构建
cargo build --target wasm32-wasi --release
# 结果 wasm 文件
cp target/wasm32-wasi/release/llama-chat.wasm .
推荐阅读: