Trackers and Metrics

Logging values and other metadata about a run is a core requirement for any ML framework. Until recently, Levanter had a hard dependency on W&B for tracking such values. We now provide a pluggable tracker interface with built‑in support for W&B, Tensorboard, and Trackio.

In the latest version, we introduce the levanter.tracker.Tracker interface, which allows you to use any tracking backend you want. The interface name is taken from the HuggingFace Accelerate framework.

Levanter ships with trackers for W&B, TensorBoard, and a lightweight JSON logger that emits structured log lines. The interface is designed to look similar to W&B's API. The methods currently exposed are:

• levanter.tracker.current_tracker: returns the current tracker instance or sets it.
• levanter.tracker.log: logs a dictionary of metrics for a given step.
• levanter.tracker.log_summary: logs a dictionary of "summary" information, analogous to W&B's version.
• levanter.tracker.get_tracker: returns a tracker with the given name.
• levanter.tracker.jit_log: a version of levanter.tracker.log that accumulates metrics inside of a jit-ted function.

A basic example of using the tracker interface is shown below:

import wandb
import levanter.tracker as tracker
from levanter.tracker.wandb import WandbTracker

with tracker.current_tracker(WandbTracker(wandb.init())):
    for step in range(100):
        tracker.log({"loss": 100 - 0.01 * step}, step=step)

    tracker.log_summary({"best_loss": 0.0})

A more typical example would be to use it in a config file, as we do with Trainer:

trainer:
  tracker:
    type: wandb
    project: my-project
    entity: my-entity

Multiple Trackers

In some cases, you may want to use multiple trackers at once. For example, you may want to use both W&B and Tensorboard.

To do this, you can use the levanter.tracker.tracker.CompositeTracker class, or, if using a config file, you can specify multiple trackers:

trainer:
  tracker:
    - type: wandb
      project: my-project
      entity: my-entity
    - type: tensorboard
      logdir: logs

Installation note: the TensorBoard tracker depends on tensorboardX. Install the profiling extra to get both TensorBoard and TensorBoardX: pip install "levanter[profiling]" (or uv sync --extra profiling).

Adding your own tracker

To add your own tracker, you need to implement the levanter.tracker.Tracker interface. You will also want to register your config with TrackerConfig as a "choice" in the choice type. Follow the pattern for Tensorboard and W&B.

TODO: expand this section.

API Reference

Core Functions

current_tracker(tracker: Optional[Tracker] = None) -> Tracker | typing.ContextManager

current_tracker() -> Tracker

current_tracker(tracker: Tracker) -> typing.ContextManager

Get or set the global tracker. Note that setting the global tracker is not thread-safe, and using a tracker from multiple threads is only supported if the tracker itself is thread-safe.

Parameters:

• tracker

  (Optional[Tracker], default: None ) –

  If provided, returns a context manager that sets the global tracker to the provided tracker when used.

Returns:

• Tracker | ContextManager –

  If no tracker is provided, returns the current global tracker.

• Tracker | ContextManager –

  If a tracker is provided, returns a context manager that sets the global tracker to the provided tracker when used.

Examples:

>>> from levanter.tracker import current_tracker, log
>>> from levanter.tracker.wandb import WandbTracker
>>> with current_tracker(WandbTracker()):
...     log({"foo": 1}, step=0)
...     current_tracker().log({"foo": 2}, step=1)

log(metrics: typing.Mapping[str, LoggableValue | Any], *, step: Optional[int], commit: Optional[bool] = None)

Log metrics to the global tracker.

Parameters:

• metrics

  (Mapping[str, LoggableValue | Any]) –

  Metrics to log. We use LoggableValues just to give you a sense of what you can log. Backends may support additional types.

• step

  (Optional[int]) –

  Step to log at. If None, uses the default for the tracker.

• commit

  (Optional[bool], default: None ) –

  Whether to commit the metrics. If None, uses the default for the tracker.

log_summary(metrics: dict[str, Any])

Log summary metrics to the global tracker.

Parameters:

• metrics

  (dict[str, Any]) –

  Metrics to log

get_tracker(name: str) -> Tracker

get_tracker(name: Literal['wandb']) -> WandbTracker

get_tracker(name: Literal['tensorboard']) -> TensorboardTracker

get_tracker(name: Literal['trackio']) -> TrackioTracker

get_tracker(name: str) -> Tracker

Lookup a tracker in the current global tracker with the provided name.

Parameters:

• name

  (str) –

  Name of the tracker to lookup

Returns:

• Tracker –

  The tracker with the provided name

Examples:

>>> from levanter.tracker import get_tracker, log
>>> from levanter.tracker.wandb import WandbTracker
>>> with current_tracker(WandbTracker()):
...     log({"foo": 1}, step=0)
...     get_tracker("wandb").log_metrics({"foo": 2}, step=1)

jit_log(metrics, *, step=None)

JAX doesn't allow tracers to escape the jit boundary, so we have to be clever about how we log metrics. In Levanter, we enable tracking inside jit with two mechanisms.

• The first, most performant way, is to use the levanter.tracker.defer_tracker_for_jit context manager, which will cause logging to go to a dictionary (that is returned by capture_logging). You can then return this dictionary from the JIT function and log it outside of the JIT.
• The second way is to just use an effect callback to log the metrics to the host.

We strongly recommend using the first method, as it is much more performant.

Trackers

Tracker

Bases: ABC

A tracker is responsible for logging metrics, hyperparameters, and artifacts. Meant to be used with the levanter.tracker.current_tracker context manager, but can also be used directly.

The name is borrowed from HF Accelerate.

Examples:

>>> from levanter.tracker import current_tracker, log
>>> from levanter.tracker.wandb import WandbTracker
>>> with current_tracker(WandbTracker()):
...     log({"foo": 1}, step=0)

Methods:

• log_hyperparameters –
• log –

  Log metrics to the tracker. Step is always required.

• log_summary –
• log_artifact –
• finish –

  Finish the tracker. This is called when the tracker is no longer needed. This can, e.g.,

Attributes:

• name (str) –

name: str instance-attribute

log_hyperparameters(hparams: dict[str, Any]) abstractmethod

log(metrics: typing.Mapping[str, typing.Any], *, step: Optional[int], commit: Optional[bool] = None) abstractmethod

Log metrics to the tracker. Step is always required.

Parameters:

• metrics

  (Mapping[str, Any]) –

  Metrics to log

• step

  (Optional[int]) –

  Step to log at

• commit

  (Optional[bool], default: None ) –

  Whether to commit the metrics. If None, uses the default for the tracker.

log_summary(metrics: dict[str, Any]) abstractmethod

log_artifact(artifact_path, *, name: Optional[str] = None, type: Optional[str] = None) abstractmethod

finish() abstractmethod

Finish the tracker. This is called when the tracker is no longer needed. This can, e.g., force a commit of all metrics.

CompositeTracker(loggers: List[Tracker])

Bases: Tracker

Methods:

• log_hyperparameters –
• log –
• log_summary –
• log_artifact –
• finish –

Attributes:

• loggers –
• name (str) –

loggers = loggers instance-attribute

name: str instance-attribute

log_hyperparameters(hparams: dict[str, Any])

log(metrics: typing.Mapping[str, Any], *, step, commit=None)

log_summary(metrics: dict[str, Any])

log_artifact(artifact_path, *, name: Optional[str] = None, type: Optional[str] = None)

finish()

NoopTracker

Bases: Tracker

Methods:

• log_hyperparameters –
• log –
• log_summary –
• log_artifact –
• finish –

Attributes:

• name (str) –

name: str = 'noop' class-attribute instance-attribute

log_hyperparameters(hparams: dict[str, Any])

log(metrics: typing.Mapping[str, Any], *, step, commit: Optional[bool] = None)

log_summary(metrics: dict[str, Any])

log_artifact(artifact_path, *, name: Optional[str] = None, type: Optional[str] = None)

finish()

TensorboardTracker(writer: SummaryWriter)

Bases: Tracker

Methods:

• log_hyperparameters –
• log –
• log_summary –
• log_artifact –
• finish –

Attributes:

• name (str) –
• writer –

name: str = 'tensorboard' class-attribute instance-attribute

writer = writer instance-attribute

log_hyperparameters(hparams: typing.Mapping[str, Any])

log(metrics: typing.Mapping[str, Any], *, step, commit=None)

log_summary(metrics: dict[str, Any])

log_artifact(artifact_path, *, name: Optional[str] = None, type: Optional[str] = None)

finish()

WandbTracker(run: Optional[WandbRun])

Bases: Tracker

Methods:

• log_hyperparameters –
• log –
• log_summary –
• log_artifact –
• finish –

Attributes:

• name (str) –
• run (WandbRun) –

name: str = 'wandb' class-attribute instance-attribute

run: WandbRun instance-attribute

log_hyperparameters(hparams: dict[str, Any])

log(metrics: typing.Mapping[str, Any], *, step, commit=None)

log_summary(metrics: typing.Mapping[str, Any])

log_artifact(artifact_path, *, name: Optional[str] = None, type: Optional[str] = None)

finish()

TrackioTracker(run: Optional[TrackioRun] = None)

Bases: Tracker

Tracker backed by trackio.

Methods:

• log_hyperparameters –
• log –
• log_summary –
• log_artifact –
• finish –

Attributes:

• name (str) –
• run (TrackioRun) –

name: str = 'trackio' class-attribute instance-attribute

run: TrackioRun instance-attribute

log_hyperparameters(hparams: dict[str, Any])

log(metrics: typing.Mapping[str, Any], *, step, commit=None)

log_summary(metrics: typing.Mapping[str, Any])

log_artifact(artifact_path, *, name: Optional[str] = None, type: Optional[str] = None)

finish()

JsonLoggerTracker(logger: Optional[logging.Logger] = None)

Bases: Tracker

Tracker that logs metrics to a Python logger as JSON lines.

Methods:

• log_hyperparameters –
• log –
• log_summary –
• log_artifact –
• finish –

Attributes:

• name (str) –
• logger –

name: str = 'json_logger' class-attribute instance-attribute

logger = logger or logging.getLogger('levanter.json_logger') instance-attribute

log_hyperparameters(hparams: dict[str, Any])

log(metrics: Mapping[str, Any], *, step: Optional[int], commit: Optional[bool] = None)

log_summary(metrics: Mapping[str, Any])

log_artifact(artifact_path, *, name: Optional[str] = None, type: Optional[str] = None)

finish()

Tracker Config

TrackerConfig

Bases: PluginRegistry, ABC

Methods:

• init –
• default_choice_name –

Attributes:

• discover_packages_path –

discover_packages_path = 'levanter.tracker' class-attribute instance-attribute

init(run_id: Optional[str]) -> Tracker abstractmethod

default_choice_name() -> Optional[str] classmethod

NoopConfig() dataclass

Bases: TrackerConfig

Methods:

• init –
• default_choice_name –

Attributes:

• discover_packages_path –

discover_packages_path = 'levanter.tracker' class-attribute instance-attribute

init(run_id: Optional[str]) -> Tracker

default_choice_name() -> Optional[str] classmethod

TensorboardConfig(logdir: str = 'tblogs', comment: Optional[str] = '', purge_step: Optional[int] = None, max_queue: Optional[int] = 10, flush_secs: Optional[int] = 120, filename_suffix: Optional[str] = '', write_to_disk: Optional[bool] = True) dataclass

Bases: TrackerConfig

Methods:

• init –
• default_choice_name –

Attributes:

• logdir (str) –
• comment (Optional[str]) –
• purge_step (Optional[int]) –
• max_queue (Optional[int]) –
• flush_secs (Optional[int]) –
• filename_suffix (Optional[str]) –
• write_to_disk (Optional[bool]) –
• discover_packages_path –

logdir: str = 'tblogs' class-attribute instance-attribute

comment: Optional[str] = '' class-attribute instance-attribute

purge_step: Optional[int] = None class-attribute instance-attribute

max_queue: Optional[int] = 10 class-attribute instance-attribute

flush_secs: Optional[int] = 120 class-attribute instance-attribute

filename_suffix: Optional[str] = '' class-attribute instance-attribute

write_to_disk: Optional[bool] = True class-attribute instance-attribute

discover_packages_path = 'levanter.tracker' class-attribute instance-attribute

init(run_id: Optional[str]) -> TensorboardTracker

default_choice_name() -> Optional[str] classmethod

WandbConfig(entity: Optional[str] = None, project: Optional[str] = 'levanter', name: Optional[str] = None, tags: List[str] = field(default_factory=list), id: Optional[str] = None, group: Optional[str] = None, mode: Optional[str] = None, resume: Optional[Union[bool, str]] = 'allow', save_code: Union[bool, str] = True, save_xla_dumps: bool = False) dataclass

Bases: TrackerConfig

Configuration for wandb.

Methods:

• init –
• default_choice_name –

Attributes:

• entity (Optional[str]) –
• project (Optional[str]) –
• name (Optional[str]) –
• tags (List[str]) –
• id (Optional[str]) –
• group (Optional[str]) –
• mode (Optional[str]) –
• resume (Optional[Union[bool, str]]) –

  Set the resume behavior. Options: "allow", "must", "never", "auto" or None.

• save_code (Union[bool, str]) –

  If string, will save code from that directory. If True, will attempt to sniff out the main directory (since we

• save_xla_dumps (bool) –

  If True, will save the XLA code to wandb (as configured by XLA_FLAGS). This is useful for debugging.

• discover_packages_path –

entity: Optional[str] = None class-attribute instance-attribute

project: Optional[str] = 'levanter' class-attribute instance-attribute

name: Optional[str] = None class-attribute instance-attribute

tags: List[str] = field(default_factory=list) class-attribute instance-attribute

id: Optional[str] = None class-attribute instance-attribute

group: Optional[str] = None class-attribute instance-attribute

mode: Optional[str] = None class-attribute instance-attribute

resume: Optional[Union[bool, str]] = 'allow' class-attribute instance-attribute

Set the resume behavior. Options: "allow", "must", "never", "auto" or None. By default, if the new run has the same ID as a previous run, this run overwrites that data. Please refer to init and resume document for more details.

save_code: Union[bool, str] = True class-attribute instance-attribute

If string, will save code from that directory. If True, will attempt to sniff out the main directory (since we typically don't run from the root of the repo).

save_xla_dumps: bool = False class-attribute instance-attribute

If True, will save the XLA code to wandb (as configured by XLA_FLAGS). This is useful for debugging.

discover_packages_path = 'levanter.tracker' class-attribute instance-attribute

init(run_id: Optional[str]) -> WandbTracker

default_choice_name() -> Optional[str] classmethod

TrackioConfig(project: str = 'levanter', name: Optional[str] = None, space_id: Optional[str] = None, dataset_id: Optional[str] = None, config: dict[str, Any] = field(default_factory=dict), mode: Optional[str] = None, resume: str = 'auto') dataclass

Bases: TrackerConfig

Configuration for Trackio.

Methods:

• init –
• default_choice_name –

Attributes:

• project (str) –
• name (Optional[str]) –
• space_id (Optional[str]) –
• dataset_id (Optional[str]) –
• config (dict[str, Any]) –
• mode (Optional[str]) –

  Controls how Trackio logs.

• resume (str) –
• discover_packages_path –

project: str = 'levanter' class-attribute instance-attribute

name: Optional[str] = None class-attribute instance-attribute

space_id: Optional[str] = None class-attribute instance-attribute

dataset_id: Optional[str] = None class-attribute instance-attribute

config: dict[str, Any] = field(default_factory=dict) class-attribute instance-attribute

mode: Optional[str] = None class-attribute instance-attribute

Controls how Trackio logs.

"online" logs normally, "offline" is synonymous with online since Trackio is local-first, and "disabled" disables logging. If None the default is "online" for the primary process and "disabled" for others.

resume: str = 'auto' class-attribute instance-attribute

discover_packages_path = 'levanter.tracker' class-attribute instance-attribute

init(run_id: Optional[str]) -> Tracker

default_choice_name() -> Optional[str] classmethod

JsonLoggerConfig(logger_name: str = 'levanter.json_logger', level: int = logging.INFO) dataclass

Bases: TrackerConfig

Configuration for :class:JsonLoggerTracker.

Methods:

• init –
• default_choice_name –

Attributes:

• logger_name (str) –
• level (int) –
• discover_packages_path –

logger_name: str = 'levanter.json_logger' class-attribute instance-attribute

level: int = logging.INFO class-attribute instance-attribute

discover_packages_path = 'levanter.tracker' class-attribute instance-attribute

init(run_id: Optional[str]) -> JsonLoggerTracker

default_choice_name() -> Optional[str] classmethod