Skip to content

[metric] feat: add weighted aggregation to Metric/reduce_metrics (#1508) - #7501

Closed
amanyagami wants to merge 1 commit into
verl-project:mainfrom
amanyagami:fix-1508-weighted-metrics
Closed

amanyagami wants to merge 1 commit into
verl-project:mainfrom
amanyagami:fix-1508-weighted-metrics

Conversation

@amanyagami

Copy link
Copy Markdown

What

Metric.aggregate() / Metric.aggregate_dp() (landed via #4778) already fixed the mean-of-means problem for metrics migrated onto the Metric API, but there was still no way to carry a per-value weight (e.g. per-microbatch sample/token count) through it, so DP ranks or microbatches with unequal counts were still averaged as if equal. Two independent contributors audited this on the issue thread (2026-07-06, 2026-07-24) and proposed the same direction, but nobody had landed it.

Also fixes a related bug flagged on the same thread: reduce_metrics's raw-list fallback picks a reduction by substring match ("min" in key), which misfires on keys like "actor/minibatch_loss".

Fixes #1508.

Changes

  • Metric now optionally accepts a weight alongside each value (append/extend/__init__), stored as a parallel self.weights list, defaulting to 1.0.
  • Metric.aggregate(): MEAN uses np.average(values, weights=weights) and SUM uses sum(value*weight) only when a non-unit weight is present — otherwise the exact original np.mean/np.sum call runs unchanged (verified byte-for-byte in tests).
  • Metric.aggregate_dp() gained an optional per-rank weights param, combining ranks via Σ(w·v)/Σw instead of a plain np.mean when supplied.
  • reduce_metrics: keys are now split on non-alphanumeric characters, and "min"/"max" must appear as a standalone token — audited every metric key currently produced across verl//recipe/ to confirm none rely on the old loose substring behavior before tightening it.

Fully backward compatible: every existing call site keeps its exact current behavior when no weights are supplied.

Testing

tests/trainer/ppo/test_metric_utils_on_cpu.py: 23 pre-existing + 16 new tests, 39/39 pass. Covers weighted mean/sum, unit-weight byte-for-byte backward compatibility, append/extend weight plumbing (scalar broadcast, length-mismatch errors), weighted aggregate_dp across unequal-sized DP ranks (asserted against the ground-truth flattened mean, contrasted with the visibly-wrong naive unweighted mean-of-means), and the reduce_metrics substring fix.

(Note: this environment's ambient torch==1.12.1 predates tensordict's torch>=2 requirement, which verl/__init__.py transitively imports — an unrelated, pre-existing environment issue confirmed to also affect unmodified main via git stash. Verified instead by loading verl/utils/metric/utils.py directly, which has no verl-internal imports.)

🤖 Generated with Claude Code

…-project#1508)

Implements the two independent proposals from the verl-project#1508 issue thread
(2026-07-06 and 2026-07-24 comments): PR verl-project#4778 introduced `Metric` and
`AggregationType` but did not carry a per-value weight through
aggregation, so DP ranks or micro-batches with unequal sample/token
counts were still averaged as if equal.

- `Metric.__init__`/`append`/`extend` now optionally accept a `weight`
  per value (stored as a parallel `self.weights` list, defaulting each
  value to 1.0 when omitted). A scalar weight passed to `extend`
  broadcasts to every value being added.
- `Metric.aggregate()` / `Metric._aggregate()`: MEAN uses
  `np.average(values, weights=weights)` and SUM uses
  `sum(value * weight)` whenever any weight differs from 1.0; when all
  weights are 1.0 (or omitted) the exact original `np.mean`/`np.sum`
  code path runs unchanged, so behavior is byte-identical for existing
  callers.
- `Metric.aggregate_dp()` now accepts an optional per-rank `weights`
  list (e.g. each rank's local sample/token count). MEAN/SUM combine
  ranks via `np.average(value_arrays, axis=0, weights=weights)`
  (Σ(w·v)/Σw) instead of a plain `np.mean` when weights are supplied;
  MIN/MAX are unaffected (still flatten-then-reduce). Omitting
  `weights` reproduces the previous unweighted behavior exactly.
- `reduce_metrics`: fixed the substring-match bug also raised on the
  issue thread, where `"min"/"max" in key` misfires on keys like
  "...minibatch..." or "...maximum...". Keys are now split on
  non-alphanumeric characters and "min"/"max" must appear as a
  standalone token (e.g. "loss/min", "loss_min") to trigger that
  reduction. Audited every metric key currently produced across
  verl/recipe (compute_data_metrics, losses.py, rollout_corr_helper.py,
  detach_utils.py, etc.) and none rely on the old substring-only
  behavior, so this tightening is safe.

Backward compatible: default weight is 1.0 everywhere, and
`aggregate()`/`aggregate_dp()` fall back to the exact unweighted
numpy calls when no non-default weights are present (verified
numerically in the new tests, e.g. test_weighted_mean_with_unit_weights_
matches_plain_mean, test_weighted_sum_with_unit_weights_matches_plain_
sum, test_aggregate_dp_unweighted_still_matches_plain_mean).

Adds CPU-only unit tests to tests/trainer/ppo/test_metric_utils_on_cpu.py
covering: weighted Metric.aggregate() vs hand-computed weighted mean,
unit-weight byte-for-byte backward compatibility, weighted SUM,
append/extend weight plumbing (including scalar-weight broadcast and
length-mismatch errors), weighted aggregate_dp across DP ranks with
unequal sample counts (asserting against the ground-truth mean of all
flattened raw samples, where the naive unweighted mean-of-means is
visibly wrong), and the reduce_metrics min/max substring-token fix.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@CLAassistant

CLAassistant commented Aug 21, 2026 •

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@wuxibin89 wuxibin89 closed this Aug 21, 2026
@amanyagami amanyagami changed the title feat(metric): add weighted aggregation to Metric/reduce_metrics (#1508) [metric] feat: add weighted aggregation to Metric/reduce_metrics (#1508) Aug 21, 2026
@amanyagami

amanyagami commented Aug 21, 2026 •

Copy link
Copy Markdown
Author

Retitled to match the repo's [module] type: description convention (originally used conventional-commit style by mistake).

Flagging since this closed without a comment, in case it was unreviewed triage: two separate contributors on #1508's thread (2026-07-06, 2026-07-24) independently audited the current Metric/aggregate_dp API and proposed this exact direction — per-value weights, defaulting to 1.0 so it's fully backward compatible — but neither had landed code. 39/39 tests pass, including an explicit byte-for-byte backward-compat proof for the unweighted case.

I don't have permission to reopen this myself (only the closer or a repo collaborator can); please reopen if you'd like it reconsidered, or let me know if there's a reason it's not wanted.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FR] Better metric reduction utility

3 participants