Skip to content

[NNX] Delete Linen (1a/5): collapse dispatch in pre-train checkpoint … #4652

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ecnal-cienet wants to merge 1 commit into from
Closed

[NNX] Delete Linen (1a/5): collapse dispatch in pre-train checkpoint … #4652

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 30 additions & 60 deletions src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@
"""

import argparse
import functools
import gc
import os
import sys
Expand All @@ -47,11 +46,7 @@
from maxtext.configs import pyconfig
from maxtext.utils.globals import MAXTEXT_PKG_DIR
from maxtext.common import checkpointing
from maxtext.common.common_types import MODEL_MODE_TRAIN
from maxtext.layers import quantizations
from maxtext.common import train_state_nnx
from maxtext.models.models import transformer_as_linen
from maxtext.optimizers import optimizers
from maxtext.utils import max_logging
from maxtext.utils import max_utils
from maxtext.utils import maxtext_utils
Expand Down Expand Up @@ -92,23 +87,15 @@ def convert(paxml_ckpt_path, maxtext_model_name, base_output_directory, run_name
devices_array = maxtext_utils.create_device_mesh(cfg)
mesh = Mesh(devices_array, cfg.mesh_axes)

if cfg.pure_nnx:
rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng)
model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs)
_, tx = train_utils.create_training_optimizer(cfg, model)
_create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh)
rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng)
model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs)
_, tx = train_utils.create_training_optimizer(cfg, model)
_create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh)

def init_state_fn():
nnx_model = _create_model_partial()
optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
return train_state_nnx.TrainStateNNX(nnx_model, optimizer)

else:
quant = quantizations.configure_quantization(cfg)
model = transformer_as_linen(cfg, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(cfg)
tx = optimizers.get_optimizer(cfg, learning_rate_schedule)
init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, cfg, True, init_rng)
def init_state_fn():
nnx_model = _create_model_partial()
optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
return train_state_nnx.TrainStateNNX(nnx_model, optimizer)

checkpoint_manager = checkpointing.create_orbax_checkpoint_manager(
cfg.checkpoint_dir,
Expand Down Expand Up @@ -201,24 +188,18 @@ 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']<rest>.value
# adam mu / nu -> ['optimizer']['opt_state'][0]['mu' | 'nu']<rest>.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),
}
# 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']<rest>.value
# adam mu / nu -> ['optimizer']['opt_state'][0]['mu' | 'nu']<rest>.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),
}

def get_layer_prefix(keystr_pax):
# different path format between decoder_layer variable
Expand All @@ -231,26 +212,15 @@ 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"['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,
)

def verify_fn(key_path, _):
keystr = jax.tree_util.keystr(key_path)
Expand Down Expand Up @@ -302,7 +272,7 @@ 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
step_value = int(converted_state.optimizer.step.value)
if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state):
max_logging.log(f"saved a checkpoint at step {step_value}")
# Upon preemption, exit when and only when all ongoing saves are complete.
Expand Down
33 changes: 14 additions & 19 deletions src/maxtext/common/checkpointing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ def _load_linen_checkpoint_into_nnx(
use_ocdbt,
use_zarr3,
):
"""Restores a Linen-layout checkpoint into an NNX state (pure_nnx resume).
"""Restores a Linen-layout checkpoint into an NNX state.

Restores a Linen-shape target that includes `nnx_aux`, then reshapes back via
`_linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when
Expand Down Expand Up @@ -273,7 +273,7 @@ def combine_sharding(sds, shardings):
else:
raise ocp_v1.errors.InvalidLayoutError(f"Unknown checkpoint layout: {source_checkpoint_layout}")
else:
# pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state.
# Checkpoints use the Linen on-disk layout; reshape it back into the NNX state.
if isinstance(abstract_unboxed_pre_state, nnx.State):
return _load_linen_checkpoint_into_nnx(
path,
Expand Down Expand Up @@ -730,12 +730,8 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
if step is not None:
actual_step = int(step)
else:
if config.pure_nnx:
# Under DiLoCo the step lives on the DiLoCoTrainState; otherwise on the optimizer.
actual_step = int(state.step if config.enable_diloco else state.optimizer.step) - 1
else:
# Linen TrainState has .step attribute
actual_step = int(state.step) - 1
# Under DiLoCo the step lives on the DiLoCoTrainState; otherwise on the optimizer.
actual_step = int(state.step if config.enable_diloco else state.optimizer.step) - 1

# Determine if a checkpoint save should be forced, overriding the usual
# `config.checkpoint_period` logic.
Expand All @@ -753,17 +749,16 @@ 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)
# Save in the Linen on-disk layout so 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)
Expand Down
6 changes: 3 additions & 3 deletions src/maxtext/common/train_state_nnx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,9 @@ def apply_gradients(self, grads: Any, **kwargs):

# On-disk checkpoint format.
#
# A pure_nnx run saves in the same on-disk layout as a Linen run, so the two are
# interchangeable. The NNX state pure dict differs from Linen's in three ways, all
# reshaped below at save time:
# NNX saves in the same on-disk layout as the old Linen format, so checkpoints
# stay interchangeable. The NNX state pure dict differs from that layout in three
# ways, all reshaped below at save time:
# 1. top-level keys: {model, optimizer:{step, opt_state}} -> {params:{params:...}, step, opt_state}
# 2. weights: model/... -> params/params/... (Linen `params` collection)
# 3. opt_state: int-keyed dict (empty entries skipped) -> list with None for EmptyState,
Expand Down
Loading
Loading