Strict checkpoint restore not working as expected

[AakashKumarNain]

I have a simple flow for my use case where I am using orbax with pure JAX to save and load models. Here is the project.

The Flow

1. Build a model (a PyTree) in JAX using a config file (e.g. size of layers, number of layers, etc.)
2. Train the model and save the checkpoint using orbax
3. Restore the checkpoint for inference or for continued training.

Problem

Say, I built a model that has 12 layers, trained it and saved the checkpoint. After a few days, I train a bigger model, say a model with 16 layers. If I forget to provide the correct checkpoint path for inference, orbax instead of complaining about the mismatch of sizes in the current model instance and the checkpoint, silently loads partial weights and never errors out. This is an undesirable behavior and can be a serious footgun for anyone who trains a lots of models and have lots of checkpoints for a use case.

Here is the code I am using to restore the checkpoint. Can anything be done about it? I want it to error out if there is a mismatch instead of silently restoring weights partially

def get_sharding_for_checkpoint(x, mesh):
    """Obtain shardings for the leaves of the pytree."""
    if hasattr(x, "ndim") and x.ndim == 0:
        return NamedSharding(mesh, P())
    if isinstance(x, jax.Array) and hasattr(x, "sharding"):
        from jax.sharding import SingleDeviceSharding
        # Ensure small optimizer leaves (e.g., Muon scalars/vectors) are replicated,
        # not left on a single device, to match model param shardings during train_step.
        if isinstance(x.sharding, SingleDeviceSharding):
            return NamedSharding(mesh, P())
        return x.sharding
    else:
        return NamedSharding(mesh, P())

def load_weights_from_checkpoint(path, params, mesh):
    print(f"Restoring params from: {path}")

    # 1. Create an abstract 'target' structure for validation
    abstract_params = jax.tree.map(ocp.utils.to_shape_dtype_struct, params)

    restore_args = jax.tree.map(
        lambda leaf: ocp.ArrayRestoreArgs(
        sharding=get_sharding_for_checkpoint(leaf, mesh), strict=True),
        params,
    )

    with ocp.PyTreeCheckpointer() as ckptr:
        restored = ckptr.restore(
            path,
            args=ocp.args.PyTreeRestore(
                # Use the abstract structure here to trigger validation
                item=abstract_params, 
                restore_args=restore_args,
            ),
        )
    return restored

[cpgaffney1]

The strict argument here is on a per-parameter basis - it enforces strict dtype / shape checking for each parameter, but does not do anything about tree structures.

restore_args = jax.tree.map(
        lambda leaf: ocp.ArrayRestoreArgs(
        sharding=get_sharding_for_checkpoint(leaf, mesh), strict=True),
        params,
    )

Tree structure checks are enforced separately, by the usage of item and/or restore_args. If either of those doesn't match the checkpoint tree structure, it will raise an error, unless you've enabled PyTreeRestore.partial_restore = True. Based on that your code snippet seems like it should work....

Is there perhaps another place in the code in which restore is called without passing restore_args / abstract_params? That would restore the checkpoint unconditionally, since no abstract state was provided as a check.

[AakashKumarNain]

If either of those doesn't match the checkpoint tree structure, it will raise an error

That's why I was surprised. I had to write extra functions to validate the tree structure, but ideally it should be baked in the checkpoint restorer.

unless you've enabled PyTreeRestore.partial_restore = True

I haven't enabled this

Is there perhaps another place in the code in which restore is called without passing restore_args / abstract_params?

This is the only code for checkpoint restoration

[cpgaffney1]

Locally I do see the error being raised as expected with this toy model (at head):

class ReproIssue2972Test(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.directory = epath.Path(
        self.create_tempdir(name='checkpointing_test').full_path
    )

  def test_mismatched_layer_count_raises_error(self):
    small_model = {
        f'layer_{i}': {
            'w': np.ones((4, 4)),
            'b': np.ones((4,)),
        }
        for i in range(12)
    }
    ckpt_dir = self.directory / 'ckpt'
    with ocp.PyTreeCheckpointer() as ckptr:
      ckptr.save(ckpt_dir, args=ocp.args.PyTreeSave(small_model))

    large_model = {
        f'layer_{i}': {
            'w': np.ones((4, 4)),
            'b': np.ones((4,)),
        }
        for i in range(16)
    }

    mesh = jax.sharding.Mesh(jax.devices(), ('devices',))
    with self.assertRaises(ValueError):
      load_weights_from_checkpoint(ckpt_dir, large_model, mesh)

There's a slight version mismatch but your project requires 0.11.28, and I don't think this logic was touched recently.

[lyglst]

Thank you for reporting this issue, this is a bit surprising as strict=True should be sufficient for catching this type of mismatch. I tested this locally and it does throw an exception as expected:

  ckpt_tree = {'params': {'w': jnp.zeros((4, 4))}}
  model_tree = {'params': {'w': jnp.zeros((8, 8))}}

  path = '/tmp/orbax_repro_shape'
  if os.path.exists(path):
    shutil.rmtree(path)

  ckptr = ocp.PyTreeCheckpointer()
  ckptr.save(path, item=ckpt_tree)

  print('Saved checkpoint')

  try:
    restored = ckptr.restore(path, item=model_tree)
    print('Restore succeeded! This is the bug. Shape was ignored.')
    print(f"Original model_tree shape: {model_tree['params']['w'].shape}")
    print(f"Restored tree shape: {restored['params']['w'].shape}")
  except Exception as e:
    print(f'Restore failed with error: {e}')

Could you provide a bit more context?

[lyglst]

Another possibility is - if the keys of your pytree match, but the shape or dtype of your Array does not, Orbax might restore successfully without throwing an exception. e.g.:

    # SAVED model
    save_model = {
        'params': {
            'layer1': jnp.ones((4, 4), dtype=jnp.float32),
            'layer2': jnp.ones((2, 2), dtype=jnp.int32),
        }
    }

    # TARGET model (mismatched sizes and types)
    target_model = {
        'params': {
            'layer1': jnp.zeros((8, 8), dtype=jnp.float32),  # Shape mismatch
            'layer2': jnp.zeros((2, 2), dtype=jnp.float32),  # Dtype mismatch
        }
    }
    abstract_target_model = jax.tree.map(
        ocp.utils.to_shape_dtype_struct, target_model
    )
    restore_args = jax.tree.map(
        lambda leaf: ocp.ArrayRestoreArgs(
            strict=True
        ),
        abstract_target_model,
    )

    # Should fail if we have strict structure/shape check?
    with self.assertRaises(ValueError):
      checkpointer.restore(
          ckpt_dir,
          args=ocp.args.PyTreeRestore(
              # Use the abstract structure here to trigger validation
              item=abstract_target_model,
              restore_args=restore_args,
          ),
      )

I'm working on a fix for this scenario, not sure if it's the culprit here.

[AakashKumarNain]

if the keys of your pytree match, but the shape or dtype of your Array does not,

This is the case, but I do not think that it's a good idea to restore. We hardly change names of the layers in an architecture, but we change the size of those pytrees frequently. IMO strict should mean that the name of the pytree along with shape and dtype should match for restoration. If not, it should throw an error. A partial restore is a foot gun

[cpgaffney1]

Partial tree restoration is useful in many cases, particularly e.g. when users wish to restore params, but not optimizer states for evals or inference. This is controlled by the partial_restore option which defaults to false.

The strict option does normally raise errors with mismatched array shapes / dtypes, but you hit a bit of an edge case since these properties were not specified in ArrayRestoreArgs, which is where the check happens. The whole item / restore_args distinction is definitely a footgun and problematic design which is addressed by the upcoming V1 API, but unfortunately it is not neatly solvable in the current API.

[AakashKumarNain]

Partial tree restoration is useful in many cases, particularly e.g. when users wish to restore params, but not optimizer states for evals or inference. This is controlled by the partial_restore option which defaults to false.

Of course, but I was talking about the "strict" context.

which is addressed by the upcoming V1 API,

Any timelines for that?

but unfortunately it is not neatly solvable in the current API.

For now, I have written a few functions on my side to deal with it.

[cpgaffney1]

Any timelines for that?

We are in the process of migrating a few major Google-owned codebases as a final step before official release, but I don't expect any substantive changes. You're welcome to take a look https://orbax.readthedocs.io/en/latest/guides/checkpoint/v1/orbax_checkpoint_101.html.

[AakashKumarNain]

Oh that's a big change! I think apart from the API, the documentation itself requires a major overhaul. A few of us are ready to help. Let us know if we have a process which we can follow for the contribution for the new API

[cpgaffney1]

Yes, we are focusing our efforts on the V1 guides and V1 API reference. Happy to hear any proposals you may have for improvements, either to the docs or the API itself. If there are sizable-enough proposals from external contributors, we can prioritize setting up the external to internal mirroring, which should not be too difficult - we just have not had time to do.

[AakashKumarNain]

Happy to hear any proposals you may have for improvements, either to the docs or the API itself.

I will come up with a clean list and put in this thread latest by the weekend. Thanks

[lyglst]

The original issue should be resolved in #2976 - please help verify and close the bug if it works!

[AakashKumarNain]

Will test tomorrow, and update it if works. Thanks @lyglst