大數據文摘受權轉載自雲鼎實驗室
2023年在AI的發展史上一定是濃墨重彩的一筆，在這一年裡出現了百模大戰、全民“煉丹”的場面，圍繞著各種模型的訓練技術和算力需求有很多討論。隨著模型的成熟以及算力市場的發展，7B、13B這類小型號的模型也出現了端上部署的需求，其中以移動設備廠商最為突出。2024年，在端上部署和應用模型也可能會成為各傢移動廠商的一個營銷熱點。
在端上部署推理模型自然就會帶來一個安全問題，如何保證模型文件的安全？LLM的架構目前都是Transformer的變體，如果拿到模型文件，借助開源的代碼很容易完成適配，然後就可以異地部署起來。LLM現在從零開始預訓練、指令微調、強化學習到最後可用，成本並不低，如果很容易被別人竊取到模型文件並被推理使用，會造成不小的損失。
為了解決上面這個安全問題，我們借助TVM編譯器幫助我們實現高效模型推理部署，並在TVM中引入Tensor加解密層。這樣模型文件在本地是處於加密狀態，在推理運行的時候才是解密狀態，以此保證模型文件不被輕易竊取。下面介紹技術細節。
TVM簡介
TVM全稱Tensor Virtual Machine，中文名稱張量虛擬機，它主要的用途是將深度學習模型編譯成最小可部署的模塊，並提供更多的後端自動生成和優化模型的基礎設施。架構如下：
它的特點是：
1. 最佳的性能：通過編譯和最小運行時，在現有的硬件上優化機器學習工作負載，進而發揮最佳性能。
2. 豐富的後端：可以在CPU、GPU、瀏覽器、微控制器、FPGA等硬件上運行，同時可自動在多種後端生成和優化張量算子。
3. 靈活可用：TVM 的靈活設計支持對區塊稀疏度、量化（1、2、4、8 位整數，posit）、隨機森林/經典 ML、內存規劃、MISRA-C 兼容性、Python 原型設計。
4. 簡單易用：在 Keras、MXNet、PyTorch、Tensorflow、CoreML、DarkNet 等的深度學習模型上編譯。用 Python 借助 TVM 進行編譯，即可直接使用 C 、Rust 或 Java 完成推理部署。
TVM工作流程
TVM屬於領域編譯器，其工作流程也遵循通用CPU編譯器的設計規則，將復雜的任務簡單化，並進行多層抽象， TVM也是采用多層IR設計，IR的轉換如下：
多層IR轉換以後，最後生成指定的後端文件，如上圖所示TVM編譯會生成三個文件：mod.json是模型的計算圖描述文件，mod.so是經過優化生成的算子，mod.param是存儲的模型參數文件，其文件格式如下：
TVM 推理運行如下：
Safe Tensor
我們提出了Safe Tensor的概念。Safe Tensor的設計原則是簡單高效易用，我們放棄了全量加密、mpc等技術，使用了傳統的對稱加解密算法，在TVM生成mod.param文件的時候，在其保存每層tensor的過程中將裡面的數據進行加密。與原始mod.param文件比較，我們隻加密了最關鍵的數據，在加解密速度上保證了最快，也將引入的加解密將損耗降到了最低。Safe Tensor采用了傳統的對稱加密算法。為了確保業務安全，密鑰的保存和使用應該遵循最佳實踐。在使用密鑰之前，需要使用非對稱加密算法進行解密。
TVM生成模型參數文件格式就變成如下：
效果如下：
Safe Tensor的方案隻是加密了模型最重要的tensor數據，這對TVM的上層是無感的，運行的方法和之前介紹的方法也相同。
總結
Safe Tensor的方案有如下優勢：
1. 相對TVM生成的推理程序，加密後業務接入無感，靈活易用。
2. 推理時性能無損失。
3. 模型文件被加密，和推理代碼一一綁定，直接拷走模型文件後不可用，從而起到保護作用。
Safe Tensor給大傢提供了一個新的模型保護思路，現階段隻是在模型文件生成和加載的時候進行加解密。如果想獲得更好的保護效果，還需要進程內存保護的技術保證內存dump相關的攻擊。隨著未來AI技術的發展，硬件設備廠商也會慢慢關註這塊市場，也會推出相關的硬件特性，屆時可以使用軟硬結合的安全保護，讓模型應用既高效又安全。
新上線一批4090/A800/H800/H100
特別適合企業級應用
