Write post-training checkpoints in MaxText's on-disk layout - #4975
Draft
ecnal-cienet wants to merge 1 commit into
Draft
Write post-training checkpoints in MaxText's on-disk layout#4975ecnal-cienet wants to merge 1 commit into
ecnal-cienet wants to merge 1 commit into
Conversation
Tunix writes its own checkpoint shape -- model_params/ and optimizer_state/ beside a custom_metadata -- while pre-training expects items/ holding params under the Linen params collection, opt_state, step and nnx_aux. A model trained by SFT, DPO, RL or distillation therefore could not be handed back to pre-training without converting it by hand. MaxTextLayoutCheckpointManager subclasses Tunix's manager and reshapes the tree on the way out: it restores the Linen collection level, strips the inject_hyperparams shell that optax wraps a scheduled optimizer in -- otherwise mu and nu land a level short of where pre-training looks for them -- and drops the adapter level a LoRA wrapper adds. install() swaps it in for the manager the base class built, closing that one so it does not leak its writer thread. Gradient clipping moves out of optax.chain for the same reason: chaining a stateless clip in front of the optimizer nests its state under an extra level that pre-training cannot read. RL keeps its actor/ level, since GRPO checkpoints an actor and a reference separately. The end-to-end scripts, the demo notebooks and the eval README follow the path from model_params to items. Three things had to work before that code was reachable. Post-training reached for transformers' AutoTokenizer directly, so a run had to name a HuggingFace repo even when the model ships a sentencepiece or tiktoken asset in this repo. The tokenizers built by build_tokenizer now serve post-training too, and the two that are ours gain a chat template renderer, since RL applies one and only HuggingFace tokenizers carried that method. The vLLM rollout had three call-site assumptions that no longer hold, each of which stops RL before a single rollout: EngineArgs dropped swap_space, and passing a field it no longer takes is a TypeError rather than something vLLM ignores; tpu_inference mutates data_parallel_size and then deletes the sharding config before calling with_hf_config, leaving the rebuilt config to fail its own device-count assertion; is_init_field raises on the fields tpu_inference injects dynamically. Distillation then ran past its step count, because Tunix stops on max_steps only when it owns the checkpoint manager and installing ours sets is_managed_externally. The iterator carries a batch budget of steps times gradient accumulation, minus whatever a resumed run already consumed, and raises StopIteration once spent. post_train_skip_checkpointing exists for runs that only want to load base weights and write nothing back.
![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 aligns post-training workflows (SFT, DPO, RL, and distillation) with MaxText's native on-disk checkpoint layout, ensuring compatibility with pre-training and inference. It introduces a post_train_skip_checkpointing option, adds Jinja2 chat template rendering to native tokenizers, and implements inline gradient clipping to preserve optimizer state structure. Feedback on these changes suggests refining the adapter-stripping logic to prevent accidental removal of legitimate 'base' keys, suppressing noisy warning logs when metadata is loaded from item-level paths, and using getattr for decode_context_parallel_size to improve robustness across different vLLM versions.
Comment on lines
+64
to
+101
| def _drop_adapter_level(tree): | ||
| """Removes the adapter level wherever it wraps a weight-shaped subtree. | ||
|
|
||
| The optimizer is built over the adapter, so its accumulators (mu, nu, acc_grads) are keyed by | ||
| the adapter's graph and carry the level even though the weights they shadow do not. | ||
|
|
||
| Args: | ||
| tree: A pure dict, typically the optimizer state. | ||
|
|
||
| Returns: | ||
| The same tree with every `{"base": subtree}` replaced by `subtree`. | ||
| """ | ||
| if isinstance(tree, dict): | ||
| if set(tree) == {_ADAPTER_CHILD}: | ||
| return _drop_adapter_level(tree[_ADAPTER_CHILD]) | ||
| return {k: _drop_adapter_level(v) for k, v in tree.items()} | ||
| if isinstance(tree, list): | ||
| return [_drop_adapter_level(v) for v in tree] | ||
| return tree | ||
|
|
||
|
|
||
| def _add_adapter_level(tree, guide): | ||
| """Inverse of `_drop_adapter_level`. | ||
|
|
||
| Args: | ||
| tree: A pure dict with the adapter level removed. | ||
| guide: The same tree before removal, giving the positions to restore. | ||
|
|
||
| Returns: | ||
| `tree` with the adapter level put back wherever `guide` carries it. | ||
| """ | ||
| if isinstance(guide, dict) and set(guide) == {_ADAPTER_CHILD}: | ||
| return {_ADAPTER_CHILD: _add_adapter_level(tree, guide[_ADAPTER_CHILD])} | ||
| if isinstance(guide, dict) and isinstance(tree, dict): | ||
| return {k: (_add_adapter_level(v, guide[k]) if k in guide else v) for k, v in tree.items()} | ||
| if isinstance(guide, list) and isinstance(tree, list) and len(guide) == len(tree): | ||
| return [_add_adapter_level(t, g) for t, g in zip(tree, guide)] | ||
| return tree |
There was a problem hiding this comment.
The recursive implementation of _drop_adapter_level will strip any dictionary that has only the key "base" at any depth. If the model itself has a legitimate module or parameter named "base" as its sole child at some level, this function will incorrectly strip it, leading to structure mismatches during restore. Restricting the stripping of the "base" key to only the root of the parameter-structured subtrees (rather than recursively descending into the unwrapped subtree) avoids this issue.
def _drop_adapter_level(tree):
"""Removes the adapter level wherever it wraps a weight-shaped subtree.
The optimizer is built over the adapter, so its accumulators (mu, nu, acc_grads) are keyed by
the adapter's graph and carry the level even though the weights they shadow do not.
Args:
tree: A pure dict, typically the optimizer state.
Returns:
The same tree with every `{"base": subtree}` replaced by `subtree`.
"""
if isinstance(tree, dict):
if set(tree) == {_ADAPTER_CHILD}:
return tree[_ADAPTER_CHILD]
return {k: _drop_adapter_level(v) for k, v in tree.items()}
if isinstance(tree, list):
return [_drop_adapter_level(v) for v in tree]
return tree
def _add_adapter_level(tree, guide):
"""Inverse of `_drop_adapter_level`.
Args:
tree: A pure dict with the adapter level removed.
guide: The same tree before removal, giving the positions to restore.
Returns:
`tree` with the adapter level put back wherever `guide` carries it.
"""
if isinstance(guide, dict) and set(guide) == {_ADAPTER_CHILD}:
return {_ADAPTER_CHILD: tree}
if isinstance(guide, dict) and isinstance(tree, dict):
return {k: (_add_adapter_level(v, guide[k]) if k in guide else v) for k, v in tree.items()}
if isinstance(guide, list) and isinstance(tree, list) and len(guide) == len(tree):
return [_add_adapter_level(t, g) for t, g in zip(tree, guide)]
return tree
Comment on lines
+834
to
+868
| def _custom_metadata_at(checkpoint_dir: epath.Path) -> dict[str, Any]: | ||
| """Reads the custom metadata stored at exactly this directory. | ||
|
|
||
| Args: | ||
| checkpoint_dir: Directory to read. | ||
|
|
||
| Returns: | ||
| The metadata dict, empty if there is none or the read fails. | ||
| """ | ||
| checkpoint_dir = epath.Path(checkpoint_dir_path) | ||
| try: | ||
| ckptr = ocp.StandardCheckpointer() | ||
| metadata = ckptr.metadata(checkpoint_dir) | ||
| metadata = ocp.StandardCheckpointer().metadata(checkpoint_dir) | ||
| return metadata.custom_metadata or {} | ||
| except Exception as e: # pylint: disable=broad-except | ||
| max_logging.log(f"Warning: Failed to load checkpoint metadata: {e}") | ||
| return {} | ||
|
|
||
|
|
||
| def load_checkpoint_metadata(checkpoint_dir_path: str) -> dict[str, Any]: | ||
| """Loads custom metadata from an Orbax checkpoint. | ||
|
|
||
| The metadata belongs to the step, so it sits at `<step>/` and not at `<step>/items/`. Callers | ||
| pass `load_parameters_path`, which points at the item, so fall back to the parent directory. | ||
|
|
||
| Args: | ||
| checkpoint_dir_path: Path to the checkpoint directory, item level or step level. | ||
|
|
||
| Returns: | ||
| A dictionary containing custom metadata, or an empty dictionary if none is | ||
| present or loading fails. | ||
| """ | ||
| checkpoint_dir = epath.Path(checkpoint_dir_path) | ||
| metadata = _custom_metadata_at(checkpoint_dir) | ||
| if not metadata and checkpoint_dir.parent != checkpoint_dir: | ||
| metadata = _custom_metadata_at(checkpoint_dir.parent) | ||
| return metadata |
There was a problem hiding this comment.
Since load_checkpoint_metadata often receives an item-level path (e.g., ending in /items) where checkpoint metadata is not directly stored, the first call to _custom_metadata_at is expected to fail. Logging a warning on this first attempt creates misleading and noisy warning logs during normal operation. Consider adding a log_warning parameter to _custom_metadata_at to suppress logging on the initial attempt.
def _custom_metadata_at(checkpoint_dir: epath.Path, log_warning: bool = True) -> dict[str, Any]:
"""Reads the custom metadata stored at exactly this directory.
Args:
checkpoint_dir: Directory to read.
log_warning: Whether to log a warning if the read fails.
Returns:
The metadata dict, empty if there is none or the read fails.
"""
try:
metadata = ocp.StandardCheckpointer().metadata(checkpoint_dir)
return metadata.custom_metadata or {}
except Exception as e: # pylint: disable=broad-except
if log_warning:
max_logging.log(f"Warning: Failed to load checkpoint metadata: {e}")
return {}
def load_checkpoint_metadata(checkpoint_dir_path: str) -> dict[str, Any]:
"""Loads custom metadata from an Orbax checkpoint.
The metadata belongs to the step, so it sits at `<step>/` and not at `<step>/items/`. Callers
pass `load_parameters_path`, which points at the item, so fall back to the parent directory.
Args:
checkpoint_dir_path: Path to the checkpoint directory, item level or step level.
Returns:
A dictionary containing custom metadata, or an empty dictionary if none is
present or loading fails.
"""
checkpoint_dir = epath.Path(checkpoint_dir_path)
metadata = _custom_metadata_at(checkpoint_dir, log_warning=False)
if not metadata and checkpoint_dir.parent != checkpoint_dir:
metadata = _custom_metadata_at(checkpoint_dir.parent, log_warning=True)
return metadata| sp = sharding_strategy.get("sequence_parallelism", 1) | ||
| attn_dp = sharding_strategy.get("attention_data_parallelism", 1) | ||
| attn_dp_ep = sharding_strategy.get("attention_data_expert_parallelism", 1) | ||
| dcp = pc.decode_context_parallel_size |
There was a problem hiding this comment.
To ensure robustness across different vLLM versions where decode_context_parallel_size might be absent or named differently, it is safer to use getattr with a default fallback value.
Suggested change
| dcp = pc.decode_context_parallel_size | |
| dcp = getattr(pc, "decode_context_parallel_size", 1) |
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Post-training writes Tunix's checkpoint shape:
Pre-training reads MaxText's:
So a model trained by SFT, DPO, RL or distillation could not be handed back to pre-training without converting it by hand. This makes post-training write the layout pre-training already reads.
MaxTextLayoutCheckpointManagersubclasses Tunix's manager and reshapes the tree on the way out. It restores the Linen collection level; strips theinject_hyperparamsshell that optax wraps a scheduled optimizer in, without which mu and nu land a level short of where pre-training looks for them; and drops the adapter level a LoRA wrapper adds.install()swaps it in for the manager the base class built, closing that one so it does not leak its writer thread. Restoring still reads the old layout, so a run resumed from an existing checkpoint converts itself on its next save.Gradient clipping moves out of
optax.chainfor the same reason: chaining a stateless clip in front of the optimizer nests its state under an extra level pre-training cannot read.add_gradient_clippingapplies the clip while leaving the optimizer's state shape alone.RL keeps its
actor/level, since GRPO checkpoints an actor and a reference separately. The end-to-end scripts, both demo notebooks and the eval README follow the path frommodel_paramstoitems.Three things had to work before that code was reachable, and they are in this PR because the change is not testable end to end without them:
AutoTokenizerdirectly, so a run had to name a HuggingFace repo even when the model ships a sentencepiece or tiktoken asset in this repo, and a gated repo needed a token for something the checkpoint already carried. The tokenizersbuild_tokenizerreturns now serve post-training too, and the two that are ours gain a chat template renderer, since RL applies one and only HuggingFace tokenizers carried that method.EngineArgsdroppedswap_space, and passing a field it no longer takes is aTypeErrorrather than something vLLM ignores;tpu_inferencemutatesdata_parallel_sizeand then deletes the sharding config before callingwith_hf_config, leaving the rebuilt config to fail its own device-count assertion;is_init_fieldraises on the fieldstpu_inferenceinjects dynamically.max_stepsonly when it owns the checkpoint manager, and installing ours setsis_managed_externally, so distillation ran past its step count until the input pipeline ran dry. The iterator now carries a batch budget of steps times gradient accumulation, minus whatever a resumed run already consumed, and raisesStopIterationonce spent.post_train_skip_checkpointingis new, for runs that only want to load base weights and write nothing back.Tests
tests/post_training/— 285 passed, 94 skipped.New unit coverage for the manager itself (
post_train_checkpointing_test.py,distillation_checkpointing_test.py), for the clipping wrapper(
optimizers_test.py::AddGradientClippingTest), and for the layout on the restore side(
checkpointing_test.py,checkpointing_nnx_load_test.py).End to end, on a v6e-8: every trainer wrote a checkpoint and every other trainer, plus
pre-training, read it back. The written trees were checked rather than the exit codes — a run can
exit zero and leave nothing usable. A converted checkpoint holds
which is what pre-training expects, including mu and nu inside the params collection.
RL was run on qwen3-0.6b, qwen2.5-1.5b and gemma3-4b, scanned and unscanned, each writing
actor/1andactor/2. llama3.1-8b and qwen3-8b run out of HBM on a single v6e-8 host at thisbatch size, in the gradient update rather than anywhere near the checkpoint code.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.