Merge pull request #3303 from abhinavgoel95:fix/gate-tpu-power-profiling-events

PiperOrigin-RevId: 878616815

Author: Google-ML-Automation

src/maxtext/common/profiler.py

@@ -51,7 +51,7 @@ def __init__(self, config, offset_step=0):
       ManagedMLDiagnostics(config)  # Initialize the MLRun instance.
 
     self.profiling_options = jax.profiler.ProfileOptions()
-    if self.mode == "xplane" and not self.managed_mldiagnostics:
+    if self.mode == "xplane" and not self.managed_mldiagnostics and config.profile_power_events:
       self.profiling_options.advanced_configuration = {
           "tpu_power_trace_level": config.xprof_tpu_power_trace_level,
           "e2e_enable_fw_throttle_event": config.xprof_e2e_enable_fw_throttle_event,

src/maxtext/configs/base.yml

@@ -916,6 +916,7 @@ xprof_tpu_power_trace_level: 0
 xprof_e2e_enable_fw_throttle_event: False
 xprof_e2e_enable_fw_power_level_event: False
 xprof_e2e_enable_fw_thermal_event: False
+profile_power_events: False # Set to True to enable TPU-specific power/thermal profiling events. Defaults to False to avoid breaking GPU xplane tracing.
 
 log_config: True # Prints the config (after defaults have been set by pyconfig logic)
 debug_sharding: False # Prints model weights sharding info

src/maxtext/configs/types.py

@@ -1383,6 +1383,7 @@ class Profiling(BaseModel):
   xprof_e2e_enable_fw_throttle_event: bool = Field(False, description="Enable FW throttle event.")
   xprof_e2e_enable_fw_power_level_event: bool = Field(False, description="Enable FW power level event.")
   xprof_e2e_enable_fw_thermal_event: bool = Field(False, description="Enable FW thermal event.")
+  profile_power_events: bool = Field(False, description="Enable TPU-specific power/thermal profiling events. Defaults to False to avoid breaking GPU xplane tracing.")
 
 
 class HloDump(BaseModel):