From 78993c5de97aafb4a1ad6aaf256c8b982dedcbc2 Mon Sep 17 00:00:00 2001 From: continuousml Date: Sat, 25 Jul 2026 16:19:33 -0700 Subject: [PATCH] Keep the checkpointable data iterator out of the load-balanced CP reorder setup_train_loop installed the context-parallel load-balance reorder as a bare map(reorder_fn, data_iterator) and passed the result everywhere the original iterator went. A map object supports iteration but loses both the MaxText iterator interface and the concrete iterator type. With ici_context_parallelism > 1 and context_parallel_load_balance (default true), affected runs fail at the first scheduled evaluation because the eval loop calls eval_data_iterator.reset(), or at the first grain checkpoint save or restore because checkpointing dispatches on the concrete iterator type (isinstance branches for RemoteIteratorWrapper, PlaceHolderDataIterator, and iterator lists in checkpointing.py, then data_iterator.local_iterator access). Wrapping the iterator in a forwarding class cannot restore the checkpoint behavior because the dispatch is by isinstance, not attribute presence. Instead, the original iterator now continues to setup_training_state and checkpointing unchanged, and only the batch consumers receive the reordered view: the DataLoader gets _ReorderedDataIterator (per element when the grain pipeline returns an iterator list for 0 < expansion_factor_real_data < 1, matching the list support in data_loader.py), and the eval iterator is wrapped directly since it is never checkpointed and only needs reset() and iteration. Checkpoint restore mutates the iterator in place, so the loader's view tracks restored state. Unit tests cover the reordered view (per-batch reorder, reset forwarding, iteration protocol) and the loader-view builder (single iterator wrapped, lists preserved with each element wrapped around the original iterator). An integration test in setup_train_loop_nnx_test.py pins both installation sites: with ici_context_parallelism=2 and context_parallel_load_balance the DataLoader holds the reordered view around the exact iterator object setup_train_loop returns for checkpointing, that returned iterator stays unwrapped, and the eval iterator is the reordered view. Local checks: DECOUPLE_GCLOUD=TRUE PYTHONPATH=src python3 -m pytest tests/unit/train_utils_test.py -q (20 passed) and tests/unit/checkpointing_test.py tests/unit/grain_utility_test.py tests/unit/multihost_dataloading_test.py -q (27 passed, 5 skipped; the env var stubs ml_goodput_measurement, absent locally and unrelated); the integration test executed end-to-end and passed on CPU with XLA_FLAGS=--xla_force_host_platform_device_count=4 in a venv providing ml_collections; pyink --pyink-indentation=2 --line-length=122 --check on all three files (unchanged); git diff --check clean. --- src/maxtext/utils/train_utils.py | 33 ++++++++++-- .../integration/setup_train_loop_nnx_test.py | 23 ++++++-- tests/unit/train_utils_test.py | 53 +++++++++++++++++++ 3 files changed, 103 insertions(+), 6 deletions(-) diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index 2cbea04aa3..c3999d721c 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -204,6 +204,30 @@ def jit_train_and_eval_step( return p_train_step, p_eval_step +class _ReorderedDataIterator: + """Applies a reorder function to each batch.""" + + def __init__(self, reorder_fn, data_iterator): + self.reorder_fn = reorder_fn + self.data_iterator = data_iterator + + def __iter__(self): + return self + + def __next__(self): + return self.reorder_fn(next(self.data_iterator)) + + def reset(self): + self.data_iterator.reset() + + +def _reorder_data_iterator_for_loader(reorder_fn, data_iterator): + """Wraps the data iterator, or each element of an iterator list, with the reorder view.""" + if isinstance(data_iterator, list): + return [_ReorderedDataIterator(reorder_fn, iterator) for iterator in data_iterator] + return _ReorderedDataIterator(reorder_fn, data_iterator) + + def setup_train_loop(config, recorder, devices=None): """Set up prerequisites for the training loop - @@ -278,6 +302,7 @@ def create_train_state_fn(): raise ValueError("Packing is only supported for load balanced ring attention with context parallelism for GPU.") # Apply reordering wrapper to data iterators if context parallelism is enabled + data_iterator_for_loader = data_iterator with jax.set_mesh(mesh): if context_parallel_size > 1 and config.context_parallel_load_balance: @@ -294,12 +319,14 @@ def create_train_state_fn(): reorder_fn = maxtext_utils.get_reorder_callable( context_parallel_size, config.shard_mode, reorder_strategy, config.hardware ) - data_iterator = map(reorder_fn, data_iterator) + # data_iterator itself stays unwrapped because checkpointing dispatches on + # its concrete type; only batch consumers receive the reordered view. + data_iterator_for_loader = _reorder_data_iterator_for_loader(reorder_fn, data_iterator) if eval_data_iterator: - eval_data_iterator = map(reorder_fn, eval_data_iterator) + eval_data_iterator = _ReorderedDataIterator(reorder_fn, eval_data_iterator) # Create data_loader AFTER reordering wrapper is applied - data_loader = create_dataloader(config, mesh, data_iterator, recorder, rampup_manager) + data_loader = create_dataloader(config, mesh, data_iterator_for_loader, recorder, rampup_manager) state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state( data_iterator, config, mesh, checkpoint_manager, init_state_fn diff --git a/tests/integration/setup_train_loop_nnx_test.py b/tests/integration/setup_train_loop_nnx_test.py index f446c8f9df..7a26cec21b 100644 --- a/tests/integration/setup_train_loop_nnx_test.py +++ b/tests/integration/setup_train_loop_nnx_test.py @@ -31,7 +31,7 @@ from maxtext.common import train_state_nnx from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT -from maxtext.utils.train_utils import setup_train_loop +from maxtext.utils import train_utils from tests.utils.test_helpers import get_test_config_path import pytest @@ -83,7 +83,7 @@ def test_pure_nnx_setup_returns_train_state_nnx(self): rampup_manager, eval_data_iterator, train_state, - ) = setup_train_loop(config, recorder=None) + ) = train_utils.setup_train_loop(config, recorder=None) # The NNX path returns a fully-merged TrainStateNNX (lines 352-354 in train_utils.py). self.assertIsInstance(train_state, train_state_nnx.TrainStateNNX) @@ -109,13 +109,30 @@ def test_pure_nnx_setup_returns_train_state_nnx(self): # flag them as unused — they're part of the public return contract. del checkpoint_manager, rampup_manager, eval_data_iterator + def test_load_balanced_cp_keeps_checkpoint_iterator_unwrapped(self): + """The reordered view goes to the DataLoader and eval; the original iterator goes to checkpointing.""" + config = _tiny_nnx_pyconfig( + ici_context_parallelism=2, + context_parallel_load_balance=True, + packing=False, + eval_interval=1, + ) + + *_, data_iterator, data_loader, _, eval_data_iterator, _ = train_utils.setup_train_loop(config, recorder=None) + + # pylint: disable=protected-access + self.assertIsInstance(data_loader.data_iterator, train_utils._ReorderedDataIterator) + self.assertIs(data_loader.data_iterator.data_iterator, data_iterator) + self.assertNotIsInstance(data_iterator, train_utils._ReorderedDataIterator) + self.assertIsInstance(eval_data_iterator, train_utils._ReorderedDataIterator) + def test_pure_nnx_setup_param_only_split_matches_model(self): """nnx.split(state.model, nnx.Param, ...) must yield a non-empty Param tree whose structure matches state_mesh_shardings.model after the same split. """ config = _tiny_nnx_pyconfig() - *_, state_mesh_shardings, model, _, _, _, _, _, _, train_state = setup_train_loop(config, recorder=None) + *_, state_mesh_shardings, model, _, _, _, _, _, _, train_state = train_utils.setup_train_loop(config, recorder=None) _, params, _ = nnx.split(train_state.model, nnx.Param, ...) _, params_shardings, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) diff --git a/tests/unit/train_utils_test.py b/tests/unit/train_utils_test.py index d7857c24b2..e0eb5d83a5 100644 --- a/tests/unit/train_utils_test.py +++ b/tests/unit/train_utils_test.py @@ -58,6 +58,59 @@ class MockConfig: use_iota_embed: bool = False +class _FakeIterator: + """Minimal resettable iterator for reorder-view tests.""" + + def __init__(self, batches): + self._batches = batches + self._index = 0 + self.reset_calls = 0 + + def __next__(self): + batch = self._batches[self._index] + self._index += 1 + return batch + + def reset(self): + self.reset_calls += 1 + self._index = 0 + + +class TestReorderedDataIterator(unittest.TestCase): + """Tests for the load-balanced CP reorder view.""" + + # pylint: disable=protected-access + + def test_reorders_batches_and_forwards_reset(self): + inner = _FakeIterator([1, 2]) + wrapped = train_utils._ReorderedDataIterator(lambda batch: batch * 10, inner) + + self.assertEqual(next(wrapped), 10) + self.assertEqual(next(wrapped), 20) + self.assertIs(iter(wrapped), wrapped) + wrapped.reset() + self.assertEqual(inner.reset_calls, 1) + self.assertEqual(next(wrapped), 10) + + def test_loader_view_wraps_single_iterator(self): + inner = _FakeIterator([1]) + view = train_utils._reorder_data_iterator_for_loader(lambda batch: batch * 10, inner) + + self.assertIsInstance(view, train_utils._ReorderedDataIterator) + self.assertIs(view.data_iterator, inner) + + def test_loader_view_wraps_each_list_element(self): + inners = [_FakeIterator([1]), _FakeIterator([2])] + view = train_utils._reorder_data_iterator_for_loader(lambda batch: batch * 10, inners) + + self.assertIsInstance(view, list) + self.assertEqual(len(view), 2) + for wrapped, inner in zip(view, inners): + self.assertIsInstance(wrapped, train_utils._ReorderedDataIterator) + self.assertIs(wrapped.data_iterator, inner) + self.assertEqual(next(view[1]), 20) + + class TestValidateTrainConfig(unittest.TestCase): """Tests for validate_train_config."""