From dd76238cc8f82c271e137fec5058a486c6f97609 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Wed, 29 Jul 2026 15:24:47 +0000 Subject: [PATCH] [NNX] Delete Linen (pre-train 2/3): collapse dispatch in sharding and muon utils - sharding.maybe_update_params_sharding_with_opt delegates to the _nnx variant; build_zero1_input_state_mesh_shardings drops its Linen early return. - muon_utils.get_muon_weight_dimension_numbers drops the isinstance(nnx.Module) test and the Linen get_abstract_param path; get_model_mdn loses its pure_nnx parameter and always builds the abstract NNX model. - run_sharding_dump loses its --pure_nnx flag. Tests follow: the Linen branch test in muon_utils_test goes away, optimizers_test drops the pure_nnx argument and the dual-shape comparison, sharding_nnx_test drops the flag from its fake config. --- src/maxtext/utils/muon_utils.py | 51 ++++++++------------------------ src/maxtext/utils/sharding.py | 24 +-------------- tests/unit/muon_utils_test.py | 32 -------------------- tests/unit/optimizers_test.py | 8 ++--- tests/unit/sharding_nnx_test.py | 5 ++-- tests/utils/run_sharding_dump.py | 11 ++----- 6 files changed, 20 insertions(+), 111 deletions(-) diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index 91d7d781c9..02c3d549eb 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/src/maxtext/utils/muon_utils.py @@ -33,8 +33,6 @@ import jax from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_PKG_DIR -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.utils import maxtext_utils, model_creation_utils from optax.contrib._muon import MuonDimensionNumbers as mdn @@ -142,26 +140,17 @@ def get_transform_tree(tree, path=()): def get_muon_weight_dimension_numbers(model, config, verbose=False): """Extract muon dimension number from model structure.""" - if isinstance(model, nnx.Module): - _, abstract_param, _ = nnx.split(model, nnx.Param, ...) + _, abstract_param, _ = nnx.split(model, nnx.Param, ...) - def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf): - # Convert jax.tree_util.KeyEntry path to Tuple[str, ...] - path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey)) - return transform_logic(path_strings) + def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf): + # Convert jax.tree_util.KeyEntry path to Tuple[str, ...] + path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey)) + return transform_logic(path_strings) - # NNX abstract_param is an nnx.State (not Linen's dict of LogicallyPartitioned leaves); - # tree_map_with_path round-trips that structure so each Param.value holds the mdn result. - muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path( - apply_transform_nnx, nnx.to_pure_dict(abstract_param) - ) - muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers) - - else: # Linen - # quickly get param structure without materialization - abstract_param = maxtext_utils.get_abstract_param(model, config) - # get muon dimension number from param - muon_weight_dimension_numbers = get_transform_tree(abstract_param) + # tree_map_with_path handles NNX's PyTree structure; result is an nnx.State with the + # same structure, where each Param's value holds the mdn result. + muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path(apply_transform_nnx, nnx.to_pure_dict(abstract_param)) + muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers) if verbose: _print_structure_debug(abstract_param, muon_weight_dimension_numbers) @@ -193,7 +182,7 @@ def get_leaf_info(leaf): print("\nIs this reasonable?") -def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False): +def get_model_mdn(model_name, scan_layers=True, verbose=False): """Initializes a model and retrieves its Muon dimension numbers. This function sets up the configuration for a given model, initializes the @@ -217,30 +206,16 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False): f"model_name={model_name}", f"scan_layers={scan_layers}", "attention=dot_product", - f"pure_nnx={pure_nnx}", "skip_jax_distributed_system=True", ] - if not pure_nnx: - argv.extend( - [ - "enable_nnx=False", - "pure_nnx_decoder=False", - ] - ) config = pyconfig.initialize(argv) # Setup model devices_array = maxtext_utils.create_device_mesh(config) mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) - quant = quantizations.configure_quantization(config) - if pure_nnx: - _, model = model_creation_utils.create_nnx_abstract_model(config, mesh) - else: - model = models.transformer_as_linen(config, mesh=mesh, quant=quant) + _, model = model_creation_utils.create_nnx_abstract_model(config, mesh) # Get dimension number muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config, verbose=verbose) - if pure_nnx: - muon_weight_dimension_numbers = {"params": nnx.to_pure_dict(muon_weight_dimension_numbers)} - return muon_weight_dimension_numbers + return {"params": nnx.to_pure_dict(muon_weight_dimension_numbers)} if __name__ == "__main__": @@ -249,4 +224,4 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False): sys.exit(1) model_name_arg = sys.argv[1] scan_layers_arg = sys.argv[2].lower() == "true" - get_model_mdn(model_name_arg, scan_layers_arg, verbose=True, pure_nnx=False) + get_model_mdn(model_name_arg, scan_layers_arg, verbose=True) diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index a1e320603f..7cb3b4760c 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -28,7 +28,6 @@ from maxtext.configs import pyconfig from maxtext.utils import max_logging from maxtext.utils import max_utils -import optax _LOGGED_ACTIVATION_SHARDINGS = set() _ACTIVATION_SHARDINGS_DUMP = [] @@ -712,26 +711,7 @@ def maybe_update_params_sharding_with_opt(config, state_mesh_shardings): - updated_state_mesh_shardings: State mesh shardings with updated params field (unchanged if shard_optimizer_over_data is False) """ - if config.pure_nnx: - return maybe_update_params_sharding_with_opt_nnx(config, state_mesh_shardings) - prev_params_shardings = state_mesh_shardings.params - if config.shard_optimizer_over_data: - if isinstance(state_mesh_shardings.opt_state, optax.ScaleByAdamState): - sharded_fp32_params = state_mesh_shardings.opt_state.mu - elif isinstance(state_mesh_shardings.opt_state, tuple) and isinstance( - state_mesh_shardings.opt_state[0], optax.ScaleByAdamState - ): - sharded_fp32_params = state_mesh_shardings.opt_state[0].mu - else: - raise NotImplementedError(f"Could not find optimizer state shardings from {type(state_mesh_shardings.opt_state)}") - if "params" not in sharded_fp32_params.keys(): # pyrefly: ignore[missing-attribute] - # When quantization=fp8 is enabled the sharded_fp32_params - # are not wrapped in `params`. Here we wrap them back. - sharded_fp32_params = {"params": sharded_fp32_params} - state_mesh_shardings = state_mesh_shardings.replace( - params=dict(prev_params_shardings, **sharded_fp32_params) # pyrefly: ignore[bad-unpacking] - ) # pyrefly: ignore[bad-unpacking] - return prev_params_shardings, state_mesh_shardings + return maybe_update_params_sharding_with_opt_nnx(config, state_mesh_shardings) def maybe_update_params_sharding_with_opt_nnx( @@ -865,8 +845,6 @@ def build_zero1_input_state_mesh_shardings(config, state_mesh_shardings, params_ """ if not config.shard_optimizer_over_data: return state_mesh_shardings - if not config.pure_nnx: - return state_mesh_shardings.replace(params=params_shardings) # nnx.State has no .replace: shallow-copy via tree_map (preserves nested container # types) and overlay params_shardings under input_state.model. input_state = jax.tree_util.tree_map( diff --git a/tests/unit/muon_utils_test.py b/tests/unit/muon_utils_test.py index 0b78056a1f..1ae5afeae7 100644 --- a/tests/unit/muon_utils_test.py +++ b/tests/unit/muon_utils_test.py @@ -19,7 +19,6 @@ import io import contextlib import unittest -from unittest import mock import jax import jax.numpy as jnp @@ -220,37 +219,6 @@ def test_nnx_verbose_path_executes_print_debug(self): self.assertIn("Muon Dimension Numbers", buf.getvalue()) -class TestGetMuonWeightDimensionNumbersLinen(unittest.TestCase): - """Covers the Linen branch of get_muon_weight_dimension_numbers.""" - - def test_linen_branch_uses_get_abstract_param(self): - """Linen models dispatch to maxtext_utils.get_abstract_param + get_transform_tree.""" - # Build a Linen nn.Module so isinstance(model, nnx.Module) is False. - - class LinenStub(nn.Module): - - @nn.compact - def __call__(self, x): - return x - - model = LinenStub() - - # Mock the heavy get_abstract_param call with a pre-shaped dict that exercises - # both a standard weight path and an excluded path. - fake_abstract_param = { - "params": { - "self_attention": {"out": object()}, - "norm": {"scale": object()}, - }, - } - - with mock.patch.object(muon_utils.maxtext_utils, "get_abstract_param", return_value=fake_abstract_param): - result = muon_utils.get_muon_weight_dimension_numbers(model, config=mock.MagicMock()) - - self.assertEqual(result["params"]["self_attention"]["out"], mdn((0, -2), (-1,))) - self.assertIsNone(result["params"]["norm"]["scale"]) - - class TestPrintStructureDebug(unittest.TestCase): """Covers both branches of get_leaf_info inside _print_structure_debug.""" diff --git a/tests/unit/optimizers_test.py b/tests/unit/optimizers_test.py index 0057dab45c..7719bd20ef 100644 --- a/tests/unit/optimizers_test.py +++ b/tests/unit/optimizers_test.py @@ -574,12 +574,8 @@ def test_model_integration(self, model_name, expected_output): Initializes the specified MaxText model and asserts that the generated Muon dimension numbers match the hardcoded reference. """ - is_pure_nnx = model_name in {"deepseek4-284b"} - actual_output = muon_utils.get_model_mdn(model_name, scan_layers=True, pure_nnx=is_pure_nnx) - if "params" in expected_output and "params" in actual_output: - self.assertEqual(actual_output["params"], expected_output["params"]) - else: - self.assertEqual(actual_output, expected_output) + actual_output = muon_utils.get_model_mdn(model_name, scan_layers=True) + self.assertEqual(actual_output, expected_output) class AdamWMaskTest(parameterized.TestCase): diff --git a/tests/unit/sharding_nnx_test.py b/tests/unit/sharding_nnx_test.py index 5e0db5e592..481e93d9d3 100644 --- a/tests/unit/sharding_nnx_test.py +++ b/tests/unit/sharding_nnx_test.py @@ -30,7 +30,6 @@ @dataclass class _Cfg: - pure_nnx: bool = True shard_optimizer_over_data: bool = False @@ -95,9 +94,9 @@ class TestMaybeUpdateParamsShardingWithOptNNX(unittest.TestCase): def setUp(self): self.model = _LinearNNX(rngs=nnx.Rngs(0)) - def test_dispatch_from_main_helper_when_pure_nnx(self): + def test_dispatch_from_main_helper(self): """maybe_update_params_sharding_with_opt should dispatch to the NNX variant.""" - cfg = _Cfg(pure_nnx=True, shard_optimizer_over_data=False) + cfg = _Cfg(shard_optimizer_over_data=False) state_mesh_shardings = _build_state_mesh_shardings(self.model, optax.adam(1e-3)) prev, updated = sharding.maybe_update_params_sharding_with_opt(cfg, state_mesh_shardings) # prev is the param-only view (no rngs / non-Param nodes) diff --git a/tests/utils/run_sharding_dump.py b/tests/utils/run_sharding_dump.py index fe371fb160..7d3156fe00 100644 --- a/tests/utils/run_sharding_dump.py +++ b/tests/utils/run_sharding_dump.py @@ -59,12 +59,9 @@ flags.DEFINE_string("topology", None, "Specific topology to dump.") flags.DEFINE_string("num_slice", None, "Specific number of slices to dump.") flags.DEFINE_string("custom_mesh_and_rule", None, "Specific custom_mesh_and_rule to dump.") -flags.DEFINE_bool("pure_nnx", False, "Use pure NNX model.") -def run_single_dump( - model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple, pure_nnx: bool = False -) -> None: +def run_single_dump(model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple) -> None: """Generate sharding json file for one specific model, topology, slice and rule.""" args = [ "python3", @@ -82,10 +79,6 @@ def run_single_dump( args.append(f"custom_mesh_and_rule={custom_mesh_and_rule}") if overrides: args.extend(overrides) - if pure_nnx: - args.append("pure_nnx=true") - else: - args.extend(["pure_nnx=False", "enable_nnx=False", "pure_nnx_decoder=False"]) subprocess.run(args, check=True) @@ -124,7 +117,7 @@ def main(argv: Sequence[str]) -> None: print(" -> Sharding files already exist. Regenerating to overwrite.") try: - run_single_dump(model_name, topology, str(num_slice), custom_mesh_and_rule, overrides, pure_nnx=FLAGS.pure_nnx) + run_single_dump(model_name, topology, str(num_slice), custom_mesh_and_rule, overrides) except subprocess.CalledProcessError: print(f"!!! FAILED: {model_name} {topology} {num_slice} {custom_mesh_and_rule} overrides={overrides}")