diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index 960aed7..7fa60f3 100644 --- a/train_omni_turn_lora.py +++ b/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,