diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 89abd56d4c..533a9ca193 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -1922,8 +1922,10 @@ def __init__(self, **kwargs): "qwen3-235b-a22b": qwen3_235b_a22b_thinking_2507_config, "qwen3-480b-a35b": qwen3_coder_480b_a35b_config, "deepseek2-16b": deepseek2_16b_config, + "deepseek3-tiny": deepseek3_671b_config, "deepseek3-671b": deepseek3_671b_config, "deepseek3.2-671b": deepseek32_671b_config, + "deepseek4-tiny": deepseek4_284b_config, "deepseek4-284b": deepseek4_284b_config, "gpt-oss-20b": gpt_oss_20b_config, "gpt-oss-120b": gpt_oss_120b_config, diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 7191190f2b..7bedc332d6 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -30,6 +30,7 @@ from flax.training import train_state from grain.experimental import ElasticIterator import jax +import jax.numpy as jnp from maxtext.checkpoint_conversion.utils.load_dynamic import load_safetensors_dynamic_state from maxtext.common import emergency_checkpointing from maxtext.common import grain_utility @@ -266,6 +267,176 @@ def _resolve_conversion_fn(checkpoint_conversion_fn): return fn +def _drop_adapter_level(tree): + if isinstance(tree, dict): + if set(tree) == {"base"}: + return _drop_adapter_level(tree["base"]) + return {k: _drop_adapter_level(v) for k, v in tree.items()} + elif isinstance(tree, (list, tuple)): + if hasattr(tree, "_fields"): # namedtuple + return type(tree)(*[_drop_adapter_level(v) for v in tree]) + return type(tree)([_drop_adapter_level(v) for v in tree]) + return tree + + +def _drop_inject_hyperparams(opt_state): + if isinstance(opt_state, dict) and {"count", "hyperparams", "hyperparams_states", "inner_state"}.issubset( + opt_state.keys() + ): + return opt_state["inner_state"] + return opt_state + + +def _add_adapter_level(tree): + if isinstance(tree, dict): + return {"base": tree} + elif isinstance(tree, (list, tuple)): + if hasattr(tree, "_fields"): # namedtuple + return type(tree)(*[_add_adapter_level(v) for v in tree]) + return type(tree)([_add_adapter_level(v) for v in tree]) + return tree + + +def _load_tunix_full_state_from_path( + path, + abstract_unboxed_pre_state, + checkpoint_storage_concurrent_gb, + use_ocdbt, + use_zarr3, + maxtext_config=None, +): + """Load and convert a full Tunix post-training checkpoint into MaxText pre-train state.""" + is_nnx = isinstance(abstract_unboxed_pre_state, (nnx.State, train_state_nnx.TrainStateNNX)) + if is_nnx: + want_params = nnx.split_state(abstract_unboxed_pre_state.model, nnx.Param, ...)[0].to_pure_dict() + want_opt = abstract_unboxed_pre_state.optimizer + else: + want_params = abstract_unboxed_pre_state.params + want_opt = abstract_unboxed_pre_state.opt_state + + path_obj = epath.Path(path) + + # For CheckpointManager, path should point to the root dir, and step should be the folder name. + step_str = path_obj.name + root_dir = path_obj.parent + step = int(step_str) if step_str.isdigit() else 0 + + item_handlers = { + "model_params": ocp.PyTreeCheckpointHandler( + restore_concurrent_gb=checkpoint_storage_concurrent_gb, + use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3, + ), + "optimizer_state": ocp.PyTreeCheckpointHandler( + restore_concurrent_gb=checkpoint_storage_concurrent_gb, + use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3, + ), + } + + mgr = ocp.CheckpointManager( + root_dir, + item_names=("model_params", "optimizer_state"), + item_handlers=item_handlers, + ) + + has_base = False + has_inject = False + try: + item_meta = mgr.item_metadata(step) + + if hasattr(item_meta, "get"): + model_meta = item_meta.get("model_params") + opt_meta = item_meta.get("optimizer_state") + else: + model_meta = getattr(item_meta, "model_params", None) + opt_meta = getattr(item_meta, "optimizer_state", None) + + if model_meta is not None: + tree = getattr(model_meta, "tree", model_meta) + if isinstance(tree, dict) and "base" in tree: + has_base = True + + if opt_meta is not None: + tree = getattr(opt_meta, "tree", opt_meta) + if isinstance(tree, dict) and "inner_state" in tree and "count" in tree: + has_inject = True + except Exception: # pylint: disable=broad-except + pass + + want_params_dict = want_params.to_pure_dict() if isinstance(want_params, nnx.State) else want_params + want_opt_dict = want_opt.to_pure_dict() if isinstance(want_opt, nnx.State) else want_opt + + wrapped_params = jax.tree.map(lambda v: {"value": v}, want_params_dict) + wrapped_opt_base = jax.tree.map(lambda v: {"value": v}, want_opt_dict) + + target_params = {"base": wrapped_params} if has_base else wrapped_params + target_opt_base = _add_adapter_level(wrapped_opt_base) if has_base else wrapped_opt_base + + if has_inject: + target_opt = { + "count": {"value": jnp.zeros((), dtype=jnp.int32)}, + "hyperparams": {}, + "hyperparams_states": {}, + "inner_state": target_opt_base, + } + else: + target_opt = target_opt_base + + restored = mgr.restore( + step, + args=ocp.args.Composite( + model_params=ocp.args.PyTreeRestore( + item=target_params, + restore_args=ocp.checkpoint_utils.construct_restore_args(target_params), + partial_restore=True, + ), + optimizer_state=ocp.args.PyTreeRestore( + item=target_opt, + restore_args=ocp.checkpoint_utils.construct_restore_args(target_opt), + partial_restore=True, + ), + ), + ) + + restored_params = restored["model_params"] + restored_opt = restored["optimizer_state"] + + if has_base: + restored_params = _drop_adapter_level(restored_params) + restored_opt = _drop_adapter_level(restored_opt) + + if has_inject: + restored_opt = _drop_inject_hyperparams(restored_opt) + + restored_params = train_state_nnx._strip_rng_state(restored_params) # pylint: disable=protected-access + restored_params = jax.tree.map( + lambda v: v["value"] if isinstance(v, dict) and "value" in v else v, + restored_params, + is_leaf=lambda x: isinstance(x, dict) and "value" in x and not isinstance(x.get("value"), dict), + ) + restored_opt = jax.tree.map( + lambda v: v["value"] if isinstance(v, dict) and "value" in v else v, + restored_opt, + is_leaf=lambda x: isinstance(x, dict) and "value" in x and not isinstance(x.get("value"), dict), + ) + + _raise_on_weight_mismatch(want_params, restored_params, config=maxtext_config) + + if is_nnx: + if isinstance(abstract_unboxed_pre_state, nnx.State): + linen_state, aux_state, ephemeral = train_state_nnx.split_for_checkpoint(abstract_unboxed_pre_state) + nnx.replace_by_pure_dict(linen_state, {"model": restored_params, "optimizer": restored_opt}) + return nnx.merge_state(linen_state, aux_state, ephemeral) + nnx.replace_by_pure_dict(abstract_unboxed_pre_state, {"model": restored_params, "optimizer": restored_opt}) + return abstract_unboxed_pre_state + else: + return abstract_unboxed_pre_state.replace( + params=restored_params, + opt_state=restored_opt, + ) + + def _load_full_state_from_path( path, abstract_unboxed_pre_state, @@ -297,6 +468,24 @@ def _load_full_state_from_path( The loaded state. """ + if source_checkpoint_layout == "orbax": + try: + if (epath.Path(path) / "model_params").exists() and (epath.Path(path) / "optimizer_state").exists(): + max_logging.log(f"Auto-detected Tunix checkpoint layout at {path}") + source_checkpoint_layout = "tunix" + except Exception: # pylint: disable=broad-except + pass + + if source_checkpoint_layout == "tunix": + return _load_tunix_full_state_from_path( + path, + abstract_unboxed_pre_state, + checkpoint_storage_concurrent_gb, + use_ocdbt, + use_zarr3, + maxtext_config, + ) + if enable_orbax_v1: if source_checkpoint_layout == "orbax": # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. @@ -712,11 +901,96 @@ def load_params_from_path( assert load_parameters_from_path, "load_parameters_from_path is not defined." max_logging.log(f"restoring params from {load_parameters_from_path}") + # Check if Tunix Layout (either pointing to step dir or model_params dir) + is_tunix = False + path_obj = epath.Path(load_parameters_from_path) + target_path = path_obj + if path_obj.name == "model_params": + is_tunix = True + target_path = path_obj.parent + else: + try: + if (path_obj / "model_params").exists(): + is_tunix = True + target_path = path_obj + except Exception: # pylint: disable=broad-except + pass + + is_nnx = isinstance(abstract_unboxed_params, nnx.State) + want = abstract_unboxed_params.to_pure_dict() if is_nnx else abstract_unboxed_params + + if is_tunix: + max_logging.log(f"Detected Tunix layout for parameters at {target_path}") + + # Use CheckpointManager instead of Checkpointer to handle v1 multi-item directories + step_str = target_path.name + root_dir = target_path.parent + + step = int(step_str) if step_str.isdigit() else 0 + + item_handlers = { + "model_params": ocp.PyTreeCheckpointHandler( + restore_concurrent_gb=checkpoint_storage_concurrent_gb, + use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3, + ) + } + + mgr = ocp.CheckpointManager( + root_dir, + item_names=("model_params",), + item_handlers=item_handlers, + ) + + has_base = False + try: + item_meta = mgr.item_metadata(step) + if hasattr(item_meta, "get"): + model_meta = item_meta.get("model_params") + else: + model_meta = getattr(item_meta, "model_params", None) + + if model_meta is not None: + tree = getattr(model_meta, "tree", model_meta) + if isinstance(tree, dict) and "base" in tree: + has_base = True + except Exception: # pylint: disable=broad-except + pass + + target_want = jax.tree.map(lambda v: {"value": v}, want) + target_params = {"base": target_want} if has_base else target_want + + restored = mgr.restore( + step, + args=ocp.args.Composite( + model_params=ocp.args.PyTreeRestore( + item=target_params, + restore_args=ocp.checkpoint_utils.construct_restore_args(target_params), + partial_restore=True, + ) + ), + ) + restored_weights = restored["model_params"] + + if has_base: + restored_weights = _drop_adapter_level(restored_weights) + + restored_weights = train_state_nnx._strip_rng_state(restored_weights) # pylint: disable=protected-access + restored_weights = jax.tree.map( + lambda v: v["value"] if isinstance(v, dict) and "value" in v else v, + restored_weights, + is_leaf=lambda x: isinstance(x, dict) and "value" in x and not isinstance(x.get("value"), dict), + ) + + _raise_on_weight_mismatch(want, restored_weights) + if is_nnx: + nnx.replace_by_pure_dict(abstract_unboxed_params, restored_weights) + return abstract_unboxed_params + return restored_weights + # On disk the weights live at `params/params/...`: an outer key naming the item, and Flax's # `params` collection inside it. A Linen TrainState.params is that collection; an NNX params # state sits one level below it (bare weights), so wrap it going in and unwrap it coming out. - is_nnx = isinstance(abstract_unboxed_params, nnx.State) - want = abstract_unboxed_params.to_pure_dict() if is_nnx else abstract_unboxed_params # Determine the restore key based on the leaf directory name to support native and custom SFT restore_key = os.path.basename(load_parameters_from_path) diff --git a/src/maxtext/integration/tunix/tunix_adapter.py b/src/maxtext/integration/tunix/tunix_adapter.py index a5bcbf62b2..6c99b29bdd 100644 --- a/src/maxtext/integration/tunix/tunix_adapter.py +++ b/src/maxtext/integration/tunix/tunix_adapter.py @@ -81,9 +81,11 @@ def __init__( ): super().__init__() self.base = base_model + hf_config = HF_MODEL_CONFIGS.get(self.base.config.model_name) + hf_dict = hf_config.to_dict() if hf_config is not None else {} self._vllm_weight_mapping = VllmWeightMapping( self.base.config.model_name, - HF_MODEL_CONFIGS[self.base.config.model_name].to_dict(), + hf_dict, use_standalone_mappings, ) self.use_no_op_mappings = use_no_op_mappings diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index a2a1403125..87b75e0c05 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -415,8 +415,18 @@ def _fix_one(path, restore_arg): if not isinstance(restore_arg, ocp.ArrayRestoreArgs): return restore_arg stored_meta = _lookup_stored_meta(path) + dtype = restore_arg.dtype + if dtype is not None and jax.dtypes.issubdtype(dtype, jax.dtypes.prng_key): + dtype = jnp.uint32 + if dtype is None and stored_meta is not None and hasattr(stored_meta, "dtype"): + dtype = stored_meta.dtype + if dtype is None: + dtype = jnp.bfloat16 + if stored_meta is None: missing_paths.append(f" {'.'.join(_key_str(k) for k in path)}") + if restore_arg.dtype != dtype: + return dataclasses.replace(restore_arg, dtype=dtype) return restore_arg if _is_orbax_array_metadata(stored_meta): stored_shape = tuple(stored_meta.shape) @@ -434,14 +444,18 @@ def _fix_one(path, restore_arg): path_str = f" {'.'.join(_key_str(k) for k in path)}: stored={stored_shape} -> model={restore_arg.global_shape}" if _stored_shape_evenly_shardable(restore_arg, stored_shape): mismatched_paths_sharded.append(path_str) - return dataclasses.replace(restore_arg, global_shape=stored_shape, shape=stored_shape) + return dataclasses.replace(restore_arg, global_shape=stored_shape, shape=stored_shape, dtype=dtype) mismatched_paths_replicated.append(path_str) return dataclasses.replace( - restore_arg, global_shape=None, shape=None, sharding=replicated, mesh=None, mesh_axes=None + restore_arg, global_shape=None, shape=None, sharding=replicated, mesh=None, mesh_axes=None, dtype=dtype ) else: found_array_count[0] += 1 + if restore_arg.dtype != dtype: + return dataclasses.replace(restore_arg, dtype=dtype) + elif restore_arg.dtype != dtype: + return dataclasses.replace(restore_arg, dtype=dtype) return restore_arg fixed = jax.tree_util.tree_map_with_path(_fix_one, restore_args, is_leaf=lambda x: isinstance(x, ocp.ArrayRestoreArgs)) @@ -941,6 +955,8 @@ def from_pretrained( with mesh: if config.load_parameters_path: + # For Tunix checkpoints (output of CheckpointManager), direct the single-item + # Checkpointer directly to the `model_params` subfolder to read PyTree metadata. ckptr = ocp.Checkpointer( ocp.PyTreeCheckpointHandler( restore_concurrent_gb=config.checkpoint_storage_concurrent_gb, @@ -950,13 +966,14 @@ def from_pretrained( ) ) - # This is a memory optimization. We don't want to restore the entire checkpoint - only the params. - # Rather than passing the entire abstract state, which could unnecessarily restore opt_state and - # waste memory, we instead restore the params field of the checkpoint (which itself may be a dictionary - # containing a key named 'params'). + load_params_path = epath.Path(config.load_parameters_path) + try: + if (load_params_path / "model_params").exists(): + load_params_path = load_params_path / "model_params" + except Exception: # pylint: disable=broad-except + pass - # Get the structure of checkpoint in `config.load_parameters_path` - metadata = ckptr.metadata(config.load_parameters_path) + metadata = ckptr.metadata(load_params_path) if metadata is None or metadata.item_metadata is None: max_logging.log( f"ERROR: No valid Orbax checkpoint found at '{config.load_parameters_path}'. " @@ -1009,9 +1026,17 @@ def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx): # serve-mode quantized models store their scale factors and integer payloads in custom AQT variable # types (e.g. `qrhs.frozen`), which are NOT subclasses of `nnx.Param`. Negative filtering with # `not isinstance(...)` safely retains all weight-like leaves while excluding transient runtime state. - param_state = sharded_state.filter( - lambda path, var: not isinstance(var, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat)) - ) + def _is_weight_param(path, var): + if isinstance(var, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat)): + return False + val = var.get_value() if hasattr(var, "get_value") else getattr(var, "value", None) + if val is not None: + dtype = getattr(val, "dtype", None) + if dtype is not None and jax.dtypes.issubdtype(dtype, jax.dtypes.prng_key): + return False + return True + + param_state = sharded_state.filter(_is_weight_param) is_nnx_checkpoint = True if ( "params" in metadata.item_metadata.tree.keys() @@ -1120,7 +1145,7 @@ def _free_device_memory(path, node): jax.tree_util.tree_map_with_path(_free_device_memory, sharded_state, is_leaf=lambda n: isinstance(n, nnx.Variable)) restored = ckptr.restore( - epath.Path(config.load_parameters_path), + epath.Path(load_params_path), item=item_to_restore, transforms={}, restore_args=restore_args, diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py index 9440d64821..fff6f667c1 100644 --- a/tests/unit/checkpointing_test.py +++ b/tests/unit/checkpointing_test.py @@ -512,5 +512,76 @@ def test_error_handler_raises_runtime_error(self): self.assertIs(cm.exception.__cause__, original_error) +class TunixCheckpointConversionTest(parameterized.TestCase): + """Unit tests for Tunix on-load checkpoint conversion and PyTree transformation helpers.""" + + # pylint: disable=protected-access + + def test_drop_adapter_level_dict(self): + """Tests stripping 'base' adapter wrapper from nested dictionaries.""" + tree = {"base": {"layer_0": {"kernel": jnp.ones((4, 4))}, "layer_1": {"bias": jnp.zeros((4,))}}} + stripped = checkpointing._drop_adapter_level(tree) + self.assertIn("layer_0", stripped) + self.assertIn("layer_1", stripped) + self.assertNotIn("base", stripped) + self.assertEqual(stripped["layer_0"]["kernel"].shape, (4, 4)) + + def test_drop_adapter_level_tuple_and_namedtuple(self): + """Tests stripping 'base' adapter wrapper from tuples and namedtuples.""" + import collections # pylint: disable=import-outside-toplevel + + OptTuple = collections.namedtuple("OptTuple", ["mu", "nu"]) + wrapped_namedtuple = OptTuple( + mu={"base": {"w": jnp.ones((2, 2))}}, + nu={"base": {"w": jnp.zeros((2, 2))}}, + ) + stripped = checkpointing._drop_adapter_level(wrapped_namedtuple) + self.assertIsInstance(stripped, OptTuple) + self.assertIn("w", stripped.mu) + np.testing.assert_allclose(stripped.mu["w"], np.ones((2, 2))) + self.assertIn("w", stripped.nu) + np.testing.assert_allclose(stripped.nu["w"], np.zeros((2, 2))) + + wrapped_list = [{"base": {"a": 1}}, {"base": {"b": 2}}] + stripped_list = checkpointing._drop_adapter_level(wrapped_list) + self.assertEqual(stripped_list, [{"a": 1}, {"b": 2}]) + + def test_drop_inject_hyperparams(self): + """Tests unwrapping Optax inject_hyperparams state structure.""" + opt_state = { + "count": jnp.zeros(()), + "hyperparams": {"learning_rate": 1e-4}, + "hyperparams_states": {}, + "inner_state": {"mu": jnp.ones((2, 2)), "nu": jnp.zeros((2, 2))}, + } + unwrapped = checkpointing._drop_inject_hyperparams(opt_state) + self.assertNotIn("inner_state", unwrapped) + self.assertIn("mu", unwrapped) + self.assertIn("nu", unwrapped) + np.testing.assert_allclose(unwrapped["mu"], np.ones((2, 2))) + np.testing.assert_allclose(unwrapped["nu"], np.zeros((2, 2))) + + # Pass-through if not inject_hyperparams + passthrough = {"mu": 123} + self.assertEqual(checkpointing._drop_inject_hyperparams(passthrough), passthrough) + + def test_add_adapter_level(self): + """Tests wrapping dictionaries and tuples with a 'base' level for restore targeting.""" + import collections # pylint: disable=import-outside-toplevel + + tree = {"layer_0": {"kernel": jnp.ones((4, 4))}} + wrapped = checkpointing._add_adapter_level(tree) + self.assertIn("base", wrapped) + self.assertIn("layer_0", wrapped["base"]) + np.testing.assert_allclose(wrapped["base"]["layer_0"]["kernel"], np.ones((4, 4))) + + OptTuple = collections.namedtuple("OptTuple", ["mu"]) + named_tree = OptTuple(mu={"w": jnp.ones((2, 2))}) + wrapped_named = checkpointing._add_adapter_level(named_tree) + self.assertIsInstance(wrapped_named, OptTuple) + self.assertIn("base", wrapped_named.mu) + np.testing.assert_allclose(wrapped_named.mu["base"]["w"], np.ones((2, 2))) + + if __name__ == "__main__": absltest.main()