nemo_automodel.components.loggers.trackio_utils

Trackio run wrapper used by the training recipes.

Module Contents

Classes

Name	Description
`TrackioLogger`	Log training metrics to a Trackio run.

Data

RESERVED_KEYS

logger

API

class nemo_automodel.components.loggers.trackio_utils.TrackioLogger(
    run: typing.Any
)

Log training metrics to a Trackio run.

Wraps the trackio.Run returned by trackio.init so recipes can log MetricsSample dictionaries directly: reserved keys are dropped, tensors are reduced to Python floats and non-numeric values are skipped. Built by nemo_automodel.components.loggers.loggers.TrackioConfig.build on rank 0 only.

Parameters:

run

Any

The trackio.Run to log to.

nemo_automodel.components.loggers.trackio_utils.TrackioLogger.finish() -> None

Finish the run and flush pending metrics.

nemo_automodel.components.loggers.trackio_utils.TrackioLogger.log_metrics(
    metrics: collections.abc.Mapping[str, typing.Any],
    step: int
) -> None

Log a metrics dictionary at step.

Parameters:

metrics

Mapping[str, Any]

Metric name to value. Values may be Python numbers or tensors; a tensor with more than one element is logged as its mean. Reserved Trackio keys (e.g. step, timestamp) and non-numeric values are skipped.

step

int

Training step used as the x-axis.

nemo_automodel.components.loggers.trackio_utils.RESERVED_KEYS = frozenset({'project', 'run', 'timestamp', 'step', 'time', 'metrics'})

nemo_automodel.components.loggers.trackio_utils.logger = logging.getLogger(__name__)