一個簡單模型就讓ChatGLM性能大幅提升
本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。
引言
自大語言模型 (LLM) 成為熱點話題以來，涌現了一大批中文大語言模型并在優化平臺中得到了積極部署。ChatGLM 正是廣受好評的主流中文大語言模型之一。
然而，由于 ChatGLM 模型尚未成為 Transformer 生態的原生模型，因此，官方 optimum 擴展庫對其仍缺乏支持。
本文提供了一種使用 OpenVINO? opset 重構該模型架構的便捷方法。
該方案包含專為 ChatGLM 定制的優化節點，且這些節點都利用英特爾? 高級矩陣擴展（Intel? Advanced Matrix Extensions，縮寫為英特爾? AMX）內聯和 MHA（Multi-Head Attention，多頭注意力）融合實現了高度優化。
請注意，本文僅介紹了通過為 ChatGLM 創建 OpenVINO? stateful模型實現優化的解決方案。本方案受平臺限制，必須使用內置了英特爾? AMX 的第四代英特爾? 至強? 可擴展處理器[1]（代號 Sapphire Rapids）。筆者不承諾對該解決方案進行任何維護。
ChatGLM 模型簡介
筆者在查看 ChatGLM 原始模型的源碼[2]時，發現 ChatGLM 與 Optimum ModelForCasualML并不兼容，而是定義了新的類 ChatGLMForConditionalGeneration[3]。
該模型的流水線回路包含 3 個主要模塊（Embedding、GLMBlock 層[4]和 lm_logits），結構如下：
△圖1 ChatGLM 模型結構
如上圖所示，整個流水線實際要求模型有兩個不同的執行圖，使用輸入提示符進行首次推理時不需要 KV 緩存作為 GLMBlock 層的輸入。從第二次迭代開始，QKV 注意力機制的上一次結果將成為當前一輪模型推理的輸入。
隨著生成符的長度不斷增加，在流水線推理過程中，模型輸入和輸出之間將存留大量的大型內存副本。
以 ChatGLM6b 默認模型配置[5]為示例，輸入和輸出陣列之間的內存副本類似于以下偽代碼，其內存拷貝的開銷由模型的參數 hidden_size 以及迭代的次數決定：
while(eos_token_id || max_seq_len){
memcpy(model_inp, model_outp, num_layer*2*sizeof(model_outp)* hidden_size)
model_outp.push_back(gen_token)
}△代碼若顯示不全，可左右滑動因此，本文要解決的兩大關鍵問題是：
- 如何優化模型推理流水線來消除模型輸入和輸出之間的內存副本
- 如何通過重新設計執行圖來優化 GLMBlock 模塊
構建 OpenVINO? stateful 模型實現顯著優化
首先，需要分析 GLMBlock 層的結構，嘗試封裝一個類并按以下工作流來調用 OpenVINO? opset。接著，將圖形數據序列化為 IR 模型 (.xml, .bin)。
△圖2 ChatGLM構建OpenVINO? stateful模型
關于如何構建 OpenVINO? stateful模型，以及如何使用OpenVINO? 提供的模型創建樣本，在 opset 構建模型，可參考文末文檔。
ChatGLM 的自定義注意力機制是本文所關注和優化的部分。
主要思路是：構建全局上下文結構體，用于在模型內部追加并保存每一輪迭代后的 pastKV 的結果，這樣減少了 pastKV 作為模型輸入輸出的拷貝開銷，同時使用內聯優化以實現 Rotary Embedding 和多頭注意力機制 (Multi-Head Attentions)。
英特爾? AMX 是內置在第四代英特爾? 至強? 可擴展處理器中的矩陣乘法加速器，能夠更快速地處理 bf16 或 int8 數據類型的矩陣乘加運算，通過加速張量處理，顯著提高推理和訓練性能。借助英特爾? AMX 內聯指令（用于加速計算的單指令多操作），實現了對 ChatGLM 模型中 Attention，Rotary Embedding 等算子的高度優化，并且使用 bf16 指令進行乘加操作，在保證浮點指數位精度的同時提高運算效率。
與此同時，本方案還使用 int8 精度來壓縮全連接層的權重，在實時計算中將使用bf16進行計算。因此，無需通過訓練后量化 (PTQ) 或量化感知訓練 (QAT) 對模型進行低精度處理。模型壓縮方法可以降低模型存儲空間，減少內存帶寬的負載，因為計算仍然使用浮點，不會造成溢出，不會對模型精度造成損失。
為 ChatGLM 創建 OpenVINO? stateful模型
請依照下方示例配置軟硬件環境，并按照以下步驟優化 ChatGLM：
硬件要求
第四代英特爾? 至強? 可擴展處理器（代號 Sapphire Rapids）或其后續的、仍內置英特爾? AMX 的產品
軟件驗證環境
Ubuntu 22.04.1 LTS
面向 OpenVINO? Runtime Python API 的 Python 3.10.11
用于構建 OpenVINO? Runtime 的 GCC 11.3.0
cmake 3.26.4
構建 OpenVINO? 源碼
- 安裝系統依賴并設置環境
- 創建并啟用 Python 虛擬環境
$ conda create -n ov_py310 pythnotallow=3.10 -y
$ conda activate ov_py310
△代碼若顯示不全，可左右滑動- 安裝 Python 依賴
$ pip install protobuf transformers==4.30.2 cpm_kernels torch>=2.0 sentencepiece pandas△代碼若顯示不全，可左右滑動- 使用 GCC 11.3.0 編譯 OpenVINO?
- 克隆 OpenVINO? 并升級子模塊
$ git clone https://github.com/luo-cheng2021/openvino.git -b luocheng/chatglm_custom
$ cd openvino && git submodule update --init --recursive
△代碼若顯示不全，可左右滑動- 安裝 Python 環境依賴，以構建 Python Wheel
$ python -m pip install -U pip
$ python -m pip install -r ./src/bindings/python/src/compatibility/openvino/requirements-dev.txt
$ python -m pip install -r ./src/bindings/python/wheel/requirements-dev.txt
△代碼若顯示不全，可左右滑動- 創建編譯目錄
$ mkdir build && cd build
△代碼若顯示不全，可左右滑動- 使用 CMake 編譯 OpenVINO?
$ cmake .. -DENABLE_LLMDNN=ON \
-DBUILD_PYTHON_TESTS=ON \
-DENABLE_CPU_DEBUG_CAPS=OFF \
-DENABLE_DEBUG_CAPS=OFF \
-DCMAKE_BUILD_TYPE=Release \
-DENABLE_INTEL_MYRIAD_COMMON=OFF \
-DENABLE_INTEL_GNA=OFF \
-DENABLE_OPENCV=OFF \
-DENABLE_CPPLINT=ON \
-DENABLE_CPPLINT_REPORT=OFF \
-DENABLE_NCC_STYLE=OFF \
-DENABLE_TESTS=ON \
-DENABLE_OV_CORE_UNIT_TESTS=OFF \
-DENABLE_INTEL_CPU=ON \
-DENABLE_INTEL_GPU=OFF \
-DENABLE_AUTO=OFF \
-DENABLE_AUTO_BATCH=OFF \
-DENABLE_MULTI=OFF \
-DENABLE_HETERO=OFF \
-DENABLE_INTEL_GNA=OFF \
-DENABLE_PROFILING_ITT=ON\
-DENABLE_SAMPLES=ON \
-DENABLE_PYTHON=ON \
-DENABLE_TEMPLATE=OFF \
-DENABLE_OV_ONNX_FRONTEND=OFF \
-DENABLE_OV_PADDLE_FRONTEND=OFF \
-DENABLE_OV_PYTORCH_FRONTEND=OFF \
-DENABLE_OV_TF_FRONTEND=OFF \
-DENABLE_OPENVINO_DEBUG=OFF \
-DENABLE_CPU_DEBUG_CAPS=ON \
-DCMAKE_INSTALL_PREFIX=`pwd`/install \
-DCMAKE_INSTALL_RPATH=`pwd`/install/runtime/3rdparty/tbb/lib:`pwd`/install/runtime/3rdparty/hddl/lib:`pwd`/install/runtime/lib/intel64 \
-Dgflags_Dir=`pwd`/../thirdparty/gflags/gflags/cmake
$ make --jobs=$(nproc --all)
$ make install
△代碼若顯示不全，可左右滑動- 安裝針對 OpenVINO? Runtime 和 openvino-dev 工具構建好的 Python Wheel
$ pip install ./install/tools/openvino*.whl
△代碼若顯示不全，可左右滑動- 檢查系統 GCC 版本和 Conda Runtime GCC 版本。如下所示，如果系統 GCC 版本高于 Conda GCC 版本，請升級 Conda GCC 至相同版本，以滿足 OpenVINO? Runtime 的需求。（可選）
##check system (OpenVINO compiling env) gcc version
$ gcc --version
gcc (Ubuntu 11.3.0-1ubuntu1~22.04.1) 11.3.0##check conda python (runtime env for OpenVINO later) gcc version
$ python
Python 3.10.11 (main, May 16 2023, 00:28:57) [GCC 11.2.0] on linux##If sys gcc ver > conda gcc ver, upgrade conda gcc ver -> sys gcc ver
$ conda install -c conda-forge gcc=11.3.0
△代碼若顯示不全，可左右滑動- 將 PyTorch 模型轉為 OpenVINO? IR
$ cd ..
$ python tools/gpt/gen_chatglm.py /path/to/pytorch/model /path/to/ov/IR
△代碼若顯示不全，可左右滑動使用 OpenVINO? Runtime API 為 ChatGLM 構建推理流水線
本文提供了使用 Transformer 和 OpenVINO? Runtime API 構建推理流水線的樣本。首先，在 test_chatglm.py 中，創建一個由 transformers.PreTrainedModel 衍生的新類。
然后，通過使用 OpenVINO? Runtime Python API 構建模型推理流水線來更新轉發函數。其他成員函數則遷移自 modeling_chatglm.py [2]的 ChatGLMForConditionalGeneration。
如此一來，即可確保輸入準備工作、set_random_seed、分詞器/連接器 (tokenizer/detokenizer) 以及余下的流水線操作能夠與原始模型的源碼保持一致。
如需啟用 int8 權重壓縮，只需設置簡單的環境變量 USE_INT8_WEIGHT=1。這是因為在模型生成階段，已使用 int8 對全連接層的權重進行了壓縮，因此模型可在之后的運行過程中直接使用 int8 權重進行推理，從而免除了通過框架或量化工具壓縮模型的步驟。
請按照以下步驟使用 OpenVINO? Runtime 流水線測試 ChatGLM：
- 運行 bf16 模型
$ python3 tools/gpt/test_chatglm.py /path/to/pytorch/model /path/to/ov/IR --use=ov△代碼若顯示不全，可左右滑動- 運行 int8 模型
$ USE_INT8_WEIGHT=1 python test_chatglm.py /path/to/pytorch/model /path/to/ov/IR --use=ov
△代碼若顯示不全，可左右滑動權重壓縮：降低內存帶寬使用率，提升推理速度
本文采用了 Vtune 對模型權重數值精度分別為 bf16 和 int8 的內存帶寬使用率（圖 3 和圖 4）以及 CPI 率進行了性能對比分析（表 1）。結果發現：當模型權重數值精度壓縮至 int8 時，可同時降低內存帶寬使用率和 CPI 率。
圖3 模型權重數值精度為 bf16 時的內存帶寬使用率
圖4 模型權重數值精度為 int8 時的內存帶寬使用率
表1 采用不同模型權重數值精度時的 CPI 率
每條指令消耗的時鐘周期 (Clockticks per Instruction Retired, CPI) 事件率，也稱為“平均指令周期數 (Cycles per Instruction)”，是基于硬件事件抽樣收集的基礎性能指標之一，在抽樣模式下也稱為“性能監控計數器 (PMC) 分析”。
該比率計算方式為：用處于非停機狀態的處理器時鐘周期數 (Clockticks) 除以已消耗指令數。每個處理器用于計算時鐘周期數和已消耗指令數的確切事件可能并不相同，但 VTune Profiler 可辨別和使用正確的數量。
CPI < 1 時，通常為采用指令密集型代碼的應用，而 CPI > 1 則可能是停滯時鐘周期密集型應用，也可能是內存密集型應用。
由此，我們可以得出結論，類似 chatGLM 等語言模型對內存帶寬的要求非常高，性能往往受到內存操作或帶寬的限制。
很多場景下，消除內存操作的負載，性能會因此獲得大幅收益。在優化此類模型時，如何在不影響精度的同時對模型進行壓縮或輕量化處理是一項不可或缺的技巧。除此之外，在異構平臺和框架上進行部署，還涉及到減少內存/設備存儲之間的數據搬運等優化思路。
因此，在壓縮模型的同時，還需要考慮對原始 pytorch 模型推理 forward/generates 等函數流水線的優化，而 OpenVINO? 在優化模型自身的同時，還將流水線的優化思路體現在修改模型結構中（將 KV cache保存在模型內部），通過優化 Optimum-intel 等框架的流水線，減少內存拷貝和數據搬運。
結論
筆者根據上述方法重新設計執行圖并優化了 GLMBlock，消除了 ChatGLM 模型輸入和輸出之間的內存副本，且模型運行高效。
隨著 OpenVINO? 的不斷升級，本方案的優化工作也將得到推廣并集成至正式發布的版本中。這將有助于擴展更多的大語言模型用例。敬請參考 OpenVINO? 官方版本[6]和 Optimum-intel OpenVINO? 后端[7]，獲取有關大語言模型的官方高效支持。
了解更多內容，請點擊文末【閱讀原文】。
作者簡介：
英特爾? OpenVINO? 開發工具客戶支持工程師趙楨和鄒文藝，英特爾? OpenVINO? 開發工具 AI 框架工程師羅成和李亭騫，都在從事 AI 軟件工具開發與優化工作。
OpenVINO? stateful模型構建：https://docs.openvino.ai/2022.3/openvino_docs_OV_UG_network_state_intro.html
通過 opset 構建模型：https://github.com/openvinotoolkit/openvino/blob/master/samples/cpp/model_creation_sample/main.cpp
