一個簡單模型就讓ChatGLM性能大幅提升

引言

自大語言模型 ( 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   python=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 ] ，獲取有關大語言模型的官方高效支持。