Skip to content

Conversation

@liyonghua0910
Copy link
Collaborator

需求描述

在 RL 场景需要交替进行训练和推理,因此每次推理结束后需要调用 /clear_load_weight 清理权重给训练腾出空间,训练结束后需要调用 /update_model_weight 时重新加载权重进行推理。现在 RL 需要开启上下文前缀缓存(Prefix Caching)实现推理加速,因此 KV 缓存也要配套该接口进行与模型权重的清除和加载。本 PR 主要实现该功能点。

主要改动

  • 新增了算子 unset_data_ipc 用于解除 set_data_ipc 对 kv cache tensor 的引用
  • 新增了 prefix_tree_status_signal 和 kv_cache_status_signal 信号,分别用于通知清除缓存树索引和缓存本身
  • 修改 kv cache 的初始化时机,现在有两种情况:(1) cache_transfer_manager 创建 cache 而 gpu_model_runner 连接,(2) gpu_model_runner 创建 cache 而 cache_transfer_manager 连接。只有非 profile 模式的 PD 分离部署为第 1 种情况
  • 扩展了 cache_ready_signal 信号的使用场景,原本用于 prefix_cache_manager 等待 cache_transfer_manager 进程创建 cache tensor,而现在 cache tensor 不一定由 cache_transfer_manager 创建,因此现在 cache_transfer_manager 和 gpu_model_runner 都有可能负责将该信号量置 1,或等待对方将该信号置 1
  • 新增了 swap_space_ready_signal 信号,用于 prefix_cache_manager 等待 cache_transfer_manager 完成 cpu 缓存的分配
  • 新增了部分 IPC 信号的状态常量的名称映射,以免开发时直接使用 0, 1, -1 等数字表示不同状态容易导致混淆
  • 给 engine_client 中的 clear/update 操作加了互斥锁,避免同时清除和加载权重
  • 新增了环境变量 FD_ENABLE_SWAP_SPACE_CLEARING 控制在清除 PrefixCache 时是否要顺带清除 CPU 缓存,默认为 0,即不会清除 CPU 缓存

@paddle-bot
Copy link

paddle-bot bot commented Sep 12, 2025

Thanks for your contribution!

self.enable_prefix_caching = False
if self.dynamic_load_weight:
self.enable_prefix_caching = False
# if self.dynamic_load_weight:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个PR合入后,再提一个PR注释代码移除

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这行会让 RL 场景的 prefix cache 强制关闭,但是在 develop 又已经被其他人移除了

@Jiang-Jia-Jun Jiang-Jia-Jun merged commit 7ccbcc5 into PaddlePaddle:release/2.2 Sep 16, 2025
23 of 27 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants