一、背景

蚂蚁集团的日常业务中，搜推广模型有着广泛的应用。而这些模型大多数是通过 Parameter Server 训练任务生成的。日常有大量推荐模型训练任务需要消耗极为可观的 CPU 资源。通常这些训练任务由用户配置后提交到集群运行。但是，相当部分提交的任务存在资源配置不当的问题，从而导致了难以忽视的问题：

• 训练任务资源配置不足，可能导致训练任务 OOM 错误或是训练性能低下。
  
• 训练任务资源配置超额，任务利用率低下，导致集群资源紧张，大量训练任务 pending。
  

蚂蚁 AI Infra 开发了 DLRover 分布式训练系统[1]，从工程上实现了对于 PS 训练任务的弹性规模调整，具体而言，DLRover 实现了 DLRover Brain Service，基于训练任务的运行状态，生成相应的资源配置优化方案，从而优化训练任务。

实现上，DLRover 将一个训练任务分解成多个阶段，然后针对每个阶段分别利用多种算法进行优化。

DLRover Brain 集成 LLM optimizer service 的架构设计

DLRover brain 将训练任务分成了 Create, Worker Initial, PS Initial 和 Running 四个阶段。每个阶段基于该阶段的特性，使用不同的优化方法(optimizer)来生成任务资源配置方案。其中，job running 阶段是作业的主要运行时间，该阶段的资源设置，决定了这个作业的资源利用率。在这个过程中，DLRover brain service 可以使用内置的规则算法进行资源动态调优，根据训练节点的统计情况，执行各种既定的调优规则。然而，实际场景中发现基于简单判断规则的调优策略的鲁棒性和泛化性能存在明显的局限性。同时既定的规则无法适应随机发生的混部集群的负载波动导致的性能波动。

针对上述问题，目前 DLRover 已经支持了基于离线强化学习(offline RL)的资源调优算法，用于优化 job running 阶段的资源调优，关于 offline RL 算法实现的细节，请阅读我们之前发布的文章：

告别资源浪费！DLRover Brain如何基于数据建模优化训练任务？

针对上述 job running 阶段的资源优化问题，本文进一步分析了现有优化方法的优劣势，结合大语言模型强大的世界知识和推理能力，构建了 DLRover LLM Agent，提出了一种全新的利用大语言模型(LLM)上下文学习的资源优化算法。

二、 问题定义

上述 DLRover 的 Job running optimizer 阶段的优化目标可以形式化表述为下面的优化问题。 设有一个分布式参数服务器(Parameter Server，PS)架构下的 TensorFlow 训练任务，该架构包括一组参数服务器 W 和一组工作节点 W 。每个参数服务器 pi∈P 和工作节点 wj∈W 均具备特定的资源配置，包括CPU核心数量和内存容量。

对于参数服务器集合 P ，其资源配置可以表示为 ，其中 、 和 分别表示第 i 个参数服务器的数量、CPU 核心数量和内存容量。对于工作节点集合 W ，其资源配置表示为 ( )，其中 n、c和 m分别代表工作节点的数量、单个节点的 CPU 核心数量和内存容量。

训练任务 T  被定义为一组固定条件约束 C ，这些条件约束描述了执行分布式训练任务所需的环境和要求，包括但不限于模型规模、数据集特性、训练批次大小等。因此，任务 T 可以表示为一个多维特征空间 T=Ck∣k∈K，其中K 表示条件约束的集合。

优化目标是在给定任务 T 和参数服务器资源配置 P 的情况下，找到最优的工作节点配置 (n^∗,c^∗,m^∗) ，以最大化 CPU 利用率与训练速度的加权和。形式上，我们寻求优化函数:

其中， 表示在给定工作节点配置和任务约束下的 CPU 利用率， 代表训练速度，而 是权重系数。

约束条件包括:

1.c≤Cmax ，每个工作节点的 CPU 核心数量不能超过集群设置上限。

2.m≤Mmax，每个工作节点的内存容量不能超过集群设置上限。

3. Umem(n,c,m;T,P)≤1，确保所有节点的内存利用率不超过100%导致OOM。

4. Vtrain(n,c,m;T,P )≥Vpre ，确保训练速度不低于历史作业阈值。

因此，优化问题可表示为寻找最优工作节点配置 (n^∗,c^∗,m^∗)  ，满足:

同时遵守上述约束条件。

三、Motivation

随着大语言模型（LLMs）在规模和基于注意力机制的架构方面的不断进步，它们展现了前所未有的通用性和处理各种任务的能力，特别是在 zero-shot 和 few-shot 的场景上。Few-shot in context learning（ICL）是大语言模型（LLM）在少量输入输出示例时执行新任务的能力。重要的是，ICL 中模型参数不必针对新任务进行微调即可解决新的任务。上下文学习（ICL）展现出来的一个重要能力是，这些模型并未直接被训练以优化元学习（meta-learning）目标，但却表现出了对多样化下游任务分布（至少是专门化的）的泛化能力[2]。

对于优化问题来说，许多优化方法都是迭代的，优化从初始解开始，然后迭代更新解以优化目标函数。

近期的一些研究已经表明，利用大语言模型 few-shot ICL 的泛化能力，可以让 LLM 执行高效的黑盒优化（black-box optimization）。使用大语言模型的 ICL 能力，不需要正式定义优化问题并使用编程求解器导出更新步骤，而是可以直接使用自然语言描述优化问题，然后指示 LLM 根据问题描述和先前找到的解决方案迭代地生成新的解决方案。LLM 的优化可以通过改变提示中的问题描述来快速适应不同的任务，并且可以通过添加指令来指定解决方案的期望属性来定制优化过程或者引导特定的优化方向[3,4]。

对于离线训练资源调优这一特定场景而言，整个混部集群的负载随着时间的变化表现出非周期性波动，使用过往集群任务数据训练的 offline RL 算法策略在领域内（in-domain）分布的训练任务的初始化上表现优秀，但是往往无法应对集群状态的非周期性变化以及集群负载突变的场景。受到 LLM 上下文学习能力的启发，我们可以通过使用一个强大的大语言模型，通过在提示中给出当前任务的负载和集群利用率的情况，让大语言模型推理出一个新的PS训练的工作节点配置（下文中，我们把这种工作节点配置称为 plan）。

基于上述思想，我们可以实现一个基于大语言模型的优化算法（下文中，我们把这种优化算法称为 LLM optimizer）。LLM optimizer 可以从集群和任务的实时统计指标中获取当前运行状态，进行实时的负载分析和总结，结合之前的历史策略表现推理出更优的 plan，在这些实时场景中，LLM optimizer 可以实现更加鲁棒的资源优化策略。

3.1 A Case of LLM Optimizer

为了更好的阐述基于大语言模型的优化范式，下面首先展示一个使用大语言模型 ICL 能力解决黑盒优化的一个例子。

假设存在一个线性回归函数 ，给出该函数的 50 组 ( ) 含噪声采样点。即从函数 中采样 50 组样本。

给定一个初始 ( ) 的估计值，根据在 50 组样本上的误差，迭代优化 ( )，使其尽可能接近真实值 ( ) 。

该问题可以直接使用上述 LLM 的 ICL 优化范式进行迭代搜索。即设计如下元提示(meta-prompt):

线性回归优化任务的 meta-prompt

其中橙色部分标识历史方案的误差估计，根据历史方案的误差“线索”，大语言模型可以直接进行合理的推断，迭代生成一个新的更加准确的 ( ) 估计。

三种不同的语言模型在线性回归问题上的搜索效率比较

上面的表格展示了不同 LLM 模型(text-bison、gpt-3.5-turbo 和 gpt-4)在优化线性回归问题中的表现。每一行表示一个 ( ) 组合，包括在真实值区域内、接近真实值区域但稍微偏离，以及远离真实值区域的情况。表格的两大列分别为每个模型搜索出正确解所需的平均步骤数，以及探索的可能的( )组合的数量及其标准差。

这个线性回归优化任务的实验结果表明，LLM 在优化过程中迭代探索的 ( ) 数量少于穷举搜索，表明这些大语言模型能够进行黑盒优化:LLM 可以基于历史优化轨迹，合理推测出下一步可能的下降方向，从而在复杂优化空间中高效地搜索。其中更强大的语言模型gpt-4 能够更快的逼近标准答案。综合来看，LLM 在无需显式定义解析表达式的情况下，通过自然语言提示即可进行黑盒优化。在小规模优化问题上，LLM 能够通过精心设计的提示找到高质量的解决方案，有时可以达到甚至超越手工设计的启发式算法。

四、DLRover LLM Agent

受到上述大语言模型迭代优化的启发，我们可以抽象出一个 LLM ICL 的优化范式(LLM optimizer)，并基于此范式，构建一个DLRover LLM Agent 用于DLRover PS 训练任务的资源调优。

4.1 LLM Optimizer 优化范式

根据上述优化例子所展示的流程，我们可以将这个过程抽象，设计一个通用的 LLM optimizer优化方案，使用大语言模型上下文学习能力的黑盒优化算法。

假设我们存在任务 ，给定策略 ，以及任务约束 ，优化目标函数为 。目标是找到最优的策略， 使得:

1. 首先初始化策略 p ，然后评估该策略下的评估指标 将这一策略观察放入历史策略池中(其被初始化为空)。

2. 迭代搜索策略:

 a.从历史策略池中，根据 的指标，召回 TopK 个历史策略以及对应的评估指标集合:

 b.使用一个任务特定的 meta-prompt 提示词模板 PT 结合召回的历史策略集合，使用 LLM 的 ICL 能力生成新的策略:

 c. 直至策略收敛或者 符合既定要求，退出迭代。

4.2 DLRover LLM Optimizer 方法介绍

遵循上述 LLM optimizer 优化范式，我们使用迭代搜索的方式，对 DLRover job running 阶段的 Worker 节点配置进行策略优化。

具体而言，整个受到 LLM 优化的训练任务从提交到完成训练被分为四个阶段:

 任务初始化:job master 拉起任务。

• Worker 初始化:初始化 Worker-0 节点作为最小的初始化设置。Worker-0的 CPU 和内存配置基于历史作业的统计值进行直接估计得到。
• PS 初始化:该阶段使用 DLRover Brain 默认的配置策略，也就是基于历史作业的统计估计来确定 PS 的数量和每个 PS 的规模配置。
• 运行时资源调优:使用上述 LLM optimizer 优化范式进行高效的策略搜索 ，从而最大化 。

DLRover 自动资源调优任务的四个阶段

利用上述的 LLM optimizer 优化范式，可以不需要任何历史数据，仅仅从 1 个 Worker 配置以及对应的观察评估指标出发，迭代搜索出新的配置(plan)，从而最优化综合 CPU 利用率的同时，维持训练速度没有明显衰减。

根据我们的结果，使用该大模型优化范式，在 Worker 资源配置优化任务上一个强大的大语言模型可以高效地进行策略搜索。在测试中，大部分任务在 10 次以内的迭代搜索后，就可以达到近似的最优策略配置点。在达到最优策略配置点后，针对该训练的所有后续资源设置会从之前的生成式策略探索任务，退化成一个最优策略检索任务，直接可以从历史策略池(下图中的 history plan database)检索出最优的配置。

LLM optimizer 工作流程图，随着历史任务的积累，优化任务从策略迭代生成任务退化成策略检索任务

4.3 任务冷启动

在实际应用场景中，如前文所述，offline RL 模型已经在 DLRover 系统中被证明能够提供超越规则性能调优策略表现。但是由于离线训练集群的混部特性，集群的利用率情况和负载情况往往是非周期性随机变化的，整个集群的负载往往不可预知。因此，在训练集群负载发生显著变化的时候，受到离群数据分布的影响，offline RL 模型如果不经过更新，往往无法很好地在新的负载场景上泛化，从而提供新环境下更加合理的策略规划。因此 offline RL 模型需要经常性地进行数据更新训练。同时，对于集群的一些突发利用率变化（例如任务负载突然增加），offline RL 由于训练数据的限制往往不会做出及时正确的响应。

另一方面，受到真实场景下策略探索的限制，过往任务的训练任务在不同资源配置下的表现无法被完全的探索和收集，因此 offline RL 实际可能会错过一些更加优秀的策略。

大语言模型的丰富的世界知识和强大的推理能力让运行时资源优化成为了可能。具体而言，LLM 可以从历史策略动作的数据库中抽取历史“记忆”，历史策略的表现可以让LLM 抽象出新的优化“经验”，同时由于历史策略表现在数据库中实时更新，因此“记忆”和“经验”可以天然地进行实时更新，从而更好地应对混部集群负载的不规则变化特性。

然而，直接使用LLM optimizer 也存在其固有的缺陷，特别是在任务冷启动（也就是当训练任务是之前从来没有出现过的新任务）的时候，LLM optimizer 往往无法有任何锚定参考，此时，直接使用 LLM optimizer 无法从历史策略中召回任何有效的观察点，从而导致潜在的远离最优策略的不合理初始化。这时候，offline RL 策略往往在最开始的任务冷启动上，可以提供更加精准优秀的初始化策略锚定。

因此，我们融合了 LLM optimizer 和 offline RL 算法的各自优势，最终构建了 DLRover LLM Agent。

4.4  DLRover LLM Agent Overview

在设计DLRover LLM Agent 时，我们借鉴了学术上主流的LLM Agent 架构理论[5]，尤其是关于记忆、反思和规划的组件设计，这使得整个架构具有更强的适应性和动态响应能力。DLRover LLM Agent 将 Offline RL 模型与 LLM optimizer 相结合，为 PS 训练任务的资源优化提供了一个灵活且自适应的框架。

记忆管理（Memory Management）

DLRover 的历史策略数据库（history plan database）类似于LLM Agent 架构中的长期记忆模块，存储了过去的策略及其性能指标。这一模块使得 Agent 能够在新任务中召回相关策略，以基于历史数据进行决策，从而实现记忆驱动的决策方式。该数据库的设计符合LLM Agent 架构中存储、合成和应用历史知识的关键原则。

反思与综合（Reflection and Synthesis）

在LLM Agent 中，反思是将记忆合成为更高级别的洞察，用以指导行为。在 DLRover LLM Agent 中，LLM Optimizer 使用历史策略数据来进行反思，即对过往经验进行泛化，从而生成针对新任务的优化计划。这种机制契合了LLM Agent 的核心功能，即通过分析过去的经验来推导出新洞察，以改进未来的行动。

计划生成（Plan Generation）

DLRover LLM Agent 的迭代优化框架模拟了LLM Agent 架构中的规划过程，即将高层次的目标转化为可执行的策略。在每次迭代中，DLRover 优化资源分配策略，旨在最大化任务的运行性能，同时适应训练环境中的动态负载条件。这种迭代方法灵感来源于 Agent 的规划框架，能够在保持连贯性的前提下，提高资源配置的适应性。

Offline RL 与冷启动的结合

在没有历史数据的新任务场景中，DLRover 引入 Offline RL 模型来提供初始计划，以确保即使在全新任务下，系统也能够有一个合理的起点。之后，LLM Optimizer 会基于实际表现进一步优化策略。Offline RL 与 LLM optimizer 的结合，为应对熟悉与陌生场景提供了强大的适应性，使整个系统更加健壮。综合以上设计，DLRover LLM Agent 形成了一个融合记忆、反思与规划的智能 Agent 架构。历史策略及其表现数据被存储在 history plan database 中，LLM Optimizer 对这些数据进行综合和优化，迭代生成符合当前任务需求的最佳资源配置。在新任务冷启动时，Offline RL 提供起始策略，随后 LLM Optimizer 进一步调整，使得 DLRover LLM Agent 能够在变化多端的训练负载下有效地进行资源优化。

这种混合架构有效结合了基于记忆的 LLM 上下文学习与离线强化学习的优势，形成了一个符合当前主流 Agent 设计范式的解决方案，最大化了资源利用率并能实时响应训练负载的变化。

随着 DLRover LLM Agent 不断运行，越来越多的作业的策略探索轨迹和对应的表现都会保存到 history plan database 中，对于高频训练任务，经过 LLM optimizer 的迭代探索，很容易将最优策略探索到，并保存在 history plan database 中，对于这类高频作业 DLRover LLM Agent 会直接总结历史 Plan 的表现，直接在任务资源初始化阶段给出最优的资源配置，并且能够在后续优化中实时响应集群状态的变化，生成新的策略来应对 out-of-domain 的场景。

DLRover LLM Agent 架构

4.5 元提示词（meta-prompt）设计

下面我们介绍最重要的针对该资源调优任务的 meta-prompt 设计。

我们的目标是使用在单个 prompt 中，完成新策略的生成，即一次性生成 Worker 的配置策略，其包括三个设置：Worker 数量、CPU 核数以及内存。

基于上述需求，我们的 prompt 分为五个部分。

• 任务说明
  
• few-shot 示例
  
• 历史经验回放
  
• 思维链（CoT）先验知识提示
  
• 策略输出约束
  

由于整体提示词超过了5000 个token，下面按照 prompt 的设计顺序依次介绍。

4.5.1 任务说明

任务说明提示词

任务说明部分阐述了问题的背景，以及在当前背景下的优化任务和目标。这里对于参数服务器架构没有展开进行解释，因为大语言模型本身就具有对于该架构的内在认知。

4.5.2 few-shot 示例

现有的研究表明，在进行大语言模型推理任务时候，few-shot 提示和思维链（CoT）都可以极大地提升大语言模型的推理能力[6]。在本任务的实验中，也发现合理且完善的 few-shot 提示可以极大地提高模型进行新策略推理的能力。因此，few-shot 的设计至关重要。

我们在这里使用了 3-shot 示例，每个示例的输入输出格式与最终推理任务的输入输出完全一致。即根据输入的一系列历史观察数据，进行分析归纳，从而推理出新的 Worker 设置策略（在该任务中也称为新的 Plan）。3-shot 示例中的输出经过精心的设计，目的是通过带有思维链的示例让语言模型能够按照合理的推理路径进行思考归纳。而在示例选择上，3-shot 中的每个例子也覆盖了三种常见的策略搜索情况，即：

• 历史策略中 Worker 数量不足，PS 利用率还未达到峰值，需要增加 Worker 数量；
  
• 历史策略逐步搜索到最优策略附近，需要在最优策略范围区间进行精细化策略搜索；
  
• 历史策略中 Speed 几乎不随着 Worker 数量的增加有明显的增加，这暗示集群可能存在通信瓶颈，不应再大幅增加 Worker 数量，需要在现有策略附近搜索收益边际点。
  

（1）示例一

few-shot 提示 example 1 模板（左侧是英文模板，右侧是中文模板），模板由三部分组成，分别是历史观察，分析，以及最后推理的新 plan

Example 1 展示了历史策略中 Worker 数量不足，PS 利用率还未达到峰值，需要增加 Worker 数量的一个决策例子。在 Example 1 的历史策略中，PS 节点的 CPU 利用率最高只有 30%，而且呈现随着 Worker 数量的增加近似线性增加的趋势，这表明 Worker 数量还远小于最优策略，因此需要大幅度增加 Worker 的数量。

few-shot 提示的示例一，其中蓝色是该示例中的过往历史观察，红色是该示例给出的新的 plan，该示例展示了历史策略中 Worker 数量不足，PS 利用率还未达到峰值，需要增加 Worker 数量

在这个例子中，首先输入部分历史策略通过一个固定的模板格式给出：

集群历史指标提示模板

从历史策略纪录（也就经验记忆）中提取的每一个Plan，都会被格式化为上述描述模板展示在提示词中。在这个模板中，首先会给出Plan 的设置，这是一个三元组，格式和最终策略的输出格式一致。然后会给出在这个 Plan（策略）下面的真实集群资源使用情况。使用情况的展示数据基于原始集群观察数据的删减和加工，只展示与优化目标相关的关键指标，同时尽可能地精简展示的数量。

目前的研究表明，大语言模型对于复杂的数字的理解能力要弱于量化后简单的数字。

在内存和CPU 指标的展示上，不仅展示了均值，还展示了最大和最小值，这有助于更好地展示集群的资源的分布区间，极值分布对于策略指定也具有重要作用。同时，在顺序排列上，我们把最终需要优化的两个效果指标，综合CPU 利用率和速度 Speed 放在了最后，这是因为目前的研究表明，大语言模型倾向于更关注开头和结尾的token。

在 Example 1 展示了三条历史策略纪录，同时他们根据对应 Plan（策略）的 Speed 顺序从小到大进行排列，这是由于当前一些研究工作表明，这样的展示顺序有助于大语言模型更好地发现和归纳历史策略中的关键因素，从而可以在归纳关键因素的基础上推理出更加合理的策略。

在该例子的输出部分，我们构造了一个详细的思维链分析，然后基于这个分析结论给出了新的策略（Plan），在给出 Plan 以后，最后又对给出的新的策略进行了一次 review 和 explaination，这样的显式二次 review 可以提高推理的合理性。在思维链分析上，我们对整个新策略生成的过程进行了步骤分解，将三个输出策略（Worker 数量，CPU 核数，内存数量）的制定按照人类专家的分析路径，依次进行思考。第一步会分析最重要的 Worker 数量的变化对于 PS 节点 CPU 利用率和训练速度的影响，这一步对整个历史策略中 Worker 数量的变化与速度和 PS 节点 CPU 利用率的相关性进行定性分析，进而归纳抽象出一般的结论。第二步对综合（Worker 和 PS 节点） CPU 利用率变动的关键影响因素进行定性分析和归纳，这种分析也对接下来产生新的策略（Plan）的 CPU 设置提供了经验指导。第三步是对内存分配的分析，基于当前 Worker 的内存利用率对内存情况进行判断。在这个部分引用了大量输入中的数字进行分析，目的是指导大语言模型能够尽量引用输入数据进行准确的数值分析（一般而言，大语言模型倾向于定性输出结论而忽略数值的定量分析）。

经过三步分析和归纳后，会给出新的策略（Plan），给出新的 Plan 后，会依次对策略中的三个值（Worker 数量、CPU 核心数、内存数量）分别进行 review 和 explaination，从而确定该策略是基于上述思维链分析归纳后的合理推理策略。

（2）示例二

Example 2 的提示设计与 Example 1 完全一致。与 Example 1 不同的是，Example 2 展示了历史策略逐步搜索到最优策略附近时，需要在最优策略范围区间进行精细化策略搜索的情况。

在 Example 2 中，Worker 数量在 20 到 30 之间的时候，Speed 和综合 CPU 利用率与 Worker 数量不再正相关，而是在该 Worker 数量范围内进行了浮动，同时此时 PS 节点的 CPU 利用率已经趋近峰值。这意味着 Worker 数量的设置已经趋近于最优策略，因此此时需要的决策是在该范围内进行进一步地精细化搜索来确定最优策略。

few-shot 提示的示例二，其中蓝色是该示例中的过往历史观察，红色是该示例给出的新的 plan，该示例展示历史策略逐步搜索到最优策略附近，需要在最优策略范围区间进行精细化策略搜索

（3）示例三

Example 3 的提示设计与 Example 1 和 2 完全一致。Example 3 展示了分布式训练速度的提升受到了集群通信瓶颈的影响（这是实际训练过程中最常见的一种性能制约）。历史策略的观察显示了随着 Worker 数量线性上升，训练速度（Speed）和 PS 的 CPU 利用率并没有显著的上升，这时候则暗示整个分布式训练受到了额外的通信瓶颈的影响。此时，不应该再大幅提高 Worker 数量去尝试继续提升 PS 节点的 CPU 利用率，而是应该小幅度变动策略，探索 Worker 数量潜在的收益边际。Example 3 的决策展示了这一探索策略。

few-shot 提示的示例三，其中蓝色是该示例中的过往历史观察，红色是该示例给出的新的 plan，该示例展示分布式训练受到了额外的通信瓶颈的影响，应该小幅度变动策略

4.5.3 历史经验回放

在给出了三个有代表决策例子后，接下来就是给出当前任务的历史观察输入，该步骤会从当前训练任务的所有有效的资源观察点中进行历史策略的召回。对所有观察点下相同策略的资源情况观察做时序平均（按照持续时间求加权均值），然后根据训练速度（Speed）召回 TopK 个训练速度最高的历史策略以及对应的资源观察数据。对召回的 TopK 个历史策略，使用与 few-shot 相同的输入格式。

当前任务历史提示模板

4.5.4 思维链（CoT）先验知识提示

思维链先验提示模板（这里省略了具体的提示词内容），此处需要加入最为关键的专家知识

在输入完成当前任务的历史策略后，还补充了一块内容来全面叙述策略需要优化的目标，以及最关键的指标之间的约束。上述叙述不仅显式地提供了一个思考路径模板，同时也注入了大量与实际业务集群相关的先验知识和约束（这些经验是基于集群真实训练数据观察所得到的经验），例如常见的通信瓶颈，以及 Worker 数量的增加存在边际收益点，这个过程中PS节点的 CPU 利用率是一个可观察的黄金指标。这些先验知识的注入，有助于提供更直接的策略搜索经验，从而帮助算法提高样本效率，进行更高效的策略搜索迭代。

在此处同样加入了收敛条件，也就是当策略的优化目标指标非常接近最优时，不应该继续探索新的策略而是直接返回现有的策略，表明迭代优化结束。

4.5.5 策略输出约束

策略输出约束提示模板

在 prompt 的末尾，会给出最重要的策略生成规范，这里包含了集群设置所需求的约束（例如 CPU 核心数量的约束和内存的约束），还有输出格式的规范。在输出格式规范中，与 few-shot 类似，使用了思维链提示（CoT）引导生成，让语言模型先对历史策略数据进行分析和归纳，根据归纳结果生成最后的策略。规范输出格式是为了让程序能够顺利通过模板匹配，从大语言模型的响应中抽取解析出最终的策略输出。

最终，大语言模型会生成与 few-shot 格式完全一致的分析和策略输出(n,c,m) 。

4.6 其他技巧

这里补充介绍在 prompt engineering 中用到的额外技巧，用于提升模型的推理准确性和稳定性。

4.6.1 多数性投票

在整个训练优化过程中，由于大语言模型生成的随机性，可能会出现一次生成的策略出现显著的逻辑偏离与推理错误（为了保证策略搜索的多样性，因此所有生成都设置为 temperature = 1），为了规避潜在的异常策略，提高策略生成的稳定性，这里使用了多次生成结果进行多数性投票的方法。

这种方法已经广泛地用于在 benchmark 中获取更好的测试效果，大量工作（例如 OpenAI o1-preview）已经表明扩展 inference time 可以显著提高模型的推理性能[9]。

具体而言，实验中会并发地重复执行 N 次 LLM 策略生成，对于生成的 N 个策略，采取多数性投票，将出现频次最多的策略作为本次策略迭代的最终结果。如果存在最高频次相同的两个策略，则取 Worker 数量最小的那个策略作为最终结果。由于成本限制，实验中N=3。

4.6.2 随机遗忘

在迭代搜索过程中，由于混部集群的影响，某些策略下的指标观察可能会出现极大的波动，导致某些策略的历史经验

存在较大偏差，这些异常的离群观察点如果频繁出现，会很大程度上影响策略生成的结果。另一方面，假设在 meta prompt 的历史经验中，出现大量的异常观察点，那么很容易引导 LLM 将策略收敛到一个次优设置上，并在次优点附近进行反复搜索而最终无法正确收敛到最优点。

在深度学习中，为了防止过拟合会经常使用 dropout 的方法。借鉴这一思想，在 LLM optimizer 的 ICL 决策中，我们会随机遗忘一些历史经验，将这些经验不加入到 meta-prompt 中，具体而言，对于召回的 TopK 个历史经验和其对应的指标观察，每个经验都有pf的概率被遗忘，每个策略 p 的遗忘概率pf  同时正相关于其所对应的观察指标 。

随机遗忘使得在策略迭代中不容易陷入某个次优点附近进行重复搜索，同时也一定程度上降低了偶发的混部集群指标异常对于策略生成的影响。

五、实验结果

下面展示 DLRover LLM Agent 在各个训练场景中的实验对比结果。

5.1 任务离线评估

首先我们对进行任务离线评估，离线评估指的是设置固定 PS 规模，对比不同算法对于 Worker 资源调节的能力。

我们使用不同能力的大语言模型（包括开源的 Qwen，闭源的 GPT-4）作为 LLM optimizer 的生成模型，在 5 个主流的搜推训练任务上评估了 LLM optimizer 的性能。

下面展示了在首页推荐任务上，PS 设置为 6 个节点，每节点 8 核 CPU 与 16GB 内存的条件下，LLM optimizer 与 DLRover 规则，以及 offline RL 算法三者的性能比较。为了公平比较，下面的所有对比都保持了 Worker-0 的配置一致，同时集群的设置也一致。

注意，LLM optimizer 没有任何的历史数据积累，也就是完全从头开始进行黑盒迭代优化的。

三种算法的训练调优动态曲线，横轴是对应训练任务的 global step（三个实验保持一致），三幅图分别展示了 LLM optimizer（左上），DLRover 规则（右上）和 offline RL（左下）的动态曲线

针对同样的任务，下图展示了更换语言模型，使用 GPT-4 作为 LLM optimizer 生成模型进行第二轮的效果测试的结果（其余设置与第一轮测试完全一致），从而展示该方法在语言模型能力进一步提升后的表现。

三种算法的训练第二轮动态曲线，横轴是对应训练任务的 global step（三个实验保持一致），三幅图分别展示了 LLM optimizer（左上），DLRover 规则（右上）和 offline RL（左下）的动态曲线

• 从对比中可以发现，LLM optimizer 相较于之前的算法产生了有竞争力的优化性能，在核心的利用率和速度指标上，都领先于 DLRover 默认规则算法。
  
• 同时 LLM optimizer 具有非常高效的策略搜索效率，即使是从无任何历史策略观察出发，仅仅通过小于 10 次的搜索，就能快速搜出一个具有竞争力的策略。同时不同语言模型的能力也会对优化效果产生影响，从结果来看，更强大的语言模型会取得更加优秀的优化效果。
  

5.2 任务线上评估

我们同样测试了算法在线上系统的效果，在线上环境中，PS 资源由资源预测的树模型（LightGBM）算法接管，其不再固定，也会动态发生变化。在这一线上设置下，评估 LLM 算法的效果。同样使用上述首页推荐训练任务。

首页推荐任务 LLM optimizer 动态调优效果展示，图表中间段数据缺失是由于模型在进行并行评估和导出，此时训练中断， DLRover 在此期间不会对其规模进行调整

上图展示了 LLM 在线上任务（PS 数量变化）时候的快速策略搜索能力，在前期 PS 资源较少（蓝色框）时快速增加 PS 趋近最优峰值利用率和速度，随着 PS 的增加，在后期（红色框）根据前期的调整经验实时学习，同样快速增加 Worker 数量，在新的 PS 数量下快速调整到 PS 峰值利用率，在不同 PS 规模下都实现了快速探索策略和经验的实时回放学习，避免了资源的浪费。

5.3 对于混部集群的适应性

下面一个训练任务展示 LLM optimizer 对于混部集群的适应性。

我们在同一个硬件集群重复运行了两次 DeepCTR 训练任务，两次运行受到混部集群的影响，导致集群总体负载发生了明显的变化，因此速度和利用差距显著。

LLM optimizer 集群发生负载变化时的快速策略调整

可以看到第二次与第一次训练的训练速度规律和利用率规律完全不一致，在这种集群情况变化导致训练动态大幅变化的情况下，LLM 算法展现了强大的适应性，在无需先验数据的情况下实时调整策略，在各种情况都通过快速调优 Worker 配置，维持了较高的 PS 利用率和合理的速度。

六、总结与讨论

本文介绍了DLRover LLM Agent，展示了基于 LLM 上下文学习能力的优化算法设计理念以及在 DLRover 资源调优上的应用方法和效果。

在 DLRover 资源调优问题上，与强化学习方法所不同的是，LLM optimizer 并没有把这一调优过程当成一个 MDP 的序列决策任务，而是直接将其抽象成了一个更加简单的优化问题。

这里实际上强化学习的做法更加符合实际情况，因为目前 DLRover 的资源调优，新的策略并不能完全生效，例如新的策略希望增加 Worker 数量，DLRover 只会新建一些节点使得总 Worker 数量符合新策略的预期，然后其并不会修改已有的 Worker 节点的配置。

因此对于每种策略的实际应用，其效果并不完全只取决于新策略的配置，还受到原有状态的影响。因此考虑过往状态的序列决策任务确实是更加严谨的建模。

不过好在随着迭代次数的降低，在 LLM optimizer 中，这一简化的影响会随着探索的深入逐步减小。

在该问题上使用强化学习最大的障碍就是样本采集的效率过于低下，受制于任务类型的原因，在 DLRover 中使用 RL 在线上环境并行采样来获取海量的数据是不现实的。而 RL 本身就是一个样本效率（data efficient）很低的算法，其依赖广泛且全面地对于状态空间的充分探索，然后依靠贝尔曼最优方程，去迭代优化得到一个更准确的值函数。

在最常见的 online RL 中，往往不会在策略探索过程中过度地引入一些先验信息从而加速探索的过程。一方面 RL 中的随机探索采用的动作空间上的概率随机，要在这一过程中引入先验信息一般采用一些间接的方法，例如在 reward 建模中引入先验从而惩罚一些明显不合理的探索状态，或者直接对于特定状态下面的策略动作进行筛选和过滤，而不会在动作生成时候直接嵌入一些先验经验约束。另一方面，引入过多的先验可能会使得探索出现预料之外的情况，特别是将知识先验引入 reward 函数的建模。因此，为了样本的收集效率，DLRover 目前使用的是 offline RL 算法。

而回到 LLM optimizer，LLM optimizer 明显具有极高的样本效率，这得益于一方面其不依赖随机探索，其每一步都是基于大量历史经验归纳以及考虑先验经验的策略生成。通过 ICL 这一途径，我们可以直接在上下文中嵌入任何自然语言可表述的先验信息和条件约束，从而使得策略搜索空间大大简化了，这是 LLM optimizer 样本效率高的本质。

对比 RL 的方案，实际上 LLM optimizer 在不更新任何参数的条件下，实现了策略的迭代。这是由于 RL 依赖历史经验数据的梯度，根据历史数据的梯度来优化策略网络的参数，而 LLM optimizer 直接靠 prompt 中的历史经验回放数据的更新，直接改变生成的策略，达到无参数更新更直接的策略迭代。当然，使用 LLM optimizer 取得上述这些显著优势实际上对于优化的任务有很高的条件约束。

• 与优化问题相关联的状态可以直接使用自然语言进行全面完善相对简短的表征（也就是状态空间相对简单）。
• 输出的动作空间是离散的，同时动作空间的规模是有限的。
• 人类的先验知识对策略探索至关重要，整个动作空间有很多无需全面探索，先验知识能够极大地提高策略搜索的效率。
  

只有在满足上述三个条件下，使用 LLM optimizer 进行优化才能发挥出合理的性能。

除了样本效率和先验知识注入这两个优势之外，LLM optimizer 还有两个显著优势：

• 对于某些 bad case，可以通过 few-shot 直接注入修正案例，这使得对一个优化算法进行某个场景和 case的策略修正显得非常简单和低成本。
  
• LLM optimizer 虽然是黑盒优化，但是其每个策略输出由于 CoT 提示的存在，都具有一定的可解释性，这种可解释性也可以指导我们对某些异常优化的性能表现直接进行针对性的优化，而不是基于概率和梯度下降方法的深度学习方法那样的完全黑盒算法，不具有可解释性。
  
• 策略探索生成到策略检索。得益于高效的样本效率，其对于新任务的适应能力远比基于离线数据训练的推理模型要好，其在没有见过的新的训练任务上可以体现强大的适应性和普遍的泛化性。
  
• 整个策略的学习和迭代是实时的，所有的策略生成能力都是随着训练过程中历史策略的积累而实时改进，相较于需要定期离线训练更新参数的方法来说，其具有显著的实时更新优势。
  

当然这一方法也有几个显著的缺点：

• 非常依赖大语言模型的基础能力。由于任务的复杂性，该方法非常需要强大的语言模型，能够适配该任务的模型目前非常稀少。
  
• 成本问题。由于大语言模型的规模、meta-prompt 的复杂性以及长文本，单次决策的成本其实是无法忽视的（当然随着整个算法和硬件的进步，这个问题会逐步缓解）。同时由于延迟问题，也会使得 DLRover brain 链路出现潜在的异常超时可能，在工程实现上，DLRover brain 因此引入了许多容错设计来保证 DLRover LLM Agent 的顺利运行。
• 局部最优与非稳定生成。虽然在其他技巧章节介绍了多数性投票和随机遗忘来优化整个策略探索效果，但是这个问题依旧无法彻底规避，同时面对非稳定生成所采用的多数性投票规则无疑也是的整个系统的运行成本翻了数倍。
  

七、未来改进方向

7.1 多目标策略生成分解

在现有的决策方案中，使用的是一个 meta-prompt一次性生成三个策略目标(n,c,m)，这导致了目前的 meta-prompt 的长度过长，单次推理时间在 20-40s 之间。实际上，这一过程可以进行分解优化。

生成目标(n,c,m)中，最重要的是 Worker 数量n，一旦生成了n 可以根据 n的变化进一步推理出每个节点的 CPU 核心数量c，而m由于并不在主要优化目标中，因此可以直接通过统计规则进行贪心设置。

分解后不仅降低了策略推导的难度，同时可以极大地减小 meta-prompt 的输入 token 长度，降低成本的同时也提高了推理速度。

7.2 基于检索的 few-shot 动态生成

在目前 meta-prompt 的设计中，我们使用了固定的 3-shot 案例设计，其覆盖了三种主流的决策情况。但是实际的历史情景会比当前的设置更复杂。当前对于 few-shot 的研究已经表明，影响 few-shot 的生成效果受到以下因素的影响[7,8]：

• Number of Demonstrations
  
• Demonstration Formatting
• Order of Demonstrations
• Diversity of Demonstrations
• Chain of Thought (CoT)
  

在目前的 3-shot 设计中，我们的 meta-prompt 在 Diversity of Demonstrations 上依旧存在提升的空间。在策略迭代过程中，与当前观察更加接近的演示（demonstrations）对于当前的新策略推理具有明显的帮助，而无关的演示对策略迭代往往没有帮助。因此如果可以维护一个演示集合，通过相似性检索出与当前输入分布最类似的决策演示，动态生成 few-shot 示例，会相较于目前静态的 few-shot 更加有帮助。

7.3 离群异常观测指标识别

由于受到混部集群的影响，同一个策略的观测指标可能会在不同时刻出现明显的异常波动，受到影响的异常观测指标如果直接进入历史经验回放中，会有较大的概率“误导”LLM 进行错误的新策略推断和经验归纳。因此在历史经验回放过程中，可以尝试引入基于统计检验的离群指标，剔除一些受到混部影响的观测数据。

7.4 跨任务的经验迁移

目前，针对该优化任务，我们采用的是根据用户训练任务的模型特征，判定是否是同一优化任务，也就是一旦用户的训练模型出现任何变化，我们都不会跨任务复用任何经验，新的模型的训练任务都会从头进行策略搜索。这无疑增加了新任务的搜索成本。实际上，很多任务的策略迭代具有高度的相似性质，因此结合前文的动态 few-shot 生成，可以尝试开发跨任务的经验迁移和回放，可以极大地提高在新任务上的搜索效率，降低搜索成本。

参考文献

[1] Wang Q, Lan T, Tang Y, et al. DLRover-RM: Resource Optimization for Deep Recommendation Models Training in the Cloud[J]. Proc. VLDB Endow., 2024.

[2] Yang C, Wang X, Lu Y, et al. Large Language Models as Optimizers[J]. arXiv preprint arXiv:2309.03409, 2023.

[3] Agarwal R, Singh A, Zhang L M, et al. Many-shot in-context learning[J]. arXiv preprint arXiv:2404.11018, 2024.

[4] Garg S, Tsipras D, Liang P S, et al. What can transformers learn in-context? a case study of simple function classes[J]. Advances in Neural Information Processing Systems, 2022, 35: 30583-30598.

[5] Park J S, O'Brien J, Cai C J, et al. Generative agents: Interactive simulacra of human behavior[C]//Proceedings of the 36th annual acm symposium on user interface software and technology. 2023: 1-22.

[6] Wei J, Wang X, Schuurmans D, et al. Chain-of-thought prompting elicits reasoning in large language models[J]. Advances in neural information processing systems, 2022, 35: 24824-24837.

[7] Luo M, Xu X, Liu Y, et al. In-context learning with retrieved demonstrations for language models: A survey[J]. arXiv preprint arXiv:2401.11624, 2024.

[8] Min S, Lyu X, Holtzman A, et al. Rethinking the role of demonstrations: What makes in-context learning work?[J]. arXiv preprint arXiv:2202.12837, 2022.

[9] Snell C, Lee J, Xu K, et al. Scaling llm test-time compute optimally can be more effective than scaling model parameters[J]. arXiv preprint arXiv:2408.03314, 2024.