diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py index 2137dd6482..d849932b00 100644 --- a/src/maxtext/common/metric_logger.py +++ b/src/maxtext/common/metric_logger.py @@ -111,6 +111,15 @@ def __init__(self, config, learning_rate_schedule): if self.config.managed_mldiagnostics: ManagedMLDiagnostics(config) # Initialize the MLRun instance. + if self.config.enable_wandb and jax.process_index() == 0: + import wandb # pylint: disable=import-outside-toplevel # lazy import: wandb is an optional dependency + + wandb.init( + project=config.wandb_project_name, + name=config.wandb_run_name, + resume="allow", + ) # Initialize wandb logger. + def reset_eval_metrics(self): """Resets the cumulative metrics dictionary for a new evaluation run.""" self.cumulative_eval_metrics = {"scalar": defaultdict(float)} @@ -133,6 +142,9 @@ def write_metrics(self, metrics, step, metric_type="train"): if self.config.managed_mldiagnostics: self.write_metrics_to_managed_mldiagnostics(metrics, step) + if self.config.enable_wandb and jax.process_index() == 0: + self.write_metrics_to_wandb(metrics, step) + if metric_type == "train": self._maybe_abort_after_write_metrics(metrics) @@ -333,6 +345,18 @@ def write_metrics_to_managed_mldiagnostics(self, metrics, step): mapped_metric_name = _METRICS_TO_MANAGED.get(metric_name, metric_name) mldiag.metrics.record(mapped_metric_name, value, step=int(step)) + def write_metrics_to_wandb(self, metrics, step): + """Write metrics to weights and biases (wandb).""" + import wandb # pylint: disable=import-outside-toplevel # lazy import: wandb is an optional dependency + + flat_metrics = {} + for key, val in metrics.get("scalar", {}).items(): + flat_metrics[key] = float(val) + for key, val in metrics.get("scalars", {}).items(): + for subkey, subval in val.items(): + flat_metrics[f"{key}/{subkey}"] = float(subval) + wandb.log(flat_metrics, step=step) + def write_setup_info_to_tensorboard(self, params): """Writes setup information like train config params, num model params, and XLA flags to TensorBoard.""" num_model_parameters = max_utils.calculate_num_params_from_pytree(params) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 82e448c530..ccaa7f14f6 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -101,6 +101,10 @@ metrics_file: "" # for testing, local file that stores scalar metrics. if empty, # if true save metrics such as loss and tflops to gcs in {base_output_directory}/{run_name}/metrics/ gcs_metrics: false +enable_wandb: False +wandb_project_name: "maxtext" +wandb_run_name: "" + # if true save config to gcs in {base_output_directory}/{run_name}/ save_config_to_gcs: false diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a8c1745..5c59601c1d 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1879,6 +1879,9 @@ class Metrics(BaseModel): False, description="Whether to enable Tunix-managed metrics measurement. The metrics will be uploaded to tensorboard.", ) + enable_wandb: bool = Field(False, description="Enable Weights & Biases logging.") + wandb_project_name: str = Field("maxtext", description="Weights & Biases project name.") + wandb_run_name: str = Field("", description="Weights & Biases run name. If empty, a default name is generated.") class ManagedMLDiagnostics(BaseModel): diff --git a/tests/unit/metric_logger_abort_test.py b/tests/unit/metric_logger_abort_test.py index cd32c25fa7..71845d7c34 100644 --- a/tests/unit/metric_logger_abort_test.py +++ b/tests/unit/metric_logger_abort_test.py @@ -36,6 +36,7 @@ def _make_logger(self, abort_on_nan_loss, abort_on_inf_loss): metrics_file="/tmp/fake_metrics.jsonl", gcs_metrics=True, managed_mldiagnostics=True, + enable_wandb=False, ) return logger @@ -99,6 +100,31 @@ def test_abort_flags_disabled_does_not_exit(self): logger.write_metrics(self._metrics(np.nan), step=1, metric_type="train") +class MetricLoggerWandbTest(unittest.TestCase): + """Tests for MetricLogger wandb metric writing.""" + + def test_write_metrics_to_wandb_flattens_scalars_and_scalars(self): + logger = MetricLogger.__new__(MetricLogger) # skip __init__ + metrics = { + "scalar": {"learning/loss": 1.5}, + "scalars": {"perf": {"step_time": 0.25, "tflops_per_device": 100.0}}, + } + + fake_wandb = mock.MagicMock() + # wandb is lazily imported inside write_metrics_to_wandb, so inject a fake module. + with mock.patch.dict("sys.modules", {"wandb": fake_wandb}): + logger.write_metrics_to_wandb(metrics, step=7) + + fake_wandb.log.assert_called_once_with( + { + "learning/loss": 1.5, + "perf/step_time": 0.25, + "perf/tflops_per_device": 100.0, + }, + step=7, + ) + + class MetricLoggerMetadataTest(unittest.TestCase): """Tests for MetricLogger metadata and setup initialization."""