执行摘要
- 一句话:对齐separate_async步长与同步训练,新增指标聚合
- 推荐动作:该 PR 为 BREAKING CHANGE,所有使用 separate_async 训练器的团队都应仔细审查,确认配置和步长语义变更不会破坏现有工作流。MetricsAggregator 的设计思路值得借鉴,可推广到其他异步训练场景。建议作者补充
num_warmup_batches 从 4 改为 1 的说明。
功能与动机
separate_async训练器原本的全局步长与同步训练器不一致,导致global_steps的含义混乱,影响学习率调度、日志记录和checkpoint对齐。本PR旨在将它们统一:一个全局步长对应一次参数同步,内部可包含多次本地更新(由parameter_sync_step控制)。
实现拆解
- 新增 MetricsAggregator 类(
verl/trainer/ppo/v1/utils.py):管理一个 parameter_sync_step 周期内每次迭代的指标,支持加权平均、求和、最大值、最小值、最后值和时间求和等聚合策略。
- 修改 trainer_base.py 的 step 方法:从单次 update 改为循环 parameter_sync_step 次
_step_once,每次调用后收集指标到 MetricsAggregator,最后用聚合后的指标更新 metrics dict。
- 调整 separate_async 权重同步时机(
trainer_separate_async.py):on_step_end 现在每次 step 都触发权重同步,而不是每 parameter_sync_step 步。同时构造函数增强 assert,要求 train_batch_size 等于 parameter_sync_step * ppo_mini_batch_size。
- 简化 ReplayBuffer 的 staleness 计算(
replay_buffer.py):移除 parameter_sync_step 内部属性,staleness 直接取 global_steps - prompt_global_steps + 1,不再除以 parameter_sync_step,因为 global_steps 已经是参数同步版本步数。
- 配套测试与配置:新增
test_metrics_aggregator_on_cpu.py 覆盖所有聚合规则;移动并重写 test_replay_buffer_on_cpu.py 从 v2 到 v1 适配新语义;更新 ppo_trainer.yaml 及四个自动生成的 yaml 配置文件,默认 num_warmup_batches 从 4 改为 1。
关键文件:
verl/trainer/ppo/v1/utils.py(模块 训练器;类别 source;类型 core-logic;符号 MetricsAggregator, init, _init_aggregation_rules, add_step_metrics): 新增 MetricsAggregator 类,实现 parameter_sync_step 周期内指标聚合的核心逻辑
tests/trainer/ppo/v1/test_metrics_aggregator_on_cpu.py(模块 CPU 测试;类别 test;类型 test-coverage;符号 test_empty_aggregator_returns_empty, test_mean_metrics_are_weighted_by_sample_count, test_default_metric_is_weighted_by_sample_count, test_max_and_min_are_reduced): 新增 MetricsAggregator 的单元测试,覆盖所有聚合规则和边界情况
verl/trainer/ppo/v1/trainer_base.py(模块 训练器;类别 source;类型 core-logic;符号 _step_once): 修改 step 方法,使用 MetricsAggregator 循环调用 _step_once 并聚合指标
tests/trainer/ppo/v1/test_replay_buffer_on_cpu.py(模块 重播缓冲区测试;类别 test;类型 rename-or-move;符号 test_drop_respects_parameter_sync_step, test_drop_uses_version_based_staleness): 从 v2 目录移动至 v1,并更新 staleness 计算为直接差值(不再除以 parameter_sync_step)
verl/trainer/ppo/v1/trainer_separate_async.py(模块 训练器;类别 source;类型 core-logic): 调整 assert 条件,on_step_end 中每次 step 都进行权重同步,并移除条件判断
verl/trainer/ppo/v1/replay_buffer.py(模块 训练器;类别 source;类型 core-logic): 移除 parameter_sync_step 属性,修改 staleness 计算为直接差值
verl/trainer/config/ppo_trainer.yaml(模块 配置;类别 config;类型 configuration): 调整 separate_async 默认配置,num_warmup_batches 从 4 改为 1
关键符号:MetricsAggregator, _step_once, add_step_metrics, get_aggregated_metrics, _get_aggregation_type, _get_metric_weight, _init_aggregation_rules, _aggregate_single_metric
关键源码片段
verl/trainer/ppo/v1/utils.py
新增 MetricsAggregator 类,实现 parameter_sync_step 周期内指标聚合的核心逻辑
class MetricsAggregator:
"""
聚合 parameter_sync_step 周期内的每个迭代训练指标。
支持加权平均、求和、最大值、最小值、最后值、时间求和等策略。
"""
def __init__(self):
self.metric_values: dict[str, list[float]] = defaultdict(list)
self.metric_weights: dict[str, list[int]] = defaultdict(list)
self.step_count = 0
self.aggregation_rules = self._init_aggregation_rules()
def _init_aggregation_rules(self) -> dict[str, list[str]]:
return {
"sum": ["training/off_policy/dropped_samples", "validation/off_policy/dropped_samples"],
"last": ["training/global_step", "training/rollout_probs_diff_valid"],
}
def add_step_metrics(self, metrics: dict[str, Any], sample_count: int = 0):
self.step_count += 1
for key, value in metrics.items():
if isinstance(value, bool):
continue
if isinstance(value, int | float | np.number):
self.metric_values[key].append(float(value))
self.metric_weights[key].append(self._get_metric_weight(key, metrics, sample_count))
elif isinstance(value, torch.Tensor) and value.numel() == 1:
self.metric_values[key].append(float(value.item()))
self.metric_weights[key].append(self._get_metric_weight(key, metrics, sample_count))
def _get_metric_weight(self, metric_name: str, metrics: dict[str, Any], sample_count: int) -> int:
# 对于 dropped_samples_staleness,使用实际丢弃样本数作为权重
if metric_name.endswith("/off_policy/dropped_samples_staleness/mean"):
prefix = metric_name.rsplit("_staleness/mean", 1)[0]
dropped_samples = metrics.get(prefix, sample_count)
if isinstance(dropped_samples, torch.Tensor):
return int(dropped_samples.item()) if dropped_samples.numel() == 1 else sample_count
if isinstance(dropped_samples, int | float | np.number):
return int(dropped_samples)
return sample_count
def _get_aggregation_type(self, metric_name: str) -> str:
for agg_type, metric_list in self.aggregation_rules.items():
if metric_name in metric_list:
return agg_type
metric_lower = metric_name.lower()
# 学习率始终取最后一个值
if metric_lower.endswith("/lr") or metric_lower.endswith("_lr") or metric_lower == "lr":
return "last"
# 时间指标累加
if "timing_s/" in metric_lower or "timing_per_token_ms/" in metric_lower:
return "time_sum"
# 包含 max 或 maximum 取最大值
if any(keyword in metric_lower for keyword in ["max", "maximum"]):
return "max"
# 包含 min 或 minimum 取最小值
if any(keyword in metric_lower for keyword in ["min", "minimum"]):
return "min"
# 包含 sum 或 total 取和
if any(keyword in metric_lower for keyword in ["sum", "total"]):
return "sum"
# 默认使用加权平均
return "weighted_avg"
verl/trainer/ppo/v1/trainer_base.py
修改 step 方法,使用 MetricsAggregator 循环调用 _step_once 并聚合指标
def step(self, metrics: dict, timing_raw: dict) -> KVBatchMeta:
train_batch_size = self.config.data.train_batch_size
assert train_batch_size % self.parameter_sync_step == 0, (
f"train_batch_size ({train_batch_size}) must be divisible by "
f"parameter_sync_step ({self.parameter_sync_step})"
)
sample_batch_size = train_batch_size // self.parameter_sync_step
# 初始化聚合器
metrics_aggregator = MetricsAggregator()
combined_keys: list = []
combined_tags: list = []
combined_partition_id = "train"
# 循环执行 parameter_sync_step 次本地更新
for _ in range(self.parameter_sync_step):
iter_metrics: dict = {}
batch = self._step_once(iter_metrics, timing_raw, sample_batch_size)
# 计算非 padding 样本数作为聚合权重
sample_count = sum(not tag.get("is_padding", False) for tag in batch.tags)
metrics_aggregator.add_step_metrics(iter_metrics, sample_count=sample_count)
combined_keys.extend(batch.keys)
combined_tags.extend(batch.tags)
combined_partition_id = batch.partition_id
# 用聚合后的指标更新整体 metrics
metrics.update(metrics_aggregator.get_aggregated_metrics())
return KVBatchMeta(partition_id=combined_partition_id, keys=combined_keys, tags=combined_tags)
评论区精华
reviewer wuxibin89 在 verl/trainer/config/ppo_trainer.yaml 第 222 行提问:“Why only warmup a mini-batch in separate_async trainer?”,指出默认 num_warmup_batches 从 4 改为 1 缺乏解释。该问题未收到作者直接回复,但 PR 后续被 approved 并合并。
- num_warmup_batches 默认值从4减为1 (question): PR 仍被 approved 并合并,未明确回应。
风险与影响
- 风险:
- BREAKING CHANGE:用户如果依赖旧的 global_steps 语义(如自定义学习率调度、checkpoint 命名等),需要适配。
- 权重同步频率:on_step_end 现在每次 step 都调用 update_weights,比原来每 parameter_sync_step 步一次更频繁,可能增加通信开销。
- 指标聚合默认行为:如果用户之前直接使用 per-iteration metrics,现在需改用聚合后的值,可能影响监控和告警。
- 配置验证更严格:train_batch_size 必须为 parameter_sync_step * ppo_mini_batch_size,否则启动报错。
- 影响:主要影响使用 separate_async 训练器的用户。他们需要更新配置(train_batch_size、learning rate 步数),理解新的 global_steps 语义,并检查自定义的指标处理代码。新增的 MetricsAggregator 和测试覆盖为新逻辑提供保障,但用户需验证自己的训练流水线不因步长语义变化而受损。
- 风险标记:步长语义 BREAKING, 权重同步频率变化, 配置不兼容
关联脉络
参与讨论