AI Agent 是時下熱門的一個方向，在 OpenAI 應用研究主管 LilianWeng 寫的萬字長文中[1]，她提出 Agent = LLM+ 記憶 + 規劃技能 + 工具使用。

圖1 Overview of a LLM-powered autonomous agent system

簡單來說，Agent 就是借助 LLM 的強大語言理解能力、邏輯推理能力調用工具來幫助人類完成任務。但同時也存在著一些挑戰，例如基礎模型的能力決定 agent 調用工具的效率，但基礎模型本身存在著大模型幻覺等問題。

本文以「輸入一段指令自動實現復雜任務拆分和函數調用」的場景為例，來構建基礎 Agent 流程，并側重講解如何通過「基礎模型選擇」、「Prompt設計」等來成功構建「任務拆分」和「函數調用」模塊。

項目地址：

https://sota.jiqizhixin.com/project/smart_agent

GitHub Repo：

https://github.com/zzlgreat/smart_agent

任務拆分&函數調用 Agent 流程

對于實現「輸入一段指令自動實現復雜任務拆分和函數調用」，項目構建的 Agent 流程如下：

• planner：根據用戶輸入的指令拆分任務。確定自己擁有的工具列表 toolkit，告訴拆分任務的大模型 planner 自己具有哪些工具，需要完成什么樣的任務，planner 把任務拆分為計劃 1，2，3...
• distributor：負責選擇適當的工具來 toolkit 執行計劃。函數調用模型需要根據計劃的不同分別選擇對應的工具。
• worker：負責調用工具箱中的任務，并且返回任務調用的結果。
• solver：整理出來的分布計劃和對應的結果組合為一個 long story，再由 solver 進行總結歸納。

圖1 《ReWOO: Decoupling Reasoning from Observations for Efficient Augmented Language Models》

為了實現上述流程，在「任務拆分」和「函數調用」模塊中，項目分別設計了兩個微調模型，來實現將復雜任務拆分并按需調用自定義函數的功能。歸納總結的模型 solver，可以和拆分任務模型相同。

微調任務拆分&函數調用模型

2.1 微調經驗總結

在「任務拆分」模塊中，大模型需要具有將復雜任務分解為簡單任務的能力。「任務拆分」成功與否，極大程度上取決于兩點：

• 基礎模型選擇：為了拆分復雜任務，選擇微調的基礎模型本身需要具備良好的理解和泛化能力，即根據 prompt 指令拆分訓練集中未見的任務。目前來講，選擇高參數的大模型更容易做到這一點。
• Prompt 設計：prompt 能否成功地調用模型的思維鏈，將任務拆分為子任務。

同時希望任務拆分模型在給定 prompt 模板下的輸出格式可以盡可能相對固定，但也不會過擬合喪失模型原本的推理和泛化能力，這里采取 lora 微調 qv 層，對原模型的結構改動盡可能地少。

在「函數調用」模塊中，大模型需要具備穩定調用工具的能力，以適應處理任務的要求:

• 損失函數調整：除選擇的基礎模型本身泛化能力、Prompt 設計外，為實現模型的輸出盡可能地固定、根據輸出穩定調用所需函數，采用「prompt loss-mask」的方法[2]進行 qlora 訓練（詳見下文闡述），并通過魔改 attention mask 的方式，在 qlora 微調中使用插入 eos token 的小技巧來穩定住模型的輸出。

此外，在算力使用方面，通過 lora/qlora 微調實現了低算力條件下大型語言模型的微調和推理，并采用量化部署的方式，進一步降低推理的門檻。

2.2 基礎模型選擇

針對「任務拆分」模型的選擇，希望模型同時具備強泛化能力和一定的思維鏈能力。這里可以參考 HuggingFace 上 Open LLM Leaderboard 來選擇模型，更多關注的是衡量文本模型多任務準確性的測試 MMLU 和綜合評分 Average。

圖2 HuggingFace  Open LLM Leaderboard （0921）

本項目選定任務拆分模型型號為：

• AIDC-ai-business/Marcoroni-70B：該模型基于 Llama2 70B 微調，負責拆分任務。根據 HuggingFace 上 Open LLM Leaderboard 顯示，該模型的 MMLU 和 average 都比較高，而且該模型的訓練過程中加入了大量的 Orca 風格的數據，適用于多輪對話，在 plan-distribute-work-plan-work……summary 的流程中效果表現會更好。

針對「函數調用」模型的選擇，meta 開源的 Llama2 版代碼編程模型 CodeLlama 原始訓練數據包含了大量的代碼數據，這樣就可以嘗試自定義腳本的 qlora 微調。針對函數調用的模型，選擇 CodeLlama 模型（34b/13b/7b 均可）作為基座。

本項目選定函數調用模型型號為：

• codellama 34b/7b：負責函數調用的模型，該模型采用大量代碼數據訓練，代碼數據中必然包含大量對函數的描述類自然語言，對于給定函數的描述具有良好的 zero-shot 能力。

針對「函數調用」模型微調，項目采用 prompt loss mask 的方式進行 qlora 訓練，以穩定處理模型輸出。損失函數調整方式如下：

• 損失掩碼 （loss_mask）：

1. loss_mask 是一個與輸入序列 input_ids 形狀相同的張量 （tensor）。每個元素都是 0 或 1，其中 1 表示對應的位置的標簽應被考慮在損失計算中，而 0 表示不應被考慮。
2. 例如，如果某些標簽是填充的（通常是因為批處理中的序列長度不同），不想在損失的計算中考慮這些填充的標簽。在這種情況下，loss_mask 為這些位置提供了一個 0，從而遮蔽掉了這些位置的損失。

• 損失計算：

1. 首先，使用了 CrossEntropyLoss 來計算未 mask 的損失。
   設置 reductinotallow='none' 來確保為序列中的每個位置都返回一個損失值，而不是一個總和或平均值。
2. 然后，使用 loss_mask 來 mask 損失。通過將 loss_mask 與 losses 相乘，得到了 masked_loss。這樣，loss_mask 中為 0 的位置在 masked_loss 中的損失值也為 0。

• 損失聚合：

1. 將所有的 masked_loss 求和，并通過 loss_mask.sum() 來歸一化。這確保了你只考慮了被 mask 為 1 的標簽的損失。為了防止除以零的情況，加一個很小的數 1e-9。
2. 如果 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 格式設計

對于任務拆分，項目采用了大型語言模型高效推理框架 ReWOO（Reasoning WithOut Observation）[2] 中 planner 設計的 Prompt 格式。這里只需把'Wikipedia[input]'等函數替換為的函數和描述即可，該 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.

對于函數調用，因為后續會進行 qlora 微調，所以直接采用 huggingface 上開源函數調用數據集 [3] 中的 prompt 樣式。請參見下文。

指令數據集準備

3.1 數據來源

• 拆任務模型：Marcoroni-70B 采用的是 alpaca 的提示模板。該模型在 Llama2 70B 上進行指令微調，為和原始模型的模板進行對齊，需采用 alpaca 格式的數據集。這里使用 rewoo 的 planner 數據集格式，但在原始數據集中只有調用 wiki 和自己的選項，所以可以套用該模板，并采用 gpt4 的接口來制作該樣式的數據集。
• 函數調用模型：盡管所選用的 HuggingFace 開源函數調用的數據集的數據量較少（55 行），但 qlora 十分有效，在該數據集中還附帶了代碼訓練模板。

3.2 數據集格式

• 任務拆分模型數據格式：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.

• 函數調用模型數據格式：

這里采取 trelis 的數據集的格式。數據集量并不多，只有 55 行。結構上其實也是 alpaca 格式。分為 systemPrompt，userPrompt，assistantResponse 分別對應 alpaca 的 Instruction，prompt，Response。示例如下：

圖3 HuggingFace 函數調用開源數據集示例

微調過程說明

4.1 微調環境

ubuntu 22.04 系統，cuda 11.8，Pytorch 2.0.1，采取了 LLaMA-Efficient-Tuning 框架。Deepspeed 0.10.4

4.2 微調步驟

1）針對 Marcoroni-70B 的 lora 微調

1. LLaMA-Efficient-Tuning 框架支持 deepspeed 集成，在訓練開始前輸入 accelerate config 進行設置，根據提示選擇 deepspeed zero stage 3，因為是 6 卡總計 144G 的 VRAM 做 lora 微調，offload optimizer states 可以選擇 none， 不卸載優化器狀態到內存。
2. offload parameters 需要設置為 cpu，將參數量卸載到內存中，這樣內存峰值占用最高可以到 240G 左右。gradient accumulation 需要和訓練腳本保持一致，這里選擇的是 4。gradient clipping 用來對誤差梯度向量進行歸一化，設置為 1 可以防止梯度爆炸。
3. zero.init 可以進行 partitioned 并轉換為半精度，加速模型初始化并使高參數的模型能夠在 CPU 內存中全部進行分配。這里也可以選 yes。

全部選擇完成后，新建一個訓練的 bash 腳本，內容如下：

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

這樣的設置需要的內存峰值最高可以到 240G，但還是保證了 6 卡 4090 可以進行訓練。開始的時候可能會比較久，這是因為 deepspeed 需要對模型進行 init。之后訓練就開始了。

 圖4 6 卡 4090 訓練帶寬速度

共計用時 8：56 小時。本次訓練中因為主板上的 NVME 插槽會和 OCULINK 共享一路 PCIE4.0 x16 帶寬。所以 6 張中的其中兩張跑在了 pcie4.0 X4 上，從上圖就可以看出 RX 和 TX 都只是 PCIE4.0 X4 的帶寬速度。這也成為了本次訓練中最大的通訊瓶頸。如果全部的卡都跑在 pcie 4.0 x16 上，速度應該是比現在快不少的。

圖5 LLaMA-Efficient-Tuning 生成 loss 曲線

以上是 LLaMA-Efficient-Tuning 自動生成的 loss 曲線，可以看到 4 個 epoch 后收斂效果還是不錯的。

2）針對 codellama 的 qlora 微調

根據上文所述的 prompt loss mask 方法重構了 trainer 類（見項目代碼倉庫 func_caller_train.py）。因為數據集本身比較小（55 行）。所以跑 4 個 epoch 只用了兩分鐘，模型很快達到了收斂。

4.3微調完成后的測試效果

在項目代碼倉庫中，提供了一個簡短可用的 toolkit 示例。里面的函數包括：

1. 必應搜索
2. 維基搜索
3. bilibili 搜索
4. 獲取當前時間
5. 保存文件
6. ...

現在有一個 70B 和一個 34B 的模型，在實際使用中，用 6 張 4090 同時以 bf16 精度運行這兩個模型是不現實的。但是可以通過量化的方法壓縮模型大小，同時提升模型推理速度。這里采用高性能 LLM 推理庫 exllamav2 運用 flash_attention 特性來對模型進行量化并推理。在項目頁面中作者介紹了一種獨特的量化方式，本文不做贅述。按照其中的轉換機制可以將 70b 的模型按照 2.5-bit 量化為 22G 的大小，這樣一張顯卡就可以輕松加載。

1）測試方法

給定一段不在訓練集中的復雜任務描述，同時在 toolkit 中添加訓練集中不包含的函數和對應描述。如果 planner 可以完成對任務進行拆分，distributor 可以調用函數，solver 可以根據整個流程對結果進行總結。

2）測試結果

任務拆分：先使用 text-generation-webui 快速測試一下任務拆分模型的效果，如圖所示：

圖6 任務拆分測試結果

這里可以寫一個簡單的 restful_api 接口，方便在 agent 測試環境下的調用（見項目代碼 fllama_api.py）。

函數調用：在項目中已經寫好了一個簡單的 planner-distributor-worker-solver 的邏輯。接下來就讓測試一下這個任務。輸入一段指令：what movies did the director of 'Killers of the Flower Moon' direct？List one of them and search it in bilibili.

「搜索 bilibili 」這個函數是不包含在項目的函數調用訓練集中的。同時這部電影也是一部還沒有上映的新電影，不確定模型本身的訓練數據有沒有包含。可以看到模型很好地將輸入指令進行拆分：

• 從維基百科上搜索該電影的導演
• 根據 1 的結果，從 bing 上搜索電影 Goodfellas 的結果
• 在 bilibili 上搜索電影 Goodfellas

同時進行函數調用得到了以下結果：點擊結果是 Goodfellas，和該部電影的導演匹配得上。

總結

本項目以「輸入一段指令自動實現復雜任務拆分和函數調用」場景為例，設計了一套基本 agent 流程：toolkit-plan-distribute-worker-solver 來實現一個可以執行無法一步完成的初級復雜任務的 agent。通過基礎模型的選型和 lora 微調使得低算力條件下一樣可以完成大模型的微調和推理。并采用量化部署的方式，進一步降低推理的門檻。最后通過該 pipeline 實現了一個搜索電影導演其他作品的示例，實現了基礎的復雜任務完成。

局限性：本文只是基于搜索和基本操作的 toolkit 設計了函數調用和任務拆分。使用的工具集非常簡單，并沒有太多設計。針對容錯機制也沒有太多考慮。通過本項目，大家也可以繼續向前一步探索 RPA 領域上的應用，進一步完善 agent 流程，實現更高程度的智能自動化提升流程的可管理性。