What's with ML software and pickles?

Nelson Elhage

機器學習軟體與 pickle 是怎麼回事?

原文由 Nelson Elhage 發布,訂閱此部落格

多年來,我身為一名軟體工程師,完全是機器學習領域的門外漢,但對它抱有一些好奇,也偶爾有些邊緣上的接觸。這段期間,每當我在 Python 的機器學習生態系中看到 Python pickle 被廣泛使用,內心反覆出現的感受就是驚愕(老實說,還有不屑)。

除了重大的安全性問題1之外,使用 pickle 來做序列化往往非常脆弱,隨著你演進程式碼、升級函式庫與 Python 版本,會帶來各種惡夢。在我身為軟體工程師的職涯中,我一度認為——在相當程度上——pickle 模組根本沒有任何正當的使用情境,它純粹是一個誘人的麻煩;是年輕、熱情卻又天真的生態系所犯下的錯誤。

不過,最近這幾年我開始在機器學習領域擔任正職工作,這讓我對這個生態系有了新的視角。我實際體會到機器學習軟體為何會長成現在這個樣子,以及 pickle 究竟在幫使用者解決哪些問題。我還是不喜歡它,但對於這個問題領域有了更多理解與同理,也不再認為這個問題是「微不足道」的,或僅僅出於無知或偷懶。這篇文章就是想分享我所學到的一點心得,並試著弭平這兩個世界之間的隔閡。

做研究,而非做軟體

歸根究柢,我最想傳達的核心觀點是:機器學習生態系中的大多數軟體,首要目的是為了支援研究人員,而非軟體工程師。此外,(至少在這個脈絡下)研究具有幾個相關的特質——甚至可以說就是由這些特質所定義的:

  • 研究的主要產出是知識,而非軟體成品。研究團隊寫程式是為了回答研究問題、增進自己、團隊乃至整個領域對某個領域的理解,而不是為了打造出軟體工具或解決方案本身。
  • 研究人員往往事前並不知道哪些想法會成功或有用,因此他們優先追求的是高效率地嘗試大量想法,並預設其中大多數最終都不會成功。

作為上述情況的必然結果,研究人員(尤其是在機器學習這類領域!)可能會寫下大量程式碼,但其中絕大多數很快就被丟棄也是完全正常的。失敗的實驗通常會被直接捨棄,即使是成功的實驗,往往在論文寫完後也會被丟掉。即使未來的研究要建立在某個成功的想法之上,常見的做法也是為了新的專案重寫程式碼,此時對問題已有更深的理解,並帶著新的問題或研究課題出發。

相較於典型的軟體工程環境,這種環境下的取捨與條件,讓天平更大幅度地傾向 pickle(而非其他序列化策略):

  • 大多數程式碼從頭到尾都只由同一小群人(研究人員及其團隊)撰寫與執行,而且是在單一環境中;通常「開發」與「正式環境」之間幾乎沒有區別,也沒有外部使用者、資料來源或互動。這些因素加起來,使得 pickle 序列化在實務上的安全風險低得多(但仍非零!)
  • 能夠直接拿起一個資料結構——甚至是一段程式碼——就將其序列化,而不必停下來思考綱要、格式或資料演進,有著巨大的好處;我們在乎的是多快能得到結果,花心思去想綱要、序列化格式或在磁碟格式與記憶體格式之間轉換,都是完全非必要的額外工作。
  • 由於大多數程式碼很快就會被丟棄,pickle 的脆弱性、以及處理程式碼變更與版本升級所帶來的挑戰,就顯得不那麼重要。

研究小記

為了讓上述的取捨更具體一點,我想以一種風格化的方式描述一位研究人員進行機器學習研究的過程,試著呈現她會在哪些具體情境下用到 pickle,以及我所看到的相關取捨,這些都是來自我自身從事機器學習研究、以及支援(比我有經驗得多!的)機器學習研究人員的經驗。

假設一位研究人員正在實驗一種新的深度學習模型架構,或是某個現有架構的變體。她的架構會有大量組態選項與超參數:層數、每一層的類型、各種向量的維度、在何處以及如何對 activation 做正規化、要使用哪種非線性函數等等。模型中許多元件會是機器學習框架提供的標準層,但研究人員也會穿插加入一些新穎的邏輯片段。

我們的研究人員需要一種方式來描述一個具體的模型——也就是這些設定的一種特定組合——並且能夠將其序列化、之後再重新載入。她基於幾個相關的原因需要這麼做:

  • 她很可能可以使用包含 GPU 或其他加速器的運算叢集來執行任務。她需要一種方式,將模型描述提交給在叢集上執行的程式,讓叢集能執行她的模型。
  • 在模型訓練期間,她需要以能夠重新載入並接續訓練的方式儲存進度快照,以防硬體故障或任務被搶占。
  • 模型訓練完成後,研究人員會想再次載入它們(可能是最終快照,以及一些訓練中途的檢查點),以便對其進行評估與實驗。

在這裡我們可以想像許多不同的解法!我們可以為模型寫一個 protobuf 或 JSON 描述,並撰寫程式碼將其轉換為實際的模型實例。我們甚至可以想像把模型表示為「一組命令列參數陣列」,再將其解析為模型實例。

然而,這裡有一個能解決所有這些問題、而且前期幾乎不需要額外成本的解法:

  • 「一個特定模型」的定義,就只是「一個實作該模型的物件」(很可能是 torch.nn.Module 的子類別)
  • 研究人員透過一個研究用的腳本,以某種即興的方式建構這些物件,並隨著進行新實驗不斷編輯、演進該腳本
  • 接著她使用 pickle 將整個物件序列化,並將其作為儲存版本,供後續分析或在遠端機器上載入執行之用。

這種做法不需要大費周章,對於研究程式碼或實作的結構也不需要特定的紀律或模式,相較於「只是建構並執行一個模型」,幾乎沒有額外的認知負擔。它能同樣輕鬆地支援標準元件與自訂元件:研究人員可以在類別階層深處的任意位置做些小幅調整,或是替換整個模型架構,兩者都不需要對序列化系統多做思考或修改。

到了這個階段,我們的研究人員大概會寫一些腳本來產生各種超參數的笛卡兒積組合,將它們全部排程到叢集上,然後就去吃午餐(或直接去度週末),讓它們慢慢訓練。

這些 pickle 檔案確實很脆弱!它們將模型的全部實作細節都序列化了,當研究人員新增設定與參數、重構實驗時,很容易就破壞舊的實驗與模型。但話說回來,大多數實驗的結果本就乏善可陳,根本不需要再回頭看。對於少數有價值的結果,只要記住每個實驗對應的 git commit,要回溯 checkout 也很容易。而如果研究人員需要更持久地保留少數特定模型,通常也可以在某處的 __setstate__ 方法中做幾個戰術性的修補來挽救。

話雖如此,如果這一系列實驗整體上算是成功的,我們需要沿著這條研究路線繼續深入、嘗試更多變體,並為了比較而更長期地保存模型,那麼這種脆弱性就開始帶來痛苦了,因此研究人員最終可能會考慮重構。一個直接的做法是將模型的「描述」與實作分離:把模型的組態/設定/超參數集中到一個 class ModelConfig: 中;這個類別才是被序列化的對象(依然透過 pickle),而研究人員則手動維護一個 build(cfg: ModelConfig) -> Model 函式,負責根據組態來實例化完整的模型。

如此一來,研究人員現在只需要擔心這個類別的向前與向後相容性,就可以相對自由地重構任何實作程式碼,只要同步重構從 ModelConfig 建構 Model 物件的那段程式碼即可。

但值得注意的是,我們的研究人員(和/或她的團隊)仍在實驗、嘗試新事物,迭代速度以及「進行實驗所需的時間」,甚至更重要的是「進行實驗的認知成本」,才是主導一切的原則。因此,雖然這個類別上的許多欄位會是漂亮、正常的整數、布林值,或許還有列舉型別,但每做一個新實驗就得為其新增支援,還是有些繁瑣。

於是,或許我們的研究人員會留下一些「逃生艙口」。舉一個簡單的例子:在這個領域,嘗試替換在全連接層之後對張量套用的非線性函數或「激活函數」,至今仍算是頗為流行的嘗試(我在這方面自己也難辭其咎)。

我們要如何在組態中儲存一個激活函數(它很可能就是一個型別為 Tensor -> Tensor 的 Python 函式)呢?

我們可以在某處用一個 dict[str, Callable] 來註冊函式,並以名稱來儲存函式……但話說回來,Python 的模組命名空間本身不就是這樣一個 dict 嗎?所以如果我們用的是 pickles,大概可以直接把函式直接存在組態上就好。

受到這個發現的鼓舞,我們的研究人員更進一步;她甚至可以儲存整個 build_xxxx: Callable[[Config], Module | None] 的覆寫,讓組態能以任意方式覆寫組態中的任意部分!

這類功能讓過去那種脆弱性問題大多又回來了;但重要的是,它們讓研究團隊得以在兩者之間取得平衡:一邊是較舊、較穩定、以「正常」參數為一等公民的核心組態,另一邊則是實驗的快速翻攪與流動,其中(幾乎)沒有什麼能長久留存,迭代速度才是一切。

範疇蔓延

以我個人的看法,上述的一切「大致上還算可以」,雖然有點令人不適、也帶有風險,只要專案真的只在可信的環境中、僅供內部研究使用

然而,更大的問題在於,一旦你開始以這種方式例行性地使用 pickle,就會習慣性地將它用在所有序列化問題上,即使是在那些使用 pickle 會帶來更大麻煩的場合。不僅研究人員會這麼做,這種模式還往往會外溢到工具與整個生態系中。

舉例來說,PyTorch 就將 pickle 檔案作為其預設序列化格式的一部分(透過 torch.loadtorch.save),儘管在大多數情況下,他們真正需要的只是一些「基本款」的資料結構——字串與 Tensor 所組成的 list 與 dict。PyTorch 甚至有一個自訂的 C++ pickle 載入器,只支援少數內建型別,可以透過載入器中的 weights_only 參數來啟用;但為了向後相容,這個旗標預設為 False,而且(在沒有攻擊者的情況下)啟用它只會造成東西壞掉;因此它很少被使用,而每一個 PyTorch 權重檔都等於公開邀請使用者去使用這個危險的介面,讓自己暴露於風險之中。

該怎麼做?

我寫這篇文章的主要目的,是想試著凸顯——某種程度上——「從內部看出去的視角」,說明為何我在機器學習生態系中看到如此大量使用 pickle,以及它為使用者解決了哪些問題。話雖如此,隨著機器學習的重要性日益提升,並在各種不同情境中獲得更多使用者與部署,我預見不安全的 pickle 使用所造成的問題,其重要性與急迫性也將與日俱增,因此想就未來的幾條出路提供一些想法。

更好的資料格式與介面

首先,我認為我們(這裡主要指的是那些具有軟體工程背景與技能,和/或在機器學習領域內外擔任以軟體工程為主的角色的人)應該持續投資於高品質的 pickle 替代方案,以便為研究人員提供更多、更好的選擇。如果做得好,在許多情況下,我們甚至可以為研究人員提供比 pickles 更好的體驗,同時具備更佳的安全性。

在此我若不指出,以 XLA 為核心建構的 Google 機器學習生態系——Tensorflowjax 及其各種函式庫與框架——在很大程度上避免了本文所描述的問題,那就有所疏漏;它們傾向於使用以 protocol-buffer 與numpy 格式為基礎的序列化格式,pickle 的濫用情況要少得多。一方面,我將它們視為這種生態系確實可行的存在證明;另一方面,我傾向認為 PyTorch 之所以能取得如此主導地位,是因為它整體上為研究人員提供了更優越的使用體驗,尤其是在上手門檻與修改/探索的便利性方面,因此我不太願意將 Google 的函式庫指為理想的解決方案。

Hugging Face 的safetensors 函式庫是這個領域進展的一個絕佳範例,它為以安全且高效的方式序列化張量資料提供了堅實的基礎。在我看來,它最好被視為一個稍微偏底層的基礎元件;它支援序列化與還原形式為 dict[str, Tensor] 的物件;但研究人員往往想要巢狀字典,或希望將中繼資料與權重一起儲存在同一個封包中;我很期待看到在 safetensors 之上發展出更符合人體工學、對研究人員更友善的格式與 API。

研究人員常常也想序列化 Python 類別樹以及張量資料;或許 Python 型別標註的出現與近期成為主流,以及建立在其上的 attrsdataclasses 等 API,為打造一個由 Python 型別標註與少量額外中繼資料所驅動的函式庫提供了契機(或許可以建構在 cattrs 之上),讓某些經過型別標註的 Python 類別樹能夠或多或少直接地被序列化與反序列化。

限制 pickle

然而不管好壞,我們今天已經有大量 pickle 檔案與使用 pickle 檔案的函式庫,而且問題仍在擴大。在我們努力提供替代方案、並鼓勵專案遷移到這些替代方案的同時,這裡有一個用於降低危害或緩解問題的推測性想法:

首先,我在想 Python 是否應該支援一個全域的「nopickle」模式來完全停用 pickle 模組。這樣的旗標至少能讓部署 Python 的維運人員輕易判斷自己在執行時載入不受信任的 pickle 是否面臨風險。

然而,如果這樣太過嚴格,那麼加入一些選項來以特定方式限制 pickle 的使用如何?

Pickle 清單

特別是,我對要求所有 pickle 的使用都必須宣告一份可用於解碼的合法二進位物件清單這個想法很感興趣,這份清單大概會以 sha256 或其他密碼學雜湊來命名。我們可以想像同時支援兩種方式:在 pickle.load(或某個新 API)的呼叫端指定預期的物件,或是在某種清單檔案中以帶外(out-of-band)的方式指定,甚至可能指定合法的(呼叫堆疊、blob ID)組合。任何未指定此資料、或不符合允許雜湊的 pickle 載入都會被拒絕。

這個方法之所以吸引我,是因為它或許能讓那些部署機器學習軟體的人——只需一次——蒐集所有正在被載入的模型權重檔或組態檔的清單,然後「凍結」這個集合,確保在執行時不會遇到任何新的 pickle。

反過來,這在很大程度上將管理 pickle 使用的問題,化約為管理第三方相依套件的問題;可以肯定的是,這並非一個已解決的問題,但這本來就是你無論如何都會面對的問題。

具體來說,為了部署一個新的第三方機器學習模型,你通常需要同時安裝模型實作(以 Python 模組或腳本的形式)以及模型權重。兩者若帶有惡意,都會構成安全風險;但 pickle 帶來的風險在某種意義上更糟,因為即使你已審查並鎖定模型程式碼的特定版本,pickle 檔案往往仍會在執行時才被下載,這意味著你依然暴露在來自系統外部的遠端程式碼執行(RCE)風險之下。

如果我們能保證那段程式碼只會載入特定的 pickle 檔案,而我們可以在最初決定信任該相依套件時就對其進行一次稽核,那麼我們就讓 pickle 的使用情境「看起來像」安裝第三方程式碼的問題——一旦我們抓取並鎖定某個版本,至少不會有額外的執行時攻擊面。

我們可以想像透過一個執行時組態系統讓這些想法變得更複雜;在極限情況下,我們可能會想像出像這樣相當複雜的 API 來控制 pickle 的使用:

with pickle.no_pickles() as token:
  pickle.loads(...) # NoPickleException

  with pickle.allow_pickles(token):
    pickle.loads(...)

  with pickle.allow_pickles(token, only_objects=[
    "sha256:some sha256",
  ]):
    pickle.loads(buf) # hashes buf, only proceeds if it matches

  # You must pass the same token returned by `no_pickles` to
  # `allow_pickle`; this prevents libraries from trivially blanket
  # calls to renable `pickle`
  with pickle.allow_pickles(bad_token):  # BadTokenException
    ...

這類機制當然永遠不會完美;函式庫總是有辦法繞過它們。在極限情況下,函式庫永遠可以自行打包一個不受這些限制的純 Python 版 pickle 實作。然而,這只是一個函式庫可能(有意或無意地)植入後門的百萬種方式之一;這項提案的目標並非解決供應鏈安全問題,而是緩解 pickle 特有的風險,並且如前所述,將風險從執行時轉移到套件安裝時。我們會期望那些公然繞過標準安全機制的函式庫,會被社群注意到並標記為不可信2

這類功能可以在一個透過 monkey patch 修改 pickle 模組的函式庫中,以 out-of-tree 的方式先做出原型,但從標準化與採用的角度來看,最終將其內建到直譯器中或許也有其價值。


  1. 十多年前,我發表了(就我所知)最早針對 pickle 系統進行攻擊的公開文章之一;但這個問題其實早已眾所周知多年,而如今甚至已經出現了一鍵式利用工具↩︎

  2. 連我自己寫下這句話時,也覺得它顯得有些天真。但話說回來,我確實認為這樣的解法在很大程度上是將問題化約為現有的供應鏈安全問題,無論好壞皆然。 ↩︎

本文章由 muse-spark-1.2-contributor 進行翻譯

留言