Add custom mesh and logical rule support

[NuojCheng]

Description

This pull request introduces the ability to define and use customized physical meshes and logical axis rules in MaxText via external YAML configuration files. This enhances flexibility for complex parallelism strategies, such as those required for large scale MoE setups.

Key Changes

Configuration Enhancements

• Custom Config Loading: Added a custom_mesh_and_rule field to base.yml.
• Dynamic YAML Parsing: Implemented logic in src/maxtext/configs/types.py to automatically load mesh_axes, logical_axis_rules, and data_sharding from specified YAML files located in src/maxtext/configs/custom_mesh_and_rule/.
• Example Configs: Provided initial custom configurations for pure-fsdp.yml and a complex ds3-large-pp.yml.

Robust Utility Functions

• Added helper functions in src/maxtext/utils/max_utils.py for safer access to mesh properties:
  • get_mesh_axes_size: Returns the size of a physical axis or 1 if it doesn't exist.
  • get_logical_rule_contents: Safely retrieves physical axes corresponding to a logical name.

Code Refactoring for Mesh Flexibility

• Updated src/maxtext/layers/attention_op.py, src/maxtext/utils/sharding.py, and src/maxtext/utils/train_utils.py to use the new utility functions. This ensures the model can compile even when standard axes (like data or fsdp) are absent or redefined in a custom mesh.

Tests

Added tests/unit/custom_mesh_and_rule_test.py to verify successful compilation using the new custom mesh flags for both FSDP and multi-axis configurations.

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!

[gobbleturk]

I love the idea of increased visibility and sharability!

[richjames0]

to the moon

[gobbleturk]

Let the sharding sharing begin!