Skip to content

Custom Metrics

About

You can implement a custom metric by implementing the Metric interface. The d9d.metric.component package provides helpers for the distributed state.

Design Guidelines

Metric implementations usually follow these rules:

  • GPU residency: Accumulate data in GPU tensors, so updates need no host sync.
  • Additive states: Do not store averages such as "current accuracy". Store raw counts such as "total correct" and "total samples". Counts stay correct when summed with all_reduce.

Helper Components

MetricAccumulator keeps a local and a synchronized copy of a metric state tensor. It supports the sum, max and min reduction operations (MetricReduceOp). There is no average operation, because averages are not additive.

Usage

This MaxMetric tracks the maximum value seen across all ranks with a MetricAccumulator.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import torch
from typing import Any

from d9d.metric import Metric
from d9d.metric.component import MetricAccumulator, MetricReduceOp
from d9d.core.dist_context import DistributedContext

class MaxMetric(Metric[torch.Tensor]):
    def __init__(self):
        self._max_val = MetricAccumulator(
            torch.tensor(float('-inf')),
            reduce_op=MetricReduceOp.max
        )

    def update(self, value: torch.Tensor):
        # Updates the local maximum, no communication.
        self._max_val.update(value)

    def sync(self, dist_context: DistributedContext):
        # Runs all_reduce across the default process group.
        self._max_val.sync()

    def compute(self) -> torch.Tensor:
        # Returns the synchronized value.
        return self._max_val.value

    def reset(self):
        self._max_val.reset()

    def to(self, device: str | torch.device | int):
        self._max_val.to(device)

    # Stateful protocol for checkpointing.
    def state_dict(self) -> dict[str, Any]:
        return {'max_val': self._max_val.state_dict()}

    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        self._max_val.load_state_dict(state_dict['max_val'])

API Reference

d9d.metric.component

Reusable components for building distributed metrics.

MetricAccumulator

Bases: Stateful

Accumulator of a distributed metric state.

It keeps two copies of the state: a local copy, updated on every step, and a synchronized copy, filled by an all-reduce in sync().

value property

The synchronized value if sync() was called after the last update, otherwise the local value.

__init__(initial_value, reduce_op=MetricReduceOp.sum)

Constructs the MetricAccumulator object.

Parameters:

Name Type Description Default
initial_value Tensor

The starting value, e.g. 0 for sum or -inf for max. It sets the device and dtype of the accumulator.

 required
reduce_op MetricReduceOp

The reduction operation for updates and synchronization.

 sum
load_state_dict(state_dict)

Restores the accumulator state from a checkpoint.

Parameters:

Name Type Description Default
state_dict dict[str, Any]

Dictionary containing state to load.

 required
reset()

Resets the accumulator to its initial state.

state_dict()

Returns the serialized state of the accumulator.

Returns:

Type Description
dict[str, Any]

Dictionary containing local and synchronized tensors and status flags.
sync()

Synchronizes the accumulator across the default process group.

Every rank must call it. The local state stays unchanged.

to(device)

Moves internal tensors to the specified device.

Parameters:

Name Type Description Default
device str | device | int

Target device.

 required
update(value)

Accumulates a value into the local state with the configured reduction operation.

After an update, value returns the local state until the next sync().

Parameters:

Name Type Description Default
value Tensor | float | bool

The value to accumulate. The max and min operations accept only tensors.

 required

Raises:

Type Description
ValueError

If value is not a tensor and the reduction operation is max or min.

MetricReduceOp

Bases: StrEnum

Reduction operation of a MetricAccumulator, used both for updates and for synchronization.

Attributes:

Name Type Description
sum

Adds values.
max

Keeps the maximum value.
min

Keeps the minimum value.