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

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

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

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

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

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

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

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

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

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

我們的初期主要合作者

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

目前，DeepSpeed4Science很榮幸地支持來自微軟研究院AI4Science、微軟WebXT/Bing、美國能源部國家實驗室和多所大學(xué)的幾個關(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)建一個統(tǒng)一的大規(guī)模基礎(chǔ)模型，以支持自然科學(xué)發(fā)現(xiàn)，支持多種輸入、多個科學(xué)領(lǐng)域（例如，藥物、材料、生物學(xué)、健康等）和計算任務(wù)。

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

ClimaX，微軟研究院AI4Science

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

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

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

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

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

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

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

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

微軟天氣，微軟WebXT/Bing

圖5：微軟降水預(yù)報（每4分鐘一次對接下來4小時進(jìn)行預(yù)測）。

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

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

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

外部合作者

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

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

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

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

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

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

合作展示

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

zhihu.com/people/deepspeed