feat: Auto-detect and convert Tunix post-training checkpoints on load - #4951
feat: Auto-detect and convert Tunix post-training checkpoints on load#4951hsuan-lun-chiang wants to merge 1 commit into
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements on-the-fly detection and conversion of Tunix-formatted checkpoints during the pre-training load process, avoiding manual checkpoint conversion. The feedback highlights several critical issues: first, setting metadata to None on exception in model_creation_utils.py will cause an AttributeError when accessing metadata.item_metadata later; second, _drop_adapter_level must recurse into lists, tuples, and namedtuples to handle nested optimizer states; and third, an _add_adapter_level helper is required to correctly wrap optimizer state fields in base when restoring adapter checkpoints to prevent structure mismatches.
| try: | ||
| metadata = ckptr.metadata(config.load_parameters_path) | ||
| except Exception: | ||
| metadata = None | ||
|
|
||
| # Get the structure of checkpoint in `config.load_parameters_path` | ||
| metadata = ckptr.metadata(config.load_parameters_path) | ||
| if metadata is None or metadata.item_metadata is None: | ||
| # Only error out if the directory doesn't even exist, defer metadata checks to the loader. | ||
| if not epath.Path(config.load_parameters_path).exists(): |
There was a problem hiding this comment.
If config.load_parameters_path points to a Tunix checkpoint, ckptr.metadata will raise an exception (since Checkpointer cannot read metadata from a multi-item CheckpointManager directory directly), and metadata will be set to None.
However, subsequent lines in from_pretrained (such as line 1019: metadata.item_metadata.tree.keys(), and lines 1088, 1089) access metadata directly without checking if it is None. This will cause an immediate AttributeError: 'NoneType' object has no attribute 'item_metadata' crash.
To fix this and seamlessly support Tunix checkpoints in from_pretrained, we should detect if the path contains a model_params subdirectory, and if so, point to it. This allows Checkpointer to successfully read the metadata and restore the parameters.
Note: You should also update line 1125 (which is outside this diff hunk) to use load_params_path instead of config.load_parameters_path for the restore call:
restored = ckptr.restore(
epath.Path(load_params_path),
...
) load_params_path = epath.Path(config.load_parameters_path)
if (load_params_path / "model_params").exists():
load_params_path = load_params_path / "model_params"
try:
metadata = ckptr.metadata(load_params_path)
except Exception:
metadata = None
# Only error out if the directory doesn't even exist, defer metadata checks to the loader.
if not load_params_path.exists():| 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()} | ||
| return tree |
There was a problem hiding this comment.
The current implementation of _drop_adapter_level only recurses into dictionaries. However, in Optax, the optimizer state (optimizer_state) is typically represented as a tuple of NamedTuples (e.g., ScaleByAdamState). Because _drop_adapter_level does not recurse into lists, tuples, or namedtuples, it will fail to drop the base adapter level from the optimizer state fields (like mu and nu), leading to a structure/shape mismatch when loading the checkpoint.
We should extend _drop_adapter_level to recurse into lists, tuples, and namedtuples so that it can correctly reach and unwrap the nested parameter-like states.
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 |
There was a problem hiding this comment.
To support wrapping the parameter-like fields of the optimizer state in base when has_base is True, we need an _add_adapter_level helper function. This function will recurse into lists, tuples, and namedtuples, and wrap any parameter-like dictionaries in {"base": ...}.
| 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 _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 |
| target_params = {"base": want_params} if has_base else want_params | ||
|
|
||
| if has_inject: | ||
| import jax.numpy as jnp | ||
| target_opt = { | ||
| "count": jnp.zeros((), dtype=jnp.int32), | ||
| "hyperparams": {}, | ||
| "hyperparams_states": {}, | ||
| "inner_state": want_opt | ||
| } | ||
| else: | ||
| target_opt = want_opt |
There was a problem hiding this comment.
When has_base is True (indicating a LoRA/adapter checkpoint), the parameter-like fields inside the optimizer state (such as mu and nu in Adam) are also wrapped in base on disk. However, target_opt is currently constructed using want_opt directly without wrapping its parameter-like fields in base. This will cause a structure mismatch and fail during the Orbax restore of optimizer_state.
To fix this, we should apply the _add_adapter_level helper to want_opt before wrapping it in inner_state (if has_inject is True).
target_params = {"base": want_params} if has_base else want_params
target_opt_base = _add_adapter_level(want_opt) if has_base else want_opt
if has_inject:
import jax.numpy as jnp
target_opt = {
"count": jnp.zeros((), dtype=jnp.int32),
"hyperparams": {},
"hyperparams_states": {},
"inner_state": target_opt_base
}
else:
target_opt = target_opt_base
0494c94 to
90aea05
Compare
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
82c484b to
41673bb
Compare
- Add auto-detection for Tunix-formatted checkpoints (checks for `model_params`) - Implement on-the-fly conversion logic (`_load_tunix_full_state_from_path`) during load - Seamlessly reshape Tunix `TrainState` in memory (strip `inject_hyperparams`, drop `base` adapter) - Support both full state and parameters-only restores
41673bb to
0097bba
Compare
1 participant
Overview
This PR updates the checkpoint loader in MaxText to auto-detect and dynamically convert Tunix post-training checkpoints on load, allowing downstream MaxText pre-training runs to restore from Tunix DPO/SFT checkpoints seamlessly.
Key Changes
model_paramscollection).inject_hyperparamswrapper states (_drop_inject_hyperparams)._drop_adapter_level) when mapping from adapter states to base pre-training states, handling dicts, lists, tuples, and namedtuples._add_adapter_level) when querying metadata or loading checkpoints withhas_base=True.model_creation_utils.from_pretrained: Directs parameter inspection and single-item PyTree restore tomodel_paramswhen loading Tunix checkpoints.gemma3-4bandgemma4-31b(Scanned DPO and SFT) with zero shape or key mismatches across multi-chip TPU topologies (v6e-32).Checklist