Loading checkpoint for Dora PEFT has an issue

halo

I encountered an issue while attempting to load a checkpoint using the Dora PEFT . Here's the specific context:

I am using the following code snippet:


load_adapter_checkpoint(
    model=tfm,
    adapter_checkpoint_path="./checkpoints/run_20240716_163900_lyo4psz3",
    lora_rank=8,
    lora_target_modules="all",
    use_dora=True,
)

However, this results in an unexpected error or behavior. Below are the details.

Restoring adapter checkpoint from /home/fch/chitsaz/Stock_FM/checkpoints/run_20241113_014723/.
WARNING:absl:No registered CheckpointArgs found for handler type: <class 'paxml.checkpoints.FlaxCheckpointHandler'>
WARNING:absl:Configured `CheckpointManager` using deprecated legacy API. Please follow the instructions at https://orbax.readthedocs.io/en/latest/api_refactor.html to migrate by May 1st, 2024.
WARNING:absl:train_state_unpadded_shape_dtype_struct is not provided. We assume `train_state` is unpadded.
WARNING:absl:A possible mismatch (could be spurious) between the saved checkpoint structure and the current one has been detected (PyTreeDef(CustomNode(TrainState[()], [*, {'params': {'core_layer': {'freq_emb': {'emb_var': *}, 'horizon_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'input_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'stacked_transformer_layer': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}}}}}, [(CustomNode(NestedMap[('count',)], [*]), CustomNode(NestedMap[('count',)], [*]), CustomNode(NestedMap[('count', 'm', 'v')], [*, {'params': {'core_layer': {'freq_emb': {'emb_var': *}, 'horizon_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'input_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'stacked_transformer_layer': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}}}}}, {'params': {'core_layer': {'freq_emb': {'emb_var': *}, 'horizon_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'input_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'stacked_transformer_layer': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}}}}}]), CustomNode(NestedMap[('count',)], [*]), CustomNode(NestedMap[('count', 'ema')], [*, {'params': {'core_layer': {'freq_emb': {'emb_var': *}, 'horizon_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'input_ff_layer': {'hidden_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'output_layer': {'bias': {'b': *}, 'linear': {'w': *}}, 'residual_layer': {'bias': {'b': *}, 'linear': {'w': *}}}, 'stacked_transformer_layer': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'w': *}, 'query': {'b': *, 'w': *}, 'value': {'b': *, 'w': *}}}}}}}]))], ()])) vs PyTreeDef(CustomNode(TrainState[()], [*, {'x_layers_0': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_1': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_10': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_11': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_12': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_13': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_14': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_15': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_16': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_17': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_18': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_19': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_2': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_3': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_4': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_5': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_6': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_7': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_8': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}, 'x_layers_9': {'ffn_layer1': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'ffn_layer2': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'key': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'post': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'query': {'dora_m': *, 'lora_a': *, 'lora_b': *}, 'value': {'dora_m': *, 'lora_a': *, 'lora_b': *}}}, [], ()]))).
ERROR:absl:For checkpoint version > 1.0, we require users to provide
          `train_state_unpadded_shape_dtype_struct` during checkpoint
          saving/restoring, to avoid potential silent bugs when loading
          checkpoints to incompatible unpadded shapes of TrainState.

{
    "name": "ValueError",
    "message": "Ellipses lengths do not match.",
    "stack": "---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[13], line 1
----> 1 load_adapter_checkpoint(
      2     model=tfm,
      3     adapter_checkpoint_path=\"/home/fch/chitsaz/Stock_FM/checkpoints/run_20241113_014723/\",
      4     lora_rank=1,
      5     lora_target_modules=\"all\",
      6     use_dora=True,
      7 
      8 )

File ~/chitsaz/Stock_FM/timesfm/src/adapter/utils.py:171, in load_adapter_checkpoint(model, adapter_checkpoint_path, lora_rank, lora_target_modules, use_dora)
    162 adapter_train_state = checkpoints.restore_checkpoint(
    163     state_global_shapes=adapter_state_local_shapes,
    164     checkpoint_dir=adapter_checkpoint_path,
   (...)
    167     step=None,
    168 )
    170 # add adapter weights to the original weights
--> 171 _merge_adapter_weights(
    172     model=model,
    173     adapter_train_state=adapter_train_state,
    174     lora_target_modules=lora_target_modules,
    175     num_layers=model._model.stacked_transformer_params_tpl.num_layers,
    176     use_dora=use_dora,
    177 )
    179 # replace back with the original model layer
    180 if lora_target_modules in [\"all\", \"mlp\"]:

File ~/chitsaz/Stock_FM/timesfm/src/adapter/utils.py:232, in _merge_adapter_weights(model, adapter_train_state, lora_target_modules, num_layers, use_dora)
    228 lora_b = params[\"lora_b\"]
    230 w = linear[\"w\"]
--> 232 lora_delta = jnp.einsum(\"...dr,...nr->...dn\", lora_a, lora_b)
    233 lora_delta = jnp.reshape(lora_delta, w.shape)
    234 w_prime = w + lora_delta

File ~/anaconda3/envs/r_ch/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:3569, in einsum(subscripts, out, optimize, precision, preferred_element_type, _dot_general, *operands)
   3567   contract_path = _poly_einsum_handlers.get(ty, _default_poly_einsum_handler)
   3568 # using einsum_call=True here is an internal api for opt_einsum... sorry
-> 3569 operands, contractions = contract_path(
   3570       *operands, einsum_call=True, use_blas=True, optimize=optimize)
   3572 contractions = tuple((a, frozenset(b), c) for a, b, c, *_ in contractions)
   3574 einsum = jit(_einsum, static_argnums=(1, 2, 3, 4), inline=True)

File ~/anaconda3/envs/r_ch/lib/python3.10/site-packages/opt_einsum/contract.py:210, in contract_path(*operands, **kwargs)
    207 use_blas = kwargs.pop('use_blas', True)
    209 # Python side parsing
--> 210 input_subscripts, output_subscript, operands = parser.parse_einsum_input(operands)
    212 # Build a few useful list and sets
    213 input_list = input_subscripts.split(',')

File ~/anaconda3/envs/r_ch/lib/python3.10/site-packages/opt_einsum/parser.py:318, in parse_einsum_input(operands)
    315     longest = ellipse_count
    317 if ellipse_count < 0:
--> 318     raise ValueError(\"Ellipses lengths do not match.\")
    319 elif ellipse_count == 0:
    320     split_subscripts[num] = sub.replace('...', '')

ValueError: Ellipses lengths do not match."
}

google-research / timesfm

Loading checkpoint for Dora PEFT has an issue #185