执行摘要
- 一句话:修复权重同步按 GPU UUID 选择载荷,支持 SP 多 GPU 引擎
- 推荐动作:值得精读。PR 展示了跨仓库数据契约的设计思路:用物理 GPU UUID 做显式配对,把“巧合正确”变成“必然正确”,并把不匹配升级为响亮失败。回退到世界秩索引的兼容策略也值得借鉴。建议后续为
_select_own_gpu_payload 补上针对标签匹配、长度校验、无标签回退和 UUID 归一化的单元测试,并确认 get_device_uuid 在各硬件平台可用。
功能与动机
并置的 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 是引擎侧的消费端配套。
实现拆解
- 数据契约扩展:在
python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py 的 UpdateWeightFromTensorReqInput 中新增可选字段 payload_gpu_uuids: list[str] | None = None,声明每个序列化载荷导出的物理 GPU UUID,旧请求不传时保持 None。
- 入口透传:
weights_api.py 的 update_weights_from_tensor HTTP 端点构造请求对象时增加 payload_gpu_uuids=body.get("payload_gpu_uuids"),把新字段从请求体带进调度器。
- 核心选择逻辑重写:
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。
- 调用点切换:
update_weights_from_tensor 改为调用 _select_own_gpu_payload,并将 payload_gpu_uuids 从请求对象传入。
- 测试配套:本 PR 未新增任何自动化测试(PR 中明确勾选
no tests added);验证依赖 8xH200 上的手工矩阵,覆盖 LoRA IPC、全量权重同步、双引擎 LTX 与旧生产者回退兼容。
关键文件:
python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py(模块 权重同步;类别 source;类型 core-logic;符号 _normalize_gpu_uuid, _select_own_gpu_payload, update_weights_from_tensor): 核心修复所在:把载荷选择从 TP 秩索引改为 GPU UUID 匹配,并保留无标签请求的世界秩回退,是本次变更的主逻辑。
python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py(模块 入口层;类别 source;类型 entrypoint;符号 update_weights_from_tensor): HTTP 入口需要把新的 payload_gpu_uuids 字段从请求体透传到请求对象,否则新字段无法到达 worker。
python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py(模块 数据契约;类别 source;类型 core-logic;符号 UpdateWeightFromTensorReqInput): 定义跨进程数据契约,新增的 payload_gpu_uuids 字段是所有后续逻辑的前提。
关键符号:_normalize_gpu_uuid, _select_own_gpu_payload, update_weights_from_tensor
关键源码片段
python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py
核心修复所在:把载荷选择从 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 评论;mickqian 两轮 APPROVED。实际设计讨论体现在 PR 正文:作者对比了 UUID 匹配与 TP 秩索引的取舍,说明回退到世界秩索引是为了兼容无标签的旧调用方,并将“找不到本机载荷”从静默错载改为显式报错。PR 也诚实指出性能收益不明确(同 GPU 0.51 ms vs 跨 GPU 2.72 ms 每 GiB fp32,但总时长被 all-gather、序列化和名称匹配主导),核心价值在正确性与故障可见性。
- 无技术性 review 讨论,两轮 Approve 与 CI 重跑 (other): 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_loader 按 tp_rank 切分。与生产者侧 radixark/miles_diffusion#81 联合构成多 GPU 权重同步闭环,后续 CFG 并行引擎支持需另行落地。整体影响程度中低,但修复了静默错载的潜在数据损坏风险。
- 风险标记:缺少测试覆盖, 回退逻辑变更, 依赖平台 UUID 接口, 显式失败可能中断训练
关联脉络
- PR #81 feat(rollout): send per-GPU weight-sync payloads to multi-GPU engines: 生产者侧配套改动:为每个 CUDA-IPC 载荷标注 payload_gpu_uuids,本 PR 是消费者侧实现;两者在 8xH200 上联调验证,PR 正文明确说明依赖关系。
参与讨论