AIAgent如何实现?6张4090 魔改Llama2:一句指令拆分任务、调用函数

AI Agent 是时下热门的一个方向，在 OpenAI 应用研究主管 LilianWeng 写的万字长文中[1]，她提出 Agent = LLM+ 记忆 + 规划技能 + 工具使用。

简单来说，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 进行总结归纳。

为了实现上述流程，在「任务拆分」和「函数调用」模块中，项目分别设计了两个微调模型，来实现将复杂任务拆分并按需调用自定义函数的功能。归纳总结的模型 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。

本项目选定任务拆分模型型号为：

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）：

loss_mask 是一个与输入序列 input_ids 形状相同的张量（tensor）。每个元素都是 0 或 1，其中 1 表示对应的位置的标签应被考虑在损失计算中，而 0 表示不应被考虑。
例如，如果某些标签是填充的（通常是因为批处理中的序列长度不同），不想在损失的计算中考虑这些填充的标签。在这种情况下，loss_mask 为这些位置提供了一个 0，从而遮蔽掉了这些位置的损失。

损失计算：

首先，使用了 CrossEntropyLoss 来计算未 mask 的损失。设置 reduction=’none’ 来确保为序列中的每个位置都返回一个损失值，而不是一个总和或平均值。
然后，使用 loss_mask 来 mask 损失。通过将 loss_mask 与 losses 相乘，得到了 masked_loss。这样，loss_mask 中为 0 的位置在 masked_loss 中的损失值也为 0。

损失聚合：

将所有的 masked_loss 求和，并通过 loss_mask.sum() 来归一化。这确保了你只考虑了被 mask 为 1 的标签的损失。为了防止除以零的情况，加一个很小的数 1e-9。
如果 loss_mask 的所有值都是 0（即 loss_mask.sum() == 0），那么直接返回一个 0 的损失值。

2.2 硬件需求：

6*4090 for Marcoroni-70B’s 16bit lora
2*4090 for codellama 34b’s qlora / 1*4090 for codellama 13/7b’s qlora

2.3 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。示例如下：

微调过程说明

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 微调

LLaMA-Efficient-Tuning 框架支持 deepspeed 集成，在训练开始前输入 accelerate config 进行设置，根据提示选择 deepspeed zero stage 3，因为是 6 卡总计 144G 的 VRAM 做 lora 微调，offload optimizer states 可以选择 none，不卸载优化器状态到内存。
offload parameters 需要设置为 cpu，将参数量卸载到内存中，这样内存峰值占用最高可以到 240G 左右。gradient accumulation 需要和训练脚本保持一致，这里选择的是 4。gradient clipping 用来对误差梯度向量进行归一化，设置为 1 可以防止梯度爆炸。
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。之后训练就开始了。

共计用时 8：56 小时。本次训练中因为主板上的 NVME 插槽会和 OCULINK 共享一路 PCIE4.0 x16 带宽。所以 6 张中的其中两张跑在了 pcie4.0 X4 上，从上图就可以看出 RX 和 TX 都只是 PCIE4.0 X4 的带宽速度。这也成为了本次训练中最大的通讯瓶颈。如果全部的卡都跑在 pcie 4.0 x16 上，速度应该是比现在快不少的。

以上是 LLaMA-Efficient-Tuning 自动生成的 loss 曲线，可以看到 4 个 epoch 后收敛效果还是不错的。

2）针对 codellama 的 qlora 微调

根据上文所述的 prompt loss mask 方法重构了 trainer 类（见项目代码仓库 func_caller_train.py）。因为数据集本身比较小（55 行）。所以跑 4 个 epoch 只用了两分钟，模型很快达到了收敛。

4.3微调完成后的测试效果

在项目代码仓库中，提供了一个简短可用的 toolkit 示例。里面的函数包括：

必应搜索
维基搜索
bilibili 搜索
获取当前时间
保存文件
…

现在有一个 70B 和一个 34B 的模型，在实际使用中，用 6 张 4090 同时以 bf16 精度运行这两个模型是不现实的。但是可以通过量化的方法压缩模型大小，同时提升模型推理速度。这里采用高性能 LLM 推理库 exllamav2 运用 flash_attention 特性来对模型进行量化并推理。在项目页面中作者介绍了一种独特的量化方式，本文不做赘述。按照其中的转换机制可以将 70b 的模型按照 2.5-bit 量化为 22G 的大小，这样一张显卡就可以轻松加载。

1）测试方法

给定一段不在训练集中的复杂任务描述，同时在 toolkit 中添加训练集中不包含的函数和对应描述。如果 planner 可以完成对任务进行拆分，distributor 可以调用函数，solver 可以根据整个流程对结果进行总结。

2）测试结果

任务拆分：先使用 text-generation-webui 快速测试一下任务拆分模型的效果，如图所示：

这里可以写一个简单的 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 流程，实现更高程度的智能自动化提升流程的可管理性。

本文章转载微信公众号@机器之心SOTA模型