Add adapter checkpoint support for continued training in Omni LoRA

train_omni_turn_lora.py

@@ -211,6 +211,14 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         help="Override target path for a single --cache-dir",
     )
     parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument(
+        "--adapter-checkpoint",
+        type=Path,
+        help=(
+            "Existing PEFT adapter directory to continue from; optimizer and "
+            "scheduler are reinitialized"
+        ),
+    )
     parser.add_argument("--head-checkpoint", type=Path)
     parser.add_argument("--max-examples", type=int, default=0)
     parser.add_argument("--max-steps", type=int, default=200)
@@ -696,6 +704,7 @@ def _load_initial_head(
 def _prepare_language_lora(
     model: torch.nn.Module,
     *,
+    adapter_checkpoint: Path | None,
     rank: int,
     alpha: int,
     dropout: float,
@@ -704,7 +713,7 @@ def _prepare_language_lora(
     verbose: bool = True,
 ) -> torch.nn.Module:
     try:
-        from peft import LoraConfig, get_peft_model
+        from peft import LoraConfig, PeftModel, get_peft_model
     except ImportError as exc:
         raise RuntimeError(
             "PEFT is required. The ms-swift 3.12.5 image should include it."
@@ -712,6 +721,39 @@ def _prepare_language_lora(
 
     model.requires_grad_(False)
     language_model = model.model
+    if adapter_checkpoint is not None:
+        adapter_checkpoint = adapter_checkpoint.expanduser().resolve()
+        config_path = adapter_checkpoint / "adapter_config.json"
+        if not config_path.is_file():
+            raise FileNotFoundError(
+                f"Missing PEFT adapter config: {config_path}"
+            )
+        language_model = PeftModel.from_pretrained(
+            language_model,
+            str(adapter_checkpoint),
+            is_trainable=True,
+        )
+        language_model.config.use_cache = False
+        if gradient_checkpointing:
+            try:
+                language_model.gradient_checkpointing_enable(
+                    gradient_checkpointing_kwargs={"use_reentrant": False}
+                )
+            except TypeError:
+                language_model.gradient_checkpointing_enable()
+        if verbose:
+            language_model.print_trainable_parameters()
+            trainable_lora = sum(
+                1
+                for name, parameter in language_model.named_parameters()
+                if parameter.requires_grad and "lora_" in name
+            )
+            print(
+                f"Continued trainable LoRA from {adapter_checkpoint}; "
+                f"trainable LoRA tensors={trainable_lora}"
+            )
+        return language_model
+
     matched = [
         name
         for name, module in language_model.named_modules()
@@ -852,6 +894,12 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         raise ValueError("--prefetch-depth must be positive")
     if args.ddp_timeout_minutes <= 0:
         raise ValueError("--ddp-timeout-minutes must be positive")
+    if args.adapter_checkpoint is not None and args.head_checkpoint is None:
+        if run.is_main:
+            print(
+                "Warning: --adapter-checkpoint was provided without "
+                "--head-checkpoint; the turn head will be reinitialized."
+            )
     if (
         args.audio_right_padding_samples is not None
         and args.audio_right_padding_samples < 0
@@ -955,6 +1003,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         raise ValueError("--lora-target-modules is empty")
     language_model = _prepare_language_lora(
         base_model,
+        adapter_checkpoint=args.adapter_checkpoint,
         rank=args.lora_rank,
         alpha=args.lora_alpha,
         dropout=args.lora_dropout,
@@ -1239,6 +1288,12 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         "schema_version": 1,
         "status": "pilot" if args.max_examples else "training",
         "model_id": args.model_id,
+        "initial_adapter_checkpoint": (
+            str(args.adapter_checkpoint.expanduser().resolve())
+            if args.adapter_checkpoint is not None
+            else None
+        ),
+        "optimizer_state_resumed": False,
         "cache_dirs": [str(path.expanduser().resolve()) for path in args.cache_dir],
         "num_examples": len(dataset),
         "initial_loss": initial_loss,