From 753dcd255f2be5f0f3bd38c22da4b46ba210485c Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Sat, 31 Jan 2026 20:04:07 +0530 Subject: [PATCH] [trl] vllm trl topk fixup (#3935) * [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 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/rl_replacements.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index ff36da125d..ab0c37a608 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -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