Logging and experiment tracking

Logging and experiment tracking#

This guide covers how Primus emits training metrics and logs, and how to wire up the supported experiment trackers—TensorBoard, Weights & Biases (WandB), and MLflow (including Databricks-hosted MLflow)—across the Megatron and TorchTitan backends. Parameters are grounded in primus/configs/modules/megatron/trainer_base.yaml, primus_megatron_module.yaml, and primus/configs/modules/torchtitan/pre_trainer.yaml.

For an operations-oriented overview, see Monitoring and logging. For required credentials/keys, see Environment variables.

1. Tracker toggles at a glance (Megatron)#

All three trackers are opt-in and disabled by default (primus/configs/modules/megatron/primus_megatron_module.yaml):

disable_tensorboard: true
disable_wandb: true
disable_mlflow: true

Set the relevant disable_* to false to enable a tracker. Primus performs validation checks at startup—e.g. it warns if WandB is enabled but WANDB_API_KEY is unset (primus/backends/megatron/patches/args/wandb_config_patches.py). MLflow logging is initialized in primus/backends/megatron/training/global_vars.py; Databricks-hosted MLflow additionally requires DATABRICKS_HOST (read by the mlflow client).

2. Console and step logging (Megatron)#

Core logging cadence and content (trainer_base.yaml, overridden by pre_trainer.yaml):

Parameter	Default	Purpose
`log_interval`	`100` (`pre_trainer.yaml` sets `1`)	Steps between log lines.
`log_throughput`	`false` (`pre_trainer.yaml` sets `true`)	Log tokens/s and TFLOP/s throughput.
`log_progress`	`false`	Log progress/ETA.
`log_params_norm`	`false`	Log parameter L2 norm.
`log_num_zeros_in_grad`	`false`	Log gradient sparsity.
`log_avg_skip_iterations`	`2`	Warmup iterations excluded from averages.
`log_avg_reset_interval`	`10`	Reset window for running averages.
`timing_log_level`	`0`	Verbosity of timer breakdowns.
`timing_log_option`	`minmax`	Timer aggregation across ranks.
`logging_level`	`null`	Python logging level override.

3. TensorBoard (Megatron)#

Enable with disable_tensorboard: false and set an output directory:

Parameter	Default	Purpose
`tensorboard_dir`	`null`	Output directory for event files (required when enabled).
`tensorboard_log_interval`	`1`	Steps between TensorBoard writes.
`tensorboard_queue_size`	`1000`	Event queue size before flush.
`log_learning_rate_to_tensorboard`	`true`	Log LR.
`log_loss_scale_to_tensorboard`	`true`	Log loss scale (mixed precision).
`log_timers_to_tensorboard`	`false`	Log per-stage timers.
`log_batch_size_to_tensorboard`	`false`	Log batch size.
`log_memory_to_tensorboard`	`false`	Log GPU memory.
`log_world_size_to_tensorboard`	`false`	Log world size.
`log_validation_ppl_to_tensorboard`	`false`	Log validation perplexity.

4. Weights and biases (Megatron)#

Enable with disable_wandb: false. Configuration:

Parameter	Default	Purpose
`wandb_project`	`null`	WandB project name.
`wandb_exp_name`	`null`	Run/experiment name.
`wandb_save_dir`	`null`	Local directory for WandB files.
`wandb_entity`	`null`	Team/entity.

Credentials (see Environment variables):

export WANDB_API_KEY=...     # required when WandB is enabled
export WANDB_PROJECT=...     # optional
export WANDB_RUN_NAME=...    # optional
export WANDB_TEAM=...        # optional (entity)

WANDB_API_KEY is on the container passthrough allowlist (runner/.primus.yaml), so it propagates into the training container.

5. MLflow (Megatron)#

Enable with disable_mlflow: false. Run identification and upload behavior:

Parameter	Default	Purpose
`mlflow_run_name`	`null`	MLflow run name.
`mlflow_experiment_name`	`null`	MLflow experiment name.
`mlflow_upload_traces`	`false`	Upload profiler trace files.
`mlflow_upload_logs`	`false`	Upload training log files.
`mlflow_upload_performance_metrics`	`false`	Upload the comprehensive perf/memory/utilization metric set (implicitly enables throughput calc).
`mlflow_upload_tracelens_report`	`false`	Generate + upload TraceLens reports (see Profiling & observability).

Credentials and endpoints:

export MLFLOW_TRACKING_URI=...     # tracking server URI
export MLFLOW_REGISTRY_URI=...     # optional model registry
# Databricks-hosted MLflow:
export DATABRICKS_HOST=...         # checked at startup when MLflow is enabled
export DATABRICKS_TOKEN=...

The mlflow_upload_* flags are designed so MLflow stays opt-in: they only take effect when disable_mlflow: false.

6. One-logger (Megatron)#

NVIDIA One-Logger telemetry is enabled by default in trainer_base.yaml:

Parameter	Default	Purpose
`enable_one_logger`	`true`	Enable One-Logger collection.
`one_logger_project`	`megatron-lm`	Project tag.
`one_logger_run_name`	`null`	Run name.
`one_logger_async`	`false`	Async upload.
`app_tag_run_name` / `app_tag_run_version`	`null` / `0.0.0`	Application tags.

7. Metrics and logging (TorchTitan)#

Configured under metrics: in primus/configs/modules/torchtitan/pre_trainer.yaml:

Parameter	Default	Purpose
`enable_tensorboard`	`false`	Enable TensorBoard logging.
`enable_wandb`	`false`	Enable WandB logging.
`log_freq`	`10`	Steps between metric logs.
`save_tb_folder`	`tb`	TensorBoard output subfolder.
`save_for_all_ranks`	`false`	Write metrics from every rank (default: rank 0 only).
`disable_color_printing`	`false`	Disable ANSI colors in console output.

TorchTitan reads WandB settings from the environment (WANDB_PROJECT, WANDB_RUN_NAME, WANDB_TEAM) via primus/backends/torchtitan/patches/wandb_patches.py and third_party/torchtitan/torchtitan/components/metrics.py.

8. Logging (MaxText)#

MaxText logging cadence is controlled by log_period (primus/configs/modules/maxtext/pre_trainer.yaml, default 100). See MaxText parameters.

9. Recommended setup#

Local-only: enable TensorBoard (disable_tensorboard: false, set tensorboard_dir)—no credentials required.
Team tracking: enable WandB (disable_wandb: false) + export WANDB_API_KEY and wandb_project/wandb_entity.
Enterprise / scaling studies: enable MLflow (disable_mlflow: false) + MLFLOW_TRACKING_URI (or Databricks host/token), and turn on mlflow_upload_performance_metrics for throughput/memory/utilization dashboards.
Keep WANDB_API_KEY and tokens out of YAML—pass them as environment variables (allowlisted for container passthrough). See Security.