diff --git a/.github/workflows/fine_tune.yaml b/.github/workflows/fine_tune.yaml
index 054fce5..34a7918 100644
--- a/.github/workflows/fine_tune.yaml
+++ b/.github/workflows/fine_tune.yaml
@@ -18,6 +18,10 @@ jobs:
       - uses: actions/checkout@v4
       - name: Set up Python 3.11
         uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: "pip"
+          cache-dependency-path: requirements.txt
       - name: Install dependencies
         run: pip install unionai
       - name: Run fine-tuning
diff --git a/config/phi_3_mini_128k_instruct.json b/config/phi_3_mini_128k_instruct.json
index 7a5b026..313b681 100644
--- a/config/phi_3_mini_128k_instruct.json
+++ b/config/phi_3_mini_128k_instruct.json
@@ -2,7 +2,7 @@
     "model_path": "microsoft/Phi-3-mini-128k-instruct",
     "output_dir": "./output",
     "checkpoint_dir": null,
-    "num_epochs": 100,
+    "num_epochs": 60,
     "max_steps": -1,
     "batch_size": 2,
     "test_size": 0.001,
diff --git a/requirements.txt b/requirements.txt
index 6fe9798..1c163a3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -19,8 +19,8 @@ tokenizers
 transformers==4.41.1
 torch==2.0.1
 wandb
-unionai==0.1.23
-flytekitplugins-envd==1.12.1rc1
-flytekitplugins-kfpytorch==1.12.1rc1
-flytekitplugins-huggingface==1.12.1rc1
-flytekitplugins-deck-standard==1.12.1rc1
+unionai==0.1.25
+flytekitplugins-envd==1.12.2
+flytekitplugins-kfpytorch==1.12.2
+flytekitplugins-huggingface==1.12.2
+flytekitplugins-deck-standard==1.12.2