add ml-diagnostic guide

Author: Sharon Yu

docs/guides/monitoring_and_debugging.md

@@ -26,4 +26,5 @@ monitoring_and_debugging/monitor_goodput.md
 monitoring_and_debugging/understand_logs_and_metrics.md
 monitoring_and_debugging/use_vertex_ai_tensorboard.md
 monitoring_and_debugging/xprof_user_guide.md
+monitoring_and_debugging/ml_workload_diagnostics.md
 ```

docs/guides/monitoring_and_debugging/ml_workload_diagnostics.md

@@ -0,0 +1,117 @@
+<!--
+ Copyright 2024 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.
+ -->
+
+# Running a workload with ML Diagnostics Enabled
+This guide provides an overview on how to enable ML Diagnostics for your MaxText workload.
+
+## Overview
+Google Cloud ML Diagnostics is an end-to-end managed platform for ML Engineers to optimize and diagnose their AI/ML workloads on Google Cloud. The product allows ML Engineers to collect and visualize all their workload metrics, configs and profiles with one single platform, all within the same UI.  This platform works for any ML workload (training, inference, etc) including working with Maxtext/Maxdiffusion as well as any orchestrator on TPU including GKE as well as custom orchestrator. The current product offering focuses on workloads running on XLA-based frameworks (JAX, Pytorch XLA, Tensorflow/Keras) on Google Cloud TPUs and GPUs. Current support is for JAX on Google Cloud TPUs only. 
+
+Google Cloud ML Diagnostics includes the following features:
+- SDK: An open source ML Diagnostics [SDK](https://github.com/AI-Hypercomputer/google-cloud-mldiagnostics?tab=readme-ov-file) to use with your ML workload in order to enable managed ML workload diagnostics experience on Google Cloud
+- Integration with JAX and Pytorch framework and libraries (only JAX supported for Preview)
+- Metrics/configs/profiles management: 
+Track workload metrics, including model quality, model performance and system metrics.
+Track workload configs including software configs, system configs as well as user-defined configs
+- Manage profile sessions 
+Managed XProf: Managed profiling with [XProf](https://openxla.org/xprof), which allows faster loading of large profiles, supports multiple users simultaneously accessing profiles and supports easy to use out-of-the-box features such as multi-host profiling and on-demand profiling. 
+- Visualization of metrics/configs/profiles in both Cluster Director and Google Kubernetes Engine on the Google Cloud console
+- Link sharing for ML runs and profiles for easy collaboration
+
+## Enabling ML Diagnostics on Maxtext Workload
+MaxText has integrated the ML Diagnostics SDK in its code. You can enable ML Diagnostics with the **managed-mldiagnostics** flag. If this is enabled, it will
+
+- create a managed MachineLearning run with all the MaxText configs
+- upload profiling traces, if the profiling is enabled by profiler="xplane".
+- upload training metrics, at the defined log_period interval.
+
+1. **Set your configuration**
+
+    ```
+    export PROJECT_ID="your-gcp-project-id"
+    export ZONE="your-gcp-zone" # e.g., us-central1-a
+    export CLUSTER_NAME="your-existing-cluster-name"
+    export BASE_OUTPUT_DIR="gs://your-output-bucket/"
+    export DATASET_PATH="gs://your-dataset-bucket/"
+    ```
+
+2.  **Configure gcloud CLI**
+
+    ```
+    gcloud config set project $PROJECT_ID
+    gcloud config set compute/zone $ZONE
+    ```
+
+### Run Maxtext via XPK
+1.   Enable ML Diagnostics to just capture Maxtext metrics and configs
+
+            xpk workload create\
+              --cluster ${CLUSTER_NAME}\
+              --workload ${USER}-tpu-job\
+              --base-docker-image maxtext_base_image\
+              --tpu-type v5litepod-256\
+              --num-slices 1\
+              --command "python3 -m MaxText.train src/MaxText/configs/base.yml run_name=${USER}-tpu-job base_output_directory=${BASE_OUTPUT_DIR} dataset_path=${DATASET_PATH} steps=100 log_period=10 managed_mldiagnostics=True"
+    
+2.   Enable ML Diagnostics to capture Maxtext metrics, configs and singlehost profiles (on the first TPU device)
+
+            xpk workload create\
+              --cluster ${CLUSTER_NAME}\
+              --workload ${USER}-tpu-job\
+              --base-docker-image maxtext_base_image\
+              --tpu-type v5litepod-256\
+              --num-slices 1\
+              --command "python3 -m MaxText.train src/MaxText/configs/base.yml run_name=${USER}-tpu-job base_output_directory=${BASE_OUTPUT_DIR} dataset_path=${DATASET_PATH} steps=100 log_period=10 profiler=xplane managed_mldiagnostics=True"
+            
+3.   Enable ML Diagnostics to capture Maxtext metrics, configs and multihost profiles (on all TPU devices)
+
+            xpk workload create\
+              --cluster ${CLUSTER_NAME}\
+              --workload ${USER}-tpu-job\
+              --base-docker-image maxtext_base_image\
+              --tpu-type v5litepod-256\
+              --num-slices 1\
+              --command "python3 -m MaxText.train src/MaxText/configs/base.yml run_name=${USER}-tpu-job base_output_directory=${BASE_OUTPUT_DIR} dataset_path=${DATASET_PATH} steps=100 log_period=10 profiler=xplane upload_all_profiler_results=True managed_mldiagnostics=True"
+### Run Maxtext via Pathways
+1.   Enable ML Diagnostics to just capture Maxtext metrics and configs
+
+            xpk workload create-pathways\
+              --cluster ${CLUSTER_NAME}\
+              --workload ${USER}-tpu-job\
+              --base-docker-image maxtext_base_image\
+              --tpu-type v5litepod-256\
+              --num-slices 1\
+              --command "python3 -m MaxText.train src/MaxText/configs/base.yml run_name=${USER}-tpu-job base_output_directory=${BASE_OUTPUT_DIR} dataset_path=${DATASET_PATH} steps=100 enable_single_controller=True log_period=10 managed_mldiagnostics=True"
+    
+2.   Enable ML Diagnostics to capture Maxtext metrics, configs and singlehost profiles (on the first TPU device)
+
+            xpk workload create-pathways\
+              --cluster ${CLUSTER_NAME}\
+              --workload ${USER}-tpu-job\
+              --base-docker-image maxtext_base_image\
+              --tpu-type v5litepod-256\
+              --num-slices 1\
+              --command "python3 -m MaxText.train src/MaxText/configs/base.yml run_name=${USER}-tpu-job base_output_directory=${BASE_OUTPUT_DIR} dataset_path=${DATASET_PATH} steps=100 enable_single_controller=True log_period=10 profiler=xplane managed_mldiagnostics=True"
+            
+3.   Enable ML Diagnostics to capture Maxtext metrics, configs and multihost profiles (on all TPU devices)
+
+            xpk workload create-pathways\
+              --cluster ${CLUSTER_NAME}\
+              --workload ${USER}-tpu-job\
+              --base-docker-image maxtext_base_image\
+              --tpu-type v5litepod-256\
+              --num-slices 1\
+              --command "python3 -m MaxText.train src/MaxText/configs/base.yml run_name=${USER}-tpu-job base_output_directory=${BASE_OUTPUT_DIR} dataset_path=${DATASET_PATH} steps=100 enable_single_controller=True log_period=10 profiler=xplane upload_all_profiler_results=True managed_mldiagnostics=True"

docs/run_maxtext/run_maxtext_via_xpk.md

@@ -225,4 +225,4 @@ For instance, to run a job across **four TPU slices**, you would change `--num-s
 
     ```
     xpk workload delete --cluster ${CLUSTER_NAME} --workload <your-workload-name>
-    ```
+    ```