Metric Collection
About
The d9d.internals.metric_collector package processes metrics without blocking the training loop. The Metric interface is synchronous. AsyncMetricCollector wraps a metric and runs its synchronization and computation on a side CUDA stream. So the training loop continues at once and does not wait for the metric all-reduce. Collecting the results makes the current stream wait for the side stream, but not the host: the results stay tensors on the GPU.
Internal API
If you use the standard d9d Trainer, you do not need to use this package directly. This page is for users who write their own training loop or logging.
API Reference
d9d.internals.metric_collector
Asynchronous synchronization and computation of metrics.
AsyncMetricCollector
Asynchronous metric collector that synchronizes and computes a metric on a side CUDA stream.
The distributed reduction and the computation run on the side stream, off the main training stream.
__init__(metric)
Constructs the AsyncMetricCollector object.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metric
|
Metric
|
The metric to collect. |
required |
bind()
Moves the metric to CUDA and creates the side stream.
collect_results()
Returns the results of the asynchronous computation and resets the metric.
The current CUDA stream waits for the computation, the host does not.
Returns:
| Type | Description |
|---|---|
PyTree[Tensor]
|
A PyTree with the structure of the metric output and tensors on the device of the metric as leaves. |
Raises:
| Type | Description |
|---|---|
RuntimeError
|
If the collector is not bound, or if |
schedule_collection(dist_context)
Schedules the metric synchronization and computation on the side stream.
In a distributed setup, the metric is synchronized across ranks before the computation.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
dist_context
|
DistributedContext
|
The distributed context to synchronize the metric with. |
required |
Raises:
| Type | Description |
|---|---|
RuntimeError
|
If the collector is not bound. |
unbind()
Releases the reference to the side stream.