Support loading Tunix post-train checkpoints in MaxText pre-training - #4945
Draft
hsuan-lun-chiang wants to merge 35 commits into
Draft
Support loading Tunix post-train checkpoints in MaxText pre-training#4945hsuan-lun-chiang wants to merge 35 commits into
hsuan-lun-chiang wants to merge 35 commits into
Conversation
Post-training saved through Tunix's checkpoint manager, which stores
nnx.state(model) verbatim under a model_params item. MaxText's on-disk
layout is the Linen one: weights in params/params, the optimizer in
opt_state and step, and NNX-only state such as rngs in nnx_aux. The two
never matched, so a checkpoint from an SFT, DPO, RL or distillation run
could not be loaded by pre-training, and loaders had started growing
branches to read the post-training shape instead.
Add MaxTextLayoutCheckpointManager, which converts in both directions and
writes the same items tree everything else in MaxText reads. It lives
under trainers/post_train rather than common/checkpointing because it
subclasses Tunix's manager, and common/checkpointing is imported by
pre-training and inference, which run without Tunix installed. Older
checkpoints are still in the Tunix layout, so maybe_restore falls back to
the base class for those, and load_params_from_path learns to read them
by restoring into the NNX state itself.
DPO and RL train through TunixMaxTextAdapter, whose base level would
otherwise reach the checkpoint. Strip it from the weights and from the
optimizer accumulators that mirror them, and put it back on restore.
Two optimizer differences also blocked a full-state resume:
- train_dpo passed gradient_accumulation_steps unconditionally, so
Tunix wrapped the optimizer in optax.MultiSteps even at 1. Pass None
below 2, as train_sft already did.
- post-training chained optax.clip_by_global_norm into the optimizer,
nesting its state a level deeper than pre-training, which clips raw
gradients in its train step. add_gradient_clipping applies the same
math inside the update and keeps the optimizer's own state tree.
A DPO checkpoint now resumes into pre-training with its weights, its
optimizer state and its step counter. RL and distillation additionally
wrap in optax.inject_hyperparams so Tunix can log the learning rate,
which keeps their full-state resume out of reach; their weights load.
The RL and SFT demo notebooks read the checkpoint back to convert it to
HuggingFace format, so update the paths they hardcode. Both now point at
the items directory rather than model_params. RL also drops the actor
level: Tunix appended it only when it owned checkpoint_root_directory,
which is now None, so the actor trainer writes straight to checkpoint_dir
as SFT and DPO do.
A checkpoint also records what the run that wrote it was configured as:
scan_layers, and the LoRA settings when there are any. verify_and_sync_scan_layers
and sync_lora_metadata read it back to fill in a value the run left at its default,
or to reject one that contradicts the checkpoint. Tunix passed none, so
post-training checkpoints arrived with an empty dict and got neither behaviour. The
manager now takes the run's config and records the same metadata pre-training does,
through one shared builder so the two cannot drift apart. All four trainers pass
their config, distillation the student's.
Reading it back was broken too, and not only for post-training. The metadata belongs
to the step, so it sits at <step>/ and not at <step>/items/, while every caller
passes load_parameters_path, which by convention points at the item. Both readers
have been getting an empty dict from pre-training checkpoints as well and silently
doing nothing with it. load_checkpoint_metadata now falls back to the parent
directory, so both spellings work.
The layout conversion is covered for scanned and unscanned models, which the
RL scripts ship both of and vLLM requires the latter of.
The SFT and DPO pipelines called transformers.AutoTokenizer directly, so a run configured with tokenizer_type=sentencepiece or tiktoken either loaded an HF tokenizer it was not asked for or failed outright when there was no HF repo behind the path. Dispatch on tokenizer_type and fall back to input_pipeline_utils.get_tokenizer for the native ones. Those tokenizers spell the pad token differently -- pad_id and unk_id rather than pad_token_id and unk_token_id -- so _get_pad_id looks for both before giving up, and tokenization() takes the encode() path for tokenizers that are not callable the way an HF one is. Two smaller fixes ride along. DPO passed max_prompt_length through unguarded, so a config that left it unset produced a negative max_response_length; it now defaults to half of max_target_length. And the checkpoint manager takes root_directory directly instead of reading it back off the manager the base class built, with a close() to match.
RL applies a chat template to build its prompts, which until now meant only HF tokenizers could drive it: the sentencepiece and tiktoken classes have no apply_chat_template. ChatTemplateMixin renders the Jinja template directly and resolves bos/eos from whichever backend the tokenizer wraps, so both native classes gain the method. HFTokenizer forwards unknown attributes to the tokenizer it wraps, so callers reaching for an HF-only attribute keep working through it. train_rl builds its tokenizer through input_pipeline.tokenizer.build_tokenizer rather than AutoTokenizer, so the tokenizer_type in the config is the one that takes effect, and reads pad_id off it instead of pad_token_id.
Resuming a run and saving at the step it resumed from makes Orbax raise StepAlreadyExistsError, which killed the job. Check all_steps() first and return False, and catch the error as a backstop for the race between the two. RL also points vLLM at vllm_hf_config_path when one is set, falling back to tokenizer_path, and passes trust_remote_code so models whose HF repo ships custom code can load.
Drives every trainer over every model and scan mode and records what loads what. Phase A trains a pre-training base and reloads it from SFT, DPO and RL. Phase B does the reverse: each trainer writes a checkpoint, reloads it itself, and then pre-training reloads it, which is the direction the layout change exists for. Results land in a CSV, per-job logs under local_logs, and checkpoints in GCS.
MaxTextLayoutCheckpointManager replaces the manager Tunix's base class builds over its own item names. The guard that closed it looked for a _checkpointer attribute, which the base class does not have, so it never fired and the replaced manager was left open with its handles and threads. Close _checkpoint_manager, which is what the base class actually sets. The regression test tracks close calls across both managers built during __init__, and checks the base class's is closed while the live one is not. It fails without the fix.
pre-commit runs over every file a PR changes, so the script fails the code quality check as it stands: no docstrings, open() without an encoding, a Popen left unclosed, and an unused import. Add the docstrings, pass encoding, close the Popen through a context manager, and use rsplit for the model suffix. The helpers inside the model loop do close over the loop variables, but each is only called within the iteration that defines it, so that warning is silenced with a note rather than restructured. No behaviour change.
tpu_platform mutates data_parallel_size and then deletes sharding_config before vLLM calls with_hf_config, so rebuilding the config trips an assertion on the device_indexes length. Recover the original value from the sharding strategy's device_indexes and put it back on the way through. DPO builds its optimizer and reference model inside the mesh and axis rules, so logical axes resolve during maybe_restore rather than at first use. The decoder reads mhc_expansion_rate defensively, since not every config defines it.
The patch replaced the attribute on the symbol this module imported, which is not the object vLLM looks up at call time. Assign it on vllm.config.vllm instead, so the wrapper is the one that actually runs.
wait_until_finished and close read _checkpoint_manager, which does not exist yet if the base class raised part-way through __init__. Reach for it with getattr so teardown on a half-built manager reports the original failure instead of an AttributeError on top of it.
RL needs a chat template, and llama3's tiktoken tokenizer has none, so every llama3 RL job in the matrix died at startup: ValueError: Tokenizer ... has no chat_template and config.chat_template / config.chat_template_path are both empty Swap in the instruct model's HF tokenizer for RL jobs only, and point vllm_hf_config_path at the same repo.
RL and distillation wrap their optimizer in optax.inject_hyperparams so Tunix can log the learning rate. That nests the optimizer state a level deeper than pre-training's, which was the last thing keeping their checkpoints from a full state resume: the weights loaded but opt_state did not line up. Drop the wrapper on save and rebuild it on restore from the freshly built optimizer, carrying the step into its count. What lands on disk is then the same opt_state pre-training writes. RL checkpoints the actor under its own subdirectory again, so the eval guide and the end-to-end RL tests move with it; they now read actor/<step>/items. And unwrap_model recurses, so a model behind more than one wrapper still resolves.
RL checkpoints the actor under its own subdirectory again, so the notebook's conversion step has to look there. It was reading the step directories straight out of checkpoint_dir and finding nothing: FileNotFoundError: .../rl_llama3_output/<run>/checkpoints/actor The eval guide and the end-to-end RL shell tests already moved with the code; this is the one reader that was left behind. The item name stays items.
train_distill sets is_managed_externally on the trainer, which it needs so Tunix
leaves the input pipeline and the checkpoint manager to MaxText. The same flag
also gates Tunix's own stopping rule:
if (
not self.is_managed_externally
and self.config.max_steps is not None
and self._train_steps >= self.config.max_steps
):
break
So max_steps never applied, and on a synthetic dataset -- which never ends --
the run never ended either. A steps=5 job was seen still training after fifteen
hours, having written more than fifteen hundred checkpoints and held the TPU the
whole time. The step count did reach the trainer; it only fed the loss schedules.
Rather than clear the flag, which would also hand back the resume and teardown
behaviour MaxText handles itself, bound the data. The loop already stops when the
iterator is exhausted, so give MaxTextToTunixIterator a batch budget: one step
consumes gradient_accumulation_steps batches, and a resumed run has spent some
already.
test_train_save_and_resume has been skipped as "Hangs indefinitely on synthetic
datasets" -- the same cause.
…estoration - Safely get _iter_steps from trainer before calculating batch_budget, preventing TypeError when MaxTextDistillationTrainer is mocked in unit tests. - Update MaxTextCheckpointManager.restore_iterator to use self._checkpoint_manager with ocp.args.Composite.
…eckpointBaseManagerTest
…ingFace AutoTokenizer
…rain_checkpointing_test
DPO runs with dataset_type=hf, and that pipeline tokenizes through transformers.AutoTokenizer. The matrix was handing it MaxText's own tokenizer assets, which AutoTokenizer cannot open: OSError: It looks like the config file at 'src/maxtext/assets/tokenizers/tokenizer.gemma3' is not a valid JSON file That worked only while hf_data_processing branched on tokenizer_type. It no longer does, so name the model's HF repo instead, the way the RL jobs already do. HF_IDS already maps every model MaxText knows to its repo.
to_checkpoint_dict moves the optimizer accumulators into the Linen `params` collection by finding mu and nu at the top of the optimizer state. RL and distillation wrap their optimizer in optax.inject_hyperparams, whose state keys sit above those, so the conversion found nothing to move and the shell was stripped afterwards, leaving them bare: pre-training / SFT / DPO opt_state/0/mu/params/decoder/decoder_norm/scale distillation / RL opt_state/0/mu/decoder/decoder_norm/scale Every distillation checkpoint written so far has this, on every model and both scan modes, and a resume that carries opt_state cannot line them up. Convert what was behind the shell after stripping it. opt_state_to_linen exposes the conversion for that, rather than reaching into a private name from another module. Nothing caught this. The end-to-end reloads go through load_parameters_path, which restores only the params item and never looks at opt_state, so the weights loaded and 96 jobs passed over it. The test compares where mu lands for an optimizer with and without the wrapper. It needs a schedule, not a plain float -- optax only builds hyperparams_states for a callable, and without it the shell has a different shape and the bug does not reproduce.
Three tokenizer problems in the validation matrix, all the same shape: a job was
handed a tokenizer it could not open.
- The qwen, olmo and gpt-oss branch built a path from the model name, and no
such asset exists. With tokenizer_type=huggingface the path was then read as
a repo id: "Repo id must be in the form 'repo_name' or 'namespace/repo_name'".
- mistral had no branch at all and silently fell back to the config default.
- RL only overrode the tokenizer for llama3, though it samples through vLLM for
every model and applies a chat template, which MaxText's own assets do not
carry. Base checkpoints are not all instruction tuned, so a template is passed
for models whose tokenizer ships none.
DPO needs the same treatment: it reads dataset_type=hf, and that pipeline
tokenizes through AutoTokenizer, which cannot open a local sentencepiece file.
HF_IDS already maps every model MaxText knows to its repo, so use it instead of
building paths or hardcoding names.
…rministic unit testing
![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 introduces an end-to-end TPU validation matrix and refactors post-training checkpointing (SFT, DPO, RL, and distillation) to use MaxText's native on-disk layout, ensuring compatibility with pre-training. It also adds gradient clipping support without nesting optimizer states, integrates Jinja2 chat template rendering for native tokenizers, and updates tokenizer padding handling. The review feedback highlights several critical improvements: ensuring the adapter-level stripping logic doesn't corrupt legitimate nested 'base' keys in models, using epath.Path to robustly handle trailing slashes in checkpoint paths, and correctly invoking pad_id and unk_id when they are callable methods in SentencePiece tokenizers to prevent downstream type errors.
Comment on lines
+64
to
+82
| 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 |
There was a problem hiding this comment.
The recursive implementation of _drop_adapter_level strips any dictionary that has exactly the single key "base". If the model itself contains a legitimate nested module or parameter named "base" as its sole child, it will be incorrectly stripped on save. This leads to checkpoint incompatibility when pre-training tries to load the checkpoint, as it expects the "base" key to be present.
To prevent this, restrict the stripping of "base" to only the root of known optimizer accumulator subtrees (e.g., mu, nu, acc_grads, ema, trace, sum_gradients_sq).
Suggested change
| 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 _drop_adapter_level(tree): | |
| """Removes the adapter level from the root of optimizer accumulator subtrees. | |
| 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. | |
| To avoid corrupting legitimate nested "base" keys in the model, we only strip the "base" | |
| key when it is the immediate child of known accumulator keys. | |
| """ | |
| accumulator_keys = {"mu", "nu", "acc_grads", "ema", "trace", "sum_gradients_sq"} | |
| if isinstance(tree, dict): | |
| new_tree = {} | |
| for k, v in tree.items(): | |
| if k in accumulator_keys and isinstance(v, dict) and list(v.keys()) == [_ADAPTER_CHILD]: | |
| new_tree[k] = v[_ADAPTER_CHILD] | |
| else: | |
| new_tree[k] = _drop_adapter_level(v) | |
| return new_tree | |
| if isinstance(tree, list): | |
| return [_drop_adapter_level(v) for v in tree] | |
| return tree |
Comment on lines
+85
to
+101
| 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.
Update _add_adapter_level to symmetrically restore the "base" key only under the same known optimizer accumulator keys, matching the safer non-recursive stripping logic.
Suggested change
| 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 | |
| def _add_adapter_level(tree, guide): | |
| """Inverse of `_drop_adapter_level`.""" | |
| accumulator_keys = {"mu", "nu", "acc_grads", "ema", "trace", "sum_gradients_sq"} | |
| if isinstance(guide, dict) and isinstance(tree, dict): | |
| new_tree = {} | |
| for k, v in tree.items(): | |
| if k in guide: | |
| g_v = guide[k] | |
| if k in accumulator_keys and isinstance(g_v, dict) and list(g_v.keys()) == [_ADAPTER_CHILD]: | |
| new_tree[k] = {_ADAPTER_CHILD: _add_adapter_level(v, g_v[_ADAPTER_CHILD])} | |
| else: | |
| new_tree[k] = _add_adapter_level(v, g_v) | |
| else: | |
| new_tree[k] = v | |
| return new_tree | |
| 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 |
| # `nnx.state(model)`, rather than the Linen on-disk layout. Post-training wrote this before it | ||
| # moved to MaxText's layout, and the training engine still does. The tree is the whole | ||
| # checkpoint, not an item inside one. | ||
| is_nnx_native = os.path.basename(load_parameters_from_path) in ("model_params", "model") |
There was a problem hiding this comment.
Using os.path.basename on load_parameters_from_path can fail to identify the native NNX checkpoint layout if the path contains a trailing slash (e.g., gs://bucket/model_params/), as os.path.basename returns an empty string in such cases. Using epath.Path(load_parameters_from_path).name is more robust and handles trailing slashes correctly.
Suggested change
| is_nnx_native = os.path.basename(load_parameters_from_path) in ("model_params", "model") | |
| is_nnx_native = epath.Path(load_parameters_from_path).name in ("model_params", "model") |
Comment on lines
+41
to
+44
| elif hasattr(tokenizer, "pad_id") and getattr(tokenizer, "pad_id", None) is not None: | ||
| pad_id = tokenizer.pad_id | ||
| elif hasattr(tokenizer, "unk_id") and getattr(tokenizer, "unk_id", None) is not None: | ||
| pad_id = tokenizer.unk_id |
There was a problem hiding this comment.
In some tokenizers (such as SentencePiece), pad_id and unk_id are methods rather than properties. Accessing them directly without calling them will store the method object itself, leading to downstream type errors. Check if they are callable and invoke them if necessary.
Suggested change
| elif hasattr(tokenizer, "pad_id") and getattr(tokenizer, "pad_id", None) is not None: | |
| pad_id = tokenizer.pad_id | |
| elif hasattr(tokenizer, "unk_id") and getattr(tokenizer, "unk_id", None) is not None: | |
| pad_id = tokenizer.unk_id | |
| elif hasattr(tokenizer, "pad_id") and getattr(tokenizer, "pad_id", None) is not None: | |
| pad_id = tokenizer.pad_id() if callable(tokenizer.pad_id) else tokenizer.pad_id | |
| elif hasattr(tokenizer, "unk_id") and getattr(tokenizer, "unk_id", None) is not None: | |
| pad_id = tokenizer.unk_id() if callable(tokenizer.unk_id) else tokenizer.unk_id |
| trainer_devices, | ||
| sampler_devices, | ||
| tokenizer_pad_id=model_tokenizer.pad_token_id, | ||
| tokenizer_pad_id=model_tokenizer.pad_id, |
There was a problem hiding this comment.
If model_tokenizer is a SentencePiece-based tokenizer, pad_id is a method rather than an integer property. Passing model_tokenizer.pad_id directly will pass the method object itself, leading to downstream type errors. Use a check to call the method if it is callable, or default to a safe value.
Suggested change
| tokenizer_pad_id=model_tokenizer.pad_id, | |
| tokenizer_pad_id=model_tokenizer.pad_id() if callable(getattr(model_tokenizer, "pad_id", None)) else getattr(model_tokenizer, "pad_id", -1), |
hsuan-lun-chiang
changed the title
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
3 participants
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
This PR implements on-the-fly checkpoint conversion and loading for Tunix / post-training checkpoints directly within MaxText pre-training loaders.
Key Changes
src/maxtext/common/checkpointing.py, added robust detection for Tunix checkpoints (model_paramson-disk layout)._build_tunix_target_tree, inspects on-disk metadata to reconstruct the exact on-disk PyTree structure, providing expected sharding/shapes for model parameters while safely capturing auxiliary RNG / dropout states.{"value": Array}wrappers so checkpoints restore seamlessly into Linen and NNX states.Phase A pre_train(save Linen checkpoint)[PASS]Phase A SFT(reload Linen checkpoint into SFT)[PASS]Phase B SFT(reload Tunix checkpoint into SFT)[PASS]Phase B pre_train(reload Tunix checkpoint into pre-training)[PASS]