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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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