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

[trl] vllm trl topk fixup (#3935)

Browse source 
* [transformers] [v5] remove unused hybridcache (#3910)

* remote unused hybridcache

* cleanup

* Fix top_k on trl GRPO

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Datta Nimmaturi 2026-01-31 20:04:07 +05:30 • committed by GitHub
commit 753dcd255f
1 changed files with 13 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

16
unsloth/models/rl_replacements.py
View file

@ -75,6 +75,16 @@ def sft_trainer_fix_untrained_tokens(call_args, extra_args):
RL_EXTRA_ARGS["sft_trainer"].append(sft_trainer_fix_untrained_tokens)
# Fix top_k for GRPO vLLM.
# https://github.com/huggingface/trl/pull/4695 with this change trl added top_k in GRPOConfig and defaults to 0
# We don't want that since vllm's all include top_k is -1 and 0 returns an error on SamplingParams creation.
def grpo_config_fix_vllm_top_k(old_RLTrainer_source, old_RLConfig_source):
return "if use_vllm and (top_k is None or top_k == 0): top_k = -1\n"
RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_config_fix_vllm_top_k)
# Remove DPO columns which might randomnly be tokenized
def dpo_trainer_fix_columns(call_args, extra_args):
if "model" in call_args and "train_dataset" in call_args:
@ -283,7 +293,7 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
re.MULTILINE,
)
replacement_text = """
replacement_text = """
if self.args.gradient_accumulation_steps % generate_every != 0 or (
self.use_vllm
):"""
@ -365,7 +375,7 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
replacement_string = """ if "image_sizes" in prompt_inputs:
output["image_sizes"] = prompt_inputs["image_sizes"]
if max_left_pad is not None:
output["max_left_pad"] = torch.tensor(prompt_ids.shape[0] * [max_left_pad]).unsqueeze(-1)
output["max_left_pad"] = torch.tensor(prompt_ids.shape[0] * [max_left_pad]).unsqueeze(-1)
try:
if self.use_vllm and getattr(self, "vllm_importance_sampling_correction", False):
output["sampling_per_token_logps"] = sampling_per_token_logps
@ -381,7 +391,7 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
replacement_string = """ if images is not None:
output["num_images"] = num_images
if max_left_pad is not None:
output["max_left_pad"] = torch.tensor(prompt_ids.shape[0] * [max_left_pad]).unsqueeze(-1)
output["max_left_pad"] = torch.tensor(prompt_ids.shape[0] * [max_left_pad]).unsqueeze(-1)
try:
if self.use_vllm and getattr(self, "vllm_importance_sampling_correction", False):
output["sampling_per_token_logps"] = sampling_per_token_logps