Add Google Cloud ML Diagnostics metrics support and documentation

README.md

@@ -95,6 +95,8 @@ MaxDiffusion supports
   - [Getting Started: Multihost development](#getting-started-multihost-development)
 - [Comparison to Alternatives](#comparison-to-alternatives)
 - [Development](#development)
+  - [Profiling](#profiling)
+  - [Metrics](#metrics)
 
 # Getting Started
 
@@ -914,4 +916,7 @@ This script will automatically format your code with `pyink` and help you identi
 The full suite of -end-to end tests is in `tests` and `src/maxdiffusion/tests`. We run them with a nightly cadance.
 
 ## Profiling
-To learn how to enable ML Diagnostics and XProf profiling for your runs, please see our [ML Diagnostics Guide](docs/profiling.md).
+To learn how to enable ML Diagnostics and XProf profiling for your runs, please see our [ML Diagnostics Guide](docs/profiling.md).
+
+## Metrics
+To learn how to enable ML Diagnostics metrics tracking for your runs, please see our [Metrics Guide](docs/metrics.md).

docs/metrics.md

@@ -0,0 +1,160 @@
+<!--
+ 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.
+-->
+
+# 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/<project_id>/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/<PROJECT_ID>/logs/ml_diagnostics_metric" AND resource.labels.namespace="loss"' \
+  --limit=5 \
+  --format="json"
+
+# Query custom metrics
+gcloud logging read 'logName="projects/<PROJECT_ID>/logs/ml_diagnostics_metric" AND resource.labels.namespace="custom/latents_mean"' \
+  --limit=5 \
+  --format="json"
+
+# Query hardware metrics
+gcloud logging read 'logName="projects/<PROJECT_ID>/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`.

src/maxdiffusion/__init__.py

@@ -47,8 +47,6 @@
     "configuration_utils": ["ConfigMixin"],
     "max_logging": [],
     "max_utils": [],
-    "maxdiffusion_google": [],
-    "maxdiffusion_google_hub": [],
     "maxdiffusion_utils": [],
     "models": [],
     "multihost_dataloading": [],

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: