Prhub

#27692 [RL] Skip rotary cache tensors in weight checker

原始 PR 作者 zhuzilin 合并时间 2026-08-06 09:22 文件变更 1 提交数 3 评论 1 代码增减 +3 / -0

执行摘要

weight checker 跳过 cos/sin 旋转缓存比较

PR body 明确说明:Rotary cos/sin cache tensors are deterministic implementation details and may differ in shape across implementations。在 RL 权重检查流程中,期望模型与实际模型的 cos/sin 缓存因实现差异而 shape 不同,导致比较必然失败并产生误报,因此需要在比较前跳过这些张量。

该 PR 改动极小、逻辑直白,适合作为工具类修复的快速阅读示例,不必精读。值得关注的唯一设计决策是采用基于命名子串的跳过方式,而不是更精确的类型或来源判断,这可能在将来引入漏检风险;若后续继续演进 weight checker,建议考虑白名单或 metadata 标记机制。

讨论亮点

该 PR 没有产生实质性的 review 讨论。两位 reviewer(b8zhong、ch-wan)直接批准,无任何代码评审评论;唯一的评论来自 gemini-code-assist 机器人,内容为每日配额限制提示,与本次变更无关。

实现拆解

  1. 变更入口:python/sglang/srt/utils/weight_checker.py 中的 _check_tensors() 函数,该函数在 RL 权重检查流程中成对消费期望与实际的张量条目(名称、是否应比较、可比较张量)。
  2. 核心改动:在 zip 循环内部、名称相等断言之前,新增对 expect_name 是否包含子串 ".cos_sin_cache" 的判断;命中时直接 continue,跳过后续的断言与权重比较。
  3. 这样改的原因:cos/sin 缓存由模型配置中的 rope 参数确定性生成,不同 attention 实现可能采用不同 shape 存储,直接比较必然报错,而该张量并非真正的模型权重,跳过可消除误报。
  4. 对后续逻辑的影响:跳过该张量后,其名称一致性断言也不会执行;其余张量的比较、量化容差处理和错误收集逻辑均保持不变。
  5. 测试与配套:本 PR 未新增任何测试、配置或文档改动。
文件 模块 状态 重要度
python/sglang/srt/utils/weight_checker.py 权重检查 modified 4.59

关键符号

_check_tensors

关键源码片段

python/sglang/srt/utils/weight_checker.py core-logic

唯一变更文件,在权重比较循环中通过名称过滤跳过 .cos_sin_cache 张量,避免 RL 权重检查因旋转位置编码缓存形状差异产生误报。

def _check_tensors(
    expect_tensors: Iterable[CheckEntry],
    actual_tensors: Iterable[CheckEntry],
    allow_quant_error: bool = False,
):
    good_names = []
    error_messages = []
    info_messages = []
​
    for (expect_name, should_compare, expect_comparable), (
        actual_name,
        actual_should_compare,
        actual_comparable,
    ) in zip(expect_tensors, actual_tensors, strict=True):
        # 跳过 cos/sin 缓存:它由 shape 与 dtype 确定性生成,
        # 不同实现可能使用不同 shape,直接比较必然误报。
        # 注意:这里跳过的是整条条目的检查,包括后面的名称一致性断言。
        if ".cos_sin_cache" in expect_name:
            continue
​
        # 名称必须严格一致,否则说明期望与实际模型的结构不匹配
        assert expect_name == actual_name, f"{expect_name=} {actual_name=}"
        assert (
            should_compare == actual_should_compare
        ), f"{should_compare=} {actual_should_compare=}"
        name = expect_name
​
        try:
            equal, max_abs_err, mean_abs_err, num_exceed = compare_weights(
                expect_comparable, actual_comparable
            )
        except Exception as e:
            # 补充张量上下文信息,便于定位失败项
            e.add_note(
                f"when handling {name=} expect={expect_comparable!r} actual={actual_comparable!r}"
            )
            raise
        if equal:
            good_names.append(name)
            continue
​
        # 汇总差异信息:绝对误差、均值误差与超限数量
        msg = (
            f"name={name} "
            f"max_abs_err={max_abs_err} "
            f"mean_abs_err={mean_abs_err} "
            f"num_exceed={num_exceed} "
            f"expect={expect_comparable!r} actual={actual_comparable!r} "
        )
        if not should_compare:
            info_messages.append(msg)
        elif allow_quant_error and num_exceed == 0:
            info_messages.append(msg + "(within quantization ULP tolerance)")
        else:
            error_messages.append(msg)

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

  1. 漏检风险:跳过逻辑基于命名约定 ".cos_sin_cache",若未来某个模型将 cos/sin 缓存作为可训练参数或需要实际加载的权重,此改动会使其完全不被检查,可能掩盖真实错误。
  2. 名称断言被绕过:由于 continue 位于 expect_name == actual_name 断言之前,两边的张量名称即使不一致也不会被发现,可能隐藏命名问题。
  3. 影响范围:改动集中在 RL 权重检查工具,不涉及推理路径、调度或内核,整体回归风险低。

影响范围限于使用 weight_checker 的 RL 训练与验证流程:消除了因不同实现(如不同 attention 后端)产生不同 shape 的 cos/sin 缓存导致的误报,使权重检查更稳定。对推理性能和用户 API 无影响。对团队的收益是减少 RL 流程中的无效失败,但需注意跳过检查可能带来的覆盖盲区。

跳过特定张量比较可能漏检 无新增测试覆盖

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论