Add Google Cloud ML Diagnostics integration for profiling and metrics#1371
Open
rapatchi wants to merge 1 commit into
Open
Add Google Cloud ML Diagnostics integration for profiling and metrics#1371rapatchi wants to merge 1 commit into
rapatchi wants to merge 1 commit into
Conversation
rapatchi
force-pushed
the
mldiag-integration
branch
5 times, most recently
from
July 10, 2026 08:28
1376906 to
551d10d
Compare
Introduce integration with `google_cloud_mldiagnostics` to support automated profiler (xprof) tracing and metric collection on Google Cloud Platform. * Add `ManagedMLDiagnostics` (`axlearn/common/managed_mldiagnostics.py`), a thread-safe singleton wrapper that maps common AXLearn metrics (such as loss, MFU, learning rate, and step time) to their corresponding GCP ML Diagnostics MetricTypes and uploads them. The run name is automatically resolved from the `AXLEARN_JOB_NAME` environment variable. * Configure GKE launcher (`axlearn/cloud/gcp/tpu.py`, `axlearn/cloud/gcp/jobset_utils.py`) to inject `AXLEARN_JOB_NAME` (containing GKE job name) into containers. * Add `MLDiagnosticsMetricsWriter` to `axlearn/common/summary_writer.py` which writes scalar metrics to the diagnostics singleton. * Add `--enable_ml_diagnostics_xprof`, `--enable_ml_diagnostics_metrics`, and `--ml_diagnostics_region` flags to `launch_trainer.py`. * In `launch_trainer.py`, automatically inject `MLDiagnosticsMetricsWriter` into the trainer and evalers' summary writers if ML Diagnostics metrics are enabled. If the configured writer is not a composite, wrap it in a `CompositeWriter` alongside the new writer. * Update `SpmdTrainer` and `SpmdEvaler` to route profiler tracing through the ML Diagnostics `xprof` module when xprof is enabled, instead of direct `jax.profiler` calls. * Pass the trainer's `ml_diagnostics` configuration to all evalers to enable profiler tracing and/or metrics on them. * Add unit tests in `managed_mldiagnostics_test.py`, `summary_writer_test.py`, `trainer_test.py`, `evaler_test.py` and `launch_trainer_test.py` to cover initialization, metric mapping, writer injection, and xprof tracing logic.
rapatchi
force-pushed
the
mldiag-integration
branch
from
July 15, 2026 04:57
551d10d to
c282ef8
Compare
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
1 participant
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.
Introduce integration with
google_cloud_mldiagnosticsto support automated profiler (xprof) tracing and metric collection on Google Cloud Platform.ManagedMLDiagnostics(axlearn/common/managed_mldiagnostics.py), a thread-safe singleton wrapper that maps common AXLearn metrics (such as loss, learning rate, and step time) to their corresponding GCP ML Diagnostics MetricTypes and uploads them.MLDiagnosticsMetricsWritertoaxlearn/common/summary_writer.pywhich writes scalar metrics to the diagnostics singleton.--enable_ml_diagnostics,--ml_diagnostics_run_name, and--ml_diagnostics_regionflags tolaunch_trainer.py.launch_trainer.py, automatically injectMLDiagnosticsMetricsWriterinto the trainer and evalers' summary writers if ML Diagnostics is enabled. If the configured writer is not a composite, wrap it in aCompositeWriteralongside the new writer.SpmdTrainerandSpmdEvalerto route profiler tracing through the ML Diagnosticsxprofmodule when enabled, instead of directjax.profilercalls.ml_diagnosticsconfiguration to all evalers to enable profiler tracing on them.managed_mldiagnostics_test.py,summary_writer_test.py,trainer_test.py,evaler_test.pyandlaunch_trainer_test.pyto cover initialization, metric mapping, writer injection, and xprof tracing logic.