MMEval: Machine Learning Evaluation Library

MMEval abstracts distributed communication requirements through the BaseDistBackend interface. This abstraction allows the library to support multiple communication libraries (like MPI) by providing a unified way to perform all-gather and broadcast operations on Python objects.

To implement a custom backend, you must inherit from mmeval.core.dist_backends.BaseDistBackend and implement the following interface:

• is_initialized (bool): Indicates if the distributed environment is ready.
• rank (int): The rank index of the current process group.
• world_size (int): The total number of processes in the current group.
• all_gather_object(obj): Performs an all-gather operation on a Pickle-serializable Python object and returns a list of objects from all processes.
• broadcast_object(obj, src=0): Broadcasts a Pickle-serializable Python object from a source process (defaulting to rank 0) to all other processes.

Note that all objects passed to these methods must be serializable by Pickle.

from mpi4py import MPI
from typing import Any, List
import os
from mmeval.core.dist_backends import BaseDistBackend

class MPI4PyDist(BaseDistBackend):
    """A distributed communication backend for mpi4py."""

    @property
    def is_initialized(self) -> bool:
        """Returns True if the distributed environment has been initialized."""
        return 'OMPI_COMM_WORLD_SIZE' in os.environ

    @property
    def rank(self) -> int:
        """Returns the rank index of the current process group."""
        comm = MPI.COMM_WORLD
        return comm.Get_rank()

    @property
    def world_size(self) -> int:
        """Returns the world size of the current process group."""
        comm = MPI.COMM_WORLD
        return comm.Get_size()

    def all_gather_object(self, obj: Any) -> List[Any]:
        """All gather the given object from the current process group and
        returns a list consisting gathered object of each process."""
        comm = MPI.COMM_WORLD
        return comm.allgather(obj)

    def broadcast_object(self, obj: Any, src: int = 0) -> Any:
        """Broadcast the given object from source process to the current
        process group."""
        comm = MPI.COMM_WORLD
        return comm.bcast(obj, root=src)

The mmeval.core.dist_backends module provides various backend implementations for handling distributed evaluation. These backends allow mmeval to scale evaluation processes across different distributed computing frameworks and hardware setups.

Available backend types include:

• Base Classes:
  • BaseDistBackend: The base class for all distributed backends.
  • TensorBaseDistBackend: A base class for backends that operate on tensors.
• Non-Distributed:
  • NonDist: Used when no distributed evaluation is required.
• Framework/Hardware Specific Backends:
  • MPI4PyDist: Uses MPI (via mpi4py) for distributed communication.
  • TorchCPUDist: Uses PyTorch for distributed evaluation on CPUs.
  • TorchCUDADist: Uses PyTorch for distributed evaluation on GPUs (CUDA).
  • TFHorovodDist: Uses TensorFlow with Horovod for distributed evaluation.
  • PaddleDist: Uses PaddlePaddle for distributed evaluation.
  • OneFlowDist: Uses OneFlow for distributed evaluation.

In distributed evaluation, datasets are often split across multiple GPUs/cards to speed up inference. Because simple averaging of metrics from different subsets is often mathematically incorrect, BaseMetric provides a mechanism to synchronize results across all processes before the final calculation.

BaseMetric follows a two-step lifecycle:

1. Data Collection: The user calls add() to store inference results or intermediate metric values into the BaseMetric._results list.
2. Synchronization and Computation: The user calls compute(). This triggers an all-gather operation across all processes to synchronize the _results list, followed by a call to the user-defined compute_metric() method to calculate the final metric.

To ensure correctness when using data samplers like PyTorch's DistributedSampler (which may pad datasets with duplicate samples to ensure even distribution across cards), BaseMetric.compute() can accept a size parameter. This parameter represents the true number of samples in the dataset. BaseMetric will use the dist_collect_mode to remove these duplicate padding samples before calling compute_metric().

Note: To successfully remove padding samples, the values stored in _results must maintain a one-to-one mapping with the original dataset samples.

class MyMetric(BaseMetric):
    def add(self, result):
        # Store results in the internal list
        self._results.append(result)

    def compute_metric(self, results):
        # Perform the actual math on the synchronized results
        return sum(results) / len(results)

# Usage flow:
# 1. metric = MyMetric(...)
# 2. metric.add(inference_result)
# 3. final_score = metric.compute(size=true_dataset_size)