From 8b1611c26c4752ba4adcc9ebf950b3d52ecabe87 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 31 May 2026 03:02:03 +0000 Subject: [PATCH 1/7] Fix full finetuning precision on V100 / no-bf16 GPUs Float16 full finetuning upcasts weights to float32, so the model dtype is float32, not bfloat16. The SFTTrainer precision template treated 'not float16' as 'bfloat16', which broke full finetuning on V100/T4: the dtype guard rejected float32 + fp16, and the auto mixed precision branch forced bf16 on hardware without it. FORCE_FLOAT32 models (Gemma3, gpt_oss, gemma3n, qwen3_5) were also only kept in float32 for LoRA, so full finetuning fell through to fp16 and produced NaNs. Distinguish bfloat16 from float32 in the guard, pick bf16 in the auto branch only when the GPU supports it (else fp16), and apply force_float32 for FORCE_FLOAT32 models in full finetuning too. bf16 GPUs are unchanged. --- unsloth/models/rl.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 359716fda4..2f806c59cd 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1027,8 +1027,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "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 models (Gemma3, gpt_oss, ...) cannot use float16; keep + # them in float32 even for full finetuning so V100/T4 never autocast to fp16. + "if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1':\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" @@ -1037,8 +1038,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "from unsloth_zoo.utils import _get_dtype\n" "dtype = _get_dtype(dtype)\n" "float16 = dtype == torch.float16\n" + "bfloat16 = dtype == torch.bfloat16\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" + "if not force_float32 and (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 force_float32:\n" " # Forced float32 training\n" " args.fp16 = False\n" @@ -1047,11 +1049,12 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): " 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 (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" + " # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16.\n" + " use_bf16_amp = (not float16) and torch.cuda.is_bf16_supported()\n" + " args.fp16 = not use_bf16_amp\n" + " args.bf16 = use_bf16_amp\n" + " os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16'\n" + " if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16'\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" From 84f76a42cbe1b033f6c8ba4ac705b45876343cd8 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 31 May 2026 03:43:12 +0000 Subject: [PATCH 2/7] Add regression tests for V100 full finetuning precision Exercise the real SFTTrainer mixed-precision template from rl.py source against mocked inputs: normal models get float32 weights + fp16 forward, FORCE_FLOAT32 models stay pure float32, no bf16 on no-bf16 hardware, and bf16 GPUs are unchanged. Covers issue #4082. --- tests/python/test_v100_fullft_precision.py | 126 +++++++++++++++++++++ 1 file changed, 126 insertions(+) create mode 100644 tests/python/test_v100_fullft_precision.py diff --git a/tests/python/test_v100_fullft_precision.py b/tests/python/test_v100_fullft_precision.py new file mode 100644 index 0000000000..cd88ff9d3f --- /dev/null +++ b/tests/python/test_v100_fullft_precision.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression tests for full finetuning precision on no-bf16 GPUs (V100/T4). + +Full finetuning upcasts trainable weights to float32, so the model dtype is +float32 (not bfloat16). The SFTTrainer mixed-precision template in +unsloth/models/rl.py must then: + - run the forward pass under float16 autocast for normal models, + - keep FORCE_FLOAT32 models (Gemma3, gpt_oss, ...) in pure float32, + - never select bf16 on hardware without bf16. + +We execute the REAL template block extracted from rl.py source (no heavy unsloth +import) against mocked inputs. See issue #4082. +""" +from __future__ import annotations + +import os +import sys +import types +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") + +RL_PY = Path(__file__).resolve().parents[2] / "unsloth" / "models" / "rl.py" + + +def _extract_mixed_precision_code() -> str: + lines = RL_PY.read_text().split("\n") + try: + start = next(i for i, l in enumerate(lines) if "mixed_precision = (" in l) + except StopIteration: + pytest.skip("mixed_precision template not found in rl.py") + body, k = [], start + 1 + while lines[k].strip() != ")": + body.append(lines[k]); k += 1 + return eval("(\n" + "\n".join(body) + "\n)") # only string literals + comments + + +CODE = _extract_mixed_precision_code() + + +def _decide(dtype, *, bf16_supported, force_float32, full_finetuning, + mixed_precision, fp16, bf16): + """Run the template block; return (args.fp16, args.bf16, ACCELERATE_MP, raised).""" + uz = types.ModuleType("unsloth_zoo"); uzu = types.ModuleType("unsloth_zoo.utils") + uzu._get_dtype = lambda x: x + sys.modules.setdefault("unsloth_zoo", uz); sys.modules["unsloth_zoo.utils"] = uzu + for k in ("UNSLOTH_FORCE_FLOAT32", "UNSLOTH_ENABLE_FULL_FINETUNING", + "UNSLOTH_MIXED_PRECISION", "ACCELERATE_MIXED_PRECISION"): + os.environ.pop(k, None) + os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" if force_float32 else "0" + os.environ["UNSLOTH_ENABLE_FULL_FINETUNING"] = "1" if full_finetuning else "0" + os.environ["UNSLOTH_MIXED_PRECISION"] = mixed_precision + orig = torch.cuda.is_bf16_supported + torch.cuda.is_bf16_supported = lambda *a, **k: bf16_supported + args = types.SimpleNamespace(fp16=fp16, bf16=bf16, mixed_precision=None) + emb = types.SimpleNamespace(weight=types.SimpleNamespace(dtype=dtype)) + model = types.SimpleNamespace( + config=types.SimpleNamespace(dtype=dtype, torch_dtype=dtype), + get_input_embeddings=lambda: emb) + raised = None + try: + exec(CODE, {"torch": torch, "os": os}, {"args": args, "model": model}) + except TypeError: + raised = "TypeError" + finally: + torch.cuda.is_bf16_supported = orig + return args.fp16, args.bf16, os.environ.get("ACCELERATE_MIXED_PRECISION"), raised + + +def test_v100_normal_fullft_fp16_explicit(): + # Normal model, full FT (weights upcast to float32), V100, fp16=True. + fp16, bf16, amp, raised = _decide( + torch.float32, bf16_supported=False, force_float32=False, + full_finetuning=True, mixed_precision="float32", fp16=True, bf16=False) + assert raised is None + assert (fp16, bf16) == (True, False) # float32 weights + fp16 forward + + +def test_v100_normal_fullft_precision_unset(): + # Same, but user left precision unset -> must pick fp16, never bf16. + fp16, bf16, amp, raised = _decide( + torch.float32, bf16_supported=False, force_float32=False, + full_finetuning=True, mixed_precision="float32", fp16=False, bf16=False) + assert raised is None + assert (fp16, bf16) == (True, False) + assert amp == "fp16" + + +def test_force_float32_model_fullft_is_pure_float32(): + # FORCE_FLOAT32 model (Gemma3, gpt_oss, ...) in full FT -> pure float32, no autocast. + fp16, bf16, amp, raised = _decide( + torch.float32, bf16_supported=False, force_float32=True, + full_finetuning=True, mixed_precision="float32", fp16=True, bf16=False) + assert raised is None + assert (fp16, bf16) == (False, False) + assert amp in (None, "no") + + +def test_no_bf16_on_volta_in_auto_branch(): + # bf16 model dtype but no bf16 HW, precision unset -> fp16, never bf16. + fp16, bf16, amp, raised = _decide( + torch.bfloat16, bf16_supported=False, force_float32=False, + full_finetuning=False, mixed_precision="float32", fp16=False, bf16=False) + assert bf16 is False + + +def test_bf16_gpu_unchanged_auto_branch(): + # Regression guard: on a bf16 GPU, a float32 model with unset precision + # still selects bf16 autocast (behavior must not change for bf16 hardware). + fp16, bf16, amp, raised = _decide( + torch.float32, bf16_supported=True, force_float32=False, + full_finetuning=True, mixed_precision="float32", fp16=False, bf16=False) + assert raised is None + assert (fp16, bf16) == (False, True) + + +def test_genuine_bf16_model_with_fp16_still_raises(): + # A real bfloat16 model on bf16 HW with fp16 requested is a genuine mismatch. + _, _, _, raised = _decide( + torch.bfloat16, bf16_supported=True, force_float32=False, + full_finetuning=False, mixed_precision="float32", fp16=True, bf16=False) + assert raised == "TypeError" From cca441f1258fb74181af6b1ae0dc2fd5482ba7a2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 31 May 2026 03:43:23 +0000 Subject: [PATCH 3/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/python/test_v100_fullft_precision.py | 101 ++++++++++++++++----- 1 file changed, 77 insertions(+), 24 deletions(-) diff --git a/tests/python/test_v100_fullft_precision.py b/tests/python/test_v100_fullft_precision.py index cd88ff9d3f..0567ee638e 100644 --- a/tests/python/test_v100_fullft_precision.py +++ b/tests/python/test_v100_fullft_precision.py @@ -13,6 +13,7 @@ unsloth/models/rl.py must then: We execute the REAL template block extracted from rl.py source (no heavy unsloth import) against mocked inputs. See issue #4082. """ + from __future__ import annotations import os @@ -35,32 +36,48 @@ def _extract_mixed_precision_code() -> str: pytest.skip("mixed_precision template not found in rl.py") body, k = [], start + 1 while lines[k].strip() != ")": - body.append(lines[k]); k += 1 + body.append(lines[k]) + k += 1 return eval("(\n" + "\n".join(body) + "\n)") # only string literals + comments CODE = _extract_mixed_precision_code() -def _decide(dtype, *, bf16_supported, force_float32, full_finetuning, - mixed_precision, fp16, bf16): +def _decide( + dtype, + *, + bf16_supported, + force_float32, + full_finetuning, + mixed_precision, + fp16, + bf16, +): """Run the template block; return (args.fp16, args.bf16, ACCELERATE_MP, raised).""" - uz = types.ModuleType("unsloth_zoo"); uzu = types.ModuleType("unsloth_zoo.utils") + uz = types.ModuleType("unsloth_zoo") + uzu = types.ModuleType("unsloth_zoo.utils") uzu._get_dtype = lambda x: x - sys.modules.setdefault("unsloth_zoo", uz); sys.modules["unsloth_zoo.utils"] = uzu - for k in ("UNSLOTH_FORCE_FLOAT32", "UNSLOTH_ENABLE_FULL_FINETUNING", - "UNSLOTH_MIXED_PRECISION", "ACCELERATE_MIXED_PRECISION"): + sys.modules.setdefault("unsloth_zoo", uz) + sys.modules["unsloth_zoo.utils"] = uzu + for k in ( + "UNSLOTH_FORCE_FLOAT32", + "UNSLOTH_ENABLE_FULL_FINETUNING", + "UNSLOTH_MIXED_PRECISION", + "ACCELERATE_MIXED_PRECISION", + ): os.environ.pop(k, None) os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" if force_float32 else "0" os.environ["UNSLOTH_ENABLE_FULL_FINETUNING"] = "1" if full_finetuning else "0" os.environ["UNSLOTH_MIXED_PRECISION"] = mixed_precision orig = torch.cuda.is_bf16_supported torch.cuda.is_bf16_supported = lambda *a, **k: bf16_supported - args = types.SimpleNamespace(fp16=fp16, bf16=bf16, mixed_precision=None) - emb = types.SimpleNamespace(weight=types.SimpleNamespace(dtype=dtype)) + args = types.SimpleNamespace(fp16 = fp16, bf16 = bf16, mixed_precision = None) + emb = types.SimpleNamespace(weight = types.SimpleNamespace(dtype = dtype)) model = types.SimpleNamespace( - config=types.SimpleNamespace(dtype=dtype, torch_dtype=dtype), - get_input_embeddings=lambda: emb) + config = types.SimpleNamespace(dtype = dtype, torch_dtype = dtype), + get_input_embeddings = lambda: emb, + ) raised = None try: exec(CODE, {"torch": torch, "os": os}, {"args": args, "model": model}) @@ -74,17 +91,29 @@ def _decide(dtype, *, bf16_supported, force_float32, full_finetuning, def test_v100_normal_fullft_fp16_explicit(): # Normal model, full FT (weights upcast to float32), V100, fp16=True. fp16, bf16, amp, raised = _decide( - torch.float32, bf16_supported=False, force_float32=False, - full_finetuning=True, mixed_precision="float32", fp16=True, bf16=False) + torch.float32, + bf16_supported = False, + force_float32 = False, + full_finetuning = True, + mixed_precision = "float32", + fp16 = True, + bf16 = False, + ) assert raised is None - assert (fp16, bf16) == (True, False) # float32 weights + fp16 forward + assert (fp16, bf16) == (True, False) # float32 weights + fp16 forward def test_v100_normal_fullft_precision_unset(): # Same, but user left precision unset -> must pick fp16, never bf16. fp16, bf16, amp, raised = _decide( - torch.float32, bf16_supported=False, force_float32=False, - full_finetuning=True, mixed_precision="float32", fp16=False, bf16=False) + torch.float32, + bf16_supported = False, + force_float32 = False, + full_finetuning = True, + mixed_precision = "float32", + fp16 = False, + bf16 = False, + ) assert raised is None assert (fp16, bf16) == (True, False) assert amp == "fp16" @@ -93,8 +122,14 @@ def test_v100_normal_fullft_precision_unset(): def test_force_float32_model_fullft_is_pure_float32(): # FORCE_FLOAT32 model (Gemma3, gpt_oss, ...) in full FT -> pure float32, no autocast. fp16, bf16, amp, raised = _decide( - torch.float32, bf16_supported=False, force_float32=True, - full_finetuning=True, mixed_precision="float32", fp16=True, bf16=False) + torch.float32, + bf16_supported = False, + force_float32 = True, + full_finetuning = True, + mixed_precision = "float32", + fp16 = True, + bf16 = False, + ) assert raised is None assert (fp16, bf16) == (False, False) assert amp in (None, "no") @@ -103,8 +138,14 @@ def test_force_float32_model_fullft_is_pure_float32(): def test_no_bf16_on_volta_in_auto_branch(): # bf16 model dtype but no bf16 HW, precision unset -> fp16, never bf16. fp16, bf16, amp, raised = _decide( - torch.bfloat16, bf16_supported=False, force_float32=False, - full_finetuning=False, mixed_precision="float32", fp16=False, bf16=False) + torch.bfloat16, + bf16_supported = False, + force_float32 = False, + full_finetuning = False, + mixed_precision = "float32", + fp16 = False, + bf16 = False, + ) assert bf16 is False @@ -112,8 +153,14 @@ def test_bf16_gpu_unchanged_auto_branch(): # Regression guard: on a bf16 GPU, a float32 model with unset precision # still selects bf16 autocast (behavior must not change for bf16 hardware). fp16, bf16, amp, raised = _decide( - torch.float32, bf16_supported=True, force_float32=False, - full_finetuning=True, mixed_precision="float32", fp16=False, bf16=False) + torch.float32, + bf16_supported = True, + force_float32 = False, + full_finetuning = True, + mixed_precision = "float32", + fp16 = False, + bf16 = False, + ) assert raised is None assert (fp16, bf16) == (False, True) @@ -121,6 +168,12 @@ def test_bf16_gpu_unchanged_auto_branch(): def test_genuine_bf16_model_with_fp16_still_raises(): # A real bfloat16 model on bf16 HW with fp16 requested is a genuine mismatch. _, _, _, raised = _decide( - torch.bfloat16, bf16_supported=True, force_float32=False, - full_finetuning=False, mixed_precision="float32", fp16=True, bf16=False) + torch.bfloat16, + bf16_supported = True, + force_float32 = False, + full_finetuning = False, + mixed_precision = "float32", + fp16 = True, + bf16 = False, + ) assert raised == "TypeError" From 690185abe155336c6d9a0288975c5244f672ae9f Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 31 May 2026 04:25:46 +0000 Subject: [PATCH 4/7] Use device-aware bf16 check so AMD/Intel are unaffected The auto mixed-precision branch now gates bf16 on unsloth_zoo's device_is_bf16_supported() (CUDA/XPU/HIP) instead of torch.cuda.is_bf16_supported(), which is only patched on CUDA. This keeps V100/T4 on fp16 while leaving AMD (HIP) and Intel (XPU) behavior unchanged. Falls back to the torch call on older unsloth_zoo. --- tests/python/test_v100_fullft_precision.py | 3 +++ unsloth/models/rl.py | 8 +++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/python/test_v100_fullft_precision.py b/tests/python/test_v100_fullft_precision.py index 0567ee638e..f15b574f23 100644 --- a/tests/python/test_v100_fullft_precision.py +++ b/tests/python/test_v100_fullft_precision.py @@ -58,8 +58,11 @@ def _decide( uz = types.ModuleType("unsloth_zoo") uzu = types.ModuleType("unsloth_zoo.utils") uzu._get_dtype = lambda x: x + uzd = types.ModuleType("unsloth_zoo.device_type") + uzd.device_is_bf16_supported = lambda: bf16_supported # device-aware signal stub sys.modules.setdefault("unsloth_zoo", uz) sys.modules["unsloth_zoo.utils"] = uzu + sys.modules["unsloth_zoo.device_type"] = uzd for k in ( "UNSLOTH_FORCE_FLOAT32", "UNSLOTH_ENABLE_FULL_FINETUNING", diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 2f806c59cd..49b755ffd3 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1036,6 +1036,12 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "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" + # device-aware bf16 check (CUDA/XPU/HIP), so V100/T4 never pick bf16 + # but AMD/Intel are unaffected; fall back on older unsloth_zoo. + "try:\n" + " from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported\n" + "except Exception:\n" + " _bf16_supported = torch.cuda.is_bf16_supported\n" "dtype = _get_dtype(dtype)\n" "float16 = dtype == torch.float16\n" "bfloat16 = dtype == torch.bfloat16\n" @@ -1050,7 +1056,7 @@ def _patch_trl_rl_trainers_impl(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. bf16 only if the GPU supports it; V100/T4 use fp16.\n" - " use_bf16_amp = (not float16) and torch.cuda.is_bf16_supported()\n" + " use_bf16_amp = (not float16) and _bf16_supported()\n" " args.fp16 = not use_bf16_amp\n" " args.bf16 = use_bf16_amp\n" " os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16'\n" From a1e2333d7c4c9532691f9275b8593d5285f060d1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 31 May 2026 12:58:16 +0000 Subject: [PATCH 5/7] Fix GRPO full finetuning returning logits instead of hidden states In full finetuning the model reaches GRPO as a plain *ForCausalLM with the stock forward, so it has no UNSLOTH_RETURN_HIDDEN_STATES branch and no support marker. The fallback then mis-targets the inner trunk (no lm_head) via _grpo_hidden_states_wrap_target, so the outer forward re-applies lm_head and the chunked log-softmax receives logits (vocab) instead of hidden states (hidden), crashing in the lm_head matmul. Prefer an lm_head passthrough: when UNSLOTH_RETURN_HIDDEN_STATES=1, short circuit lm_head to return its input so the forward yields logits == hidden and the vocab projection is skipped (memory efficient). Also harden the forward-wrapper fallback for models without a discoverable lm_head: do not descend past the lm_head owner, and tolerate an injected leading module arg from accelerate. No-op for LoRA/QLoRA and when the flag is unset. --- unsloth/models/rl.py | 78 +++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 74 insertions(+), 4 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 359716fda4..8374920e75 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -580,6 +580,31 @@ _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR = "_unsloth_grpo_hidden_states_forward_ _UNSLOTH_GRPO_HIDDEN_STATES_WARNING_ATTR = "_unsloth_grpo_hidden_states_warning_issued" +def _grpo_owns_lm_head(module): + # Does this module apply lm_head itself (i.e. its forward emits `.logits`)? + if module is None: + return False + if getattr(module, "lm_head", None) is not None: + return True + get_output_embeddings = getattr(module, "get_output_embeddings", None) + if callable(get_output_embeddings): + try: + return get_output_embeddings() is not None + except Exception: + return False + return False + + +def _grpo_causal_head(model): + # The module that owns lm_head / output embeddings (whose forward emits `.logits`). + get_base_model = getattr(model, "get_base_model", None) + if callable(get_base_model): + base_model = get_base_model() + if base_model is not None: + return base_model + return model + + def _grpo_hidden_states_wrap_target(model): if model is None: return None @@ -588,10 +613,14 @@ def _grpo_hidden_states_wrap_target(model): base_model = get_base_model() if base_model is not None and base_model is not model: return base_model - for attr in ("base_model", "model"): - child = getattr(model, attr, None) - if child is not None and child is not model and hasattr(child, "forward"): - return child + # Only descend into a child when `model` does not own lm_head itself. GRPO consumes the + # `.logits` of the module that applies lm_head; wrapping the inner trunk (no lm_head) would + # let the outer forward re-apply lm_head and leak logits into the chunked log-softmax (#708). + if not _grpo_owns_lm_head(model): + for attr in ("base_model", "model"): + child = getattr(model, attr, None) + if child is not None and child is not model and hasattr(child, "forward"): + return child return model @@ -686,12 +715,49 @@ def _replace_outputs_logits(outputs, hidden_states): ) +def _install_grpo_lm_head_passthrough(model): + # Preferred hidden-states path for a plain *ForCausalLM (e.g. full finetuning, where the model + # is not PEFT-wrapped, keeps the stock HF forward, and so has no RETURN_HIDDEN_STATES branch or + # support marker). Short-circuit lm_head to return its input (the hidden states) when + # UNSLOTH_RETURN_HIDDEN_STATES=1; the forward then yields `.logits == hidden`, which the GRPO + # log-prob path projects in chunks itself, and the full vocab projection is skipped. The lm_head + # weight is untouched, and the accelerate-managed top-level forward is not wrapped, so there is + # no bound-self collision. No-op when the flag is 0. + head = _grpo_causal_head(model) + lm_head = getattr(head, "lm_head", None) + if lm_head is None: + get_output_embeddings = getattr(head, "get_output_embeddings", None) + if callable(get_output_embeddings): + try: lm_head = get_output_embeddings() + except Exception: lm_head = None + if lm_head is None or getattr(lm_head, "_unsloth_grpo_passthrough", False): + return False + + original_lm_head_forward = lm_head.forward + def passthrough_forward(*args, **kwargs): + if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1": + return args[0] if args else next(iter(kwargs.values())) + return original_lm_head_forward(*args, **kwargs) + lm_head.forward = passthrough_forward + lm_head._unsloth_grpo_passthrough = True + setattr(model, _UNSLOTH_RETURN_HIDDEN_STATES_SUPPORT_MARKER, True) + setattr(head, _UNSLOTH_RETURN_HIDDEN_STATES_SUPPORT_MARKER, True) + return True + + def _install_grpo_hidden_states_forward_wrapper(model): if model is None or getattr(model, _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR, False): return False if _model_supports_unsloth_return_hidden_states(model): return False + # Preferred: short-circuit lm_head (robust for a plain full-FT CausalLM, skips the vocab + # projection, and avoids wrapping the accelerate-managed top-level forward). Fall back to the + # forward wrapper only when no lm_head can be found. + if _install_grpo_lm_head_passthrough(model): + setattr(model, _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR, True) + return True + target_model = _grpo_hidden_states_wrap_target(model) if getattr(target_model, _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR, False): setattr(model, _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR, True) @@ -702,6 +768,10 @@ def _install_grpo_hidden_states_forward_wrapper(model): model_name = type(target_model).__name__ def wrapped_forward(*args, **kwargs): + # Tolerate being invoked as a bound method: accelerate / nn.Module __call__ can inject + # `self` as the first positional arg once the wrapper lives on the outer CausalLM. + if len(args) > 0 and args[0] is target_model: + args = args[1:] if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") != "1": return original_forward(*args, **kwargs) From ca826bc3d8f712eaf24a0fe698285746ad167e62 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 31 May 2026 12:59:12 +0000 Subject: [PATCH 6/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/rl.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 8374920e75..82ef7b48d2 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -728,16 +728,20 @@ def _install_grpo_lm_head_passthrough(model): if lm_head is None: get_output_embeddings = getattr(head, "get_output_embeddings", None) if callable(get_output_embeddings): - try: lm_head = get_output_embeddings() - except Exception: lm_head = None + try: + lm_head = get_output_embeddings() + except Exception: + lm_head = None if lm_head is None or getattr(lm_head, "_unsloth_grpo_passthrough", False): return False original_lm_head_forward = lm_head.forward + def passthrough_forward(*args, **kwargs): if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1": return args[0] if args else next(iter(kwargs.values())) return original_lm_head_forward(*args, **kwargs) + lm_head.forward = passthrough_forward lm_head._unsloth_grpo_passthrough = True setattr(model, _UNSLOTH_RETURN_HIDDEN_STATES_SUPPORT_MARKER, True) From e0becf489f135464b0a1f430424bf98f3645a7ca Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 31 May 2026 13:53:37 +0000 Subject: [PATCH 7/7] Allow bf16 full finetuning for FORCE_FLOAT32 models on bf16 GPUs FORCE_FLOAT32 models (Gemma3, gpt_oss, qwen3_5) cannot use float16. Forcing float32 is required on GPUs without bf16 (V100/T4) so they never autocast to fp16, but on a bf16-capable GPU full finetuning can use bf16 autocast (the master weights stay float32), which is faster and uses less memory. Gate the force-float32 switch to skip full finetuning when bf16 is supported; LoRA/QLoRA still go to float32 when forced. --- unsloth/models/rl.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 49b755ffd3..6b076df23b 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1027,21 +1027,24 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "use_fp16 = getattr(args, 'fp16', False)\n" "if type(use_fp16) is not bool: use_fp16 = False\n" "force_float32 = False\n" - # FORCE_FLOAT32 models (Gemma3, gpt_oss, ...) cannot use float16; keep - # them in float32 even for full finetuning so V100/T4 never autocast to fp16. - "if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1':\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" # device-aware bf16 check (CUDA/XPU/HIP), so V100/T4 never pick bf16 # but AMD/Intel are unaffected; fall back on older unsloth_zoo. "try:\n" " from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported\n" "except Exception:\n" " _bf16_supported = torch.cuda.is_bf16_supported\n" + # FORCE_FLOAT32 models (Gemma3, gpt_oss, ...) cannot use float16. On a GPU without + # bf16 (V100/T4) keep them in float32 so they never autocast to fp16. On a bf16 GPU, + # full finetuning can still use bf16 autocast (master weights stay float32), which is + # faster and uses less memory; LoRA/QLoRA keep float32 when forced. + "full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1'\n" + "if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' and not (full_finetuning and _bf16_supported()):\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" "bfloat16 = dtype == torch.bfloat16\n"