diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 72a6d7d10f..9a13b6fd4c 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -1200,8 +1200,9 @@ class FastModel(FastBaseModel): redirector = contextlib.redirect_stdout(open(os.devnull, "w")) model_types = ["siglip"] + model_types - # Set forced float32 env flag - os.environ["UNSLOTH_FORCE_FLOAT32"] = "0" + # Preserve explicit user override if set prior to model loading. + user_forced_float32 = os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1" + os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" if user_forced_float32 else "0" do_forced_float32 = False for model_type_arch in model_types: if model_type_arch != "siglip": diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index fbb32b2c50..26c406832a 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -715,19 +715,22 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): "if type(use_bf16) is not bool: use_bf16 = False\n" "use_fp16 = getattr(args, 'fp16', False)\n" "if type(use_fp16) is not bool: use_fp16 = False\n" - "force_float32 = False\n" - "full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1'\n" - "if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'):\n" + "force_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'\n" + "if force_float32:\n" " print('Unsloth: Switching to float32 training since model cannot work with float16')\n" - " force_float32 = True\n" "mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32')\n" "dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None)\n" "if dtype is None: dtype = model.get_input_embeddings().weight.dtype\n" "from unsloth_zoo.utils import _get_dtype\n" "dtype = _get_dtype(dtype)\n" - "float16 = dtype == torch.float16\n" - "if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n" - "if not force_float32 and (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" + "is_float16 = dtype == torch.float16\n" + "is_bfloat16 = dtype == torch.bfloat16\n" + "is_float32 = dtype == torch.float32\n" + "if not force_float32 and (is_float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n" + "if not force_float32 and (is_bfloat16 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 force_float32 and (is_float32 and use_fp16):\n" + " print('Unsloth: Model is float32 but fp16 was requested. Switching to float32 training.')\n" + " force_float32 = True\n" "if force_float32:\n" " # Forced float32 training\n" " args.fp16 = False\n" @@ -737,10 +740,22 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): " # args.mixed_precision is a new argument which needs to be set now\n" "elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':\n" " # Mixed precision training\n" - " args.fp16 = float16\n" - " args.bf16 = not float16\n" - " os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16'\n" - " if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16'\n" + " if is_float16:\n" + " args.fp16 = True\n" + " args.bf16 = False\n" + " os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16'\n" + " if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16'\n" + " elif is_bfloat16:\n" + " args.fp16 = False\n" + " args.bf16 = True\n" + " os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16'\n" + " if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16'\n" + " else:\n" + " # Float32 params should not force bf16/fp16 mixed precision.\n" + " args.fp16 = False\n" + " args.bf16 = False\n" + " os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n" + " if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n" " # args.mixed_precision is a new argument which needs to be set now\n" "elif mixed_precision_dtype == 'bfloat16':\n" " # Both False since bfloat16 full finetuning doesn't do any autocasting.\n"