通过新的 Olive CLI 实现 AI 模型优化的普及化

作者

Jambay Kinley, Hitesh Shah, Xiaoyu Zhang, Devang Patel, Sam Kemp

2024 年 11 月 11 日

👋 简介

Build 2023 大会上,微软发布了 Olive (ONNX Live):一个先进的模型优化工具包,旨在简化使用 ONNX 运行时优化用于部署的 AI 模型的过程。 正如以下图表所示,Olive 可以从 PyTorch 或 Hugging Face 等框架中获取模型,并输出针对特定部署目标量身定制的优化 ONNX 模型。

Olive workflow. Olive 高级工作流程。 这些硬件目标可以包括主要硬件供应商(如 Qualcomm、AMD、Nvidia 和 Intel)提供的各种 AI 加速器(GPU、CPU)

Olive 通过结构化的工作流程运行,该工作流程由一系列模型优化任务组成,这些任务称为passes(步骤)。 这些步骤可以包括模型压缩、图捕获、量化和图优化。 每个步骤都有可调整的参数,可以对其进行调整以实现最佳指标,例如准确性和延迟,这些指标由相应的评估器进行评估。 该工具利用搜索策略,采用算法来自动调整单个步骤或成组的步骤,从而确保部署目标的最佳性能。

虽然 Olive 中使用的工作流程范例非常灵活,但对于不熟悉模型优化流程的 AI 开发人员来说,学习曲线可能具有挑战性。 为了使模型优化更易于上手,我们为常见场景策划了一组 Olive 工作流程,并将它们作为 新的易于使用的 Olive CLI 中的一个简单命令公开。

Olive Commands. 新的 Olive CLI 命令与执行的相关 Olive 工作流程的映射。

在本博客中,我们将向您展示如何使用 Olive CLI 为 ONNX Runtime 准备模型。

🚀 Olive CLI 入门

首先,使用 pip 安装 Olive

pip install olive-ai[cpu,finetune]

🪄 自动优化器

安装 Olive 后,尝试自动优化器 (olive auto-opt)。 在单个命令中,Olive 将会

  1. 从 Hugging Face 下载模型
  2. 将模型结构捕获到 ONNX 图中,并将权重转换为 ONNX 格式。
  3. 优化 ONNX 图(例如,融合)
  4. 将模型权重量化为 int4

在 CPU 设备上为 Llama-3.2-1B-Instruct 模型运行自动优化器的命令是

olive auto-opt \
    --model_name_or_path meta-llama/Llama-3.2-1B-Instruct \
    --trust_remote_code \ 
    --output_path optimized-model \
    --device cpu \
    --provider CPUExecutionProvider \
    --precision int4 \
    --use_model_builder True \
    --log_level 1

提示: 如果要以以下目标为目标

  • CUDA GPU,然后将 --device 更新为 gpu,并将 --provider 更新为 CUDAExecutionProvider
  • Windows DirectML,然后将 --device 更新为 gpu,并将 --provider 更新为 DmlExecutionProvider

Olive 将应用特定于设备和提供程序的优化。

使用 auto-opt 命令,您可以将输入模型更改为 Hugging Face 上可用的模型 - 例如,HuggingFaceTB/SmolLM-360M-Instruct - 或本地磁盘上的模型。 应该注意的是,olive auto-opt 中的 --trust_remote_code 参数仅对于 Hugging Face 中需要在您的机器上运行代码的自定义模型是必需的 - 有关更多详细信息,请阅读 Hugging Face 关于 trust_remote_code 的文档。 Olive 将经历相同的自动转换(为 ONNX)、优化图和量化权重的过程。

🧪 尝试不同的量化算法

Olive CLI 允许您尝试许多不同的量化算法 - 例如 AWQ、GPTQ 和 QuaRot - 以及这些算法的不同实现。 例如,要使用 激活感知量化 (AWQ) 量化 Llama-3.2-1B-Instruct

注意: 您的计算机将需要安装 CUDA GPU 设备和相关驱动程序才能运行 AWQ、GPTQ 和 QuaRot 量化。 此外,您应该使用以下命令安装 AutoAWQ 包

pip install autoawq

olive quantize \
    --model_name_or_path meta-llama/Llama-3.2-1B-Instruct \
    --algorithm awq \
    --output_path quantized-model \
    --log_level 1

使用 AWQ 方法时,quantize 命令将输出 PyTorch 模型,如果您打算在 ONNX Runtime 上使用该模型,则可以使用以下命令将其转换为 ONNX

olive capture-onnx-graph \
    --model_name_or_path quantized-model/model \
    --use_ort_genai True \
    --log_level 1 \

🎚️ 微调

Olive CLI 还提供了使用 LoRA 或 QLoRA 在我们自己的数据上为特定任务微调 AI 模型的工具。 以下示例将微调 Llama-3.2-1B-Instruct 以进行短语分类(给定一个英语短语,它将从 joy/sad/fear/surprised 中输出该短语的类别)。

olive finetune \
    --model_name_or_path meta-llama/Llama-3.2-1B-Instruct \
    --trust_remote_code \
    --output_path models/llama3.2/ft \
    --data_name xxyyzzz/phrase_classification \
    --text_template "<|start_header_id|>user<|end_header_id|>\n{phrase}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n{tone}" \
    --method qlora \
    --max_steps 30 \
    --log_level 1 \

finetune 命令将输出 Hugging Face PEFT 适配器,您可以使用以下命令为 ONNX 运行时准备该适配器

# Step 1 - capture the ONNX graph of the base model and adapter
olive capture-onnx-graph \
    --model_name_or_path models/llama3.2/ft/model \
    --adapter_path models/llama3.2/ft/adapter \
    --use_ort_genai \
    --output_path models/llama3.2/onnx \
    --log_level 1

# Step 2 - Extract adapter weights from ONNX model and store in separate file for ORT
 olive generate-adapter \
    --model_name_or_path models/llama3.2/onnx \
    --output_path adapter-onnx \
    --log_level 1       

🤝 使用 ONNX Runtime 的 Generate API 推理您优化的 AI 模型

以下 Python 代码创建了一个简单的基于控制台的聊天界面,该界面使用 ONNX 运行时的 Generate API 推理您优化的模型。

提示: 其他语言绑定 - 例如 C#、C/C++、Java - 更多语言绑定即将推出。 有关最新列表,请访问 ONNX Runtime Github 页面的 Generate API

import onnxruntime_genai as og
import numpy as np
import os

model_folder = "optimized-model/model"

# Load the base model and tokenizer
model = og.Model(model_folder)
tokenizer = og.Tokenizer(model)
tokenizer_stream = tokenizer.create_stream()

# Set the max length to something sensible by default,
# since otherwise it will be set to the entire context length
search_options = {}
search_options['max_length'] = 200
search_options['past_present_share_buffer'] = False

chat_template = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>

You are a helpful assistant<|eot_id|><|start_header_id|>user<|end_header_id|>

{input}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
""" 

text = input("Input: ")

# Keep asking for input phrases
while text != "exit":
    if not text:
        print("Error, input cannot be empty")
        exit

    # generate prompt (prompt template + input)
    prompt = f'{chat_template.format(input=text)}'

    # encode the prompt using the tokenizer
    input_tokens = tokenizer.encode(prompt)

    params = og.GeneratorParams(model)
    params.set_search_options(**search_options)
    params.input_ids = input_tokens
    generator = og.Generator(model, params)

    print("Output: ", end='', flush=True)
    # stream the output
    try:
        while not generator.is_done():
            generator.compute_logits()
            generator.generate_next_token()

            new_token = generator.get_next_tokens()[0]
            print(tokenizer_stream.decode(new_token), end='', flush=True)
    except KeyboardInterrupt:
        print("  --control+c pressed, aborting generation--")

    print()
    text = input("Input: ")

结论

在本博客中,我们演示了如何使用新的 Olive CLI 为 ONNX Runtime 组合模型,然后使用 ONNX Runtime 的 Generate API 推理这些模型。 Olive CLI 命令为您执行策划的 Olive 工作流程,这意味着您将继续获得以下所有好处

  • 减少挫败感和时间,避免通过反复试验手动尝试不同的图优化、压缩和量化技术。 定义您的质量和性能约束,让 Olive 自动为您找到最佳模型。
  • 40 多个内置模型优化组件,涵盖量化、压缩、图优化和微调方面的尖端技术。
  • 支持创建模型,以便可以使用 Multi LoRA 范例来服务它们。
  • Hugging FaceAzure AI 集成。
  • 内置 缓存 机制,以节省成本并 增强团队协作。 正如我们在之前的博客文章中分享的那样,Olive 还支持 共享缓存