Add wandb support

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)

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
 

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):

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."""