Prhub

#32685 [diffusion] fix: update weight from tensor detects device by uuid

原始 PR 作者 Rockdu 合并时间 2026-08-09 16:11 文件变更 3 提交数 5 评论 5 代码增减 +47 / -16

执行摘要

修复权重同步按 GPU UUID 选择载荷,支持 SP 多 GPU 引擎

并置的 RL 训练器在其引擎的 GPU 跨度内为每个训练秩导出一个 CUDA-IPC 载荷,载荷列表长度为 tp_size*sp_degree;旧代码按 TP 秩索引,只接受长度 1 或 tp_size 的载荷,导致 tp_size=1、sp_degree=2 的引擎每次同步都失败 serialized_named_tensors size must be 1 or tp_size (1), got 2。即便长度碰巧符合,TP 秩索引也只是“巧合地”同 GPU:跨度排序与生产者一致时才对上同驻秩,否则会静默导入其他 GPU 的缓冲。关联 Issue radixark/miles_diffusion#81 在生产者侧为每个载荷附带 payload_gpu_uuids,本 PR 是引擎侧的消费端配套。

值得精读。PR 展示了跨仓库数据契约的设计思路:用物理 GPU UUID 做显式配对,把“巧合正确”变成“必然正确”,并把不匹配升级为响亮失败。回退到世界秩索引的兼容策略也值得借鉴。建议后续为 _select_own_gpu_payload 补上针对标签匹配、长度校验、无标签回退和 UUID 归一化的单元测试,并确认 get_device_uuid 在各硬件平台可用。

讨论亮点

没有逐行 review 评论;mickqian 两轮 APPROVED。实际设计讨论体现在 PR 正文:作者对比了 UUID 匹配与 TP 秩索引的取舍,说明回退到世界秩索引是为了兼容无标签的旧调用方,并将“找不到本机载荷”从静默错载改为显式报错。PR 也诚实指出性能收益不明确(同 GPU 0.51 ms vs 跨 GPU 2.72 ms 每 GiB fp32,但总时长被 all-gather、序列化和名称匹配主导),核心价值在正确性与故障可见性。

实现拆解

  1. 数据契约扩展:在 python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.pyUpdateWeightFromTensorReqInput 中新增可选字段 payload_gpu_uuids: list[str] | None = None,声明每个序列化载荷导出的物理 GPU UUID,旧请求不传时保持 None
  2. 入口透传weights_api.pyupdate_weights_from_tensor HTTP 端点构造请求对象时增加 payload_gpu_uuids=body.get("payload_gpu_uuids"),把新字段从请求体带进调度器。
  3. 核心选择逻辑重写gpu_worker_post_training_mixin.py 移除 _select_rank_scoped_payload,新增 _normalize_gpu_uuid(去掉 NVML 的 GPU-/MIG- 前缀并转小写)与 _select_own_gpu_payload。带标签时先校验标签数与载荷数一致,再用 current_platform.get_device_uuid(self.local_rank) 得到本机 UUID 并匹配,找不到则返回显式错误;无标签时单载荷直接取第一个,多载荷回退到世界秩索引 world_group.rank_in_group
  4. 调用点切换update_weights_from_tensor 改为调用 _select_own_gpu_payload,并将 payload_gpu_uuids 从请求对象传入。
  5. 测试配套:本 PR 未新增任何自动化测试(PR 中明确勾选 no tests added);验证依赖 8xH200 上的手工矩阵,覆盖 LoRA IPC、全量权重同步、双引擎 LTX 与旧生产者回退兼容。
文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py 权重同步 modified 7.64
python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py 入口层 modified 4.56
python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py 数据契约 modified 4.49

关键符号

_normalize_gpu_uuid _select_own_gpu_payload update_weights_from_tensor

关键源码片段

python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py core-logic

核心修复所在:把载荷选择从 TP 秩索引改为 GPU UUID 匹配,并保留无标签请求的世界秩回退,是本次变更的主逻辑。

def _normalize_gpu_uuid(uuid: str) -> str:
    # NVML 导出的 UUID 带 GPU-/MIG- 前缀,torch 设备属性不带前缀,统一成小写无前缀格式再比较
    return uuid.removeprefix("MIG-").removeprefix("GPU-").lower()
​
​
class GPUWorkerPostTrainingMixin:
    def _select_own_gpu_payload(
        self,
        payloads: list,
        payload_gpu_uuids: list[str] | None,
    ) -> tuple[object | None, str | None]:
        # 基础校验:载荷列表必须是合法的非空列表
        if not isinstance(payloads, list):
            return None, "serialized_named_tensors must be a list"
        if not payloads:
            return None, "serialized_named_tensors is required"
​
        if payload_gpu_uuids is None:
            # 无标签回退:单载荷直接取第一个,保持旧调用方行为
            if len(payloads) == 1:
                return payloads[0], None
            # 多载荷按世界秩索引:TP-only 时世界秩 == TP 秩,等价于旧逻辑
            world_group = get_world_group()
            if len(payloads) != world_group.world_size:
                return None, (
                    f"serialized_named_tensors size must be 1 or world_size "
                    f"({world_group.world_size}), got {len(payloads)}"
                )
            return payloads[world_group.rank_in_group], None
​
        # 标签列表必须与载荷一一对应,否则拒绝更新
        if len(payload_gpu_uuids) != len(payloads):
            return None, (
                f"payload_gpu_uuids needs one entry per payload, "
                f"got {len(payload_gpu_uuids)} for {len(payloads)} payloads"
            )
​
        # 用本机 GPU UUID 精确匹配:找不到本机导出的载荷必须显式报错,而不是静默取别人的
        own_gpu_uuid = _normalize_gpu_uuid(
            current_platform.get_device_uuid(self.local_rank)
        )
        normalized_uuids = [_normalize_gpu_uuid(uuid) for uuid in payload_gpu_uuids]
        if own_gpu_uuid not in normalized_uuids:
            return None, (
                f"no payload was exported from this worker's GPU {own_gpu_uuid}, "
                f"got {payload_gpu_uuids}"
            )
        return payloads[normalized_uuids.index(own_gpu_uuid)], None

评论区精华

无技术性 review 讨论,两轮 Approve 与 CI 重跑 other

mickqian 两次标记 APPROVED,无逐行评论;Rockdu 与 mickqian 多次发送 /tag-and-rerun-ci;gemini-code-assist[bot] 发布了服务终止通知。设计取舍主要在 PR 正文中通过验证矩阵说明。

结论:mickqian 批准合并;PR Test 显示通过,PR Test (Extra) 显示失败,经重跑命令后关闭。 · 已解决

风险与影响

缺少单元测试:_select_own_gpu_payload 是纯逻辑函数,非常适合单测,但本次未覆盖,后续改动容易回归。回退语义变更:无标签多载荷路径由 TP 秩校验改为世界秩校验,当老调用方按 TP 数量发送载荷且 world_size 大于 tp_size 时,会从原有静默选择变为报错;TP-only 场景下世界秩与 TP 秩等价,兼容性成立。平台接口依赖:current_platform.get_device_uuid 在非 NVIDIA 平台(如 NPU)是否可用未验证,若未实现可能抛异常。显式失败的服务中断:找不到本机载荷会直接拒绝整个权重更新,比静默错误安全,但训练循环会中断,要求生产者侧标签配置完全正确。

影响面集中在 diffusion 后训练/RL 权重同步路径,不触及常规 serving 与推理。此前无法用 tp=1、sp=2 等多 GPU 配置跑通权重同步的引擎现在可以按 GPU UUID 精确取载荷,TP×SP 组合下 worker 拿到同驻 GPU 的完整张量后由 weight_loadertp_rank 切分。与生产者侧 radixark/miles_diffusion#81 联合构成多 GPU 权重同步闭环,后续 CFG 并行引擎支持需另行落地。整体影响程度中低,但修复了静默错载的潜在数据损坏风险。

缺少测试覆盖 回退逻辑变更 依赖平台 UUID 接口 显式失败可能中断训练

关联 Issue

#81 feat(rollout): send per-GPU weight-sync payloads to multi-GPU engines

完整报告

参与讨论