一、從 2024 年諾貝爾化學獎談起

2024 年諾貝爾化學獎得主都不是來自化學專業。其中 David Baker 從事多年蛋白質設計研究，包括一些模型和傳統生物信息工具，類似于現在的生成式場景。另外兩位得主來自谷歌旗下的 DeepMind 團隊，該團隊主要專注于蛋白質生成領域，其另一重要成就是之前在圍棋比賽中戰勝人類的 AlphaGo。

蛋白質的業務價值非常大，幾乎所有生物公司都無法繞開這個領域，都會做一些蛋白質相關的應用或者模型二次開發。同時我們也發現這個領域對于計算資源的消耗是非常大的，一個模型就需要消耗多個 CPU 或者多張卡來處理一個請求，其并發延遲遠超傳統的推薦搜索模型。舉個例子，一個蛋白質結構生成預測，如果需要預測一個 100*1000 個序列的混合物，就得 30 分鐘，一天只能計算幾十的。蛋白質的序列生成就更加耗時，可能一個顯卡設計一個就需要幾個小時。所以在蛋白質研究領域，計算能力有很大的提高空間，接下來就將分兩部分來介紹這個場景的優化。

二、加速蛋白質結構預測性能

首先來介紹 DeepMind 團隊關于加速蛋白質結構預測的工作，其主要思想是基于 Ray 的 Workflow，實現高效異構調度。

AlphaFold 有很多版本，這里介紹一個比較顛覆傳統的版本 v2.3。上圖中第一幅圖，輸入是人體的氨基酸序列，輸出是預測的結構，中間模型經過多次轉化，比如第一步是預處理，之后進入模型，類似 transformer。第二幅圖中是精簡后的架構，用了一些生物學工具，效果很好但是效率低，非常慢。通常人們喜歡多個模型疊加拿到更好的效果，但是這樣就會更慢。得到結果后，結果可能與生物理解不一致，需要微調，去掉完全不符合生物學屬性的蛋白質坐標。

我們深入細節，發現第一步是最慢的，主要是來自于一個傳統聲訊工具 MSA，其需要消耗 1T-2T 內存。最后一步是 Relax，是一個混合計算，既需要 GPU 也需要 CPU。GPU 需要幾百兆顯存，CPU 時高時低，如果不做優化，是非常浪費時間的，模型結束后給到 Relax，GPU 利用率就會較低，因此需要一個異構的分布式調度，以充分利用資源。

DLmind 是谷歌的一個云原生解決方案，業界使用廣泛。其核心思想是用 Kubeflow 的 pipeline 把整個鏈路全部異構起來，每個是一個單獨的節點處理，這樣預處理不需要 GPU，中間的串行可以解耦整個模塊，并且利用其擴展性，分布更為靈活。其中的串行調度，類似批處理模式，但是整個延遲還是很嚴重的，有一定的局限性，比如吞吐很慢，基本無彈性能力。當有多個請求，多個 batch 時資源也是固定的，并且這些 batch 都必須在第一步預處理結束后才能觸發（第一步耗時最多）就導致了下游有更多 GPU 是無法利用的。

圖 3 中的紅色部分是其中的關鍵節點。a 和 b （串行，資源固定）前面已經介紹，下面講一下 c 數據預處理，其與業務場景高度相關，這與傳統推薦搜索場景是完全不一樣的。這個數據預處理過程需要 30 多分鐘，這是因為該過程是一個傳統的生信工具沒有太多深度優化，需要預處理之后，將分布式的數據以及MLE數據轉入到內存中才能加速后面的處理。傳統的 Kuberflow 沒有辦法進行這樣的預處理，但是有些方法可以繞開這個限制，比如部署一個 MLE Server，但是這種方案復雜度比較高。所以我們想是否可以應用 Ray 來提供一種高效率、快吞吐的方案，因為 Ray 在離線計算已經有較好的應用。

上圖展示了我們設計的架構，采用 Ray 的 Workflow 方案。這個架構有兩大特點，一是流式調度，二是高效構圖。我們將其拆分成幾個節點，都是對應的 flow 的 node 節點，可以靈活構圖，靈活構圖的好處就是每個節點均可插拔，每個節點可以無縫替換。

第二個核心設計是，考慮到 MSA 的 node 預處理非常慢，因此設計為 Actor Pool 初始化一個節點，預處理做好，推理時其已經是一個預熱好的節點，這樣可以從 30 分鐘優化到 2 分鐘，效果非常可觀。

第三步就到了一個 GPU 節點，類似傳統語言模型，將其作為 model node，如果機器足夠多，可以自動擴縮容，無需人為定義資源。

最后就是 Relax 生信工具涉及到 CPU 和 GPU 混合運算，優化方案有兩種。一是可以將其任務拆分很細，把 GPU 和 CPU 運算進行分離，但是這種方案需要較多的深度開發，開發難度較大。利用 Ray 支持小而一的調度，所以在每個 GPU 節點我們拆分更細，不用做較大改動就可以大幅提高性能。

結果輸出本身就是端到端，Ray 支持通用節點不會導致 OOM（超內存，out of memory）。節點和節點間，請求和請求之間都是可以同步進行的。另外生信場景數據交流均是到 G 級別的，這種如果用傳統解決方案，只能使用分布式存儲系統，頻繁的 I/O 就會有一定的瓶頸。這里我們用到 Ray 的一個共享 Object 之間傳輸有一些傳統架構，就不會有 I/O 的瓶頸，整個吞吐就會非常高效。

Ray 在 AI 時代之所以應用很廣，一個原因就是其 Python 友好，能接入 Python 對庫，很多算子優化均可以用 Python 程序進行封裝完美接入，模型也可以做更多的優化。最后真個過程可以從 30 分鐘減少到 60 秒，這個在業界是比較領先的。

回到業務的核心，利用 Ray 可提升執行效率，并且由于 Ray 的可擴展性，再加上 Ray 整個架構是一個非常好的解耦架構，因此可以降低運維成本，提升合作開發的效率。另外其底層還是 K8S，我們不需要關注 GPU 和 CPU 節點的情況，對于創業公司（人員不足的情況下）是非常友好的。

核心就是 Workflow 的屬性，解決了延遲和吞吐的問題。

三、加速蛋白質生成設計性能

下面介紹蛋白質生成場景的應用，這里用到了 Ray 的另外一個屬性，Ray data，這是一個非常高階且實用的屬性。如上圖左側所示，生成場景主要是包含一個模型的擴散，每一次都會將一個模型擴散成多個模型，一步一步擴散下來，就會需要非常多的處理時間。從一個模型可以擴散到上千級別的，生命科學和其他領域不一樣的就會有傳統生信工具，需要去掉不符合生物學特征的數據。如果不做任何優化，進行一次設計，就需要 2 個小時。最后的 Relax 場景，看似很快，但其實是一個單核場景，堆積起來，1000 個模型需要 1000 個核就會導致整個處理非常慢。

如果可以做到右邊所示的調度流程，就可以完美 overlap 并行運算，是理想中的最優結果，這個方案非常完美，但如果想要用自建的方式來實現還是比較難的，會涉及到多卡多 CPU，需要關注各種分布式的通訊調度異常，執行起來難度非常大。

所以我們引入 Ray 的解決方案，Ray data。Ray data 是一個 high level API，它是一個簡單高效的執行器，對于處理串行計算是一個非常不錯的選擇，但是不適合結構預測有多個分支的任務。

上圖中給出了一個示例代碼，第一步是結構預測，結果出來后用 Ray data 將所有流程串起來。這個代碼是一個非常優雅的解決方案，少量代碼即可實現。實際運行時有一定的問題，主要是第一步會耗時很久。所以我們做了一些升級，第一個是利用典型的流式輸出，完美的 overlap。

第二步是 Relax filter 不需要完整的一張卡，我們利用 Ray 會自動管理底層資源，并行度范圍可以自行設置，最小 1 張卡。Ray data 會根據數據量自動擴容，可大幅減少運維成本，卡更多就可以處理更多的請求。

實際業務中，數據輸出量過大就容易導致 OOM，而數據量過小，則會過于碎片化，都不是完美的解決方案。在這種融合計算架構中使用 Ray 的接口可以有效避免這些問題。

整個運行時間 Baseline 是 2 個小時級別，優化后是 1-2 分鐘，對生命科學領域加速模型處理的意義是十分重大的。

生命科學不僅僅包含蛋白質的處理，還會有 RNA、DNA 等。此外除了離線 batch 任務，還有在線任務。我們也有自己的大語言模型底座，需要微調出來的模型就又不一樣，所以特點是是業務多，模型多。實際問題非常復雜，效率優化是非常重要的。

整個過程非常復雜，需要不斷模型調優，加上創業公司人員不足，不同模型使用的語言，接口都不同，會需要很多重復建設。同時性能也有一定的問題，如果每個人都有自己的模塊，就無法復用，無法滿足高效執行，低吞吐。所以我們希望設計一個新的架構，可以同時減少運維操作，并提高性能。

四、Ray 融合計算架構

基于上述背景，我們設計了基于 Ray 的融合計算架構，將所有事情都在 Ray 中完成，每個接口可插拔，底層可以統一，優化就可以同步進行，具體架構如下圖所示。

融合架構的理念就是所有事情都在 Ray 中進行。最下面的組件大部分是一樣的。向上一層是私有化云原生的部署，上面做了一個封裝使得 Ray 上不會感知到是私有化部署還是云原生，這里我們做了一個 Ray 的抽象。這里面的場景其實很豐富，每一個小的模塊包含幾十個模型。在此之上我們做了一層封裝，將相似模型做一個統一接口，做成 task 或 actor，稱為統一融合引擎。我們也做了一些調度，比如 Ray server （在線服務），Ray data （串行 pipeline），對于自定義等更復雜的場景就用 Ray 的 workflow，僅需要用原生 Python 語言去嵌入各個節點。

很多場景下并非一次調整就能得到理想結果，而是需要基于反饋反復調整，進行多模型優化。

基于上述基礎架構，可以實現基于 Ray 進行積木化組裝模型應用。

基于 Ray 可以實現：

• 高效構建：Python 友好，可以統一編程語言；分門別類，統一接口；統一調度，減少構建 pipeline 成本。
• 高性能執行：可以彈性自動擴縮容；Stream overlap 執行；融合單節點、單模型優化。
• 低成本運維：既能解決私有化也能解決云原生，并且對業務屏蔽，甚至無需了解 Ray 就可以進行模型推理。