NNX: fail fast for pipeline parallelism under pure_nnx

The NNX decoder has no pipeline path yet, so under pure_nnx the scanned-layers
axis is sharded by 'stage' and dies with a cryptic IndivisibleError at state
init. Raise a clear NotImplementedError at config validation pointing users to
ici_pipeline_parallelism=1 or the Linen path. NNX pipeline support is tracked
as PR11.5.

Author: ecnal-cienet

src/maxtext/common/checkpointing.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer."""
 
 import time
@@ -844,9 +845,7 @@ def map_to_pspec(data):
           (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager),
       ):
         checkpoint_path = str(checkpoint_manager.directory / str(step) / "items")
-        with handle_checkpoint_mismatch(
-            "restore NNX checkpoint", checkpoint_path
-        ):
+        with handle_checkpoint_mismatch("restore NNX checkpoint", checkpoint_path):
           restored_nnx = _load_linen_checkpoint_into_nnx(
               checkpoint_path,
               abstract_unboxed_pre_state,
@@ -882,9 +881,7 @@ def map_to_pspec(data):
                   EmergencyReplicatorCheckpointManager,
               ),
           ):
-            restored = checkpoint_manager.restore(
-                step, args=Composite(state=checkpoint_args)
-            ).state
+            restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state
             _assert_no_shaped_dtype_struct(restored)
             return (
                 restored,
@@ -912,9 +909,7 @@ def map_to_pspec(data):
           # Case 3: Default/Fallback case.
           # This case acts as a wildcard ('_') and matches if none of the preceding cases were met.
           case _:
-            restored = checkpoint_manager.restore(
-                step, args=Composite(items=checkpoint_args)
-            )
+            restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args))
             _assert_no_shaped_dtype_struct(restored)
             return (restored, None)
 
@@ -924,9 +919,7 @@ def map_to_pspec(data):
     else:
       params = abstract_unboxed_pre_state.params
 
-    with handle_checkpoint_mismatch(
-        "load parameters", load_parameters_from_path
-    ):
+    with handle_checkpoint_mismatch("load parameters", load_parameters_from_path):
       restored_params = load_params_from_path(
           load_parameters_from_path,
           params,
@@ -938,9 +931,7 @@ def map_to_pspec(data):
     return None, restored_params
   elif load_full_state_from_path != "":
     max_logging.log(f"Loading full state from path: {load_full_state_from_path}")
-    with handle_checkpoint_mismatch(
-        "load full state", load_full_state_from_path
-    ):
+    with handle_checkpoint_mismatch("load full state", load_full_state_from_path):
       restored_state = _load_full_state_from_path(
           path=load_full_state_from_path,
           abstract_unboxed_pre_state=abstract_unboxed_pre_state,

src/maxtext/configs/types.py

@@ -2784,6 +2784,15 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
 
     self.using_pipeline_parallelism = self.ici_pipeline_parallelism > 1 or self.dcn_pipeline_parallelism > 1
     if self.using_pipeline_parallelism:
+      if self.pure_nnx:
+        # The NNX decoder has no pipeline path yet, so the scanned-layers axis ends up
+        # sharded by 'stage' and fails with a cryptic IndivisibleError at state init.
+        # Fail fast with a clear message instead. NNX pipeline support is tracked as PR11.5.
+        raise NotImplementedError(
+            "Pipeline parallelism is not yet supported on the NNX path. Set "
+            "ici_pipeline_parallelism=1 and dcn_pipeline_parallelism=1, or use the Linen path "
+            "(pure_nnx=False enable_nnx=False)."
+        )
       num_stages = int(self.ici_pipeline_parallelism * self.dcn_pipeline_parallelism)
       if self.num_pipeline_repeats == -1:
         num_pipeline_repeats, remainder = divmod(

tests/unit/train_state_nnx_checkpoint_test.py

@@ -351,7 +351,7 @@ def _build_linen_state(self, num_steps):
   def _invoke_maybe_save(self, state, pure_nnx):
     """Call maybe_save_checkpoint with save_checkpoint patched, return {step, state} captured."""
     # checkpoint_period=1 keeps force_ckpt_save False regardless of actual_step.
-    config = SimpleNamespace(pure_nnx=pure_nnx, checkpoint_period=1, async_checkpointing=False)
+    config = SimpleNamespace(pure_nnx=pure_nnx, checkpoint_period=1, async_checkpointing=False, enable_diloco=False)
     mgr = mock.MagicMock()
     mgr.reached_preemption.return_value = False
 
@@ -417,7 +417,7 @@ def test_maybe_save_checkpoint_skips_if_already_saved(self):
     state = self._build_nnx_state(self.N_STEPS)
     actual_step = self.N_STEPS - 1
 
-    config = SimpleNamespace(pure_nnx=True, checkpoint_period=1, async_checkpointing=False)
+    config = SimpleNamespace(pure_nnx=True, checkpoint_period=1, async_checkpointing=False, enable_diloco=False)
     mgr = mock.MagicMock()
     mgr.reached_preemption.return_value = False
     # Mock latest_step to return the same actual_step
@@ -436,7 +436,7 @@ def test_maybe_save_checkpoint_saves_if_not_already_saved(self):
     state = self._build_nnx_state(self.N_STEPS)
     actual_step = self.N_STEPS - 1
 
-    config = SimpleNamespace(pure_nnx=True, checkpoint_period=1, async_checkpointing=False)
+    config = SimpleNamespace(pure_nnx=True, checkpoint_period=1, async_checkpointing=False, enable_diloco=False)
     mgr = mock.MagicMock()
     mgr.reached_preemption.return_value = False
     # Mock latest_step to return a different step (or None)
@@ -460,9 +460,7 @@ def _nnx_pure(self):
     return {
         "model": {
             "decoder": {"norm": {"scale": jnp.ones((3,))}},
-            "dropout": {
-                "rngs": {"default": {"key": jnp.ones((2,), dtype=jnp.uint32)}}
-            },  # NNX-only
+            "dropout": {"rngs": {"default": {"key": jnp.ones((2,), dtype=jnp.uint32)}}},  # NNX-only
         },
         "optimizer": {
             "step": jnp.asarray(7, dtype=jnp.uint32),
@@ -481,9 +479,7 @@ def test_to_linen_layout(self):
     linen = train_state_nnx.to_linen_checkpoint_dict(self._nnx_pure())
     self.assertEqual(set(linen.keys()), {"params", "step", "opt_state"})
     self.assertIn("params", linen["params"])  # params/params/ collection wrap
-    self.assertNotIn(
-        "dropout", linen["params"]["params"]
-    )  # NNX-only rngs/dropout stripped
+    self.assertNotIn("dropout", linen["params"]["params"])  # NNX-only rngs/dropout stripped
     self.assertEqual(linen["step"].dtype, jnp.int32)  # Linen step is int32
     # opt_state is a list with None for the EmptyState slot, mu/nu wrapped under params.
     self.assertIsInstance(linen["opt_state"], list)
@@ -493,19 +489,11 @@ def test_to_linen_layout(self):
 
   def test_round_trip_preserves_values(self):
     nnx_pure = self._nnx_pure()
-    back = train_state_nnx.from_linen_checkpoint_dict(
-        train_state_nnx.to_linen_checkpoint_dict(nnx_pure)
-    )
+    back = train_state_nnx.from_linen_checkpoint_dict(train_state_nnx.to_linen_checkpoint_dict(nnx_pure))
     self.assertEqual(set(back.keys()), {"model", "optimizer"})
-    self.assertEqual(
-        back["optimizer"]["step"].dtype, jnp.uint32
-    )  # NNX step back to uint32
-    self.assertEqual(
-        set(back["optimizer"]["opt_state"].keys()), {0, 2}
-    )  # int-keyed dict, EmptyState dropped
-    self.assertNotIn(
-        "params", back["optimizer"]["opt_state"][0]["mu"]
-    )  # mu/nu unwrapped
+    self.assertEqual(back["optimizer"]["step"].dtype, jnp.uint32)  # NNX step back to uint32
+    self.assertEqual(set(back["optimizer"]["opt_state"].keys()), {0, 2})  # int-keyed dict, EmptyState dropped
+    self.assertNotIn("params", back["optimizer"]["opt_state"][0]["mu"])  # mu/nu unwrapped
     self.assertTrue(
         jnp.array_equal(
             nnx_pure["model"]["decoder"]["norm"]["scale"],