Wire total_weights into ML Diagnostics metrics

docs/metrics.md

@@ -34,7 +34,7 @@ MaxDiffusion integrates with Google Cloud ML Diagnostics to provide real-time te
 
 ### Predefined Metrics
 
-MaxDiffusion automatically translates internal scalar keys to canonical `MetricType` enums expected by the Control Plane UI:
+MaxDiffusion automatically translates internal scalar keys to canonical metric names expected by the Control Plane UI:
 
 - **Loss** (`loss`): Training loss value per step (mapped from `learning/loss`).
 - **Learning Rate** (`learning_rate`): Current optimizer learning rate (mapped from `learning/current_learning_rate`).
@@ -80,12 +80,16 @@ Inside the trainer's `training_loop()`:
 ```python
 from maxdiffusion import train_utils
 
-# Record standard step metrics (and any custom metrics in train_metric["scalar"]):
+# Optional: Add any custom metrics directly to the scalar dictionary
+train_metric["scalar"]["custom/my_metric"] = my_metric_value
+
+# Record standard step metrics:
 train_utils.record_scalar_metrics(
     train_metric,
     step_time_delta,
     self.per_device_tflops,
     learning_rate_scheduler(step),
+    total_weights=num_model_parameters,
 )
 
 if self.config.write_metrics:

src/maxdiffusion/tests/dreambooth_trainer_test.py

@@ -0,0 +1,92 @@
+"""Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import unittest
+from unittest.mock import MagicMock, patch
+from maxdiffusion.trainers.dreambooth_trainer import DreamboothTrainer
+
+UNET_PARAMS = 1000
+TEXT_ENCODER_PARAMS = 500
+
+
+class MockConfig:
+
+  def __init__(self, **kwargs):
+    for k, v in kwargs.items():
+      setattr(self, k, v)
+
+
+class DreamboothTrainerTest(unittest.TestCase):
+
+  @patch("maxdiffusion.trainers.dreambooth_trainer.train_utils")
+  @patch("maxdiffusion.trainers.dreambooth_trainer.max_utils")
+  @patch("maxdiffusion.trainers.dreambooth_trainer.jax")
+  @patch("maxdiffusion.trainers.dreambooth_trainer.os")
+  def test_training_loop_total_weights(self, mock_os, mock_jax, mock_max_utils, mock_train_utils):
+    """total_weights includes the text encoder only when it is trained."""
+    mock_jax.process_index.return_value = 0
+    mock_jax.random.split.return_value = ("dummy1", "dummy2")
+    mock_os.environ = {"LIBTPU_INIT_ARGS": ""}
+    mock_max_utils.profiler_enabled.return_value = False
+    mock_max_utils.calculate_num_params_from_pytree.side_effect = lambda params: {
+        "unet_params": UNET_PARAMS,
+        "text_encoder_params": TEXT_ENCODER_PARAMS,
+    }[params]
+    mock_train_utils.get_first_step.return_value = 0
+
+    unet_state = MagicMock()
+    unet_state.params = "unet_params"
+    text_encoder_state = MagicMock()
+    text_encoder_state.params = "text_encoder_params"
+    train_states = {"unet_state": unet_state, "text_encoder_state": text_encoder_state}
+
+    p_train_step = MagicMock()
+    p_train_step.return_value = (unet_state, text_encoder_state, {}, "rngs")
+
+    for train_text_encoder, expected_total_weights in (
+        (False, UNET_PARAMS),
+        (True, UNET_PARAMS + TEXT_ENCODER_PARAMS),
+    ):
+      with self.subTest(train_text_encoder=train_text_encoder):
+        mock_train_utils.record_scalar_metrics.reset_mock()
+        config = MockConfig(
+            train_text_encoder=train_text_encoder,
+            max_train_steps=1,
+            per_device_batch_size=1,
+            checkpoint_every=-1,
+            write_metrics=False,
+            metrics_file=None,
+            gcs_metrics=None,
+            skip_first_n_steps_for_profiler=999,
+            profiler_steps=10,
+        )
+
+        with patch("maxdiffusion.trainers.dreambooth_trainer.BaseStableDiffusionTrainer.__init__", return_value=None):
+          trainer = DreamboothTrainer(config)
+        trainer.config = config
+        trainer.total_train_batch_size = 1
+        trainer.per_device_tflops = 1.0
+        trainer.rng = "rng"
+        trainer.checkpoint_manager = MagicMock()
+        trainer.save_checkpoint = MagicMock()
+
+        trainer.training_loop(p_train_step, None, None, train_states, MagicMock(), MagicMock())
+
+        kwargs = mock_train_utils.record_scalar_metrics.call_args.kwargs
+        self.assertEqual(kwargs.get("total_weights"), expected_total_weights)
+
+
+if __name__ == "__main__":
+  unittest.main()

src/maxdiffusion/tests/metrics_test.py

@@ -146,11 +146,13 @@ def test_write_metrics_mld_dispatch_master(self, mock_process_index, mock_mld_me
     mock_mld_metrics.record_metrics.assert_called_once()
     records = mock_mld_metrics.record_metrics.call_args[0][0]
 
-    # Verify records contain translated names and float values
+    # Verify records contain translated string names and float values
     record_dict = {r["metric_name"]: r["value"] for r in records}
-    self.assertAlmostEqual(record_dict[train_utils._METRICS_TO_MANAGED["learning/loss"]], 0.42, places=4)
-    self.assertAlmostEqual(record_dict[train_utils._METRICS_TO_MANAGED["learning/current_learning_rate"]], 0.0001, places=6)
-    self.assertAlmostEqual(record_dict[train_utils._METRICS_TO_MANAGED["learning/total_weights"]], 1000000.0, places=1)
+    self.assertAlmostEqual(record_dict["loss"], 0.42, places=4)
+    self.assertAlmostEqual(record_dict["learning_rate"], 0.0001, places=6)
+    self.assertAlmostEqual(record_dict["total_weights"], 1000000.0, places=1)
+    self.assertAlmostEqual(record_dict["step_time"], 1.0, places=4)
+    self.assertAlmostEqual(record_dict["tflops"], 50.0, places=4)
     self.assertAlmostEqual(record_dict["custom/accuracy"], 0.95, places=4)
 
   @patch("maxdiffusion.train_utils.mld_metrics")

src/maxdiffusion/tests/stable_diffusion_trainer_test.py

@@ -0,0 +1,140 @@
+"""Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import unittest
+from unittest.mock import MagicMock, patch
+from maxdiffusion.trainers.stable_diffusion_trainer import StableDiffusionTrainer
+
+
+class MockConfig:
+
+  def __init__(self, **kwargs):
+    for k, v in kwargs.items():
+      setattr(self, k, v)
+
+
+class StableDiffusionTrainerTest(unittest.TestCase):
+
+  @patch("maxdiffusion.trainers.stable_diffusion_trainer.train_utils")
+  @patch("maxdiffusion.trainers.stable_diffusion_trainer.max_utils")
+  @patch("maxdiffusion.trainers.stable_diffusion_trainer.jax")
+  @patch("maxdiffusion.trainers.stable_diffusion_trainer.os")
+  def test_training_loop_total_weights(self, mock_os, mock_jax, mock_max_utils, mock_train_utils):
+    # Setup mocks
+    mock_jax.process_index.return_value = 0
+    mock_jax.random.split.return_value = ("dummy1", "dummy2")
+    mock_os.environ = {"LIBTPU_INIT_ARGS": ""}
+
+    mock_max_utils.profiler_enabled.return_value = False
+
+    def fake_calc_params(pytree):
+      if pytree == "unet_params":
+        return 1000
+      elif pytree == "text_encoder_params":
+        return 500
+      return 0
+
+    mock_max_utils.calculate_num_params_from_pytree.side_effect = fake_calc_params
+
+    # We want the loop to hit exactly 1 step then exit.
+    mock_train_utils.get_first_step.return_value = 0
+
+    unet_state = MagicMock()
+    unet_state.params = "unet_params"
+
+    vae_state = MagicMock()
+
+    text_encoder_state = MagicMock()
+    text_encoder_state.params = "text_encoder_params"
+
+    train_states = {
+        "unet_state": unet_state,
+        "vae_state": vae_state,
+        "text_encoder_state": text_encoder_state,
+    }
+
+    p_train_step = MagicMock()
+    # p_train_step returns: unet_state, text_encoder_state, train_metric, train_rngs
+    p_train_step.return_value = (unet_state, text_encoder_state, {}, "rngs")
+
+    data_iterator = MagicMock()
+    lr_scheduler = MagicMock()
+    lr_scheduler.return_value = 0.001
+
+    # Instance of trainer
+    with patch("maxdiffusion.trainers.stable_diffusion_trainer.BaseStableDiffusionTrainer.__init__") as mock_init:
+      mock_init.return_value = None
+
+      # Test train_text_encoder = False
+      config_false = MockConfig(
+          train_text_encoder=False,
+          max_train_steps=1,
+          per_device_batch_size=1,
+          checkpoint_every=-1,
+          write_metrics=False,
+          metrics_file=None,
+          gcs_metrics=None,
+          skip_first_n_steps_for_profiler=999,
+          profiler_steps=10,
+      )
+      trainer = StableDiffusionTrainer(config_false)
+      trainer.config = config_false
+      trainer.total_train_batch_size = 1
+      trainer.per_device_tflops = 1.0
+      trainer.rng = "rng"
+      trainer.checkpoint_manager = MagicMock()
+      trainer.checkpoint_manager.reached_preemption.return_value = False
+      trainer.save_checkpoint = MagicMock()
+
+      trainer.training_loop(p_train_step, None, None, train_states, data_iterator, lr_scheduler)
+
+      # Verify total_weights recorded
+      mock_train_utils.record_scalar_metrics.assert_called()
+      kwargs = mock_train_utils.record_scalar_metrics.call_args.kwargs
+      self.assertEqual(kwargs.get("total_weights"), 1000)
+
+      # Reset mocks
+      mock_train_utils.record_scalar_metrics.reset_mock()
+
+      # Test train_text_encoder = True
+      config_true = MockConfig(
+          train_text_encoder=True,
+          max_train_steps=1,
+          per_device_batch_size=1,
+          checkpoint_every=-1,
+          write_metrics=False,
+          metrics_file=None,
+          gcs_metrics=None,
+          skip_first_n_steps_for_profiler=999,
+          profiler_steps=10,
+      )
+      trainer_true = StableDiffusionTrainer(config_true)
+      trainer_true.config = config_true
+      trainer_true.total_train_batch_size = 1
+      trainer_true.per_device_tflops = 1.0
+      trainer_true.rng = "rng"
+      trainer_true.checkpoint_manager = MagicMock()
+      trainer_true.checkpoint_manager.reached_preemption.return_value = False
+      trainer_true.save_checkpoint = MagicMock()
+
+      trainer_true.training_loop(p_train_step, None, None, train_states, data_iterator, lr_scheduler)
+
+      mock_train_utils.record_scalar_metrics.assert_called()
+      kwargs = mock_train_utils.record_scalar_metrics.call_args.kwargs
+      self.assertEqual(kwargs.get("total_weights"), 1500)
+
+
+if __name__ == "__main__":
+  unittest.main()

src/maxdiffusion/train_utils.py

@@ -77,30 +77,18 @@ def _validate_gcs_bucket_name(bucket_name, config_var):
 
 
 try:
-  from google_cloud_mldiagnostics import metrics as mld_metrics, metric_types
+  from google_cloud_mldiagnostics import metrics as mld_metrics
 except ImportError:
   mld_metrics = None
-  metric_types = None
-
-
-if metric_types is not None:
-  _METRICS_TO_MANAGED = {
-      "learning/loss": metric_types.MetricType.LOSS,
-      "learning/current_learning_rate": metric_types.MetricType.LEARNING_RATE,
-      "learning/grad_norm": metric_types.MetricType.GRADIENT_NORM,
-      "learning/total_weights": metric_types.MetricType.TOTAL_WEIGHTS,
-      "perf/step_time_seconds": metric_types.MetricType.STEP_TIME,
-      "perf/per_device_tflops_per_sec": metric_types.MetricType.TFLOPS,
-  }
-else:
-  _METRICS_TO_MANAGED = {
-      "learning/loss": "loss",
-      "learning/current_learning_rate": "learning_rate",
-      "learning/grad_norm": "gradient_norm",
-      "learning/total_weights": "total_weights",
-      "perf/step_time_seconds": "step_time",
-      "perf/per_device_tflops_per_sec": "tflops",
-  }
+
+_METRICS_TO_MANAGED = {
+    "learning/loss": "loss",
+    "learning/current_learning_rate": "learning_rate",
+    "learning/grad_norm": "gradient_norm",
+    "learning/total_weights": "total_weights",
+    "perf/step_time_seconds": "step_time",
+    "perf/per_device_tflops_per_sec": "tflops",
+}
 
 
 def record_scalar_metrics(metrics, step_time_delta, per_device_tflops, lr, total_weights=None):

src/maxdiffusion/trainers/base_wan_trainer.py

@@ -340,7 +340,11 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
             self._profiler.stop()
 
         train_utils.record_scalar_metrics(
-            train_metric, last_step_completion - start_step_time, per_device_tflops, learning_rate_scheduler(step)
+            train_metric,
+            last_step_completion - start_step_time,
+            per_device_tflops,
+            learning_rate_scheduler(step),
+            total_weights=num_model_parameters,
         )
         if self.config.write_metrics:
           train_utils.write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)

src/maxdiffusion/trainers/dreambooth_trainer.py

@@ -188,6 +188,8 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
     text_encoder_state = train_states["text_encoder_state"]
 
     num_model_parameters = max_utils.calculate_num_params_from_pytree(unet_state.params)
+    if self.config.train_text_encoder:
+      num_model_parameters += max_utils.calculate_num_params_from_pytree(text_encoder_state.params)
     max_utils.add_text_to_summary_writer("number_model_parameters", str(num_model_parameters), writer)
     max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
     max_utils.add_config_to_summary_writer(self.config, writer)
@@ -222,7 +224,11 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
       new_time = datetime.datetime.now()
 
       train_utils.record_scalar_metrics(
-          train_metric, new_time - last_step_completion, self.per_device_tflops, learning_rate_scheduler(step)
+          train_metric,
+          new_time - last_step_completion,
+          self.per_device_tflops,
+          learning_rate_scheduler(step),
+          total_weights=num_model_parameters,
       )
       if self.config.write_metrics:
         train_utils.write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)

src/maxdiffusion/trainers/flux_trainer.py

@@ -441,7 +441,11 @@ def training_loop(
       new_time = datetime.datetime.now()
 
       record_scalar_metrics(
-          train_metric, new_time - last_step_completion, self.per_device_tflops, unet_learning_rate_scheduler(step)
+          train_metric,
+          new_time - last_step_completion,
+          self.per_device_tflops,
+          unet_learning_rate_scheduler(step),
+          total_weights=num_model_parameters,
       )
       if self.config.write_metrics:
         write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)

src/maxdiffusion/trainers/sdxl_trainer.py

@@ -259,7 +259,11 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
         difference_in_ms = time_difference.total_seconds() * 1000
         max_logging.log(f"Step time {difference_in_ms}ms")
         record_scalar_metrics(
-            train_metric, last_step_completion - start_step_time, self.per_device_tflops, unet_learning_rate_scheduler(step)
+            train_metric,
+            last_step_completion - start_step_time,
+            self.per_device_tflops,
+            unet_learning_rate_scheduler(step),
+            total_weights=num_model_parameters,
         )
         if self.config.write_metrics:
           write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)