ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

Agent Lightning 异步训练详解:Collocated Asynchronous Rollout 的开启、原理与调优

Agent Lightning 异步训练详解:Collocated Asynchronous Rollout 的开启、原理与调优 Agent Lightning 异步训练详解Collocated Asynchronous Rollout 的开启、原理与调优【免费下载链接】agent-lightningThe absolute trainer to light up AI agents.项目地址: https://gitcode.com/GitHub_Trending/ag/agent-lightning长时程 Agent 的 rollout 时长差异很大同步训练中一个慢 rollout 就会拖住整个更新步。Agent Lightning v1.0 提供的collocated 异步训练让 rollout 生成与模型更新共享同一 GPU 池未完成的任务组自动结转carry over到后续步骤。读完本文你将掌握如何用两行配置开启异步收集、理解“暂停/排空 Gateway”的底层机制、读懂 WB 上报的六项异步指标并配置 rollout correction 修正策略失配staleness。为什么需要异步训练在同步模式下一次参数更新必须等待该步所有 rollout 组全部完成更新步的耗时等于最慢的那个组。对于需要多轮工具调用、检索或沙箱执行的任务rollout 时长方差极大GPU 会大量空转等待。Agent Lightning 的解法是collocated共置异步推理与训练不再分置不同设备而是共享 GPU 池训练侧只要凑够一批完成的组就立即更新其余组继续跑同时把新的 prompt 组补进活跃队列。开启异步训练通过agentlightning.async_rollout.enabled打开异步收集并且必须显式设置async_train_batch_sizeagentlightning: async_rollout: enabled: true async_train_batch_size: 64async_train_batch_size是保持活跃active的 prompt 组数量它必须严格大于data.train_batch_size——后者是单次更新消费的已完成组数量。完整的配置示例data: train_batch_size: 32 agentlightning: async_rollout: enabled: true async_train_batch_size: 64这个约束在 Trainer 初始化时会被强制校验trainer.py 中若async_train_batch_size train_batch_size会直接抛出ValueError。仓库自带的默认配置见 config.yaml其中enabled: false、async_train_batch_size: null即默认走同步模式。一个实用的起步值$$B_{async} 2 B_{train}$$即async_train_batch_size取train_batch_size的两倍。仓库中的官方示例也是这么做的train_calc_agent.py 在--async模式下自动将async_train_batch_size设为train_batch_size * 2train_gsm8k_agent.py、train_sw_agent.py、train_smith_agent.py 等示例同样在配置中预置了该项。调优方向沿用官方文档的建议调大async_train_batch_size当 rollout 时长方差显著、且承载 Agent 的资源本地进程或 Kubernetes Job容量充足时更大的活跃窗口能更充分地掩盖慢任务。调小当活跃的 Agent 进程 / K8s Job 消耗过多 CPU 或内存资源时应收缩窗口以控制并发压力。工作机制活跃窗口、结转与 Gateway 排空官方文档将异步收集描述为六个步骤下面结合源码逐步印证。Trainer 维持至多async_train_batch_size个活跃 prompt 组。见 trainer.py 的_next_train_batch_dict_for_rollout每步先统计上一轮结转的组数n_carry_over再按n_new async_train_batch_size - n_carry_over从 dataloader 取新组从而把活跃窗口恒定在B_async。Controller 在本地进程或 Kubernetes Job 中启动 Agent 执行。rollout 通过AglAsyncRolloutManageragl_rollout_manager.py批量提交给 Controller由后者按local.agent_class或k8s.job_template_path配置的模板拉起。凑满data.train_batch_size个完成的组即用于更新不等所有活跃组。核心是AglAsyncRolloutManager.enqueue_and_wait_until_group_completedagl_rollout_manager.py它把结转 rollouts 与新 rollouts 合并为active_rollouts按data_id分组后轮询状态while len(completed_group_keys) target_finished_group_num循环直到完成组数达到train_batch_sizeTrainer 侧在 trainer.py 以target_finished_group_numself.config.data.train_batch_size传入。未完成的组保持活跃并结转到下一步。循环结束后所有未完成的组打包为new_carry_over_rollouts返回Trainer 将其存入self._carry_over_rolloutstrainer.py供下一轮补位。更新权重前Gateway 暂停新请求并等待在途请求排空。完成组收集后_rollout会调用_pause_and_drain_gatewaytrainer.py先POST /proxy/pause携带原因超时放宽到 300 秒拿到暂停瞬间的inflight计数然后每 0.25 秒轮询GET /proxy/state直到inflight 0并记录排空耗时。服务端实现在 routes/proxy.py 的/pause、/resume、/state三个管理端点状态由 proxy.py 中的ProxyPauseStatepaused、inflight、reason等字段带asyncio.Lock维护。共享 GPU 执行模型更新随后推理恢复进入下一个 rollout 阶段。下一次_rollout开始时Trainer 会先_resume_gatewaytrainer.py幂等可重试恢复转发。组完整性GRPO/RLOO 的统计前提每个 prompt 组必须保持完整例如actor_rollout_ref.rollout.n为 4 时该 prompt 的 4 个兄弟 rollout必须全部结束这个组才能被优化器消费。源码中这一点有硬断言——AglAsyncRolloutManager按data_id分组后检查assert len(group) self._train_rollout_nagl_rollout_manager.py并只在组内所有 rollout 都到达终态时才计入completed_group_keysL503-L516。这样做保证了组内兄弟样本来自同一 promptGRPO/RLOO 等依赖组内对比的优势估计不被破坏。Agent 侧要求使用带重试的客户端当 Agent 的请求恰好落在 Gateway 暂停窗口时forward_request 会立即返回一个可重试响应HTTP 429附Retry-After头默认 5 秒见ProxyPauseState.retry_after_seconds与X-Agl-Paused: true标记body 为{error: gateway paused, reason: ...}。因此官方建议 Agent 使用带重试的 OpenAI 或 HTTP 客户端暂停期间到达的请求收到 429 后自动退避推理恢复后即可继续整个 Agent 流程无需感知训练更新。监控WB 异步指标Trainer 每步向 WB 上报一组training/async/前缀的指标其计算逻辑见_compute_async_rollout_metrics加上暂停/排空阶段的两个指标trainer.py指标含义training/async/n_prev_carry_over_rollouts从上一步继承过来的 rollout 数training/async/n_completed_rollouts本步被消费完成的 rollout 数training/async/n_new_carry_over_rollouts结转给下一步的未完成 rollout 数training/async/new_carry_over_age_max_steps当前结转 rollout 中“最老”的已跨过的优化器步数training/async/proxy_inflight_at_pauseGateway 暂停时刻仍在执行的请求数training/async/proxy_drain_seconds等待在途请求全部结束所花的时间判读建议n_new_carry_over_rollouts长期逼近窗口上限、new_carry_over_age_max_steps持续增大说明慢任务积压可考虑调大async_train_batch_size或收紧 rollout 超时rollout_timeout_seconds默认 1800 秒见 config.yamlproxy_drain_seconds反映每次更新前排空在途请求的代价通常应保持在一个较小量级。此外 Trainer 还会上报timing/rollout_*系列排队/运行耗时聚合trainer.py可用于观察 K8s 批量拉起 Pod 造成的队列等待。处理 Stalenessrollout correction异步 rollouts 由较旧版本的模型生成可能在被用于训练前就已“过期”stale带来策略失配。对此官方方案是启用verl的rollout correction并推荐token 级重要性采样TIS、裁剪阈值设为2algorithm: rollout_correction: rollout_is: token rollout_is_threshold: 2该配置直接并入 verl 的 PPO trainer 配置与agentlightning.async_rollout相互独立前者解决“用旧策略采的样本如何纠正”后者解决“如何不停机地持续采样”。适用前提与限制小结适用 Agent Lightningv1.0当前仓库主干v1.0 之前的旧版架构v0.x 分支不适用本文配置项。必须是共置部署rollout 生成与模型更新共享同一 GPU 池暂停/排空机制依赖 API Gateway 的/pause与/state端点。开启后必须同时设置async_train_batch_size且严格大于train_batch_size否则 Trainer 初始化即报错。Agent 侧应使用可处理 429 Retry-After的重试客户端否则会在暂停窗口内中断。建议同时评估 staleness 处理rollout 越多、跨步结转越久越应启用rollout_correction。配套阅读Trainer Configurationverl 集成与 trace 聚合、API Gateway ConfigurationGateway 与模型代理配置、Controller Configuration本地与 K8s 执行器行为与约束的交叉引用见 4-trainer-configuration.md 末尾对本文档的回链。【免费下载链接】agent-lightningThe absolute trainer to light up AI agents.项目地址: https://gitcode.com/GitHub_Trending/ag/agent-lightning创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表