Split logical rule for eval from for train by NuojCheng · Pull Request #4523 · AI-Hypercomputer/maxtext

NuojCheng · 2026-07-17T18:26:05Z

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.

This PR only supports auto sharding. For explicit sharding, the logical rule split is way more complexed. I would defer the explicit sharding support to 2nd PR.

FIXES: b/535639231

Tests

Example 1: base rule for train, custom rule for eval

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

Example 2: custom rule for train, another custom rule for eval

Run on v5p-8: we use EP-as-DP 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  custom_mesh_and_rule=ep-as-dp sharding_tolerance=0.5 override_model_config=true base_num_decoder_layers=7

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 · 2026-07-17T18:26:14Z

Caution

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

codecov · 2026-07-17T18:33:46Z

Codecov Report

❌ Patch coverage is 92.59259% with 2 lines in your changes missing coverage. Please review.

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

📢 Thoughts on this report? Let us know!

gemini-code-assist · 2026-07-17T23:01:12Z

Caution

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

github-actions · 2026-07-20T20:27:15Z

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions

📋 Review Summary

This pull request introduces the ability to split logical axis rules between training and evaluation phases, allowing specialized sharding strategies (such as CP/TP) specifically for evaluation under fractional batch sizes. The overall architecture is very well thought out, but several correctness bugs in the fallback validation, active rule detection, and evaluation input sharding should be addressed to ensure seamless operation.

🔍 General Feedback

Input Sharding Consistency: Ensure that other data loaders or input processing utilities also correctly pick up the active logical_axis_rules_for_eval during evaluation to prevent layout mismatches.
Model-Level Rule Lookups: In addition to Qwen3, double check that other models (like DeepSeek or custom layer architectures) have been migrated to call get_logical_axis_rules() dynamically rather than accessing self.config.logical_axis_rules directly, to ensure they respect the eval-specific sharding rules.
Test Coverage: Consider adding a unit test validating that when custom_mesh_and_rule_for_eval is left as default, it successfully falls back and matches the main training logical rules as expected.

github-actions · 2026-07-20T20:31:39Z

+  """
+  if config.using_pipeline_parallelism or config.eval_interval > 1:
+    return None
+  return config.logical_axis_rules


🟠 Checking config.eval_interval > 1 is a brittle proxy for determining whether evaluation and training rule sets differ. If a user runs evaluation with eval_interval = 1, this check will evaluate to False and return config.logical_axis_rules (the training rules), overriding and ignoring the custom evaluation rules set in the active context manager during evaluation.

Comparing the rule lists directly or checking if evaluation rules are defined is much more robust and correct.

Suggested change

return config.logical_axis_rules

if config.using_pipeline_parallelism or (

config.logical_axis_rules_for_eval is not None

and config.logical_axis_rules_for_eval != config.logical_axis_rules

):

I definitly agree that eval_interval > 1 is not the best way for selecting logical rules. I have tried using flax meta data to collect logical rules but it gets some unit test failures. I will be working on b/536927795 to deprecate this usage as a followup step.

Are you trying to distinguish if eval is turned on? If so, wondering if eval_per_device_batch_size a better indicator. I see default is 0:

maxtext/src/maxtext/configs/base.yml

Line 696 in 2ce3866

eval_per_device_batch_size: 0.0

.

In MLPerf 6.0 runs, we set eval_interval=1. Or maybe you could check if eval_interval as -1:

maxtext/src/maxtext/configs/base.yml

Line 986 in 2ce3866

eval_interval: -1 # the specific number of train step between eval_step

Thank you for the suggestion! Updated using eval_per_device_batch_size > 0 now.

The logic in

maxtext/src/maxtext/configs/types.py

Line 2811 in bff4616

if getattr(self, "eval_per_device_batch_size", 0.0) == 0.0:

sets eval pdb same with pdb. I will use eval_interval != -1 instead.

github-actions · 2026-07-20T20:31:39Z

      if 0 < eval_steps <= eval_step_count:
        break
-      with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules):
+      with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules_for_eval):


🟠 During evaluation, the batch eval_batch is sharded on line 728 via sharding.get_input_data_sharding(config, mesh), which uses config.logical_axis_rules (the training rules) by default. Since the evaluation loop runs inside the context manager nn_partitioning.axis_rules(logical_axis_rules_for_eval), this mismatch will cause costly resharding/communication overhead or compilation failures.

To resolve this, update sharding.get_input_data_sharding to accept an optional rules=None parameter:

def get_input_data_sharding(config, mesh, rules=None): if rules is None: rules = config.logical_axis_rules # ... Use rules instead of config.logical_axis_rules

And then pass logical_axis_rules_for_eval here:

eval_batch = jax.device_put( eval_batch, sharding.get_input_data_sharding(config, mesh, rules=logical_axis_rules_for_eval) )

Good point. Updated

I feel like the Gemini PR review shows better quality now.

Yes I agree!

github-actions · 2026-07-20T20:31:39Z

+# Set as true a different logical rule is used for eval step
+custom_mesh_and_rule_for_eval: ""


🟢 The comment says # Set as true a different logical rule is used for eval step, but custom_mesh_and_rule_for_eval is a string identifier for custom rules (like "ep-as-cp") rather than a boolean. Clarify the comment to prevent user confusion.

Suggested change

# Set as true a different logical rule is used for eval step

custom_mesh_and_rule_for_eval: ""

# Customized mesh and logical rules for evaluation (e.g. "ep-as-cp")

custom_mesh_and_rule_for_eval: ""

github-actions · 2026-07-20T20:31:39Z

+    # Handle eval custom mesh and rule
+    if self.custom_mesh_and_rule_for_eval is CustomRule.DEFAULT:
+      # Fallback to primary rule if eval is DEFAULT
+      self.custom_mesh_and_rule_for_eval = self.custom_mesh_and_rule
+    else:
+      eval_config = self._load_mesh_config_from_yaml(self.custom_mesh_and_rule_for_eval.value)
+      if "logical_axis_rules" in eval_config:
+        self.logical_axis_rules_for_eval = eval_config["logical_axis_rules"]


🔴 The fallback logic is incomplete when custom_mesh_and_rule_for_eval is DEFAULT. While self.custom_mesh_and_rule_for_eval is correctly updated to copy self.custom_mesh_and_rule, self.logical_axis_rules_for_eval is never updated and remains empty []. This will cause evaluation steps to use empty sharding rules instead of falling back to the training sharding rules, leading to compilation failures or incorrect device sharding.

Suggested change

# Handle eval custom mesh and rule

if self.custom_mesh_and_rule_for_eval is CustomRule.DEFAULT:

# Fallback to primary rule if eval is DEFAULT

self.custom_mesh_and_rule_for_eval = self.custom_mesh_and_rule

else:

eval_config = self._load_mesh_config_from_yaml(self.custom_mesh_and_rule_for_eval.value)

if "logical_axis_rules" in eval_config:

self.logical_axis_rules_for_eval = eval_config["logical_axis_rules"]

# Handle eval custom mesh and rule

if self.custom_mesh_and_rule_for_eval is CustomRule.DEFAULT:

# Fallback to primary rule if eval is DEFAULT

self.custom_mesh_and_rule_for_eval = self.custom_mesh_and_rule

self.logical_axis_rules_for_eval = self.logical_axis_rules

else:

eval_config = self._load_mesh_config_from_yaml(self.custom_mesh_and_rule_for_eval.value)

if "logical_axis_rules" in eval_config:

self.logical_axis_rules_for_eval = eval_config["logical_axis_rules"]

else:

self.logical_axis_rules_for_eval = self.logical_axis_rules

RissyRan

Thanks for the change! A few minor comments for naming.

I didn't see you test feature for logical_axis_rules_for_eval. Do you forsee potential conflict when applying separately?

RissyRan · 2026-07-20T23:43:08Z

  config = immutable_data["config"]  # for helpers
-  logical_axis_rules = immutable_data["logical_axis_rules"]
+  logical_axis_rules_train = immutable_data["logical_axis_rules_train"]
+  logical_axis_rules_for_eval = immutable_data["logical_axis_rules_for_eval"]


Shall we align the name from logical_axis_rules_for_eval to logical_axis_rules_eval?

I updated the names to

logical_axis_rules_for_train

logical_axis_rules_for_eval
for consistency. Good point!

RissyRan · 2026-07-20T23:44:12Z

      if 0 < eval_steps <= eval_step_count:
        break
-      with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules):
+      with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules_for_eval):


I feel like the Gemini PR review shows better quality now.

RissyRan · 2026-07-20T23:47:41Z

    train_step_partial = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings)
    train_step = diloco.build_diloco_train_step(config, train_step_partial, mesh=mesh)
-  data_sharding = sharding.get_input_data_sharding(config, mesh)
+  train_data_sharding = sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules)


I see some misalignment with logical_axis_rules_train and logical_axis_rules. Do you think we could let them aligned? i.e. we could use logical_axis_rules for training stage if that makes things simple?

Agree! Now they are

data_sharding_for_train

data_sharding_for_eval

RissyRan · 2026-07-20T23:51:08Z

+  """
+  if config.using_pipeline_parallelism or config.eval_interval > 1:
+    return None
+  return config.logical_axis_rules


Are you trying to distinguish if eval is turned on? If so, wondering if eval_per_device_batch_size a better indicator. I see default is 0:

maxtext/src/maxtext/configs/base.yml

Line 696 in 2ce3866

eval_per_device_batch_size: 0.0

.

In MLPerf 6.0 runs, we set eval_interval=1. Or maybe you could check if eval_interval as -1:

maxtext/src/maxtext/configs/base.yml

Line 986 in 2ce3866

eval_interval: -1 # the specific number of train step between eval_step

NuojCheng · 2026-07-22T00:46:51Z

Regarding to question

I didn't see you test feature for logical_axis_rules_for_eval. Do you forsee potential conflict when applying separately?

I added an integration test protecting this feature. Thank you for bringing it up!

NuojCheng force-pushed the chengnuojin-train-eval-mesh branch 3 times, most recently from 1972c3b to 7e78b58 Compare July 17, 2026 22:17

NuojCheng marked this pull request as ready for review July 17, 2026 23:01

NuojCheng requested review from RissyRan, gagika, michelle-yooh, richjames0, shralex and suexu1025 as code owners July 17, 2026 23:01

NuojCheng requested review from A9isha, SurbhiJainUSC, abhinavclemson, aireenmei, bvandermoon, darisoy, dipannita08, gobbleturk, hengtaoguo, huytransformer, igorts-git, jiangjy1982, khatwanimohit, parambole, shuningjin, vipannalla and xibinliu as code owners July 17, 2026 23:01

NuojCheng force-pushed the chengnuojin-train-eval-mesh branch from 7e78b58 to 77bb651 Compare July 20, 2026 18:28

RissyRan added the gemini-review label Jul 20, 2026

github-actions Bot reviewed Jul 20, 2026

View reviewed changes

NuojCheng force-pushed the chengnuojin-train-eval-mesh branch from 77bb651 to bd59acc Compare July 20, 2026 21:59

RissyRan reviewed Jul 20, 2026

View reviewed changes

NuojCheng force-pushed the chengnuojin-train-eval-mesh branch from bd59acc to 94ca83b Compare July 21, 2026 00:25

gobbleturk approved these changes Jul 21, 2026

View reviewed changes

NuojCheng force-pushed the chengnuojin-train-eval-mesh branch from 94ca83b to 583b6a6 Compare July 21, 2026 20:37

RissyRan approved these changes Jul 22, 2026

View reviewed changes

github-actions Bot added the pull ready label Jul 22, 2026

add separate logical rule for eval

de63b4a

NuojCheng force-pushed the chengnuojin-train-eval-mesh branch from 583b6a6 to de63b4a Compare July 22, 2026 00:45

copybara-service Bot merged commit 085068d into main Jul 22, 2026
83 of 87 checks passed

copybara-service Bot deleted the chengnuojin-train-eval-mesh branch July 22, 2026 17:00

-  return config.logical_axis_rules
+  if config.using_pipeline_parallelism or (
+      config.logical_axis_rules_for_eval is not None
+      and config.logical_axis_rules_for_eval != config.logical_axis_rules
+  ):

		# Set as true a different logical rule is used for eval step
		custom_mesh_and_rule_for_eval: ""

Uh oh!

Conversation

NuojCheng commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Example 1: base rule for train, custom rule for eval

Example 2: custom rule for train, another custom rule for eval

Checklist

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

codecov Bot commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

github-actions Bot commented Jul 20, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

📋 Review Summary

🔍 General Feedback

Uh oh!

github-actions Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NuojCheng Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

RissyRan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NuojCheng commented Jul 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

NuojCheng commented Jul 17, 2026 •

edited

Loading

codecov Bot commented Jul 17, 2026 •

edited

Loading

NuojCheng Jul 21, 2026 •

edited

Loading