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

Update rl.py

Browse source
This commit is contained in:
Daniel Han 2025-03-10 04:57:47 -07:00
commit 11868359e2
1 changed files with 16 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

16
unsloth/models/rl.py
View file

@ -106,6 +106,8 @@ import torch
import numpy as np
from contextlib import nullcontext
from torch.nn import functional as F
from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling
torch_compile_options = {{
"epilogue_fusion" : True,
"max_autotune" : False,
@ -337,6 +339,20 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
extra_args += training_check
pass
# Check data collator if it's correct!
if "data_collator" in call_args and "train_dataset" in call_args:
data_collator_check = \
"if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names):\n"\
" print('Unsloth: Changing data collator to `DataCollatorForLanguageModeling` since `labels` not found.)\n"\
" data_collator = DataCollatorForLanguageModeling("\
"tokenizer = processing_class if 'processing_class' in locals() else tokenizer, mlm = False)\n"\
"elif isinstance(data_collator, DataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names):\n"\
" print('Unsloth: Changing data collator to `DataCollatorForSeq2Seq` since `labels` found.)\n"\
" data_collator = DataCollatorForSeq2Seq("\
"tokenizer = processing_class if 'processing_class' in locals() else tokenizer)\n"\
extra_args += data_collator_check
pass
# Check NEFTune
if "model" in call_args:
neftune_check = \