隨着大型語言模型的不斷蓬勃發展，相關新模型，新應用和新範式也在不斷湧現，自4月發佈以來，FATE-LLM已經迭代發佈了多個版本，不斷完善大語言模型在聯邦學習場景下的支持，以解決構建、使用大模型時的數據隱私保護問題以及公域數據短缺，可用數據不足的問題，在社區中得到了廣泛的關注。
由VMware AI Labs團隊發起並貢獻的KubeFATE項目也在最近的多個版本中增強了對FATE-LLM在雲原生環境下的支持，特別是針對FATE-LLM任務的專有需求，KubeFATE在包括容器構建、GPU調度與使用、模型存儲等方面加入了專門的設計。本篇文章將給出一個基於KubeFATE v1.11.2和FATE-LLM v1.2的聯邦大模型微調任務實例，並從任務設定、環境部署、所用技術、實驗結果與分析等角度進行一個全方位的完整介紹。
任務設定與環境部署
在本文中，我們以FATE-LLM v1.2的GPT2教程爲基礎，設定一個兩方的橫向聯邦學習場景，兩方Party ID分別爲9999和10000，我們使用一個文本情感分類的任務來微調一個預訓練的GPT2模型。該任務使用的數據集爲IMDB影評數據集，與原教程的少量數據的示例不同，我們將訓練數據集中全部25000條記錄平均分給兩方作爲各方的訓練數據，同時使用測試數據集的25000條記錄作爲評估數據。
對於實驗環境，我們創建了兩個Kubernetes集羣，分別代表聯邦任務的兩個參與方。每個Kubernetes集羣都包含一個vSphere虛擬機作爲GPU節點，各虛擬機使用PCI Passthrough技術分別與一塊Nvidia V100 GPU集成。我們使用Docker和cri-docker，以及nvidia-container-runtime和k8s-device-plugin等組件以在Kubernetes集羣中使用該GPU。本實驗所有關鍵依賴項的具體版本如下：
我們可以按照KubeFATE項目GitHub倉庫中的K8s環境使用指南來部署KubeFATE和FATE集羣。要使用 FATE-LLM，我們需要在用於部署FATE集羣的cluster.yaml文件中顯式指定某些設置。首先我們應該將algorithm參數設置爲“ALL”，並將device參數設置爲“GPU”，這表示我們要使用包含FATE-LLM和相關模塊，以及帶有GPU驅動和庫的FATE容器鏡像。此外，我們需要在python組件即FATE-Flow容器的配置中，在resources資源部分加入GPU資源的請求，例如本文使用了“nvidia.com/gpu: 1”。以下是這些設置的示例：
需要注意的是，本文示例中FATE-LLM的訓練任務是由FATE-Flow容器執行的，因此GPU資源分配給了名爲python的容器。對於使用了DeepSpeed的FATE-LLM任務，我們需要爲nodemanager組件配置該GPU資源設定。此外，我們建議爲FATE集羣開啓持久化，從而使KubeFATE能夠保存FATE-LLM任務中的預訓練模型、微調模型、任務記錄、日誌等，防止這些關鍵文件在容器發生重啓後被清理。
爲了驗證FATE集羣部署後的環境和設置，我們可以使用kubectl exec進入fateflow pod並使用nvidia-smi命令檢查其中可用的GPU資源。在FATE集羣部署好後，我們就可以開始發起FATE-LLM任務了，FATE-LLM v1.2中使用HuggingFace生態的peft庫來支持多種高效的參數微調方法，在本文的實踐中，我們將採用LoRA、Prefix Tuning、Prompt Tuning和P-Tuning的方法進行比較實驗。
採用的PEFT方法
爲了使用大語言模型來完成特定的下游任務，如果需要對模型的全部參數進行微調，會造成了巨大的時間開銷與數據存儲、傳輸成本。PEFT方法在保持模型性能相當的前提下，通過減少微調的參數量以降低計算、存儲、傳輸數據的成本。在本節中，我們將簡要介紹採用的四種PEFT方法。
-
LoRA（Low-Rank Adaptation）：LoRA的核心思想是在適應新任務時，權重矩陣的更新不必與原始權重矩陣具有相同的“秩”。因此對每個權重矩陣，我們可以引入兩個低秩矩陣模塊去代表該更新，微調時則僅訓練更新這些低秩的矩陣。在實踐中，LoRA一般被應用到attention模塊中。
-
Prefix Tuning：該方法在每個transformer層的輸入之前構造一段任務相關的virtual tokens作爲prefix，在訓練時只更新prefix部分的參數，而固定transformer中的其他部分參數。此外，爲了防止直接更新prefix的參數導致訓練不穩定的情況，prefix層前可以加入了MLP（Multi-layer Perception）結構，即將prefix分解爲更小維度的input與MLP的組合。
-
Prompt Tuning：該方法可以看作是Prefix Tuning的簡化版本，它只在輸入層加入prompt tokens，而並不需要加入MLP進行調整。
-
P-Tuning：該方法的思想也是在輸入層添加可訓練的prompt參數，由於prompt的詞之間是彼此關聯的，離散的prompt對於連續的神經網絡並不是最優解，我們需要採用某種方法將它們關聯起來。於是P-Tuning將一些僞prompt輸入至LSTM或類似模型中，利用LSTM的輸出向量來替代原始的prompt token，最後一起輸入至預訓練語言模型中。
運行FATE-LLM任務
接下來我們便可以着手運行FATE-LLM任務。首先，guest與host兩方需將本地數據上傳至FATE系統，即將預處理後的CSV文件分別放入各方的FATE-Flow容器中。然後我們可以在KubeFATE提供的Jupyter Notebook中使用以下代碼將上傳的數據綁定到FATE中。（對於host方，需要綁定data_host數據）
隨後我們便可以在guest方提交聯邦學習任務，具體流程與FATE-LLM倉庫中的GPT2微調教程基本一致，以下是需額外注意的幾點：
-
因爲本文加入了用於評估的測試數據，我們可以新建一個Reader組件讀取該數據並作爲NN組件的validate_date；
-
若任務出現超時或Pin memory thread退出等異常，可以嘗試將“save_to_local_dir=True, pin_memory=False”添加到TrainerParam；
-
對於本文使用的二分類任務，我們爲其配置了Evaluation組件用於評估訓練後的模型的性能。
我們使用了多種不同的參數設置方式運行了若干個FATE-LLM任務來評估它們對訓練過程與模型性能的影響：
-
設置TrainerParam中的CUDA參數來比較訓練時使用CPU和GPU的差異
-
設定peft_type與peft_config參數來採用不同的PEFT方法以及配置
-
設置aggregate_every_n_epoch參數比較不同的本地訓練輪次對訓練效果的影響
-
調整TrainerParam中的epoch參數找到模型效果最好的輪次（最大爲10）
-
我們實現了非PEFT的GPT2 Full Fine-tuning方法，並以此作爲對照進行了對比實驗
實驗結果分析
我們對使用的設備類型、PEFT方法以及運行的epoch次數均進行了多元化的設置。模型性能通過AUC、F1-score等機器學習指標進行評估。除此以外，我們還列舉了若干實驗本身的數據，包括聚合過程的數據傳輸量，任務運行時間等，便於進行更深入的分析。在使用PEFT方法對模型訓練過程進行優化時，我們儘量使用peft庫默認的參數配置。以下是我們實驗的部分結果數據：
通過對比表格中與模型性能相關的指標，我們可以看到LoRA、Prefix Tuning方法與Full Fine-tuning對照實驗的模型性能相當。與此同時，它們的數據傳輸量、整體訓練時間也顯著低於Full Fine-tuning，更進一步，我們觀察到設定更大一點的模型聚合間的本地訓練輪次（同時總輪次不變），能夠進一步降低數據傳輸成本，同時不影響模型性能，這證實了聯邦學習與PEFT方法之於微調大語言模型是可行且高效的。而對於另兩種PEFT方法，模型性能則有較大的差距，我們猜測這與可訓練參數過少或相關PEFT組件參數設置有關，導致模型泛化能力受限。我們在接下來的實驗中也將嘗試其他方法和參數組合以進一步驗證。
總結與展望
本文介紹了使用KubeFATE從部署、配置到發起運行FATE-LLM任務的完整流程。我們展示了不同配置下的實驗結果，並進行了分析。實驗數據證實了將大語言模型與聯邦學習架構結合的合理性，並體現出PEFT方法的強大性能。其在顯著降低通信成本、數據存儲成本的同時，將模型性能維持在一個相當可觀的水平。
本文提供的實例基於FATE-LLM v1.2版本，可作爲橫向聯邦場景下，使用雲原生基礎設施對同構的大語言模型進行PEFT的微調訓練的一個參考。除GPT2示例以外，KubeFATE也支持FATE-LLM v1.2中的其他模型和訓練方法，例如使用DeepSpeed進行ChatGLM-6B、Llama等模型的訓練。而在最新發布的v1.3版本中，FATE-LLM項目引入了FTL-LLM面向大語言模型的聯邦遷移學習範式，並實現了聯邦Offsite-Tuning框架。我們會在後續的KubeFATE迭代中進一步加強相關的支持，也歡迎社區開發者和用戶的參與和關注。
內容來源｜公衆號：VMware 中國研發中心
有任何疑問，歡迎掃描下方公衆號聯繫我們哦～
