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

Use device-aware bf16 check so AMD/Intel are unaffected

Browse source 
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.
This commit is contained in:
danielhanchen 2026-05-31 04:25:46 +00:00
commit 690185abe1
2 changed files with 10 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
tests/python/test_v100_fullft_precision.py
View file

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

8
unsloth/models/rl.py
View file

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