FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Add Google Cloud ML Diagnostics metrics support and documentation by richaguptaa17 · Pull Request #459 · AI-Hypercomputer/maxdiffusion · GitHub

/ maxdiffusion Public

Add Google Cloud ML Diagnostics metrics support and documentation - #459

Closed
richaguptaa17 wants to merge 1 commit into
AI-Hypercomputer:mainfrom
richaguptaa17:mldiagnostics-metrics
ConversationCommits (1)ChecksFiles changed
Closed

Add Google Cloud ML Diagnostics metrics support and documentation#459
richaguptaa17 wants to merge 1 commit into
AI-Hypercomputer:mainfrom
richaguptaa17:mldiagnostics-metrics

Conversation

richaguptaa17 commented Aug 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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 metrics names to canonical MetricType enums (loss, learning_rate, gradient_norm, total_weights, step_time, tflops) with automatic pass-through for custom metrics.
  • max_utils.py: configured region=None for GCP cluster auto-discovery, and enabled background hardware metric collection (log_system_metrics=True).
  • metrics.md: Created comprehensive guide for capturing metrics using google-cloud-mldiagnostics integration in training scripts.

Tested:

  • Ran a multi-host distributed training run on TPU v6e cluster. Verified successful metric ingestion in Cloud Logging for predefined, custom, and hardware utilization metrics.

richaguptaa17 requested a review from entrpn as a code owner August 13, 2026 16:21
gemini-code-assist Bot reviewed Aug 13, 2026
View reviewed changes

gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

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.

Comment thread src/maxdiffusion/train_utils.py Outdated
Comment thread src/maxdiffusion/train_utils.py Outdated
richaguptaa17 force-pushed the mldiagnostics-metrics branch 2 times, most recently from 8f59dbe to ed938eb Compare August 13, 2026 16:37
richaguptaa17 changed the title Add Google Cloud ML Diagnostics metrics support and developer guide. Add Google Cloud ML Diagnostics metrics support and documentation Aug 13, 2026
richaguptaa17 force-pushed the mldiagnostics-metrics branch from 74bb624 to 7d3579f Compare August 13, 2026 16:44
entrpn previously approved these changes Aug 13, 2026
View reviewed changes
mbohlool reviewed Aug 13, 2026
View reviewed changes
Comment thread src/maxdiffusion/max_utils.py Outdated
Comment thread src/maxdiffusion/train_utils.py Outdated

Perseus14 commented Aug 14, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the PR @richaguptaa17. Could you please do the following?

  • Run linting commands ruff check . and pyink src/maxdiffusion --pyink-indentation=2 --line-length=125 (ruff==0.15.12 and pyink==23.10.0)
  • Squash commits

richaguptaa17 force-pushed the mldiagnostics-metrics branch 4 times, most recently from c42939a to 36cfd7e Compare August 14, 2026 15:31

richaguptaa17 commented Aug 14, 2026

Copy link
Copy Markdown
Contributor Author

Ran the linting commands and squashed the commits.

entrpn previously approved these changes Aug 14, 2026
View reviewed changes
richaguptaa17 force-pushed the mldiagnostics-metrics branch 4 times, most recently from b27b0ab to 8457cc7 Compare August 18, 2026 13:55
Perseus14 previously approved these changes Aug 18, 2026
View reviewed changes
richaguptaa17 force-pushed the mldiagnostics-metrics branch from 8457cc7 to b51352e Compare August 20, 2026 21:50
Add Google Cloud ML Diagnostics metrics support and developer guide.
d91454d 
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 b51352e to d91454d Compare August 21, 2026 14:13
richaguptaa17 requested a review from entrpn August 21, 2026 18:34
mbohlool requested changes Aug 21, 2026
View reviewed changes
Comment thread src/maxdiffusion/__init__.py
"configuration_utils": ["ConfigMixin"],
"max_logging": [],
"max_utils": [],
"maxdiffusion_google": [],

mbohlool Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

why you removed these two lines?

richaguptaa17 Aug 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

These are Google3-internal modules that do not exist in the open-source GitHub repository.

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
Show hidden characters
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

mbohlool mbohlool approved these changes
entrpn entrpn approved these changes
Perseus14 Perseus14 left review comments
+1 more reviewer
gemini-code-assist[bot] gemini-code-assist[bot] left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

Back | FazBrowse Home | New Git URL