Fix MaxText 22b.sh AOT compilation OOM on v4-128

[darisoy]

Description

This PR resolves the TPU HBM Out-of-Memory (OOM) error during Ahead-of-Time (AOT) compilation for the v4-128 topology with per_device_batch_size=13 by enabling Vocabulary Tiling and forcing Flash Attention in the 22b.sh configuration.

Details

• The default configuration of 22b.sh fails with an HBM OOM during compilation.
• Enabling vocabulary tiling (num_vocab_tiling=8) and forcing flash attention (attention=flash) reduces peak memory usage and compiler fragmentation, allowing the compilation to succeed.

FIXED: b/517329766
BUGS: b/517329766

Tests

The fix was verified by running the AOT compilation script directly on a TPU VM targeting the 2-slice v4-128 topology.

Command to Reproduce

bash src/maxtext/configs/tpu/v4/22b.sh \
  EXECUTABLE=train_compile \
  M_COMPILE_TOPOLOGY=v4-128 \
  M_COMPILE_TOPOLOGY_NUM_SLICES=2 \
  DATASET_PATH=dummy-dataset \
  OUTPUT_PATH=dummy-output-dir \
  RUN_PREFLIGHT=false

Results

Compilation completed successfully:

Jitting and compilation complete!
Finished train_compile.py successfully!

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

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!