Add Google Cloud ML Diagnostics metrics support and documentation - #459
Open
richaguptaa17 wants to merge 1 commit into
Open
Add Google Cloud ML Diagnostics metrics support and documentation#459richaguptaa17 wants to merge 1 commit into
richaguptaa17 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.
richaguptaa17
force-pushed
the
mldiagnostics-metrics
branch
2 times, most recently
from
August 13, 2026 16:37
8f59dbe to
ed938eb
Compare
richaguptaa17
changed the title
richaguptaa17
force-pushed
the
mldiagnostics-metrics
branch
from
August 13, 2026 16:44
74bb624 to
7d3579f
Compare
entrpn
previously approved these changes
Aug 13, 2026
mbohlool
reviewed
Aug 13, 2026
Collaborator
|
Thanks for the PR @richaguptaa17. Could you please do the following?
|
richaguptaa17
force-pushed
the
mldiagnostics-metrics
branch
4 times, most recently
from
August 14, 2026 15:31
c42939a to
36cfd7e
Compare
Contributor
Author
|
Ran the linting commands and squashed the commits. |
entrpn
previously approved these changes
Aug 14, 2026
richaguptaa17
force-pushed
the
mldiagnostics-metrics
branch
3 times, most recently
from
August 18, 2026 09:33
3af4abf to
b27b0ab
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.
richaguptaa17
force-pushed
the
mldiagnostics-metrics
branch
from
August 18, 2026 13:55
b27b0ab to
8457cc7
Compare
Perseus14
approved these changes
Aug 18, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Integrate Google Cloud ML Diagnostics SDK (google-cloud-mldiagnostics) into MaxDiffusion to automatically record training, system, and performance metrics.
Key Changes:
Tested: