add parameter efficient finetuning pipeline

Why? this commit adds a generic finetuning pipeline with LoRA and DoRA support

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+checkpoints/
+wandb/
+datasets/
+results/

environment.yml

@@ -1,4 +1,4 @@
-name: tfm_env
+name: ok_tfm_env
 
 channels:
   - conda-forge
@@ -16,3 +16,5 @@ dependencies:
     - jax[cuda12]==0.4.26
     - einshape
     - scikit-learn
+    - typer
+    - wandb

environment_cpu.yml

@@ -16,3 +16,5 @@ dependencies:
     - jax[cpu]==0.4.26
     - einshape
     - scikit-learn
+    - typer
+    - wandb

peft/finetune.py

@@ -0,0 +1,410 @@
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Finetune pipeline.
+"""
+import gc
+import logging
+import warnings
+from datetime import datetime
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+import typer
+import wandb
+from jax import numpy as jnp
+from paxml import checkpoint_types, checkpoints, learners, tasks_lib, trainer_lib
+from praxis import optimizers, pax_fiddle, py_utils, schedules
+from rich import print
+from tqdm import tqdm
+from typing_extensions import Annotated
+
+from adapter.utils import get_adapter_params, load_adapter_layer
+from timesfm import TimesFm, data_loader, patched_decoder
+
+NestedMap = py_utils.NestedMap
+
+
+warnings.filterwarnings("ignore")
+cmdstanpy_logger = logging.getLogger("cmdstanpy")
+absl_logger = logging.getLogger("absl")
+cmdstanpy_logger.disabled = True
+absl_logger.disabled = True
+
+"""
+TimesFM model config. These are fixed since pre-training was done 
+with this configuration.
+"""
+INPUT_PATCH_LEN = 32
+OUTPUT_PATCH_LEN = 128
+NUM_LAYERS = 20
+MODEL_DIMS = 1280
+
+QUANTILES = list(np.arange(1, 10) / 10.0)
+EPS = 1e-7
+RANDOM_SEED = 1234
+
+
+def get_forecasts(model, past: np.ndarray, freq: int) -> np.ndarray:
+    """Get forecasts."""
+    lfreq = [freq] * past.shape[0]
+    _, out = model.forecast(list(past), lfreq)
+    out = out[:, :, 5]
+    return out
+
+
+def finetune(
+    *,
+    checkpoint_path: Annotated[
+        str, typer.Option(help="The path to the model checkpoint.")
+    ] = None,
+    model_name: Annotated[
+        str, typer.Option(help="Specify the name of the huggingface model.")
+    ] = "google/timesfm-1.0-200m",
+    datetime_col: Annotated[str, typer.Option(help="Column having datetime.")] = "ds",
+    ts_cols: Annotated[
+        list[str], typer.Option(help="Columns of time-series features.")
+    ] = [],
+    normalize: Annotated[
+        bool, typer.Option(help="Normalize data for eval or not")
+    ] = True,
+    context_len: Annotated[int, typer.Option(help="Length of the context window")],
+    horizon_len: Annotated[int, typer.Option(help="Prediction length.")],
+    freq: Annotated[
+        str,
+        typer.Option(
+            ...,
+            help="Frequency Map Str",
+        ),
+    ],
+    data_path: Annotated[str, typer.Option(help="Path to dataset csv")],
+    boundaries: Annotated[
+        Tuple[int, int, int],
+        typer.Option(
+            help="boundaries of dataset to train, val, test",
+        ),
+    ] = (0, 0, 0),
+    backend: Annotated[str, typer.Option(help="Backend device: cpu, gpu, tpu")],
+    batch_size: Annotated[
+        int, typer.Option(help="Batch size for the randomly sampled batch")
+    ] = 16,
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs")],
+    learning_rate: Annotated[float, typer.Option(help="adam optimizer learning rate")],
+    adam_epsilon: Annotated[float, typer.Option(help="adam optimizer epsilon")],
+    adam_clip_threshold: Annotated[
+        float, typer.Option(help="adam optimizer clip threshold")
+    ],
+    cos_initial_decay_value: Annotated[
+        float, typer.Option(help="cosine initial decay value")
+    ],
+    cos_final_decay_value: Annotated[
+        float, typer.Option(help="cosine final decay value")
+    ],
+    cos_decay_steps: Annotated[int, typer.Option(help="Number of cosine decay steps")],
+    ema_decay: Annotated[float, typer.Option(help="Exponential moving average decay")],
+    early_stop_patience: Annotated[
+        int, typer.Option(..., help="Early stopping patience")
+    ] = 5,
+    use_lora: Annotated[
+        bool,
+        typer.Option(
+            help="Train low rank adapters. Freeze all other params in model",
+        ),
+    ] = False,
+    lora_rank: Annotated[
+        int,
+        typer.Option(
+            help="LoRA Rank",
+        ),
+    ] = 8,
+    lora_target_modules: Annotated[
+        str,
+        typer.Option(help="LoRA target modules. Allowed values: [all, attention, mlp]"),
+    ] = "all",
+    use_dora: Annotated[
+        bool,
+        typer.Option(
+            help="Apply DoRA strategy along with LoRA.",
+        ),
+    ] = False,
+    use_linear_probing: Annotated[
+        bool,
+        typer.Option(
+            help="Linear Probing. Train only input/output and embedding params. Freeze params in self attention modules.",
+        ),
+    ] = False,
+    checkpoint_dir: Annotated[
+        str, typer.Option(help="Checkpoint directory")
+    ] = "./checkpoints",
+    wandb_project: Annotated[
+        str, typer.Option(help="Weights & Biases project name")
+    ] = "google_timesfm_finetune",
+) -> None:
+    key = jax.random.PRNGKey(seed=RANDOM_SEED)
+    wandb.init(project=wandb_project, config=locals())
+
+    data_df = pd.read_csv(open(data_path, "r"))
+
+    if boundaries == (0, 0, 0):
+        # Default boundaries: train 60%, val 20%, test 20%
+        boundaries = [
+            int(len(data_df) * 0.6),
+            int(len(data_df) * 0.8),
+            len(data_df) - 1,
+        ]
+
+    ts_cols = [col for col in data_df.columns if col != datetime_col]
+
+    dtl = data_loader.TimeSeriesdata(
+        data_path=data_path,
+        datetime_col=datetime_col,
+        num_cov_cols=None,
+        cat_cov_cols=None,
+        ts_cols=np.array(ts_cols),
+        train_range=[0, boundaries[0]],
+        val_range=[boundaries[0], boundaries[1]],
+        test_range=[boundaries[1], boundaries[2]],
+        hist_len=context_len,
+        pred_len=horizon_len,
+        batch_size=batch_size,
+        freq=freq,
+        normalize=normalize,
+        epoch_len=None,
+        holiday=False,
+        permute=False,
+    )
+
+    train_batches = dtl.tf_dataset(mode="train", shift=1).batch(batch_size)
+    val_batches = dtl.tf_dataset(mode="val", shift=horizon_len)
+
+    for tbatch in tqdm(train_batches.as_numpy_iterator()):
+        pass
+
+    tfm = TimesFm(
+        context_len=context_len,
+        horizon_len=horizon_len,
+        input_patch_len=INPUT_PATCH_LEN,
+        output_patch_len=OUTPUT_PATCH_LEN,
+        num_layers=NUM_LAYERS,
+        model_dims=MODEL_DIMS,
+        backend=backend,
+        per_core_batch_size=batch_size,
+        quantiles=QUANTILES,
+    )
+
+    if checkpoint_path:
+        tfm.load_from_checkpoint(
+            checkpoint_path=checkpoint_path,
+            checkpoint_type=checkpoints.CheckpointType.FLAX,
+        )
+    else:
+        tfm.load_from_checkpoint(
+            repo_id=model_name,
+            checkpoint_type=checkpoints.CheckpointType.FLAX,
+        )
+
+    model = pax_fiddle.Config(
+        patched_decoder.PatchedDecoderFinetuneModel,
+        name="patched_decoder_finetune",
+        core_layer_tpl=tfm.model_p,
+    )
+
+    if use_lora:
+        load_adapter_layer(
+            mdl_vars=tfm._train_state.mdl_vars,
+            model=model.core_layer_tpl,
+            lora_rank=lora_rank,
+            lora_target_modules=lora_target_modules,
+            use_dora=use_dora,
+        )
+
+    @pax_fiddle.auto_config
+    def build_learner() -> learners.Learner:
+        bprop_variable_inclusion = None
+        bprop_variable_exclusion = None
+        if use_lora:
+            bprop_variable_inclusion = [r"^.*lora.*$"]
+            if use_dora:
+                bprop_variable_inclusion.append(r"^.*dora.*$")
+        elif use_linear_probing:
+            bprop_variable_exclusion = [".*/stacked_transformer_layer/.*"]
+
+        return pax_fiddle.Config(
+            learners.Learner,
+            name="learner",
+            loss_name="avg_qloss",
+            optimizer=optimizers.Adam(
+                epsilon=adam_epsilon,
+                clip_threshold=adam_clip_threshold,
+                learning_rate=learning_rate,
+                lr_schedule=pax_fiddle.Config(
+                    schedules.Cosine,
+                    initial_value=cos_initial_decay_value,
+                    final_value=cos_final_decay_value,
+                    total_steps=cos_decay_steps,
+                ),
+                ema_decay=ema_decay,
+            ),
+            bprop_variable_exclusion=bprop_variable_exclusion,
+            bprop_variable_inclusion=bprop_variable_inclusion,
+        )
+
+    task_p = tasks_lib.SingleTask(
+        name="ts-learn",
+        model=model,
+        train=tasks_lib.SingleTask.Train(
+            learner=build_learner(),
+        ),
+    )
+
+    task_p.model.ici_mesh_shape = [1, 1, 1]
+    task_p.model.mesh_axis_names = ["replica", "data", "mdl"]
+
+    DEVICES = np.array(jax.devices()).reshape([1, 1, 1])
+    jax.sharding.Mesh(DEVICES, ["replica", "data", "mdl"])
+
+    num_devices = jax.local_device_count()
+    print(f"num_devices: {num_devices}")
+    print(f"device kind: {jax.local_devices()[0].device_kind}")
+
+    jax_task = task_p
+    key, init_key = jax.random.split(key)
+
+    def process_train_batch(batch):
+        past_ts = batch[0].reshape(batch_size * len(ts_cols), -1)
+        actual_ts = batch[3].reshape(batch_size * len(ts_cols), -1)
+        return NestedMap(input_ts=past_ts, actual_ts=actual_ts)
+
+    def process_eval_batch(batch):
+        past_ts = batch[0]
+        actual_ts = batch[3]
+        return NestedMap(input_ts=past_ts, actual_ts=actual_ts)
+
+    jax_model_states, _ = trainer_lib.initialize_model_state(
+        jax_task,
+        init_key,
+        process_train_batch(tbatch),
+        checkpoint_type=checkpoint_types.CheckpointType.GDA,
+    )
+    jax_model_states.mdl_vars["params"]["core_layer"] = tfm._train_state.mdl_vars[
+        "params"
+    ]
+    gc.collect()
+
+    jax_task = task_p
+
+    def train_step(states, prng_key, inputs):
+        return trainer_lib.train_step_single_learner(jax_task, states, prng_key, inputs)
+
+    def eval_step(states, prng_key, inputs):
+        states = states.to_eval_state()
+        return trainer_lib.eval_step_single_learner(jax_task, states, prng_key, inputs)
+
+    key, train_key, eval_key = jax.random.split(key, 3)
+    train_prng_seed = jax.random.split(train_key, num=jax.local_device_count())
+    eval_prng_seed = jax.random.split(eval_key, num=jax.local_device_count())
+
+    p_train_step = jax.pmap(train_step, axis_name="batch")
+    p_eval_step = jax.pmap(eval_step, axis_name="batch")
+
+    replicated_jax_states = trainer_lib.replicate_model_state(jax_model_states)
+
+    def reshape_batch_for_pmap(batch, num_devices):
+        def _reshape(input_tensor):
+            bsize = input_tensor.shape[0]
+            residual_shape = list(input_tensor.shape[1:])
+            nbsize = bsize // num_devices
+            return jnp.reshape(input_tensor, [num_devices, nbsize] + residual_shape)
+
+        return jax.tree.map(_reshape, batch)
+
+    patience = 0
+    best_eval_loss = 1e7
+    checkpoint_dir = (
+        f"{checkpoint_dir}/run_ignore_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+    )
+    for epoch in range(num_epochs):
+        print(f"Epoch: {epoch + 1}")
+        train_its = train_batches.as_numpy_iterator()
+        train_losses = []
+        for batch in tqdm(train_its):
+            if patience >= early_stop_patience:
+                print("Early stopping.")
+                break
+            tbatch = process_train_batch(batch)
+            tbatch = reshape_batch_for_pmap(tbatch, num_devices)
+            replicated_jax_states, step_fun_out = p_train_step(
+                replicated_jax_states, train_prng_seed, tbatch
+            )
+            train_losses.append(step_fun_out.loss[0])
+            wandb.log({"train_step_loss": step_fun_out.loss[0]})
+
+        avg_train_loss = np.mean(train_losses)
+
+        print("Starting eval.")
+        val_its = val_batches.as_numpy_iterator()
+        eval_losses = []
+        for ev_batch in tqdm(val_its):
+            ebatch = process_eval_batch(ev_batch)
+            ebatch = reshape_batch_for_pmap(ebatch, num_devices)
+            _, step_fun_out = p_eval_step(replicated_jax_states, eval_prng_seed, ebatch)
+            eval_losses.append(step_fun_out.loss[0])
+            wandb.log({"eval_step_loss": step_fun_out.loss[0]})
+
+        avg_eval_loss = np.mean(eval_losses)
+
+        print(f"Train Loss: {avg_train_loss}, Val Loss: {avg_eval_loss}")
+
+        wandb.log(
+            {
+                "epoch": epoch + 1,
+                "avg_train_loss": avg_train_loss,
+                "avg_val_loss": avg_eval_loss,
+            }
+        )
+
+        if avg_eval_loss < best_eval_loss or np.isnan(avg_eval_loss):
+            best_eval_loss = avg_eval_loss
+            print("Saving checkpoint.")
+            jax_state_for_saving = py_utils.maybe_unreplicate_for_fully_replicated(
+                replicated_jax_states
+            )
+            if use_lora:
+                adapter_params = get_adapter_params(
+                    params=jax_state_for_saving.mdl_vars,
+                    lora_target_modules=lora_target_modules,
+                    num_layers=NUM_LAYERS,
+                    use_dora=use_dora,
+                )
+                jax_state_for_saving.mdl_vars["params"] = adapter_params
+
+            checkpoints.save_checkpoint(
+                jax_state_for_saving, checkpoint_dir, overwrite=True
+            )
+
+            patience = 0
+            del jax_state_for_saving
+            gc.collect()
+        else:
+            patience += 1
+            print(f"patience: {patience}")
+    print("Fine-tuning completed.")
+
+
+if __name__ == "__main__":
+    typer.run(finetune)

src/adapter/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TimesFM init file."""
+
+from .dora_layers import DoraAttentionProjection, DoraCombinedQKVProjection, DoraLinear
+from .lora_layers import LoraAttentionProjection, LoraCombinedQKVProjection, LoraLinear

src/adapter/dora_layers.py

@@ -0,0 +1,205 @@
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from jax import numpy as jnp
+from praxis import base_layer, pytypes
+from praxis.layers.attentions import AttentionProjection, CombinedQKVProjectionLayer
+from praxis.layers.linears import Linear
+
+WeightInit = base_layer.WeightInit
+template_field = base_layer.template_field
+WeightHParams = base_layer.WeightHParams
+JTensor = pytypes.JTensor
+
+
+class DoraTheta(base_layer.Theta):
+    def __init__(self, module):
+        self.module = module
+
+    def _dora_initialized(self):
+        if (
+            self.module.has_variable("params", "lora_a")
+            and self.module.has_variable("params", "lora_b")
+            and self.module.has_variable("params", "dora_m")
+            and "lora_a" in self.module._weight_hparams
+            and "lora_b" in self.module._weight_hparams
+            and "dora_m" in self.module._weight_hparams
+        ):
+            return True
+        else:
+            return False
+
+    def _dorafy_var(self, var):
+        lora_a = super().__getattr__("lora_a")
+        lora_b = super().__getattr__("lora_b")
+        dora_m = super().__getattr__("dora_m")
+
+        new_var = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
+        new_var = jnp.reshape(new_var, var.shape)
+
+        new_var += var
+
+        column_norm = jnp.linalg.norm(new_var, ord=2, axis=0, keepdims=True)
+        norm_adapted = new_var / column_norm
+        w = dora_m * norm_adapted
+        return w
+
+    def __getattr__(self, k):
+        var = super().__getattr__(k)
+        if not self._dora_initialized():
+            return var
+
+        if k == "w":
+            return self._dorafy_var(var)
+
+        return var
+
+    def __getitem__(self, k):
+        var = super().__getattr__(k)
+        if not self._dora_initialized():
+            return var
+
+        if k == "w":
+            return self._dorafy_var(var)
+
+        return var
+
+
+class DoraThetaDescriptor:
+    """Dot syntax accession descriptor."""
+
+    def __get__(self, obj, objtype=None):
+        return DoraTheta(obj)
+
+
+class DoraLinear(Linear):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = DoraThetaDescriptor()
+
+    def setup(self) -> None:
+        lora_init = self.lora_init if self.lora_init else self.weight_init
+
+        super().setup()
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[self.input_dims, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[self.output_dims, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+        self.create_variable(
+            "dora_m",
+            WeightHParams(
+                shape=[1, self.output_dims],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+
+
+class DoraAttentionProjection(AttentionProjection):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = DoraThetaDescriptor()
+
+    def setup(self) -> None:
+        super().setup()
+        w_weight_params = self._weight_hparams["w"]
+        lora_init = self.lora_init if self.lora_init else w_weight_params.init
+
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[self.input_dim, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[
+                    None,
+                    None,
+                ],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[self.dim_per_head * self.num_heads, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[
+                    None,
+                    None,
+                ],
+            ),
+        )
+        self.create_variable(
+            "dora_m",
+            WeightHParams(
+                shape=[1, self.num_heads, self.dim_per_head],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )
+
+
+class DoraCombinedQKVProjection(CombinedQKVProjectionLayer):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = DoraThetaDescriptor()
+
+    def setup(self) -> None:
+        super().setup()
+        w_weight_params = self._weight_hparams["w"]
+        lora_init = self.lora_init if self.lora_init else w_weight_params.init
+
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[3, self.input_dim, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[3, self.dim_per_head * self.num_heads, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )
+        self.create_variable(
+            "dora_m",
+            WeightHParams(
+                shape=[3, 1, self.num_heads, self.dim_per_head],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None, None],
+            ),
+        )

src/adapter/lora_layers.py

@@ -0,0 +1,170 @@
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from jax import numpy as jnp
+from praxis import base_layer, pax_fiddle, pytypes
+from praxis.layers.attentions import AttentionProjection, CombinedQKVProjectionLayer
+from praxis.layers.linears import Linear
+
+WeightInit = base_layer.WeightInit
+LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]
+template_field = base_layer.template_field
+WeightHParams = base_layer.WeightHParams
+JTensor = pytypes.JTensor
+
+
+class LoraTheta(base_layer.Theta):
+    def __init__(self, module):
+        self.module = module
+
+    def _lora_initialized(self):
+        if (
+            self.module.has_variable("params", "lora_a")
+            and self.module.has_variable("params", "lora_b")
+            and "lora_a" in self.module._weight_hparams
+            and "lora_b" in self.module._weight_hparams
+        ):
+            return True
+        else:
+            return False
+
+    def _lorafy_var(self, var):
+        lora_a = super().__getattr__("lora_a")
+        lora_b = super().__getattr__("lora_b")
+        new_var = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
+        new_var = jnp.reshape(new_var, var.shape)
+        new_var += var
+        return new_var
+
+    def __getattr__(self, k):
+        var = super().__getattr__(k)
+        if not self._lora_initialized():
+            return var
+
+        if k == "w":
+            return self._lorafy_var(var)
+
+        return var
+
+    def __getitem__(self, k):
+        var = super().__getattr__(k)
+        if not self._lora_initialized():
+            return var
+
+        if k == "w":
+            return self._lorafy_var(var)
+
+        return var
+
+
+class LoraThetaDescriptor:
+    """Dot syntax accession descriptor."""
+
+    def __get__(self, obj, objtype=None):
+        return LoraTheta(obj)
+
+
+class LoraLinear(Linear):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = LoraThetaDescriptor()
+
+    def setup(self) -> None:
+        lora_init = self.lora_init if self.lora_init else self.weight_init
+
+        super().setup()
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[self.input_dims, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[self.output_dims, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+
+
+class LoraAttentionProjection(AttentionProjection):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = LoraThetaDescriptor()
+
+    def setup(self) -> None:
+        super().setup()
+        w_weight_params = self._weight_hparams["w"]
+        lora_init = self.lora_init if self.lora_init else w_weight_params.init
+
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[self.input_dim, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[
+                    None,
+                    None,
+                ],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[self.dim_per_head * self.num_heads, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[
+                    None,
+                    None,
+                ],
+            ),
+        )
+
+
+class LoraCombinedQKVProjection(CombinedQKVProjectionLayer):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = LoraThetaDescriptor()
+
+    def setup(self) -> None:
+        super().setup()
+        w_weight_params = self._weight_hparams["w"]
+        lora_init = self.lora_init if self.lora_init else w_weight_params.init
+
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[3, self.input_dim, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[3, self.dim_per_head * self.num_heads, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )

src/adapter/utils.py

@@ -0,0 +1,411 @@
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+
+import jax
+import jax.numpy as jnp
+from paxml import checkpoints, tasks_lib
+from paxml.train_states import TrainState
+from praxis import pax_fiddle
+
+from adapter.dora_layers import (
+    DoraAttentionProjection,
+    DoraCombinedQKVProjection,
+    DoraLinear,
+)
+from adapter.lora_layers import (
+    LoraAttentionProjection,
+    LoraCombinedQKVProjection,
+    LoraLinear,
+)
+from timesfm import TimesFm
+
+
+def get_adapter_params(
+    params: dict, lora_target_modules: str, num_layers: int, use_dora: bool = False
+) -> dict:
+    adapter_params = {}
+    for i in range(num_layers):
+        layer_key = f"x_layers_{i}"
+        adapter_params[layer_key] = {}
+
+        if lora_target_modules in ["all", "mlp"]:
+            for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
+                linear = params["params"]["core_layer"]["stacked_transformer_layer"][
+                    layer_key
+                ]["ff_layer"][ff_layer_key]["linear"]
+
+                lora_a = linear["lora_a"]
+                lora_b = linear["lora_b"]
+
+                adapter_params[layer_key][ff_layer_key] = {
+                    "lora_a": lora_a,
+                    "lora_b": lora_b,
+                }
+
+                if use_dora:
+                    adapter_params[layer_key][ff_layer_key]["dora_m"] = linear["dora_m"]
+
+        if lora_target_modules in ["all", "attention"]:
+            attention = params["params"]["core_layer"]["stacked_transformer_layer"][
+                layer_key
+            ]["self_attention"]
+
+            for component in ["key", "query", "value", "post"]:
+                lora_a = attention[component]["lora_a"]
+                lora_b = attention[component]["lora_b"]
+
+                adapter_params[layer_key][component] = {
+                    "lora_a": lora_a,
+                    "lora_b": lora_b,
+                }
+
+                if use_dora:
+                    adapter_params[layer_key][component]["dora_m"] = attention[
+                        component
+                    ]["dora_m"]
+    return adapter_params
+
+
+def load_adapter_checkpoint(
+    model: TimesFm,
+    adapter_checkpoint_path: str,
+    lora_rank: int,
+    lora_target_modules: str,
+    use_dora: bool,
+) -> None:
+    """
+    currently loading and initializing the model with adapter layers first and then merging the
+    adapter weights to original weights and replacing the adapter layers back to original layer.
+    # NOTE: refactor this. there should be a better way to load the LoRA checkpoint.
+    """
+    model._logging(f"Restoring adapter checkpoint from {adapter_checkpoint_path}.")
+    start_time = time.time()
+    original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl = (
+        load_adapter_layer(
+            mdl_vars=model._train_state.mdl_vars,
+            model=model._model,
+            lora_rank=lora_rank,
+            lora_target_modules=lora_target_modules,
+            use_dora=use_dora,
+        )
+    )
+
+    var_weight_hparams = model._model.abstract_init_with_metadata(
+        model._get_sample_inputs(), do_eval=True
+    )
+
+    adapter_weight_hparams = _get_adapter_weight_params(
+        var_weight_hparams=var_weight_hparams,
+        lora_target_modules=lora_target_modules,
+        num_layers=model._model.stacked_transformer_params_tpl.num_layers,
+        use_dora=use_dora,
+    )
+
+    adapter_state_partition_specs = tasks_lib.create_state_partition_specs(
+        adapter_weight_hparams,
+        mesh_shape=model.mesh_shape,
+        mesh_axis_names=model.mesh_name,
+        discard_opt_states=True,
+        learners=None,
+    )
+    adapter_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
+        adapter_weight_hparams,
+        discard_opt_states=True,
+        learners=None,
+    )
+    adapter_train_state = checkpoints.restore_checkpoint(
+        state_global_shapes=adapter_state_local_shapes,
+        checkpoint_dir=adapter_checkpoint_path,
+        checkpoint_type=checkpoints.CheckpointType.FLAX,
+        state_specs=adapter_state_partition_specs,
+        step=None,
+    )
+
+    # add adapter weights to the original weights
+    _merge_adapter_weights(
+        model=model,
+        adapter_train_state=adapter_train_state,
+        lora_target_modules=lora_target_modules,
+        num_layers=model._model.stacked_transformer_params_tpl.num_layers,
+        use_dora=use_dora,
+    )
+
+    # replace back with the original model layer
+    if lora_target_modules in ["all", "mlp"]:
+        model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl = (
+            original_linear_tpl
+        )
+
+    if lora_target_modules in ["all", "attention"]:
+        model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl = (
+            original_attn_tpl
+        )
+        model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
+            original_combined_qkv_tpl
+        )
+    model._logging(
+        f"Restored adapter checkpoint in {time.time() - start_time:.2f} seconds."
+    )
+
+    # jit compile the model
+    model.jit_decode()
+
+
+def _merge_adapter_weights(
+    model: TimesFm,
+    adapter_train_state: TrainState,
+    lora_target_modules: str,
+    num_layers: int,
+    use_dora: bool,
+) -> None:
+    for i in range(num_layers):
+        layer_key = f"x_layers_{i}"
+
+        if lora_target_modules in ["all", "mlp"]:
+            for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
+                linear = model._train_state.mdl_vars["params"][
+                    "stacked_transformer_layer"
+                ][layer_key]["ff_layer"][ff_layer_key]["linear"]
+
+                params = adapter_train_state.mdl_vars[layer_key][ff_layer_key]
+                lora_a = params["lora_a"]
+                lora_b = params["lora_b"]
+
+                var = linear["w"]
+
+                new_var = jnp.einsum("...dr,...nr->...dn", lora_a, lora_b)
+                new_var = jnp.reshape(new_var, var.shape)
+                new_var += var
+
+                if use_dora:
+                    dora_m = params["dora_m"]
+                    column_norm = jnp.linalg.norm(new_var, ord=2, axis=0, keepdims=True)
+                    norm_adapted = new_var / column_norm
+                    calc_weights = dora_m * norm_adapted
+                    linear["w"] = calc_weights
+                    del linear["dora_m"]
+
+                else:
+                    linear["w"] = new_var
+
+                del linear["lora_a"]
+                del linear["lora_b"]
+
+        if lora_target_modules in ["all", "attention"]:
+            attention = model._train_state.mdl_vars["params"][
+                "stacked_transformer_layer"
+            ][layer_key]["self_attention"]
+
+            for component in ["key", "query", "value", "post"]:
+                params = adapter_train_state.mdl_vars[layer_key][component]
+                lora_a = params["lora_a"]
+                lora_b = params["lora_b"]
+
+                var = attention[component]["w"]
+
+                new_var = jnp.einsum("...dr,...nr->...dn", lora_a, lora_b)
+                new_var = jnp.reshape(new_var, var.shape)
+                new_var += var
+
+                if use_dora:
+                    m = params["dora_m"]
+                    column_norm = jnp.linalg.norm(new_var, ord=2, axis=0, keepdims=True)
+                    norm_adapted = new_var / column_norm
+                    calc_weights = m * norm_adapted
+                    attention[component]["w"] = calc_weights
+                    del attention[component]["dora_m"]
+
+                else:
+                    attention[component]["w"] = new_var
+
+                del attention[component]["lora_a"]
+                del attention[component]["lora_b"]
+
+
+def _get_adapter_weight_params(
+    var_weight_hparams: dict, lora_target_modules: str, num_layers: int, use_dora: bool
+) -> dict:
+    adapter_params = {}
+    for i in range(num_layers):
+        layer = f"x_layers_{i}"
+        adapter_params[layer] = {}
+
+        if lora_target_modules in ["all", "mlp"]:
+            for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
+                adapter_weight_params = var_weight_hparams["params"][
+                    "stacked_transformer_layer"
+                ][layer]["ff_layer"][ff_layer_key]["linear"]
+                adapter_params[layer][ff_layer_key] = {
+                    "lora_a": adapter_weight_params["lora_a"],
+                    "lora_b": adapter_weight_params["lora_b"],
+                }
+
+                if use_dora:
+                    adapter_params[layer][ff_layer_key]["dora_m"] = (
+                        adapter_weight_params["dora_m"]
+                    )
+
+        if lora_target_modules in ["all", "attention"]:
+            for component in ["key", "value", "query", "post"]:
+                adapter_weight_params = var_weight_hparams["params"][
+                    "stacked_transformer_layer"
+                ][layer]["self_attention"][component]
+                adapter_params[layer][component] = {
+                    "lora_a": adapter_weight_params["lora_a"],
+                    "lora_b": adapter_weight_params["lora_b"],
+                }
+
+                if use_dora:
+                    adapter_params[layer][component]["dora_m"] = adapter_weight_params[
+                        "dora_m"
+                    ]
+
+    return adapter_params
+
+
+def load_adapter_layer(
+    mdl_vars: dict,
+    model: pax_fiddle.Config,
+    lora_rank: int,
+    lora_target_modules: str,
+    use_dora: bool = False,
+) -> tuple[pax_fiddle.Config, pax_fiddle.Config]:
+    """
+    update self attention modules with LoRA/DoRA layers
+    """
+    original_linear_tpl = original_attn_tpl = original_combined_qkv_tpl = None
+    if lora_target_modules in ["all", "mlp"]:
+        original_linear_tpl = (
+            model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl
+        )
+        adapter_linear_tpl = (
+            pax_fiddle.Config(
+                DoraLinear,
+                rank=lora_rank,
+            )
+            if use_dora
+            else pax_fiddle.Config(
+                LoraLinear,
+                rank=lora_rank,
+            )
+        )
+        adapter_linear_tpl.copy_fields_from(original_linear_tpl)
+        model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl = (
+            adapter_linear_tpl
+        )
+
+    if lora_target_modules in ["all", "attention"]:
+        original_attn_tpl = (
+            model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl
+        )
+
+        adapter_attn_tpl = (
+            pax_fiddle.Config(DoraAttentionProjection, rank=lora_rank)
+            if use_dora
+            else pax_fiddle.Config(LoraAttentionProjection, rank=lora_rank)
+        )
+        adapter_attn_tpl.copy_fields_from(original_attn_tpl)
+
+        original_combined_qkv_tpl = (
+            model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl
+        )
+
+        adapter_combined_qkv_tpl = (
+            pax_fiddle.Config(DoraCombinedQKVProjection, rank=lora_rank)
+            if use_dora
+            else pax_fiddle.Config(LoraCombinedQKVProjection, rank=lora_rank)
+        )
+        adapter_combined_qkv_tpl.copy_fields_from(original_combined_qkv_tpl)
+
+        model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl = (
+            adapter_attn_tpl
+        )
+        model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
+            adapter_combined_qkv_tpl
+        )
+
+    # initialize and add adapter weights
+    _initialize_adapter_params(
+        mdl_vars=mdl_vars,
+        num_layers=model.stacked_transformer_params_tpl.num_layers,
+        lora_rank=lora_rank,
+        lora_target_modules=lora_target_modules,
+        use_dora=use_dora,
+    )
+
+    return original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl
+
+
+def _initialize_adapter_params(
+    mdl_vars: dict,
+    num_layers,
+    lora_rank: int,
+    lora_target_modules: str,
+    use_dora: bool = False,
+    seed: int = 1234,
+) -> dict:
+    """
+    initialize and add LoRA params in self attention
+    """
+    for i in range(num_layers):
+        layer_key = f"x_layers_{i}"
+        if lora_target_modules in ["all", "mlp"]:
+            for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
+                linear = mdl_vars["params"]["stacked_transformer_layer"][layer_key][
+                    "ff_layer"
+                ][ff_layer_key]["linear"]
+                original_w = linear["w"]
+                input_dim, output_dim = original_w.shape
+                std_dev = 1 / jnp.sqrt(lora_rank)
+
+                normal_initializer = jax.nn.initializers.normal(std_dev)
+                lora_a = normal_initializer(
+                    jax.random.key(seed), (input_dim, lora_rank), jnp.float32
+                )
+                lora_b = jnp.zeros((output_dim, lora_rank))
+
+                linear["lora_a"] = lora_a
+                linear["lora_b"] = lora_b
+
+                if use_dora:
+                    norm = jnp.linalg.norm(original_w, ord=2, axis=0, keepdims=True)
+                    linear["dora_m"] = norm
+
+        if lora_target_modules in ["all", "attention"]:
+            attention = mdl_vars["params"]["stacked_transformer_layer"][layer_key][
+                "self_attention"
+            ]
+
+            for component in ["key", "query", "value", "post"]:
+                original_w = attention[component]["w"]
+                w_dim = original_w.shape[0]
+                std_dev = 1 / jnp.sqrt(lora_rank)
+
+                normal_initializer = jax.nn.initializers.normal(std_dev)
+                lora_a = normal_initializer(
+                    jax.random.key(seed), (w_dim, lora_rank), jnp.float32
+                )
+                lora_b = jnp.zeros((w_dim, lora_rank))
+
+                attention[component]["lora_a"] = lora_a
+                attention[component]["lora_b"] = lora_b
+
+                if use_dora:
+                    norm = jnp.linalg.norm(
+                        original_w, ord=2, axis=0, keepdims=True
+                    ).astype(jnp.float32)
+                    attention[component]["dora_m"] = norm
+    return mdl_vars