ONNX:生成式AI运行时

有时一个小模型就能解决问题。而且你并不总是需要 GPU。有理由直接在 CPU 上运行"工具"任务。有时你根本无法拥有 GPU。或者你想保持数据本地化。或者只是想保持架构简单。

这就是 ONNX Runtime GenAI 发挥作用的地方。它使你能够在想要的地方运行模型:如果有 GPU 就用 GPU,没有就用 CPU。所有这些都不需要更改一行代码。在这篇文章中,我将展示它如何工作,包括在 Google Cloud Run 上托管你的模型。你可以在 onnx-inference 存储库中找到所有示例代码。

1、起源:可移植的 ML 模型

这一切始于 2017 年 9 月。行业领导者意识到 ML 工具的碎片化正在伤害每个人,因此 ONNX 作为一个共享标准启动了。其推广很简单:模型的"通用翻译器"。你可以在首选框架中训练,导出到通用标准,并在各种硬件目标上高效运行模型。

行业没有浪费时间。到年底,更多公司加入进来,2018 年微软开源了 ONNX Runtime,这个引擎旨在在任何硬件上高效运行这些模型。它于 2019 年毕业成为 Linux Foundation AI 项目,巩固了其作为开放标准的地位。

2、方向:可移植的 LLMs

当 LLM 爆发在 2023 年左右袭来时,ONNX 面临新的挑战。传统模型是无状态的:输入进去,预测出来。LLMs 不同。它们健谈。它们有记忆。它们逐个 token 生成文本,需要管理"KV 缓存"来记住对话上下文。

标准的 ONNX 运行时并不是为这个循环构建的。

所以 2024 年,社区发布了 onnxruntime-genai。它用 LLMs 需要的所有特定逻辑包装了核心运行时:分词、生成循环、搜索策略如束搜索,以及状态管理。

快进到 2026 年,我们现在在 Hugging Face 上有预量化 ONNX 模型库。你可以从货架上拉出模型并立即运行,无需任何训练或格式转换。

3、使用库

onnxruntime-genai 的好处是它为你处理生成。如果你尝试使用原始 ONNX Runtime 来做这件事,你需要编写手动循环将输出 token 作为输入反馈。

现在它看起来是这样的:

import onnxruntime_genai as og
# 加载模型(此路径自动适用于 CPU、GPU 或移动设备)
model = og.Model('path/to/model')
tokenizer = og.Tokenizer(model)
# 配置你想要的搜索方式
params = og.GeneratorParams(model)
params.set_search_options(max_length=256, batch_size=1)
# 生成循环
generator = og.Generator(model, params)
generator.append_tokens(tokenizer.encode(prompt))
while True:
    generator.generate_next_token()
    if generator.is_done():
        break

    # 边生成边解码和打印
    token = generator.get_next_tokens()[0]
    print(tokenizer.decode(token), end='', flush=True)

它在这里做了很多繁重的工作:处理 KV 缓存、应用你的搜索策略(贪婪、top-p 等),并将操作路由到最佳可用硬件(CUDA、CoreML 或 CPU)。

4、硬件、模型和量化

自从 LLMs 的早期以来,一些事情发生了变化。处理器变快了,模型在更小的情况下仍然具有令人难以置信的能力。

然后是量化。我们不再受限于以完整的 32 位精度运行模型。像 INT4 量化这样的技术可以大大压缩权重,而对精度的影响令人惊讶地小。

我用 onnx-inference 测试了很多模型。注意这些小模型能够识别结构,但不适合知识密集型任务。

对于非常简单的任务,你可以使用像 SmolLM2–135M 这样的微小模型。它非常适合基本完成或分类。

对于更复杂的任务,你会想要一个更大的模型。Qwen 3–0.6B 为多几亿个参数提供了更多的能力。

当你在 CPU 上考虑超过 500M 参数的模型时,你需要仔细监控上下文窗口、最大 token 和你可以容忍的延迟。

5、构建可移植服务器

将推理逻辑包装在轻量级服务器中使其普遍可访问。无论语言如何,都可以从任何服务轻松调用。我的存储库中的代码提供了一个 FastAPI 服务器,正是这样做的。

主类是 OnnxTextGenerator,它处理推理逻辑:

from inference import OnnxTextGenerator

# 自动检测硬件
generator = OnnxTextGenerator()

# 简单运行
result = generator.generate(
    prompt="Explain quantum computing like I'm five:",
    max_new_tokens=100,
    temperature=0.7
)

print(result['generated_text'])

对于实时应用程序,你不能等待整个答案。你可以流式传输它:

for chunk, metadata in generator.stream_generate(
    prompt="Write a haiku about Docker:",
    max_new_tokens=50,
    temperature=0.8
):
    print(chunk, end='', flush=True)

每个函数都作为端点暴露在 FastAPI 中:

@app.post("/generate")
async def generate(request: GenerateRequest):
    result = generator.generate(
        prompt=request.prompt,
        max_new_tokens=request.max_new_tokens,
        temperature=request.temperature
    )
    return {
        "generated_text": result["generated_text"],
        "tokens_generated": result["tokens_generated"],
        "finish_reason": result["finish_reason"]
    }

我想要一个"就能用"的东西,所以初始化器自动寻找最佳硬件执行提供程序:CUDA(NVIDIA GPU)→ CoreML(Apple Silicon)→ CPU(通用回退)。

6、容器化策略

对于像前面列出的那些较小模型,一个方便的模式是将模型直接烘焙到镜像中。这使得它们在部署后立即可用,而无需等待下载。

这是我使用的 Dockerfile 的简化视图:

FROM python:3.12-slim

COPY requirements.txt .
RUN pip install -r requirements.txt

ARG MODEL_ID=onnx-community/SmolLM2-135M-Instruct-ONNX
RUN hf download ${MODEL_ID} --local-dir /app/model

COPY . /app
WORKDIR /app

CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8080"]

当你加载此容器时,模型已经在磁盘上了。它立即启动,在没有互联网访问的情况下工作。

7、在 Google Cloud Run 上实现无服务器

Cloud Run 非常适合使用小模型的应用程序。当你不使用它时,它扩展到零,所以你停止为资源付费。由于我们使用 CPU,我们也不需要提供 GPU 实例。

7.1 从源代码部署

我们可以直接从源代码部署。以下命令使用 Google Cloud Build 构建容器,并在一步中将其部署到 Cloud Run。

我们故意在这里设置了一些特定标志。我们分配 2 个 CPU,因为推理是计算密集型的,以及 4Gi 内存 以舒适地容纳小模型和 KV 缓存。我们还将并发设置为 4,这允许实例处理一些同时请求而不会颠簸缓存。

gcloud run deploy onnx-inference \
    --allow-unauthenticated \
    --concurrency 4 \
    --cpu 2 \
    --labels dev-tutorial=onnx-inference \
    --memory 4Gi \
    --region us-central1 \
    --source .

7.2 测试

部署完成后,我们需要获取新服务的安全 URL。然后我们可以使用简单的 curl 命令进行测试。

SERVICE_URL=$(gcloud run services describe onnx-inference \
    --region $REGION \
    --format 'value(status.url)')

curl -X POST "$SERVICE_URL/generate" \
    -H "Content-Type: application/json" \
    -d '{"prompt": "Why is efficient AI important?", "max_new_tokens": 50}'

7.3 故障排除

一些事情可能会绊倒你:

  • 缺失 genai_config.json:并非所有 Hugging Face ONNX 模型都包含 GenAI 库的配置。我的库尝试在缺失时推断配置,但最好使用已有配置的模型。
  • 执行提供程序:示例当前包括 CUDA、CoreML 和 CPU,但添加其他提供程序(如 TensorRTOpenVINO)是直截了当的。
  • 参数:当你扩大 max_new_tokens 时,KV 缓存增长,注意力机制执行更多计算。密切关注内存使用和延迟。

8、结束语

较小的模型已经走了很长的路。通过 ONNX Runtime GenAI 和适度量化,你可以在几年前似乎不可能的地方运行有能力的 LLMs。

这开辟了一整套新的应用程序:完全私有的本地助手、智能边缘设备,以及几乎没有成本维护的无服务器 API。

如果你想尝试一下,大约需要 5 分钟开始。你可以从 GitHub 上的 onnx-inference 存储库 获取代码,查看官方 ONNX Runtime GenAI 文档 了解更深入的详细信息,或浏览 Hugging Face ONNX Community 找到你的下一个模型。


原文链接: Run LLMs anywhere: Local and CPU inference with ONNX Runtime GenAI

汇智网翻译整理,转载请标明出处