[maxtext] improve profiling in decode

Limit the profiling of decode to `config.profiler_steps` steps. User may request generate many tokens, but since each generation is identical, it's not necessary to profile all of them. Profiling only the first few should be sufficient.

Also added a method in Profiler to do post-pocessing on the collected trace.

PiperOrigin-RevId: 773929904

Author: maxtext authors

MaxText/decode.py

@@ -157,12 +157,20 @@ def main(argv: Sequence[str]) -> None:
     decode_state = engine.insert(prefill_result_list[i], decode_state, slot=i)
 
   # Generate
+  prof_deactivated = False
   steps = range(config.max_prefill_predict_length, config.max_target_length)
   sampled_tokens_list.append(_batch_first_result_token(first_token_list, batch_size))
   for i in steps:
     rng, rng_generate = jax.random.split(rng)
     with jax.profiler.StepTraceAnnotation("generate", step=i):
       decode_state, sampled_tokens = engine.generate(params, decode_state, rng=rng_generate)
+
+    # Automatically deactivate profiler after profiler_steps steps
+    if i > config.max_prefill_predict_length + config.profiler_steps:
+      jax.block_until_ready(sampled_tokens)
+      prof.deactivate()
+      prof_deactivated = True
+
     sampled_tokens_list.append(sampled_tokens)
 
   # Get results
@@ -176,7 +184,11 @@ def main(argv: Sequence[str]) -> None:
   ), f"generated text mismatch {output=}, {config.autoregressive_decode_assert=}"
 
   # Deactivate profiler
-  prof.deactivate()
+  if not prof_deactivated:
+    jax.block_until_ready(output)
+    prof.deactivate()
+
+  prof.post_process()
 
 def _validate_config(config):
   assert config.load_full_state_path == "", (

MaxText/profiler.py

@@ -92,3 +92,6 @@ def should_activate_periodic_profile(self, step):
 
   def should_deactivate_periodic_profile(self, step):
     return self.profile_period > 0 and (step - self.finished_initial_profile_step) % self.profile_period == 0
+
+  def post_process(self):
+    pass