diff --git a/tests/python/test_v100_fullft_precision.py b/tests/python/test_v100_fullft_precision.py new file mode 100644 index 0000000000..c8ca769d45 --- /dev/null +++ b/tests/python/test_v100_fullft_precision.py @@ -0,0 +1,193 @@ +# 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 _restore(mapping, saved): + """Restore a dict-like to its saved snapshot: pop keys that were absent.""" + for k, v in saved.items(): + if v is None: + mapping.pop(k, None) + else: + mapping[k] = v + + +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). + + Stubs (sys.modules, env vars, torch.cuda.is_bf16_supported) are restored on + exit so a decision can't leak into later tests in the same process. + """ + 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 + + env_keys = ( + "UNSLOTH_FORCE_FLOAT32", + "UNSLOTH_ENABLE_FULL_FINETUNING", + "UNSLOTH_MIXED_PRECISION", + "ACCELERATE_MIXED_PRECISION", + ) + mod_keys = ("unsloth_zoo", "unsloth_zoo.utils", "unsloth_zoo.device_type") + saved_env = {k: os.environ.get(k) for k in env_keys} + saved_mods = {k: sys.modules.get(k) for k in mod_keys} + orig_bf16 = torch.cuda.is_bf16_supported + try: + sys.modules.setdefault("unsloth_zoo", types.ModuleType("unsloth_zoo")) + sys.modules["unsloth_zoo.utils"] = uzu + sys.modules["unsloth_zoo.device_type"] = uzd + for k in env_keys: + 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 + 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" + return args.fp16, args.bf16, os.environ.get("ACCELERATE_MIXED_PRECISION"), raised + finally: + torch.cuda.is_bf16_supported = orig_bf16 + _restore(os.environ, saved_env) + _restore(sys.modules, saved_mods) + + +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" diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index a85c1d08a4..53668d14d8 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -994,8 +994,18 @@ 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" + # 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 not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '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" @@ -1004,8 +1014,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" @@ -1014,11 +1025,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 _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"