背景
当前 XTuner 使用 LMDeploy 作为 rollout backend 时,LMDeploy 返回的 routed_experts 通常不是专家索引真实数据,而是一
个 Ray SharedStore 中的 key。
目前 XTuner 的处理链路如下:
LMDeploy routed_experts key
-> rollout worker 根据 key 从 LMDeploy SharedStore 取出真实 ndarray
-> rollout worker 再 ray.put(ndarray),得到 XTuner 自己的 ObjectRef
-> RolloutState.routed_experts 保存这个新的 ObjectRef
-> trainer worker 后续 ray.get(ObjectRef),取出真实 routed_experts 数据
相关代码位置:
- xtuner/xtuner/v1/rl/rollout/lmdeploy.py::_decode_routed_experts
- xtuner/xtuner/v1/rl/rollout/worker.py::_safe_handle_response
- xtuner/xtuner/v1/rl/trainer/worker.py::_add_rollout_routed_experts
- lmdeploy/lmdeploy/pytorch/engine/engine_instance.py::SharedStore
问题
当前链路会在 rollout worker 阶段提前 materialize routed_experts,然后再次 ray.put 到 Ray object store。训练阶段真正
消费数据时,trainer worker 又需要 ray.get。
对于 MoE 模型,routed_experts 可能比较大,形状通常类似:
(num_tokens - 1, num_layers, top_k)
因此当前链路可能引入一次不必要的数据中转和 object-store 写入,尤其在 rollout worker、LMDeploy SharedStore actor、
trainer worker 不在同一节点时,可能带来额外跨节点数据传输。
建议优化方向
方案 1:XTuner 延迟解析 LMDeploy routed_experts key
可以考虑不在 rollout worker 中立即根据 key 取真实数据,而是把 LMDeploy 返回的 key 保存到 RolloutState.routed_experts
中,等 trainer worker 真正需要 routed_experts 时,再根据 key 到 LMDeploy SharedStore 拉取真实数据。
优化后的链路:
LMDeploy routed_experts key
-> RolloutState.routed_experts 保存 key 或结构化 ref
-> trainer worker 根据 key 从 LMDeploy SharedStore 获取真实数据
建议不要直接保存裸字符串,可以定义一个结构化引用类型,例如:
@DataClass
class LMDeployRoutedExpertsRef:
key: str
actor_name: str = "shared_store"
namespace: str = "lmdeploy"
这样可以避免把普通字符串和 LMDeploy shared-store key 混淆。
方案 2:LMDeploy 新增接口返回真实数据的 ObjectRef
另一种方式是修改 LMDeploy SharedStore 的实现,新增接口根据 key 返回真实数据对应的 Ray ObjectRef,而不是直接
ray.get(ref) 返回 ndarray。
这样 XTuner 可以直接把该 ObjectRef 保存到 RolloutState.routed_experts,避免 rollout worker 侧先取 ndarray 再 ray.put
一次。
优化后的链路:
LMDeploy routed_experts key
-> LMDeploy SharedStore 返回对应 ObjectRef
-> RolloutState.routed_experts 保存该 ObjectRef
-> trainer worker ray.get(ObjectRef)
该方案可能需要 LMDeploy 侧明确 ObjectRef 生命周期和引用计数语义。
需要注意的问题
-
LMDeploy 当前 SharedStore.get(key) 是 destructive read,会 pop 掉 key,因此延迟解析时需要确保 key 只被消费一次,或
者 LMDeploy 提供 peek/free 等更明确的生命周期接口。
-
partial rollout 当前可能需要在 rollout worker 中拼接历史和当前 routed_experts,这类场景可能仍然需要 eager
resolve。
-
RolloutState.routed_experts 当前主要保存 ndarray / Ray ObjectRef,如果引入 LMDeploy key 或结构化 ref,需要同步更新
trainer 侧解析逻辑。
预期收益
- 减少 rollout worker 侧一次 routed_experts materialization。
- 减少一次 ray.put(ndarray)。
- 降低大规模 MoE rollout 场景下的 object store 写入和可能的跨节点数据传输。
- 让 routed_experts 更接近按需消费的数据流。
背景
当前 XTuner 使用 LMDeploy 作为 rollout backend 时,LMDeploy 返回的
routed_experts通常不是专家索引真实数据,而是一个 Ray
SharedStore中的 key。目前 XTuner 的处理链路如下:
相关代码位置:
问题
当前链路会在 rollout worker 阶段提前 materialize routed_experts,然后再次 ray.put 到 Ray object store。训练阶段真正
消费数据时,trainer worker 又需要 ray.get。
对于 MoE 模型,routed_experts 可能比较大,形状通常类似:
(num_tokens - 1, num_layers, top_k)
因此当前链路可能引入一次不必要的数据中转和 object-store 写入,尤其在 rollout worker、LMDeploy SharedStore actor、
trainer worker 不在同一节点时,可能带来额外跨节点数据传输。
建议优化方向
方案 1:XTuner 延迟解析 LMDeploy routed_experts key
可以考虑不在 rollout worker 中立即根据 key 取真实数据,而是把 LMDeploy 返回的 key 保存到 RolloutState.routed_experts
中,等 trainer worker 真正需要 routed_experts 时,再根据 key 到 LMDeploy SharedStore 拉取真实数据。
优化后的链路:
LMDeploy routed_experts key
-> RolloutState.routed_experts 保存 key 或结构化 ref
-> trainer worker 根据 key 从 LMDeploy SharedStore 获取真实数据
建议不要直接保存裸字符串,可以定义一个结构化引用类型,例如:
@DataClass
class LMDeployRoutedExpertsRef:
key: str
actor_name: str = "shared_store"
namespace: str = "lmdeploy"
这样可以避免把普通字符串和 LMDeploy shared-store key 混淆。
方案 2:LMDeploy 新增接口返回真实数据的 ObjectRef
另一种方式是修改 LMDeploy SharedStore 的实现,新增接口根据 key 返回真实数据对应的 Ray ObjectRef,而不是直接
ray.get(ref) 返回 ndarray。
这样 XTuner 可以直接把该 ObjectRef 保存到 RolloutState.routed_experts,避免 rollout worker 侧先取 ndarray 再 ray.put
一次。
优化后的链路:
LMDeploy routed_experts key
-> LMDeploy SharedStore 返回对应 ObjectRef
-> RolloutState.routed_experts 保存该 ObjectRef
-> trainer worker ray.get(ObjectRef)
该方案可能需要 LMDeploy 侧明确 ObjectRef 生命周期和引用计数语义。
需要注意的问题
LMDeploy 当前 SharedStore.get(key) 是 destructive read,会 pop 掉 key,因此延迟解析时需要确保 key 只被消费一次,或
者 LMDeploy 提供 peek/free 等更明确的生命周期接口。
partial rollout 当前可能需要在 rollout worker 中拼接历史和当前 routed_experts,这类场景可能仍然需要 eager
resolve。
RolloutState.routed_experts 当前主要保存 ndarray / Ray ObjectRef,如果引入 LMDeploy key 或结构化 ref,需要同步更新
trainer 侧解析逻辑。
预期收益