hf_rollout.py — HuggingFace Rollout

文件概述

基于 HuggingFace Transformers 的简单推理实现，直接使用 model.generate() 进行文本生成。

核心类

HFRollout(BaseRollout)

class HFRollout(BaseRollout):
    def generate_sequences(self, prompts: DataProto) -> DataProto:
        """使用 HuggingFace 的 generate API 生成序列

        特点:
        - 最简单的推理方式，无需额外依赖
        - 性能较低（无 PagedAttention、无连续批处理）
        - 适合小规模实验和调试
        """
        output = model.generate(
            input_ids=prompt_ids,
            attention_mask=attention_mask,
            max_new_tokens=response_length,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
        )

适用场景

• 快速原型验证
• 小模型调试
• 不想安装 vLLM/SGLang 时的备选方案

与其他模块的关系

• 继承 base.py 的 BaseRollout
• 不需要额外的推理引擎依赖

小结

HFRollout 是最简单的推理实现，适合入门和调试，但生产环境推荐使用 vLLM 或 SGLang。