[worker, training_utils] fix: Engine Metric Aggregation

[JacobHelwig]

What does this PR do?

Because some metrics are scaled by global_bsz/global_tokens in workers.utils.losses.ppo_loss, the mean in reduce_metrics adds an extra scaling of the metric by the number of gradient accumulation steps (see examples in Test sec):
https://github.com/volcengine/verl/blob/c191c5eb5c9499dca6666a52bc5f360bfd4bbf4f/verl/utils/metric/utils.py#L53

Aggregation of the loss metric handles this by taking sum: https://github.com/volcengine/verl/blob/c191c5eb5c9499dca6666a52bc5f360bfd4bbf4f/verl/workers/engine_workers.py#L143-L144
Depending on how metrics are handled in workers.utils.losses.ppo_loss, it may not be correct to aggregate all of them using sum (as in #4785). For example, actor/pg_loss and actor/kl_loss are scaled by global batch sizes/ token counts, and should be aggregated using sum, while the pg_metrics pg_clipfrac, ppo_kl, and pg_clipfrac_lower are scaled by local token counts and should be aggregated using mean.

This PR introduces a metric management class to allow flexibility in deciding the aggregation type on a per-metric basis.

Test

This test demonstrates the scaling of metrics with the number of gradient accumulation steps, as well as how this is resolved on this branch. The command for running is below.

gsm8k_train_path=$DATA_DIR/gsm8k/train.parquet
gsm8k_test_path=$DATA_DIR/gsm8k/test.parquet

train_files="['$gsm8k_train_path']"
test_files="['$gsm8k_test_path']"

ppo_micro_batch_size_per_gpu=2
ppo_micro_batch_size_per_gpu=8

branch=main
branch=fixEngineMetrics

python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.dataloader_num_workers=0 \
    data.return_full_prompt=True \
    data.train_files="$train_files" \
    data.val_files="$test_files" \
    data.train_batch_size=8 \
    data.max_prompt_length=512 \
    data.max_response_length=1024 \
    data.filter_overlong_prompts=True \
    data.truncation='error' \
    actor_rollout_ref.model.path=Qwen/Qwen3-0.6B \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.ppo_mini_batch_size=8 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$ppo_micro_batch_size_per_gpu \
    actor_rollout_ref.actor.use_kl_loss=True \
    actor_rollout_ref.actor.fsdp_config.param_offload=True \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
    actor_rollout_ref.rollout.n=5 \
    trainer.logger='["console","wandb"]' \
    trainer.project_name='fixEngineMetrics' \
    trainer.experiment_name="$branch/ppo_micro_batch_size_per_gpu$ppo_micro_batch_size_per_gpu" \
    trainer.n_gpus_per_node=2 \
    trainer.nnodes=1 \
    trainer.save_freq=400 \
    trainer.test_freq=40 \
    trainer.use_legacy_worker_impl=disable \
    trainer.total_epochs=2 \
    trainer.total_training_steps=10 \
    trainer.resume_mode=disable \
    actor_rollout_ref.actor.use_torch_compile=False \
    actor_rollout_ref.actor.fsdp_config.use_torch_compile=False \
    trainer.val_before_train=False \
    actor_rollout_ref.rollout.enforce_eager=True \
    actor_rollout_ref.ref.fsdp_config.use_torch_compile=False

Design & Code Changes

Adds a Metric class which tracks metric values and aggregation type.

[gemini-code-assist]

Code Review

This pull request introduces a MetricList class to provide more flexible aggregation of metrics, which is a good improvement. The changes correctly adapt the ppo_loss function to use this new abstraction.

However, I've found a few critical issues that will cause runtime errors:

1. In verl/utils/metric/utils.py, there's a bug in type checking using isinstance with a typing.Union, which will raise a TypeError.
2. In verl/workers/engine_workers.py, the logic for aggregating metrics from distributed processes doesn't correctly handle the new MetricValue objects, which will also lead to a TypeError.

I've also included some high-severity suggestions to improve code quality and maintainability, such as avoiding wildcard imports and correcting inconsistent type hints. Please review the detailed comments for suggestions on how to fix these issues.

[gemini-code-assist]

Using wildcard imports (import *) is generally discouraged as it can pollute the namespace and make it unclear which names are being imported. It's better to explicitly import the names you need from the utils module. This improves code readability and maintainability.

Suggested change

	from .utils import *
	from .utils import AggregationType, MetricList, MetricValue, reduce_metrics

[gemini-code-assist]

The type hint value: float | MetricValue is inconsistent with the method's implementation, which also handles torch.Tensor and int values. To improve type safety and clarity, the hint should be expanded to include all supported types. Using the Numeric alias, which is Union[int, float, torch.Tensor], would be appropriate here.

Suggested change

	def append(self, value: float | MetricValue) -> None:
	def append(self, value: Union[Numeric, MetricValue]) -> None:

[JacobHelwig]

/gemini

[gemini-code-assist]

Hello @JacobHelwig! It looks like you've invoked me, but you need to specify a command after /gemini.

For example, you can use /gemini review to request a code review or /gemini summary to get a summary of the pull request.

Here's the full list of available commands for your reference:

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

[JacobHelwig]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a flexible metric aggregation system using a MetricList class, which is a valuable improvement for handling different metric scaling behaviors. The implementation is mostly solid, but I've identified two significant issues that could lead to runtime errors. One is a critical IndexError from accessing an empty list, and the other is a TypeError due to improper handling of mixed list and MetricList types for the same metric. My review includes specific code suggestions to resolve these bugs.

[gemini-code-assist]

There's a potential TypeError here. If data[new_key] is a standard Python list (which can happen if an empty list [] was appended for this key previously) and val is a MetricList, the call data[new_key].extend(val) will fail because MetricList is not iterable by default for list.extend. You should handle this case by 'upgrading' the list to a MetricList before extending to ensure type consistency.

        if isinstance(data[new_key], list) and isinstance(val, (MetricValue, MetricList)):
            # Upgrade list to MetricList if it's not already one
            new_list = val.init_list()
            new_list.extend(data[new_key])
            data[new_key] = new_list

        if isinstance(val, (list, MetricList)):
            data[new_key].extend(val)
        else:
            data[new_key].append(val)

[JacobHelwig]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a Metric class to allow for flexible aggregation of metrics, which is a solid approach to address the incorrect scaling issue. The implementation is largely well-executed. However, I've identified a few critical issues that could lead to runtime crashes. These include potential IndexError and AttributeError exceptions from improper handling of empty lists and incorrect method invocation, respectively. Furthermore, the Metric.aggregate method is susceptible to a ValueError when dealing with empty metrics using MIN or MAX aggregation. My specific comments provide detailed suggestions to resolve these critical problems.

[gemini-code-assist]

The init_list method is an instance method of the Metric class, not a class method. Calling Metric.init_list(val) will raise an AttributeError because the Metric class itself does not have this method. You should call it on the Metric instance val, as in val.init_list().

Suggested change

	data[new_key] = Metric.init_list(val) if isinstance(val, Metric) else []
	data[new_key] = val.init_list() if isinstance(val, Metric) else []

[wuxibin89]

@JacobHelwig Please address gemini's review, init_list is not a classmethod.

[wuxibin89]

@JacobHelwig Please format code according to: https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting

[wuxibin89]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a Metric class to provide flexible aggregation strategies for different metrics, addressing an issue where metrics were being incorrectly scaled. The implementation is mostly solid and includes comprehensive tests. However, I've found a critical issue that could lead to a crash due to an IndexError when handling an empty list of metrics. I've also identified a couple of high-severity bugs related to the new Metric class logic that could lead to incorrect behavior or unexpected errors. Please address these points to ensure the stability and correctness of the new metric management system.

[gemini-code-assist]

The expression isinstance(val[0], Metric) will raise an IndexError if val is an empty list. This will crash the worker. You should first check if the list is not empty before accessing its first element.

Suggested change

	Metric.chain(val) if isinstance(val[0], Metric) else list(chain.from_iterable(val))
	Metric.chain(val) if val and isinstance(val[0], Metric) else list(chain.from_iterable(val))

[gemini-code-assist]

The __init__ method does not correctly handle initialization with a list of values. The type hint for value is Optional[Numeric | list[Numeric]], but if a list is passed, self.append(value) is called. This will raise a ValueError because append is designed for single numeric values or Metric objects, not lists. You should check if value is a list and call self.extend(value) in that case.

Suggested change

	if value is not None:
	self.append(value)
	if value is not None:
	if isinstance(value, list):
	self.extend(value)
	else:
	self.append(value)

[gemini-code-assist]

When chain is called with an empty list of metrics, it returns a new Metric with AggregationType.MEAN by default. This arbitrary default can lead to incorrect behavior. For example, if a list of metrics intended for summation is empty, chaining them should result in a metric that aggregates to 0 (the sum of an empty set), not NaN (the mean of an empty set). It would be safer to raise a ValueError if the list is empty, forcing the caller to handle this case explicitly.

Suggested change

	if len(metric_lists) == 0:
	return cls(aggregation=AggregationType.MEAN)
	if not metric_lists:
	raise ValueError("Cannot chain an empty list of metrics.")

[wuxibin89]

@JacobHelwig Docstring Coverage ci failed, please add doc string for Metric.

[JacobHelwig]

Added!

[pengwu22]

Hi @JacobHelwig, thanks for catching the metrics issue!

I think the fix needs one more adjustment. While the summation over micro-batches is correct, the reduction across data-parallel ranks should use mean (not sum) in reduce_metrics at ray_trainer.py, since the actual loss values are already correctly scaled by dp_size in agg_loss.

Recommended approach:

1. Keep Metric.chain in train_mini_batch to sum over micro-batches, OR explicitly aggregate metrics there without the Metric class
2. Use mean reduction in reduce_metrics at ray_trainer.py

This ensures the final metric correctly represents the global per-token average without being inflated by dp_size.

[JacobHelwig]

Nice catch, @pengwu22. Opened #5203 to resolve this.