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/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/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

@@ -222,7 +222,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)

src/maxdiffusion/trainers/stable_diffusion_trainer.py

@@ -217,7 +217,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, 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:
         train_utils.write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config)