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 训练效率影响很大。