Add Google Cloud ML Diagnostics metrics support and documentation - #459
Add Google Cloud ML Diagnostics metrics support and documentation#459richaguptaa17 wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request integrates the Google Cloud ML Diagnostics SDK into MaxDiffusion to enable real-time telemetry, including workload, system, and accelerator metrics. Key changes include adding documentation, filtering non-JSON serializable keys from the configuration, and recording scalar metrics like gradient norm and total weights in the training loops. The review feedback highlights two important issues: first, workload metrics recording should be restricted to the master node (process 0) to prevent duplicate logs in multi-host setups; second, the MetricType enum resolution should be made safer to avoid potential AttributeError crashes during module import.
8f59dbe to
ed938eb
Compare
richaguptaa17
changed the title
74bb624 to
7d3579f
Compare
| return "enable_ml_diagnostics" in config.get_keys() and config.enable_ml_diagnostics | ||
|
|
||
|
|
||
| _ml_diagnostics_profiler_enabled = ml_diagnostics_enabled |
There was a problem hiding this comment.
Removed. It was an unnecessary alias from renaming the function.
There was a problem hiding this comment.
I can still see it, maybe you forgot to push?
| metric_types = None | ||
|
|
||
|
|
||
| _metric_type_enum = getattr(metric_types, "MetricType", None) if metric_types else None |
There was a problem hiding this comment.
do we expect MetricType attribute not present if metric_types is not None? Lots of silence failure here. If ml diagnostic exists, we should not ignore any failure after that. So if I install ml_diagnostic dependencies on a job, I expect it to work and if there is any incompatibilities/failures, the job should fail rather than ignore it and disable ml_diagnostics.
There was a problem hiding this comment.
Agreed. If google_cloud_mldiagnostics is installed, metric_types.MetricType should be present. Replaced fallback with direct mapping.
Made a change to raise an explicit ImportError if a user sets enable_ml_diagnostics: True but google_cloud_mldiagnostics is not installed or initialized.
|
Thanks for the PR @richaguptaa17. Could you please do the following?
|
141c0a7 to
c42939a
Compare
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.
c42939a to
36cfd7e
Compare
|
Ran the linting commands and squashed the commits. |
4 participants
Integrate Google Cloud ML Diagnostics SDK (google-cloud-mldiagnostics) into MaxDiffusion to automatically record training, system, and performance metrics.
Key Changes:
Tested: