Prhub

#6977 [BREAKING][trainer] fix: separate_async should use the same step granularity with other trainers

原始 PR 作者 Begunner 合并时间 2026-07-08 17:44 文件变更 11 提交数 5 评论 1 代码增减 +375 / -63

执行摘要

对齐 separate_async 步长与同步训练,新增指标聚合

separate_async训练器原本的全局步长与同步训练器不一致,导致global_steps的含义混乱,影响学习率调度、日志记录和checkpoint对齐。本PR旨在将它们统一:一个全局步长对应一次参数同步,内部可包含多次本地更新(由parameter_sync_step控制)。

该 PR 为 BREAKING CHANGE,所有使用 separate_async 训练器的团队都应仔细审查,确认配置和步长语义变更不会破坏现有工作流。MetricsAggregator 的设计思路值得借鉴,可推广到其他异步训练场景。建议作者补充 num_warmup_batches 从 4 改为 1 的说明。

讨论亮点

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 并合并。

实现拆解

  1. 新增 MetricsAggregator 类verl/trainer/ppo/v1/utils.py):管理一个 parameter_sync_step 周期内每次迭代的指标,支持加权平均、求和、最大值、最小值、最后值和时间求和等聚合策略。
  2. 修改 trainer_base.py 的 step 方法:从单次 update 改为循环 parameter_sync_step 次 _step_once,每次调用后收集指标到 MetricsAggregator,最后用聚合后的指标更新 metrics dict。
  3. 调整 separate_async 权重同步时机trainer_separate_async.py):on_step_end 现在每次 step 都触发权重同步,而不是每 parameter_sync_step 步。同时构造函数增强 assert,要求 train_batch_size 等于 parameter_sync_step * ppo_mini_batch_size。
  4. 简化 ReplayBuffer 的 staleness 计算replay_buffer.py):移除 parameter_sync_step 内部属性,staleness 直接取 global_steps - prompt_global_steps + 1,不再除以 parameter_sync_step,因为 global_steps 已经是参数同步版本步数。
  5. 配套测试与配置:新增 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 训练器 modified 8.56
tests/trainer/ppo/v1/test_metrics_aggregator_on_cpu.py CPU 测试 added 7.97
verl/trainer/ppo/v1/trainer_base.py 训练器 modified 7.38
tests/trainer/ppo/v1/test_replay_buffer_on_cpu.py 重播缓冲区测试 renamed 6.78
verl/trainer/ppo/v1/trainer_separate_async.py 训练器 modified 6.06
verl/trainer/ppo/v1/replay_buffer.py 训练器 modified 5.67
verl/trainer/config/ppo_trainer.yaml 配置 modified 3.41

关键符号

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 core-logic

新增 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 core-logic

修改 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)

评论区精华

num_warmup_batches 默认值从 4 减为 1 question

reviewer wuxibin89 在 ppo_trainer.yaml 第 222 行提问:"Why only warmup a mini-batch in separate_async trainer?",指出默认值从 4 改为 1 缺乏解释。作者未直接回复。

结论:PR 仍被 approved 并合并,未明确回应。 · 已解决

风险与影响

  1. BREAKING CHANGE:用户如果依赖旧的 global_steps 语义(如自定义学习率调度、checkpoint 命名等),需要适配。
  2. 权重同步频率:on_step_end 现在每次 step 都调用 update_weights,比原来每 parameter_sync_step 步一次更频繁,可能增加通信开销。
  3. 指标聚合默认行为:如果用户之前直接使用 per-iteration metrics,现在需改用聚合后的值,可能影响监控和告警。
  4. 配置验证更严格:train_batch_size 必须为 parameter_sync_step * ppo_mini_batch_size,否则启动报错。

主要影响使用 separate_async 训练器的用户。他们需要更新配置(train_batch_size、learning rate 步数),理解新的 global_steps 语义,并检查自定义的指标处理代码。新增的 MetricsAggregator 和测试覆盖为新逻辑提供保障,但用户需验证自己的训练流水线不因步长语义变化而受损。

步长语义 BREAKING 权重同步频率变化 配置不兼容

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论