跳转至

rollout.py — Rollout 配置

文件概述

定义推理引擎(Rollout)的配置,是配置项最多的配置文件之一,涵盖采样参数、服务器配置、检查点引擎等。

核心配置类

RolloutConfig

@dataclass
class RolloutConfig:
    # 基本推理参数
    temperature: float = 1.0              # 采样温度
    top_p: float = 1.0                    # nucleus 采样
    top_k: int = -1                       # top-k 采样
    response_length: int = 512            # 最大生成长度
    prompt_length: int = 512              # prompt 最大长度

    # 并行配置
    tensor_model_parallel_size: int = 1   # 推理 TP 大小
    data_parallel_size: int = 1           # 推理 DP 大小
    expert_parallel_size: int = 1         # 专家并行大小

    # 推理引擎配置
    gpu_memory_utilization: float = 0.5   # GPU 显存利用率
    max_num_seqs: int = 128               # 最大并发序列数
    max_model_len: int = None             # 最大序列长度
    enable_prefix_caching: bool = False   # 前缀缓存
    enable_chunked_prefill: bool = True   # 分块预填充
    enable_sleep_mode: bool = True        # 休眠模式(Hybrid Engine)
    free_cache_engine: bool = True        # 是否释放缓存引擎

    # 引擎类型
    load_format: str = "dummy"            # dummy=不加载权重, auto=自动加载
    dtype: str = "bfloat16"               # 推理精度

SamplingConfig

采样参数的独立配置。

ServerConfig

HTTP 服务器配置(用于异步推理模式):

@dataclass
class ServerConfig:
    timeout: float = 60.0          # 请求超时时间
    max_attempts: int = 3          # 最大重试次数
    retry_delay: float = 2.0       # 重试延迟
    max_connections: int = 2000    # 最大连接数

CheckpointEngineConfig

权重同步引擎配置:

@dataclass
class CheckpointEngineConfig:
    type: str = "naive"                    # naive/nccl/nixl
    update_weights_bucket_megabytes: int = 512  # 权重传输缓冲区大小

MultiTurnConfig

多轮对话配置。

AgentLoopConfig

Agent 循环配置(用于 agent 场景的工具调用等)。

与其他模块的关系

  • 被 rollout/ 下所有推理引擎使用
  • enable_sleep_mode 和 free_cache_engine 控制 Hybrid Engine 行为
  • checkpoint_engine 控制权重同步方式

小结

RolloutConfig 是配置项最丰富的配置类,控制推理引擎的方方面面。对 RLHF 训练效率影响很大。