打包库和权重

当我们想要使用 MLC LLM 构建 LLM 应用程序(例如 iOS/Android 应用程序)时,通常需要构建静态模型库和应用程序绑定库,有时还需要将模型权重打包到应用程序中。MLC LLM 提供了工具用于快速打包模型库和权重:mlc_llm package

本页面简要介绍了如何使用 mlc_llm package 进行打包。教程 iOS Swift SDKAndroid SDK 包含了使用此打包工具进行 iOS 和 Android 部署的详细示例和说明。


简介

要使用 mlc_llm package,必须克隆 MLC LLM 的源代码并 安装 MLC LLM 和 TVM Unity 包。根据构建的应用程序,可能还需要一些其他依赖项,这些在相应的 iOSAndroid 教程中有描述。

克隆完成后,mlc_llm package 的基本用法如下。

export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm
cd /path/to/app  # The app root directory which contains "mlc-package-config.json".
                 # E.g., "ios/MLCChat" or "android/MLCChat"
mlc_llm package

打包命令从当前目录下的 JSON 文件 mlc-package-config.json 中读取配置。 该命令的输出是目录 dist/,其中包含打包好的模型库(位于 dist/lib/ 下)和权重(位于 dist/bundle/ 下)。该目录包含了应用程序构建所需的所有数据。根据构建的应用程序,dist/lib/ 的内部结构可能会有所不同。

dist
├── lib
│   └── ...
└── bundle
    └── ...

输入文件 mlc-package-config.json 指定了以下内容:

  • 要打包模型库和权重的设备(例如 iPhone 或 Android),

  • 要打包的模型列表。

以下是示例 mlc-package-config.json 文件:

{
    "device": "iphone",
    "model_list": [
        {
            "model": "HF://mlc-ai/Mistral-7B-Instruct-v0.2-q3f16_1-MLC",
            "model_id": "Mistral-7B-Instruct-v0.2-q3f16_1",
            "estimated_vram_bytes": 3316000000,
            "bundle_weight": true,
            "overrides": {
                "context_window_size": 512
            }
        },
        {
            "model": "HF://mlc-ai/gemma-2b-it-q4f16_1-MLC",
            "model_id": "gemma-2b-q4f16_1",
            "estimated_vram_bytes": 3000000000,
            "overrides": {
                "prefill_chunk_size": 128
            }
        }
    ]
}

这个示例 mlc-package-config.json 指定了目标设备为 "iphone"。在 model_list 中,

  • model 指向包含预转换模型权重的 Hugging Face 仓库。应用程序将从 Hugging Face URL 下载模型权重。

  • model_id 是唯一的模型标识符。

  • estimated_vram_bytes 是对模型在运行时占用的 vRAM 的估计。

  • "bundle_weight": true 表示在构建应用程序时,模型的权重将被打包到应用程序中。

  • overrides 指定了一些模型配置参数的覆盖值。

以下是 mlc-package-config.json 文件的更详细规范。JSON 文件中的 "model_list" 每个条目包含以下字段:

model

(必填)要构建到应用程序中的 MLC 转换模型的路径。

通常它是 Hugging Face URL(例如 "model": "HF://mlc-ai/phi-2-q4f16_1-MLC"),其中包含预转换的模型权重。对于 iOS,它也可以是一个包含转换后模型权重的本地模型目录路径(例如 "model": "../dist/gemma-2b-q4f16_1")。如果你想将本地模型构建到应用程序中,请查看 转换模型权重

model_id

(必填)用于标识模型的唯一本地标识符。它可以是任意的。

estimated_vram_bytes

(必填)运行模型所需的 vRAM 估计值。

bundle_weight

(可选)布尔标志,指示是否将模型权重打包到应用程序中。如果此字段设置为 true,mlc_llm package 命令将把模型权重复制到 dist/bundle/$model_id

overrides

(可选)字典,用于覆盖默认的模型上下文窗口大小(以限制 KV 缓存大小)和预填充块大小(以限制模型临时执行内存)。例如:

{
   "device": "iphone",
   "model_list": [
      {
            "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC",
            "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1",
            "estimated_vram_bytes": 2960000000,
            "overrides": {
               "context_window_size": 512,
               "prefill_chunk_size": 128
            }
      }
   ]
}
model_lib

(可选)字符串,指定用于模型的系统库前缀。通常在你想要将具有相同架构的多个模型变体构建到应用程序中时使用。此字段不会影响任何应用程序功能。 下面介绍的 "model_lib_path_for_prepare_libs" 也与此相关。例如:

{
   "device": "iphone",
   "model_list": [
      {
            "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC",
            "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1",
            "estimated_vram_bytes": 2960000000,
            "model_lib": "gpt_neox_q4f16_1"
      }
   ]
}

除了 MLCChat/mlc-package-config.json 中的 model_list 外,你还可以 选择性地 指定 "model_lib_path_for_prepare_libs" 字典,如果你想使用手动编译的模型库。此字典的键应为模型列表中指定的 model_lib,而值应为手动编译的模型库的路径(绝对路径或相对路径)。在运行 mlc_llm package 时,"model_lib_path_for_prepare_libs" 中指定的模型库将被构建到应用程序中。例如:

{
   "device": "iphone",
   "model_list": [
      {
            "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC",
            "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1",
            "estimated_vram_bytes": 2960000000,
            "model_lib": "gpt_neox_q4f16_1"
      }
   ],
   "model_lib_path_for_prepare_libs": {
      "gpt_neox_q4f16_1": "../../dist/lib/RedPajama-INCITE-Chat-3B-v1-q4f16_1-iphone.tar"
   }
}

编译缓存

mlc_llm package 利用本地 JIT 缓存来避免重复编译相同的输入。它还利用本地缓存从远程下载权重。这些缓存在整个项目中共享。有时,当有新的编译器更新或缓存库出现问题时,强制重新构建会很有帮助。你可以通过设置环境变量 MLC_JIT_POLICY=REDO 来实现这一点。

MLC_JIT_POLICY=REDO mlc_llm package

mlc_llm package 的参数

命令 mlc_llm package 可以选择性地接受以下参数:

--package-config

指向包含设备和模型规范的 mlc-package-config.json 的路径。默认情况下,它是当前目录下的 mlc-package-config.json

--mlc-llm-source-dir

MLC LLM 源代码的路径(从 https://github.com/mlc-ai/mlc-llm 克隆)。默认情况下,它是 $MLC_LLM_SOURCE_DIR 环境变量。如果既没有指定 $MLC_LLM_SOURCE_DIR 也没有指定 --mlc-llm-source-dir,则会报错。

--output / -o

mlc_llm package 命令的输出目录。默认情况下,它是当前目录下的 dist/

总结与下一步

在本页面中,介绍了用于快速打包模型库和权重的 mlc_llm package 命令。

  • 它接受输入文件 mlc-package-config.json,其中包含用于打包的设备和模型规范。

  • 它输出目录 dist/,其中包含打包好的库(位于 dist/lib/ 下)和模型权重(位于 dist/bundle/ 下)。

接下来,请随意查看 iOSAndroid 教程,了解使用 mlc_llm package 的详细示例。