Prhub

#22360 [diffusion] fix: fix loading multiple ckpts with different precision for a same module

原始 PR 作者 mickqian 合并时间 2026-04-09 02:44 文件变更 4 提交数 9 评论 2 代码增减 +349 / -17

执行摘要

修复扩散模型加载多精度检查点时因重复权重文件导致的不一致性问题。

根据PR body,动机是'Sana (and some other models, including mova) consistency flakiness was caused by duplicate transformer safetensors variants being loaded together, which made final weights depend on loader/streamer ordering.'

建议工程师精读transformer_load_utils.py和weight_utils.py的变更,重点关注如何处理重复精度变体的设计决策,以及快速失败检查的实现细节,以学习确保加载确定性的最佳实践。

讨论亮点

review中只有一个评论,由gemini-code-assist[bot]提出,讨论了scheduler的warmup过程在多节点环境中的问题,但与当前PR的核心变更无关。评论指出使用临时目录在多节点部署中会失败,建议更健壮的方案。此评论可能指向一个潜在的设计问题,但未在当前PR的代码中体现。

实现拆解

实现分为两个关键部分:

  1. 在transformer_load_utils.py中添加_filter_duplicate_precision_variant_safetensors函数,使用正则表达式识别并过滤如foo.safetensors和foo.fp16.safetensors的重复精度变体文件,优先保留非精度后缀的规范文件。
  2. 在weight_utils.py中添加_raise_if_duplicate_safetensors_keys函数,在safetensors加载前检测跨文件的重复tensor键,并快速失败以避免不确定行为。此外,更新了consistency_threshold.json中的多个模型测试阈值,并新增test_transformer_quant.py单元测试来覆盖过滤逻辑。
文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py diffusion/loader modified 8.0
python/sglang/multimodal_gen/runtime/loader/weight_utils.py diffusion/loader modified 7.0
python/sglang/multimodal_gen/test/unit/test_transformer_quant.py test added 5.0
python/sglang/multimodal_gen/test/server/consistency_threshold.json test modified 4.0

关键符号

_filter_duplicate_precision_variant_safetensors _raise_if_duplicate_safetensors_keys

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

warmup 同步的多节点兼容性问题 设计

reviewer gemini-code-assist[bot] 提到 scheduler 的 warmup 过程在多节点环境中可能失败,因为使用临时目录路径在其他节点上不存在,建议更健壮的方案如固定路径或共享存储。

结论:未在当前 PR 中解决,但指出了潜在的设计改进点。 · 待处理

风险与影响

风险包括:过滤逻辑可能错误地过滤掉必要的精度变体文件,导致加载失败或精度不匹配;快速失败检查可能过于严格,误报重复键而中断合法加载场景。但PR新增了单元测试,提供了部分覆盖,需确保测试充分。

该变更直接影响扩散模型的权重加载过程,修复了因文件加载顺序导致的非确定性行为,提升了模型输出的一致性。对于用户而言,减少Sana等模型的性能波动;对于系统,增强扩散模块的稳定性和可靠性。

过滤逻辑复杂性 潜在误过滤

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论