mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-24 08:13:59 +00:00
a335c00605
* Export ACCELERATE_MIXED_PRECISION when bf16/fp16 is set explicitly
In the default path the only branch that exports ACCELERATE_MIXED_PRECISION
is gated on `(not use_bf16 and not use_fp16)`, i.e. on the user having set
neither flag. Downstream readers resolve the autocast dtype from that
variable with a hardcoded 'fp16' default:
trainer._autocast_dtype = torch.float16 \
if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' \
else torch.bfloat16
(unsloth_zoo/rl_replacements.py, and unsloth/models/rl_replacements.py)
So passing bf16=True left an initially unset variable unset, the reader
fell back to 'fp16', and a float16 autocast was wrapped around a bfloat16
model. GRPO then crashed on the first step inside matmul_lora:
RuntimeError: self and mat2 must have the same dtype,
but got Half and BFloat16
Paradoxically the default configuration worked (the automatic branch runs
and exports 'bf16'), while explicitly selecting bf16 -- the configuration
recommended in #4891 -- crashed. The same failure also reproduces with
16-bit loading, so 4-bit quantization is not required.
Add an `elif use_bf16 or use_fp16` branch that exports the resolved
precision, restoring the invariant the readers assume. Placed last so the
force_float32 and UNSLOTH_MIXED_PRECISION branches keep precedence.
Fixes #4891
* Add regression tests for explicit bf16/fp16 precision export
Covers the #4891 gap and the precedence that must not change:
- explicit bf16=True exports ACCELERATE_MIXED_PRECISION='bf16'
- explicit fp16=True exports 'fp16'
- force_float32 still wins over the explicit-flag branch
- UNSLOTH_MIXED_PRECISION='bfloat16' still wins over it
The two export tests fail against the pre-fix template and pass with it;
the two precedence tests pass either way. Uses the existing _decide()
harness, which executes the real template block extracted from rl.py.
* Merge main and tighten the mixed-precision comment for PR #7534
Names the actual root cause: transformers exported ACCELERATE_MIXED_PRECISION
itself until 5.x dropped the write.
---------
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
262 lines
8.7 KiB
Python
262 lines
8.7 KiB
Python
# 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(encoding = "utf-8").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"
|
|
|
|
|
|
def test_explicit_bf16_exports_bf16_env():
|
|
# Issue #4891: bfloat16 model, user explicitly sets bf16=True. The env var
|
|
# must still be exported: downstream readers (unsloth_zoo/rl_replacements.py
|
|
# and unsloth/models/rl_replacements.py) default ACCELERATE_MIXED_PRECISION
|
|
# to 'fp16' when it is unset, which wraps a bfloat16 model in a float16
|
|
# autocast and crashes GRPO inside matmul_lora.
|
|
fp16, bf16, amp, raised = _decide(
|
|
torch.bfloat16,
|
|
bf16_supported = True,
|
|
force_float32 = False,
|
|
full_finetuning = False,
|
|
mixed_precision = "float32",
|
|
fp16 = False,
|
|
bf16 = True,
|
|
)
|
|
assert raised is None
|
|
assert (fp16, bf16) == (False, True)
|
|
assert amp == "bf16"
|
|
|
|
|
|
def test_explicit_fp16_exports_fp16_env():
|
|
# float16 model with fp16 explicitly requested must export 'fp16'.
|
|
fp16, bf16, amp, raised = _decide(
|
|
torch.float16,
|
|
bf16_supported = False,
|
|
force_float32 = False,
|
|
full_finetuning = False,
|
|
mixed_precision = "float32",
|
|
fp16 = True,
|
|
bf16 = False,
|
|
)
|
|
assert raised is None
|
|
assert (fp16, bf16) == (True, False)
|
|
assert amp == "fp16"
|
|
|
|
|
|
def test_force_float32_beats_explicit_bf16():
|
|
# force_float32 must keep precedence over the explicit-flag export.
|
|
fp16, bf16, amp, raised = _decide(
|
|
torch.float32,
|
|
bf16_supported = True,
|
|
force_float32 = True,
|
|
full_finetuning = False,
|
|
mixed_precision = "float32",
|
|
fp16 = False,
|
|
bf16 = True,
|
|
)
|
|
assert raised is None
|
|
assert (fp16, bf16) == (False, False)
|
|
assert amp == "no"
|
|
|
|
|
|
def test_unsloth_mixed_precision_bfloat16_beats_explicit_bf16():
|
|
# UNSLOTH_MIXED_PRECISION='bfloat16' (pure bf16 full FT, no autocast) must
|
|
# keep precedence over the explicit-flag export.
|
|
fp16, bf16, amp, raised = _decide(
|
|
torch.bfloat16,
|
|
bf16_supported = True,
|
|
force_float32 = False,
|
|
full_finetuning = True,
|
|
mixed_precision = "bfloat16",
|
|
fp16 = False,
|
|
bf16 = True,
|
|
)
|
|
assert raised is None
|
|
assert (fp16, bf16) == (False, False)
|
|
assert amp == "no"
|