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

Autocast

Browse source
This commit is contained in:
Daniel Han 2025-03-13 06:02:33 -07:00
commit 0288db2fa7
3 changed files with 27 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

10
unsloth/models/_utils.py
View file

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

16
unsloth/models/rl.py
View file

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

7
unsloth/models/vision.py
View file

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