diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 513594ad57..47abd80881 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -98,10 +98,7 @@ def get_functional_train_with_signature( """Get the shardings (both state and data) for `train_step`.""" functional_train = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings) functional_train.__name__ = "train_step" # pyrefly: ignore[missing-attribute] - if config.pure_nnx: - in_shardings = (state_mesh_shardings, data_sharding) # State, batch - else: - in_shardings = (state_mesh_shardings, data_sharding, None) # State, batch, rng + in_shardings = (state_mesh_shardings, data_sharding) # State, batch out_shardings = (state_mesh_shardings, None) # State, metrics static_argnums = () # We partial out the static argnums of model and config donate_argnums = 0 # This is the index of the state - we allow the compiler to make use of this memory. @@ -112,10 +109,7 @@ def get_functional_eval_with_signature(eval_step, data_sharding, state_mesh_shar """Get the shardings (both state and data) for `eval_step`.""" functional_eval = functools.partial(eval_step, model, config) functional_eval.__name__ = "eval_step" # pyrefly: ignore[missing-attribute] - if config.pure_nnx: - in_shardings = (state_mesh_shardings, data_sharding) # State, batch (NNX: no rng) - else: - in_shardings = (state_mesh_shardings, data_sharding, None) # State, batch, rng + in_shardings = (state_mesh_shardings, data_sharding) # State, batch out_shardings = None # metrics static_argnums = () # We partial out the static argnums of model, config donate_argnums = () # state will be kept instead of being donated in eval_step @@ -279,11 +273,7 @@ def get_train_input_output_trees(func, input_args, input_kwargs): serialized_compiled = load_serialized_compiled(config.compiled_trainstep_file) shaped_batch = get_shaped_batch(config) - if config.pure_nnx: - shaped_input_args = (state, shaped_batch) - else: - example_rng = jax.random.PRNGKey(0) - shaped_input_args = (state, shaped_batch, example_rng) + shaped_input_args = (state, shaped_batch) shaped_input_kwargs = {} in_tree, out_tree = get_train_input_output_trees(partial_train, shaped_input_args, shaped_input_kwargs) p_train_step = deserialize_and_load(serialized_compiled, in_tree, out_tree, execution_devices=execution_devices) @@ -1879,51 +1869,8 @@ def get_logical_annotations(config, mesh, init_state_fn): def get_abstract_state(config, mesh, init_state_fn, is_training=True): - """Get a shaped abstraction of the state (including optimizer)""" - if config.pure_nnx: - return get_abstract_state_nnx(config, mesh, init_state_fn, is_training) - - init_state_partial = init_state_fn - - with nn_partitioning.axis_rules(config.logical_axis_rules): - abstract_state = jax.eval_shape(init_state_partial) - - state_logical_annotations = nn.get_partition_spec(abstract_state) - - state_mesh_shardings = nn.logical_to_mesh_sharding(state_logical_annotations, mesh, config.logical_axis_rules) - if is_training and config.shard_optimizer_over_data: - # Add data to sharding for optimizer state - state_mesh_shardings = state_mesh_shardings.replace( - opt_state=jax.tree.map_with_path( - functools.partial(sharding.add_data_to_sharding, mesh), - max_utils.unbox_logicallypartioned(abstract_state).opt_state, - state_mesh_shardings.opt_state, - ) - ) - if is_training and config.optimizer_memory_host_offload: - opt_state = jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind="pinned_host"), state_mesh_shardings.opt_state) - state_mesh_shardings = state_mesh_shardings.replace(opt_state=opt_state) - if is_training and config.parameter_memory_host_offload: - assert config.param_scan_axis == 0, "You must set the scan axis 0 to enable parameter offloading." - - def move(path, x): - max_logging.log(f"max_utils.py: Moving {path} to host") - return x.with_memory_kind(kind="pinned_host") - - params = jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params) - state_mesh_shardings = state_mesh_shardings.replace(params=params) - - abstract_sharded_state = jax.jit(init_state_partial, in_shardings=None, out_shardings=state_mesh_shardings).eval_shape() - - unboxed_abstract_sharded_state = max_utils.unbox_logicallypartioned(abstract_sharded_state) - # Initialization - with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): - state_mesh_annotations = nn.logical_to_mesh(state_logical_annotations) - return ( - unboxed_abstract_sharded_state, - state_mesh_annotations, - state_mesh_shardings, - ) + """Get a shaped abstraction of the state (including optimizer).""" + return get_abstract_state_nnx(config, mesh, init_state_fn, is_training) def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=True): diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index 41f74a21b6..e3e7b4c26f 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -20,7 +20,6 @@ import optax import functools import orbax.checkpoint.pathways as ocp_pathways -from functools import partial from flax import nnx from flax.linen import partitioning as nn_partitioning @@ -252,32 +251,25 @@ def setup_train_loop(config, recorder, devices=None): from maxtext.input_pipeline.input_pipeline_interface import create_data_iterator with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT): - is_training = True init_rng = jax.random.PRNGKey(config.init_weights_seed) mesh = maxtext_utils.get_mesh_from_config(config, devices) context_parallel_size = mesh.shape.get(config.context_sharding, 1) - if config.pure_nnx: - # Create abstract NNX model. - _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices) - else: - model = model_creation_utils.from_config(config, devices) + # Create abstract NNX model. + _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices) learning_rate_schedule, tx = create_training_optimizer(config, model) - if config.pure_nnx: - # For NNX, the train state is wrapped in the TrainStateNNX module. - def create_train_state_fn(): - model = _create_model_partial() - wrt = ( - getattr(nnx, "LoRAParam", nnx.Param) - if getattr(getattr(config, "lora", None), "enable_lora", False) - else nnx.Param - ) - optimizer = nnx.Optimizer(model, tx, wrt=wrt) - return train_state_nnx.TrainStateNNX(model, optimizer) - - init_state_fn = create_train_state_fn - else: - init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, is_training, init_rng) + # The train state is wrapped in the TrainStateNNX module. + def create_train_state_fn(): + model = _create_model_partial() + wrt = ( + getattr(nnx, "LoRAParam", nnx.Param) + if getattr(getattr(config, "lora", None), "enable_lora", False) + else nnx.Param + ) + optimizer = nnx.Optimizer(model, tx, wrt=wrt) + return train_state_nnx.TrainStateNNX(model, optimizer) + + init_state_fn = create_train_state_fn checkpoint_manager = create_checkpoint_manager(config, mesh, init_state_fn) if checkpoint_manager is not None: checkpoint_step = checkpointing.latest_step(checkpoint_manager) @@ -340,24 +332,23 @@ def create_train_state_fn(): state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state( data_iterator, config, mesh, checkpoint_manager, init_state_fn ) - if config.pure_nnx: - if getattr(getattr(config, "lora", None), "enable_lora", False) and getattr(config.lora, "lora_restore_path", None): - # Restore standalone LoRA adapter weights onto the base model state after initialization. - target_model_state = ( - state["model"] - if (isinstance(state, (nnx.State, dict)) and "model" in state) - else getattr(state, "model", state) - ) - # pyrefly: ignore[bad-argument-type] - lora_utils.restore_lora_from_path(target_model_state, config) - _, _, state_mesh_shardings = maxtext_utils.get_abstract_state_nnx(config, mesh, init_state_fn, True) - with nn_partitioning.axis_rules(config.logical_axis_rules): - # We only need the graphdef here; it's merged with state below. Avoid - # nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable - # under jax.set_mesh(mesh) and rejects any logical name missing from - # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes - # without a mesh skips sharding resolution, so it avoids the crash. - state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn)) + if getattr(getattr(config, "lora", None), "enable_lora", False) and getattr(config.lora, "lora_restore_path", None): + # Restore standalone LoRA adapter weights onto the base model state after initialization. + target_model_state = ( + state["model"] + if (isinstance(state, (nnx.State, dict)) and "model" in state) + else getattr(state, "model", state) + ) + # pyrefly: ignore[bad-argument-type] + lora_utils.restore_lora_from_path(target_model_state, config) + _, _, state_mesh_shardings = maxtext_utils.get_abstract_state_nnx(config, mesh, init_state_fn, True) + with nn_partitioning.axis_rules(config.logical_axis_rules): + # We only need the graphdef here; it's merged with state below. Avoid + # nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable + # under jax.set_mesh(mesh) and rejects any logical name missing from + # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes + # without a mesh skips sharding resolution, so it avoids the crash. + state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn)) if isinstance(state, diloco.DiLoCoTrainState): state_params = state.params @@ -365,13 +356,10 @@ def create_train_state_fn(): _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) else: state_mesh_shardings_params = state_mesh_shardings.params - elif config.pure_nnx: + else: with nn_partitioning.axis_rules(config.logical_axis_rules): _, state_params, _ = nnx.split(state.model, nnx.Param, ...) _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) - else: - state_params = state.params - state_mesh_shardings_params = state_mesh_shardings.params if config.enable_diloco: with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): @@ -409,28 +397,21 @@ def create_train_state_fn(): # print weights sharding info under debug sharding mode if config.debug_sharding: - if config.pure_nnx: - # TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we - # probably already lost the logical partition info at this moment. - logical_annotations_params = None - else: - logical_annotations = maxtext_utils.get_logical_annotations(config, mesh, init_state_fn) - logical_annotations_params = logical_annotations.params + # TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we + # probably already lost the logical partition info at this moment. + logical_annotations_params = None max_utils.print_non_trivial_mesh_axis(model.mesh) # pyrefly: ignore[missing-attribute] maxtext_utils.print_shardings_params(state_params, state_mesh_shardings_params, mesh, logical_annotations_params) - if config.pure_nnx: - if config.enable_diloco: - # Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner - # train step needs that graphdef as jit_model; the wrapper passes through as state. - train_state = state - model = state_graphdef # pyrefly: ignore[unbound-name] - else: - train_state = nnx.merge(state_graphdef, state) # pyrefly: ignore[unbound-name] - model = train_state.model - else: + if config.enable_diloco: + # Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner + # train step needs that graphdef as jit_model; the wrapper passes through as state. train_state = state + model = state_graphdef # pyrefly: ignore[unbound-name] + else: + train_state = nnx.merge(state_graphdef, state) # pyrefly: ignore[unbound-name] + model = train_state.model return ( init_rng, diff --git a/tests/integration/setup_train_loop_nnx_test.py b/tests/integration/setup_train_loop_nnx_test.py index 7a26cec21b..a3b9f31a91 100644 --- a/tests/integration/setup_train_loop_nnx_test.py +++ b/tests/integration/setup_train_loop_nnx_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Integration test for setup_train_loop with pure_nnx=True. +"""Integration test for setup_train_loop on the NNX path. setup_train_loop wires together create_nnx_abstract_model, the training optimizer, @@ -43,7 +43,6 @@ def _tiny_nnx_pyconfig(**overrides): "enable_checkpointing": False, "dataset_type": "synthetic", "model_name": "default", - "pure_nnx": True, "per_device_batch_size": 1.0, "base_emb_dim": 8, "base_num_query_heads": 4, @@ -68,7 +67,7 @@ def _tiny_nnx_pyconfig(**overrides): class SetupTrainLoopNNXIntegrationTest(unittest.TestCase): """End-to-end check that setup_train_loop returns a usable TrainStateNNX.""" - def test_pure_nnx_setup_returns_train_state_nnx(self): + def test_setup_returns_train_state_nnx(self): config = _tiny_nnx_pyconfig() ( @@ -126,7 +125,7 @@ def test_load_balanced_cp_keeps_checkpoint_iterator_unwrapped(self): 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): + def test_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. diff --git a/tests/unit/correctness_tests_nnx_dispatch_test.py b/tests/unit/correctness_tests_nnx_dispatch_test.py index 163a5a836e..627980bc6a 100644 --- a/tests/unit/correctness_tests_nnx_dispatch_test.py +++ b/tests/unit/correctness_tests_nnx_dispatch_test.py @@ -18,7 +18,7 @@ exercises the changed dispatch code that otherwise has no CPU coverage: - `mt.from_config` is exported (the GRPO trainer calls it) - the SFT correctness test's `setup_maxtext_model` / `get_maxtext_logits` run on - both paths (pure_nnx=True -> NNX, pure_nnx=False -> Linen) and stay finite + the NNX path and stay finite The GRPO NNX building blocks the other dispatch helpers call (`compute_log_probs_nnx`, `grpo_loss_fn_nnx`) are already covered by grpo_nnx_test. @@ -54,15 +54,12 @@ } -def _sft_config(pure_nnx): +def _sft_config(): return pyconfig.initialize( [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs/post_train", "sft.yml")], - run_name=f"unit-sft-{pure_nnx}", + run_name="unit-sft-nnx", model_name="default", enable_checkpointing=False, - pure_nnx=pure_nnx, - enable_nnx=pure_nnx, - pure_nnx_decoder=pure_nnx, **_SMALL, ) @@ -84,12 +81,7 @@ def test_from_config_is_exported(self): self.assertTrue(hasattr(mt, "from_config")) def test_sft_logits_nnx_path(self): - config = _sft_config(pure_nnx=True) - logits = sft.get_maxtext_logits(config, _fake_data(config)) - self.assertTrue(bool(jnp.isfinite(logits).all())) - - def test_sft_logits_linen_path(self): - config = _sft_config(pure_nnx=False) + config = _sft_config() logits = sft.get_maxtext_logits(config, _fake_data(config)) self.assertTrue(bool(jnp.isfinite(logits).all())) diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 3256230694..cb56501df3 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -16,7 +16,6 @@ from collections.abc import Callable from dataclasses import dataclass, field -import functools from types import SimpleNamespace from typing import Any, Sequence import unittest @@ -31,11 +30,9 @@ import jax.numpy as jnp from jax.sharding import AxisType, Mesh, NamedSharding, PartitionSpec from maxtext.common import train_state_nnx -from maxtext.common.common_types import DecoderBlockType, MODEL_MODE_TRAIN, ShardMode from maxtext.configs import pyconfig +from maxtext.common.common_types import DecoderBlockType, ShardMode from maxtext.inference import inference_utils -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import maxtext_utils_nnx @@ -47,8 +44,6 @@ import optax import pytest -Transformer = models.transformer_as_linen - class TestGradientClipping(unittest.TestCase): """test class for gradient clipping""" @@ -363,50 +358,29 @@ def setUp(self): self.config = pyconfig.initialize([None, get_test_config_path()], enable_checkpointing=False) devices_array = maxtext_utils.create_device_mesh(self.config) self.mesh = Mesh(devices_array, self.config.mesh_axes) - quant = quantizations.configure_quantization(self.config) - if self.config.pure_nnx: - self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - else: - self.model = models.transformer_as_linen(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) def test_setup_decode_state(self): - rng = random.PRNGKey(0) - if self.config.pure_nnx: + def create_train_state_fn(): + nnx_model = self._create_model_partial() + return train_state_nnx.TrainStateNNX(nnx_model, None) - def create_train_state_fn(): - nnx_model = self._create_model_partial() - return train_state_nnx.TrainStateNNX(nnx_model, None) - - init_state_fn = create_train_state_fn - else: - init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng) + init_state_fn = create_train_state_fn state, _ = maxtext_utils.setup_decode_state(self.config, self.mesh, None, init_state_fn) - if self.config.pure_nnx: - self.assertNotIn("optimizer", state) - else: - self.assertEqual(state.tx, None) - self.assertEqual(state.opt_state, {}) + self.assertNotIn("optimizer", state) def test_setup_initial_state(self): - rng = random.PRNGKey(0) tx = optax.adam(learning_rate=0.001) - if self.config.pure_nnx: - def create_train_state_fn(): - nnx_model = self._create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + def create_train_state_fn(): + nnx_model = self._create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - init_state_fn = create_train_state_fn - else: - init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, tx, self.config, True, rng) + init_state_fn = create_train_state_fn state, _, _, _, was_restored = maxtext_utils.setup_initial_state(None, self.config, self.mesh, None, init_state_fn) self.assertFalse(was_restored) - if self.config.pure_nnx: - self.assertIsNotNone(state.optimizer) - else: - self.assertEqual(state.tx, tx) - self.assertNotEqual(state.opt_state, {}) + self.assertIsNotNone(state.optimizer) class MaxUtilsPpAsDp(unittest.TestCase): @@ -1045,9 +1019,8 @@ def train_step(_model, _config, _state_shardings, _params_shardings, state, _bat return train_step - def _make_mock_config(self, pure_nnx=False): + def _make_mock_config(self): cfg = MagicMock() - cfg.pure_nnx = pure_nnx return cfg def test_returns_five_tuple(self): @@ -1064,20 +1037,11 @@ def test_functional_train_has_correct_name(self): ) self.assertEqual(fn.__name__, "train_step") - def test_linen_in_shardings_includes_rng(self): - """pure_nnx=False: in_shardings should be (state, batch, rng).""" - step = self._make_mock_step() - _, in_shardings, _, _, _ = maxtext_utils.get_functional_train_with_signature( - step, "data_sharding", "state_shardings", "model", self._make_mock_config(pure_nnx=False) - ) - self.assertEqual(len(in_shardings), 3) - self.assertIsNone(in_shardings[2]) # rng sharding is None - def test_nnx_in_shardings_excludes_rng(self): - """pure_nnx=True: in_shardings should be (state, batch) — no rng slot.""" + """in_shardings should be (state, batch) — no rng slot.""" step = self._make_mock_step() _, in_shardings, _, _, _ = maxtext_utils.get_functional_train_with_signature( - step, "data_sharding", "state_shardings", "model", self._make_mock_config(pure_nnx=True) + step, "data_sharding", "state_shardings", "model", self._make_mock_config() ) self.assertEqual(len(in_shardings), 2) @@ -1113,9 +1077,8 @@ def eval_step(_model, _config, _state, _batch, _rng=None): return eval_step - def _make_mock_config(self, pure_nnx=False): + def _make_mock_config(self): cfg = MagicMock() - cfg.pure_nnx = pure_nnx return cfg def test_returns_five_tuple(self): @@ -1143,21 +1106,13 @@ def test_donate_argnums_is_empty(self): self.assertEqual(donate_argnums, ()) def test_nnx_in_shardings_excludes_rng(self): - """pure_nnx=True: in_shardings should be (state, batch) — no rng slot.""" + """in_shardings should be (state, batch) — no rng slot.""" step = self._make_mock_eval_step() _, in_shardings, _, _, _ = maxtext_utils.get_functional_eval_with_signature( - step, "batch_sharding", "state_sharding", "model", self._make_mock_config(pure_nnx=True) + step, "batch_sharding", "state_sharding", "model", self._make_mock_config() ) self.assertEqual(len(in_shardings), 2) - def test_linen_in_shardings_includes_rng(self): - """pure_nnx=False: in_shardings should be (state, batch, rng).""" - step = self._make_mock_eval_step() - _, in_shardings, _, _, _ = maxtext_utils.get_functional_eval_with_signature( - step, "batch_sharding", "state_sharding", "model", self._make_mock_config(pure_nnx=False) - ) - self.assertEqual(len(in_shardings), 3) - class TestGetShapedBatch(unittest.TestCase): """Tests for get_shaped_batch.""" @@ -1472,39 +1427,20 @@ def setUp(self): self.config = pyconfig.initialize([None, get_test_config_path()], enable_checkpointing=False) devices_array = maxtext_utils.create_device_mesh(self.config) self.mesh = Mesh(devices_array, self.config.mesh_axes) - quant = quantizations.configure_quantization(self.config) - if self.config.pure_nnx: - self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - else: - self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) def test_setup_training_state_returns_train_state(self): - rng = jax.random.PRNGKey(0) tx = optax.adam(learning_rate=0.001) - if self.config.pure_nnx: - def create_train_state_fn(): - nnx_model = self._create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + def create_train_state_fn(): + nnx_model = self._create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - init_state_fn = create_train_state_fn - else: - init_state_fn = functools.partial( - maxtext_utils.init_initial_state, - self.model, - tx, - self.config, - True, - rng, - ) + init_state_fn = create_train_state_fn state, _, _, _, was_restored = maxtext_utils.setup_training_state(None, self.config, self.mesh, None, init_state_fn) self.assertFalse(was_restored) - if self.config.pure_nnx: - self.assertIsNotNone(state.optimizer) - else: - self.assertEqual(state.tx, tx) - self.assertNotEqual(state.opt_state, {}) + self.assertIsNotNone(state.optimizer) class TestGetLogicalAnnotations(unittest.TestCase): @@ -1514,36 +1450,20 @@ def setUp(self): self.config = pyconfig.initialize([None, get_test_config_path()], enable_checkpointing=False) devices_array = maxtext_utils.create_device_mesh(self.config) self.mesh = Mesh(devices_array, self.config.mesh_axes) - quant = quantizations.configure_quantization(self.config) - if self.config.pure_nnx: - self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - else: - self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) self.rng = jax.random.PRNGKey(0) self.tx = optax.adam(learning_rate=0.001) def test_returns_partition_spec_tree(self): - if self.config.pure_nnx: - - def create_train_state_fn(): - nnx_model = self._create_model_partial() - optimizer = nnx.Optimizer(nnx_model, self.tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - - init_state_fn = create_train_state_fn - annotations = maxtext_utils_nnx.get_partition_spec_nnx( - maxtext_utils.get_abstract_state(self.config, self.mesh, init_state_fn, True)[2] - ) - else: - init_state_fn = functools.partial( - maxtext_utils.init_initial_state, - self.model, - self.tx, - self.config, - True, - self.rng, - ) - annotations = maxtext_utils.get_logical_annotations(self.config, self.mesh, init_state_fn) + def create_train_state_fn(): + nnx_model = self._create_model_partial() + optimizer = nnx.Optimizer(nnx_model, self.tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + + init_state_fn = create_train_state_fn + annotations = maxtext_utils_nnx.get_partition_spec_nnx( + maxtext_utils.get_abstract_state(self.config, self.mesh, init_state_fn, True)[2] + ) # Result should be a pytree with PartitionSpec leaves leaves = jax.tree_util.tree_leaves(annotations) self.assertGreater(len(leaves), 0) diff --git a/tests/unit/sharding_compare_test.py b/tests/unit/sharding_compare_test.py index aaefcb9020..2331d54cf3 100644 --- a/tests/unit/sharding_compare_test.py +++ b/tests/unit/sharding_compare_test.py @@ -12,325 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Compare expected sharding of models with actual sharding of models.""" +"""Compare expected sharding of models with actual sharding of models. -import functools -import hashlib -import json -import os -import jax -import jax.numpy as jnp -from maxtext.configs import pyconfig -from maxtext.utils import maxtext_utils -from maxtext.utils.sharding import clear_input_shardings_dump -# import optax - -from maxtext.layers import quantizations -from maxtext.models import models -from maxtext.optimizers import optimizers -from maxtext.trainers.pre_train.train_compile import get_shaped_inputs, get_topology_mesh, validate_config -from tests.utils.sharding_dump import TEST_CASES, load_json, input_sharding_to_json, named_shardings_to_json, partition_specs_to_json -from tests.utils.test_helpers import get_test_config_path -import pytest - -Transformer = models.transformer_as_linen - - -def compute_checksum(d: dict) -> str: - """Compute a checksum (SHA256) of a dictionary.""" - # Serialize the dictionary into a JSON string (ensuring consistent ordering of keys) - json_str = json.dumps(d, sort_keys=True) - - # Compute the SHA256 checksum of the serialized string - checksum = hashlib.sha256(json_str.encode("utf-8")).hexdigest() - - return checksum - - -def compare_sharding_jsons(json1: dict, model1_name: str, json2: dict, model2_name: str) -> bool: - """Compare two json files and print the differences if any.""" - keys1 = set(json1.keys()) - keys2 = set(json2.keys()) - - only_in_1 = keys1 - keys2 - only_in_2 = keys2 - keys1 - shared_keys = keys1 & keys2 - - has_diff = False - - if only_in_1: - print(f"Keys only in {model1_name}:") - for k in sorted(only_in_1): - print(f" {k}") - has_diff = True - - if only_in_2: - print(f"Keys only in {model2_name}:") - for k in sorted(only_in_2): - print(f" {k}") - has_diff = True - - for key in sorted(shared_keys): - entry1 = json1[key] - entry2 = json2[key] - - if isinstance(entry1, dict) and isinstance(entry2, dict): - mesh1 = entry1.get("mesh", {}) - mesh2 = entry2.get("mesh", {}) - - spec1 = entry1.get("partition_spec", []) - spec2 = entry2.get("partition_spec", []) - - shape1 = entry1.get("shape") - shape2 = entry2.get("shape") - - if mesh1 != mesh2: - print(f"\nMesh mismatch at '{key}':") - print(f" {model1_name}: {mesh1}") - print(f" {model2_name}: {mesh2}") - has_diff = True - - if spec1 != spec2: - print(f"\nPartitionSpec mismatch at '{key}':") - print(f" {model1_name}: {spec1}") - print(f" {model2_name}: {spec2}") - has_diff = True - - if shape1 != shape2: - print(f"\nShape mismatch at '{key}':") - print(f" {model1_name}: {shape1}") - print(f" {model2_name}: {shape2}") - has_diff = True - - else: - print(f"\nFormat mismatch at '{key}':") - print(f" {model1_name} type: {type(entry1)}") - print(f" {model2_name} type: {type(entry2)}") - has_diff = True - - return has_diff - - -# Requires JAX TPU support to generate the simulated TPU topology. -@pytest.mark.tpu_backend -@pytest.mark.parametrize("model_name, topology, num_slice, custom_mesh_and_rule, overrides", TEST_CASES) -def test_sharding_dump_for_model( - model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple -) -> None: - """ - Test sharding configurations from train_compile.get_shaped_inputs. - This test verifies that the sharding configurations for various models and topologies remain consistent with golden files. - """ - params = [ - "/deps/MaxText/tests/unit/sharding_compare_test", - get_test_config_path(), - f"compile_topology={topology}", - f"compile_topology_num_slices={num_slice}", - f"model_name={model_name}", - "log_config=false", - "debug_sharding=true", # for input sharding dump - "pure_nnx=False", - "enable_nnx=False", - "pure_nnx_decoder=False", - ] - if custom_mesh_and_rule: - params.append(f"custom_mesh_and_rule={custom_mesh_and_rule}") - if overrides: - params.extend(overrides) - - root_dir = "tests/utils/sharding_info" - rule_name = f"rule_{custom_mesh_and_rule}" if custom_mesh_and_rule else "rule_default" - if overrides: - rule_name += "_" + "_".join(overrides) - base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}", rule_name) - - named_json_path = os.path.join(base_path, "named_shardings.json") - logical_json_path = os.path.join(base_path, "logical_shardings.json") - input_json_path = os.path.join(base_path, "input_shardings.json") - - if not os.path.exists(named_json_path): - pytest.skip(f"Missing named_shardings.json for {model_name} {topology} slice {num_slice}") - return - if not os.path.exists(logical_json_path): - pytest.skip(f"Missing logical_shardings.json for {model_name} {topology} slice {num_slice}") - return - if not os.path.exists(input_json_path): - pytest.skip(f"Missing input_shardings.json for {model_name} {topology} slice {num_slice}") - return - - config = pyconfig.initialize(params) - validate_config(config) - - clear_input_shardings_dump() - topology_mesh = get_topology_mesh(config) - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - optimizers.get_optimizer(config, learning_rate_schedule) - shaped_train_args, _, state_mesh_shardings, logical_shardings, _ = get_shaped_inputs(topology_mesh, config) - - error_messages = [] - - # 1. Compare Named Shardings - actual_named = named_shardings_to_json(state_mesh_shardings, shaped_train_args[0]) - expected_named = load_json(named_json_path) - # calculate checksum - actual_named_sum = compute_checksum(actual_named) - expected_named_sum = compute_checksum(expected_named) - named_match = actual_named_sum == expected_named_sum - - if not named_match: - print(f"\n[FAIL] Physical Sharding Mismatch: {model_name} {topology} slice {num_slice}", flush=True) - compare_sharding_jsons(expected_named, "Expected (Physical)", actual_named, "Actual (Physical)") - error_messages.append(f" Physical sharding mismatch for {model_name} on {topology} slice {num_slice}") - - # 2. Compare Logical Shardings - actual_logical = partition_specs_to_json(logical_shardings, shaped_train_args[0]) - expected_logical = load_json(logical_json_path) - # calculate checksum - actual_logical_sum = compute_checksum(actual_logical) - expected_logical_sum = compute_checksum(expected_logical) - logical_match = actual_logical_sum == expected_logical_sum - - if not logical_match: - print(f"\n[FAIL] Logical Sharding Mismatch: {model_name} {topology} slice {num_slice}", flush=True) - compare_sharding_jsons(expected_logical, "Expected (Logical)", actual_logical, "Actual (Logical)") - error_messages.append(f"Logical sharding mismatch for {model_name} on {topology} slice {num_slice}") - - # 3. Compare Input Shardings - actual_input = input_sharding_to_json() - expected_input = load_json(input_json_path) - # calculate checksum - actual_input_sum = compute_checksum(actual_input) - expected_input_sum = compute_checksum(expected_input) - - input_match = actual_input_sum == expected_input_sum - - if not input_match: - print(f"\n[FAIL] Input Sharding Mismatch: {model_name} {topology} slice {num_slice}", flush=True) - # compare_sharding_jsons(expected_input, "Expected (Input)", actual_input, "Actual (Input)") - error_messages.append(f"Input sharding mismatch for {model_name} on {topology} slice {num_slice}") - - assert not error_messages, "\n".join(error_messages) - - -@pytest.fixture( - scope="module", - params=[pytest.param(case, id=f"{case[0]}-{case[1]}-{case[2]}-{case[3]}-{''.join(case[4])}") for case in TEST_CASES], -) -def abstract_state_and_shardings(request): - """Pytest fixture to set up model, config, and generate abstract state once per test case.""" - model_name, topology, num_slice, custom_mesh_and_rule, overrides = request.param - print( - f"Testing model: {model_name}, topology: {topology}, num_slices: {num_slice}, " - "rule: {custom_mesh_and_rule}, overrides: {overrides}", - flush=True, - ) - params = [ - "/deps/MaxText/tests/unit/sharding_compare_test", - get_test_config_path(), - f"compile_topology={topology}", - f"compile_topology_num_slices={num_slice}", - f"model_name={model_name}", - "weight_dtype=float32", - "pure_nnx=False", - "enable_nnx=False", - "pure_nnx_decoder=False", - ] - if custom_mesh_and_rule: - params.append(f"custom_mesh_and_rule={custom_mesh_and_rule}") - if overrides: - params.extend(overrides) - config = pyconfig.initialize(params) - validate_config(config) - - topology_mesh = get_topology_mesh(config) - quant = quantizations.configure_quantization(config) - model = Transformer(config, mesh=topology_mesh, quant=quant) - - learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - # tx = optax.adam(learning_rate=learning_rate_schedule) - tx = optimizers.get_optimizer(config, learning_rate_schedule) - rng = jax.random.PRNGKey(0) - - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, rng) - - # Get abstract state and physical shardings from maxtext_utils - abstract_state, _, state_mesh_shardings = maxtext_utils.get_abstract_state( - config, topology_mesh, init_state_fn, is_training=True - ) - - # Get logical shardings from maxtext_utils - logical_shardings = maxtext_utils.get_logical_annotations(config, topology_mesh, init_state_fn) - - return ( - model_name, - topology, - num_slice, - custom_mesh_and_rule, - overrides, - abstract_state, - state_mesh_shardings, - logical_shardings, - ) - - -@pytest.mark.tpu_backend -class TestGetAbstractState: - """Test class for get_abstract_state function and sharding comparison.""" - - # Requires JAX TPU support to generate the simulated TPU topology. - def test_get_abstract_state_sharding(self, abstract_state_and_shardings): # pylint: disable=redefined-outer-name - """Tests that get_abstract_state returns a state with the correct abstract structure and compares sharding.""" - - ( - model_name, - topology, - num_slice, - custom_mesh_and_rule, - overrides, - abstract_state, - state_mesh_shardings, - logical_shardings, - ) = abstract_state_and_shardings - - assert hasattr(abstract_state, "params") - assert hasattr(abstract_state, "opt_state") - param_leaf = jax.tree_util.tree_leaves(abstract_state.params)[0] - assert isinstance(param_leaf, jax.ShapeDtypeStruct) - assert param_leaf.dtype == jnp.float32 - - root_dir = "tests/utils/sharding_info" # Or your target directory - rule_name = f"rule_{custom_mesh_and_rule}" if custom_mesh_and_rule else "rule_default" - if overrides: - rule_name += "_" + "_".join(overrides) - base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}", rule_name) - os.makedirs(base_path, exist_ok=True) # Ensure directory exists for saving actual - - error_messages = [] - - # 1. Compare Physical/Named Shardings - named_json_path = os.path.join(base_path, "named_shardings.json") - if not os.path.exists(named_json_path): - pytest.skip(f"Missing named_shardings.json for {model_name} {topology} slice {num_slice}") - return - - # Use state_mesh_shardings from the fixture - actual_named = named_shardings_to_json(state_mesh_shardings, abstract_state) - expected_named = load_json(named_json_path) - - if compare_sharding_jsons(expected_named, "Expected (Physical)", actual_named, "Actual (Physical)"): - error_messages.append(f"Physical sharding mismatch for {model_name} on {topology} slice {num_slice}") - - # 2. Compare Logical Shardings - logical_json_path = os.path.join(base_path, "logical_shardings.json") - if not os.path.exists(logical_json_path): - pytest.skip(f"Missing logical_shardings.json for {model_name} {topology} slice {num_slice}") - return - - # Use logical_shardings from the fixture - actual_logical = partition_specs_to_json(logical_shardings, abstract_state) - expected_logical = load_json(logical_json_path) - - if compare_sharding_jsons(expected_logical, "Expected (Logical)", actual_logical, "Actual (Logical)"): - error_messages.append(f"Logical sharding mismatch for {model_name} on {topology} slice {num_slice}") - - assert not error_messages, "\n".join(error_messages) +The sharding-comparison tests in this file relied on Linen golden files and +Linen TrainState structure, which no longer exist after the NNX-only migration. +They were removed; this module is intentionally left without tests. +""" diff --git a/tests/unit/state_dtypes_test.py b/tests/unit/state_dtypes_test.py index 3d640cc62d..a92394d99c 100644 --- a/tests/unit/state_dtypes_test.py +++ b/tests/unit/state_dtypes_test.py @@ -14,25 +14,20 @@ """Test that all weights are expected dtype (default float32)""" -from functools import partial import unittest from flax import nnx import jax import jax.numpy as jnp from jax.sharding import Mesh + from maxtext.common import train_state_nnx -from maxtext.common.common_types import MODEL_MODE_TRAIN from maxtext.configs import pyconfig -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.optimizers import optimizers from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils from tests.utils.test_helpers import get_test_config_path -Transformer = models.transformer_as_linen - class StateDtypes(unittest.TestCase): """Tests that state has expected dtypes, e.g. weights default to float32""" @@ -41,43 +36,30 @@ def get_state(self, argv): """Gets model state including weights and optimizer state""" # Setup necessary inputs to build a model state config = pyconfig.initialize(argv) - quant = quantizations.configure_quantization(config) devices_array = maxtext_utils.create_device_mesh(config) mesh = Mesh(devices_array, config.mesh_axes) - if config.pure_nnx: - _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh) - else: - model = Transformer(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh) learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) tx = optimizers.get_optimizer(config, learning_rate_schedule, model) - _, example_rng = jax.random.split(jax.random.PRNGKey(0), 2) - - if config.pure_nnx: - def create_train_state_fn(): - nnx_model = _create_model_partial() - optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(nnx_model, optimizer) + def create_train_state_fn(): + nnx_model = _create_model_partial() + optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param) + return train_state_nnx.TrainStateNNX(nnx_model, optimizer) - init_state_fn = create_train_state_fn - else: - init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, example_rng) + init_state_fn = create_train_state_fn abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, True) - return abstract_state, config.pure_nnx + return abstract_state def get_weights(self, argv): - state, is_nnx = self.get_state(argv) - if is_nnx: - return state.model - return state.params + state = self.get_state(argv) + return state.model def get_mu(self, argv): - state, is_nnx = self.get_state(argv) - if is_nnx: - return state.optimizer.opt_state[0]["mu"] - return state.opt_state[0].mu + state = self.get_state(argv) + return state.optimizer.opt_state[0]["mu"] def assert_pytree_is_dtype(self, weights, expected_dtype): """Asserts that all valid parameter arrays within the PyTree match the expected dtype.""" diff --git a/tests/unit/train_compile_test.py b/tests/unit/train_compile_test.py index 4719e90119..faa5c120b1 100644 --- a/tests/unit/train_compile_test.py +++ b/tests/unit/train_compile_test.py @@ -828,27 +828,13 @@ def test_deepseek32(self): @parameterized.named_parameters( { - "testcase_name": "linen_scanned_dot_product", + "testcase_name": "scanned_dot_product", "scan_layers": "true", - "enable_nnx": "False", "attention": "dot_product", }, { - "testcase_name": "linen_scanned_flash", + "testcase_name": "scanned_flash", "scan_layers": "true", - "enable_nnx": "False", - "attention": "flash", - }, - { - "testcase_name": "nnx_scanned_dot_product", - "scan_layers": "true", - "enable_nnx": "True", - "attention": "dot_product", - }, - { - "testcase_name": "nnx_scanned_flash", - "scan_layers": "true", - "enable_nnx": "True", "attention": "flash", }, ) @@ -856,11 +842,10 @@ def test_deepseek32(self): def test_deepseek4( self, scan_layers, - enable_nnx, attention="dot_product", ): - # test deepseek4 compile across Linen and NNX - compiled_trainstep_file = f"/tmp/test_deepseek4_{scan_layers}_{enable_nnx}_{attention}.pickle" + # test deepseek4 compile. + compiled_trainstep_file = f"/tmp/test_deepseek4_{scan_layers}_{attention}.pickle" train_compile_main( ( "", @@ -885,9 +870,6 @@ def test_deepseek4( "sa_block_kv_dq=128", "dtype=bfloat16", "weight_dtype=bfloat16", - f"enable_nnx={enable_nnx}", - f"pure_nnx={enable_nnx}", - f"pure_nnx_decoder={enable_nnx}", "routed_bias=False", "override_model_config=True", ) @@ -1198,13 +1180,7 @@ def test_zero1_optimizer_sharding(self): ) def test_vocab_tiling_bf16_nnx(self): - """AOT compile vocab tiling on the NNX path (vocab_tiling_nnx_loss + custom_vjp). - - Sets `pure_nnx`/`enable_nnx`/`pure_nnx_decoder` explicitly so the NNX AOT - path is covered regardless of the default values. Once those defaults flip - to True, `test_vocab_tiling_bf16` above will already exercise this same - path via defaults. - """ + """AOT compile vocab tiling on the NNX path (vocab_tiling_nnx_loss + custom_vjp).""" compiled_trainstep_file = "/tmp/test_vocab_tiling_bf16_nnx.pickle" train_compile_main( ( @@ -1218,9 +1194,6 @@ def test_vocab_tiling_bf16_nnx(self): "max_target_length=1024", "num_vocab_tiling=4", "weight_dtype=bfloat16", - "pure_nnx=true", - "enable_nnx=true", - "pure_nnx_decoder=true", ) ) @@ -1246,9 +1219,6 @@ def test_envy(self, scan_layers): "attention=dot_product", "dtype=bfloat16", "weight_dtype=bfloat16", - "enable_nnx=True", - "pure_nnx=True", - "pure_nnx_decoder=True", "override_model_config=True", ) )