在接下來(lái)的十年中，深度學(xué)習(xí)可能會(huì)徹底改變自然科學(xué)，增強(qiáng)我們對(duì)自然現(xiàn)象進(jìn)行建模和預(yù)測(cè)的能力。

這可能預(yù)示著科學(xué)探索的新時(shí)代，為從藥物開(kāi)發(fā)到可再生能源的各個(gè)領(lǐng)域帶來(lái)重大進(jìn)展。

對(duì)此，微軟DeepSpeed團(tuán)隊(duì)啟動(dòng)了一個(gè)名為DeepSpeed4Science的新計(jì)劃，旨在通過(guò)AI系統(tǒng)技術(shù)創(chuàng)新幫助領(lǐng)域?qū)＜医怄i當(dāng)今最大的科學(xué)之謎。

DeepSpeed系統(tǒng)是由微軟開(kāi)發(fā)的業(yè)界領(lǐng)先的開(kāi)源AI系統(tǒng)框架，它為各種AI硬件上的深度學(xué)習(xí)訓(xùn)練和推理提供了前所未有的規(guī)模和速度。

圖1：DeepSpeed4Science方法概述：專(zhuān)為加速科學(xué)發(fā)現(xiàn)和應(yīng)對(duì)其復(fù)雜性而量身定制的AI系統(tǒng)技術(shù)開(kāi)發(fā)。

圖1展示了我們對(duì)DeepSpeed4Science這一新計(jì)劃的基本方法。

通過(guò)利用DeepSpeed當(dāng)前的技術(shù)方案（訓(xùn)練、推理和壓縮）作為基礎(chǔ)技術(shù)推動(dòng)器，DeepSpeed4Science將創(chuàng)建一套專(zhuān)為加速科學(xué)發(fā)現(xiàn)而量身定制的AI系統(tǒng)技術(shù)，以應(yīng)對(duì)其獨(dú)特的復(fù)雜性，超越用于加速通用大型語(yǔ)言模型（LLMs）的常見(jiàn)技術(shù)方法。

在這篇博客中，我們展示了DeepSpeed4Science如何幫助解決結(jié)構(gòu)生物學(xué)研究中的兩個(gè)關(guān)鍵AI系統(tǒng)挑戰(zhàn)：

（1）解決了以Evoformer為中心的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)模型中的內(nèi)存爆炸問(wèn)題，以及

（2）為更好地理解引發(fā)大流行的疾病的進(jìn)化提供AI模型長(zhǎng)序列支持。

我們的初期主要合作者

DeepSpeed4Science的新系統(tǒng)技術(shù)可以用于很多推動(dòng)科學(xué)邊界的標(biāo)志性模型，賦能AI驅(qū)動(dòng)的科學(xué)發(fā)現(xiàn)。

目前，DeepSpeed4Science很榮幸地支持來(lái)自微軟研究院AI4Science、微軟WebXT/Bing、美國(guó)能源部國(guó)家實(shí)驗(yàn)室和多所大學(xué)的幾個(gè)關(guān)鍵科學(xué)模型。

內(nèi)部合作伙伴

科學(xué)基礎(chǔ)模型（Scientific Foundation Model，SFM），微軟研究院AI4Science

圖2：科學(xué)基礎(chǔ)模型（Scientific Foundation Model，SFM）及其當(dāng)前探索：Distributional Graphormer

科學(xué)基礎(chǔ)模型（SFM）旨在創(chuàng)建一個(gè)統(tǒng)一的大規(guī)?；A(chǔ)模型，以支持自然科學(xué)發(fā)現(xiàn)，支持多種輸入、多個(gè)科學(xué)領(lǐng)域（例如，藥物、材料、生物學(xué)、健康等）和計(jì)算任務(wù)。

DeepSpeed4Science合作伙伴關(guān)系將為SFM團(tuán)隊(duì)提供新的訓(xùn)練和推理技術(shù)，以支持他們的新生成AI方法（例如Distributional Graphormer）這樣的項(xiàng)目進(jìn)行持續(xù)研究。

ClimaX，微軟研究院AI4Science

圖3：ClimaX是第一個(gè)設(shè)計(jì)用于執(zhí)行各種天氣和氣候建模任務(wù)的基礎(chǔ)模型

我們的氣候正在發(fā)生變化，導(dǎo)致極端天氣事件的頻率增加。為了減輕負(fù)面影響，預(yù)測(cè)這些事件將發(fā)生的地方變得越來(lái)越重要。

ClimaX是第一個(gè)設(shè)計(jì)用于執(zhí)行各種天氣和氣候建模任務(wù)的基礎(chǔ)模型。它可以吸收許多具有不同變量和分辨率的數(shù)據(jù)集以提高天氣預(yù)報(bào)的準(zhǔn)確性。

DeepSpeed4Science正在為ClimaX創(chuàng)建新的系統(tǒng)支持和加速策略，以高效地預(yù)訓(xùn)練/微調(diào)更大的基礎(chǔ)模型，同時(shí)處理非常大的高分辨率圖像數(shù)據(jù)（例如，數(shù)十到數(shù)百PB）和長(zhǎng)序列。

AI驅(qū)動(dòng)的第一性原理分子動(dòng)力學(xué)（AI Powered Ab Initio Molecular Dynamics，AI2MD），微軟研究院AI4Science

圖4：一百萬(wàn)步的分子動(dòng)力學(xué)模擬：RBD-蛋白（RBD-protein）與蛋白抑制劑（protein inhibitor）相互作用。

這個(gè)項(xiàng)目模擬了使用AI驅(qū)動(dòng)的力場(chǎng)模型進(jìn)行近似第一性原理計(jì)算精度的大型（百萬(wàn)原子）分子系統(tǒng)的動(dòng)態(tài)模擬，同時(shí)保持了經(jīng)典分子動(dòng)力學(xué)的效率和可擴(kuò)展性。這些模擬足夠高效，可以生成足夠長(zhǎng)的軌跡來(lái)觀察化學(xué)上有意義的事件。

通常，這個(gè)過(guò)程需要數(shù)百萬(wàn)甚至數(shù)十億的推理步驟。這對(duì)優(yōu)化圖神經(jīng)網(wǎng)絡(luò)（GNN）+ LLM模型的推理速度提出了重大挑戰(zhàn)，DeepSpeed4Science將為此提供新的加速策略。

微軟天氣，微軟WebXT/Bing

圖5：微軟降水預(yù)報(bào)（每4分鐘一次對(duì)接下來(lái)4小時(shí)進(jìn)行預(yù)測(cè)）。

微軟天氣提供精確的天氣信息，幫助用戶(hù)為他們的生活方式、健康、工作和活動(dòng)做出更好的決策——包括每小時(shí)多次更新的準(zhǔn)確的10天全球天氣預(yù)報(bào)。

此前，微軟天氣受益于DeepSpeed技術(shù)，加速了他們的多GPU訓(xùn)練環(huán)境。

現(xiàn)在，DeepSpeed4Science正在與微軟WebXT天氣預(yù)報(bào)團(tuán)隊(duì)合作，進(jìn)一步增強(qiáng)微軟天氣預(yù)報(bào)服務(wù)的最新功能和改進(jìn)。

外部合作者

DeepSpeed4Science的旅程始于兩個(gè)開(kāi)創(chuàng)性的基于LLM的結(jié)構(gòu)生物學(xué)研究AI模型：來(lái)自哥倫比亞大學(xué)的OpenFold，一個(gè)開(kāi)源的高保真蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)模型；以及來(lái)自阿貢國(guó)家實(shí)驗(yàn)室的GenSLMs，一個(gè)獲得ACM戈登貝爾獎(jiǎng)的用于學(xué)習(xí)SARS-CoV-2（COVID-19）基因組的進(jìn)化的語(yǔ)言模型。

作為此次發(fā)布的特色展示，它們代表了當(dāng)今AI驅(qū)動(dòng)的結(jié)構(gòu)生物學(xué)研究面臨的兩個(gè)常見(jiàn)AI系統(tǒng)挑戰(zhàn)。我們將在下一節(jié)中討論DeepSpeed4Science如何賦能這些科學(xué)研究。

此外，DeepSpeed4Science最近擴(kuò)大了其范圍，以支持更多樣的科學(xué)模型。

例如，在我們與阿貢國(guó)家實(shí)驗(yàn)室合作訓(xùn)練Aurora Exascale系統(tǒng)上的萬(wàn)億參數(shù)科學(xué)模型的工作中，DeepSpeed4Science技術(shù)將幫助他們達(dá)到這一關(guān)鍵任務(wù)所需的性能要求和可擴(kuò)展性。

此外，通過(guò)與橡樹(shù)嶺國(guó)家實(shí)驗(yàn)室和國(guó)家癌癥研究所（NCI）合作進(jìn)行癌癥監(jiān)測(cè)，DeepSpeed4Science將幫助從非結(jié)構(gòu)化的臨床文本中高保真地提取和分類(lèi)信息，以供MOSSAIC項(xiàng)目使用。

Brookhaven國(guó)家實(shí)驗(yàn)室還將采用DeepSpeed4Science技術(shù)，支持使用LLMs開(kāi)發(fā)大型數(shù)字雙胞胎模型，以便為清潔能源研究產(chǎn)生更真實(shí)的模擬數(shù)據(jù)。您可以在deepspeed4science.ai上找到有關(guān)我們外部合作者及其科學(xué)任務(wù)的更多詳細(xì)信息。

合作展示

展示（I）：DeepSpeed4Science通過(guò)DS4Sci_EvoformerAttention消除以Evoformer為中心的結(jié)構(gòu)生物學(xué)模型的內(nèi)存爆炸問(wèn)題

圖6：在訓(xùn)練過(guò)程中OpenFold對(duì)PDB鏈7B3A_A的預(yù)測(cè)

OpenFold是DeepMind的AlphaFold2的開(kāi)源社區(qū)再現(xiàn)，使其可以在新數(shù)據(jù)集上訓(xùn)練或微調(diào)AlphaFold2。

研究人員已經(jīng)使用它從頭開(kāi)始重新訓(xùn)練AlphaFold2，生成新的模型參數(shù)集，研究AlphaFold2的早期訓(xùn)練階段（圖6），并開(kāi)發(fā)新的蛋白質(zhì)折疊系統(tǒng)。

圖7：在OpenFold中，對(duì)多序列比對(duì)（MSA）Attention內(nèi)核（包含偏差）變體的訓(xùn)練峰值內(nèi)存需求。（左）使用在AlphaFold2中的EvoformerAttention的原始OpenFold實(shí)現(xiàn)。對(duì)于這些類(lèi)型的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)模型，在訓(xùn)練/推理中的內(nèi)存爆炸問(wèn)題是常見(jiàn)的。最先進(jìn)的FlashAttention無(wú)法有效支持這樣的Attention變體。（右）DeepSpeed4Science的一種新解決方案DS4Sci_EvoformerAttention在不影響模型品質(zhì)的條件下顯著地減少了OpenFold的訓(xùn)練峰值內(nèi)存需求（最多13倍）。

盡管OpenFold有使用最先進(jìn)的系統(tǒng)技術(shù)進(jìn)行性能和內(nèi)存優(yōu)化，但從頭開(kāi)始訓(xùn)練AlphaFold2仍然在計(jì)算上很昂貴。目前階段的模型參數(shù)很小，只有9300萬(wàn)個(gè)參數(shù)，但它包含了幾個(gè)需要非常大的中間內(nèi)存的特殊Attention變體。

在標(biāo)準(zhǔn)AlphaFold2訓(xùn)練的「微調(diào)」階段，只是這些變體中的其中一個(gè)在半精度下就生成了超過(guò)12GB的張量，使其峰值內(nèi)存要求遠(yuǎn)遠(yuǎn)超過(guò)了相同大小的語(yǔ)言模型。

即使使用像activation checkpointing和DeepSpeed ZeRO優(yōu)化這樣的技術(shù)，這種內(nèi)存爆炸問(wèn)題仍然嚴(yán)重限制了可訓(xùn)練模型的序列長(zhǎng)度和MSA深度。

此外，近似策略可能會(huì)顯著影響模型的準(zhǔn)確性和收斂性，同時(shí)仍然導(dǎo)致內(nèi)存爆炸，如圖7左側(cè)（橙色）所示。

為了應(yīng)對(duì)結(jié)構(gòu)生物學(xué)研究（例如，蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)和平衡分布預(yù)測(cè)）中的這一常見(jiàn)系統(tǒng)挑戰(zhàn)，DeepSpeed4Science通過(guò)為這類(lèi)科學(xué)模型中廣泛出現(xiàn)的注意力變體（即EvoformerAttention）設(shè)計(jì)定制的精確注意力內(nèi)核來(lái)解決這一內(nèi)存效率問(wèn)題。

具體來(lái)說(shuō)，我們?cè)O(shè)計(jì)了一套由復(fù)雜的融合/矩陣分塊策略和動(dòng)態(tài)內(nèi)存減少方法而組成的高內(nèi)存效率DS4Sci_EvoformerAttention內(nèi)核，作為高質(zhì)量機(jī)器學(xué)習(xí)模塊供更廣泛的生物學(xué)研究社區(qū)使用。

通過(guò)整合到OpenFold中，這些定制內(nèi)核在訓(xùn)練期間提供了顯著的加速，并顯著減少了模型的訓(xùn)練和推理的峰值內(nèi)存需求。

這使得OpenFold可以用更大、更復(fù)雜的模型，使用更長(zhǎng)的序列在更廣泛的硬件上進(jìn)行實(shí)驗(yàn)。關(guān)于這項(xiàng)技術(shù)的詳細(xì)信息可以在這里找到。

展示（II）：DeepSpeed4Science通過(guò)系統(tǒng)和算法方法為基因組基礎(chǔ)模型（例如，GenSLMs）提供長(zhǎng)序列支持

圖8：GenSLMs：獲2022年ACM 戈登貝爾獎(jiǎng)的COVID基因組模型（基于GPT-NeoX的25B/33B模型）。它用于學(xué)習(xí)描述SARS-CoV-2基因組生物學(xué)意義的潛在空間。這個(gè)GIF展示了一個(gè)重要的蛋白質(zhì)家族蘋(píng)果酸脫氫酶（malate dehydrogenase）的根據(jù)重要特征（如序列長(zhǎng)度和GC含量（核酸鳥(niǎo)嘌呤和胞嘧啶的含量與腺嘌呤和胸腺嘧啶的比率。它測(cè)量DNA鏈抵抗熱的能力））著色的潛在空間的投影。

GenSLMs，一個(gè)來(lái)自阿貢國(guó)家實(shí)驗(yàn)室的2022年ACM 戈登貝爾獎(jiǎng)獲獎(jiǎng)的基因組模型，可以通過(guò)大型語(yǔ)言模型（LLMs）的基因組數(shù)據(jù)訓(xùn)練來(lái)學(xué)習(xí)SARS-CoV-2（COVID-19）基因組的進(jìn)化。它旨在改變?nèi)绾巫R(shí)別和分類(lèi)引發(fā)大流行的病毒（特別是SARS-CoV-2）的新變種。

GenSLMs代表了第一批可以泛化到其他預(yù)測(cè)任務(wù)的基因組基礎(chǔ)模型。對(duì)潛在空間的良好理解可以幫助GenSLMs處理超出僅僅是病毒序列的新領(lǐng)域，并擴(kuò)展它們模擬細(xì)菌病原體甚至真核生物的能力（例如，理解功能、途徑成員資格和進(jìn)化關(guān)系等事物）。

為了實(shí)現(xiàn)這一科學(xué)目標(biāo)，GenSLMs和類(lèi)似的模型需要非常長(zhǎng)的序列支持用于訓(xùn)練和推理，這超出了像FlashAttention這樣的通用LLM的長(zhǎng)序列策略。

通過(guò)DeepSpeed4Science的新設(shè)計(jì)，科學(xué)家現(xiàn)在可以構(gòu)建和訓(xùn)練具有顯著更長(zhǎng)的上下文窗口的模型，允許他們探索以前無(wú)法訪問(wèn)的關(guān)系。

圖9：由不同框架在不同規(guī)模下支持的兩個(gè)GenSLMs模型的最大序列長(zhǎng)度。使用NVIDIA DGX，每個(gè)節(jié)點(diǎn)有八個(gè)40G A100 GPU

具體在系統(tǒng)層面，我們發(fā)布了包括長(zhǎng)序列支持和其他新優(yōu)化的最新的Megatron-DeepSpeed框架。

科學(xué)家現(xiàn)在可以通過(guò)我們新添加的內(nèi)存優(yōu)化技術(shù)（如注意力掩碼異步處理和位置碼分割）、張量并行、流水線并行、序列并行、基于ZeRO的數(shù)據(jù)并行和模型狀態(tài)異步處理等技術(shù)的協(xié)同組合，用更長(zhǎng)的序列訓(xùn)練他們的GenSLMs等大型科學(xué)模型。

圖9展示了我們的新版本使GenSLMs的25B和33B模型的最長(zhǎng)序列長(zhǎng)度分別比之前的Megatron-DeepSpeed版本增加了12倍和14倍。

在支持的序列長(zhǎng)度方面，這個(gè)新Megatron-DeepSpeed框架也顯著地超過(guò)了NVIDIA的Megatron-LM（對(duì)于25B和33B模型分別高達(dá)9.8倍和9.1倍）。

例如，阿貢實(shí)驗(yàn)室團(tuán)隊(duì)的GenSLMs 25B模型在64個(gè)GPU上的原始序列長(zhǎng)度為42K，而現(xiàn)在可以用512K的核苷酸序列進(jìn)行訓(xùn)練。這在不損失準(zhǔn)確性的條件下大大提高了模型質(zhì)量和科學(xué)發(fā)現(xiàn)的范圍。

對(duì)于那些更喜歡相對(duì)位置編碼技術(shù)這樣的算法策略的領(lǐng)域科學(xué)家，這個(gè)新版本也進(jìn)行了集成。

轉(zhuǎn)載自微軟DeepSpeed組官方知乎賬號(hào)：

zhihu.com/people/deepspeed