diff --git a/trl/experimental/bco/bco_trainer.py b/trl/experimental/bco/bco_trainer.py index a5a10725ae0..d807e41b8f2 100644 --- a/trl/experimental/bco/bco_trainer.py +++ b/trl/experimental/bco/bco_trainer.py @@ -377,7 +377,7 @@ class BCOTrainer(_BaseTrainer): The optimizer and scheduler to use for training. preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): The function to use to preprocess the logits before computing the metrics. - peft_config (`dict`, defaults to `None`): + peft_config ([`~peft.PeftConfig`], *optional*): The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model. compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): @@ -421,7 +421,7 @@ def __init__( callbacks: list[TrainerCallback] | None = None, optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - peft_config: dict | None = None, + peft_config: "PeftConfig | None" = None, compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, model_adapter_name: str | None = None, ref_adapter_name: str | None = None, diff --git a/trl/experimental/cpo/cpo_trainer.py b/trl/experimental/cpo/cpo_trainer.py index 432c0bfd070..6d59f8e660e 100644 --- a/trl/experimental/cpo/cpo_trainer.py +++ b/trl/experimental/cpo/cpo_trainer.py @@ -101,7 +101,7 @@ class CPOTrainer(_BaseTrainer): The optimizer and scheduler to use for training. preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): The function to use to preprocess the logits before computing the metrics. - peft_config (`dict`, defaults to `None`): + peft_config ([`~peft.PeftConfig`], *optional*): The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model. compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): @@ -142,7 +142,7 @@ def __init__( callbacks: list[TrainerCallback] | None = None, optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - peft_config: dict | None = None, + peft_config: "PeftConfig | None" = None, compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, ): if train_dataset is None: diff --git a/trl/experimental/nash_md/nash_md_trainer.py b/trl/experimental/nash_md/nash_md_trainer.py index f8b10b3392d..fd2f7c816c3 100644 --- a/trl/experimental/nash_md/nash_md_trainer.py +++ b/trl/experimental/nash_md/nash_md_trainer.py @@ -42,7 +42,7 @@ if is_peft_available(): - from peft import PeftModel + from peft import PeftConfig, PeftModel class GeometricMixtureWrapper(GenerationMixin): @@ -133,7 +133,7 @@ class NashMDTrainer(OnlineDPOTrainer): Processing class used to process the data. If provided, will be used to automatically process the inputs for the model, and it will be saved along the model to make it easier to rerun an interrupted training or reuse the fine-tuned model. - peft_config (`dict`): + peft_config ([`~peft.PeftConfig`], *optional*): The peft config to use for training. compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to @@ -177,7 +177,7 @@ def __init__( | FeatureExtractionMixin | ProcessorMixin | None = None, - peft_config: dict | None = None, + peft_config: "PeftConfig | None" = None, compute_metrics: Callable[[EvalPrediction], dict] | None = None, callbacks: list[TrainerCallback] | None = None, optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), diff --git a/trl/experimental/orpo/orpo_trainer.py b/trl/experimental/orpo/orpo_trainer.py index 22e6a81fbf6..eb8378b02db 100644 --- a/trl/experimental/orpo/orpo_trainer.py +++ b/trl/experimental/orpo/orpo_trainer.py @@ -112,7 +112,7 @@ class ORPOTrainer(_BaseTrainer): The optimizer and scheduler to use for training. preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): The function to use to preprocess the logits before computing the metrics. - peft_config (`dict`, defaults to `None`): + peft_config ([`~peft.PeftConfig`], *optional*): The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model. compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): @@ -151,7 +151,7 @@ def __init__( callbacks: list[TrainerCallback] | None = None, optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - peft_config: dict | None = None, + peft_config: "PeftConfig | None" = None, compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, ): if train_dataset is None: diff --git a/trl/experimental/prm/prm_trainer.py b/trl/experimental/prm/prm_trainer.py index c6bf17ad453..89cac4ddd44 100644 --- a/trl/experimental/prm/prm_trainer.py +++ b/trl/experimental/prm/prm_trainer.py @@ -127,7 +127,7 @@ class PRMTrainer(_BaseTrainer): The optimizer and scheduler to use for training. preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): The function to use to preprocess the logits before computing the metrics. - peft_config (`dict`, defaults to `None`): + peft_config ([`~peft.PeftConfig`], *optional*): The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in a PEFT model. """ @@ -167,7 +167,7 @@ def __init__( None, ), preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - peft_config: dict | None = None, + peft_config: "PeftConfig | None" = None, ): if train_dataset is None: raise ValueError("`train_dataset` is required") diff --git a/trl/experimental/xpo/xpo_trainer.py b/trl/experimental/xpo/xpo_trainer.py index f88125e2916..29fe8d7a6a6 100644 --- a/trl/experimental/xpo/xpo_trainer.py +++ b/trl/experimental/xpo/xpo_trainer.py @@ -41,7 +41,7 @@ if is_peft_available(): - from peft import PeftModel + from peft import PeftConfig, PeftModel class XPOTrainer(OnlineDPOTrainer): @@ -74,7 +74,7 @@ class XPOTrainer(OnlineDPOTrainer): Processing class used to process the data. If provided, will be used to automatically process the inputs for the model, and it will be saved along the model to make it easier to rerun an interrupted training or reuse the fine-tuned model. - peft_config (`dict`): + peft_config ([`~peft.PeftConfig`], *optional*): The peft config to use for training. compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to @@ -117,7 +117,7 @@ def __init__( | ProcessorMixin | None = None, reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, - peft_config: dict | None = None, + peft_config: "PeftConfig | None" = None, compute_metrics: Callable[[EvalPrediction], dict] | None = None, callbacks: list[TrainerCallback] | None = None, optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),