當前位置: 妍妍網 > 資訊

在家就能訓練 700 億大模型!Answer.AI 開源計畫沖上 HN 熱榜

2024-03-12資訊

整理 | 王啟隆

出品 | AI 科技大本營(ID:rgznai100)

自從 Meta 入局開源了 LLaMA,大模型圈的創業浪潮似乎就沒停下來過,數以萬計的初創公司嶄露頭角,頗有當年個人電腦(PC)行業在 Linux 系統開源化之後的繁榮景象。

就在去年 12 月,曾經創立了 fast.ai 的 Jeremy Howard 和 Eric Ries 建立了一個名為 Answer.AI 的新人工智慧研發實驗室。該實驗室得到了 Decibel VC 的 1000 萬美元資金支持,旨在透過專註於利用現有技術建立實用的人工智慧應用程式。

Jeremy Howard

準確地說, Answer.AI 對自己的定位是「 獨特的營利性質研發實 驗室 」,其精神更接近 19 世紀的「 電力實驗室 」,和當前如火如荼的 AI 初創浪潮區分開來。

3 月 6 日, Answer.AI 終於在自己的官網部落格上正式公開了自己的第一個計畫:一個完全 開源 的系統,首次能夠高效地在配備 兩個或多個標準遊戲顯卡 (如老黃家的 RTX 3090 或 4090)的普通台式電腦上訓練一個 70b( 700 億參數 )的大規模語言模型 。目前,Hacker News 對該計畫已經討論得熱火朝天。

在 GitHub Trending 上,這一計畫今天也沖上了 Jupyter Notebook 的熱度第一。

「端側大模型」是今年相當熱門的話題,之前開源了 7b 模型的法國公司 Mistral 就用一個小模型擊敗了 LLaMA,還得到了微軟的青睞; 公眾 號曾分享過的 更是從底層掀翻了大模型牌局。 Answer.AI 的這個計畫基於 FSDP 和 QLoRA 技術,是 Answer.AI、Tim Dettmers、以及 Hugging Face 團隊的 Titus von Koeller 和 Sourab Mangrulkar 共同合作的成果。

接下來,讓我們一起來看看 A nswer.AI 究竟用了哪些手段實作「在家訓練大模型」,一同展望「人人都能建立個人化模型」的時刻還有多久到來。

原文連結:https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html#a-first-step

用遊戲顯卡訓練大模型!

2024 年的今天,訓練深度學習模型使用的硬體通常分為兩大陣營。

一派是數據中心級別的高端硬體,比如 H100s 和 A100s,售價高達數十萬美元;另一派則是配備有遊戲顯卡的普通台式電腦,例如一對 RTX 4090 顯卡,總價不足 1 萬美元(甚至可以采用二手部件拼裝,俗稱「撿垃圾」,成本不到預裝整機的一半)。

關鍵在於, 盡管遊戲顯卡的價格僅為數據中心級顯卡的十分之一,但它們的效能卻相當接近 。設想若能有效利用這類成本低廉的遊戲顯卡訓練大語言模型,無疑將是一個巨大的進步。

然而,現實情況是, 頂級的數據中心級顯卡擁有 80 GB 的 RAM,而遊戲顯卡的最大 RAM 則為 24 GB 。鑒於只有大型乃至超大型模型才能取得最優結果,迄今為止,構建此類頂尖模型對於大多數個人和小型團隊來說都是遙不可及的。

Answer.AI 的團隊認識到,理論上存在使用低成本遊戲顯卡進行大規模模型訓練的固有限制。關鍵是如何找到合適的方法,在滿足顯卡記憶體限制的前提下,合理地輸入模型和數據。於是,自然產生了一個疑問: 為何至今仍未有人實作這一目標呢?

學術界不乏擅長解決難題的傑出人才。但對於解決這個特定問題,學術界尚未找到答案。這是因為大學研究人員很難證明花費時間在這類工作上的合理性。將現有工具和技術組合在一起通常被認為不夠新穎, 不足以在高影響力期刊上發表論文。

大型科技公司也滿載著解決難題的天才人士。但對於使用消費級GPU訓練模型這個特定問題,他們並不需要解決——因為 各大行業巨頭實驗室已擁有昂貴得多的硬體裝置,因此缺乏動力去攻克這個問題。

創始人 Eric Ries 用一句話 解釋了這一困境: 今天的金融市場迫使企業優先考慮短期收益而非其他一切 」。

簡而言之,這活兒只能讓初創公司甚至個人開發者去幹。

Answer.AI 的核心創新觀念很直接: 探尋如何利用這些價格親民、記憶體相對較小的遊戲顯卡訓練出一流的開源大模型。 具體目標設定為:僅用遊戲顯卡訓練一個擁有 700 億參數(70b)的模型,即訓練用的單塊 GPU 記憶體不超過 24 GB。這項任務頗具挑戰性,因為每個參數通常占用 16 位(2 字節),因此即便是儲存模型權重就需要至少 140 GB 的空間,還未計入諸如啟用值、梯度以及最佳化狀態等其它相關數據。

兩項技術釋出為 Answer.AI 開辟了新的道路,即 QLoRA (由 Tim Dettmers 等研究人員提出)及 FSDP (來自 Meta 的 PyTorch 團隊)。

QLoRA:在單個 GPU 上訓練更大的模型

QLoRA 巧妙地結合了現代神經網路領域的兩項關鍵技術突破:量化技術(quantization)和 LoRA。

量化技術是將原本需用 16 位乃至 32 位表示的神經網路權重縮減至僅用 4 位(甚至更少)表示。盡管 4 位數位僅有 16 種可能的取值,但 Tim Dettmers 等人證明對於當前廣泛套用的大型語言模型而言,這一點已足夠。

此外,Tim Dettmers 開發的 bits and bytes 庫使得建立 4 位「量化」模型變得簡便易行,同時 Hugging Face 團隊也開始參與對該庫的維護和文件編寫工作。

然而,量化後的模型面臨一個重大難題,即無法用常規方法繼續訓練——由於可供表達的數值有限,梯度下降法在幾乎整個權重空間內都無法獲得有效的梯度,從而無法對量化權重進行更新。 這意味著量化技術目前只能套用於推理階段,而非模型的預訓練或微調階段。 盡管推理套用具有重要意義,但未來的建模不應該是單向的消費行為,而是人人都能參與共創的過程。

解決這一困境的關鍵在於引入 LoR A 技術,即「Low-Rank Adaptation of Large Language Models」。LoRA 並不直接對整個大型模型進行訓練,而是添加體積非常小(通常不到整個模型體積的 1%)的「介面卡」,在維持模型主體不 變的前提下單獨訓練這些介面卡。像 Stable Diffusion 這樣的模型常常會使用此類介面卡,這也是為何它們能以較小體積快速下載的原因。

Tim Dettmers 發現量化技術與 LoRA 可以有效結合:采用一個固定的量化基線模型,附加不可量化的、可訓練的 LoRA 介面卡,這就形成了 QLoRA 技術。借助這項技術,他的團隊首次在一塊 48 GB GPU 上成功訓練了一個(未經量化時)超過 GPU 記憶體容量的 65b 模型(未量化時大小為 130 GB)。

Hugging Face 團隊再度出手,建立了 PEFT 庫,大大簡化了 LoRA 訓練流程,並將其與 bits and bytes 庫直接整合,使得使用者只需寥寥幾行程式碼就能使用 QLoRA 技術。 QLoRA 雖未能完全攻克 Answer.AI 在 24 GB 顯卡上訓練 70b 模型的目標,但它已經取得了前所未有的進展。 若將 70b 模型量化至 4 位(即 0.5 字節),則占用 35 GB 空間,超過了 24 GB 遊戲顯卡的可用記憶體。

QLoRA 尚存在其他局限。首先,48 GB 的高端顯卡價格不菲,而且僅能勉強放下 65b 模型的訓練任務。除此之外,在訓練過程中,還需要考慮儲存大量額外數據,包括啟用值、梯度以及模型的最佳化狀態,這就要求有足夠的剩余記憶體。若載入完模型權重後記憶體所剩無幾,那麽模型訓練時的執行記憶體就會捉襟見肘。

舉例來說,語言模型的一大優點是可以模擬對話或者分析長篇文本。為了讓模型能夠處理較長的序列,訓練時就需要提供長序列樣本。訓練中使用的最長序列稱為「序列長度」。在 48 GB 顯卡上訓練 65b QLoRA 模型時,如果選擇過長的序列長度,將會因記憶體不足而導致訓練失敗,因為大部份記憶體都被用於儲存模型權重,無法承載更多序列資訊。

此外,如果模型每次只能處理一個序列,則完成整個訓練數據集的叠代周期將極為漫長。為了加速訓練,通常我們會采取「批次處理」的方式,一次性處理多個序列,這裏的「批次大小」指的就是一次處理的序列數目。當 GPU 記憶體緊巴巴地裝入模型權重之後,僅能使用很小的批次大小,進而導致訓練速度極度緩慢。

FSDP:在多個 GPU 上擴充套件訓練規模

面對消費級 GPU 記憶體限制問題,一種直觀的解決方案是采用多塊 GPU 進行訓練。

在開源社群中,常見做法是將模型的部份層分布在每張顯卡上。例如,在訓練過程中,首先在第一塊 GPU 上執行前幾層,然後在第二塊 GPU 上處理接下來的層,以此類推。如此一來,一個 70b(140 GB)的大規模模型可以分布到 8 塊 24 GB GPU上,每塊 GPU 占用約 17.5 GB 記憶體。

因此,需要尋求一種既 能分 割模型至多個 GPU 上,又能充分利用所有 GPU 進行平行計算的方法。Meta 公司研發的 FSDP 庫恰好填補了這一空白。FSDP 透過將大型模型參數跨多個 GPU 進行分片,使得所有 GPU 能夠同時得到充分使用。在神經網路某一層於特定 GPU 上進行訓練時,會臨時復制必要的參數分片至該 GPU 進行計算,計算完成後刪除這些臨時數據。

舉例來說,一個未量化的 70b(700 億參數)模型需要 140 GB 記憶體,即使是最頂級的 NVIDIA H100 顯卡(售價約為 4 萬美元一塊,配備 80 GB 記憶體)也無法滿足需求。然而,借助 FSDP 技術,四塊 H100 GPU 組合起來能提供總計 320 GB 的記憶體空間。

值得註意的是,構建這樣一台搭載多塊 H100 GPU 的機器可能需要花費大約 15 萬美元左右。

找到方向後,Answer.AI 的首要任務之一便是探索如何結合 QLoRA 技術,使得普通桌上型電腦上的遊戲級 GPU 也能高效地訓練 70b 模型。

合二為一!

設想一下:透過 QLoRA 將模型大小壓縮約 400%,使其適應 35 GB 記憶體,再利用 FSDP 將其分片分配到兩個或更多 24 GB 消費者級別的顯卡上,從而騰出足夠的記憶體用於模型訓練。

第一步是在 2023 年末, Answer.AI 創始人 Jeremy Tim Dettmers 討論了將 FSDP 和 QLoRA 相結合的可能性。 Tim Dettmers Jerem y 引薦給了 Hugging Face 的 Titus von Koeller,兩人共同合作研究、理解和記錄了整合這兩個庫時遇到的問題。

經過原型設計及偵錯,團隊發現並解決了 FSDP 未能正確復制各分片所需的模型量化資訊的問題。他們註意到,由於 FSDP 對同步至各 GPU 間的數據子集有嚴格的規範,必須確保量化狀態(包含參數(去)量化所必需的資訊)從參數轉移至模型層內部,以避免在 FSDP 移動分片時遺失相關資訊。

在這些問題得以解決後,研究人員成功運用 FSDP 和量化模型完成了第一批數據的訓練,並將這些改進、相應的測試以及重構內容整理成針對 bits and bytes 計畫的拉取請求。值得慶幸的是,bits and bytes 計畫維護者們迅速響應,協助這些改動順利透過審查流程。

隨著計畫接近尾聲,研究人員意識到實際情況下仍然難以直接載入大於單個 GPU 容量的大型量化模型,因為模型載入和量化過程本身就需要先將整個模型置於單個 GPU 上。

Jeremy 又花費數周時間仔細研究了 Meta 公司的 Llama-Recipes 計畫,該計畫是他找到的最佳 FSDP 微調完整實作案例。研究團隊調整後發現,可以逐層進行載入和離散化操作,從而避免在整個模型一次性位於單個 GPU 上的需求。與此同時,研究人員還編寫了一種客製版 LoRA 演算法實作,使之能與以上改進協同工作。

最終, Answer.AI 實作了在雙 3090 遊戲 GPU 上成功微調 70b 模型的目標。

如何上手?

Answer.AI 的官方部落格有更詳盡的辦法,此處只進行簡單介紹。

原文連結:https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html#a-first-step

首先,你需要擁有不止一塊 GPU。當然,現在也有不少社群雲服務都提供這類機器的租賃。

其次,需要安裝最新版本的 Transformers、PEFT 和 bitsandbytes。接著,複制 Answer.AI 的倉庫並遵循其中的 README 檔操作。

倉庫連結:https://github.com/AnswerDotAI/fsdp_qlora/tree/main

Answer.AI 的這一計畫振奮人心,他們不僅證明了 在低成本遊戲級 GPU 上擴充套件資源高效的 QLoRA 訓練規模是可行的,還引起了更多人關註降低模型訓練成本這一問題。希望未來的 AI 能 變得更易於存取,讓更多人不僅能夠消費模型,也能構建有價值的模型。

4 月 25 ~ 26 日,由 CSDN 和高端 IT 咨詢和教育平台 Boolan 聯合主辦的「 全球機器學習技術大會 」將在上海環球港凱悅酒店舉行,特邀近 50 位技術領袖和行業套用專家,與 1000+ 來自電商、金融、汽車、智慧制造、通訊、工業互聯網、醫療、教育等眾多行業的精英參會聽眾,共同探討人工智慧領域的前沿發展和行業最佳實踐。 歡迎所有開發者朋友存取官網 http://ml-summit.org、點選「閱讀原文」或掃描下方海報中的二維碼 ,進一步了解詳情。