423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-08-22 13:42:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #4045 from enji-zhou/feature/add_kto

Browse Source 
fix KTO Trainer Sampler

Former-commit-id: 30a538e2dbb79752479e11205dbb539df956f205
This commit is contained in:
hoshi-hiyouga 2024-06-03 22:09:25 +08:00 committed by GitHub
commit b2c224de69
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/llamafactory/train/kto/trainer.py
View File

@ -13,6 +13,7 @@ from ..utils import create_custom_optimzer, create_custom_scheduler
if TYPE_CHECKING: if TYPE_CHECKING:
import torch.utils.data
from transformers import PreTrainedModel, ProcessorMixin from transformers import PreTrainedModel, ProcessorMixin
from ...hparams import FinetuningArguments from ...hparams import FinetuningArguments
@ -84,6 +85,12 @@ class CustomKTOTrainer(KTOTrainer):
create_custom_scheduler(self.args, num_training_steps, optimizer) create_custom_scheduler(self.args, num_training_steps, optimizer)
return super().create_scheduler(num_training_steps, optimizer) return super().create_scheduler(num_training_steps, optimizer)
def _get_train_sampler(self) -> Optional["torch.utils.data.Sampler"]:
r"""
Replaces the sequential sampler of KTO Trainer created by trl with the random sampler.
"""
return Trainer._get_train_sampler(self)
def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, "torch.Tensor"]] = None) -> None: def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, "torch.Tensor"]] = None) -> None:
super()._save(output_dir, state_dict) super()._save(output_dir, state_dict)
if self.processor is not None: if self.processor is not None: