LoRA and DoRA PEFT support for Fine-Tuning TimesFM

.gitignore

@@ -1,3 +1,8 @@
 .venv/
 dist/
-**__pycache__/** */
+__pycache__/
+checkpoints/
+wandb/
+datasets/
+results/
+timesfm_jax.egg-info/

environment.yml

@@ -16,3 +16,6 @@ dependencies:
     - jax[cuda12]==0.4.26
     - einshape
     - scikit-learn
+    - typer
+    - wandb
+    - pytest

environment_cpu.yml

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

peft/dora.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+export TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false
+
+python3 finetune.py \
+    --model-name="google/timesfm-1.0-200m" \
+    --backend="gpu" \
+    --horizon-len=128 \
+    --context-len=512 \
+    --freq="15min" \
+    --data-path="../datasets/ETT-small/ETTm1.csv" \
+    --num-epochs=100 \
+    --learning-rate=1e-3 \
+    --adam-epsilon=1e-7 \
+    --adam-clip-threshold=1e2 \
+    --early-stop-patience=10 \
+    --datetime-col="date" \
+    --boundaries=34560 46080 57600 \
+    --use-lora \
+    --lora-rank=1 \
+    --lora-target-modules="all" \
+    --use-dora \
+    --cos-initial-decay-value=1e-4 \
+    --cos-decay-steps=40000 \
+    --cos-final-decay-value=1e-5 \
+    --ema-decay=0.9999

peft/fft.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+
+export TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false
+
+python3 finetune.py \
+    --model-name="google/timesfm-1.0-200m" \
+    --backend="gpu" \
+    --horizon-len=128 \
+    --context-len=512 \
+    --freq="15min" \
+    --data-path="../datasets/ETT-small/ETTm1.csv" \
+    --num-epochs=1 \
+    --learning-rate=1e-3 \
+    --adam-epsilon=1e-7 \
+    --adam-clip-threshold=1e2 \
+    --early-stop-patience=10 \
+    --datetime-col="date" \
+    --boundaries=34560 46080 57600 \
+    --cos-initial-decay-value=1e-4 \
+    --cos-decay-steps=40000 \
+    --cos-final-decay-value=1e-5 \
+    --ema-decay=0.9999