Fix full finetuning precision on V100 / no-bf16 GPUs (#5880)
--------- Co-authored-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
ba41e798d6
commit
de3c745fab
2 changed files with 212 additions and 7 deletions
193
tests/python/test_v100_fullft_precision.py
Normal file
193
tests/python/test_v100_fullft_precision.py
Normal file
|
|
@ -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"
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue