diff --git a/docs/metrics.md b/docs/metrics.md index 44299f8db..3cccd6d11 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -34,11 +34,10 @@ 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`). -- **Gradient Norm** (`gradient_norm`): Global L2 norm of model gradients (mapped from `learning/grad_norm`). - **Total Weights** (`total_weights`): Total trainable model parameter count (mapped from `learning/total_weights`). - **Step Time** (`step_time`): Duration of each training step in seconds (mapped from `perf/step_time_seconds`). - **TFLOPS** (`tflops`): Hardware compute throughput per accelerator in TFLOP/s (mapped from `perf/per_device_tflops_per_sec`). @@ -86,6 +85,7 @@ train_utils.record_scalar_metrics( step_time_delta, self.per_device_tflops, learning_rate_scheduler(step), + total_weights=num_model_parameters, ) if self.config.write_metrics: diff --git a/src/maxdiffusion/tests/metrics_test.py b/src/maxdiffusion/tests/metrics_test.py index e87d06afe..a4fcd4617 100644 --- a/src/maxdiffusion/tests/metrics_test.py +++ b/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") diff --git a/src/maxdiffusion/train_utils.py b/src/maxdiffusion/train_utils.py index b3cf98e29..a558841a8 100644 --- a/src/maxdiffusion/train_utils.py +++ b/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/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): diff --git a/src/maxdiffusion/trainers/base_wan_trainer.py b/src/maxdiffusion/trainers/base_wan_trainer.py index 1d2c543dc..63a3e8f1f 100644 --- a/src/maxdiffusion/trainers/base_wan_trainer.py +++ b/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) diff --git a/src/maxdiffusion/trainers/dreambooth_trainer.py b/src/maxdiffusion/trainers/dreambooth_trainer.py index 07e5a26b6..259540f96 100644 --- a/src/maxdiffusion/trainers/dreambooth_trainer.py +++ b/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) diff --git a/src/maxdiffusion/trainers/flux_trainer.py b/src/maxdiffusion/trainers/flux_trainer.py index 1c20aec1a..747b62930 100644 --- a/src/maxdiffusion/trainers/flux_trainer.py +++ b/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) diff --git a/src/maxdiffusion/trainers/sdxl_trainer.py b/src/maxdiffusion/trainers/sdxl_trainer.py index 932d9be84..e625d6268 100644 --- a/src/maxdiffusion/trainers/sdxl_trainer.py +++ b/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) diff --git a/src/maxdiffusion/trainers/stable_diffusion_trainer.py b/src/maxdiffusion/trainers/stable_diffusion_trainer.py index cda485de7..5d237a886 100644 --- a/src/maxdiffusion/trainers/stable_diffusion_trainer.py +++ b/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)