AI Agent 如何實現(xiàn)?6張4090 魔改Llama2:一句指令拆分任務(wù)、調(diào)用函數(shù)
AI Agent 是時下熱門的一個方向,在 OpenAI 應(yīng)用研究主管 LilianWeng 寫的萬字長文中[1],她提出 Agent = LLM+ 記憶 + 規(guī)劃技能 + 工具使用。
圖1 Overview of a LLM-powered autonomous agent system
簡單來說,Agent 就是借助 LLM 的強(qiáng)大語言理解能力、邏輯推理能力調(diào)用工具來幫助人類完成任務(wù)。但同時也存在著一些挑戰(zhàn),例如基礎(chǔ)模型的能力決定 agent 調(diào)用工具的效率,但基礎(chǔ)模型本身存在著大模型幻覺等問題。
本文以「輸入一段指令自動實現(xiàn)復(fù)雜任務(wù)拆分和函數(shù)調(diào)用」的場景為例,來構(gòu)建基礎(chǔ) Agent 流程,并側(cè)重講解如何通過「基礎(chǔ)模型選擇」、「Prompt設(shè)計」等來成功構(gòu)建「任務(wù)拆分」和「函數(shù)調(diào)用」模塊。
項目地址:
https://sota.jiqizhixin.com/project/smart_agent
GitHub Repo:
https://github.com/zzlgreat/smart_agent
任務(wù)拆分&函數(shù)調(diào)用 Agent 流程
對于實現(xiàn)「輸入一段指令自動實現(xiàn)復(fù)雜任務(wù)拆分和函數(shù)調(diào)用」,項目構(gòu)建的 Agent 流程如下:
- planner:根據(jù)用戶輸入的指令拆分任務(wù)。確定自己擁有的工具列表 toolkit,告訴拆分任務(wù)的大模型 planner 自己具有哪些工具,需要完成什么樣的任務(wù),planner 把任務(wù)拆分為計劃 1,2,3...
- distributor:負(fù)責(zé)選擇適當(dāng)?shù)墓ぞ邅?toolkit 執(zhí)行計劃。函數(shù)調(diào)用模型需要根據(jù)計劃的不同分別選擇對應(yīng)的工具。
- worker:負(fù)責(zé)調(diào)用工具箱中的任務(wù),并且返回任務(wù)調(diào)用的結(jié)果。
- solver:整理出來的分布計劃和對應(yīng)的結(jié)果組合為一個 long story,再由 solver 進(jìn)行總結(jié)歸納。
圖1 《ReWOO: Decoupling Reasoning from Observations for Efficient Augmented Language Models》
為了實現(xiàn)上述流程,在「任務(wù)拆分」和「函數(shù)調(diào)用」模塊中,項目分別設(shè)計了兩個微調(diào)模型,來實現(xiàn)將復(fù)雜任務(wù)拆分并按需調(diào)用自定義函數(shù)的功能。歸納總結(jié)的模型 solver,可以和拆分任務(wù)模型相同。
微調(diào)任務(wù)拆分&函數(shù)調(diào)用模型
2.1 微調(diào)經(jīng)驗總結(jié)
在「任務(wù)拆分」模塊中,大模型需要具有將復(fù)雜任務(wù)分解為簡單任務(wù)的能力。「任務(wù)拆分」成功與否,極大程度上取決于兩點:
- 基礎(chǔ)模型選擇:為了拆分復(fù)雜任務(wù),選擇微調(diào)的基礎(chǔ)模型本身需要具備良好的理解和泛化能力,即根據(jù) prompt 指令拆分訓(xùn)練集中未見的任務(wù)。目前來講,選擇高參數(shù)的大模型更容易做到這一點。
- Prompt 設(shè)計:prompt 能否成功地調(diào)用模型的思維鏈,將任務(wù)拆分為子任務(wù)。
同時希望任務(wù)拆分模型在給定 prompt 模板下的輸出格式可以盡可能相對固定,但也不會過擬合喪失模型原本的推理和泛化能力,這里采取 lora 微調(diào) qv 層,對原模型的結(jié)構(gòu)改動盡可能地少。
在「函數(shù)調(diào)用」模塊中,大模型需要具備穩(wěn)定調(diào)用工具的能力,以適應(yīng)處理任務(wù)的要求:
- 損失函數(shù)調(diào)整:除選擇的基礎(chǔ)模型本身泛化能力、Prompt 設(shè)計外,為實現(xiàn)模型的輸出盡可能地固定、根據(jù)輸出穩(wěn)定調(diào)用所需函數(shù),采用「prompt loss-mask」的方法[2]進(jìn)行 qlora 訓(xùn)練(詳見下文闡述),并通過魔改 attention mask 的方式,在 qlora 微調(diào)中使用插入 eos token 的小技巧來穩(wěn)定住模型的輸出。
此外,在算力使用方面,通過 lora/qlora 微調(diào)實現(xiàn)了低算力條件下大型語言模型的微調(diào)和推理,并采用量化部署的方式,進(jìn)一步降低推理的門檻。
2.2 基礎(chǔ)模型選擇
針對「任務(wù)拆分」模型的選擇,希望模型同時具備強(qiáng)泛化能力和一定的思維鏈能力。這里可以參考 HuggingFace 上 Open LLM Leaderboard 來選擇模型,更多關(guān)注的是衡量文本模型多任務(wù)準(zhǔn)確性的測試 MMLU 和綜合評分 Average。
圖2 HuggingFace Open LLM Leaderboard (0921)
本項目選定任務(wù)拆分模型型號為:
- AIDC-ai-business/Marcoroni-70B:該模型基于 Llama2 70B 微調(diào),負(fù)責(zé)拆分任務(wù)。根據(jù) HuggingFace 上 Open LLM Leaderboard 顯示,該模型的 MMLU 和 average 都比較高,而且該模型的訓(xùn)練過程中加入了大量的 Orca 風(fēng)格的數(shù)據(jù),適用于多輪對話,在 plan-distribute-work-plan-work……summary 的流程中效果表現(xiàn)會更好。
針對「函數(shù)調(diào)用」模型的選擇,meta 開源的 Llama2 版代碼編程模型 CodeLlama 原始訓(xùn)練數(shù)據(jù)包含了大量的代碼數(shù)據(jù),這樣就可以嘗試自定義腳本的 qlora 微調(diào)。針對函數(shù)調(diào)用的模型,選擇 CodeLlama 模型(34b/13b/7b 均可)作為基座。
本項目選定函數(shù)調(diào)用模型型號為:
- codellama 34b/7b:負(fù)責(zé)函數(shù)調(diào)用的模型,該模型采用大量代碼數(shù)據(jù)訓(xùn)練,代碼數(shù)據(jù)中必然包含大量對函數(shù)的描述類自然語言,對于給定函數(shù)的描述具有良好的 zero-shot 能力。
針對「函數(shù)調(diào)用」模型微調(diào),項目采用 prompt loss mask 的方式進(jìn)行 qlora 訓(xùn)練,以穩(wěn)定處理模型輸出。損失函數(shù)調(diào)整方式如下:
- 損失掩碼 (loss_mask):
- loss_mask 是一個與輸入序列 input_ids 形狀相同的張量 (tensor)。每個元素都是 0 或 1,其中 1 表示對應(yīng)的位置的標(biāo)簽應(yīng)被考慮在損失計算中,而 0 表示不應(yīng)被考慮。
- 例如,如果某些標(biāo)簽是填充的(通常是因為批處理中的序列長度不同),不想在損失的計算中考慮這些填充的標(biāo)簽。在這種情況下,loss_mask 為這些位置提供了一個 0,從而遮蔽掉了這些位置的損失。
- 損失計算:
- 首先,使用了 CrossEntropyLoss 來計算未 mask 的損失。
設(shè)置 reductinotallow='none' 來確保為序列中的每個位置都返回一個損失值,而不是一個總和或平均值。 - 然后,使用 loss_mask 來 mask 損失。通過將 loss_mask 與 losses 相乘,得到了 masked_loss。這樣,loss_mask 中為 0 的位置在 masked_loss 中的損失值也為 0。
- 損失聚合:
- 將所有的 masked_loss 求和,并通過 loss_mask.sum() 來歸一化。這確保了你只考慮了被 mask 為 1 的標(biāo)簽的損失。為了防止除以零的情況,加一個很小的數(shù) 1e-9。
- 如果 loss_mask 的所有值都是 0(即 loss_mask.sum() == 0),那么直接返回一個 0 的損失值。
2.3 硬件需求:
- 6*4090 for Marcoroni-70B’s 16bit lora
- 2*4090 for codellama 34b’s qlora / 1*4090 for codellama 13/7b’s qlora
2.4 Prompt 格式設(shè)計
對于任務(wù)拆分,項目采用了大型語言模型高效推理框架 ReWOO(Reasoning WithOut Observation)[2] 中 planner 設(shè)計的 Prompt 格式。這里只需把'Wikipedia[input]'等函數(shù)替換為的函數(shù)和描述即可,該 prompt 示例如下:
For the following tasks, make plans that can solve the problem step-by-step.
For each plan, indicate which external tool together with tool input to retrieve evidence.
You can store the evidence into a variable #E that can be called by later tools.
(Plan, #E1, Plan, #E2, Plan, ...) Tools can be one of the following:
Wikipedia[input]: Worker that search for similar page contents from Wikipedia.
Useful when you need to get holistic knowledge about people, places, companies, historical events, or other subjects.
The response are long and might contain some irrelevant information. Input should be a search query.
LLM[input]: A pretrained LLM like yourself. Useful when you need to act with general world knowledge and common sense.
Prioritize it when you are confident in solving the problem yourself. Input can be any instruction.
對于函數(shù)調(diào)用,因為后續(xù)會進(jìn)行 qlora 微調(diào),所以直接采用 huggingface 上開源函數(shù)調(diào)用數(shù)據(jù)集 [3] 中的 prompt 樣式。請參見下文。
指令數(shù)據(jù)集準(zhǔn)備
3.1 數(shù)據(jù)來源
- 拆任務(wù)模型:Marcoroni-70B 采用的是 alpaca 的提示模板。該模型在 Llama2 70B 上進(jìn)行指令微調(diào),為和原始模型的模板進(jìn)行對齊,需采用 alpaca 格式的數(shù)據(jù)集。這里使用 rewoo 的 planner 數(shù)據(jù)集格式,但在原始數(shù)據(jù)集中只有調(diào)用 wiki 和自己的選項,所以可以套用該模板,并采用 gpt4 的接口來制作該樣式的數(shù)據(jù)集。
- 函數(shù)調(diào)用模型:盡管所選用的 HuggingFace 開源函數(shù)調(diào)用的數(shù)據(jù)集的數(shù)據(jù)量較少(55 行),但 qlora 十分有效,在該數(shù)據(jù)集中還附帶了代碼訓(xùn)練模板。
3.2 數(shù)據(jù)集格式
- 任務(wù)拆分模型數(shù)據(jù)格式:alpaca
### Instruction:
<prompt> (without the <>)
### Response:
如:
### Instruction:
For the following tasks, make plans that can solve the problem step-by-step. For each plan, indicate which external tool together with tool input to retrieve evidence. You can store the evidence into a variable #E that can be called by later tools. (Plan, #E1, Plan, #E2, Plan, ...) Tools can be one of the following:get_current_time[input]: Get the current time for now. search_bing[input]: Search the web for content on Bing, and return the top 10 results. Remember, the results may not as good as Wikipedia. If you think the result will not exist in wikipedia, use it. search_wiki[input]: Search the web for content on Wikipedia, and return the top 3 results. If you think the result can exist in wikipedia, use it. search_bilibili[input]: Search for videos on Bilibili, which is a website full of ACG. It can return top 10 search results. save_file[input]: save a file in the target dir which contains the data. query_mysql[input]: Execute the sql in the local mysql database. And return the result.
What movies did the director of 'Oppenheim' direct? List the top 10 best.
### Response:
Step 1: Identify the director of "Oppenheimer" movie.
#E1 = search_wiki("Director of Oppenheimer")
Step 2: Retrieve information about the director's other works.
If the retrieved name exists in Wikipedia, proceed with searching their filmography using search_wiki. Otherwise, try searching through Bing or another source like IMDb.
For example, if Christopher Nolan is found as the director of "Oppenheimer":
#E2 = search_bing("Filmography of Christopher Nolan")
Step 3: Extract the list of films directed by this person from the retrieved information.
From the returned results, extract the titles of the top 10 best movies according to critics/audience ratings or personal preference.
- 函數(shù)調(diào)用模型數(shù)據(jù)格式:
這里采取 trelis 的數(shù)據(jù)集的格式。數(shù)據(jù)集量并不多,只有 55 行。結(jié)構(gòu)上其實也是 alpaca 格式。分為 systemPrompt,userPrompt,assistantResponse 分別對應(yīng) alpaca 的 Instruction,prompt,Response。示例如下:
圖3 HuggingFace 函數(shù)調(diào)用開源數(shù)據(jù)集示例
微調(diào)過程說明
4.1 微調(diào)環(huán)境
ubuntu 22.04 系統(tǒng),cuda 11.8,Pytorch 2.0.1,采取了 LLaMA-Efficient-Tuning 框架。Deepspeed 0.10.4
4.2 微調(diào)步驟
1)針對 Marcoroni-70B 的 lora 微調(diào)
- LLaMA-Efficient-Tuning 框架支持 deepspeed 集成,在訓(xùn)練開始前輸入 accelerate config 進(jìn)行設(shè)置,根據(jù)提示選擇 deepspeed zero stage 3,因為是 6 卡總計 144G 的 VRAM 做 lora 微調(diào),offload optimizer states 可以選擇 none, 不卸載優(yōu)化器狀態(tài)到內(nèi)存。
- offload parameters 需要設(shè)置為 cpu,將參數(shù)量卸載到內(nèi)存中,這樣內(nèi)存峰值占用最高可以到 240G 左右。gradient accumulation 需要和訓(xùn)練腳本保持一致,這里選擇的是 4。gradient clipping 用來對誤差梯度向量進(jìn)行歸一化,設(shè)置為 1 可以防止梯度爆炸。
- zero.init 可以進(jìn)行 partitioned 并轉(zhuǎn)換為半精度,加速模型初始化并使高參數(shù)的模型能夠在 CPU 內(nèi)存中全部進(jìn)行分配。這里也可以選 yes。
全部選擇完成后,新建一個訓(xùn)練的 bash 腳本,內(nèi)容如下:
accelerate launch src/train_bash.py \
--stage sft \
--model_name_or_path your_model_path \
--do_train \
--dataset rewoo \
--template alpaca \
--finetuning_type lora \
--lora_target q_proj,v_proj \
--output_dir your_output_path \
--overwrite_cache \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 4 \
--lr_scheduler_type cosine \
--logging_steps 10 \
--save_steps 1000 \
--learning_rate 5e-6 \
--num_train_epochs 4.0 \
--plot_loss \
--flash_attn \
--bf16
這樣的設(shè)置需要的內(nèi)存峰值最高可以到 240G,但還是保證了 6 卡 4090 可以進(jìn)行訓(xùn)練。開始的時候可能會比較久,這是因為 deepspeed 需要對模型進(jìn)行 init。之后訓(xùn)練就開始了。
圖4 6 卡 4090 訓(xùn)練帶寬速度
共計用時 8:56 小時。本次訓(xùn)練中因為主板上的 NVME 插槽會和 OCULINK 共享一路 PCIE4.0 x16 帶寬。所以 6 張中的其中兩張跑在了 pcie4.0 X4 上,從上圖就可以看出 RX 和 TX 都只是 PCIE4.0 X4 的帶寬速度。這也成為了本次訓(xùn)練中最大的通訊瓶頸。如果全部的卡都跑在 pcie 4.0 x16 上,速度應(yīng)該是比現(xiàn)在快不少的。
圖5 LLaMA-Efficient-Tuning 生成 loss 曲線
以上是 LLaMA-Efficient-Tuning 自動生成的 loss 曲線,可以看到 4 個 epoch 后收斂效果還是不錯的。
2)針對 codellama 的 qlora 微調(diào)
根據(jù)上文所述的 prompt loss mask 方法重構(gòu)了 trainer 類(見項目代碼倉庫 func_caller_train.py)。因為數(shù)據(jù)集本身比較?。?5 行)。所以跑 4 個 epoch 只用了兩分鐘,模型很快達(dá)到了收斂。
4.3微調(diào)完成后的測試效果
在項目代碼倉庫中,提供了一個簡短可用的 toolkit 示例。里面的函數(shù)包括:
- 必應(yīng)搜索
- 維基搜索
- bilibili 搜索
- 獲取當(dāng)前時間
- 保存文件
- ...
現(xiàn)在有一個 70B 和一個 34B 的模型,在實際使用中,用 6 張 4090 同時以 bf16 精度運(yùn)行這兩個模型是不現(xiàn)實的。但是可以通過量化的方法壓縮模型大小,同時提升模型推理速度。這里采用高性能 LLM 推理庫 exllamav2 運(yùn)用 flash_attention 特性來對模型進(jìn)行量化并推理。在項目頁面中作者介紹了一種獨(dú)特的量化方式,本文不做贅述。按照其中的轉(zhuǎn)換機(jī)制可以將 70b 的模型按照 2.5-bit 量化為 22G 的大小,這樣一張顯卡就可以輕松加載。
1)測試方法
給定一段不在訓(xùn)練集中的復(fù)雜任務(wù)描述,同時在 toolkit 中添加訓(xùn)練集中不包含的函數(shù)和對應(yīng)描述。如果 planner 可以完成對任務(wù)進(jìn)行拆分,distributor 可以調(diào)用函數(shù),solver 可以根據(jù)整個流程對結(jié)果進(jìn)行總結(jié)。
2)測試結(jié)果
任務(wù)拆分:先使用 text-generation-webui 快速測試一下任務(wù)拆分模型的效果,如圖所示:
圖6 任務(wù)拆分測試結(jié)果
這里可以寫一個簡單的 restful_api 接口,方便在 agent 測試環(huán)境下的調(diào)用(見項目代碼 fllama_api.py)。
函數(shù)調(diào)用:在項目中已經(jīng)寫好了一個簡單的 planner-distributor-worker-solver 的邏輯。接下來就讓測試一下這個任務(wù)。輸入一段指令:what movies did the director of 'Killers of the Flower Moon' direct?List one of them and search it in bilibili.
「搜索 bilibili 」這個函數(shù)是不包含在項目的函數(shù)調(diào)用訓(xùn)練集中的。同時這部電影也是一部還沒有上映的新電影,不確定模型本身的訓(xùn)練數(shù)據(jù)有沒有包含。可以看到模型很好地將輸入指令進(jìn)行拆分:
- 從維基百科上搜索該電影的導(dǎo)演
- 根據(jù) 1 的結(jié)果,從 bing 上搜索電影 Goodfellas 的結(jié)果
- 在 bilibili 上搜索電影 Goodfellas
同時進(jìn)行函數(shù)調(diào)用得到了以下結(jié)果:點擊結(jié)果是 Goodfellas,和該部電影的導(dǎo)演匹配得上。
總結(jié)
本項目以「輸入一段指令自動實現(xiàn)復(fù)雜任務(wù)拆分和函數(shù)調(diào)用」場景為例,設(shè)計了一套基本 agent 流程:toolkit-plan-distribute-worker-solver 來實現(xiàn)一個可以執(zhí)行無法一步完成的初級復(fù)雜任務(wù)的 agent。通過基礎(chǔ)模型的選型和 lora 微調(diào)使得低算力條件下一樣可以完成大模型的微調(diào)和推理。并采用量化部署的方式,進(jìn)一步降低推理的門檻。最后通過該 pipeline 實現(xiàn)了一個搜索電影導(dǎo)演其他作品的示例,實現(xiàn)了基礎(chǔ)的復(fù)雜任務(wù)完成。
局限性:本文只是基于搜索和基本操作的 toolkit 設(shè)計了函數(shù)調(diào)用和任務(wù)拆分。使用的工具集非常簡單,并沒有太多設(shè)計。針對容錯機(jī)制也沒有太多考慮。通過本項目,大家也可以繼續(xù)向前一步探索 RPA 領(lǐng)域上的應(yīng)用,進(jìn)一步完善 agent 流程,實現(xiàn)更高程度的智能自動化提升流程的可管理性。