vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Restore dropped FP8 weight_scale_inv tensors on load (#6978)

Browse source 
* Restore dropped FP8 weight_scale_inv tensors on load

Some block-scale FP8 checkpoints (for example Qwen3.6-27B-FP8, issue #6200) load
with transformers leaving an mlp.gate_proj as a plain bf16 Linear instead of an
fp8 module. Its raw quantized values are read into the bf16 weight and the
weight_scale_inv is dropped as an unexpected key, so the weight is used un-scaled
and the base model is garbage (perplexity around 2 million).

After load, for every checkpoint weight_scale_inv whose live weight is not fp8,
dequantize the orphaned weight in place using the block scale from the checkpoint
index. Modules that were converted correctly keep an fp8 weight and are skipped,
so healthy checkpoints and single-file checkpoints are a no-op.

Verified on Qwen3.6-27B-FP8: 64 gate_proj scales restored, perplexity 2028902 to
8.9. No-op on Qwen3-8B-FP8 (all scales already live).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden FP8 weight_scale_inv restore from review

- Skip restore when the model has no fp8 weights, so an intentionally
  dequantized load (load_in_16bit) is never re-scaled and corrupted.
- Thread revision, subfolder and cache_dir through the index and shard
  downloads so scales come from the same snapshot as the weights.
- Cover unsharded single-file model.safetensors checkpoints (no index).
- Handle transposed block-scale layouts and skip on a true grid mismatch
  instead of applying a wrong scale.
- Match text-only VLM loads where the language_model prefix was stripped.
- Restore on the FastLanguageModel text path too, not only vision.
- Handle a scalar weight_block_size; per-tensor error handling so one bad
  tensor cannot abort the rest or hide a partial mutation.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address second review round on FP8 scale restore

- Bound peak memory: dequantize block views in place with the fp32 scale
  broadcast instead of materializing a full expanded scale and fp32 copy,
  so a near-VRAM-limit load is not pushed into OOM by the repair.
- Restore on the sequence-classification load path too.
- Cover more VLM key remappings (language_model.model.* to
  model.language_model.*) when matching modules.
- Skip the restore for variant loads (variant=...) rather than risk
  applying default-checkpoint scales to variant weights.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Align FP8 scale restore revision with the loaded weights and warn on disk-offloaded layers

In llama.py the CausalLM/SequenceClassification weight loads resolve model_name on its
default branch (revision is not forwarded there), so read the dropped weight_scale_inv
tensors from the same default branch instead of the requested revision, avoiding rescaling
default-branch weights with scales from another revision.

In loader_utils.py a disk-offloaded layer keeps its weight on the meta device until the
offload hook materializes it, so the scale cannot be applied in place. Skip such layers
explicitly and print a warning rather than silently leaving them unscaled.

* Tighten comments in the FP8 scale restore path

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-09 06:44:44 -07:00 • committed by GitHub
parent fb5dc91bb4
commit d4fbc81d3a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 684 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

358
tests/test_fp8_restore_dropped_scale.py Normal file
View file

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

30
unsloth/models/llama.py
View file

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

281
unsloth/models/loader_utils.py
View file

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

17
unsloth/models/vision.py
View file

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