Gate TPU power profiling events behind profile_power_events flag

On GPU, the CUPTI tracer does not recognize TPU-specific keys in
advanced_configuration (tpu_power_trace_level, e2e_enable_fw_*_event),
causing an INVALID_ARGUMENT error that aborts the GPU xplane trace and
leaves only CPU traces behind.

Introduce a profile_power_events flag (default False) that gates the
TPU-specific advanced_configuration block. GPU runs are unaffected by
default; TPU users who want power/thermal tracing can opt in with
profile_power_events=True.

Author: abhinavgoel95

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

@@ -911,6 +911,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

@@ -1380,6 +1380,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):