Prhub

#2598 fix(metrics): propagate and correctly aggregate speculative decoding counters

原始 PR 作者 nanjiangwill 合并时间 2026-09-12 06:10 文件变更 20 提交数 5 评论 15 代码增减 +572 / -55

执行摘要

统一 SGLang 投机解码计数并修复跨会话聚合

PR body 直接点明两个根因:"The old Miles field names no longer matched the SGLang response contract, and averaging per-request ratios gives each request equal weight regardless of its number of drafts or verification steps. The session samples codec also dropped SpecInfo before it reached rollout metrics." 即旧字段名(spec_accept_token_num 等)在 SGLang 新响应中已不存在,导致计数器读不到值;逐请求比率平均会稀释大 draft 数请求的贡献;session codec 在指标到达 rollout 前丢弃了 SpecInfo。session v2 还有一个额外边界:树会物化多片共享生成节点的样本,"Summing leaf-local SpecInfo would count shared nodes more than once",因此需要服务端发布一份权威聚合、驱动侧按 session id 去重。

值得精读。该 PR 展示了跨模块数据契约演进的标准姿势:类型层 fallback 兼容 → 传输层打通 → 权威聚合 → 驱动校验 → 聚合消费,每一步都有测试锚定。特别值得借鉴的是 session v2 的『服务端聚合一次 + 按 session_id 去重』模式,它解决的是树状采样结构下共享节点重复计数的通用问题,可推广到其他以树/图结构回传样本的指标场景。阅读建议聚焦 miles/utils/types.py 的 SpecInfo 与 miles/ray/rollout/metrics.py 的 _compute_spec_metrics 两处实现。

讨论亮点

核心交锋是 guapisolo 的 codex 审查评论(P2)指出的混合批次漏计问题,以及 guapisolo 对 session v1/v2 测试覆盖的补强。审查确认了 use_session_server 只是默认生成器选择而非样本来源,训练中 Sample.generate_function_path 可把样本覆盖到 single_turn.generate,这些样本有 spec_info 但无 session carrier,原有只聚合 carrier 的实现会让全普通批次的两个指标输出 0.0。nanjiangwill 用 ac647b0e0 提交修复聚合器,并补充混合批次与全普通批次测试。另外 guapisolo 主动补上了 session v1/v2 的 wire 传输测试,并在收尾时判定 tito 检查失败为既有 CI 失败,不阻塞合并。

实现拆解

  1. 统一数据契约:miles/utils/types.py 中 Sample.SpecInfo 字段重命名为 SGLang 规范命名(spec_num_correct_drafts / spec_num_proposed_drafts / spec_verify_ct / completion_tokens),add() 增加 spec_verify_ct <= 0 守卫(未经历投机验证的普通生成不计数),to_dict() / from_dict() 成对更新且 from_dict() 保留旧字段 fallback,历史序列化样本仍可读取。配套同步 miles/rollout/generate_utils/sample_utils.py 与 miles/utils/test_utils/mock_sglang_server.py 的字段读取点。

  2. 打通 session v1/v2 共享 wire:miles/rollout/session/samples/codec.py 在共享 samples wire 上序列化/反序列化 SpecInfo,此前 session 服务端已汇总计数但 codec 在编码时丢弃,导致驱动侧解出来全是 0。分别由 tests/fast/router/test_session_samples_op.py 与 test_session_samples_op_v2.py 锚定 v1/v2 的精确计数器往返。

  3. session v2 权威聚合一次:新增 miles/rollout/session/v2/metrics.py,build_session_rollout_metrics(session_id, nodes) 对所有已提交生成节点(session.tree.nodes)汇总一份 SpecInfo;SessionCoreV2.collect_samples 在投机解码配置开启时、hooks 处理完 metadata 之后的 wire 边界注入 SESSION_ROLLOUT_METRICS_KEY。测试覆盖共享前缀树与 max_seq_len 截断排除节点仍计入聚合。

  4. 驱动侧校验与替换:miles/rollout/generate_hub/agentic_tool_call.py 的 generate 在开启投机解码且 session v2 时:成功收集后校验 carrier 的 session_id 与 tracer.session_id 一致、metrics 非空,再以权威 carrier 覆盖每片叶子 metadata(清除旧叶子携带的陈旧 carrier);收集失败或空回复返回 ABORTED 时移除陈旧 carrier,避免悬空归属。

  5. 聚合层重写与配套:miles/ray/rollout/metrics.py 的 _compute_spec_metrics 重写为:无 carrier 的普通样本(如训练中经 generate_function_path 覆盖到 single_turn 生成器的样本)直接贡献本地 spec_info;有 carrier 的按 session_id 去重(同一 session 多片叶子只取一次);所有计数先求和再算 spec_accept_rate 与 spec_accept_length。配套更新 dashboard dump_reader 字段读取、mock SGLang server 输出,并在 11 个测试文件补齐用例(types、metrics、agentic、codec、dashboard、mock、fixtures 等),PR 验证为 224 passed / 1 skipped。

文件 模块 状态 重要度
miles/utils/types.py 类型层 modified 6.9
miles/ray/rollout/metrics.py 指标聚合 modified 6.64
miles/rollout/session/v2/metrics.py 会话聚合 added 7.06
miles/rollout/session/v2/core.py 会话核心 modified 5.91
miles/rollout/generate_hub/agentic_tool_call.py 生成调度 modified 6.32
miles/rollout/session/samples/codec.py 样本编解码 modified 5.7
tests/fast/router/test_session_samples_op_v2.py 会话测试 modified 7.38
tests/fast/rollout/generate_hub/test_agentic_v2.py 生成测试 modified 6.8
tests/fast/ray/rollout/test_metrics.py 指标测试 modified 6.45
tests/fast/utils/test_types.py 类型测试 modified 6.21

关键符号

Sample.SpecInfo.add Sample.SpecInfo.from_dict Sample.SpecInfo.to_dict Sample.SpecInfo.spec_accept_rate _compute_spec_metrics build_session_rollout_metrics agentic_tool_call.generate SessionCoreV2.collect_samples

关键源码片段

miles/utils/types.py core-logic

投机解码计数器的数据契约核心:SpecInfo 字段名对齐 SGLang 规范命名,add() 增加无验证不计数守卫,from_dict() 保留旧字段 fallback,是整个 PR 的根基。

@dataclass
class SpecInfo:
    # 字段名与 SGLang 响应契约对齐,取代旧的 spec_accept_token_num /
    # spec_draft_token_num / completion_token_num 命名(旧名字在新版
    # SGLang 响应中已不存在,导致此前指标读不到值)。
    spec_num_correct_drafts: int = 0
    spec_num_proposed_drafts: int = 0
    spec_verify_ct: int = 0
    completion_tokens: int = 0
​
    @property
    def spec_accept_rate(self) -> float:
        # 比率属性只负责单条 SpecInfo 的换算;跨样本聚合时由
        # _compute_spec_metrics 先求和再算比值,避免小样本被平均稀释。
        if self.spec_num_proposed_drafts == 0:
            return 0.0
        return self.spec_num_correct_drafts / self.spec_num_proposed_drafts
​
    @property
    def spec_accept_length(self) -> float:
        return self.completion_tokens / self.spec_verify_ct if self.spec_verify_ct > 0 else 0.0
​
    def add(self, meta_info: dict) -> None:
        # spec_verify_ct 缺失或为 0 说明该响应未经历投机验证;
        # 只携带 completion_tokens 的普通生成不应污染聚合计数。
        spec_verify_ct = meta_info.get("spec_verify_ct") or 0
        if spec_verify_ct <= 0:
            return
        self.spec_num_correct_drafts += meta_info.get("spec_num_correct_drafts", 0)
        self.spec_num_proposed_drafts += meta_info.get("spec_num_proposed_drafts", 0)
        self.spec_verify_ct += spec_verify_ct
        self.completion_tokens += meta_info.get("completion_tokens", 0)
​
    def to_dict(self) -> dict:
        # 持久化与 wire 传输统一使用 SGLang 规范命名。
        return {
            "spec_num_correct_drafts": self.spec_num_correct_drafts,
            "spec_num_proposed_drafts": self.spec_num_proposed_drafts,
            "spec_verify_ct": self.spec_verify_ct,
            "completion_tokens": self.completion_tokens,
        }
​
    @staticmethod
    def from_dict(data: dict) -> "SpecInfo":
        # 历史序列化样本使用旧 Miles 字段名,这里保留 fallback 映射,
        # 保证旧 dump 仍可被正确读取(新数据则直接命中规范键)。
        info = Sample.SpecInfo()
        info.spec_num_correct_drafts = data.get(
            "spec_num_correct_drafts", data.get("spec_accept_token_num", 0)
        )
        info.spec_num_proposed_drafts = data.get(
            "spec_num_proposed_drafts", data.get("spec_draft_token_num", 0)
        )
        info.spec_verify_ct = data.get("spec_verify_ct", 0)
        info.completion_tokens = data.get("completion_tokens", data.get("completion_token_num", 0))
        return info
miles/ray/rollout/metrics.py dependency-wiring

指标聚合核心:从逐样本平均改为 carrier 去重 + 全局计数,修复混合批次中普通生成样本被漏计的问题(review 后补齐)。

def _compute_spec_metrics(args, all_samples: list[Sample]):
    # 未配置投机解码时完全不产出 spec 相关指标。
    if args.sglang_speculative_algorithm is None:
        return {}
​
    carriers = {}
    spec_infos = []
    for sample in all_samples:
        carrier = sample.metadata.get(SESSION_ROLLOUT_METRICS_KEY)
        if carrier is None:
            # 普通生成样本(如训练中经 generate_function_path 覆盖到
            # single_turn 生成器)没有 session carrier,直接贡献本地计数,
            # 避免混合批次里这部分生成被漏掉。
            spec_infos.append(sample.spec_info)
        else:
            # session v2 的多片叶子共享同一个权威 carrier:
            # 按 session_id 去重,保证共享生成节点只计一次。
            carriers[carrier["session_id"]] = carrier
    spec_infos.extend(
        Sample.SpecInfo.from_dict(carrier["metrics"]["spec_info"])
        for carrier in carriers.values()
    )
​
    # 先对所有计数器求和、再计算比值,而不是平均每个样本的比率;
    # 这样每个 draft / 验证步骤的权重与其实际数量成正比。
    num_correct_drafts = sum(info.spec_num_correct_drafts for info in spec_infos)
    num_proposed_drafts = sum(info.spec_num_proposed_drafts for info in spec_infos)
    spec_verify_ct = sum(info.spec_verify_ct for info in spec_infos)
    completion_tokens = sum(info.completion_tokens for info in spec_infos)
    return {
        "spec_accept_rate": num_correct_drafts / num_proposed_drafts if num_proposed_drafts > 0 else 0.0,
        "spec_accept_length": completion_tokens / spec_verify_ct if spec_verify_ct > 0 else 0.0,
    }
miles/rollout/session/v2/metrics.py core-logic

新增的 session 侧权威聚合模块:定义 SESSION_ROLLOUT_METRICS_KEY 与 build_session_rollout_metrics,解决树共享节点重复计数问题。

# 该 key 是 session 服务端发布的权威 rollout 指标载体,
# 驱动侧(agentic_tool_call)与聚合侧(_compute_spec_metrics)共用。
SESSION_ROLLOUT_METRICS_KEY = "session_rollout_metrics"
​
​
def build_session_rollout_metrics(session_id: str, nodes: list[TrajectoryNode]) -> dict:
    # 在服务端对所有已提交生成节点汇总一次,而不是让每片返回的
    # 叶子各自累加:v2 树中多个样本会共享同一段生成节点,叶子局部
    # 求和会把共享节点重复计数。
    spec_info = Sample.SpecInfo()
    for node in nodes:
        spec_info.add(node.record.response["choices"][0]["meta_info"])
    return {
        "session_id": session_id,
        "metrics": {"spec_info": spec_info.to_dict()},
    }

评论区精华

混合批次中普通生成样本的 spec 计数被遗漏 正确性

guapisolo 的 codex 审查(P2)指出:全局 use_session_server=v2 下,训练样本仍可能通过 Sample.generate_function_path 覆盖为 single_turn.generate,这些样本有 spec_info 但无 session carrier;只聚合 carrier 会漏计,全普通批次时 spec_accept_rate 与 spec_accept_length 均输出 0.0(mock 复现期望 0.6 / 2.5,实际全为 0.0)。

结论:nanjiangwill 在 ac647b0e0 提交修复:_compute_spec_metrics 按 session_id 去重 carrier,同时纳入无 carrier 样本的 spec_info,并新增混合批次与全普通批次测试。 · 已解决

session v1/v2 的 SpecInfo 传输缺失测试覆盖 测试

guapisolo 评论 "session server v1 / v2 not covered. my agent is fixing. will push to this pr",随后发布 843cf4c 提交携带 Sample.spec_info 穿过 v1/v2 共享 wire,候选套件 171 passed / 10 skipped。

结论:codec 补齐 SpecInfo 序列化,v1/v2 分别新增 test_spec_info_crosses_samples_wire 精确计数器往返测试。 · 已解决

CI 失败是否为既有问题 other

多次触发 /run-ci-sglang 后,guapisolo 收尾评论 "failed at tito checks, existing ci failure",判定为既有 CI 失败。

结论:非本 PR 引入,合并未受阻。 · 已解决

风险与影响

  1. 数据契约变更:miles/utils/types.py 的 from_dict 有旧字段 fallback,历史 dump 可读;但新序列化数据字段名变化会影响所有读写路径,dashboard dump_reader 与 mock server 已同步,外部消费 dump 的脚本需留意。
  2. 新增硬校验:miles/rollout/generate_hub/agentic_tool_call.py 在 spec 开启且 session v2 成功返回时,carrier 缺失 / session id 不匹配 / metrics 为空都会抛异常(KeyError / ValueError)。这是有意的 fail-fast,但要求 session server 与 driver 同步部署,旧 server + 新 driver 组合会直接失败。
  3. 聚合语义变化:从逐样本平均改为全局计数后,spec_accept_rate 数值会系统性变化,依赖旧基线的监控与看板需要重新校准解读。
  4. 测试依赖 mock:数百个测试与复现均基于 mock SGLang 与手工构造的树节点,未覆盖真实 GPU 端到端;真实 SGLang 响应若字段缺失,SpecInfo.add 的 spec_verify_ct <= 0 守卫可防崩,但计数会偏低。

对用户与系统:投机解码指标从"恒 0 或失真"变为可置信的全局计数,直接改善训练与评估监控的可观测性;session v1/v2 行为一致,wire 上新增 session_rollout_metrics 权威载体。对团队:session server 与 driver 需同步升级,dashboard 展示字段随之更新;树状采样下"服务端聚合一次 + 按 session_id 去重"的语义具有通用借鉴价值。影响范围是 rollout 指标、session 服务、dashboard 与 mock 测试基建,属跨模块但边界清晰的指标正确性修复。

数据契约变更 跨模块联动 新增硬校验 测试依赖 mock

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论