Add jaxpr dump support

[NuojCheng]

Description

This Pull Request introduces support for dumping Jaxpr (JAX intermediate representation) during the training and compilation processes. This functionality allows developers to inspect and verify the JAX traces of the training step, which is useful for debugging and ensuring consistency between different execution modes.

Key Changes

1. Configuration Options

New configuration flags have been added to src/MaxText/configs/base.yml and validated in src/MaxText/configs/types.py to control the Jaxpr dumping behavior:

• dump_jaxpr: Enables or disables Jaxpr dumping.
• dump_jaxpr_local_dir: Specifies the local directory where the .jaxpr files are initially saved.
• dump_jaxpr_gcs_dir: An optional GCS directory for uploading the dumps.
• dump_jaxpr_delete_local_after: Determines whether to delete local copies after a successful GCS upload.

2. Implementation of Jaxpr Dumping

A new utility function maybe_dump_jaxpr has been implemented in src/MaxText/maxtext_utils.py. This function:

• Uses jax.make_jaxpr to trace the jitted training step.
• Writes the resulting jaxpr to a file named train_step.jaxpr.
• Optionally uploads the dump to GCS using gcs_utils.upload_dump.

3. Integration into Training and Compilation

• Training Loop: maybe_dump_jaxpr is called within the main training loop in src/MaxText/train.py.
• AOT Compilation: It is also integrated into src/MaxText/train_compile.py to capture jaxpr during the ahead-of-time compilation process.

4. Test Refactoring and New Jaxpr Verification

The existing HLO identity tests in tests/integration/aot_hlo_identical_test.py have been replaced by a more comprehensive test suite in tests/integration/aot_identical_test.py. This new file includes:

• AotHloIdenticalTest: Updated verification for HLO graphs.
• AotJaxprIdenticalTest: A new test class that ensures the jaxpr generated during AOT compilation matches the jaxpr from a real training run.

Tests

https://paste.googleplex.com/5644648066449408

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.

[codecov]

Codecov Report

❌ Patch coverage is 73.07692% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/MaxText/maxtext_utils.py	68.18%	7 Missing ⚠️

📢 Thoughts on this report? Let us know!

[richjames0]

lgtm