加速框架学习

src/rocket/runners.py代码解读

这份文档用于解释 Rocket 框架里最核心的运行器文件：

Rocket/src/rocket/runners.py

这个文件的职责不是读取数据集、计算指标或写结果文件，而是把同一批 prompt 交给不同推理后端执行，并把 HF、vLLM、SGLang 的输出统一成同一种格式，方便上层 evaluator 做公平对比。

1. 文件定位

runners.py是 Rocket 的“推理后端适配层”。它向上对接evaluator.py，向下调用具体推理引擎：

evaluator.py -> build_backend(config) -> HfBackend -> VllmTextBackend -> SglangTextBackend -> backend.load() -> backend.generate(...) -> backend.runtime_metadata() -> backend.shutdown()

因此，runners.py解决的是一个工程问题：不同推理框架的参数、加载方式、输出格式都不一样，但实验表格需要同一套指标。这个文件把这些差异包起来，让上层只关心“输入 prompt、得到文本和 token 数”。

2. 核心数据结构

2.1Generation

@dataclass(slots=True)classGeneration:text:strgenerated_tokens:int

这是所有后端统一返回的单条生成结果。

  • text：模型生成的文本。
  • generated_tokens：本次生成的 token 数，用于计算吞吐。

HF、vLLM、SGLang 原生返回结构都不同：

  • HF 返回output_idstensor。
  • vLLM 返回 request output 对象。
  • SGLang 返回 dict/list/string 等不同版本格式。

Rocket 最终都转换成Generation，这样 evaluator 不需要理解每个框架的内部结构。

2.2BackendConfig

@dataclass(slots=True)classBackendConfig:model_path:strmethod:MethodSpec batch_size:intmax_model_len:intmax_new_tokens:inttemperature:floattop_p:floatseed:intdtype:strmax_num_batched_tokens:intgpu_memory_utilization:floatsglang_mem_fraction_static:floatinput_mode:str="auto"torch_compile_cache_size_limit:int=16torch_compile_mode:str="reduce-overhead"

这是所有后端共享的运行配置。

这里的设计重点是“变量统一”：同一个实验方法用同一批基础参数，比如模型路径、最大生成长度、batch size、采样温度、top-p 等。后端内部再把这些参数翻译成自己框架能理解的形式。

几个关键字段：

  • method：来自methods.py的MethodSpec，决定使用 HF/vLLM/SGLang、是否开 Rocket bucket、是否开 n-gram。
  • max_model_len：模型总上下文长度上限。
  • max_new_tokens：最大生成 token 数。
  • max_num_batched_tokens：vLLM/SGLang 的 chunked prefill 相关参数。
  • gpu_memory_utilization：vLLM 预留显存比例。
  • sglang_mem_fraction_static：SGLang 静态显存比例。
  • input_mode：auto/text/vision，用于 HF 路径区分纯文本模型和视觉语言模型。
  • torch_compile_cache_size_limit：HF + Rocket 使用torch.compile时的图缓存上限。

3. 后端工厂函数

defbuild_backend(config:BackendConfig):ifconfig.method.engine=="hf":returnHfBackend(config)ifconfig.method.engine=="vllm":returnVllmTextBackend(config)ifconfig.method.engine=="sglang":returnSglangTextBackend(config)

这个函数根据method.engine创建具体后端。

例如：

hf_sdpa -> HfBackend hf_rocket -> HfBackend vllm_fa -> VllmTextBackend vllm_rocket_flash -> VllmTextBackend sglang_flashinfer -> SglangTextBackend

注意：Rocket 不做“自动选择最快后端”的路由。这里的每个 method 都是固定后端、固定配置，用于公平实验对比。

4. HF 运行时兼容补丁

函数：

_patch_hf_runtime_for_installed_packages(attn_implementation)

这个函数只服务 HF 后端，目的是让当前环境里的 Transformers 能正常加载模型。

它不安装依赖，不卸载依赖，也不修改外部包文件，只在当前 Python 进程里做运行时兼容。

4.1 为什么需要 patchtorchao

当前环境里transformers 5.8.1可能会尝试导入 torchao quantizer。问题是安装的 torchao 版本依赖更高版本 PyTorch 里的：

torch.utils._pytree.register_consta