ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix SFTtraining for new trl (#2647)

Browse source 
* fix sft training with trl>0.15.2 with trl DataCollator

* Update fix to accomodate both trl and transformers DataCollatorForLanguageModeling
This commit is contained in:
DoubleMathew 2025-05-28 13:55:48 -05:00 • committed by GitHub
commit ebfa3df9c6
1 changed files with 4 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/models/rl.py
View file

@ -107,7 +107,7 @@ import torch
import numpy as np
from contextlib import nullcontext
from torch.nn import functional as F
from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling
from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling
torch_compile_options = {{
"epilogue_fusion" : True,
@ -358,8 +358,8 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
"from unsloth_zoo.vision_utils import UnslothVisionDataCollator\n"\
"if not isinstance(data_collator, UnslothVisionDataCollator):\n"\
" if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names:\n"\
" data_collator = DataCollatorForLanguageModeling(__tokenizer, mlm = False)\n"\
" elif isinstance(data_collator, DataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names:\n"\
" data_collator = TransformersDataCollatorForLanguageModeling(__tokenizer, mlm = False)\n"\
" elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names:\n"\
" data_collator = DataCollatorForSeq2Seq(__tokenizer)\n"\
"else:\n"\
" if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False\n"\
@ -374,7 +374,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
" if isinstance(data_collator, DataCollatorForSeq2Seq):\n"\
" data_collator = DataCollatorForSeq2Seq(__tokenizer.tokenizer)\n"\
" else:\n"\
" data_collator = DataCollatorForLanguageModeling(__tokenizer.tokenizer, mlm = False)\n"
" data_collator = TransformersDataCollatorForLanguageModeling(__tokenizer.tokenizer, mlm = False)\n"
extra_args += pad_check
pass