一、從 2024 年諾貝爾化學(xué)獎?wù)勂?/span>

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

蛋白質(zhì)的業(yè)務(wù)價值非常大，幾乎所有生物公司都無法繞開這個領(lǐng)域，都會做一些蛋白質(zhì)相關(guān)的應(yīng)用或者模型二次開發(fā)。同時我們也發(fā)現(xiàn)這個領(lǐng)域?qū)τ谟嬎阗Y源的消耗是非常大的，一個模型就需要消耗多個 CPU 或者多張卡來處理一個請求，其并發(fā)延遲遠(yuǎn)超傳統(tǒng)的推薦搜索模型。舉個例子，一個蛋白質(zhì)結(jié)構(gòu)生成預(yù)測，如果需要預(yù)測一個 100*1000 個序列的混合物，就得 30 分鐘，一天只能計算幾十的。蛋白質(zhì)的序列生成就更加耗時，可能一個顯卡設(shè)計一個就需要幾個小時。所以在蛋白質(zhì)研究領(lǐng)域，計算能力有很大的提高空間，接下來就將分兩部分來介紹這個場景的優(yōu)化。

二、加速蛋白質(zhì)結(jié)構(gòu)預(yù)測性能

首先來介紹 DeepMind 團(tuán)隊關(guān)于加速蛋白質(zhì)結(jié)構(gòu)預(yù)測的工作，其主要思想是基于 Ray 的 Workflow，實現(xiàn)高效異構(gòu)調(diào)度。

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

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

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

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

上圖展示了我們設(shè)計的架構(gòu)，采用 Ray 的 Workflow 方案。這個架構(gòu)有兩大特點(diǎn)，一是流式調(diào)度，二是高效構(gòu)圖。我們將其拆分成幾個節(jié)點(diǎn)，都是對應(yīng)的 flow 的 node 節(jié)點(diǎn)，可以靈活構(gòu)圖，靈活構(gòu)圖的好處就是每個節(jié)點(diǎn)均可插拔，每個節(jié)點(diǎn)可以無縫替換。

第二個核心設(shè)計是，考慮到 MSA 的 node 預(yù)處理非常慢，因此設(shè)計為 Actor Pool 初始化一個節(jié)點(diǎn)，預(yù)處理做好，推理時其已經(jīng)是一個預(yù)熱好的節(jié)點(diǎn)，這樣可以從 30 分鐘優(yōu)化到 2 分鐘，效果非?？捎^。

第三步就到了一個 GPU 節(jié)點(diǎn)，類似傳統(tǒng)語言模型，將其作為 model node，如果機(jī)器足夠多，可以自動擴(kuò)縮容，無需人為定義資源。

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

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

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

回到業(yè)務(wù)的核心，利用 Ray 可提升執(zhí)行效率，并且由于 Ray 的可擴(kuò)展性，再加上 Ray 整個架構(gòu)是一個非常好的解耦架構(gòu)，因此可以降低運(yùn)維成本，提升合作開發(fā)的效率。另外其底層還是 K8S，我們不需要關(guān)注 GPU 和 CPU 節(jié)點(diǎn)的情況，對于創(chuàng)業(yè)公司（人員不足的情況下）是非常友好的。

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

三、加速蛋白質(zhì)生成設(shè)計性能

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

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

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

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

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

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

整個運(yùn)行時間 Baseline 是 2 個小時級別，優(yōu)化后是 1-2 分鐘，對生命科學(xué)領(lǐng)域加速模型處理的意義是十分重大的。

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

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

四、Ray 融合計算架構(gòu)

基于上述背景，我們設(shè)計了基于 Ray 的融合計算架構(gòu)，將所有事情都在 Ray 中完成，每個接口可插拔，底層可以統(tǒng)一，優(yōu)化就可以同步進(jìn)行，具體架構(gòu)如下圖所示。

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

很多場景下并非一次調(diào)整就能得到理想結(jié)果，而是需要基于反饋反復(fù)調(diào)整，進(jìn)行多模型優(yōu)化。

基于上述基礎(chǔ)架構(gòu)，可以實現(xiàn)基于 Ray 進(jìn)行積木化組裝模型應(yīng)用。

基于 Ray 可以實現(xiàn)：

• 高效構(gòu)建：Python 友好，可以統(tǒng)一編程語言；分門別類，統(tǒng)一接口；統(tǒng)一調(diào)度，減少構(gòu)建 pipeline 成本。
• 高性能執(zhí)行：可以彈性自動擴(kuò)縮容；Stream overlap 執(zhí)行；融合單節(jié)點(diǎn)、單模型優(yōu)化。
• 低成本運(yùn)維：既能解決私有化也能解決云原生，并且對業(yè)務(wù)屏蔽，甚至無需了解 Ray 就可以進(jìn)行模型推理。