From 690185abe155336c6d9a0288975c5244f672ae9f Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 31 May 2026 04:25:46 +0000 Subject: [PATCH] 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"