Provide a supported way to convert str keys from orbax checkpoints back to int keys

[pcmoritz]

It would be awesome if there was an officially supported function to convert keys that orbax sadly converted from int to str back to int. I saw #4317, but it is not sufficient for my use case, which needs a little more involved loading of checkpoints. See NovaSky-AI/SkyRL#461 for the use case.

The restore_int_paths function from https://github.com/google/flax/blob/main/flax/nnx/statelib.py works great, would it be possible to make it public (or to provide a public version of that function)?

Since the conversion from int to str is part of the public interface of flax.training.checkpoints, it would make sense to also have a public interface to convert back. Another even better option would be to do the conversion automatically in checkpoints.restore_checkpoint.

Thanks for your help :)

[vfdev-5]

@pcmoritz thanks for the report. Can you please provide a small reproducer to see the problem in detail?

About restore_int_paths function, it is true that it is not exposed in nnx directly, but it is a public function (as there is no underscore prefix).

[pcmoritz]

For sure, a simple repo is this:

import jax
from flax import nnx
from flax.training import checkpoints

class Linear(nnx.Module):
  def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
    self.w = nnx.Param(rngs.params.uniform((din, dout)))
    self.b = nnx.Param(jnp.zeros((dout,)))
  def __call__(self, x: jax.Array):
    return x @ self.w + self.b

class Layers(nnx.Module):
  def __init__(self, *, rngs: nnx.Rngs):
    self.layers = nnx.List([Linear(2, 5, rngs=rngs)])

model = Layers(rngs=nnx.Rngs(params=0))
graph, state = nnx.split(model)
checkpoints.save_checkpoint(
    target={"model": nnx.to_pure_dict(state)},
    ckpt_dir="/tmp/ckpt",
    step=0,
    prefix="checkpoint_",
    overwrite=True,
)

# Now we want to reload the checkpoint
restored_data = checkpoints.restore_checkpoint(ckpt_dir="/tmp/ckpt", target=None, prefix="checkpoint_")

# If you inspect restored_data, you already see it looks like
# {'model': {'layers': {'0': {'b': Array([0., 0., 0., 0., 0.], dtype=float32), 'w': Array([[0.8423141 , 0.18237865, 0.2271781 , 0.12072563, 0.19181347],
#        [0.722015  , 0.7654456 , 0.15254045, 0.9517063 , 0.02931046]],      dtype=float32)}}}}
# So the layer number like 0 is a string '0', which will cause problems below

# For example say we wanted to sum together the initial state with the loaded state
def compute(old: jax.Array, new: jax.Array):
    return old + new
jax.tree.map(compute, nnx.to_pure_dict(state), restored_data["model"])
# It gives the error ValueError: Dict key mismatch; expected keys: [0]; present keys: ['0'].

# Now we can work around that using
jax.tree.map(compute, nnx.to_pure_dict(state), nnx.statelib.restore_int_paths(restored_data["model"]))
# but ideally the nnx.statelib.restore_int_paths wouldn't be needed and checkpoints.restore_checkpoint would automatically already convert '0' to 0

Hope that makes it clearer, let me know if you have any questions about it :)

[pcmoritz]

I considered the nnx.statelib.restore_int_paths function private, because it is not part of the documentation https://flax.readthedocs.io/ (also the ty type checker complains it is not found, it seems to me nnx.statelib is supposed to be a private module). Especially given that it is not part of the docs, it seems likely that it could be removed in the future.

[vfdev-5]

I considered the nnx.statelib.restore_int_paths function private, because it is not part of the documentation https://flax.readthedocs.io/ (also the ty type checker complains it is not found, it seems to me nnx.statelib is supposed to be a private module). Especially given that it is not part of the docs, it seems likely that it could be removed in the future.

This makes sense. I'm not sure whether it was forgotten or left on purpose. I'll ask internally.

On the other hand, it is used in an integration test here:

flax/tests/nnx/integration_test.py

Lines 290 to 300 in fd118b1

	# checkpointer.save(ckpt_dir / 'state', state)
	checkpointer.save(ckpt_dir / 'pure_dict', pure_dict_state)
	# Restore as a pure dictionary.
	restored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')
	restored_pure_dict = nnx.statelib.restore_int_paths(restored_pure_dict)
	model = nnx.eval_shape(lambda: MLPs(4, rngs=nnx.Rngs(0)))
	nnx.update(model, restored_pure_dict)
	assert model(x).shape == (3, 4) # The model still works!

Personally, I would assume it public.

[cgarciae]

We should do a couple of things:

• Add docstrings to restore_int_paths.
• Expose it in the documentation.
• Expose statelib in nnx/__init__.py so ty is happy (seems to work fine in mypy and pyright though)
• Add an example of restoring a pure orbax checkpoint (no target pytree) to the Scale up on multiple devices guide. There is a test that uses this function that shows how to do this.

[pcmoritz]

Thanks @cgarciae, this is great! Have you considered just running restore_int_paths internally in restore_checkpoint before returning the result to the user? That would be the ideal solution for me, I would rather not want to deal with this conversion at all :)

If you are concerned about compatibility, you could add a flag for this behavior. But it seems very natural, since flax also controls how the keys look like (e.g. that they are ints for the layer indices), so assuming that information and doing the right thing seems like a great way to handle it :)

[vfdev-5]

@pcmoritz I have not spotted in your reproducer code that you are using flax.training.restore_checkpoint. I wonder why you can't set target argument? Something like this, so you wont need to convert str index to int index neither:

...
# Now we want to reload the checkpoint
new_model = Layers(rngs=nnx.Rngs(params=0))
target_state = nnx.state(model)
restored_state = checkpoints.restore_checkpoint(ckpt_dir="/tmp/ckpt", target=target_state, prefix="checkpoint_")
nnx.update(new_model, restored_state)

# check saved state vs restored state:
fstate1 = nnx.to_flat_state(state)
fstate2 = nnx.to_flat_state(restored_state)

import numpy as np

for (k1, v1), (k2, v2) in zip(fstate1, fstate2, strict=True):
    assert k1 == k2, (k1, k2)
    np.testing.assert_allclose(v1, v2, err_msg=f"{k1}, {k2}")

[pcmoritz]

I can't do that, because in the real example NovaSky-AI/SkyRL#461, I'm not loading the checkpoint into a given state but the checkpoint is a subset of the state (the state is the whole model, and I'm only loading one LoRA adapter into the state). In the toy example I posted above, this would correspond to compute doing slicing, like

def compute(old: jax.Array, new: jax.Array):
    return old[2, :] = new[:]
jax.tree.map(compute, nnx.to_pure_dict(state), restored_data["model"])

[vfdev-5]

@pcmoritz I curious whether in SkyRL process_save_weights you could store the checkpoint as

def process_save_weights(...):
    ...
    checkpoints.save_checkpoint(
            target={
                "lora_weights": adapter_lora_params,
                "optimizer_state": optimizer_params,
                "lora_config": self.models[model_id].lora_config.model_dump(),
            },

and then when loading in process_load_weights, we could construct the same target?

def process_load_weights(...):

  adapter_lora_params = extract_adapter_params(adapter_index, self.lora_params, self.non_lora_params)
  optimizer_params = extract_adapter_params(adapter_index, nnx.state(self.optimizer), self.non_lora_params)
  target = {
    "lora_weights": adapter_lora_params,
    "optimizer_state": optimizer_params,
    "lora_config": self.models[model_id].lora_config.model_dump(),    
  }
  restored_data = checkpoints.restore_checkpoint(ckpt_dir=checkpoint_dir, target=target, prefix="checkpoint_")

This may look a bit as an overkill but on the other hand we ensure to load an expected structure...

[pcmoritz]

@vfdev-5 Thanks a lot for suggesting that, this is a smart way to do it, but I tried it and it doesn't seem to work. I suspect the variables returned by extract_adapter_params are copies and not references to the original data -- it is defined as

def extract_adapter_state(
    adapter_index: int, lora_params: nnx.GraphState, non_lora_params: nnx.GraphState
) -> nnx.GraphState:
    "Helper function to extract the adapter parameters for a specific adapter index."
    flat_params = dict(nnx.to_flat_state(non_lora_params))

    def extract_state(path: tuple, p: jnp.ndarray):
        if path[-2].key not in {"lora_A", "lora_B"}:
            return p
        rank = flat_params[get_rank_path(path, path[-2].key)][adapter_index]
        assert p.ndim in {3, 4}, f"LoRA parameters must have 3 or 4 dimensions, got shape {p.shape}"
        if path[-2].key == "lora_A":
            return p[adapter_index, ..., :, :rank]
        if path[-2].key == "lora_B":
            return p[adapter_index, ..., :rank, :]

    return jax.tree.map_with_path(extract_state, lora_params)

Do you know if it is possible to return mutable references to the data? That might make it work :)

[pcmoritz]

I tried to call nnx.to_refs on the result btw, but then during the checkpoint restore I'm getting ValueError: Unsupported type: <class 'jax._src.core.Ref'> :)

[vfdev-5]

@pcmoritz I'm not sure to fully understand what does not work exactly in your code. The idea is to provide target structure for the output restored_data to be structured the same as target. Finally, we still will need to call insert_adapter_state to update the state of self.lora_params etc.

def process_load_weights(...):

    adapter_lora_params = extract_adapter_params(adapter_index, self.lora_params, self.non_lora_params)
    optimizer_params = extract_adapter_params(adapter_index, nnx.state(self.optimizer), self.non_lora_params)
    target = {
      "lora_weights": adapter_lora_params,
      "optimizer_state": optimizer_params,
      "lora_config": self.models[model_id].lora_config.model_dump(),    
    }
    restored_data = checkpoints.restore_checkpoint(ckpt_dir=checkpoint_dir, target=target, prefix="checkpoint_")
    # Update both LoRA weights and optimizer state
    insert_adapter_state(adapter_index, self.lora_params, self.non_lora_params, restored_data["lora_weights"])
    insert_adapter_state(
      adapter_index, nnx.state(self.optimizer), self.non_lora_params, restored_data["optimizer_state"]
    )
   # update lora_config

Here is the example of that: https://gist.github.com/vfdev-5/2e2ceb3cc27a976d9289bc269d2eb19b

Let me know what do you think.