freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
ee7462b558aaeb76922fb166b21feed375ee546d
timesfm/peft/usage.ipynb
T
tanmayshishodia 71d980267d add parameter efficient finetuning pipeline 
Why? this commit adds a generic finetuning pipeline with LoRA and DoRA support
2024-07-16 01:26:36 +00:00

352 lines
62 KiB
Plaintext
Raw Blame History

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load Base Model"
]
},
{
"cell_type": "code",
"execution_count": 1,
"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<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Constructing model weights.\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"WARNING:absl:No registered CheckpointArgs found for handler type: <class 'paxml.checkpoints.FlaxCheckpointHandler'>\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"
]
}
],
"source": [
"from timesfm import TimesFm, freq_map, data_loader\n",
"from adapter.utils import load_adapter_checkpoint\n",
"from tqdm import tqdm\n",
"import numpy as np\n",
"import pandas as pd\n",
"\n",
"\n",
"tfm = TimesFm(\n",
" context_len=512,\n",
" horizon_len=128,\n",
" input_patch_len=32,\n",
" output_patch_len=128,\n",
" num_layers=20,\n",
" model_dims=1280,\n",
" backend=\"cpu\",\n",
")\n",
"tfm.load_from_checkpoint(repo_id=\"google/timesfm-1.0-200m\")"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"DATA_DICT = {\n",
" \"ettm2\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm2.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"ettm1\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm1.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"etth2\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh2.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"etth1\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh1.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"elec\": {\n",
" \"boundaries\": [18413, 21044, 26304],\n",
" \"data_path\": \"../datasets/electricity/electricity.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"traffic\": {\n",
" \"boundaries\": [12280, 14036, 17544],\n",
" \"data_path\": \"../datasets/traffic/traffic.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"weather\": {\n",
" \"boundaries\": [36887, 42157, 52696],\n",
" \"data_path\": \"../datasets/weather/weather.csv\",\n",
" \"freq\": \"10min\",\n",
" },\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(tfm._train_state.mdl_vars)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load Adapter Checkpoint\n",
"\n",
"Specify the adapter checkpoint path, rank and the modules used to train the adapters and whether dora was employed or not."
]
},
{
"cell_type": "code",
"execution_count": 3,
"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: <class 'paxml.checkpoints.FlaxCheckpointHandler'>\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"
]
}
],
"source": [
"load_adapter_checkpoint(\n",
" model=tfm,\n",
" adapter_checkpoint_path=\"/home/ubuntu/tanmay/timesfm/adapter_finetuning/checkpoints/run_ignore_20240715_234116\",\n",
" lora_rank=1,\n",
" lora_target_modules=\"all\",\n",
" use_dora=True,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Test Performance"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"dataset = \"ettm1\"\n",
"data_path = DATA_DICT[dataset][\"data_path\"]\n",
"freq = DATA_DICT[dataset][\"freq\"]\n",
"int_freq = freq_map(freq)\n",
"boundaries = DATA_DICT[dataset][\"boundaries\"]\n",
"\n",
"data_df = pd.read_csv(open(data_path, \"r\"))\n",
"\n",
"ts_cols = [col for col in data_df.columns if col != \"date\"]\n",
"num_cov_cols = None\n",
"cat_cov_cols = None\n",
"\n",
"context_len = 512\n",
"pred_len = 96\n",
"\n",
"num_ts = len(ts_cols)\n",
"batch_size = 16\n",
"\n",
"dtl = data_loader.TimeSeriesdata(\n",
" data_path=data_path,\n",
" datetime_col=\"date\",\n",
" num_cov_cols=num_cov_cols,\n",
" cat_cov_cols=cat_cov_cols,\n",
" ts_cols=np.array(ts_cols),\n",
" train_range=[0, boundaries[0]],\n",
" val_range=[boundaries[0], boundaries[1]],\n",
" test_range=[boundaries[1], boundaries[2]],\n",
" hist_len=context_len,\n",
" pred_len=pred_len,\n",
" batch_size=num_ts,\n",
" freq=\"15min\",\n",
" normalize=True,\n",
" epoch_len=None,\n",
" holiday=False,\n",
" permute=True,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"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"
]
}
],
"source": [
"test_batches = dtl.tf_dataset(mode=\"test\", shift=pred_len)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"120it [01:20, 1.50it/s]\n"
]
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">MAE: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.32633697986602783</span>\n",
"</pre>\n"
],
"text/plain": [
"MAE: \u001b[1;36m0.32633697986602783\u001b[0m\n"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"mae_losses = []\n",
"for batch in tqdm(test_batches.as_numpy_iterator()):\n",
" past = batch[0]\n",
" actuals = batch[3]\n",
" _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n",
" forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n",
" mae_losses.append(np.abs(forecasts - actuals).mean())\n",
"\n",
"print(f\"MAE: {np.mean(mae_losses)}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "tanmay_tfm_env",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.14"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Reference in New Issue View Git Blame Copy Permalink