diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py index 2f3e26db7b..138d5be4a5 100644 --- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py +++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py @@ -41,6 +41,7 @@ import sys from flax import nnx +from flax import serialization import jax from jax import random from jax.sharding import Mesh @@ -118,6 +119,12 @@ def init_state_fn(): ) state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn) + if cfg.pure_nnx: + state = train_state_nnx.to_checkpoint_dict(state) + state.pop("nnx_aux", None) + else: + state = serialization.to_state_dict(state) + max_logging.log("start") max_utils.print_mem_stats("After params initialized") @@ -201,24 +208,10 @@ def init_state_fn(): "['decoder']['decoder_norm']['bias']": (".params.lm.final_ln.bias", None), } - if cfg.pure_nnx: - # NNX state-tree paths after `nnx.split(TrainStateNNX)`. The state is a - # nested `nnx.State` (dict-like Mapping) with `nnx.Variable` leaves, so - # `jax.tree_util.keystr` produces dict-style entries (`['key']`) plus - # `.value` for the Variable leaf, plus `[idx]` for the optax tuple: - # model params -> ['model'].value - # adam mu / nu -> ['optimizer']['opt_state'][0]['mu' | 'nu'].value - # step -> ['optimizer']['step'].value - # opt count -> ['optimizer']['opt_state'][0]['count'].value - state_map = { - "['optimizer']['step'].value": ("step", None), - "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None), - } - else: - state_map = { - ".step": ("step", None), - ".opt_state.count": ("opt_states_0.no_prefix_0.count", None), - } + state_map = { + "['step']": ("step", None), + "['opt_state']['count']": ("opt_states_0.no_prefix_0.count", None), + } def get_layer_prefix(keystr_pax): # different path format between decoder_layer variable @@ -231,26 +224,18 @@ def get_layer_prefix(keystr_pax): for keystr_maxtext, (keystr_pax, transform_fn) in keystr_map.items(): prefix_pax_opt_state = get_layer_prefix(keystr_pax) - if cfg.pure_nnx: - state_map[f"['model']{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn) - state_map[f"['optimizer']['opt_state'][0]['mu']{keystr_maxtext}.value"] = ( - f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", - transform_fn, - ) - state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = ( - f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", - transform_fn, - ) - else: - state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn) - state_map[f".opt_state.mu['params']{keystr_maxtext}"] = ( - f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", - transform_fn, - ) - state_map[f".opt_state.nu['params']{keystr_maxtext}"] = ( - f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", - transform_fn, - ) + state_map[f"['params']['params']{keystr_maxtext}"] = ( + f"mdl_vars{keystr_pax}", + transform_fn, + ) + state_map[f"['opt_state']['mu']['params']{keystr_maxtext}"] = ( + f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}", + transform_fn, + ) + state_map[f"['opt_state']['nu']['params']{keystr_maxtext}"] = ( + f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}", + transform_fn, + ) def verify_fn(key_path, _): keystr = jax.tree_util.keystr(key_path) @@ -302,9 +287,9 @@ def map_fn(key_path, value): max_logging.log("converted state finished") max_utils.print_mem_stats("converted state finished") - step_value = int(converted_state.optimizer.step.value) if cfg.pure_nnx else converted_state.step - if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state): - max_logging.log(f"saved a checkpoint at step {step_value}") + step_value = int(converted_state["step"]) if isinstance(converted_state, dict) else int(converted_state.step) + if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state, config=cfg): + max_logging.log(f"Saved a checkpoint at step {step_value}.") # Upon preemption, exit when and only when all ongoing saves are complete. if checkpointing.reached_preemption(checkpoint_manager, step_value): checkpointing.wait_until_finished(checkpoint_manager) @@ -322,7 +307,12 @@ def map_fn(key_path, value): default="gs://mlperf-llm-public2/gpt3_spmd1x64x24_tpuv4-3072_v84_20221101/checkpoints/checkpoint_00004000", required=True, ) - parser.add_argument("--maxtext-model-name", choices=["gpt3-175b", "gpt3-52k"], type=str, required=True) + parser.add_argument( + "--maxtext-model-name", + choices=["gpt3-175b", "gpt3-52k"], + type=str, + required=True, + ) parser.add_argument("--base-output-directory", type=str, required=True) parser.add_argument("--run-name", type=str, required=True) @@ -330,4 +320,9 @@ def map_fn(key_path, value): if not args.paxml_ckpt_path.startswith("gs://"): raise ValueError("--paxml-ckpt-path should be a gcs path starting with gs://") - convert(args.paxml_ckpt_path, args.maxtext_model_name, args.base_output_directory, args.run_name) + convert( + args.paxml_ckpt_path, + args.maxtext_model_name, + args.base_output_directory, + args.run_name, + ) diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 4827cbc1a4..0d51300a9d 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -753,18 +753,6 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.") return - if config.pure_nnx: - # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable. - if config.enable_diloco: - # DiLoCoTrainState: persist the synchronized global model (outer params). - # The per-replica inner optimizer / outer-momentum state is not checkpointed. - step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step - state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) - else: - # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout - # stream continues across resumes instead of resetting to a base key. - state = train_state_nnx.to_checkpoint_dict(state) - try: checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save) if checkpoint_saved: @@ -793,12 +781,26 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=None, force=False): """Wrapper for saving checkpoint.""" - if config and config.enable_checkpointing: + if not isinstance(state, (dict, nnx.State, train_state.TrainState)): + if isinstance(state, train_state_nnx.TrainStateNNX): + state = nnx.state(state) + elif not isinstance(state, (dict, nnx.State)): + state = {} + + if config and getattr(config, "pure_nnx", False) and isinstance(state, nnx.State): + # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable. + if getattr(config, "enable_diloco", False): + step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step + state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) + else: + state = train_state_nnx.to_checkpoint_dict(state) + + if config and getattr(config, "enable_checkpointing", False): if ( force - or (step % config.checkpoint_period == 0 and not config.enable_continuous_checkpointing) + or (step % config.checkpoint_period == 0 and not getattr(config, "enable_continuous_checkpointing", False)) or (_uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0) - or (config.enable_autocheckpoint and reached_preemption(checkpoint_manager, step)) + or (getattr(config, "enable_autocheckpoint", False) and reached_preemption(checkpoint_manager, step)) ): blocking_until_ready_start = time.time() max_logging.log(f"Waiting for step {step} to finish before checkpoint...") @@ -812,7 +814,13 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= # specify chunk_byte_size to force orbax to control maximum file size in checkpoint chunk_byte_size = ( - config.checkpoint_storage_target_data_file_size_bytes if config else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE + getattr( + config, + "checkpoint_storage_target_data_file_size_bytes", + DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE, + ) + if config + else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE ) checkpoint_args = ocp.args.PyTreeSave( @@ -822,7 +830,11 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= ) save_args_composite = {"items": checkpoint_args} - if config and config.dataset_type == "grain" and not isinstance(data_iterator, PlaceHolderDataIterator): + if ( + config + and getattr(config, "dataset_type", None) == "grain" + and not isinstance(data_iterator, PlaceHolderDataIterator) + ): if isinstance(data_iterator, RemoteIteratorWrapper): # Pass the wrapper directly; GrainCheckpointHandler will call save_state with the step save_args_composite["iter"] = grain_utility.GrainCheckpointSave( diff --git a/tests/unit/train_state_nnx_checkpoint_test.py b/tests/unit/train_state_nnx_checkpoint_test.py index ef19ac427e..46ee0b73e9 100644 --- a/tests/unit/train_state_nnx_checkpoint_test.py +++ b/tests/unit/train_state_nnx_checkpoint_test.py @@ -434,13 +434,31 @@ def test_nnx_and_linen_agree_on_actual_step(self): ) def test_nnx_state_is_saved_in_linen_layout(self): - """For pure_nnx=True, maybe_save_checkpoint reshapes the NNX state to the Linen on-disk layout.""" + """For pure_nnx=True, save_checkpoint reshapes the NNX state to the Linen on-disk layout.""" state = self._build_nnx_state(self.N_STEPS) self.assertIsInstance(state, nnx.State) # precondition: NNX train_step returns an nnx.State - captured = self._invoke_maybe_save(state, pure_nnx=True) + config = self._config(pure_nnx=True, enable_checkpointing=True, checkpoint_period=1) + mgr = mock.MagicMock() + mgr.reached_preemption.return_value = False + + captured = {} - # save_checkpoint should receive a plain dict in Linen layout, not the nnx.State. + def fake_save(_step, *args, **kwargs): + composite = kwargs.get("args") + if composite: + for key in ["items", "state"]: + if hasattr(composite, "_items") and key in composite._items: # pylint: disable=protected-access + val = composite[key] + if val is not None and hasattr(val, "item"): + captured["state"] = val.item + break + return True + + mgr.save.side_effect = fake_save + checkpointing.save_checkpoint(mgr, self.N_STEPS - 1, state, config=config, force=True) + + # save_checkpoint should pass a plain dict in Linen layout to Orbax, not the nnx.State. self.assertIsInstance(captured["state"], dict) self.assertNotIsInstance(captured["state"], nnx.State) # Linen layout: {params: {params: ...}, step, opt_state}; not the NNX {model, optimizer}. @@ -558,16 +576,12 @@ def test_maybe_save_checkpoint_allows_local_checkpoint_period(self): mgr.reached_preemption.return_value = False save_checkpoint_mock = mock.MagicMock(return_value=False) - with ( - mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock), - mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock, - ): + with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock): checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5) mgr.latest_step.assert_called_once_with() mgr.reached_preemption.assert_called_once_with(5) mgr.wait_until_finished.assert_not_called() - to_checkpoint_dict_mock.assert_called_once_with(state) save_checkpoint_mock.assert_called_once() def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy( @@ -587,16 +601,12 @@ def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy( mgr.reached_preemption.return_value = False save_checkpoint_mock = mock.MagicMock(return_value=False) - with ( - mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock), - mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock, - ): + with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock): checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5) mgr.should_save.assert_called_once_with(5) mgr.latest_step.assert_called_once_with() mgr.reached_preemption.assert_called_once_with(5) - to_checkpoint_dict_mock.assert_called_once_with(state) save_checkpoint_mock.assert_called_once() def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self): @@ -610,12 +620,10 @@ def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self): with ( mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock), - mock.patch.object(train_state_nnx, "to_checkpoint_dict", return_value={}) as to_checkpoint_dict_mock, mock.patch.object(checkpointing.elastic_utils, "maybe_elastic_scale_up") as mock_maybe_scale_up, ): checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5) - to_checkpoint_dict_mock.assert_called_once_with(state) save_checkpoint_mock.assert_called_once() mock_maybe_scale_up.assert_called_once_with(config, mgr)