From 6b9e0fd93bd08c6f0290993090d0606c70efe5ef Mon Sep 17 00:00:00 2001 From: Richa Gupta Date: Thu, 13 Aug 2026 15:41:04 +0000 Subject: [PATCH] Add Google Cloud ML Diagnostics metrics support and developer guide. Integrate Google Cloud ML Diagnostics SDK (google-cloud-mldiagnostics) into MaxDiffusion to automatically record training, system, and performance metrics. Key Changes: - train_utils.py: Added _METRICS_TO_MANAGED mapping table converting MaxDiffusion keys to canonical MetricType enums (loss, learning_rate, gradient_norm, total_weights, step_time, tflops) with automatic pass-through for custom metrics. Added batch metric logging to write_metrics(). - max_utils.py: Added _clean_config_dict() to sanitize non-JSON serializable hyperparameters, configured region=None for GCP cluster auto-discovery, and enabled background hardware metric collection (log_system_metrics=True). - docs/metrics.md: Created comprehensive integration, architecture, and verification guide for developers adding new model trainers. Tested: - Ran a 500-step multi-host distributed training run on TPU v6e cluster richa-maxdiffusion-test (JobSet richa-metrics-test-v11). Verified successful metric ingestion in Cloud Logging (ml_diagnostics_metric) for predefined, custom, and hardware utilization metrics. --- docs/metrics.md | 160 ++++++++++++++++++++++ src/maxdiffusion/max_utils.py | 32 ++++- src/maxdiffusion/tests/metrics_test.py | 175 ++++++++++++++++++++++++ src/maxdiffusion/tests/profiler_test.py | 10 +- src/maxdiffusion/train_utils.py | 43 +++++- 5 files changed, 407 insertions(+), 13 deletions(-) create mode 100644 docs/metrics.md create mode 100644 src/maxdiffusion/tests/metrics_test.py diff --git a/docs/metrics.md b/docs/metrics.md new file mode 100644 index 000000000..dc6de848b --- /dev/null +++ b/docs/metrics.md @@ -0,0 +1,160 @@ + + +# Metrics Collection and Monitoring with Google Cloud ML Diagnostics + +This guide describes how to capture, monitor, and visualize training, system, and performance metrics in **MaxDiffusion** using the **Google Cloud ML Diagnostics SDK** (`google-cloud-mldiagnostics`). + +--- + +## 1. Overview + +MaxDiffusion integrates with Google Cloud ML Diagnostics to provide real-time telemetry during training runs on TPU and GPU accelerators: +- **Workload Metrics **: In multi-host JAX jobs, step-level metrics (loss, step time, learning rate, gradient norm, parameter weights, custom activations) are buffered and dispatched from master node to prevent duplicate logs. +- **System & Accelerator Metrics **: The SDK automatically runs background daemon threads on all worker hosts to capture hardware utilization (`tpu_duty_cycle`, `hbm_utilization`, `host_cpu_utilization`, `host_memory_utilization`). +- **Cloud Logging Sink**: Metrics are written to Google Cloud Logging (`projects//logs/ml_diagnostics_metric`) +- **Control Plane UI**: The Diagnostics Console renders standard metric plots + +--- + +## 2. Metric Types + +### Predefined Metrics + +MaxDiffusion automatically translates internal scalar keys to canonical `MetricType` enums 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`). + +### Custom Metrics + +Any key in `metrics["scalar"]` that is not part of `_METRICS_TO_MANAGED` is treated as a **Custom Metric**: +- Retains its raw string name (e.g., `"custom/latents_mean"`, `"snr_loss_weight"`, `"cross_attn_entropy"`). +- Are dynamically discovered by the Control Plane UI and rendered in dedicated chart cards (`Over Time` and `Over Steps`). + +### Automated System & Accelerator Metrics + +When `enable_ml_diagnostics=True` and `log_system_metrics=True` are set, background daemon threads automatically emit: +- `tpu_duty_cycle` / `gpu_utilization`: Core accelerator compute utilization percentage. +- `hbm_utilization`: High Bandwidth Memory consumed percentage. +- `host_cpu_utilization`: Host CPU usage percentage. +- `host_memory_utilization`: Host system RAM usage percentage. + +--- + +## 3. Integration Guide for Training Scripts + +Metric mapping and dispatch are centralized in `train_utils.py` and `max_utils.py`. Authors of new training scripts can integrate metrics using two steps: + +### Step 1: Initialize MachineLearningRun + +Initialize the run at the start of training: + +```python +from maxdiffusion import max_utils + +# Automatically cleans config and discovers cluster region: +max_utils.ensure_machinelearning_job_runs(config) +``` + +### Step 2: Record Scalar Metrics in the Training Loop + +Inside the trainer's `training_loop()`: + +```python +from maxdiffusion import max_utils, train_utils + +# Calculate total model parameters: +num_model_parameters = max_utils.calculate_num_params_from_pytree(unet_state.params) + +# Record standard step metrics (and any custom metrics in train_metric["scalar"]): +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: + train_utils.write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config) +``` + +--- + +## 4. Configuration + +Enable ML Diagnostics via YAML configuration files or command-line flags: + +```yaml +# configs/base_2_base.yml +run_name: "my-training-run" +enable_ml_diagnostics: True +write_metrics: True +log_period: 10 +profiler_gcs_path: "gs://my-bucket/profiler" +enable_ondemand_xprof: True +``` + +Run command: + +```bash +python train.py configs/base_2_base.yml \ + run_name=my-training-run \ + output_dir=gs://my-bucket/output \ + enable_ml_diagnostics=True \ + write_metrics=True \ + profiler_gcs_path=gs://my-bucket/profiler \ + enable_ondemand_xprof=True +``` + +--- + +## 5. Verification + +### Google Cloud Logging + +Inspect metric logs directly using `gcloud`: + +```bash +# Query loss metrics +gcloud logging read 'logName="projects//logs/ml_diagnostics_metric" AND resource.labels.namespace="loss"' \ + --limit=5 \ + --format="json" + +# Query custom metrics +gcloud logging read 'logName="projects//logs/ml_diagnostics_metric" AND resource.labels.namespace="custom/latents_mean"' \ + --limit=5 \ + --format="json" + +# Query hardware metrics +gcloud logging read 'logName="projects//logs/ml_diagnostics_metric" AND resource.labels.namespace="hbm_utilization"' \ + --limit=5 \ + --format="json" +``` + +### Google Cloud Console + +1. Open Google Cloud Console and navigate to **Hypercompute Clusters** $\rightarrow$ **Diagnostics**. +2. Select your cluster and active `MachineLearningRun`. +3. Inspect: + - **Model Metrics**: View predefined plots for `loss`, `learning_rate`, `gradient_norm`, and `total_weights`. + - **Custom Metrics**: View dynamically generated charts for all `custom/*` metrics over time and steps. + - **Performance**: View `step_time`, `tflops`, `tpu_duty_cycle`, and `hbm_utilization`. diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index b7ee9a5d1..cb15df879 100644 --- a/src/maxdiffusion/max_utils.py +++ b/src/maxdiffusion/max_utils.py @@ -88,28 +88,50 @@ def _jax_profiler_enabled(config): return "enable_profiler" in config.get_keys() and config.enable_profiler and jax.process_index() == 0 -def _ml_diagnostics_profiler_enabled(config): +def ml_diagnostics_enabled(config): return "enable_ml_diagnostics" in config.get_keys() and config.enable_ml_diagnostics +def _clean_config_dict(config): + """Filter out non-JSON serializable keys from hyperparameter config.""" + cleaned = {} + keys_dict = config.get_keys() if hasattr(config, "get_keys") else config + if isinstance(keys_dict, dict): + for k, v in keys_dict.items(): + try: + json.dumps(v, allow_nan=False) + cleaned[k] = v + except (TypeError, ValueError, OverflowError): + continue + return cleaned + + def profiler_enabled(config): - return _jax_profiler_enabled(config) or _ml_diagnostics_profiler_enabled(config) + return _jax_profiler_enabled(config) or ml_diagnostics_enabled(config) def ensure_machinelearning_job_runs(config): """Ensures that a MachineLearningJobRun is active, and if not creates one.""" global _ml_run - if _ml_run is not None or not _ml_diagnostics_profiler_enabled(config) or machinelearning_run is None: + if _ml_run is not None or not ml_diagnostics_enabled(config): return + if machinelearning_run is None: + raise ImportError( + "enable_ml_diagnostics is True, but google_cloud_mldiagnostics is not installed. " + "Please install it via 'pip install google-cloud-mldiagnostics'." + ) + logging.getLogger("google_cloud_mldiagnostics").setLevel(logging.WARNING) _ml_run = machinelearning_run( name=config.run_name, gcs_path=config.profiler_gcs_path, - configs=config.get_keys(), + configs=_clean_config_dict(config), on_demand_xprof=config.enable_ondemand_xprof, + log_system_metrics=True, + region=None, ) @@ -128,7 +150,7 @@ def __init__(self, config, session_name=None): self._active = None # "mld" | "jax" | None def start(self): - use_mld = _ml_diagnostics_profiler_enabled(self.config) and xprof is not None + use_mld = ml_diagnostics_enabled(self.config) use_jax = _jax_profiler_enabled(self.config) if use_mld and use_jax: diff --git a/src/maxdiffusion/tests/metrics_test.py b/src/maxdiffusion/tests/metrics_test.py new file mode 100644 index 000000000..e87d06afe --- /dev/null +++ b/src/maxdiffusion/tests/metrics_test.py @@ -0,0 +1,175 @@ +""" +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 datetime +import unittest +from unittest.mock import MagicMock, patch +import numpy as np +from maxdiffusion import max_utils, train_utils + + +class MockConfig: + + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + def get_keys(self): + return self.__dict__ + + +class MetricsTest(unittest.TestCase): + + def setUp(self): + max_utils._ml_run = None + train_utils._buffered_step = None + train_utils._buffered_metrics = None + + def test_ml_diagnostics_enabled(self): + config_enabled = MockConfig(enable_ml_diagnostics=True) + config_disabled = MockConfig(enable_ml_diagnostics=False) + config_missing = MockConfig() + + self.assertTrue(max_utils.ml_diagnostics_enabled(config_enabled)) + self.assertFalse(max_utils.ml_diagnostics_enabled(config_disabled)) + self.assertFalse(max_utils.ml_diagnostics_enabled(config_missing)) + + def test_clean_config_dict(self): + config = MockConfig( + run_name="test_run", + learning_rate=0.001, + infinity_val=float("inf"), + nan_val=float("nan"), + batch_size=16, + ) + cleaned = max_utils._clean_config_dict(config) + self.assertEqual(cleaned["run_name"], "test_run") + self.assertEqual(cleaned["learning_rate"], 0.001) + self.assertEqual(cleaned["batch_size"], 16) + self.assertNotIn("infinity_val", cleaned) + self.assertNotIn("nan_val", cleaned) + + @patch("maxdiffusion.max_utils.machinelearning_run", None) + def test_ensure_machinelearning_job_runs_raises_import_error(self): + config = MockConfig(enable_ml_diagnostics=True) + with self.assertRaises(ImportError) as ctx: + max_utils.ensure_machinelearning_job_runs(config) + self.assertIn("enable_ml_diagnostics is True", str(ctx.exception)) + + @patch("maxdiffusion.max_utils.machinelearning_run") + def test_ensure_machinelearning_job_runs_success(self, mock_ml_run): + config = MockConfig( + enable_ml_diagnostics=True, + run_name="my_run", + profiler_gcs_path="gs://my-bucket/profiler", + enable_ondemand_xprof=True, + ) + max_utils.ensure_machinelearning_job_runs(config) + mock_ml_run.assert_called_once_with( + name="my_run", + gcs_path="gs://my-bucket/profiler", + configs=max_utils._clean_config_dict(config), + on_demand_xprof=True, + log_system_metrics=True, + region=None, + ) + + def test_record_scalar_metrics(self): + metrics = {"scalar": {}} + step_delta = datetime.timedelta(seconds=2.5) + train_utils.record_scalar_metrics( + metrics=metrics, + step_time_delta=step_delta, + per_device_tflops=100.0, + lr=0.0001, + total_weights=1500000000, + ) + scalars = metrics["scalar"] + self.assertEqual(scalars["perf/step_time_seconds"], 2.5) + self.assertEqual(scalars["perf/per_device_tflops"], 100.0) + self.assertEqual(scalars["perf/per_device_tflops_per_sec"], 40.0) + self.assertEqual(scalars["learning/current_learning_rate"], 0.0001) + self.assertEqual(scalars["learning/total_weights"], 1500000000.0) + + @patch("maxdiffusion.train_utils.mld_metrics") + @patch("jax.process_index", return_value=0) + def test_write_metrics_mld_dispatch_master(self, mock_process_index, mock_mld_metrics): + config = MockConfig( + enable_ml_diagnostics=True, + metrics_file=False, + gcs_metrics=False, + log_period=10, + tensorboard_dir="/tmp/tensorboard", + ) + mock_writer = MagicMock() + + # Step 0 (buffers metric) + metrics_step_0 = { + "scalar": { + "learning/loss": np.array(0.42), + "custom/accuracy": 0.95, + } + } + train_utils.record_scalar_metrics( + metrics=metrics_step_0, + step_time_delta=datetime.timedelta(seconds=1.0), + per_device_tflops=50.0, + lr=0.0001, + total_weights=1000000, + ) + train_utils.write_metrics(mock_writer, None, None, metrics_step_0, 0, config) + mock_mld_metrics.record_metrics.assert_not_called() + + # Step 1 (flushes buffered step 0 metrics) + metrics_step_1 = {"scalar": {"learning/loss": np.array(0.38)}} + train_utils.record_scalar_metrics( + metrics=metrics_step_1, + step_time_delta=datetime.timedelta(seconds=1.0), + per_device_tflops=50.0, + lr=0.0001, + ) + train_utils.write_metrics(mock_writer, None, None, metrics_step_1, 1, config) + + 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 + 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["custom/accuracy"], 0.95, places=4) + + @patch("maxdiffusion.train_utils.mld_metrics") + @patch("jax.process_index", return_value=1) + def test_write_metrics_mld_skipped_on_worker(self, mock_process_index, mock_mld_metrics): + config = MockConfig( + enable_ml_diagnostics=True, + metrics_file=False, + gcs_metrics=False, + log_period=10, + ) + mock_writer = MagicMock() + + metrics_0 = {"scalar": {"learning/loss": 0.5}} + train_utils.write_metrics(mock_writer, None, None, metrics_0, 0, config) + train_utils.write_metrics(mock_writer, None, None, metrics_0, 1, config) + + mock_mld_metrics.record_metrics.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/profiler_test.py b/src/maxdiffusion/tests/profiler_test.py index 763be4ca4..ff1b48cf1 100644 --- a/src/maxdiffusion/tests/profiler_test.py +++ b/src/maxdiffusion/tests/profiler_test.py @@ -66,9 +66,7 @@ def test_ml_diagnostics_profiler(self, mock_process_index, mock_xprof, mock_ml_r @patch("maxdiffusion.max_utils.machinelearning_run") @patch("maxdiffusion.max_utils.xprof") @patch("jax.process_index", return_value=1) - def test_ml_diagnostics_profiler_non_master_host( - self, mock_process_index, mock_xprof, mock_ml_run - ): + def test_ml_diagnostics_profiler_non_master_host(self, mock_process_index, mock_xprof, mock_ml_run): """Tests that ML Diagnostics profiler is also enabled on non-master hosts (process_index != 0).""" config = MockConfig( enable_ml_diagnostics=True, @@ -79,7 +77,7 @@ def test_ml_diagnostics_profiler_non_master_host( tensorboard_dir="/tmp/fake_tensorboard", ) - self.assertTrue(max_utils._ml_diagnostics_profiler_enabled(config)) + self.assertTrue(max_utils.ml_diagnostics_enabled(config)) with max_utils.Profiler(config, session_name="test_session"): mock_xprof.return_value.start.assert_called_once_with("test_session") @@ -90,9 +88,7 @@ def test_ml_diagnostics_profiler_non_master_host( @patch("maxdiffusion.max_utils.machinelearning_run") @patch("maxdiffusion.max_utils.xprof") @patch("jax.process_index", return_value=0) - def test_both_profilers_enabled_prioritizes_mld( - self, mock_process_index, mock_xprof, mock_ml_run, mock_start_trace - ): + def test_both_profilers_enabled_prioritizes_mld(self, mock_process_index, mock_xprof, mock_ml_run, mock_start_trace): """Tests that when both ML Diagnostics and JAX profiler are enabled, ML Diagnostics is prioritized and JAX profiler is skipped.""" config = MockConfig( enable_ml_diagnostics=True, diff --git a/src/maxdiffusion/train_utils.py b/src/maxdiffusion/train_utils.py index e82174e92..87e1634bd 100644 --- a/src/maxdiffusion/train_utils.py +++ b/src/maxdiffusion/train_utils.py @@ -76,12 +76,34 @@ def _validate_gcs_bucket_name(bucket_name, config_var): ) -def record_scalar_metrics(metrics, step_time_delta, per_device_tflops, lr): +try: + from google_cloud_mldiagnostics import metrics as mld_metrics, metric_types +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 = {} + + +def record_scalar_metrics(metrics, step_time_delta, per_device_tflops, lr, total_weights=None): """Records scalar metrics to be written to tensorboard""" metrics["scalar"].update({"perf/step_time_seconds": step_time_delta.total_seconds()}) metrics["scalar"].update({"perf/per_device_tflops": per_device_tflops}) metrics["scalar"].update({"perf/per_device_tflops_per_sec": per_device_tflops / step_time_delta.total_seconds()}) metrics["scalar"].update({"learning/current_learning_rate": lr}) + if total_weights is not None: + metrics["scalar"].update({"learning/total_weights": float(total_weights)}) _metrics_queue = queue.Queue() @@ -133,6 +155,25 @@ def write_metrics(writer, local_metrics_file, running_gcs_metrics, metrics, step if config.gcs_metrics and jax.process_index() == 0: running_gcs_metrics = max_utils.write_metrics_for_gcs(_buffered_metrics, _buffered_step, config, running_gcs_metrics) + if mld_metrics is not None and max_utils.ml_diagnostics_enabled(config) and jax.process_index() == 0: + if "scalar" in _buffered_metrics: + scalars = _buffered_metrics["scalar"] + metric_records = [] + for key, raw_val in scalars.items(): + val = float(raw_val.item() if hasattr(raw_val, "item") else raw_val) + metric_name = _METRICS_TO_MANAGED.get(key, key) + metric_records.append({ + "metric_name": metric_name, + "value": val, + "step": int(_buffered_step), + }) + + if metric_records: + mld_metrics.record_metrics(metric_records) + + _buffered_step = None + _buffered_metrics = None + _buffered_step = step _buffered_metrics = metrics