diff --git a/tests/test_fp8_restore_dropped_scale.py b/tests/test_fp8_restore_dropped_scale.py new file mode 100644 index 000000000..a85b24daf --- /dev/null +++ b/tests/test_fp8_restore_dropped_scale.py @@ -0,0 +1,358 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Restoring dropped block-fp8 `weight_scale_inv` tensors on load (#6200). + +Some block-scale fp8 checkpoints leave a Linear (e.g. `mlp.gate_proj`) unconverted, so its raw +quantized values land in a plain bf16 weight and its `weight_scale_inv` is dropped, producing a +garbage un-scaled weight. `_restore_dropped_fp8_scales` dequantizes such orphaned weights in place +using the scale from the checkpoint. Runs offline on CPU with synthetic checkpoints. +""" + +import json +import os +import tempfile +from types import SimpleNamespace + +import torch +from torch import nn +from safetensors.torch import save_file + +# Import unsloth first to set UNSLOTH_IS_PRESENT env var. +import unsloth +from unsloth.models.loader_utils import _restore_dropped_fp8_scales, _FP8_DTYPES + + +_SHARD = "model-00001-of-00001.safetensors" +_FP8 = _FP8_DTYPES[0] if _FP8_DTYPES else None + + +def _write_checkpoint( + path, + tensors, + filename = _SHARD, + include_index = True, +): + save_file(tensors, os.path.join(path, filename)) + if include_index: + weight_map = {name: filename for name in tensors} + with open(os.path.join(path, "model.safetensors.index.json"), "w") as f: + json.dump({"weight_map": weight_map}, f) + + +def _fp8_config(block = (2, 2)): + return SimpleNamespace( + quantization_config = { + "quant_method": "fp8", + "weight_block_size": list(block), + } + ) + + +def _fp8_anchor(): + """A module carrying a real fp8 weight, so the model looks like a genuine fp8 load.""" + m = nn.Linear(2, 2, bias = False) + m.weight = nn.Parameter(torch.randn(2, 2).to(_FP8), requires_grad = False) + return m + + +def _bf16_linear(out_f, in_f, raw): + m = nn.Linear(in_f, out_f, bias = False).to(torch.bfloat16) + with torch.no_grad(): + m.weight.copy_(raw) + return m + + +def _expand(scale, block, shape): + bs0, bs1 = block + expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1) + return expanded[: shape[0], : shape[1]] + + +def test_restore_dequantizes_orphaned_scale(): + """A plain bf16 weight whose scale was dropped is dequantized in place.""" + if _FP8 is None: + return + torch.manual_seed(0) + raw = torch.randn(4, 4, dtype = torch.bfloat16) + scale = torch.rand(2, 2, dtype = torch.float32) + 0.1 + + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.layer = _bf16_linear(4, 4, raw) + + with tempfile.TemporaryDirectory() as d: + _write_checkpoint( + d, + { + "layer.weight": raw.to(torch.float32), + "layer.weight_scale_inv": scale, + }, + ) + restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True) + + assert restored == 1 + expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16) + assert torch.equal(model.layer.weight.data, expected) + + +def test_skips_already_fp8_weight(): + """A correctly converted fp8 weight is skipped, never double-scaled.""" + if _FP8 is None: + return + weight = torch.randn(4, 4).to(_FP8) + before = weight.clone() + + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.layer = nn.Linear(4, 4, bias = False) + model.layer.weight = nn.Parameter(weight, requires_grad = False) + + with tempfile.TemporaryDirectory() as d: + _write_checkpoint(d, {"layer.weight_scale_inv": torch.rand(2, 2)}) + restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True) + + assert restored == 0 and skipped == 1 + assert torch.equal(model.layer.weight.data.float(), before.float()) + + +def test_skips_offloaded_meta_weight(): + """A disk-offloaded layer (weight on the meta device) is skipped without error or restore.""" + if _FP8 is None: + return + raw = torch.randn(4, 4, dtype = torch.bfloat16) + scale = torch.rand(2, 2, dtype = torch.float32) + 0.1 + + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.layer = nn.Linear(4, 4, bias = False) + # Simulate an offloaded weight living on the meta device. + model.layer.weight = nn.Parameter( + torch.empty(4, 4, dtype = torch.bfloat16, device = "meta"), requires_grad = False + ) + + with tempfile.TemporaryDirectory() as d: + _write_checkpoint( + d, + { + "layer.weight": raw.to(torch.float32), + "layer.weight_scale_inv": scale, + }, + ) + restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True) + + assert restored == 0 + assert model.layer.weight.device.type == "meta" + + +def test_noop_when_fully_dequantized(): + """If the model has no fp8 weights at all (e.g. load_in_16bit dequantize), do not rescale.""" + raw = torch.randn(4, 4, dtype = torch.bfloat16) + scale = torch.rand(2, 2, dtype = torch.float32) + 0.1 + + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.layer = _bf16_linear(4, 4, raw) # no fp8 anchor -> looks dequantized + + with tempfile.TemporaryDirectory() as d: + _write_checkpoint(d, {"layer.weight_scale_inv": scale}) + restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True) + + assert (restored, skipped) == (0, 0) + assert torch.equal(model.layer.weight.data, raw) # untouched + + +def test_non_block_divisible_shape(): + """Block scale is expanded then sliced to a non-divisible weight shape.""" + if _FP8 is None: + return + raw = torch.randn(3, 4, dtype = torch.bfloat16) + scale = torch.rand(2, 2, dtype = torch.float32) + 0.1 + + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.layer = _bf16_linear(3, 4, raw) # weight shape [3, 4] + + with tempfile.TemporaryDirectory() as d: + _write_checkpoint(d, {"layer.weight_scale_inv": scale}) + restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True) + + assert restored == 1 + expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (3, 4))).to(torch.bfloat16) + assert torch.equal(model.layer.weight.data, expected) + + +def test_transposed_scale_layout(): + """A scale stored in the transposed block grid is transposed before use.""" + if _FP8 is None: + return + raw = torch.randn(4, 2, dtype = torch.bfloat16) # weight [4, 2] -> grid (2, 1) + scale_correct = torch.rand(2, 1, dtype = torch.float32) + 0.1 + scale_stored = scale_correct.t().contiguous() # stored transposed as (1, 2) + + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.layer = _bf16_linear(4, 2, raw) + + with tempfile.TemporaryDirectory() as d: + _write_checkpoint(d, {"layer.weight_scale_inv": scale_stored}) + restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True) + + assert restored == 1 + expected = (raw.to(torch.float32) * _expand(scale_correct, (2, 2), (4, 2))).to(torch.bfloat16) + assert torch.equal(model.layer.weight.data, expected) + + +def test_single_file_checkpoint_without_index(): + """Unsharded model.safetensors (no index) is still scanned for dropped scales.""" + if _FP8 is None: + return + raw = torch.randn(4, 4, dtype = torch.bfloat16) + scale = torch.rand(2, 2, dtype = torch.float32) + 0.1 + + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.layer = _bf16_linear(4, 4, raw) + + with tempfile.TemporaryDirectory() as d: + _write_checkpoint( + d, {"layer.weight_scale_inv": scale}, filename = "model.safetensors", include_index = False + ) + restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True) + + assert restored == 1 + expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16) + assert torch.equal(model.layer.weight.data, expected) + + +def test_scalar_block_size_config(): + """A scalar weight_block_size (not a list) is handled without error.""" + if _FP8 is None: + return + raw = torch.randn(4, 4, dtype = torch.bfloat16) + scale = torch.rand(2, 2, dtype = torch.float32) + 0.1 + + model = nn.Module() + model.config = SimpleNamespace( + quantization_config = {"quant_method": "fp8", "weight_block_size": 2} + ) + model.anchor = _fp8_anchor() + model.layer = _bf16_linear(4, 4, raw) + + with tempfile.TemporaryDirectory() as d: + _write_checkpoint(d, {"layer.weight_scale_inv": scale}) + restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True) + + assert restored == 1 + + +def test_text_only_prefix_mapping(): + """Checkpoint keys with a language_model prefix match the stripped text-only module names.""" + if _FP8 is None: + return + raw = torch.randn(2, 2, dtype = torch.bfloat16) + scale = torch.rand(1, 1, dtype = torch.float32) + 0.1 + + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.model = nn.Module() + model.model.gate_proj = _bf16_linear(2, 2, raw) # module lacks the language_model prefix + + with tempfile.TemporaryDirectory() as d: + # checkpoint key carries the language_model wrapper the text-only load stripped + _write_checkpoint(d, {"model.language_model.gate_proj.weight_scale_inv": scale}) + restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True) + + assert restored == 1 + expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16) + assert torch.equal(model.model.gate_proj.weight.data, expected) + + +def test_skips_variant_load(): + """A variant load (variant="fp8") is skipped to avoid applying default-checkpoint scales.""" + if _FP8 is None: + return + raw = torch.randn(4, 4, dtype = torch.bfloat16) + scale = torch.rand(2, 2, dtype = torch.float32) + 0.1 + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.layer = _bf16_linear(4, 4, raw) + with tempfile.TemporaryDirectory() as d: + _write_checkpoint(d, {"layer.weight_scale_inv": scale}) + result = _restore_dropped_fp8_scales(model, d, local_files_only = True, variant = "fp8") + assert result == (0, 0) + assert torch.equal(model.layer.weight.data, raw) # untouched + + +def test_vlm_language_model_model_alias(): + """A checkpoint key language_model.model.* matches a model.language_model.* module.""" + if _FP8 is None: + return + raw = torch.randn(2, 2, dtype = torch.bfloat16) + scale = torch.rand(1, 1, dtype = torch.float32) + 0.1 + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.model = nn.Module() + model.model.language_model = nn.Module() + model.model.language_model.gate_proj = _bf16_linear( + 2, 2, raw + ) # -> model.language_model.gate_proj + with tempfile.TemporaryDirectory() as d: + _write_checkpoint(d, {"language_model.model.gate_proj.weight_scale_inv": scale}) + restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True) + assert restored == 1 + expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16) + assert torch.equal(model.model.language_model.gate_proj.weight.data, expected) + + +def test_noop_without_scale_keys(): + if _FP8 is None: + return + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16)) + with tempfile.TemporaryDirectory() as d: + _write_checkpoint(d, {"layer.weight": torch.randn(4, 4)}) + assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0) + + +def test_noop_without_index_or_single_file(): + if _FP8 is None: + return + model = nn.Module() + model.config = _fp8_config((2, 2)) + model.anchor = _fp8_anchor() + model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16)) + with tempfile.TemporaryDirectory() as d: + assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0) + + +def test_noop_when_not_block_fp8(): + """A non-fp8 (or non-block) quantization config is ignored.""" + scale = torch.rand(2, 2) + model = nn.Module() + model.config = SimpleNamespace(quantization_config = {"quant_method": "compressed-tensors"}) + model.layer = nn.Linear(4, 4, bias = False) + with tempfile.TemporaryDirectory() as d: + _write_checkpoint(d, {"layer.weight_scale_inv": scale}) + assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 05523bc27..defacadfe 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -28,7 +28,11 @@ from ._utils import ( is_bfloat16_supported, get_quant_type, ) -from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings +from .loader_utils import ( + _exclude_rope_inv_freq_from_ddp, + _get_fp8_mode_and_check_settings, + _restore_dropped_fp8_scales, +) from ..utils.packing import ( get_packed_info_from_kwargs, mask_packed_sequence_boundaries, @@ -2678,6 +2682,18 @@ class FastLlamaModel: offload_embedding = False, fast_inference = fast_inference, ) + # Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200). + _restore_dropped_fp8_scales( + model, + model_name, + local_files_only = kwargs.get("local_files_only", False), + token = token, + # Weights load from the default branch (revision not forwarded), so read scales from there too. + revision = None, + subfolder = kwargs.get("subfolder"), + cache_dir = kwargs.get("cache_dir"), + variant = kwargs.get("variant"), + ) elif not fast_inference: if user_config is not None: # Transformers 5.x @strict model init rejects extra kwargs next @@ -2716,6 +2732,18 @@ class FastLlamaModel: offload_embedding = False, fast_inference = False, ) + # Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200). + _restore_dropped_fp8_scales( + model, + model_name, + local_files_only = kwargs.get("local_files_only", False), + token = token, + # Weights load from the default branch (revision not forwarded), so read scales from there too. + revision = None, + subfolder = kwargs.get("subfolder"), + cache_dir = kwargs.get("cache_dir"), + variant = kwargs.get("variant"), + ) model.fast_generate = make_fast_generate_wrapper(model.generate) model.fast_generate_batches = None else: diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index fa6282bcf..1e9cc641e 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -367,6 +367,287 @@ def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str): pass +_FP8_DTYPES = tuple( + dtype + for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None)) + if dtype is not None +) + + +def _fp8_block_size_from_config(model): + """Return the [block_out, block_in] block size of an fp8 checkpoint, or None if not block-fp8.""" + config = getattr(model, "config", None) + quant = getattr(config, "quantization_config", None) + if quant is None: + return None + if hasattr(quant, "to_dict"): + quant = quant.to_dict() + if not isinstance(quant, dict): + return None + if quant.get("quant_method") != "fp8": + return None + block = quant.get("weight_block_size") + if not block: + return None + if isinstance(block, (int, float)): + block = [block, block] + elif isinstance(block, (list, tuple)): + if len(block) == 1: + block = [block[0], block[0]] + elif len(block) < 2: + return None + else: + return None + return [int(block[0]), int(block[1])] + + +def _load_fp8_weight_map( + model_name, + local_files_only, + token, + revision = None, + subfolder = None, + cache_dir = None, +): + """Return the checkpoint's tensor->file map, using the same snapshot the load used. + + Prefers the sharded `model.safetensors.index.json`; falls back to a single `model.safetensors` + (every tensor maps to that one file) so unsharded checkpoints are covered too. + """ + + def _local_path(filename): + return ( + os.path.join(model_name, subfolder, filename) + if subfolder + else os.path.join(model_name, filename) + ) + + def _remote_path(filename): + from huggingface_hub import hf_hub_download + return hf_hub_download( + model_name, + filename, + revision = revision, + subfolder = subfolder, + cache_dir = cache_dir, + local_files_only = local_files_only, + token = token, + ) + + index_file = "model.safetensors.index.json" + single_file = "model.safetensors" + is_local = os.path.isdir(model_name) + + # Sharded checkpoint. + if is_local and os.path.exists(_local_path(index_file)): + index_path = _local_path(index_file) + elif not is_local: + try: + index_path = _remote_path(index_file) + except Exception: + index_path = None + else: + index_path = None + if index_path is not None: + import json + with open(index_path, "r") as f: + return json.load(f).get("weight_map", None) + + # Unsharded single file: map every tensor to it. + try: + if is_local and os.path.exists(_local_path(single_file)): + single_path = _local_path(single_file) + elif not is_local: + single_path = _remote_path(single_file) + else: + return None + from safetensors import safe_open + with safe_open(single_path, framework = "pt") as f: + return {key: single_file for key in f.keys()} + except Exception: + return None + + +def _resolve_fp8_shard( + model_name, + shard, + local_files_only, + token, + revision = None, + subfolder = None, + cache_dir = None, +): + """Resolve a checkpoint shard filename to a local path (repo id or local dir).""" + if os.path.isdir(model_name): + return ( + os.path.join(model_name, subfolder, shard) + if subfolder + else os.path.join(model_name, shard) + ) + from huggingface_hub import hf_hub_download + + return hf_hub_download( + model_name, + shard, + revision = revision, + subfolder = subfolder, + cache_dir = cache_dir, + local_files_only = local_files_only, + token = token, + ) + + +def _match_fp8_module(module_by_name, base): + """Resolve a checkpoint module name to a live module, allowing for VLM key remappings. + + VLM loads can name the text tower differently from the checkpoint keys: `text_only=True` + strips the `language_model.` wrapper (so `model.language_model.layers.*` -> `model.layers.*`), + and full VLM loads may expose `model.language_model.*` while the checkpoint stores + `language_model.model.*`. Try the raw key first, then a few safe remappings. + """ + if base in module_by_name: + return module_by_name[base] + candidates = [] + if "language_model." in base: + candidates.append(base.replace("language_model.", "", 1)) # text-only: drop wrapper + if "language_model.model." in base: + candidates.append(base.replace("language_model.model.", "model.language_model.", 1)) + if base.startswith("language_model."): + candidates.append("model." + base) # add model. prefix + for candidate in candidates: + if candidate in module_by_name: + return module_by_name[candidate] + return None + + +def _restore_dropped_fp8_scales( + model, + model_name, + *, + local_files_only = False, + token = None, + revision = None, + subfolder = None, + cache_dir = None, + variant = None, +): + """Re-apply block-fp8 `weight_scale_inv` tensors that transformers dropped on load. + + On some block-scale fp8 checkpoints (e.g. Qwen3.6-27B-FP8, issue #6200) transformers fails to + convert a Linear (such as `mlp.gate_proj`) to an fp8 module, loading the raw quantized values + into a plain bf16 weight and discarding its `weight_scale_inv` as an unexpected key. The weight + is then used un-scaled, producing a garbage model. For every checkpoint scale whose live weight + is not fp8, dequantize the orphaned weight in place. Modules that were converted correctly keep + an fp8 weight and are skipped, so a healthy checkpoint is a no-op. Returns (restored, skipped). + """ + try: + block = _fp8_block_size_from_config(model) + if block is None or not _FP8_DTYPES: + return (0, 0) + # A variant load reads variant-named files; skip to avoid applying default scales to them. + if variant: + return (0, 0) + # No fp8 params means the checkpoint was dequantized on purpose (e.g. load_in_16bit); + # re-applying a scale would corrupt those already-correct 16bit weights, so do nothing. + if not any(p.dtype in _FP8_DTYPES for p in model.parameters()): + return (0, 0) + weight_map = _load_fp8_weight_map( + model_name, local_files_only, token, revision, subfolder, cache_dir + ) + if not weight_map: + return (0, 0) + + scale_keys = {k: v for k, v in weight_map.items() if k.endswith(".weight_scale_inv")} + if not scale_keys: + return (0, 0) + + module_by_name = dict(model.named_modules()) + bs0, bs1 = block + restored = 0 + skipped = 0 + failed = 0 + offloaded = 0 + shard_cache = {} + for scale_key, shard in scale_keys.items(): + base = scale_key[: -len(".weight_scale_inv")] + module = _match_fp8_module(module_by_name, base) + if module is None: + continue + weight = getattr(module, "weight", None) + if not isinstance(weight, torch.Tensor) or weight.ndim != 2: + continue + if weight.device.type == "meta": + # Disk-offloaded layer: weight lives on meta until forward, so it cannot be + # scaled in place here. Count and warn rather than silently leave it unscaled. + offloaded += 1 + continue + if weight.dtype in _FP8_DTYPES: + # Correctly converted fp8 module: the fp8 path already handles the scale. + skipped += 1 + continue + + # Errors after this point are per-tensor: warn and continue, never abort or hide them. + try: + if shard not in shard_cache: + from safetensors import safe_open + shard_path = _resolve_fp8_shard( + model_name, + shard, + local_files_only, + token, + revision, + subfolder, + cache_dir, + ) + shard_cache[shard] = safe_open(shard_path, framework = "pt") + scale = shard_cache[shard].get_tensor(scale_key).to(torch.float32) + + out_features, in_features = weight.shape + out_blocks = (out_features + bs0 - 1) // bs0 + in_blocks = (in_features + bs1 - 1) // bs1 + if tuple(scale.shape) == (out_blocks, in_blocks): + pass + elif tuple(scale.shape) == (in_blocks, out_blocks) and out_blocks != in_blocks: + # Transposed block layout: same handling as the fp8 forward path. + scale = scale.t().contiguous() + else: + # Shape does not match the block grid: skip rather than apply a wrong scale. + continue + scale = scale.to(weight.device) + with torch.no_grad(): + if out_features % bs0 == 0 and in_features % bs1 == 0: + # Memory-frugal path: multiply block views in place against the broadcast + # fp32 scale, avoiding a full expanded scale and fp32 copy that could OOM. + # The in-place multiply promotes to fp32, matching the fallback exactly. + module.weight.data.view(out_blocks, bs0, in_blocks, bs1).mul_( + scale[:, None, :, None] + ) + else: + scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave( + bs1, dim = 1 + )[:out_features, :in_features] + module.weight.data = (weight.to(torch.float32) * scale_expanded).to( + weight.dtype + ) + restored += 1 + except Exception: + failed += 1 + continue + + if restored > 0: + print(f"Unsloth: Restored {restored} dropped FP8 weight_scale_inv tensor(s) on load") + if failed > 0: + print(f"Unsloth: {failed} dropped FP8 weight_scale_inv tensor(s) could not be restored") + if offloaded > 0: + print( + f"Unsloth: {offloaded} dropped FP8 weight_scale_inv tensor(s) skipped because the " + "layer is disk-offloaded; load without disk offload so the scales can be restored" + ) + return (restored, skipped) + except Exception: + return (0, 0) + + def check_and_disable_bitsandbytes_loading( model_config, load_in_4bit = True, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 0235d80e9..3b5af39e6 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -41,7 +41,11 @@ from ._utils import ( set_task_config_attr, ) from ._utils import * -from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings +from .loader_utils import ( + _exclude_rope_inv_freq_from_ddp, + _get_fp8_mode_and_check_settings, + _restore_dropped_fp8_scales, +) from ..save import patch_saving_functions from ..models.loader_utils import is_distributed from unsloth_zoo.gradient_checkpointing import ( @@ -1192,6 +1196,17 @@ class FastBaseModel: offload_embedding = offload_embedding, fast_inference = fast_inference, ) + # Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200). + _restore_dropped_fp8_scales( + model, + model_name, + local_files_only = local_files_only, + token = token, + revision = kwargs.get("revision"), + subfolder = kwargs.get("subfolder"), + cache_dir = kwargs.get("cache_dir"), + variant = kwargs.get("variant"), + ) if hasattr(model, "generate"): model.fast_generate = make_fast_generate_wrapper(model.generate) model.fast_generate_batches = error_out_no_vllm