diff --git a/docs/metrics.md b/docs/metrics.md index 44299f8db..dc3efe8c0 100644 --- a/docs/metrics.md +++ b/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: diff --git a/src/maxdiffusion/tests/dreambooth_trainer_test.py b/src/maxdiffusion/tests/dreambooth_trainer_test.py new file mode 100644 index 000000000..d3548b95c --- /dev/null +++ b/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() 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/tests/stable_diffusion_trainer_test.py b/src/maxdiffusion/tests/stable_diffusion_trainer_test.py new file mode 100644 index 000000000..6d7f83746 --- /dev/null +++ b/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() diff --git a/src/maxdiffusion/train_utils.py b/src/maxdiffusion/train_utils.py index b3cf98e29..b4306043c 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/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): 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..716eab6f3 100644 --- a/src/maxdiffusion/trainers/dreambooth_trainer.py +++ b/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) 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..ee2b579ee 100644 --- a/src/maxdiffusion/trainers/stable_diffusion_trainer.py +++ b/src/maxdiffusion/trainers/stable_diffusion_trainer.py @@ -177,6 +177,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) @@ -217,7 +219,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)