diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index b54ceaf842..5cb486601e 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -1331,8 +1331,18 @@ 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 force-float32 requests without leaking prior + # auto-detected force state into later model loads in the same process. + user_forced_float32 = ( + type(float32_mixed_precision) is bool and float32_mixed_precision + ) + if not user_forced_float32: + user_forced_float32 = ( + os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1" + and os.environ.get("UNSLOTH_AUTO_FORCE_FLOAT32", "0") != "1" + ) + os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" if user_forced_float32 else "0" + os.environ["UNSLOTH_AUTO_FORCE_FLOAT32"] = "0" do_forced_float32 = False for model_type_arch in model_types: if model_type_arch != "siglip": @@ -1346,6 +1356,7 @@ class FastModel(FastBaseModel): or disable_name.lower() in model_types_all ) and ((dtype == torch.float16) or not SUPPORTS_BFLOAT16): os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" + os.environ["UNSLOTH_AUTO_FORCE_FLOAT32"] = "1" dtype = torch.bfloat16 # Change to bfloat16 loading break # Apply gradient checkpointing with smart heuristics diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index a05186eee5..181286a974 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -715,39 +715,71 @@ 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" + "original_dtype = getattr(model, '_unsloth_original_dtype', None)\n" + "if original_dtype is None and hasattr(model, 'model'): original_dtype = getattr(model.model, '_unsloth_original_dtype', None)\n" + "if original_dtype is None: original_dtype = dtype\n" + "else: original_dtype = _get_dtype(original_dtype)\n" + "is_float16 = original_dtype == torch.float16\n" + "is_bfloat16 = original_dtype == torch.bfloat16\n" + "is_float32 = original_dtype == torch.float32\n" + "explicit_mixed_precision = None\n" + "if use_bf16: explicit_mixed_precision = 'bf16'\n" + "elif use_fp16: explicit_mixed_precision = 'fp16'\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 (dtype == torch.float32) and (explicit_mixed_precision == 'fp16') and not is_float16:\n" + " print('Unsloth: Model is float32 but fp16 was requested. Switching to float32 training.')\n" + " force_float32 = True\n" + "if not force_float32 and (dtype == torch.float32) and (explicit_mixed_precision == 'bf16') and not is_bfloat16:\n" + " print('Unsloth: Model is float32 but bf16 was requested. Switching to float32 training.')\n" + " force_float32 = True\n" "if force_float32:\n" " # Forced float32 training\n" + " os.environ['UNSLOTH_FORCE_FLOAT32'] = '1'\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 explicit_mixed_precision is not None:\n" + " args.fp16 = explicit_mixed_precision == 'fp16'\n" + " args.bf16 = explicit_mixed_precision == 'bf16'\n" + " os.environ['ACCELERATE_MIXED_PRECISION'] = explicit_mixed_precision\n" + " if hasattr(args, 'mixed_precision'): args.mixed_precision = explicit_mixed_precision\n" + " # args.mixed_precision is a new argument which needs to be set now\n" + "elif mixed_precision_dtype == '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" + " # 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" - " # 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" - " 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" + " 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" "\n" ) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 805086324e..4e67f22669 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -25,6 +25,7 @@ import re import torch import inspect import linecache +from contextlib import nullcontext from collections import defaultdict from unsloth_zoo.rl_replacements import ( RL_REPLACEMENTS, @@ -244,10 +245,10 @@ def grpo_trainer__prepare_inputs(function_name, function): function = function.replace( "with torch.inference_mode():", "with torch.inference_mode(), " - "torch.amp.autocast(device_type = 'cuda', " - "dtype = ((torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16) " - "if not torch.is_autocast_enabled('cuda') else nullcontext())" - "if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '0' else torch.float16):", + "(nullcontext() if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' else " + "(torch.amp.autocast(device_type = 'cuda', " + "dtype = (torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16)) " + "if not torch.is_autocast_enabled('cuda') else nullcontext())):", ) function = function.replace( "self.accelerator.unwrap_model(self.model)", @@ -624,17 +625,21 @@ def grpo_trainer__get_per_token_logps(function_name, function): if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0': return None # Unsloth efficient GRPO # Otherwise, calculate normally: - if not hasattr(self, "_autocast_dtype"): - self._autocast_dtype = ( - torch.float16 - if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16" - else torch.bfloat16 + if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": + autocaster = nullcontext() + else: + if not hasattr(self, "_autocast_dtype"): + self._autocast_dtype = ( + torch.float16 + if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16" + else torch.bfloat16 + ) + autocaster = torch.amp.autocast( + device_type = DEVICE_TYPE, dtype = self._autocast_dtype ) - if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": - self._autocast_dtype = torch.float16 os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - with torch.amp.autocast(device_type = DEVICE_TYPE, dtype = self._autocast_dtype): + with autocaster: # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded logits = model( input_ids = input_ids, @@ -690,14 +695,25 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): if compute_efficient: return None, None else: - if not hasattr(self, "_autocast_dtype"): - self._autocast_dtype = ( - torch.float16 - if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16" - else torch.bfloat16 + if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": + autocaster = nullcontext() + dtype_bytes = 32 + else: + if not hasattr(self, "_autocast_dtype"): + self._autocast_dtype = ( + torch.float16 + if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") + == "fp16" + else torch.bfloat16 + ) + autocaster = torch.amp.autocast( + device_type = DEVICE_TYPE, dtype = self._autocast_dtype + ) + dtype_bytes = ( + 16 + if self._autocast_dtype in [torch.float16, torch.bfloat16] + else 32 ) - if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": - self._autocast_dtype = torch.float16 pixel_values, image_grid_thw = ( kwargs.get("pixel_values", None), @@ -714,9 +730,6 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): lm_head = self.model.get_output_embeddings().weight - dtype_bytes = ( - 16 if self._autocast_dtype in [torch.float16, torch.bfloat16] else 32 - ) total_rows = input_ids.shape[0] seq_len = input_ids.shape[1] hidden_dim = lm_head.shape[1] @@ -843,9 +856,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): pixel_attention_mask_chunk, image_sizes_chunk, ) in zipped_inputs: - with torch.amp.autocast( - device_type = "cuda", dtype = self._autocast_dtype - ): + with autocaster: if pixel_values is None: logits_chunk = unwrapped_model( input_ids = input_ids_chunk,