add profiling option to compare_training_engine.py

[A9isha]

Description

This PR adds JAX/XProf profiling support to the gradient accumulation hardware benchmarking suite in compare_training_engine.py.

Context & Problem Solved

Profiling Gradient Accumulation Parity & Throughput:
When analyzing performance between baseline monolithic accumulation (jax.lax.scan) and MaxTextTrainingEngine (fwd_bwd + update), fine-grained TPU execution timelines and HLO trace comparisons are needed. Previously, benchmark_gradient_accumulation_performance() only measured wall-clock execution time without capturing XProf trace profiles.

Specific Changes

• tests/end_to_end/tpu/compare_training_engine.py:
  • Added profile: bool = True and profile_dir: str | None = None support to benchmark_gradient_accumulation_performance().
  • Added jax.profiler.start_trace() / jax.profiler.stop_trace() blocks around warm-up-separated baseline and engine execution loops.
  • Added jax.profiler.StepTraceAnnotation and jax.profiler.TraceAnnotation tags (baseline_cycle, engine_cycle, fwd_bwd_microstep_{idx}, engine_update) to allow detailed step inspection in the XProf Trace Viewer.
  • Added CLI flag parsing in run_all_verifications() for profile=, profile_dir=, m_steps=, and n_iterations=.
  • Used object.__setattr__ for setting use_tunix_gradient_accumulation on immutable HyperParameters config instances.

Tests

Tested end-to-end on Cloud TPU VM (v5p-8, 4 TPU chips) with Llama 3.1 8B topology:

1. Benchmark with Profiling Enabled (profile=true):

   python3 tests/end_to_end/tpu/compare_training_engine.py \
       src/maxtext/configs/post_train/rl.yml \
       model_name=llama3.1-8b \
       use_pathways=false \
       chips_per_vm=4 \
       ici_fsdp_parallelism=-1 \
       batch_size=8 \
       max_target_length=32 \
       convert_checkpoint_if_possible=false \
       test_suite=benchmark \
       profile=true \
       m_steps=5 \
       n_iterations=10 \
       profile_dir=<>
   

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces profiling capabilities to the gradient accumulation benchmark in compare_training_engine.py by integrating JAX profiler traces and step annotations, as well as adding CLI arguments to configure profiling. Feedback is provided regarding the use of object.__setattr__ to bypass the read-only design of HyperParameters, which can lead to an inconsistent state and silent bugs; it is recommended to revert to the original assertion check instead.

[gemini-code-assist]

Bypassing the read-only design of HyperParameters using object.__setattr__ is highly discouraged.\n\nHyperParameters stores its configuration values in internal dictionaries (self._flat_config and self._pydantic_config). Directly setting the attribute on the instance dictionary via object.__setattr__ will make cfg.use_tunix_gradient_accumulation return True due to standard attribute lookup, but any code calling cfg.get_keys() or copying/serializing the config will still see the old value (or False). This creates an inconsistent state and can lead to silent, hard-to-debug issues.\n\nInstead, please revert to the original assertion check to ensure the caller passes a correctly configured HyperParameters object.

  else:\n    assert (\n        cfg.use_tunix_gradient_accumulation\n    ), \