Convert Tunix-layout post-training checkpoints for pre-training by ecnal-cienet · Pull Request #4976 · AI-Hypercomputer/maxtext

ecnal-cienet · 2026-08-23T23:38:16Z

Description

Stacked on #4975. Retarget to main once that merges.

Post-training now writes MaxText's layout, and the trainers still restore the old one, so a run resumed from an existing checkpoint rewrites itself on its next save. Nothing does that for the checkpoints nobody is going to resume, and pre-training cannot read them: it looks for weights in the Linen params collection, and the old layout keeps the whole model state under model_params instead.

maxtext.checkpoint_conversion.tunix_to_maxtext rewrites them in place of a training run:

python -m maxtext.checkpoint_conversion.tunix_to_maxtext src/maxtext/configs/base.yml \
  --source gs://bucket/run/checkpoints \
  model_name=llama3.1-8b scan_layers=True \
  base_output_directory=gs://bucket/run_converted \
  hardware=cpu skip_jax_distributed_system=True

It finds every step still in the old layout, including the actor/<step> level RL writes, and skips any that already has an items/. --dry_run reports what each step would hold without writing.

The split between params and nnx_aux is why this needs a model rather than being a rename: weights and the rng counters that drive dropout sit mixed together under model_params, and only the model says which is which. Weights are never materialised — an abstract model supplies the classification and the arrays move through as numpy, so it runs on CPU.

It refuses a model that does not match the checkpoint, by name and by shape. Both mistakes are easy to make and neither announces itself: splitting by the wrong model routes the leaves it does not recognise into nnx_aux, and a config that differs only in size has every name right and fails much later, in the middle of a read. Both happened while developing this, which is why the check is there.

Tests

tests/post_training/unit/tunix_to_maxtext_test.py — 9 passed. Built on a checkpoint the test writes in the old layout, so it depends on nothing external. It covers that every leaf survives the split, that rng state stays out of the params collection, that the inject_hyperparams shell is stripped, that RL's actor/<step> steps are found, that already-converted steps are skipped, and that a mismatched model is refused.

The test that matters most loads a converted checkpoint back through checkpointing.load_params_from_path — pre-training's own loader — and checks the values.

Also run against a real 2026-06 post-training checkpoint on GCS (gpt3-52k, scanned): 46 arrays split into 30 aux and 16 weights, written, and loaded back through the pre-training loader as 16 tensors totalling 548,080 parameters, no NaNs.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

Trainers restore the old layout already, so a run resumed from one rewrites itself on its next save. Nothing does that for the checkpoints nobody is going to resume, and pre-training cannot read them: it looks for weights in the Linen params collection, and the old layout keeps the whole model state under model_params instead. The split between params and nnx_aux is why this needs a model rather than a rename. Weights and the rng counters that drive dropout sit mixed together under model_params, and only the model says which is which. An abstract model supplies the classification, and the arrays move through as numpy, so the conversion runs on CPU without materialising the model. It refuses a model that does not match the checkpoint, by name and by shape. Both mistakes are easy to make and neither announces itself: splitting by the wrong model routes the leaves it does not recognise into nnx_aux, and a config that differs only in size has every name right and fails much later, in the middle of a read.

gemini-code-assist

Code Review

This pull request updates the post-training trainers (SFT, DPO, RL, and distillation) to write checkpoints in MaxText's Linen-based on-disk layout instead of Tunix's layout, ensuring compatibility with pre-training and inference. It introduces a layout conversion script, a custom checkpoint manager, inline gradient clipping to preserve optimizer state structure, and vLLM integration fixes. The review feedback highlights several critical robustness improvements: ensuring optimizer states are always converted to the Linen layout during conversion, recursing into tuples and lists when unwrapping variables or dropping adapter levels, defensively handling callable tokenizer properties (like bos_id, eos_id, pad_id, and unk_id), and suppressing noisy warnings during metadata fallback reads.

gemini-code-assist · 2026-08-23T23:40:06Z

+    inner = _drop_inject_hyperparams(opt_state)
+    if inner is not opt_state:
+      inner = train_state_nnx.opt_state_to_linen(inner)
+    items["opt_state"] = inner


The check if inner is not opt_state: skips converting the optimizer state to the Linen layout if it was not wrapped in inject_hyperparams. However, since the old checkpoint was written in the Tunix (NNX) layout, the optimizer state is always in the NNX layout and must be converted to the Linen layout regardless of whether it was wrapped. We should always call train_state_nnx.opt_state_to_linen(inner).

Suggested change

inner = _drop_inject_hyperparams(opt_state)

if inner is not opt_state:

inner = train_state_nnx.opt_state_to_linen(inner)

items["opt_state"] = inner

inner = _drop_inject_hyperparams(opt_state)

inner = train_state_nnx.opt_state_to_linen(inner)

items["opt_state"] = inner

gemini-code-assist · 2026-08-23T23:40:06Z

+  if hasattr(tokenizer, "pad_token_id") and tokenizer.pad_token_id is not None:
    pad_id = tokenizer.pad_token_id
-  elif tokenizer.unk_token_id is not None:
+  elif hasattr(tokenizer, "unk_token_id") and tokenizer.unk_token_id is not None:
    pad_id = tokenizer.unk_token_id
+  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


In some tokenizers (e.g., SentencePiece), pad_id and unk_id are methods rather than properties. If they are methods, getattr(tokenizer, "pad_id", None) will return the method object itself, which will cause downstream type errors. We should check if they are callable and call them if so.

Suggested change

if hasattr(tokenizer, "pad_token_id") and tokenizer.pad_token_id is not None:

pad_id = tokenizer.pad_token_id

elif tokenizer.unk_token_id is not None:

elif hasattr(tokenizer, "unk_token_id") and tokenizer.unk_token_id is not None:

pad_id = tokenizer.unk_token_id

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

if hasattr(tokenizer, "pad_token_id") and tokenizer.pad_token_id is not None:

pad_id = tokenizer.pad_token_id() if callable(tokenizer.pad_token_id) else tokenizer.pad_token_id

elif hasattr(tokenizer, "unk_token_id") and tokenizer.unk_token_id is not None:

pad_id = tokenizer.unk_token_id() if callable(tokenizer.unk_token_id) else tokenizer.unk_token_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

gemini-code-assist · 2026-08-23T23:40:06Z

+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


In JAX/Optax, optimizer states are often represented as tuples of states (e.g., in optax.chain). Since _drop_adapter_level and _add_adapter_level only check for list and dict, they will completely skip recursing into tuples. This means the adapter level will not be dropped from chained optimizer states, leading to incorrect checkpoint layouts. We should support recursing into tuples as well.

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, tuple)): res = [_drop_adapter_level(v) for v in tree] return tuple(res) if isinstance(tree, tuple) else res 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, tuple)) and isinstance(tree, (list, tuple)) and len(guide) == len(tree): res = [_add_adapter_level(t, g) for t, g in zip(tree, guide)] return tuple(res) if isinstance(guide, tuple) else res return tree

gemini-code-assist · 2026-08-23T23:40:06Z

+def _unwrap_variables(tree):
+  """Turns the `{'value': array}` a serialised nnx.Variable becomes back into the array."""
+  if isinstance(tree, dict):
+    if set(tree) == {"value"}:
+      return tree["value"]
+    return {k: _unwrap_variables(v) for k, v in tree.items()}
+  return tree


The _unwrap_variables function only handles dictionaries, but optimizer states or other PyTree structures can contain lists or tuples. If nnx.Variable objects are nested inside lists or tuples, they will not be unwrapped. We should support recursing into lists and tuples as well.

Suggested change

def _unwrap_variables(tree):

"""Turns the `{'value': array}` a serialised nnx.Variable becomes back into the array."""

if isinstance(tree, dict):

if set(tree) == {"value"}:

return tree["value"]

return {k: _unwrap_variables(v) for k, v in tree.items()}

return tree

def _unwrap_variables(tree):

"""Turns the `{'value': array}` a serialised nnx.Variable becomes back into the array."""

if isinstance(tree, dict):

if set(tree) == {"value"}:

return tree["value"]

return {k: _unwrap_variables(v) for k, v in tree.items()}

if isinstance(tree, (list, tuple)):

res = [_unwrap_variables(v) for v in tree]

return tuple(res) if isinstance(tree, tuple) else res

return tree

gemini-code-assist · 2026-08-23T23:40:06Z

+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


When load_checkpoint_metadata is called with an item-level path (which is standard now), the first call to _custom_metadata_at will fail and log a noisy warning Warning: Failed to load checkpoint metadata: ... before successfully falling back to the parent directory. We should suppress this warning on the first try if a fallback is possible, and only log it if the fallback also fails.

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) can_fallback = checkpoint_dir.parent != checkpoint_dir metadata = _custom_metadata_at(checkpoint_dir, log_warning=not can_fallback) if not metadata and can_fallback: metadata = _custom_metadata_at(checkpoint_dir.parent, log_warning=True) return metadata

gemini-code-assist · 2026-08-23T23:40:06Z

+    if hasattr(self, "_tokenizer_model"):  # SentencePiece
+      if self.bos_id is not None and self.bos_id >= 0:
+        bos_token = self._tokenizer_model.IdToPiece(self.bos_id)
+      if self.eos_id is not None and self.eos_id >= 0:
+        eos_token = self._tokenizer_model.IdToPiece(self.eos_id)
+    elif hasattr(self, "model"):  # TikToken
+      if self.bos_id is not None and self.bos_id >= 0:
+        bos_token = self.decode([self.bos_id])
+      if self.eos_id is not None and self.eos_id >= 0:
+        eos_token = self.decode([self.eos_id])


If self.bos_id or self.eos_id are methods in some tokenizers, comparing them directly with self.bos_id >= 0 will raise a TypeError. We should resolve them defensively by calling them if they are callable.

bos_id = self.bos_id() if callable(getattr(self, "bos_id", None)) else getattr(self, "bos_id", None) eos_id = self.eos_id() if callable(getattr(self, "eos_id", None)) else getattr(self, "eos_id", None) if hasattr(self, "_tokenizer_model"): if bos_id is not None and bos_id >= 0: bos_token = self._tokenizer_model.IdToPiece(bos_id) if eos_id is not None and eos_id >= 0: eos_token = self._tokenizer_model.IdToPiece(eos_id) elif hasattr(self, "model"): if bos_id is not None and bos_id >= 0: bos_token = self.decode([bos_id]) if eos_id is not None and eos_id >= 0: eos_token = self.decode([eos_id])

gemini-code-assist · 2026-08-23T23:40:06Z

+def _patched_eos_id(self):
+  # pylint: disable=protected-access
+  if hasattr(self._tokenizer, "eos_id") and not callable(self._tokenizer.eos_id):
+    return self._tokenizer.eos_id
+  return _old_eos_id(self)
+
+
+TokenizerAdapter.eos_id = _patched_eos_id
+
+_old_bos_id = TokenizerAdapter.bos_id
+
+
+def _patched_bos_id(self):
+  # pylint: disable=protected-access
+  if hasattr(self._tokenizer, "bos_id") and not callable(self._tokenizer.bos_id):
+    return self._tokenizer.bos_id
+  return _old_bos_id(self)
+
+
+TokenizerAdapter.bos_id = _patched_bos_id
+
+_old_pad_id = TokenizerAdapter.pad_id
+
+
+def _patched_pad_id(self):
+  # pylint: disable=protected-access
+  if hasattr(self._tokenizer, "pad_id") and not callable(self._tokenizer.pad_id):
+    pad_id = self._tokenizer.pad_id
+    if pad_id is None or pad_id < 0:
+      return self.eos_id()
+    return pad_id
+  return _old_pad_id(self)


The monkey-patched methods for TokenizerAdapter assume that eos_id, bos_id, and pad_id are either non-callable properties or should fall back to the old implementation. However, if they are callable methods on the tokenizer, we should call them to get the correct ID rather than falling back to the HF-specific _old_* methods which might fail. Let's handle both callable and non-callable cases robustly.

def _patched_eos_id(self): # pylint: disable=protected-access if hasattr(self._tokenizer, "eos_id"): return self._tokenizer.eos_id() if callable(self._tokenizer.eos_id) else self._tokenizer.eos_id return _old_eos_id(self) TokenizerAdapter.eos_id = _patched_eos_id _old_bos_id = TokenizerAdapter.bos_id def _patched_bos_id(self): # pylint: disable=protected-access if hasattr(self._tokenizer, "bos_id"): return self._tokenizer.bos_id() if callable(self._tokenizer.bos_id) else self._tokenizer.bos_id return _old_bos_id(self) TokenizerAdapter.bos_id = _patched_bos_id _old_pad_id = TokenizerAdapter.pad_id def _patched_pad_id(self): # pylint: disable=protected-access if hasattr(self._tokenizer, "pad_id"): pad_id = self._tokenizer.pad_id() if callable(self._tokenizer.pad_id) else self._tokenizer.pad_id if pad_id is None or pad_id < 0: return self.eos_id() return pad_id return _old_pad_id(self)

codecov · 2026-08-23T23:45:48Z

Codecov Report

❌ Patch coverage is 19.44444% with 116 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
.../maxtext/checkpoint_conversion/tunix_to_maxtext.py	19.44%	115 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

gemini-code-assist Bot reviewed Aug 23, 2026

View reviewed changes

ecnal-cienet changed the base branch from main to feat/post-train-writes-maxtext-layout August 23, 2026 23:43

Navigation Menu

Uh oh!

Convert Tunix-layout post-training checkpoints for pre-training - #4976

Convert Tunix-layout post-training checkpoints for pre-training#4976
ecnal-cienet wants to merge 1 commit into
feat/post-train-writes-maxtext-layoutfrom
feat/tunix-checkpoint-converter

ecnal-cienet commented Aug 23, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Uh oh!

codecov Bot commented Aug 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

ecnal-cienet commented Aug 23, 2026

Description

Tests

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Aug 23, 2026

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant