diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 4ce3b2e1e9..9b58a493c4 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib import triton import ctypes @@ -35,7 +36,7 @@ import functools import torch torch_Tensor = torch.Tensor -from packaging.version import Version +from unsloth_zoo.utils import Version if DEVICE_TYPE == "xpu" and Version(torch.__version__) < Version("2.6.0"): raise RuntimeError( @@ -55,7 +56,6 @@ if DEVICE_TYPE == "xpu": # tl.math.tanh now is libdevice.tanh -from packaging.version import Version import triton import triton.language as tl @@ -211,6 +211,22 @@ torch_float16 = torch.float16 torch_bfloat16 = torch.bfloat16 +# Check whether torchao can be imported to get Float8Tensor +if importlib.util.find_spec("torchao") is not None: + try: + from torchao.quantization import Float8Tensor + except: + import torchao + + if Version(torchao.__version__) >= Version("0.15.0"): + print( + f"Unsloth: `from torchao.quantization import Float8Tensor` failed on version={torchao.__version__}" + ) + Float8Tensor = type(None) +else: + Float8Tensor = type(None) + + def QUANT_STATE(W): return getattr(W, "quant_state", None) @@ -335,6 +351,13 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: @torch.inference_mode def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False): # TODO: After adding XPU BNB support, check this function + if isinstance(W, Float8Tensor): + # TorchAO Float8Tensor + # In the backward pass, rowwise scaled becomes colwise scaled after we + # transpose the weight tensor. Use this case to detect backward + assert W.ndim == 2 + if W.block_size[0] == W.shape[0] and W.block_size[1] == 1: + return W.dequantize() if quant_state is None: return W if W.dtype == torch.float8_e4m3fn: @@ -441,6 +464,13 @@ elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM: @torch.inference_mode def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False): + if isinstance(W, Float8Tensor): + # TorchAO Float8Tensor + # In the backward pass, rowwise scaled becomes colwise scaled after we + # transpose the weight tensor. Use this case to detect backward + assert W.ndim == 2 + if W.block_size[0] == W.shape[0] and W.block_size[1] == 1: + return W.dequantize() if quant_state is None: return W if W.dtype == torch.float8_e4m3fn: @@ -551,6 +581,13 @@ else: @torch.inference_mode def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False): + if isinstance(W, Float8Tensor): + # TorchAO Float8Tensor + # In the backward pass, rowwise scaled becomes colwise scaled after we + # transpose the weight tensor. Use this case to detect backward + assert W.ndim == 2 + if W.block_size[0] == W.shape[0] and W.block_size[1] == 1: + return W.dequantize() if quant_state is None: return W if W.dtype == torch.float8_e4m3fn: @@ -987,8 +1024,8 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None): if W.dtype == torch.float8_e4m3fn: out = fp8_linear(X, W, W_quant) else: - W = fast_dequantize(W.t(), W_quant, use_global_buffer = True) - out = torch_matmul(X, W, out = out) + W = fast_dequantize(W, W_quant, use_global_buffer = True) + out = torch_matmul(X, W.t(), out = out) if W_quant is not None: del W diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index d63986db49..78291aebef 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -71,6 +71,7 @@ __all__ = [ "dequantize_module_weight", "patch_hf_quantizer", "verify_fp8_support_if_applicable", + "_get_inference_mode_context_manager", ] import torch @@ -2056,7 +2057,7 @@ except: @dataclass class TorchAOConfig: - qat_scheme: str = "int4" + qat_scheme: Optional[str] = "int4" # Each (config, filter_fn) pair defines a quantization rule base_config_and_filter_fns: List[ @@ -2306,3 +2307,22 @@ def verify_fp8_support_if_applicable(model_config): raise ValueError( f"Unsloth: FP8 quantization is only supported on L4 and higher GPUs with compute capability 8.9 or higher. You are using {torch.cuda.get_device_name()}. Refer to https://developer.nvidia.com/cuda-gpus for more details." ) + + +def _get_inference_mode_context_manager(model: torch.nn.Module): + """ + If the state dict was quantized using torchao, we will run into + the following error when calling ops like aten.t() in inference mode. + This is a bug in PyTorch that affects all tensor subclasses. + + Cannot set version_counter for inference tensor + + For now, we work around this issue by using `torch.no_grad()` in this case. + See https://github.com/pytorch/pytorch/issues/164872 for more details. + Otherwise, just return `torch.inference_mode()`. + """ + torchao_config = getattr(model, "torchao_config", None) + if torchao_config is not None and torchao_config.qat_scheme is None: + return torch.no_grad() + else: + return torch.inference_mode() diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index dde1ff98df..ad548dddb2 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -21,7 +21,10 @@ from ._utils import * from ._utils import patch_unsloth_smart_gradient_checkpointing from ._utils import __version__, importlib_version from ._utils import move_to_device -from ._utils import _prepare_model_for_qat +from ._utils import ( + _get_inference_mode_context_manager, + _prepare_model_for_qat, +) from torch.nn.functional import scaled_dot_product_attention from transformers import __version__ as transformers_version from unsloth_zoo.utils import Version, _get_dtype @@ -2030,7 +2033,7 @@ def unsloth_fast_generate( # Mixed precision autocast with ( - torch.inference_mode(), + _get_inference_mode_context_manager(self), torch.autocast(device_type = DEVICE_TYPE_TORCH, dtype = dtype), ): output = self._old_generate(*args, **kwargs) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 9d8a2ce9fa..8f7284ca70 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -31,7 +31,12 @@ from .cohere import FastCohereModel from transformers import AutoConfig from transformers import __version__ as transformers_version from peft import PeftConfig, PeftModel -from .loader_utils import get_model_name +from .loader_utils import ( + _check_load_in_fp8_settings, + _offline_quantize_to_fp8, + _tag_model_with_fp8_torchao_config, + get_model_name, +) import os, contextlib, sys try: @@ -140,6 +145,7 @@ class FastLanguageModel(FastLlamaModel): max_lora_rank = 64, disable_log_stats = True, qat_scheme = None, + load_in_fp8 = False, # fp8 LoRA *args, **kwargs, ): @@ -183,6 +189,7 @@ class FastLanguageModel(FastLlamaModel): max_lora_rank = max_lora_rank, disable_log_stats = disable_log_stats, qat_scheme = qat_scheme, + load_in_fp8 = load_in_fp8, *args, **kwargs, ) @@ -212,9 +219,23 @@ class FastLanguageModel(FastLlamaModel): ) load_in_4bit = False + if load_in_fp8: + _check_load_in_fp8_settings( + fast_inference, + full_finetuning, + load_in_4bit, + load_in_8bit, + load_in_16bit, + use_exact_model_name, + ) + old_model_name = model_name if not use_exact_model_name: - model_name = get_model_name(model_name, load_in_4bit) + if load_in_fp8: + model_name = _offline_quantize_to_fp8(model_name) + else: + model_name = get_model_name(model_name, load_in_4bit) + # Check if pre-quantized models are allowed # For eg AMD GPUs need blocksize = 128, but our pre-quants are blocksize = 64 if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith( @@ -476,6 +497,8 @@ class FastLanguageModel(FastLlamaModel): random_state = random_state, max_lora_rank = max_lora_rank, disable_log_stats = disable_log_stats, + qat_scheme = qat_scheme, + load_in_fp8 = load_in_fp8, *args, **kwargs, ) @@ -554,6 +577,9 @@ class FastLanguageModel(FastLlamaModel): } model.config.update({"quantization_config": quantization_config}) + if load_in_fp8: + _tag_model_with_fp8_torchao_config(model) + if is_peft: # From https://github.com/huggingface/peft/issues/184 # Now add PEFT adapters @@ -634,6 +660,7 @@ class FastModel(FastBaseModel): max_lora_rank = 64, disable_log_stats = True, qat_scheme = None, + load_in_fp8 = False, # fp8 LoRA *args, **kwargs, ): @@ -694,9 +721,23 @@ class FastModel(FastBaseModel): ) load_in_4bit = False + if load_in_fp8: + _check_load_in_fp8_settings( + fast_inference, + full_finetuning, + load_in_4bit, + load_in_8bit, + load_in_16bit, + use_exact_model_name, + ) + old_model_name = model_name if not use_exact_model_name: - model_name = get_model_name(model_name, load_in_4bit) + if load_in_fp8: + model_name = _offline_quantize_to_fp8(model_name) + else: + model_name = get_model_name(model_name, load_in_4bit) + # Check if pre-quantized models are allowed # For eg AMD GPUs need blocksize = 128, but our pre-quants are blocksize = 64 if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith( @@ -1130,6 +1171,9 @@ class FastModel(FastBaseModel): } model.config.update({"quantization_config": quantization_config}) + if load_in_fp8: + _tag_model_with_fp8_torchao_config(model) + if is_peft: # From https://github.com/huggingface/peft/issues/184 # Now add PEFT adapters diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 631f614cb4..f5504be2c8 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -12,11 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib +import os +import re +import tempfile from .mapper import INT_TO_FLOAT_MAPPER, FLOAT_TO_INT_MAPPER, MAP_TO_UNSLOTH_16bit # https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading! from packaging.version import Version -from transformers import __version__ as transformers_version +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + TorchAoConfig, + __version__ as transformers_version, +) +from unsloth.models._utils import TorchAOConfig +from unsloth_zoo.utils import Version +import torch transformers_version = Version(transformers_version) SUPPORTS_FOURBIT = transformers_version >= Version("4.37") @@ -144,3 +156,128 @@ def get_model_name(model_name, load_in_4bit = True): 'pip install --upgrade --no-cache-dir "git+https://github.com/unslothai/unsloth-zoo.git"\n' ) return new_model_name if new_model_name is not None else model_name + + +def _get_torchao_fp8_config(): + """ + Return a `torchao.quantization.Float8DynamicActivationFloat8WeightConfig` + to be used for `load_in_fp8=True`. + """ + from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerRow + + return Float8DynamicActivationFloat8WeightConfig( + granularity = PerRow(), + activation_value_lb = 1e-12, + ) + + +def _offline_quantize_to_fp8(model_name: str) -> str: + """ + Quantizes the model to fp8 using torchao and saving the quantized model to a + temporary location. Return the path to the quantized model. + + Note: Once on-the-fly quantization is added in vllm in + https://github.com/vllm-project/vllm/pull/26327, we should + dynamically quantize the model there instead: + + llm = LLM( + ... + hf_overrides={"quantization_config_file": "torchao_config.json"}, + ) + """ + temp_dir = tempfile.gettempdir() + new_model_name = model_name.split("/")[-1] + "-fp8" + new_model_name = os.path.join(temp_dir, new_model_name) + print( + f"Quantizing '{model_name}' to fp8, using model_name='{new_model_name}' instead" + ) + if not os.path.isdir(new_model_name): + qconfig = _get_torchao_fp8_config() + qconfig = TorchAoConfig(qconfig) + model = AutoModelForCausalLM.from_pretrained( + model_name, + torch_dtype = "auto", + device_map = "auto", + quantization_config = qconfig, + ) + tokenizer = AutoTokenizer.from_pretrained(model_name) + model.save_pretrained(new_model_name, safe_serialization = False) + tokenizer.save_pretrained(new_model_name) + return new_model_name + + +def _tag_model_with_fp8_torchao_config(model: torch.nn.Module): + """ + Tag a model with a `TorchAOConfig` so downstream callers will know what to do with it. + """ + base_config = _get_torchao_fp8_config() + model.torchao_config = TorchAOConfig( + qat_scheme = None, + base_config_and_filter_fns = [(base_config, None)], + ) + + +def _check_load_in_fp8_settings( + fast_inference: bool, + full_finetuning: bool, + load_in_4bit: bool, + load_in_8bit: bool, + load_in_16bit: bool, + use_exact_model_name: bool, +): + """ + Assuming `load_in_fp8=True`, raise appropriate errors on incompatible settings + and environment. Currently this feature requires: + 1. H100 GPUs or after + 2. torchao 0.15.0+ (or nightly) + 3. torch 2.9.0+ + 4. If fbgemm_gpu_genai is installed, require 1.4.1+ + """ + if not fast_inference: + raise ValueError( + "Unsloth: `load_in_fp8` is only supported for `fast_inference` for now" + ) + if full_finetuning: + raise ValueError( + "Unsloth: `load_in_fp8` is not compatible with full finetuning" + ) + if load_in_4bit or load_in_8bit or load_in_16bit: + raise ValueError( + "Unsloth: `load_in_fp8` is not compatible with `load_in_4bit`, `load_in_8bit` or `load_in_16bit`", + ) + if use_exact_model_name: + raise ValueError("Unsloth: `load_in_fp8` requires `use_exact_model_name=False`") + + # Check if this is Hopper or above + if not ( + torch.cuda.is_available() + and torch.version.cuda + and torch.cuda.get_device_capability() >= (9, 0) + ): + raise ValueError("Unsloth: `load_in_fp8` requires H100 GPUs or after") + + # Check if torch >= 2.9.0 + if Version(torch.__version__) < Version("2.9.0"): + raise ValueError("Unsloth: `load_in_fp8` requires torch 2.9.0+") + + # Check if torchao has this PR: https://github.com/pytorch/ao/pull/3158, + # which will be released in 0.15.0. + error_message = "Unsloth: `load_in_fp8` requires torchao 0.15.0+ (or nightly)" + if importlib.util.find_spec("torchao") is None: + raise ValueError(error_message) + import torchao + + if Version(torchao.__version__) < Version("0.15.0"): + raise ValueError(error_message) + + # If fbgemm_gpu_genai is installed, check if it's >= 1.4.1 + if ( + importlib.util.find_spec("fbgemm_gpu") is not None + and importlib.util.find_spec("fbgemm_gpu.experimental") is not None + ): + import fbgemm_gpu.experimental.gen_ai + + if Version(fbgemm_gpu.__version__) < Version("1.4.1"): + raise ValueError( + "Unsloth: `load_in_fp8` is only compatible with fbgemm_gpu_genai 1.4.1+" + ) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 8746399afa..38973fe30e 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -35,6 +35,7 @@ from ..device_type import ( ALLOW_PREQUANTIZED_MODELS, ) import textwrap +from ._utils import _get_inference_mode_context_manager RL_EXTRA_ARGS = defaultdict(list) RL_FUNCTIONS = defaultdict(list) @@ -536,7 +537,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): ) with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): - with torch.inference_mode(): + with _get_inference_mode_context_manager(model): if pixel_values is None: attention_mask = input_ids != self.processing_class.pad_token_id attention_mask = attention_mask.to(attention_mask.dtype) @@ -603,6 +604,9 @@ RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(UnslothEfficientGRPO)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss)) RL_PRE_ITEMS["grpo_trainer"].append(grpo_compute_loss_slow) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_update_SamplingParams)) +RL_PRE_ITEMS["grpo_trainer"].append( + inspect.getsource(_get_inference_mode_context_manager) +) # Edit _get_per_token_logps to handle mixed precision