ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Combine GRPO mixed precision fixes

Browse source 
Fold the explicit bf16/fp16 propagation from #4197 into the float32/full-finetuning safety work from #4114, and remove the remaining force-float32 autocast mismatch in GRPO helper replacements.

Co-authored-by: BenjaminBruenau <benjamin@bruenau.de>
This commit is contained in:
Daniel Han 2026-03-11 12:18:54 +00:00
commit d031935bbd
3 changed files with 94 additions and 43 deletions
Download patch file Download diff file Expand all files Collapse all files

15
unsloth/models/loader.py
View file

@ -1317,8 +1317,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":
@ -1332,6 +1342,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

66
unsloth/models/rl.py
View file

@ -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"
)

58
unsloth/models/rl_replacements.py
View file

@ -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, left_pack_padding
from unsloth_zoo.utils import Version
@ -240,10 +241,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)",
@ -620,17 +621,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,
@ -686,14 +691,22 @@ 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),
@ -710,9 +723,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]
@ -839,9 +849,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,