Split logical rule for eval from for train

[NuojCheng]

Description

We might want different sharding logical rule for eval compared with train, since the global train batch size and eval batch size are different. This PR enables this feature by introducing two new flags:

• custom_mesh_and_rule_for_eval
• logical_axis_rules_for_eval

When custom_mesh_and_rule_for_eval is specified (non-default value), an alternative logical rule is used for the eval pipeline. A typical use case would be using fractional batch size for eval, so that CP/TP like shardings are enabled (only for eval). Same axis is used as other shardings, e.g. FSDP/EP, for the train step.

Currently, we still use the same mesh for both train and eval steps. A potential future optimization would be split the meshes for train and eval, which is in theory possible since train and eval functions are in different jit.It may required careful checks.

FIXES: b/535639231

Tests

Run on v5p-8: we use FSDP+EP for train step, and EP-as-CP custom rule for eval,

smoke_train model_name=deepseek3-test per_device_batch_size=1 eval_per_device_batch_size=0.25 eval_interval=3 eval_steps=1 ici_expert_parallelism=4 custom_mesh_and_rule_for_eval=ep-as-cp debug_sharding=false profiler=xplane use_ring_of_experts=true  max_target_length=4096

log

xprof

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]

Caution

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

[codecov]

Codecov Report

❌ Patch coverage is 94.44444% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/pre_train/train.py	80.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Caution

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