思考一下,聯(lián)邦學習可以訓練大語言模型嗎?
1、概述
隨著大語言模型(Large Language Model,LLM)的火速發(fā)展,關于大語言模型對人工智能產(chǎn)業(yè)發(fā)展的影響引發(fā)了越來越多的討論。一種觀點認為,大語言模型的發(fā)展摧毀了人工智能初創(chuàng)企業(yè)的發(fā)展之路,因為大語言模型參數(shù)多,所需要的算力規(guī)模大,所依賴的訓練數(shù)據(jù)規(guī)模也大。大模型、大參數(shù)、大數(shù)據(jù)實際都集中在大的人工智能企業(yè),從而導致初創(chuàng)企業(yè)的機會越來越少。另外一種觀點則相反,他們認為,大語言模型的發(fā)展一定程度成促進了人工智能在多個領域中的廣泛發(fā)展,例如可以直接在大語言模型的基礎上利用私有數(shù)據(jù)搭建一些垂直領域的大語言模型,可以直接將大語言模型應用在不同的業(yè)務場景中等等。
我們在這篇文章中不討論上述兩種觀點,我們關注的是第二種觀點中提到的應用方式:如何在保證數(shù)據(jù)隱私性的前提下,利用私有數(shù)據(jù)訓練大語言模型,從而滿足在垂直領域中的應用需求?這一點也將是 LLM 促進人工智能初創(chuàng)企業(yè)發(fā)展的關鍵。具體的,我們關注聯(lián)邦學習是否可以用來訓練大語言模型?
1.1 大語言模型結構及需要的資源
由于 Transformer 架構具有出色的可并行性和容量,它已經(jīng)成為開發(fā)各種 LLM 事實上的主干網(wǎng)絡,使得語言模型的規(guī)模達到數(shù)千或數(shù)億的參數(shù)成為可能。一般來說,現(xiàn)有 LLM 的主流架構可以大致分為三大類型,即編碼器 - 解碼器(encoder-decoder)、因果解碼器(causal decoder)和前綴解碼器(prefix decoder),如圖 1 所示 [8]。
圖片
圖 1. 三種主流架構中的注意力模式比較。其中,藍色、綠色、黃色和灰色的圓角矩形分別表示前綴標記之間的注意力、前綴和目標標記之間的注意力、目標標記之間的注意力和掩碼的注意力。
- 編碼器 - 解碼器架構。Vanilla Transformer 模型是建立在編碼器 - 解碼器結構上的,它由兩個堆疊的 Transformer 塊組成,分別作為編碼器和解碼器。編碼器采用堆疊的多頭自注意力層對輸入序列進行編碼,以生成其潛在的表征,而解碼器則對這些表征進行交叉注意力,并自動生成目標序列。編碼器 - 解碼器(如 T5 和 BART)已經(jīng)在各種 NLP 任務上顯示出有效性。到目前為止,只有少數(shù)的 LLM 是基于編碼器 - 解碼器架構建立的,如 Flan-T5。
- 因果解碼器架構。因果解碼器的結構包含了單向的注意力掩碼,以保證每個輸入標記只能注意到過去的標記和自己。輸入和輸出標記以同樣的方式通過解碼器進行處理。作為該架構的代表性語言模型,GPT 系列模型是基于因果解碼器架構開發(fā)的。特別是 GPT-3 成功地證明了這一架構的有效性,也顯示了 LLM 驚人的內涵學習能力。有趣的是,GPT-1 和 GPT-2 并沒有表現(xiàn)出像 GPT-3 那樣優(yōu)越的能力,看來縮放在提高這個模型架構的模型容量方面發(fā)揮了重要作用。到目前為止,因果解碼器已被各種現(xiàn)有的 LLM 廣泛采用為 LLM 的架構,如 OPT、BLOOM 和 Gopher。
- 前綴解碼器架構。前綴解碼器架構(又稱非因果解碼器)修改了因果解碼器的掩碼機制,以便能夠對前綴標記進行雙向關注,而只對生成的標記進行單向關注。這樣,像編碼器 - 解碼器結構一樣,前綴解碼器可以對前綴序列進行雙向編碼,并逐一自動預測輸出標記,在編碼和解碼期間共享相同的參數(shù)。與其從頭開始預訓練,一個實用的建議是不斷訓練因果解碼器,然后將其轉換為前綴解碼器,以加速收斂,例如,U-PaLM 是由 PaLM 衍生出來的?,F(xiàn)有的基于前綴解碼器的代表性 LLM 包括 GLM130B 和 U-PaLM。
但是,在大語言模型發(fā)展過程中,越來越多的人提出了疑慮,主要集中在以下幾個方面:
- 隱私不應受到損害。最先進的模型只能通過通用的黑盒 API 訪問。這些 API 在向集中式 LLM 供應商傳輸數(shù)據(jù)時,會引起公司的隱私問題。例如,據(jù)報道,三星通過 ChatGPT 泄露了自己的秘密,突出了與此類 API 相關的風險。
- 專有的和訓練有素的模型是重要的知識產(chǎn)權。對于企業(yè)來說,最能從人工智能中受益的問題和數(shù)據(jù)集往往是敏感和專有的。這就排除了對公共模型的微調,因此業(yè)務最好是利用內部部署的模型進行微調,或在企業(yè)的防火墻內部署模型。
- 成本問題。目前,LLM 的訓練需要數(shù)以千計的 GPU 節(jié)點,成本高達十億級美元。模型可能需要平衡不同的權衡,如略微影響模型的性能以減少云的成本。擁有模型使客戶可以方便地對其進行微調和再訓練。例如,他們可以刪除多余的模型參數(shù),以實現(xiàn)特定領域任務的高性能,同時最大限度地降低云成本。
由此,引發(fā)出我們這篇文章討論的方法:FL+LLM,即引入聯(lián)邦學習來訓練大語言模型,從而為企業(yè)商業(yè)用戶提供眾多優(yōu)勢,在模型規(guī)模和性能、隱私、效率、云計算成本和勞動力成本方面大大增強企業(yè)使用大型模型的能力。
1.2 聯(lián)邦學習回顧
聯(lián)邦學習(Federated learning,F(xiàn)L)是一種機器學習環(huán)境,在這種環(huán)境下,多個客戶端(如移動設備或整個組織)在中央服務器(如服務提供商)的協(xié)調下協(xié)同訓練一個模型,同時保持訓練數(shù)據(jù)的分散性。FL 體現(xiàn)了集中數(shù)據(jù)收集和最小化的原則,可以減輕傳統(tǒng)的集中式機器學習和數(shù)據(jù)科學方法帶來的許多隱私、安全性風險和成本。因此,F(xiàn)L 是一種有效的高性能計算范式,也被看作是滿足數(shù)據(jù)隱私性要求的分布式訓練方法。目前,F(xiàn)L 已經(jīng)在一些應用場景中成功使用,包括一些消費端設備,例如 GBoard、關鍵詞發(fā)現(xiàn),以及制藥、醫(yī)學研究、金融、制造業(yè)等等。此外,也出現(xiàn)了大量 FL 開發(fā)工具和平臺,例如 Tensorow Federated, LEAF, PaddleFL 和 PySy,進一步推動 FL 的發(fā)展。
在聯(lián)邦學習框架中,中央服務器保存初始化可共享的全局數(shù)據(jù)。各個客戶端(參與者、邊緣設備)保存本地數(shù)據(jù),并根據(jù)本地數(shù)據(jù)訓練本地機器學習模型??蛻舳烁鶕?jù)一定的通信機制向中央服務器傳輸模型參數(shù)等數(shù)據(jù)(不會傳輸完整的客戶端原始數(shù)據(jù)),中央服務器匯聚各客戶端上載數(shù)據(jù)后訓練構建全局模型,各個客戶端在整個聯(lián)邦學習機制中身份和地位相同。聯(lián)邦學習有效解決了兩方或多方數(shù)據(jù)使用實體(客戶端)在不貢獻出數(shù)據(jù)的情況下的數(shù)據(jù)共同使用問題,解決了數(shù)據(jù)孤島問題。此外,在各個客戶端數(shù)據(jù)特征對齊的前提下,聯(lián)邦學習的全局模型能夠獲得與數(shù)據(jù)集中式存儲相同的建模效果。
以下是四個非常重要的聯(lián)邦學習算法:
1. FedAvg(Federated Averaging)[1]:FedAvg 是最經(jīng)典的聯(lián)邦學習算法,由 Google 于 2016 年提出。該算法采用梯度下降的方式對模型參數(shù)進行迭代更新。在 FedAvg 中,每個客戶端(例如,手機或其他設備)首先在本地用自己的數(shù)據(jù)訓練模型,然后將本地模型權重發(fā)送至中央服務器。中央服務器收集所有客戶端的權重更新,計算權重的平均值,并將更新后的權重分發(fā)回各個客戶端。這個過程在多輪中進行,直到達到收斂。
2. FedProx(Federated Proximal)[2]:FedProx 是一種改進的聯(lián)邦學習算法,用于解決聯(lián)邦學習中的非獨立同分布(non-IID)數(shù)據(jù)和設備異質性問題。這些問題會導致訓練質量下降,尤其是在訓練數(shù)據(jù)不均勻分布的情況下。FedProx 通過在優(yōu)化目標中添加一個 proximal term 來解決這個問題,這使得模型在更新時更加關注與全局模型權重接近的本地模型權重。這樣可以減輕非獨立同分布數(shù)據(jù)和設備異質性帶來的負面影響,提高模型性能。
3. SCAFFOLD(Federated Learning via Stochastic Controlled Averaging)[3]:SCAFFOLD 是一種高效的聯(lián)邦學習算法,主要關注梯度稀疏性和設備間的通信。SCAFFOLD 通過在每輪中維護一個控制變量(control variate)來減少通信量。這個控制變量可以根據(jù)全局模型的梯度變化進行調整,從而使每個設備之間的通信更為緊湊。SCAFFOLD 在通信方面更加高效,可以在不犧牲模型性能的情況下減少通信成本。
4. FedNova(Federated Learning via Normalized Averaging)[4]:FedNova 的核心思想是采用歸一化平均方法消除目標不一致性,同時保持快速的誤差收斂。在每輪迭代中,在各個客戶端設備上進行本地訓練,并對本地模型參數(shù)進行歸一化,然后將歸一化后的參數(shù)發(fā)送至中央服務器。中央服務器在收集所有客戶端的參數(shù)更新后,對其進行歸一化平均,從而減少目標不一致性。之后,中央服務器將更新后的模型參數(shù)分發(fā)回各個客戶端,以便在下一輪迭代中繼續(xù)訓練。FedNova 的優(yōu)點在于,它能夠在保持快速誤差收斂的同時,消除由于數(shù)據(jù)異質性導致的目標不一致性。這使得 FedNova 算法在處理非獨立同分布數(shù)據(jù)和設備異質性方面表現(xiàn)出較好的性能。
2、聯(lián)邦學習結構來訓練大語言模型的難點
近年來,LLM 和 FL 的研究都取得了一定的進展。但是,在目前的技術和資源條件下,直接 FL+LLM 還面臨一系列的問題。
2.1 數(shù)據(jù)方面
LLM 的訓練依托的是超大規(guī)模的數(shù)據(jù)量。Chinchilla 探討了訓練 LLM 需要多少數(shù)據(jù) [9]。2020 年 5 月,OpenAI 展示了他們的 LLM data scaling laws(也稱為 Kaplan scaling laws):300B tokens 可用于訓練參數(shù)量大小為 175B LLM。2022 年 9 月,DeepMind 為 “數(shù)據(jù)最優(yōu)” LLM 找到了新 data scaling laws(也稱為 Chinchilla 或 Hoffman scaling laws):應該使用 1,400B (1.4T) tokens 來訓練參數(shù)量大小為 70B LLM 最佳。
圖片
表 1. 與 Chinchilla 數(shù)據(jù)優(yōu)化模型一致所需的數(shù)據(jù)集大小 [9]
由表 1 中 Chinchilla 給出的數(shù)據(jù)可以看出,應用聯(lián)邦學習結構來訓練 LLM 勢必要將海量的數(shù)據(jù)量分布在不同的客戶端設備中,如果客戶端數(shù)量較少,那就造成每個客戶端中存儲、承載的數(shù)據(jù)量依然較大,對每個客戶端的計算性能要求都非常高,這與聯(lián)邦學習架構的設計初衷是不一致的;如果客戶端數(shù)量非常龐大,每個客戶端中存儲的數(shù)據(jù)量適中,那對于大量客戶端之間的協(xié)調調用機制、多方計算能力和資源分配不均的資源不均衡問題就會非常突出,不均衡情況下的客戶端激勵機制也尤為關鍵,會對聯(lián)邦學習的中央模型效果、訓練效率都有很大的影響。
2.2 客戶端設備方面
聯(lián)邦學習架構下的分布式系統(tǒng)中,單個客戶端如何容納整個 LLM 模型?在已有的聯(lián)邦學習平臺中,特別是一些實用場景下的聯(lián)邦學習平臺,分布式的客戶端設備大多為手機、平板等,這些設備根本不可能容納整個 LLM 模型或為模型訓練提供足夠的內存資源。如果將這些設備完全排除在訓練過程之外,也與聯(lián)邦學習架構的設計初衷不一致。但在客戶端無法容納整個模型、又不能將客戶端完全排除在訓練過程之外的情況下,如何對 LLM 進行并行式的訓練呢?
微軟在一篇文章中對 LLaMA 目前對設備內存的占用量進行了分析 [10]。在模型訓練過程中,大部分的內存開銷用于 model states (模型參數(shù)、優(yōu)化器的狀態(tài)、梯度)。剩下的內存開銷則來自于 residual states(激活值、臨時區(qū)緩存和內存碎片)。
以 LLaMA 7B 模型為例,模型參數(shù)所需的內存是 (參數(shù)量 * 每個參數(shù)的內存):fp32 精度: 7*4 bytes = 28GB;fp16 精度:7*2 bytes = 14GB;int8 精度:7*1 bytes = 7GB;混合精度(fp16/32)訓練:存儲 fp16 精度 + fp32 精度 = 14GB +28GB = 42GB。
梯度所需的內存(計算方式同模型參數(shù)一樣,參數(shù)量 * 每個參數(shù)的內存):fp32 精度: 7*4 bytes = 28GB;fp16 精度:7*2 bytes= 14GB;int8 精度:7*1 bytes = 7GB;混合精度(fp16/32)訓練:只存儲 fp16 精度 = 14GB。
優(yōu)化器所需的內存(參數(shù)量 * 每個參數(shù)的內存 * 2,以 Adam 為例):fp32 精度: 7*4 bytes * 2 = 56GB;fp16 精度:7*2 bytes * 2 = 28GB;int8 精度:7*1 bytes * 2 = 14GB;混合精度(fp16/32)訓練:存儲 fp32 精度 = 56 G。
總的來說,訓練 LLaMa 模型總的內存需求為,fp32 精度: 28+28+56 = 112GB;fp16 精度:14+14+28 = 56GB;int8 精度:7+7+14 = 28G;混合精度(fp16/32)訓練: 42+14+56 = 112GB。
LLaMA 7B 的規(guī)模相當于目前的 LLM 來說非常小了,但即使 LLaMA 7B 的內存消耗量對一些手機、平板來說也是難以滿足的,這就對 FL 架構下的客戶端設備選擇構成了巨大的約束。
2.3 訓練算法方面
已有的聯(lián)邦學習訓練算法對于 LLM 來說是低效的。如上所述,聯(lián)邦學習的工作過程是:“各個客戶端根據(jù)本地數(shù)據(jù)訓練本地模型??蛻舳烁鶕?jù)一定的通信機制向中央服務器傳輸模型參數(shù)等數(shù)據(jù),中央服務器匯聚各客戶端上載數(shù)據(jù)后訓練構建全局模型”。由于 LLM 的參數(shù)量和模型結構非常龐大,這種傳輸、匯聚的過程又受到通信條件的影響,所以整個工作過程的效率是非常低的。
3、FL+LLM 已有哪些工作進展
盡管 FL+LLM 還面臨諸多問題,但是不少研究人員已經(jīng)開始探索相關方向的技術可能性并取得了一定的進展。我們在這一章節(jié)中介紹一下相關的工作進展情況。
3.1 FedLLM:在專有數(shù)據(jù)上建立你自己的大型語言模型
FedLLM 是一個 MLOps 支持的訓練 pipeline,能夠實現(xiàn)在專有數(shù)據(jù)上建立企業(yè)自己的大語言模型 [7]。代碼已公開:https://github.com/FedML-AI/FedML/tree/master/python/app/fedllm 。FedLLM 代表了 "LLM 的基礎生態(tài)系統(tǒng)設計",而不僅僅是 "LLM 的聯(lián)邦學習"。
FedLLM 是 FedML.Inc. 推出的 FL+LLM 框架。FedML, Inc. (https://FedML.ai ) 是由華人主導的國際化團隊,他們起源于美國南加州大學 (University of Southern California),是全球范圍內研究該技術的早期機構之一。過去幾年,F(xiàn)edML 起初作為博士生主導的科研開源項目,服務于多個科研基金,輔助所在實驗室發(fā)表了 50 多篇相關頂級論文。2023 年,F(xiàn)edML 宣布完成 600 萬美元種子輪和 Pre-Seed 輪融資,Camford Capital 領投,Plug and Play Ventures、AimTop Ventures、Acequia Capital、LDV Partners 等投資者參投。
圖 2. FedLLM 的訓練過程
如圖 2,F(xiàn)edLLM 實現(xiàn)了數(shù)據(jù)協(xié)作、計算協(xié)作和模型協(xié)作,并支持在集中式和地理分布式的 GPU 集群上進行訓練,以及數(shù)據(jù)孤島的聯(lián)邦學習。FedLLM 與 HuggingFace 和 DeepSpeed 等流行的 LLM 庫兼容,并旨在提高效率和安全 / 隱私。
- 數(shù)據(jù)協(xié)作 — 在特定領域的專有數(shù)據(jù)上實現(xiàn) LLM 的訓練。通用的 LLMs ChatGPT 和 GPT-4 是在人類發(fā)布和注釋的大量文本上訓練的。在許多垂直領域(如醫(yī)療保健、金融科技、法律和汽車行業(yè)),ChatGPT 可能無法很好地發(fā)揮作用。使企業(yè)能夠在其專有數(shù)據(jù)上訓練他們的模型,可以在保護隱私的同時實現(xiàn)更好的性能。在某些情況下,當數(shù)據(jù)集分散在各個數(shù)倉中時,F(xiàn)edML 聯(lián)邦學習訓練管道 "Train on the Edge" 將以安全和可擴展的方式處理。
- 計算協(xié)作 — 利用分散的計算資源。由于計算能力的高投資,只有少數(shù)大型科技公司(如微軟、谷歌)才能負擔得起 LLM 預訓練。處于不同垂直領域的企業(yè)無法負擔花費數(shù)十億美元的數(shù)千個 GPU 節(jié)點。充分利用 GPU 資源的一個更具成本效益的方法是建立一個跨組織的共享機制。FedML 通過 "Train on the cloud" 平臺實現(xiàn)了這一點,它可以將訓練工作安排到地理分布的 CPU/GPU 網(wǎng)絡上。這種計算方面的合作可以減少在一個組織中購買大量 GPU 節(jié)點的財務負擔。
- 模型協(xié)作 — 以聯(lián)合的方式服務于模型。服務大型基礎模型也是一種挑戰(zhàn)。與其他 MLOps 相比,F(xiàn)edML 率先提出了通過地理分布的云計算資源進行聯(lián)邦模型推理的想法。當推理請求被發(fā)送到推理端點時,主節(jié)點將把請求路由到由 GPU 供應商托管的分散的邊緣節(jié)點,他們可以共享 GPU 的空閑時間。這樣的模式服務平臺可以提供更好的服務可靠性和更便宜的云計算成本。
對于特定的孤島式 GPU 集群,F(xiàn)edLLM 利用現(xiàn)有的開源 LLM 和流行框架進行本地訓練:
- 模型定義和預訓練的權重來自 EleutherAI 的 pythia 的 2.8B、7B 和 12B 版本。將 Hugging Face 的 transformer 作為參考實現(xiàn)(https://github.com/huggingface/peft?ref=blog.fedml.ai)。FedLLM 也是一個模型無關的框架,企業(yè)和開發(fā)者可以插入任何 LLM。
- FedLLM 支持參數(shù)高效的訓練方法,如 LoRA。其參考實現(xiàn)來自 HuggingFace 的 peft(https://huggingface.co/docs/transformers/index?ref=blog.fedml.ai)。
- 對于一個獨立的訓練器,其訓練 / 評估代碼是基于 Transformers 庫(https://huggingface.co/docs/transformers/index?ref=blog.fedml.ai)中的 Trainer。
- 特定 GPU 集群上的分布式訓練框架由 DeepSpeed(https://www.deepspeed.ai/?ref=blog.fedml.ai)處理。還啟用 Zero3 來減少單個 GPU 的內存成本。通過加強 FedML 和 DeepSpeed 之間的兼容性,F(xiàn)edLLM 可以在不同的物理空間分布式集群中運行訓練作業(yè)。
FedML 也給出了一個應用 FedLLM 的實例。如圖 3,本地數(shù)據(jù)保持在本地,只有模型參數(shù)或權重數(shù)據(jù)從中央服務器流向和流回。這個特殊的實例假設了一個聊天類型的應用,它結合了本地數(shù)據(jù)的訓練,以及中央服務器通過利用其他設備的訓練而建立的微調訓練的好處。聊天應用程序只是一個例子,可以被其他利用 LLM 模型的應用程序所取代。
圖 3. FedLLM 應用實例
FedLLM 從一定程度上能夠實現(xiàn)利用聯(lián)邦學習訓練 LLM。但是,回顧我們在第二章中討論的,利用聯(lián)邦學習架構訓練 LLM 還面臨很多問題,從公開的技術內容來看,F(xiàn)edLLM 并沒有能夠解決這些問題。FedLLM 架構中假設的計算節(jié)點都是具備一定計算能力的 GPU 節(jié)點,F(xiàn)edLLM 能夠解決的問題是 “不用大量購置 GPU”。但真實狀況下的聯(lián)邦學習,客戶端節(jié)點很多是手機、平板,F(xiàn)edLLM 并不能適配這些情況。同樣的問題也存在于客戶端節(jié)點數(shù)據(jù)量的情況,即 FedLLM 假設客戶端是足夠容納訓練本地 LLM 所需要的數(shù)據(jù)量的。最后,F(xiàn)edLLM 并沒有討論什么樣的聚合算法適合 FL for LLM,也沒有討論是否需要改進客戶端訓練 LLM 的算法,這對于真正推廣應用 FL+LLM 還是需要進一步細化解決的。
3.2 DeepSpeed:通過提高規(guī)模、速度、成本和可用性來推進大型模型訓練,釋放出訓練 1000 億個參數(shù)模型的能力
LLM 在訓練時往往需要大量內存來存儲中間激活、權重等參數(shù),百億模型甚至無法在單個 GPU 上進行訓練,使得模型訓練在某些情況下非常低效和不可能。前期工作已經(jīng)有研究人員專門聚焦在多節(jié)點分布式訓練問題中,主要采用數(shù)據(jù)并行(模型的不同實例在不同的 GPU 和不同批數(shù)據(jù)上運行)和模型并行(將模型拆分到多個 GPU 上進行訓練)兩種方式。
DeepSpeed 是微軟開發(fā)的一種深度學習訓練優(yōu)化工具,旨在通過分布式訓練和混合精度技術來提高訓練速度和節(jié)省資源。它是一個開源的 Python 庫,可以在多個平臺上運行。與傳統(tǒng)的深度學習框架,例如 TensorFlow、Pytorch、Keras 等相比,DeepSpeed 通過將模型參數(shù)拆散分布到各個 GPU 上,以實現(xiàn)大模型的計算,使得可以用更少的 GPU 訓練更大的模型,而且不受限于顯存。
數(shù)據(jù)并行會在所有 worker 之間進行模型和優(yōu)化器的復制,因此顯存效率不高。DeepSpeed 開發(fā)了 ZeRO ,它是一系列用于提高數(shù)據(jù)并行的顯存效率的優(yōu)化器。針對模型并行,DeepSpeed 借用了英偉達的 Megatron-LM 來為基于 Transformer 的語言模型提供大規(guī)模模型并行功能。
具體來說,DeepSpeed (https://www.deepspeed.ai) 能夠通過一次點擊實現(xiàn)類似 ChatGPT 的模型訓練,與 SOTA RLHF 系統(tǒng)相比,提供了 15 倍的速度提升,并在所有規(guī)模上實現(xiàn)了前所未有的成本降低。DeepSpeed 有三個創(chuàng)新點:
- DeepSpeed-Training:DeepSpeed 提供了一系列的系統(tǒng)創(chuàng)新,這些創(chuàng)新使得大規(guī)模的深度學習訓練變得有效和高效,極大地改善了使用的便利性,并且在可能的規(guī)模方面重新定義了深度學習訓練領域。這些創(chuàng)新,例如 ZeRO、3D-Parallelism、DeepSpeed-MoE、ZeRO-Infinity 等,都屬于 DeepSpeed-Training 的核心。
- DeepSpeed inference:DeepSpeed inference 是一個用于 Transformer 模型推理的全面系統(tǒng)解決方案。DeepSpeed inference 包括:(1) 一個多 GPU 推理解決方案,以最大限度地減少延遲,同時最大限度地提高密集和稀疏 Transformer 模型的吞吐量,當它們適合聚合 GPU 內存時;(2) 一個異構推理解決方案,除了 GPU 內存和計算外,還利用 CPU 和 NVMe 內存,以實現(xiàn)不適合聚合 GPU 內存的大模型的高推理吞吐量。
- DeepSpeed-Compression:為了進一步提高推理效率,DeepSpeed 為研究人員和從業(yè)人員提供了易于使用和靈活組成的壓縮技術,以壓縮他們的模型,同時提供更快的速度、更小的模型尺寸,并大大降低壓縮成本。此外,SoTA 在壓縮方面的創(chuàng)新,如 ZeroQuant 和 XTC,都包含在 DeepSpeed-Compression。
具體來說,DeepSpeed Inference 的兩個核心組件介紹如下:
- DeepSpeed Transformer:DeepSpeed Transformer 是一個僅適用于 GPU 的解決方案,旨在最小化延遲,同時最大限度地提高密集和稀疏 Transformer 模型的吞吐量。它為各種規(guī)模的 Transformer 模型實現(xiàn)了最先進的延遲和吞吐量,并支持在單個 GPU 上運行或擴展到數(shù)百個 GPU 以推理數(shù)百萬參數(shù)的模型。DeepSpeed Transformer 解決方案是一個三層的系統(tǒng)架構,包括:i) 單個 GPU 轉化器內核,在小批規(guī)模時針對內存帶寬利用率進行優(yōu)化,在大批規(guī)模時針對高吞吐量進行優(yōu)化;ii) 多個 GPU 密集 Transformer 層,利用張量切割和推理優(yōu)化的管道并行性在 GPU 上擴展密集 Transformer 模型。海量 GPU 規(guī)模的稀疏 Transformer 層,旨在利用并行技術和通信優(yōu)化策略的組合將 MoE Transformer 層擴展到數(shù)百個 GPU,同時還利用優(yōu)化的稀疏內核將單 GPU 稀疏計算的開銷降到最低。通過采取這種分層方法,每一層都解決了延遲挑戰(zhàn)的一個獨特方面:批大小、密集模型的擴展和稀疏模型的擴展,但又相互兼容并建立在彼此之上,DeepSpeed Transformer 是一個綜合系統(tǒng),盡管在批大小、模型規(guī)模和模型特征方面存在異質性,但它能夠以前所未有的規(guī)模為密集和稀疏 Transformer 模型實現(xiàn)最優(yōu)的延遲和吞吐量。
- ZeRO-Inference:ZeRO-Inference 是一個基于異構 GPU+CPU+NVMe 的解決方案,是一種用于大規(guī)模分布式深度學習的新型內存優(yōu)化技術。通過用最少的 GPU 資源實現(xiàn)大規(guī)模的模型推理,來解決內存的挑戰(zhàn)。與 DeepSpeed Transformer 相比,對于那些對延遲不太敏感但資源有限的應用,只要有足夠的 CPU 或 NVMe 內存來存儲模型參數(shù),ZeRO-Inference 就可以在單個或多個 GPU 上推斷具有數(shù)千億個參數(shù)的模型。
圖 4. DeepSpeed 壓縮庫
DeepSpeed Compression 提出了一個無縫 pipeline 來解決壓縮合成性的挑戰(zhàn),如圖 4 所示。DeepSpeed Compression 的核心部分是一個叫做 compression composer 的組件,它包括幾個重要的功能:
1. 它提供了多種尖端的壓縮方法,包括極端量化、頭 / 行 / 通道修剪和知識提煉,可以有效地減少模型大小和推理成本。隨著不斷地整合更多的最先進的壓縮方法,這個列表還會擴大。
2. 它提供了一個易于使用的 API,自動處理組裝不同壓縮技術的復雜問題,以提供多種壓縮方法的復合效益。例如,XTC 需要組成輕量級的減層、二值化和知識蒸餾。然而,將它們組合在一起并不困難。有了 compression composer,應用極端壓縮就像添加兩個新的 API 調用來啟用壓縮和清理壓縮模型一樣簡單。
3. 它是以模塊化的方式設計的,這樣用戶就可以很容易地添加新的壓縮方案。例如,額外的壓縮方法可以通過自定義壓縮層添加,通過向 compression composer 注冊,新方法可以與已經(jīng)由 composer 管理的現(xiàn)有方法組成。
4. 它與 DeepSpeed 庫無縫銜接。這有兩個好處。首先,DeepSpeed 壓縮可以通過 JSON 文件以與 DeepSpeed 訓練和推理相同的方式進行指定和啟用,其中啟用不同的壓縮技術組合只需要在 JSON 文件中修改幾行。其次,一旦配置了壓縮方案,壓縮 composer 就會自動修改模型層和訓練以啟用壓縮過程,不需要用戶對模型結構或訓練程序進行額外的修改。
3.3 相關算法基礎
為了應對 LLM 大規(guī)模預訓練的成本問題,從算法層面,研究人員提出了 LoRA( LOW-RANK ADAPTATION ),即大語言模型的低秩適應方法,一種輕量級訓練大語言模型的方法。通過凍結預訓練的模型權重,并將可訓練的秩分解矩陣注入到 Transformer 架構的每一層,極大地減少了下游任務的可訓練參數(shù)的數(shù)量,有效提升預訓練模型在下游任務上的 finetune 效率。
圖片
項目地址:https://github.com/microsoft/LoRA [6]
為解決微調大語言模型到不同領域和任務的挑戰(zhàn),已有多種方案,比如部分微調、使用 adapters 和 prompting。但這些方法存在如下問題:
- Adapters 引入額外的推理延遲 (由于增加了模型層數(shù))
- Prefix-Tuning 難于訓練,且預留給 prompt 的序列擠占了下游任務的輸入序列空間,影響模型性能
Lora 提出的思路是:雖然模型的參數(shù)眾多,但其實模型主要依賴低秩維度的內容 (low intrinsic dimension)。假設在適應過程中,權重的更新也具有較低的 "內在等級(intrinsic rank)"。對于一個預先訓練好的權重矩陣 W_0∈R^d×k,通過用低秩分解代表后者來約束其更新:
圖片
圖片
圖片
在訓練期間,凍結 W_0,不接受梯度更新,而 A 和 B 包含可訓練參數(shù)。W_0 和?W = BA 都與相同的輸入相乘,它們各自的輸出向量按坐標相加。對于 h=W_0x,修改后的正向傳遞得到了:
圖片
圖片
圖 5. Reparameterized 只訓練 A 和 B
如圖 5,對 A 使用隨機高斯初始化,對 B 使用零初始化,所以?W = BA 在訓練開始時是零。然后,用 α/r 來調整?Wx,其中 α 是 r 中的一個常數(shù)。當用 Adam 進行優(yōu)化時,如果適當?shù)卣{整初始化,那么調整 α 與調整學習率大致相同。因此,只需將 α 設置為嘗試的第一個 r,而不對其進行調整。這種比例有助于減少在改變 r 時重新調整超參數(shù)的需要。
圖片
圖 6. 常規(guī)微調和 LoRA 圖示
如圖 6,r 是超參數(shù),指定用于自適應的低秩矩陣的秩。r 越小,低秩矩陣越簡單,在自適應過程中需要學習的參數(shù)越少,訓練就更快,計算需求會相應的減少。然而,r 變小的弊端是,低秩矩陣捕獲任務特定信息的能力降低。這可能導致較低的自適應質量,并且與較高的 r 相比,模型在新任務上可能表現(xiàn)不佳??傊?LoRA 中確定 r 的取值,需要在模型復雜性、自適應能力和擬合不足或擬合過度的風險之間進行權衡。因此,重要的是用不同的 r 值進行實驗,以找到正確的平衡,從而在新任務中滿足所需的性能。
矩陣秩 (r) 可以非常低,比如對于 GPT-3 175B 模型,rank-1 或 rank-2 就能基本對齊原始 rank-12288 的效果:
- 與 Adam 微調的 GPT-3 175B 相比,LoRA 可訓練參數(shù)數(shù)量減少了 1 萬倍,GPU 內存需求減少了 3 倍
- 在 RoBERTa、DeBERTa、GPT-2 和 GPT-3 等大語言模型上,LoRA 在模型質量方面的表現(xiàn)與微調相當或更好,盡管它具有更少的可訓練參數(shù)、更高的訓練吞吐量,并且與適配器不同,沒有額外的推斷延遲
3.4 相關硬件基礎
最后,我們來看看硬件方面的工作。
Grace Hopper 是第九代英偉達數(shù)據(jù)中心 GPU,與前幾代英偉達安培 GPU 相比,旨在為大規(guī)模 AI 和 HPC 應用提供數(shù)量級的改進。線程塊集群和線程塊重新配置改善了空間和時間上的數(shù)據(jù)定位,再加上新的異步執(zhí)行引擎,使應用程序能夠始終保持所有單元的忙碌。NVIDIA Grace Hopper 通過 NVIDIA NVLink C2C 將 NVIDIA Grace CPU 和 NVIDIA Hopper GPU 融合到一個超級芯片中,NVLink C2C 是一個總帶寬為 900 GB/s 的芯片間互連。NVLink C2C 內存一致性使得 Grace CPU 超級芯片和 Grace Hopper 超級芯片都可以通過統(tǒng)一的編程模型進行編程。
圖片
圖 7. 英偉達 Grace Hopper 超級芯片邏輯概覽
2023 年 6 月,英偉達宣布其新的 GH200 Grace Hopper "超級芯片"— 專門為大規(guī)模人工智能應用創(chuàng)建的 CPU 和 GPU 組合 — 已經(jīng)進入全面生產(chǎn)。它有 528 個 GPU 張量核心,支持高達 480GB 的 CPU 內存和 96GB 的 GPU 內存,并擁有高達 4TB / 秒的 GPU 內存帶寬。
圖片
圖 8. 英偉達的 GH200 "Grace Hopper"AI 超級芯片 [5]
GH200 以 "Hopper" 為基礎,與 Nvidia 的 "Grace"CPU 平臺(兩者均以計算機先驅 Grace Hopper 命名)相結合,通過 Nvidia 的 NVLink 芯片對芯片(C2C)互連技術將其納入一個芯片。Nvidia 預計,這一組合將極大地加速人工智能和機器學習應用的訓練(創(chuàng)建模型)和推理(運行模型)。
"生成式人工智能正在迅速改變企業(yè),在醫(yī)療、金融、商業(yè)服務和更多行業(yè)釋放新的機會并加速發(fā)現(xiàn),"Nvidia 加速計算副總裁 Ian Buck 在一份新聞稿中說。"隨著 Grace Hopper 超級芯片的全面生產(chǎn),全球制造商將很快提供企業(yè)所需的加速基礎設施,以建立和部署利用其獨特的專有數(shù)據(jù)的生成式人工智能應用程序。" 據(jù)該公司稱,GH200 的主要特點包括一個新的 900GB/s 相干(共享)內存接口,比 PCIe Gen5 快 7 倍。GH200 還為 GPU 提供了 30 倍的系統(tǒng)內存總帶寬。此外,GH200 可以運行所有 Nvidia 軟件平臺,包括 Nvidia HPC SDK、Nvidia AI 和 Nvidia Omniverse。Nvidia 還宣布,它將把這種 CPU/GPU 組合芯片構建成一個新的超級計算機,稱為 DGX GH200,它可以利用 256 個 GH200 芯片的綜合能力,作為一個單一的 GPU 來執(zhí)行,提供 1 exaflop 的性能和 144 兆字節(jié)的共享內存,比上一代 Nvidia DGX A100 的內存多近 500 倍。DGX GH200 將能夠為生成語言應用、推薦系統(tǒng)和數(shù)據(jù)分析訓練巨大的下一代人工智能模型。Nvidia 目前并未公開 DGX GH200 的售價,但如果根據(jù)去年出貨的 DGX H100 作為參考,一臺搭載 8 組 H100 GPU 的 8U GPU 服務器機柜約為 20 萬美元,考慮到 DGX GH200 擁有最多 256 個 Grace Hopper,其價格可能會高于這個范圍。根據(jù) Anandtech 的報道,一臺 DGX GH200 計算機的價格 "很容易達到 8 位數(shù)(美元)的水平"。
由于 Nvidia 和 Cerebras 等供應商的持續(xù)硬件進步,高端云人工智能模型可能會隨著時間的推移繼續(xù)變得更有能力,處理更多的數(shù)據(jù),并且比以前做得更快。根據(jù)內部測試,在處理需要大量內存的 AI 工作負載方面,DGX GH200 系統(tǒng)顯示出比 DGX H100 更好表現(xiàn),平均可獲得 2-6 倍不等的性能提升,例如,在使用 1TB 的內存容量用于 GPT3 模型訓練時,DGX GH200 的處理速度快上 1 倍,若以 40TB 內存處理深度學習推薦模型(DLRM)時,速度則有多達 4 倍的提升,甚至在圖神經(jīng)網(wǎng)絡處理方面速度更快上許多,足足有 5 倍之多。
單純從 GH200 的角度看,能夠在一定程度上支撐對 LLM 的分布式訓練,使 FL+LLM 變得更加可行。但是,回到文初我們的討論,F(xiàn)L 的初衷是利用大量分散的一般性設備分布式訓練一個中央模型,一方面可以有效利用分散的客戶端資源,另一方面也滿足各個客戶端的數(shù)據(jù)隱私需要。要求這些客戶端都裝有 GH200 顯然是不現(xiàn)實的,這樣的 GPU 的成本消耗與 FL 的初衷并不一致。
此外,也有一些研究人員認為 GH200 自身也不會對大模型的應用推廣有太大影響 [11]。他們分析,算力層面,單顆 GH 芯片和 H100 芯片 FP8 浮點算力沒有差異。內部測試中 GH200 更好的表現(xiàn)是由于其存儲。DGX GH200 內部 GPU 和 CPU 的連接方式與 DGX H100 不同,其高速觸達的存儲容量可以大幅提升。然而,集群性能的提升主要有算力本身、網(wǎng)絡、存儲三大要素。因此,傳統(tǒng) GPT3、GPT4 等主流大模型下,DGX H100 集群(NVLink 聯(lián)網(wǎng))與 DGX GH200 系統(tǒng)并不會有明顯差異。
關于 H100,已有報道分析了其大模型訓練的性能 [12]。兩項 MLPerf 基準評測的最新數(shù)據(jù)中,英偉達 H100 芯片組在人工智能算力表現(xiàn)的測試中,刷新了所有組別的紀錄,也是唯一一個能夠跑完所有測試的硬件平臺。英偉達和 GPU 云算力平臺 CoreWeave 提交的數(shù)據(jù)為這項測試設定了業(yè)界標準。在 896 個英特爾至強 8462Y + 處理器和 3584 個英偉達 H100 芯片的聯(lián)合作用下,僅僅花了 10.94 分鐘就完成了基于 GPT-3 的大語言模型訓練任務。在 BERT-Large 模型訓練中,H100 和 CoreWeave 將數(shù)據(jù)刷到了極端的 0.13 分鐘,在 64 卡的情況下,測試數(shù)據(jù)也達到了 0.89 分鐘。
4、后續(xù)發(fā)展討論
我們在這篇文章中討論了與 FL+LLM 相關的一些方法,包括算法方面的改進、硬件方面的研發(fā),以及分布式訓練和聯(lián)邦學習的架構。其中,F(xiàn)edLLM 應當是最符合 FL+LLM 的一篇工作,盡管其距離實用、技術完備的 FL+LLM 還有很大差距。
DeepSpeed 是通過分布式訓練和混合精度技術來提高訓練速度和節(jié)省資源的模型訓練架構,其面向的是分布式訓練的應用場景,一般來說,分布式學習會將訓練數(shù)據(jù)均勻分布在不同節(jié)點,節(jié)點大多是專用機房中的計算節(jié)點,且節(jié)點通常處于統(tǒng)一地理位置,節(jié)點間通信條件較好。而對于 FL 來說,每個計算節(jié)點擁有的數(shù)據(jù)量與設備自身有關,很難保證不同計算節(jié)點擁有相近的數(shù)據(jù)量。此外,F(xiàn)L 的計算節(jié)點可能是手機、平板等,一般與中央服務器處于遠程連接狀態(tài),連接穩(wěn)定性和通信代價的情況都比較差。所以目前 DeepSpeed 的工作直接應用到 FL+LLM 中還有很大難度。
LoRA 是目前非常流行的微調方法,更多優(yōu)秀的微調方法也在不斷提出。算法上的改進,給 FL 各個客戶端節(jié)點基于本地數(shù)據(jù)訓練模型帶來了可能。而硬件目前的發(fā)展我們在前面也討論過,硬件自身的性能在不斷提高,但對于 FL 這個應用場景來說,令手機、平板具備這樣的硬件條件難度非常大。
由對目前與 FL+LLM 相關的一些工作的分析,我們感覺聯(lián)邦學習支撐 LLM 還有很多需要解決的問題。包括 FedLLM,實際上也沒有談及如何應對存儲容量小、處理性能差、網(wǎng)絡條件不佳的客戶端設備,而這些,恰好是 FL 最典型的應用場景。我們期待有更多的研究人員可以關注 FL+LLM 問題,給出更多的技術細節(jié)和可落地的解決方案。
作者介紹
Jiying,工學博士,畢業(yè)于北京交通大學,曾分別于香港中文大學和香港科技大學擔任助理研究員和研究助理,現(xiàn)從事電子政務領域信息化新技術研究工作。主要研究方向為模式識別、計算機視覺,愛好科研,希望能保持學習、不斷進步。