Support TPU advanced profiling options and upload_all_profiler_results in MicroStepProfiler

src/maxtext/training_engine/micro_step_profiler.py

@@ -21,6 +21,7 @@
 
 from absl import logging
 from maxtext.configs import pyconfig
+from maxtext.configs import types
 
 
 class MicroStepProfiler:
@@ -29,7 +30,11 @@ class MicroStepProfiler:
   def __init__(self, config: pyconfig.HyperParameters) -> None:
     """Initializes the profiler."""
     self.do_not_profile = False
-    if jax.process_index() != 0 or config.profiler_steps == 0:
+    if (
+        config.profiler != types.ProfilerType.XPLANE
+        or config.profiler_steps == 0
+        or (not config.upload_all_profiler_results and jax.process_index() != 0)
+    ):
       self.do_not_profile = True
       return
 
@@ -45,10 +50,28 @@ def __init__(self, config: pyconfig.HyperParameters) -> None:
       )
       return
     self.profiling_options = jax.profiler.ProfileOptions()
+    advanced_config: dict[str, Any] = {}
+    if config.profile_power_events:
+      advanced_config.update(
+          {
+              "tpu_power_trace_level": int(config.xprof_tpu_power_trace_level),
+              "e2e_enable_fw_throttle_event": config.xprof_e2e_enable_fw_throttle_event,
+              "e2e_enable_fw_power_level_event": config.xprof_e2e_enable_fw_power_level_event,
+              "e2e_enable_fw_thermal_event": config.xprof_e2e_enable_fw_thermal_event,
+          }
+      )
+    if config.enable_tpu_profiling_options:
+      advanced_config.update(
+          {
+              "tpu_num_chips_to_profile_per_task": config.tpu_num_chips_to_profile_per_task,
+              "tpu_num_sparse_core_tiles_to_trace": config.tpu_num_sparse_core_tiles_to_trace,
+              "tpu_num_sparse_cores_to_trace": config.tpu_num_sparse_cores_to_trace,
+          }
+      )
     if config.enable_continuous_profiling:
-      self.profiling_options.advanced_configuration = {
-          "enable_continuous_profiling": True,
-      }
+      advanced_config["enable_continuous_profiling"] = True
+    if advanced_config:
+      self.profiling_options.advanced_configuration = advanced_config
     self.profile_active: bool = False
     # profiler_period is the distance between successive profiling window,
     # profiler_steps=3, profiler_period=4 gives [1,3], [5,7], [9,11].

tests/post_training/integration/maxtext_engine_e2e_test.py

@@ -145,6 +145,7 @@ def setup_config(self):
         "async_checkpointing": True,
         # Disable scan_layers to prevent prepare_weight_sync from trying to unscan layers on DummyNNXModel
         "scan_layers": False,
+        "profiler": "xplane",
         "skip_first_n_steps_for_profiler": 1,
         "profiler_steps": 2,
     }

tests/post_training/unit/maxtext_engine_profiling_test.py

@@ -23,8 +23,10 @@
 import jax
 import jax.numpy as jnp
 from maxtext.configs import pyconfig
+from maxtext.configs import types
 from maxtext.training_engine import abstract_engine
 from maxtext.training_engine import maxtext_engine
+from maxtext.training_engine import micro_step_profiler
 from maxtext.utils import maxtext_utils
 import pytest
 from tests.utils.test_helpers import get_test_config_path
@@ -105,6 +107,7 @@ def test_the_engine_opens_and_closes_the_configured_window(self):
     """`fwd_bwd` starts the profile and `update` stops it, without an outer loop."""
     cfg = _tiny_config(
         base_output_directory=self.create_tempdir().full_path,
+        profiler="xplane",
         skip_first_n_steps_for_profiler=1,
         profiler_steps=1,
     )
@@ -128,6 +131,7 @@ def test_the_engine_opens_and_closes_the_configured_window(self):
   def test_close_rescues_a_profile_left_open_by_an_interrupted_step(self):
     cfg = _tiny_config(
         base_output_directory=self.create_tempdir().full_path,
+        profiler="xplane",
         skip_first_n_steps_for_profiler=0,
         profiler_steps=5,
     )
@@ -142,6 +146,86 @@ def test_close_rescues_a_profile_left_open_by_an_interrupted_step(self):
 
     self.stop_trace.assert_called_once()
 
+  def test_profiler_disabled_when_profiler_mode_is_not_xplane(self):
+    cfg = _tiny_config(
+        base_output_directory=self.create_tempdir().full_path,
+        skip_first_n_steps_for_profiler=0,
+        profiler_steps=5,
+    )
+    self.assertEqual(cfg.profiler, types.ProfilerType.NONE)
+    prof = micro_step_profiler.MicroStepProfiler(cfg)
+    self.assertTrue(prof.do_not_profile)
+    prof.maybe_activate(total_micro_steps=0, current_micro_step=0, train_step=0)
+    self.start_trace.assert_not_called()
+
+  def test_upload_all_profiler_results_enables_non_zero_process_profiling(self):
+    cfg_host0_only = _tiny_config(
+        base_output_directory=self.create_tempdir().full_path,
+        profiler="xplane",
+        skip_first_n_steps_for_profiler=0,
+        profiler_steps=1,
+        upload_all_profiler_results=False,
+    )
+    cfg_all_hosts = _tiny_config(
+        base_output_directory=self.create_tempdir().full_path,
+        profiler="xplane",
+        skip_first_n_steps_for_profiler=0,
+        profiler_steps=1,
+        upload_all_profiler_results=True,
+    )
+    with patch("jax.process_index", return_value=1):
+      prof_host0_only = micro_step_profiler.MicroStepProfiler(cfg_host0_only)
+      self.assertTrue(prof_host0_only.do_not_profile)
+      prof_host0_only.maybe_activate(total_micro_steps=0, current_micro_step=0, train_step=0)
+      self.start_trace.assert_not_called()
+
+      prof_all_hosts = micro_step_profiler.MicroStepProfiler(cfg_all_hosts)
+      self.assertFalse(prof_all_hosts.do_not_profile)
+      prof_all_hosts.maybe_activate(total_micro_steps=0, current_micro_step=0, train_step=0)
+      self.start_trace.assert_called_once()
+      prof_all_hosts.maybe_deactivate(total_micro_steps=0, current_micro_step=0, train_step=0)
+      self.stop_trace.assert_called_once()
+
+  def test_tpu_and_power_profiling_options_populated_in_advanced_configuration(self):
+    cfg = _tiny_config(
+        base_output_directory=self.create_tempdir().full_path,
+        profiler="xplane",
+        skip_first_n_steps_for_profiler=0,
+        profiler_steps=1,
+        enable_tpu_profiling_options=True,
+        tpu_num_chips_to_profile_per_task=2,
+        tpu_num_sparse_cores_to_trace=4,
+        tpu_num_sparse_core_tiles_to_trace=8,
+        profile_power_events=True,
+        xprof_tpu_power_trace_level=1,
+        xprof_e2e_enable_fw_throttle_event=True,
+        xprof_e2e_enable_fw_power_level_event=True,
+        xprof_e2e_enable_fw_thermal_event=True,
+        enable_continuous_profiling=True,
+    )
+    with patch("jax.profiler.ProfileOptions") as mock_options_cls:
+      prof = micro_step_profiler.MicroStepProfiler(cfg)
+      mock_options_cls.assert_called_once()
+      self.assertEqual(
+          prof.profiling_options.advanced_configuration,
+          {
+              "tpu_power_trace_level": 1,
+              "e2e_enable_fw_throttle_event": True,
+              "e2e_enable_fw_power_level_event": True,
+              "e2e_enable_fw_thermal_event": True,
+              "tpu_num_chips_to_profile_per_task": 2,
+              "tpu_num_sparse_core_tiles_to_trace": 8,
+              "tpu_num_sparse_cores_to_trace": 4,
+              "enable_continuous_profiling": True,
+          },
+      )
+      self.assertIs(type(prof.profiling_options.advanced_configuration["tpu_power_trace_level"]), int)
+      prof.maybe_activate(total_micro_steps=0, current_micro_step=0, train_step=0)
+      self.start_trace.assert_called_once_with(
+          cfg.tensorboard_dir,
+          profiler_options=prof.profiling_options,
+      )
+
 
 if __name__ == "__main__":
   absltest.main()