From 42a576a220402e9d21ca515f6b4ccbb7d13819e2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 14 Feb 2025 16:03:13 -0800 Subject: [PATCH] Update rl_replacements.py --- unsloth/models/rl_replacements.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 682a35ed1c..2925bd5b77 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -248,18 +248,14 @@ RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer_compute_loss) # https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L356 # TRL warns if batch size is not a multiple of num_generations -> fix this. def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source): - if "divisible by the number of generations" not in RLTrainer_source: - print(RLTrainer_source) - return "" - if "num_generations" not in RLConfig_source: - print(RLConfig_source) - return "" + if "divisible by the number of generations" not in RLTrainer_source: return "" + if "num_generations" not in RLConfig_source: return "" check_batch_size = \ "div = per_device_train_batch_size // num_generations\n"\ "if div * num_generations != per_device_train_batch_size:\n"\ - " print('Unsloth: We know expect `per_device_train_batch_size` to be a multiple of `num_generations`.\\n'\\"\ - " 'We will change the batch size of ' + per_device_train_batch_size + ' to the `num_generations` of ' + num_generations')\n"\ + " print('Unsloth: We know expect `per_device_train_batch_size` to be a multiple of `num_generations`.\\n"\ + "We will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations)')\n"\ " per_device_train_batch_size = num_generations\n" return check_batch_size pass