Stories

Detail Return Return

可信數據流通開發者必看，深度解讀隱語密態計算設備SPU（Secretflow Processing Unit） - Stories Detail

打開鏈接點亮社區Star，照亮技術的前進之路。每一個點贊，都是社區技術大佬前進的動力

Github 地址： https://github.com/secretflow

SPU是_Secretflow Processing Unit_的簡稱，它作為隱語平台的密態計算單元，為隱語提供安全的計算服務。

image

1.SPU概念理解

密態計算單元這個概念聽起來比較晦澀，我們用一個實際的例子介紹一下SPU的作用。

假設要用 JAX 寫邏輯迴歸（SPU不依賴JAX，選擇JAX因為簡單），代碼如下：

import jax
import jax.numpy as jnp

def sigmoid(x):
    return 1 / (1 + jnp.exp(-x))

def loss(x, y, w):
    pred = sigmoid(jnp.dot(x, w))
    label_prob = pred * y + (1 - pred) * (1 - y)
    return -jnp.sum(jnp.log(label_prob))

def logit_regression(x, y, epochs=3, step_size=0.1):
    w = jnp.zeros(x.shape[1])
    for _ in range(epochs):
        grad = jax.grad(loss, 2)(x, y, w)
        w -= grad * step_size
    return w

x, y = load_full_dataset()
w = fit(x, y)

【注】L15 使用了JAX的自動求導功能，對loss function進行了求導

對於上述程序，JAX提供了jit（全稱Just In Time，是編譯技術的一種）方法，在不用任何算法改動的前提下將上述程序編譯到GPU上，從而加速執行。

jax.jit(logit_regression)(x, y)  # jax.jit將logit_regression翻譯到GPU執行

JAX 本質解決了兩個問題

  1. 降低開發成本，在用户無感的情況下，利用GPU/TPU進行加速
  2. 降低學習成本，通過兼容numpy API 並提供自動求到來吸引開發者

沿着這個思路，對安全計算做個類比

  • JAX可以利用並行設備進行計算加速，我們可否用安全設備進行安全加固？
  • JAX可以複用numpy API，我們是否可以複用其他AI框架API？

帶着這兩個問題，SPU的核心API就是這樣一個函數，將AI模型翻譯到安全設備上執行

spu.jit(logit_regression)(x, y)

為了實現這個函數，我們需要兩個子模塊

  • SPU (Jit) Compiler: 將原生的AI 程序翻譯成 SPU字節碼
  • SPU VM：一個帶安全語義的虛擬機，解釋和執行SPU字節碼

實際中實現稍微複雜一些

  • 將Tensorflow/PyTorch/JAX 翻譯成SPU字節碼本身是個複雜繁瑣的工作
  • SPU字節碼使用的是密文類型系統，譬如沒有f32/f64等，類型系統需要重新設計
  • SPU後端是MPC，本質上是個分佈式系統，需要處理分佈式系統的通信和協作
  • 上述程序中，變量x, y可能由不同的參與方提供，所以IO模塊有些特別

細節會在後續部分進行簡單介紹，在此不再複述。

2.SPU的功能作用

介紹完SPU是什麼，我們再來理一下為什麼。

市面上的隱私計算框架有很多，比如 TFE，CrypTen，MP-SPDZ 等，為什麼我們要重新造一套輪子呢？

領域之間的距離

安全機器學習是一個交叉領域，其實AI和安全之間有相當的距離。比如

  • 安全開發者更關注基礎算子，比如加減乘除的安全性
  • AI開發者更關注高階算子，比如conv，tensordot

高階算子和基礎算子之間，有很大一段距離，譬如：

機器學習編譯器處理的 lowering/tiling/fusing等

運行時處理的的調度，併發等

無論是基於AI框架（TFE/CrypTen），還是從安全計算出發的框架（SPDZ），都有自己的問題。前者往往難部署，難做安全領域特定的優化。後者往往會需要寫一些Toy AI框架，學習成本高。

SPU試圖縫合這兩個領域之間的間隙，使得：

  • 向上，SPU原生對接AI框架（TF/JAX/PyTorch)，降低AI開發者的學習和開發成本。
  • 向下，SPU提供純粹安全語義接口，只需要實現很少的安全協議（比如加乘與或）就能跑起來複雜的模型，讓目光更聚焦安全本身。

算力和需求的距離

近些年，密態計算（MPC/HE）在算力上都巨大的進步，但是密態算力和AI的算法需求還是難以匹配。

在算力無法匹配算法的時候，一個直觀的想法就是“明密文混合”，用來做安全和性能的tradeoff。比如聯邦學習，將算法的某一個子步驟使用安全計算實現，犧牲局部安全性以換取更高的性能。

隱語提供了非常自由的明密文混合編程範式，我們不限制明文的引擎，也不限制密文引擎，開發者可以用他自己熟悉的框架開發，然後標記其中的某一部分用明文引擎跑，另一部分用SPU跑。

比如：

【注】圖中MPC Device就是SPU實現的

作為對比，從安全和性能這種的角度，無論TFE/CrypTen/SPDZ等都很難進行這種tradeoff。

理論和落地的距離

多方安全計算天然是一個分佈式的系統，部署模式非常多樣，比如：

  • 論文中經常假設計算方和數據提供方分離（outsourcing）
  • 真正進行業務落地時，數據提供方往往同時也是計算方（colocated）
  • 在一個複雜的隱私計算網絡裏，計算方和數據方可以是任意組合的（hybrid）

如圖，我們用三角形表示計算節點，圓形表示數據提供節點（不同顏色表示互不信任）

SPU被設計成部署模式透明的，不用修改任何一行代碼，你的模型都可以在上述任何一種部署場景上被安全且正確的執行。

並且（相對於基於AI平台的隱私計算框架）SPU運行時非常的輕量級，不需要Python runtime，可以方便的進行部署和集成。

所以，Why SPU？

作為AI開發者，你不需要任何安全背景，就可以將你現有的模型安全的應用到多方數據上。

作為安全開發者，你不需要任何AI背景，僅僅實現安全計算的基本算子，就可以支持多種前端框架。

並且，你可以方便的部署和運維，在安全和性能之間折中，找到最佳的落地方案。

3.SPU的架構

介紹完是什麼和為什麼之後，我們簡單介紹一下SPU的架構。

SPU上層對接了XLA-HLO，然後利用MLIR將HLO翻譯成SPU IR，最後交給SPU VM進行解釋執行，如圖：

Workflow

基本工作流程是：

  1. 開發人員用自己熟悉的框架建模，然後用AI compiler將模型編譯成 XLA IR。
  2. SPU compiler將XLA IR編譯成SPU IR（SPU字節碼）
  3. 參與方（Alice/Bob/Charlie) 將數據 infeed 給SPU VM
  4. SPU VM執行字節碼，接收輸入，安全計算，並且產生輸出
  5. 參與方協商將結果解密輸出到某處

Why XLA？

這裏對於XLA不熟悉的同學進行一個簡單介紹，XLA 是一種針對特定領域的線性代數編譯器，是tensorflow內部實現的一個子模塊，使用編譯器相關技術用來加速模型的執行。

我們可以將SPU理解成一個帶安全語義的Backend，對接XLA的理由是Tensorflow/Jax/PyTorch計算圖最終都可以翻譯到XLA IR，只要SPU可以解釋和執行XLA IR，理論上就可以原生支持多種AI前端。

所以理論上並不是選擇了XLA編譯器，而是選擇了XLA IR 作為AI和MPC的橋樑。

Arch

下面簡單介紹一下SPU的編譯和執行過程

  • 前端，我們依賴AI前端將python代碼翻譯成XLA IR
  • 編譯器，我們使用MLIR技術棧對HLO進行優化並翻譯成PPHLO（SPU字節碼）
  • 運行時，我們逐漸將Tensor ops拆解，經過SPU HAL(硬件抽象層，處理fxp/int)，最終dispatch到協議層
  • 協議層只需要實現Ring or Field上的基礎運算即可

最終，通過編譯時和運行時的層層翻譯，SPU將AI前端和MPC後端解耦，使得在SPU中擴展的任何安全協議都可以無感的支持多種前端。

Optimization

SPU完全自主研發，所以我們可以針對安全計算的特點進行優化，比如

  • 針對MPC延時敏感的計算類型進行併發調度
  • 針對特定協議設計特殊的VM指令集
  • 基於C++語言，開發者可以得到Low-level access，更有效的榨取性能

具體細節在此不再複述。

SPU依然在一個高速迭代過程中，隱語期待更多的AI專家、編譯專家、安全專家參與共建。

加入隱語社區，探索可信數據流通的未來！

隱語社區是一個面向融合可信數據空間、隱私計算、數據元件等多項數據流通利用基建技術設施在內的開源社區，致力於推動前沿技術探索、技術標準體系共建與產業應用場景共創，促進高質量數據資源流通利用與價值釋放。

關注我們，獲取最新資訊

Github倉庫地址：

https://github.com/secretflow

官網：

https://www.secretflow.org.cn

嗶哩嗶哩：

https://space.bilibili.com/2073575923

聯繫小助手：

SecretFlow04

開源可信隱私計算框架—— SecretFlow

支持MPC、FL、TEE等主流隱私計算技術，融合產學研生態共創能力，助力隱私計算更廣泛應用到AI、數據分析等場景中，解決千行百業隱私保護和數據孤島等應用痛點。

開源安全可信系統軟件棧——星綻 Asterinas

聚焦安全可信技術軟件棧的開源社區，面向通用計算和機密計算兩大場景提供堅實的安全底座，為不斷涌現的安全攸關和隱私敏感的應用保駕護航。

Github倉庫地址：

https://github.com/asterinas

官網：

https://asterinas.github.io/

本文由隱語社區統一發布，歡迎大家點 Star
安全 , 區塊鏈
user avatar Evsward Avatar tecdat Avatar xuri Avatar
Favorites 3 users favorite the story!
Favorites

Add a new Comments

Some HTML is okay.