diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 4d5b4d7a4d..a41ddc8f07 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -235,6 +235,60 @@ pass RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__prepare_inputs) +# Fix incorrect special tokens handling and truncation in older TRL versions +def grpo_trainer__generate_and_score_completions(function_name, function): + if function_name != "_prepare_inputs": return function + + # TRL 0.19.0 did skip_special_tokens = True which should be False + function = function.replace( + "prompt_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False", + "prompt_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False", + ) + + # Always between max_prompt_length and use_vllm + found = re.findall( + r"(\n([\s]{1,})if self\.max_prompt_length is not None:.*?"\ + r"\2if self\.use_vllm:)", + a, + flags = re.DOTALL | re.MULTILINE, + ) + if len(found) != 0: + replace_part, spacing = found[0] + removed_comments = re.sub(r"\#[^\n]{1,}", "", replace_part) + splits = removed_comments.split("\n") + if sum(re.search(rf"^{spacing}[^\s]", x) is not None for x in splits) == 2 and \ + len(spacing) == 8: + replace_part + + new_replacement = spacing + """if self.max_prompt_length is not None: + # If max_prompt_length is set, we trim the prompt to keep only the last `max_prompt_length` tokens. + # Then we decode those tokens back into text. We manually remove leading pad tokens from the decoded text, + # because we can't use `skip_special_tokens=True` (some special tokens are still needed for generation). + prompt_ids = prompt_ids[:, -self.max_prompt_length :] + prompt_mask = prompt_mask[:, -self.max_prompt_length :] + prompts_text = self.processing_class.batch_decode( + prompt_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False + ) + pad_token = self.processing_class.pad_token + def strip_leading_tokens(text): + while text.startswith(pad_token): + text = text.removeprefix(pad_token) + return text + + if pad_token is not None: + prompts_text = [ + strip_leading_tokens(text) for text in prompts_text + ] + + # Generate completions using either vLLM or regular generation + if self.use_vllm:""" + function = function.replace(replace_part, new_replacement) + pass + return function +pass +RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__generate_and_score_completions) + + # Remove _move_model_to_vllm def grpo_trainer__move_model_to_vllm(function_name, function): if function_name != "_move_model_to_vllm": return function