Write post-training checkpoints in MaxText's on-disk layout

docs/tutorials/posttraining/lora.md

@@ -190,7 +190,7 @@ After completing the fine-tuning process, your LoRA weights are stored in MaxTex
 ```sh
 python3 -m maxtext.checkpoint_conversion.to_huggingface \
     model_name="${MODEL_NAME?}" \
-    lora.lora_restore_path="${BASE_OUTPUT_DIRECTORY?}/${RUN_NAME?}/checkpoints/<STEPS>/model_params" \
+    lora.lora_restore_path="${BASE_OUTPUT_DIRECTORY?}/${RUN_NAME?}/checkpoints/<STEPS>/items" \
     base_output_directory="${BASE_OUTPUT_DIRECTORY?}/hf_lora_adapter"
 ```
 

src/maxtext/common/checkpointing.py

@@ -150,7 +150,7 @@ def _raise_on_weight_mismatch(want, have, config=None):
   )
 
 
-def _linen_items_to_nnx(restored_linen, abstract_nnx_state):
+def linen_items_to_nnx(restored_linen, abstract_nnx_state):
   """Reshapes a restored Linen-layout `items` dict into an NNX state.
 
   The inverse of `to_checkpoint_dict`, over the same `split_for_checkpoint` partition. The Linen
@@ -184,7 +184,7 @@ def _load_linen_checkpoint_into_nnx(
   """Restores a Linen-layout checkpoint into an NNX state (pure_nnx resume).
 
   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
+  `linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when
   present, else keep their fresh init value. A genuinely-missing weight raises.
   """
   max_logging.log(f"Restoring Linen-layout checkpoint into NNX state at {path}")
@@ -222,7 +222,7 @@ def _restored_linen_to_nnx(restored_linen, abstract_nnx_state, config=None):
   itself is the Linen one, since pure_nnx reads and writes the Linen on-disk layout.
   """
   _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored_linen), config=config)
-  return _linen_items_to_nnx(restored_linen, abstract_nnx_state)
+  return linen_items_to_nnx(restored_linen, abstract_nnx_state)
 
 
 def _abstract_params(abstract_unboxed_pre_state):
@@ -701,6 +701,25 @@ def setup_checkpoint_logger(config) -> Any | None:  # pytype: disable=attribute-
   return orbax_cloud_logger
 
 
+def _nnx_native_wrapper_key(ckptr, path):
+  """Returns the wrapper key an NNX-native checkpoint nests its weights under, if any.
+
+  Post-training wrote DPO and RL checkpoints through `TunixMaxTextAdapter`, whose only child
+  module is `base`, so their weights sit one level deeper. SFT and the training engine save the
+  model directly and do not.
+
+  Args:
+    ckptr: Checkpointer used to read the checkpoint metadata.
+    path: Path to the checkpoint.
+
+  Returns:
+    "base" if the checkpoint nests its weights under it, else None.
+  """
+  tree = ckptr.metadata(epath.Path(path)).item_metadata
+  tree = getattr(tree, "tree", tree)
+  return "base" if tree is not None and "base" in tree else None
+
+
 def load_params_from_path(
     load_parameters_from_path,
     abstract_unboxed_params,
@@ -718,15 +737,16 @@ def load_params_from_path(
   is_nnx = isinstance(abstract_unboxed_params, nnx.State)
   want = abstract_unboxed_params.to_pure_dict() if is_nnx else abstract_unboxed_params
 
-  # Determine the restore key based on the leaf directory name to support native and custom SFT
-  restore_key = os.path.basename(load_parameters_from_path)
-  if restore_key not in ("model_params", "model"):
-    restore_key = "params"
-
-  if restore_key in ("model_params", "model"):
-    params_collection = want
-  else:
-    params_collection = {"params": want} if is_nnx else want
+  # A path ending in `model_params` (or `model`) holds an NNX state written straight from
+  # `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")
+  if is_nnx_native and not is_nnx:
+    raise ValueError(
+        f"'{load_parameters_from_path}' holds an NNX state, which only restores into an NNX params "
+        "state. Point load_parameters_path at a checkpoint saved in the Linen on-disk layout instead."
+    )
 
   # *_concurrent_gb should be set for large models, the default is 96.
   max_logging.log(f"Creating checkpoint manager with ocdbt={use_ocdbt} and zarr3={use_zarr3}")
@@ -743,18 +763,31 @@ def load_params_from_path(
   # Rather than pass the entire abstract state, which could unnecessarily restore opt_state and such and waste
   # memory, we instead specify here that we are just restoring the params field of the checkpoint
   # (which itself may be a dictionary containing a key named 'params' or 'model').
-  restore_args = ocp.checkpoint_utils.construct_restore_args(params_collection)
-  restored = ckptr.restore(
-      epath.Path(load_parameters_from_path),
-      item={restore_key: params_collection},
-      transforms={},
-      restore_args={restore_key: restore_args},
-  )
-  restored_collection = restored[restore_key]
-
-  if restore_key in ("model_params", "model"):
-    restored_weights = restored_collection
+  if is_nnx_native:
+    wrapper_key = _nnx_native_wrapper_key(ckptr, load_parameters_from_path)
+    # Restore into the NNX state itself rather than a pure dict. Flax registers a Variable as a
+    # pytree holding its array under `value`, matching what `nnx.state(model)` wrote, so save and
+    # restore agree without reshaping either side by hand.
+    item = {wrapper_key: abstract_unboxed_params} if wrapper_key else abstract_unboxed_params
+    restored = ckptr.restore(
+        epath.Path(load_parameters_from_path),
+        item=item,
+        transforms={},
+        restore_args=ocp.checkpoint_utils.construct_restore_args(item),
+    )
+    # No `params` collection in this layout, so the weights are the whole collection.
+    restored_weights = nnx.to_pure_dict(restored[wrapper_key] if wrapper_key else restored)
+    restored_collection = restored_weights
   else:
+    params_collection = {"params": want} if is_nnx else want
+    restore_args = ocp.checkpoint_utils.construct_restore_args(params_collection)
+    restored = ckptr.restore(
+        epath.Path(load_parameters_from_path),
+        item={"params": params_collection},
+        transforms={},
+        restore_args={"params": restore_args},
+    )
+    restored_collection = restored["params"]
     restored_weights = restored_collection["params"] if is_nnx else restored_collection
 
   # `transforms={}` lets Orbax return an unmaterialized leaf for a weight the checkpoint lacks,
@@ -777,26 +810,64 @@ def save_params_to_path(checkpoint_dir, params, use_ocdbt=True, use_zarr3=True):
   print(f"Quantized params checkpoint saved at: {checkpoint_dir}")
 
 
-def load_checkpoint_metadata(checkpoint_dir_path: str) -> dict[str, Any]:
-  """Loads custom metadata from an Orbax checkpoint.
+def checkpoint_custom_metadata(config) -> dict[str, Any]:
+  """Returns the metadata a checkpoint stores alongside its state.
+
+  `verify_and_sync_scan_layers` and `lora_utils.sync_lora_metadata` read these back on load.
+  Post-training saves through its own manager, so it calls this too.
 
   Args:
-    checkpoint_dir_path: Path to the checkpoint directory.
+    config: The run's config, or None.
 
   Returns:
-    A dictionary containing custom metadata, or an empty dictionary if none is
-    present or loading fails.
+    The metadata dict, empty if there is no config to read it from.
+  """
+  custom_metadata = {}
+  if config:
+    if hasattr(config, "scan_layers"):
+      custom_metadata["scan_layers"] = config.scan_layers
+    if hasattr(config, "lora") and config.lora and getattr(config.lora, "lora_rank", 0) > 0:
+      custom_metadata["lora"] = config.lora.model_dump()
+  return custom_metadata
+
+
+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
+
+
 def _uses_local_checkpoint_period(config):
   return config.enable_emergency_checkpoint or config.enable_multi_tier_checkpointing
 
@@ -1030,12 +1101,7 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
           item=grain_iters_to_save
       )  # pyrefly: ignore[bad-assignment]
 
-  custom_metadata = {}
-  if config:
-    if hasattr(config, "scan_layers"):
-      custom_metadata["scan_layers"] = config.scan_layers
-    if hasattr(config, "lora") and config.lora and getattr(config.lora, "lora_rank", 0) > 0:
-      custom_metadata["lora"] = config.lora.model_dump()
+  custom_metadata = checkpoint_custom_metadata(config)
 
   match (checkpoint_manager, config, data_iterator):
     case (checkpoint_manager, _, _) if isinstance(

src/maxtext/common/train_state_nnx.py

@@ -120,6 +120,22 @@ def _as_chain_index(key):
   return None
 
 
+def opt_state_to_linen(opt_state):
+  """Reshapes an optimizer state into the Linen on-disk layout.
+
+  `to_checkpoint_dict` already does this for the state it is given. This is for a caller that
+  has unwrapped something the conversion could not see through -- `optax.inject_hyperparams`
+  hides mu and nu behind its own keys -- and needs the inner state converted on its own.
+
+  Args:
+    opt_state: The optimizer state to reshape.
+
+  Returns:
+    The same state in the Linen layout.
+  """
+  return _opt_state_to_linen(opt_state)
+
+
 def _opt_state_to_linen(opt_state):
   """Reshapes the NNX opt_state to Linen's on-disk layout.
 

src/maxtext/configs/base.yml

@@ -55,6 +55,8 @@ save_checkpoint_on_start: true
 save_checkpoint_on_completion: true
 async_checkpointing: true
 checkpoint_period: 10_000
+# Post-training only: skip installing a checkpoint manager, so nothing is written back.
+post_train_skip_checkpointing: false
 max_num_checkpoints_to_keep: None
 enable_continuous_checkpointing: false
 # enables one replica to read the ckpt then broadcast to the rest

src/maxtext/configs/types.py

@@ -348,6 +348,17 @@ class Checkpointing(BaseModel):
   load_checkpoint_only_once: bool = Field(False, description="If True, deep copy the reference model to the actor model.")
   async_checkpointing: bool = Field(True, description="If True, uses an asynchronous checkpointer for performance.")
   checkpoint_period: int = Field(10_000, description="The frequency (in steps) at which to save checkpoints.")
+  post_train_skip_checkpointing: bool = Field(
+      False,
+      description=(
+          "If True, post-training trainers do not install a checkpoint manager and write nothing."
+          " Base weights still load through load_parameters_path. Exists because enable_checkpointing"
+          " cannot express this: it is validated as required whenever load_parameters_path is set,"
+          " so it cannot be turned off by a run that has to load base weights. Setting"
+          " checkpoint_period to 0 is not an alternative either, since step % checkpoint_period is"
+          " evaluated on the shared path and raises ZeroDivisionError."
+      ),
+  )
   max_num_checkpoints_to_keep: int | None = Field(None, description="Maximum number of checkpoints to keep.")
   enable_single_replica_ckpt_restoring: bool = Field(
       False, description="One replica reads and broadcasts the checkpoint."

src/maxtext/eval/README.md

@@ -186,7 +186,7 @@ Example (Qwen3-30B-A3B, v6e-8):
 STEP=244
 MODEL=qwen3-30b-a3b
 HF_PATH=Qwen/Qwen3-30B-A3B
-CHECKPOINT=gs://<bucket>/run/checkpoints/actor/${STEP}/model_params
+CHECKPOINT=gs://<bucket>/run/checkpoints/actor/${STEP}/items
 OUTPUT=gs://<bucket>/eval/
 
 python -m maxtext.eval.runner.run \

src/maxtext/examples/rl_llama3_demo.ipynb

@@ -300,13 +300,13 @@
     "# Define the output directory for the Hugging Face checkpoint\n",
     "hf_output_directory = epath.Path(BASE_OUTPUT_DIRECTORY) / \"hf_checkpoint\"\n",
     "\n",
-    "# Find the latest MaxText checkpoint\n",
-    "checkpoint_dir = epath.Path(config.checkpoint_dir) / 'actor'\n",
+    "# Find the latest MaxText checkpoint. RL checkpoints the actor under its own subdirectory.\n",
+    "checkpoint_dir = epath.Path(config.checkpoint_dir) / \"actor\"\n",
     "step_dirs = [d.name for d in checkpoint_dir.iterdir() if d.name.isdigit() and d.is_dir()]\n",
     "if not step_dirs:\n",
     "    raise ValueError(f\"No checkpoint found in {checkpoint_dir}\")\n",
     "latest_step = max(step_dirs, key=int)\n",
-    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"model_params\"\n",
+    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"items\"\n",
     "\n",
     "print(f\"Converting MaxText checkpoint from: {maxtext_checkpoint_path}\")\n",
     "print(f\"Saving Hugging Face checkpoint to: {hf_output_directory}\")\n",

src/maxtext/examples/sft_llama3_demo_tpu.ipynb

@@ -341,7 +341,7 @@
     "if not step_dirs:\n",
     "    raise ValueError(f\"No checkpoint found in {checkpoint_dir}\")\n",
     "latest_step = max(step_dirs, key=int)\n",
-    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"model_params\"\n",
+    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"items\"\n",
     "\n",
     "print(f\"Converting MaxText checkpoint from: {maxtext_checkpoint_path}\")\n",
     "print(f\"Saving Hugging Face checkpoint to: {hf_output_directory}\")\n",

src/maxtext/inference/vllm_decode.py

@@ -110,6 +110,24 @@ def decode_with_vllm(config: Config) -> None:
       },
   }
 
+  # vllm_additional_config exists in the config but was never passed through, so keys tpu_inference
+  # reads straight off additional_config -- skip_quantization among them -- had no way to be set.
+  # Merged at the top level so it cannot clobber the maxtext_config and sharding blocks built above.
+  if config.vllm_additional_config:
+    for key, value in dict(config.vllm_additional_config).items():
+      if key in vllm_args["additional_config"]:
+        max_logging.log(f"Ignoring vllm_additional_config[{key!r}]: it would overwrite a generated block.")
+      else:
+        vllm_args["additional_config"][key] = value
+
+  # The worker rebuilds its own MaxTextConfig from inference/vllm.yml, which pins attention to
+  # vllm_rpa, and nothing carried the caller's choice across. Some decoder blocks reject that
+  # outright -- deepseek4 raises "DeepSeek4 decoder block currently only supports dot_product
+  # attention" -- and there was no way to say otherwise. Only forwarded when explicitly set, so
+  # models that are happy on vllm_rpa keep it.
+  if config.attention != "autoselected":
+    vllm_args["additional_config"]["maxtext_config"]["attention"] = config.attention
+
   if config.load_parameters_path:
     vllm_args["additional_config"]["maxtext_config"]["load_parameters_path"] = config.load_parameters_path
   else:

src/maxtext/input_pipeline/hf_data_processing.py

@@ -33,10 +33,15 @@
 
 
 def _get_pad_id(tokenizer):
-  if tokenizer.pad_token_id is not None:
+  """Returns the pad token id from the tokenizer, or -1 if not found."""
+  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
   else:
     pad_id = -1
   return pad_id

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -316,13 +316,25 @@ def apply_chat_template(example, tokenizer_model, data_column_name):
 
 def tokenization(example, hf_tokenizer, truncation, max_length, column_names):
   """Tokenize a HuggingFace dataset"""
+  is_hf = callable(hf_tokenizer)
   for column_name in column_names:
     if isinstance(example[column_name], list):
-      example[column_name] = [
-          hf_tokenizer(x, truncation=truncation, max_length=max_length)["input_ids"] for x in example[column_name]
-      ]
+      if is_hf:
+        example[column_name] = [
+            hf_tokenizer(x, truncation=truncation, max_length=max_length)["input_ids"] for x in example[column_name]
+        ]
+      else:
+        example[column_name] = [
+            hf_tokenizer.encode(x)[:max_length] if truncation else hf_tokenizer.encode(x) for x in example[column_name]
+        ]
     elif isinstance(example[column_name], str):
-      example[column_name] = hf_tokenizer(example[column_name], truncation=truncation, max_length=max_length)["input_ids"]
+      if is_hf:
+        example[column_name] = hf_tokenizer(example[column_name], truncation=truncation, max_length=max_length)[
+            "input_ids"
+        ]
+      else:
+        ids = hf_tokenizer.encode(example[column_name])
+        example[column_name] = ids[:max_length] if truncation else ids
   return example