diff --git a/.github/workflows/gpu_ci_test.yml b/.github/workflows/gpu_ci_test.yml
index 41358fe1e..ccb0c68fc 100644
--- a/.github/workflows/gpu_ci_test.yml
+++ b/.github/workflows/gpu_ci_test.yml
@@ -14,7 +14,7 @@ concurrency:
 jobs:
   Test:
     name: Test
-    runs-on: [self-hosted, ernie-8gpu]
+    runs-on: [self-hosted, ernie-8gpu-1]
     steps:
       - name: Start Docker 
         run: |
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index ee001cb82..2a9149add 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -13,7 +13,7 @@ env:
 jobs:
   Lint:
     name: Lint
-    runs-on: [self-hosted, ernie-cpu]
+    runs-on: [self-hosted, ernie-cpu-01]
     permissions:
       pull-requests: write
       contents: read
diff --git a/erniekit/train/ocr_vl_sft/trainer.py b/erniekit/train/ocr_vl_sft/trainer.py
index cf606b996..0e59e20ea 100644
--- a/erniekit/train/ocr_vl_sft/trainer.py
+++ b/erniekit/train/ocr_vl_sft/trainer.py
@@ -39,7 +39,7 @@
 
 from distutils.util import strtobool
 
-from paddleformers.peft import LoRAModel, PrefixModelForCausalLM
+from paddleformers.peft import LoRAModel
 from paddleformers.trainer import (
     speed_metrics,
 )
@@ -833,9 +833,7 @@ def fused_allreduce_gradients_no_sync(paramlist, hcg):
             logger.info(
                 f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric})."
             )
-            if isinstance(self.model, LoRAModel) or isinstance(
-                self.model, PrefixModelForCausalLM
-            ):
+            if isinstance(self.model, LoRAModel):
                 self._load_best_model_from_peft_checkpoint()
             else:
                 weight_name = PADDLE_WEIGHTS_NAME
diff --git a/erniekit/train/vl_sft/trainer.py b/erniekit/train/vl_sft/trainer.py
index 3e3f79564..7aeba8a90 100644
--- a/erniekit/train/vl_sft/trainer.py
+++ b/erniekit/train/vl_sft/trainer.py
@@ -41,7 +41,7 @@
 
 from setuptools._distutils.util import strtobool
 
-from paddleformers.peft import LoRAModel, PrefixModelForCausalLM
+from paddleformers.peft import LoRAModel
 from paddleformers.trainer import (
     speed_metrics,
 )
@@ -829,9 +829,7 @@ def fused_allreduce_gradients_no_sync(paramlist, hcg):
             logger.info(
                 f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric})."
             )
-            if isinstance(self.model, LoRAModel) or isinstance(
-                self.model, PrefixModelForCausalLM
-            ):
+            if isinstance(self.model, LoRAModel):
                 self._load_best_model_from_peft_checkpoint()
             else:
                 weight_name = PADDLE_WEIGHTS_NAME