From e03e2a0fb93117d2c9318d65ecf4b3b7703a1d11 Mon Sep 17 00:00:00 2001 From: chengnuojin Date: Fri, 4 Sep 2026 17:48:30 +0000 Subject: [PATCH 1/2] Trace the engine's eval kernel under the mesh and the axis rules `fwd_bwd` and `update` wrap their kernel calls in `_sharding_ctx()`; `eval_step` was the one step path that did not, so the eval kernel was traced with no mesh in context and an empty logical axis rule set. Under `shard_mode=auto` that only costs partitioning quality, which is why it went unnoticed. Under explicit axis types it is fatal: the MaxText layers call `jax.sharding.reshard(x, P(...))`, and a bare `PartitionSpec` outside a mesh context is an error rather than a no-op, so `eval_step` raises ValueError: Using PartitionSpec when you are not under a mesh context is not allowed. Reproduced on qwen3.5-35b-a3b (v7-8, ep=8, shard_mode=explicit): eval_step died on the first call after training. With the context entered, the same run completes and eval costs 139.9 ms median per step -- against 1327.4 ms for the identical run under `auto`, since the eval kernel now gets the partitioning the training kernels already had. Wrapped around the call rather than around `_compile_eval_for_batch`, because `jax.jit` is lazy: the trace happens at the call, which is the pattern `fwd_bwd` and `update` follow. --- src/maxtext/training_engine/maxtext_engine.py | 8 +++-- .../post_training/unit/maxtext_engine_test.py | 36 +++++++++++++++++++ 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 9153180a85..3ac4600b27 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -1554,9 +1554,13 @@ def eval_step(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> N signature = _batch_signature(dynamic_batch, static_batch) if self._compiled_eval is None or self._needs_recompile(signature, self._compiled_eval_signature): self._compile_eval_for_batch(dynamic_batch, static_batch) - loss, aux = self._compiled_eval(params, rest, dynamic_batch) + # Around the call, as `fwd_bwd` and `update` do: `jax.jit` is lazy, so this is where + # the eval kernel is traced and where the mesh and the axis rules have to be live. + with self._sharding_ctx(): + loss, aux = self._compiled_eval(params, rest, dynamic_batch) else: - loss, aux = self._eval_kernel(params, rest, batch) + with self._sharding_ctx(): + loss, aux = self._eval_kernel(params, rest, batch) # No metrics attached: eval metrics are buffered by `_eval_metrics_recorder` and written # in EVAL mode when `eval_context` exits. diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py index 60d937218a..154b580ae1 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -1117,6 +1117,42 @@ def test_eval_step_records_eval_metrics_and_mutates_no_training_state(self): # The recorder is drained, so a later pass cannot re-write this one's numbers. self.assertEmpty(t._eval_metrics_recorder.get_metrics_history(clear_cache=False)) + def test_eval_step_traces_under_the_mesh_and_axis_rules(self): + """The eval kernel is traced under the same context every training kernel gets. + + `fwd_bwd` and `update` wrap their kernel calls in `_sharding_ctx`; eval was the one step + path that did not. Under `shard_mode=auto` that only costs partitioning quality, so it + goes unnoticed, but the MaxText layers call `jax.sharding.reshard(x, P(...))` and a bare + `PartitionSpec` with no mesh in context is an error rather than a no-op under explicit + axis types -- there, an eval kernel traced outside the context raises instead of running. + + Both branches: `compile()` picks which of the two eval paths a run takes, and they were + missing the context independently. + """ + for compiled in (False, True): + with self.subTest(compiled=compiled): + seen = {} + + def loss_fn(model, *_args, _seen=seen, **_kwargs): + # Read from inside the kernel, which is the only place that matters: `jax.jit` is + # lazy, so a context entered around the compile call and not the kernel call would + # still leave the trace bare. + _seen["mesh"] = jax.sharding.get_abstract_mesh() + _seen["rules"] = maxtext_engine.nn_partitioning.get_axis_rules() + return ( + abstract_engine.WeightedMetric(unreduced_sum=jnp.sum(model.weights.value), denominator=jnp.array(1.0)), + {}, + ) + + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + t.with_loss_fn(loss_fn) + if compiled: + t.compile(DummyPayload()) + t.eval_step(DummyPayload()) + + self.assertFalse(seen["mesh"].empty, "the eval kernel was traced with no mesh in context") + self.assertTrue(seen["rules"], "the eval kernel was traced with an empty logical axis rule set") + def test_get_metrics_returns_one_buffer_and_a_sentinel_when_empty(self): """`get_metrics` returns a single buffer, matching both ABCs. From 4926b9000ff2092cc7be78734c3dfc2583c85248 Mon Sep 17 00:00:00 2001 From: chengnuojin Date: Fri, 4 Sep 2026 17:52:29 +0000 Subject: [PATCH 2/2] Lift the eval-context probe out of the loop to satisfy pylint W0102: binding the observation dict as a default argument to dodge the loop-closure warning trips dangerous-default-value instead. A helper method that runs one case and returns what it saw has neither problem. --- .../post_training/unit/maxtext_engine_test.py | 42 ++++++++++--------- 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py index 154b580ae1..8c7d006d82 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -1117,6 +1117,28 @@ def test_eval_step_records_eval_metrics_and_mutates_no_training_state(self): # The recorder is drained, so a later pass cannot re-write this one's numbers. self.assertEmpty(t._eval_metrics_recorder.get_metrics_history(clear_cache=False)) + def _sharding_ctx_seen_by_eval(self, compiled: bool) -> dict[str, Any]: + """Runs one `eval_step` and reports the context its kernel was actually traced under.""" + seen = {} + + def loss_fn(model, *_args, **_kwargs): + # Read from inside the kernel, which is the only place that matters: `jax.jit` is lazy, + # so a context entered around the compile call and not the kernel call would still + # leave the trace bare. + seen["mesh"] = jax.sharding.get_abstract_mesh() + seen["rules"] = maxtext_engine.nn_partitioning.get_axis_rules() + return ( + abstract_engine.WeightedMetric(unreduced_sum=jnp.sum(model.weights.value), denominator=jnp.array(1.0)), + {}, + ) + + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + t.with_loss_fn(loss_fn) + if compiled: + t.compile(DummyPayload()) + t.eval_step(DummyPayload()) + return seen + def test_eval_step_traces_under_the_mesh_and_axis_rules(self): """The eval kernel is traced under the same context every training kernel gets. @@ -1131,25 +1153,7 @@ def test_eval_step_traces_under_the_mesh_and_axis_rules(self): """ for compiled in (False, True): with self.subTest(compiled=compiled): - seen = {} - - def loss_fn(model, *_args, _seen=seen, **_kwargs): - # Read from inside the kernel, which is the only place that matters: `jax.jit` is - # lazy, so a context entered around the compile call and not the kernel call would - # still leave the trace bare. - _seen["mesh"] = jax.sharding.get_abstract_mesh() - _seen["rules"] = maxtext_engine.nn_partitioning.get_axis_rules() - return ( - abstract_engine.WeightedMetric(unreduced_sum=jnp.sum(model.weights.value), denominator=jnp.array(1.0)), - {}, - ) - - t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) - t.with_loss_fn(loss_fn) - if compiled: - t.compile(DummyPayload()) - t.eval_step(DummyPayload()) - + seen = self._sharding_ctx_seen_by_eval(compiled) self.assertFalse(seen["mesh"].empty, "the eval kernel was traced with no mesh in context") self.assertTrue(seen["rules"], "the eval kernel was traced with an empty logical axis rule set")