Skip to content

优化 LMDeploy routed_experts 传输路径,避免 rollout worker 中不必要的数据中转 #1941

Description

@wyeth93

背景

当前 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 更接近按需消费的数据流。

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions