From 0288db2fa705a3d4e20e5ac1a8b63da9fd542f8c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 13 Mar 2025 06:02:33 -0700 Subject: [PATCH] Autocast --- unsloth/models/_utils.py | 10 ++++++++-- unsloth/models/rl.py | 16 ++++++++++++++-- unsloth/models/vision.py | 7 +++++-- 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index c13b2286f3..a3fc12f6d0 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -71,6 +71,7 @@ from typing import Union, Optional, List, Any, Callable, Tuple from platform import system as platform_system platform_system = platform_system() import numpy as np +import contextlib import warnings, subprocess, re, inspect, psutil, os, math from unsloth_zoo.utils import Version @@ -986,8 +987,13 @@ def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs): "Read more on gradient accumulation issues here: https://unsloth.ai/blog/gradient" ) pass - # with torch.autocast(device_type = "cuda", dtype = torch.float32): - outputs = self._old_compute_loss(model, inputs, *args, **kwargs) + + if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "0": + autocaster = contextlib.nullcontext() + else: + autocaster = torch.autocast(device_type = "cuda", dtype = torch.float32) + with autocaster: + outputs = self._old_compute_loss(model, inputs, *args, **kwargs) return outputs pass diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 020ce85e5e..f59892dcd4 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -236,6 +236,11 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): mixed_precision = \ "use_bf16 = getattr(args, 'bf16', False)\n"\ "use_fp16 = getattr(args, 'fp16', False)\n"\ + "force_float32 = False\n"\ + "if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1':\n"\ + " if use_bf16 or use_fp16:\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"\ "dtype = getattr(model.config, 'torch_dtype', None)\n"\ "if dtype is None: dtype = model.get_input_embeddings().dtype\n"\ @@ -244,7 +249,11 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): "float16 = dtype == torch.float16\n"\ "if 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 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 use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':\n"\ + "if force_float32:\n"\ + " args.fp16 = False\n"\ + " args.bf16 = False\n"\ + " os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n"\ + "elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':\n"\ " args.fp16 = float16\n"\ " args.bf16 = not float16\n"\ " os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16'\n" @@ -287,7 +296,10 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): "bf16_full_eval = getattr(args, 'bf16_full_eval', False)\n"\ "if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True\n"\ "if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False\n"\ - "if os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16':\n"\ + "if force_float32:\n"\ + " args.bf16_full_eval = False\n"\ + " args.fp16_full_eval = False\n"\ + "elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16':\n"\ " args.bf16_full_eval = True\n"\ " args.fp16_full_eval = False\n"\ "elif not bf16_full_eval and not fp16_full_eval:\n"\ diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 26e9edffd3..efdf67a95e 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -183,11 +183,14 @@ class FastBaseModel: global FORCE_FLOAT32 os.environ["UNSLOTH_FORCE_FLOAT32"] = "0" + bnb_compute_dtype = dtype for disable_name in FORCE_FLOAT32: if disable_name.lower() == model_type_arch.lower() and dtype == torch.float16: print(f"Unsloth: Using float16 precision for {model_type_arch} won't work! Using float32.") os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" + bnb_compute_dtype = torch.float32 break + pass bnb_config = None if full_finetuning and (load_in_4bit or load_in_8bit): @@ -203,7 +206,7 @@ class FastBaseModel: load_in_4bit = True, bnb_4bit_use_double_quant = True, bnb_4bit_quant_type = "nf4", - bnb_4bit_compute_dtype = dtype, + bnb_4bit_compute_dtype = bnb_compute_dtype, llm_int8_skip_modules = SKIP_QUANTIZATION_MODULES.copy(), ) elif load_in_8bit: @@ -218,7 +221,7 @@ class FastBaseModel: load_in_4bit = True, bnb_4bit_use_double_quant = True, bnb_4bit_quant_type = "nf4", - bnb_4bit_compute_dtype = dtype, + bnb_4bit_compute_dtype = bnb_compute_dtype, llm_int8_skip_modules = SKIP_QUANTIZATION_MODULES.copy(), ) pass