diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py
index ab6af89..14a2fdd 100644
--- a/recml/core/utils/keras_utils.py
+++ b/recml/core/utils/keras_utils.py
@@ -309,7 +309,7 @@ class KerasOrbaxCheckpointManagerV3(ocp.CheckpointManager):
   def __init__(
       self,
       checkpoint_dir: str,
-      max_to_keep: int = 5,
+      max_to_keep: int | None = 5,
       save_interval_epochs: int = 1,
       choose_store_cell: bool = True,
   ):
@@ -317,7 +317,8 @@ def __init__(
 
     Args:
       checkpoint_dir: The directory to save checkpoints to.
-      max_to_keep: The maximum number of checkpoints to keep.
+      max_to_keep: The maximum number of checkpoints to keep, or None to keep
+        every checkpoint.
       save_interval_epochs: The interval (in epochs) to save checkpoints.
       choose_store_cell: Whether to dynamically select the CNS2 store cell.
     """
@@ -561,7 +562,7 @@ def _is_v3_checkpoint_path(checkpoint_path: str) -> bool:
     version_checkpointer.close()
 
 
-def _detect_checkpoint_version(checkpoint_path: str) -> CheckpointVersion:
+def detect_checkpoint_version(checkpoint_path: str) -> CheckpointVersion:
   """Detects the version of the checkpoint at the given path.
 
   Detection order and discriminator criteria:
@@ -592,7 +593,7 @@ def _detect_checkpoint_version(checkpoint_path: str) -> CheckpointVersion:
 def is_v3_checkpoint(checkpoint_dir: str, epoch: int | None = None) -> bool:
   """Checks if a checkpoint is in V3 format."""
   checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch)
-  return _detect_checkpoint_version(checkpoint_path) == CheckpointVersion.V3
+  return detect_checkpoint_version(checkpoint_path) == CheckpointVersion.V3
 
 
 def _validate_v3_checkpoint(
@@ -823,9 +824,10 @@ def _prepare_v3_restore(
     checkpointer.close()
 
   # Load index path metadata for non-trainable and optimizer variables if
-  # needed.
+  # present in `abstract_state` (partial restore omits these keys, whereas full
+  # restore includes them even when the dict is empty `{}`).
   non_trainable_paths = None
-  if model is not None and abstract_state.get(NON_TRAINABLE_VARIABLES_KEY):
+  if model is not None and NON_TRAINABLE_VARIABLES_KEY in abstract_state:
     nt_path = epath.Path(checkpoint_path) / NON_TRAINABLE_PATHS_KEY
     checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
     try:
@@ -839,7 +841,7 @@ def _prepare_v3_restore(
   if (
       model is not None
       and restore_optimizer_vars
-      and abstract_state.get(OPTIMIZER_VARIABLES_KEY)
+      and OPTIMIZER_VARIABLES_KEY in abstract_state
   ):
     opt_path = epath.Path(checkpoint_path) / OPTIMIZER_PATHS_KEY
     checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
@@ -858,7 +860,9 @@ def _prepare_v3_restore(
   state_transforms = {}
 
   for key in keys:
-    if not abstract_state.get(key):
+    # Keep empty dicts (e.g. `{}` when a model has no non-trainable variables)
+    # in `filtered_abstract_state` so the restored PyTree structure matches.
+    if key not in abstract_state:
       continue
 
     filtered_abstract_state[key] = {}
@@ -1132,7 +1136,7 @@ def restore_keras_checkpoint(
 
   checkpoint_path, epoch = resolve_orbax_checkpoint_path(checkpoint_dir, epoch)
 
-  version = _detect_checkpoint_version(checkpoint_path)
+  version = detect_checkpoint_version(checkpoint_path)
   if version == CheckpointVersion.V1:
     raise ValueError(
         f"The checkpoint in {checkpoint_dir} is in V1 format (list-based)"
@@ -1298,7 +1302,7 @@ def restore_partial_checkpoint(
   """
   checkpoint_path, _ = resolve_orbax_checkpoint_path(checkpoint_dir, epoch)
 
-  version = _detect_checkpoint_version(checkpoint_path)
+  version = detect_checkpoint_version(checkpoint_path)
   if version != CheckpointVersion.V3:
     raise ValueError(
         "restore_partial_checkpoint only supports V3 (dictionary-based)"
diff --git a/recml/core/utils/keras_utils_test.py b/recml/core/utils/keras_utils_test.py
index d130f55..cc32429 100644
--- a/recml/core/utils/keras_utils_test.py
+++ b/recml/core/utils/keras_utils_test.py
@@ -276,6 +276,38 @@ def test_keras_orbax_checkpointer_v3(
     # Ensures predictions are identical.
     np.testing.assert_allclose(preds, preds_after_restoration)
 
+  def test_keras_orbax_checkpointer_v3_no_non_trainable_variables(self):
+    model = keras.Sequential(
+        [keras.layers.Dense(5, name="dense")], name="model"
+    )
+    model.build((2, 10))
+    model.compile(optimizer=keras.optimizers.Adam(0.01), loss="mse")
+    model.optimizer.build(model.trainable_variables)
+    self.assertEmpty(model.non_trainable_variables)
+
+    checkpoint_dir = self.create_tempdir().full_path
+    checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3(
+        checkpoint_dir, max_to_keep=None
+    )
+    checkpoint_manager.save_model_variables(model, 0)
+    checkpoint_manager.wait_until_finished()
+
+    new_model = keras.Sequential(
+        [keras.layers.Dense(5, name="dense")], name="model"
+    )
+    new_model.build((2, 10))
+    new_model.compile(optimizer=keras.optimizers.Adam(0.01), loss="mse")
+    new_model.optimizer.build(new_model.trainable_variables)
+
+    checkpoint_manager.restore_model_variables(new_model, 0)
+    checkpoint_manager.close()
+
+    for expected, observed in zip(
+        [v.value for v in model.trainable_variables],
+        [v.value for v in new_model.trainable_variables],
+    ):
+      np.testing.assert_allclose(observed, expected)
+
   def test_restore_keras_checkpoint(self):
     dummy_inputs = _create_dummy_inputs()
     bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs))