From 800536774afdc5e253e64762d4ee7cf0e9cbb1e7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 11 Feb 2025 16:20:14 -0800 Subject: [PATCH] Update rl.py --- unsloth/models/rl.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index c41b45f1b1..4c488187c9 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -72,6 +72,7 @@ pass RLTrainer_replacement = ''' +import os from typing import * from dataclasses import dataclass, field from packaging.version import Version @@ -188,7 +189,8 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): "if not float16 and use_fp16: raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n"\ "if not use_bf16 and not use_fp16:\n"\ " args.fp16 = float16\n"\ - " args.bf16 = not float16\n" + " args.bf16 = not float16\n"\ + " os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16'\n" extra_args += mixed_precision pass @@ -244,8 +246,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Add statistics as well! extra_args += \ "from unsloth_zoo.logging_utils import PatchRLStatistics\n"\ - f"PatchRLStatistics('{trainer_file}')\n"\ - "print(args)\n" + f"PatchRLStatistics('{trainer_file}')\n" # Create RLTrainer args extra_args = extra_args.split("\n")