From 845661dfa32036d62960a13226c255a925ee5a4e Mon Sep 17 00:00:00 2001 From: tanmayshishodia Date: Tue, 16 Jul 2024 18:24:50 +0000 Subject: [PATCH] update usage tutorial --- peft/usage.ipynb | 170 +++-------------------------------------------- 1 file changed, 11 insertions(+), 159 deletions(-) diff --git a/peft/usage.ipynb b/peft/usage.ipynb index 58ae1c0..ab5177a 100644 --- a/peft/usage.ipynb +++ b/peft/usage.ipynb @@ -9,75 +9,9 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2024-07-16 00:36:17.861915: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n", - "2024-07-16 00:36:25.276917: W external/xla/xla/service/gpu/nvptx_compiler.cc:718] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.5.40). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.\n" - ] - }, - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "935a82ff713149179c103451e4837c27", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Fetching 5 files: 0%| | 0/5 [00:00\n", - "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.\n", - "WARNING:absl:train_state_unpadded_shape_dtype_struct is not provided. We assume `train_state` is unpadded.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Constructed model weights in 2.84 seconds.\n", - "Restoring checkpoint from /home/ubuntu/.cache/huggingface/hub/models--google--timesfm-1.0-200m/snapshots/8775f7531211ac864b739fe776b0b255c277e2be/checkpoints.\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "ERROR:absl:For checkpoint version > 1.0, we require users to provide\n", - " `train_state_unpadded_shape_dtype_struct` during checkpoint\n", - " saving/restoring, to avoid potential silent bugs when loading\n", - " checkpoints to incompatible unpadded shapes of TrainState.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Restored checkpoint in 0.86 seconds.\n", - "Jitting decoding.\n", - "Jitted decoding in 16.44 seconds.\n" - ] - } - ], + "outputs": [], "source": [ "from timesfm import TimesFm, freq_map, data_loader\n", "from adapter.utils import load_adapter_checkpoint\n", @@ -100,7 +34,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -143,15 +77,6 @@ "}" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(tfm._train_state.mdl_vars)" - ] - }, { "cell_type": "markdown", "metadata": {}, @@ -163,44 +88,13 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Restoring adapter checkpoint from /home/ubuntu/tanmay/timesfm/adapter_finetuning/checkpoints/run_ignore_20240715_234116.\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "WARNING:absl:No registered CheckpointArgs found for handler type: \n", - "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.\n", - "WARNING:absl:train_state_unpadded_shape_dtype_struct is not provided. We assume `train_state` is unpadded.\n", - "WARNING:absl:A possible mismatch (could be spurious) between the saved checkpoint structure and the current one has been detected (PyTreeDef(CustomNode(TrainState[()], [*, {'params': {'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': *}}}}, [(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': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_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': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_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': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_1': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_10': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_11': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_12': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_13': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_14': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_15': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_16': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_17': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_18': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_19': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_2': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_3': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_4': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_5': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_6': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_7': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_8': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}}, 'x_layers_9': {'ff_layer': {'ffn_layer1': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'ffn_layer2': {'bias': {'b': *}, 'linear': {'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}}, 'layer_norm': {'bias': *, 'scale': *}}, 'layer_norm': {'scale': *}, 'self_attention': {'key': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'per_dim_scale': {'per_dim_scale': *}, 'post': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'query': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_b': *, 'w': *}, 'value': {'b': *, 'dora_m': *, 'lora_a': *, 'lora_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': *}}}, [], ()]))).\n", - "ERROR:absl:For checkpoint version > 1.0, we require users to provide\n", - " `train_state_unpadded_shape_dtype_struct` during checkpoint\n", - " saving/restoring, to avoid potential silent bugs when loading\n", - " checkpoints to incompatible unpadded shapes of TrainState.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Restored adapter checkpoint in 8.77 seconds.\n", - "Jitting decoding.\n", - "Jitted decoding in 14.77 seconds.\n" - ] - } - ], + "outputs": [], "source": [ "load_adapter_checkpoint(\n", " model=tfm,\n", - " adapter_checkpoint_path=\"/home/ubuntu/tanmay/timesfm/adapter_finetuning/checkpoints/run_ignore_20240715_234116\",\n", + " adapter_checkpoint_path=\"./checkpoints/run_20240716_163900_lyo4psz3\",\n", " lora_rank=1,\n", " lora_target_modules=\"all\",\n", " use_dora=True,\n", @@ -216,7 +110,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -260,53 +154,18 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2024-07-16 00:37:10.346547: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n", - "2024-07-16 00:37:10.346626: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory\n", - "2024-07-16 00:37:10.346677: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory\n", - "2024-07-16 00:37:10.346724: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcufft.so.10'; dlerror: libcufft.so.10: cannot open shared object file: No such file or directory\n", - "2024-07-16 00:37:10.350324: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory\n", - "2024-07-16 00:37:10.350358: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1850] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n", - "Skipping registering GPU devices...\n" - ] - } - ], + "outputs": [], "source": [ "test_batches = dtl.tf_dataset(mode=\"test\", shift=pred_len)" ] }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "120it [01:20, 1.50it/s]\n" - ] - }, - { - "data": { - "text/html": [ - "
MAE: 0.32633697986602783\n",
-       "
\n" - ], - "text/plain": [ - "MAE: \u001b[1;36m0.32633697986602783\u001b[0m\n" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], + "outputs": [], "source": [ "mae_losses = []\n", "for batch in tqdm(test_batches.as_numpy_iterator()):\n", @@ -318,13 +177,6 @@ "\n", "print(f\"MAE: {np.mean(mae_losses)}\")" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": {