强化学习在推理调度中的探索基于 PPO 的 Adaptive Batch 策略设计一、固定 batch_size 的无奈为什么最佳并发数在变传统推理调度的 batch_size 是一个静态参数——设为 32不改了。但在真实负载下最优 batch_size 是动态变化的凌晨低负载时batch_size4 可降低单请求延迟白天高峰期batch_size32 可最大化吞吐晚间模型更新后新的量化方案可能需要 batch_size16 才能平衡精度与速度。问题不在于没有最优解而在于最优解在实时变化。强化学习RL恰恰擅长处理这种动态环境下的序列决策问题——调度器根据当前 GPU 负载、请求队列长度和 Token 长度分布在每个时间步决定下一个 batch 的大小。二、PPO 调度器的状态空间与奖励函数设计flowchart TD subgraph 环境 S[状态 S_t:br/- 队列长度br/- GPU 利用率br/- Token 长度分布br/- KV Cache 使用率] end subgraph PPO策略网络 S -- A[策略网络 π_θ] A -- D{动作 a_t:br/batch_size 调整} end D --|增大 batch: 32→48| E1[吞吐上升br/延迟可能上升] D --|减小 batch: 32→16| E2[延迟下降br/吞吐可能下降] E1 -- R1[奖励 w1×吞吐 - w2×P99延迟br/- w3×超时率] E2 -- R1 R1 -- Update[更新策略网络br/最大化累计奖励] Update -- S状态空间观测到的环境信息等待队列长度排队中的请求数GPU 利用率当前 SM 占用百分比KV Cache 使用率显存压力指标当前 batch 中 Token 长度的 P50/P95/P99动作空间增大 batch_size、减小 batch_size、维持不变3 种离散动作奖励函数reward w1 × throughput_norm - w2 × p99_latency_norm - w3 × timeout_rate - w4 × gpu_idle_penalty关键是w1w4的权重设置——如果吞吐权重过高调度器会无限增大 batch_size如果延迟权重过高调度器会把 batch_size 压到 1浪费 GPU。三、PPO 调度器的训练与在线推理import numpy as np import torch import torch.nn as nn import torch.optim as optim # PPO 策略网络 class BatchSchedulerPolicy(nn.Module): 输入环境状态 (队列长度, GPU利用率, KV使用率, TokenP50, TokenP99) 输出三种动作的 log 概率 (增大batch, 减小batch, 不变) def __init__(self, state_dim: int 5, hidden_dim: int 64, action_dim: int 3): super().__init__() # 共享的特征提取层 self.feature_net nn.Sequential( nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), ) # Actor 头输出动作概率分布 self.actor nn.Linear(hidden_dim, action_dim) # Critic 头输出状态价值 V(s) self.critic nn.Linear(hidden_dim, 1) def forward(self, state): features self.feature_net(state) action_logits self.actor(features) state_value self.critic(features) return action_logits, state_value # 在线推理调度循环 class AdaptiveBatchScheduler: 基于训练好的 PPO 策略进行在线调度 def __init__(self, policy: BatchSchedulerPolicy, min_batch: int 4, max_batch: int 64): self.policy policy self.policy.eval() self.current_batch 16 # 初始 batch size self.min_batch min_batch self.max_batch max_batch def observe_state(self, vllm_metrics: dict) - np.ndarray: 从 vLLM Metrics 提取状态特征 return np.array([ vllm_metrics.get(num_requests_waiting, 0) / 50.0, # 归一化 vllm_metrics.get(gpu_utilization, 0) / 100.0, vllm_metrics.get(gpu_cache_usage_perc, 0) / 100.0, vllm_metrics.get(token_length_p50, 0) / 4096.0, vllm_metrics.get(token_length_p99, 0) / 4096.0, ], dtypenp.float32) def step(self, state: np.ndarray) - int: 根据当前状态决定 batch_size 调整动作 返回新的 batch_size with torch.no_grad(): state_tensor torch.FloatTensor(state).unsqueeze(0) action_logits, _ self.policy(state_tensor) # 选择概率最高的动作 action torch.argmax(action_logits.squeeze()).item() # 动作映射 if action 0 and self.current_batch self.max_batch: self.current_batch 8 # 增大 elif action 1 and self.current_batch self.min_batch: self.current_batch - 8 # 减小 # action 2: 保持 return self.current_batch def compute_reward(self, throughput: float, p99_latency: float, timeout_rate: float) - float: 奖励函数平衡吞吐量与延迟 - throughput: tokens/s越高越好 - p99_latency: ms越低越好 - timeout_rate: 超时比例越低越好 w_throughput 0.01 # 吞吐量权重 w_latency -0.5 # 延迟惩罚 w_timeout -100.0 # 超时重罚 # 归一化 throughput_norm min(throughput / 3000.0, 1.0) # 3000 tokens/s 为满分 latency_norm min(p99_latency / 2000.0, 1.0) # 延迟 2s 可接受 reward (w_throughput * throughput_norm w_latency * latency_norm w_timeout * timeout_rate) return reward # 主调度循环 scheduler AdaptiveBatchScheduler(policyload_trained_policy()) while True: # 1. 采集当前 vLLM 状态 metrics fetch_vllm_metrics() state scheduler.observe_state(metrics) # 2. PPO 策略决策 new_batch scheduler.step(state) # 3. 应用新的 batch_size apply_batch_size(new_batch) # 4. 等待 N 秒后采集效果 time.sleep(10) metrics_after fetch_vllm_metrics() # 5. 计算奖励作为监控指标 reward scheduler.compute_reward( throughputmetrics_after[throughput], p99_latencymetrics_after[p99_latency], timeout_ratemetrics_after[timeout_rate], ) log_metrics(reward, new_batch)四、RL 调度器的工程可靠性边界离线训练 vs 在线适配PPO 模型是在历史负载数据上离线训练的。当流量模式发生根本性变化时如新模型上线离线训练的模型可能给出次优决策。需配合安全策略Safe Policy——当 PPO 建议的 batch_size 导致延迟超过 SLA 阈值的 120% 时回退到固定保守策略。仿真环境与真实环境的 GapPPO 训练需要一个推理环境仿真器来快速迭代。如果仿真器的延迟/吞吐模型与实际 vLLM 有 20% 的偏差训练出的策略在实际部署时会表现不一致。建议先在历史 traces 上验证仿真精度。动作空间的粒度batch_size 调整步长设为 8 而非 1是在策略灵活性与模型复杂度之间的权衡——步长 1 会导致动作空间爆炸60 种可能步长 8 将空间从 60 压缩到 7。五、总结PPO 在推理调度中的应用本质上是用 RL 替代人工的调参直觉。核心是奖励函数的设计——它直接决定了调度器是吞吐优先还是延迟优先。落地建议先在仿真环境中训练 PPO 模型基于历史负载 traces在num_envs4的并行环境中训练 1000 episodes然后用 A/B Test 比对 PPO 调度 vs 固定 batch_size 的效果最终以 10% 的流量灰度上线通过 PPO 的惩罚权重微调来平衡激进和保守。
强化学习在推理调度中的探索基于 PPO 的 Adaptive Batch 策略设计一、固定 batch_size 的无奈为什么最佳并发数在变传统推理调度的 batch_size 是一个静态参数——设为 32不改了。但在真实负载下最优 batch_size 是动态变化的凌晨低负载时batch_size4 可降低单请求延迟白天高峰期batch_size32 可最大化吞吐晚间模型更新后新的量化方案可能需要 batch_size16 才能平衡精度与速度。问题不在于没有最优解而在于最优解在实时变化。强化学习RL恰恰擅长处理这种动态环境下的序列决策问题——调度器根据当前 GPU 负载、请求队列长度和 Token 长度分布在每个时间步决定下一个 batch 的大小。二、PPO 调度器的状态空间与奖励函数设计flowchart TD subgraph 环境 S[状态 S_t:br/- 队列长度br/- GPU 利用率br/- Token 长度分布br/- KV Cache 使用率] end subgraph PPO策略网络 S -- A[策略网络 π_θ] A -- D{动作 a_t:br/batch_size 调整} end D --|增大 batch: 32→48| E1[吞吐上升br/延迟可能上升] D --|减小 batch: 32→16| E2[延迟下降br/吞吐可能下降] E1 -- R1[奖励 w1×吞吐 - w2×P99延迟br/- w3×超时率] E2 -- R1 R1 -- Update[更新策略网络br/最大化累计奖励] Update -- S状态空间观测到的环境信息等待队列长度排队中的请求数GPU 利用率当前 SM 占用百分比KV Cache 使用率显存压力指标当前 batch 中 Token 长度的 P50/P95/P99动作空间增大 batch_size、减小 batch_size、维持不变3 种离散动作奖励函数reward w1 × throughput_norm - w2 × p99_latency_norm - w3 × timeout_rate - w4 × gpu_idle_penalty关键是w1w4的权重设置——如果吞吐权重过高调度器会无限增大 batch_size如果延迟权重过高调度器会把 batch_size 压到 1浪费 GPU。三、PPO 调度器的训练与在线推理import numpy as np import torch import torch.nn as nn import torch.optim as optim # PPO 策略网络 class BatchSchedulerPolicy(nn.Module): 输入环境状态 (队列长度, GPU利用率, KV使用率, TokenP50, TokenP99) 输出三种动作的 log 概率 (增大batch, 减小batch, 不变) def __init__(self, state_dim: int 5, hidden_dim: int 64, action_dim: int 3): super().__init__() # 共享的特征提取层 self.feature_net nn.Sequential( nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), ) # Actor 头输出动作概率分布 self.actor nn.Linear(hidden_dim, action_dim) # Critic 头输出状态价值 V(s) self.critic nn.Linear(hidden_dim, 1) def forward(self, state): features self.feature_net(state) action_logits self.actor(features) state_value self.critic(features) return action_logits, state_value # 在线推理调度循环 class AdaptiveBatchScheduler: 基于训练好的 PPO 策略进行在线调度 def __init__(self, policy: BatchSchedulerPolicy, min_batch: int 4, max_batch: int 64): self.policy policy self.policy.eval() self.current_batch 16 # 初始 batch size self.min_batch min_batch self.max_batch max_batch def observe_state(self, vllm_metrics: dict) - np.ndarray: 从 vLLM Metrics 提取状态特征 return np.array([ vllm_metrics.get(num_requests_waiting, 0) / 50.0, # 归一化 vllm_metrics.get(gpu_utilization, 0) / 100.0, vllm_metrics.get(gpu_cache_usage_perc, 0) / 100.0, vllm_metrics.get(token_length_p50, 0) / 4096.0, vllm_metrics.get(token_length_p99, 0) / 4096.0, ], dtypenp.float32) def step(self, state: np.ndarray) - int: 根据当前状态决定 batch_size 调整动作 返回新的 batch_size with torch.no_grad(): state_tensor torch.FloatTensor(state).unsqueeze(0) action_logits, _ self.policy(state_tensor) # 选择概率最高的动作 action torch.argmax(action_logits.squeeze()).item() # 动作映射 if action 0 and self.current_batch self.max_batch: self.current_batch 8 # 增大 elif action 1 and self.current_batch self.min_batch: self.current_batch - 8 # 减小 # action 2: 保持 return self.current_batch def compute_reward(self, throughput: float, p99_latency: float, timeout_rate: float) - float: 奖励函数平衡吞吐量与延迟 - throughput: tokens/s越高越好 - p99_latency: ms越低越好 - timeout_rate: 超时比例越低越好 w_throughput 0.01 # 吞吐量权重 w_latency -0.5 # 延迟惩罚 w_timeout -100.0 # 超时重罚 # 归一化 throughput_norm min(throughput / 3000.0, 1.0) # 3000 tokens/s 为满分 latency_norm min(p99_latency / 2000.0, 1.0) # 延迟 2s 可接受 reward (w_throughput * throughput_norm w_latency * latency_norm w_timeout * timeout_rate) return reward # 主调度循环 scheduler AdaptiveBatchScheduler(policyload_trained_policy()) while True: # 1. 采集当前 vLLM 状态 metrics fetch_vllm_metrics() state scheduler.observe_state(metrics) # 2. PPO 策略决策 new_batch scheduler.step(state) # 3. 应用新的 batch_size apply_batch_size(new_batch) # 4. 等待 N 秒后采集效果 time.sleep(10) metrics_after fetch_vllm_metrics() # 5. 计算奖励作为监控指标 reward scheduler.compute_reward( throughputmetrics_after[throughput], p99_latencymetrics_after[p99_latency], timeout_ratemetrics_after[timeout_rate], ) log_metrics(reward, new_batch)四、RL 调度器的工程可靠性边界离线训练 vs 在线适配PPO 模型是在历史负载数据上离线训练的。当流量模式发生根本性变化时如新模型上线离线训练的模型可能给出次优决策。需配合安全策略Safe Policy——当 PPO 建议的 batch_size 导致延迟超过 SLA 阈值的 120% 时回退到固定保守策略。仿真环境与真实环境的 GapPPO 训练需要一个推理环境仿真器来快速迭代。如果仿真器的延迟/吞吐模型与实际 vLLM 有 20% 的偏差训练出的策略在实际部署时会表现不一致。建议先在历史 traces 上验证仿真精度。动作空间的粒度batch_size 调整步长设为 8 而非 1是在策略灵活性与模型复杂度之间的权衡——步长 1 会导致动作空间爆炸60 种可能步长 8 将空间从 60 压缩到 7。五、总结PPO 在推理调度中的应用本质上是用 RL 替代人工的调参直觉。核心是奖励函数的设计——它直接决定了调度器是吞吐优先还是延迟优先。落地建议先在仿真环境中训练 PPO 模型基于历史负载 traces在num_envs4的并行环境中训练 1000 episodes然后用 A/B Test 比对 PPO 调度 vs 固定 batch_size 的效果最终以 10% 的流量灰度上线通过 PPO 的惩罚权重微调来平衡激进和保守。
！原来鼠标右键可以这么强大！)
