Add forward_pass_compile_test script

[snehalv2002]

Description

This PR introduces a standalone script to compile the forward pass logit checker offline in MaxText, enabling mathematical verification of HBM limits without incurring actual array data overhead on remote TPUs.

Features additions:

• Offline CPU Compilation Script: Added tests/utils/forward_compile.py to allow users to mock the parameter tree structure exactly as expected by the golden checker.
• HBM Profiling Validation: Incorporates XLA static compilation lower().compile() on the model application footprint, which yields explicit CompiledMemoryStats including exact byte allocations for code size, temporaries, and IO arguments, making OOM triaging infinitely easier for large parameter meshes.

Tests

• Verified functionality directly on high-RAM CPU VMs simulating a v6e-256 TPU mesh using deepseek4-284b parameterized layers to prove that the .lower().compile() safely projects temporary size execution vectors matching hardware capacity exactly.

Command Executed:

python3 tests/utils/forward_compile.py src/maxtext/configs/base.yml \
  base_output_directory=gs://snehalv-data/forward_compile_test \
  run_name=ds4-parity scan_layers=False attention=dot_product \
  per_device_batch_size=1 model_name=deepseek4-284b max_target_length=1024 \
  ici_fsdp_parallelism=1 ici_expert_parallelism=-1 weight_dtype=bfloat16 \
  dtype=bfloat16 activations_in_float32=false matmul_precision=highest \
  float32_logits=false float32_qk_product=false override_model_config=True \
  indexer_topk=4 enable_nnx=false pure_nnx=false pure_nnx_decoder=false \
  compile_topology=v6e-256 compile_topology_num_slices=1

Output:

Jitting and compiling forward_and_gather...
/home/snehalv_google_com/maxtext/tests/utils/forward_compile.py:82: DeprecationWarning: `with mesh:` context manager has been deprecated. Please use `with jax.set_mesh(mesh):` instead.
  with jax.set_mesh(topology_mesh), topology_mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
Jitting and compilation complete!
Memory analysis: CompiledMemoryStats(generated_code_size_in_bytes=1627936768, argument_size_in_bytes=5260023808, output_size_in_bytes=529530880, alias_size_in_bytes=0, temp_size_in_bytes=4745168512, host_generated_code_size_in_bytes=0, host_argument_size_in_bytes=0, host_output_size_in_bytes=0, host_alias_size_in_bytes=0, host_temp_size_in_bytes=0)

Checklist

• 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 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]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.