From 87fc2af6c1afd9e015c9952e508d619594441732 Mon Sep 17 00:00:00 2001 From: SurbhiJainUSC Date: Thu, 15 Oct 2026 19:56:36 +0000 Subject: [PATCH] Support TPU advanced profiling options and upload_all_profiler_results in MicroStepProfiler --- .../training_engine/micro_step_profiler.py | 31 ++++++- .../integration/maxtext_engine_e2e_test.py | 1 + .../unit/maxtext_engine_profiling_test.py | 84 +++++++++++++++++++ 3 files changed, 112 insertions(+), 4 deletions(-) diff --git a/src/maxtext/training_engine/micro_step_profiler.py b/src/maxtext/training_engine/micro_step_profiler.py index dd57e879c3..52190a5d71 100644 --- a/src/maxtext/training_engine/micro_step_profiler.py +++ b/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]. diff --git a/tests/post_training/integration/maxtext_engine_e2e_test.py b/tests/post_training/integration/maxtext_engine_e2e_test.py index c4d6c722b0..d7e1c2c998 100644 --- a/tests/post_training/integration/maxtext_engine_e2e_test.py +++ b/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, } diff --git a/tests/post_training/unit/maxtext_engine_profiling_test.py b/tests/post_training/unit/maxtext_engine_profiling_test.py index f4291abc5d..e75b8c3031 100644 --- a/tests/post_training/unit/maxtext_engine_profiling_test.py +++ b/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()