How to save the parameters non-destructively?

[giganttheo]

How to save the parameters non-destructively?

[davisyoshida]

The easiest thing to use is just pickle, the reason it won't work with FLAX's msgpack utilities out of the box is that you have to manually handle non-default types (see here).

If you specifically need to use msgpack, you can replace the LoraWeight instances with tuples by doing something like this:

TAG = '_LoraWeight'

def _lora_to_tuple(lora_weight):
    return TAG, x.w, x.a, x.b, x.alpha

def _value_is_lora_tuple(x):
    return isinstance(x, tuple) and x[0] == TAG

def _lora_tuple_to_original(tup):
    _, w, a, b, alpha = tup
    return LoraWeight(w=w, a=a, b=b, alpha=alpha)

def pytree_loras_to_tuple(pytree):
    return jax.tree_map(
        lambda x: _lora_to_tuple(x) if isinstance(x, LoraWeight) else x,
        pytree,
        is_leaf=lambda x: isinstance(x, LoraWeight)
    )

def pytree_loras_to_tuple(pytree):
    return jax.tree_map(
        lambda x: _lora_to_tuple(x) if isinstance(x, LoraWeight) else x,
        pytree,
        is_leaf=lambda x: isinstance(x, LoraWeight)
    )

def pytree_tuples_to_loras(pytree):
    return jax.tree_map(
        lambda x: _lora_tuple_to_original(x) if _value_is_lora_tuple(x) else x,
        is_leaf=_value_is_lora_tuple
    )

Another option would be to entirely flatten the pytree and serialize the resulting list, then unflatten it after deserialization.

[giganttheo]

ok, thank you for the ideas, pickle is indeed quite easier to use than msgpack here