From 0fd71835ac0f95cc6b3ebd3ea4fcf65fd88f463a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 8 Jul 2026 12:31:10 +0000
Subject: [PATCH 1/6] 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).
---
 tests/test_fp8_restore_dropped_scale.py | 176 ++++++++++++++++++++++++
 unsloth/models/loader_utils.py          | 115 ++++++++++++++++
 unsloth/models/vision.py                |   7 +-
 3 files changed, 297 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_fp8_restore_dropped_scale.py

diff --git a/tests/test_fp8_restore_dropped_scale.py b/tests/test_fp8_restore_dropped_scale.py
new file mode 100644
index 00000000000..c43fa91f721
--- /dev/null
+++ b/tests/test_fp8_restore_dropped_scale.py
@@ -0,0 +1,176 @@
+# 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
+
+
+_SHARD = "model-00001-of-00001.safetensors"
+
+
+def _write_checkpoint(path, tensors, include_index = True):
+    save_file(tensors, os.path.join(path, _SHARD))
+    if include_index:
+        weight_map = {name: _SHARD 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 _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."""
+    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.layer = nn.Linear(4, 4, bias = False).to(torch.bfloat16)
+    with torch.no_grad():
+        model.layer.weight.copy_(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 and skipped == 0
+    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."""
+    fp8_dtype = getattr(torch, "float8_e4m3fn", None)
+    if fp8_dtype is None:
+        return
+    weight = torch.randn(4, 4).to(fp8_dtype)
+    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_non_block_divisible_shape():
+    """Block scale is expanded then sliced to a non-divisible weight shape."""
+    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.layer = nn.Linear(4, 3, bias = False).to(torch.bfloat16)  # weight shape [3, 4]
+    with torch.no_grad():
+        model.layer.weight.copy_(raw)
+
+    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_noop_without_scale_keys():
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.layer = nn.Linear(4, 4, bias = False)
+    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():
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.layer = nn.Linear(4, 4, bias = False)
+    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)
+
+
+def test_nested_submodule_restored():
+    """Nested module names (language_model.layers.N...) resolve via get_submodule."""
+    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.language_model = nn.Module()
+    model.language_model.gate_proj = nn.Linear(2, 2, bias = False).to(torch.bfloat16)
+    with torch.no_grad():
+        model.language_model.gate_proj.weight.copy_(raw)
+
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"language_model.gate_proj.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), (2, 2))).to(torch.bfloat16)
+    assert torch.equal(model.language_model.gate_proj.weight.data, expected)
diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py
index d6d6ce877e0..3e1c37e3367 100644
--- a/unsloth/models/loader_utils.py
+++ b/unsloth/models/loader_utils.py
@@ -362,6 +362,121 @@ 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 len(block) == 1:
+        block = [block[0], block[0]]
+    return [int(block[0]), int(block[1])]
+
+
+def _load_fp8_weight_map(model_name, local_files_only, token):
+    """Return the safetensors index weight_map for a checkpoint, or None if there is no index."""
+    index_file = "model.safetensors.index.json"
+    if os.path.isdir(model_name):
+        index_path = os.path.join(model_name, index_file)
+        if not os.path.exists(index_path):
+            return None
+    else:
+        from huggingface_hub import hf_hub_download
+        try:
+            index_path = hf_hub_download(
+                model_name, index_file,
+                local_files_only = local_files_only, token = token,
+            )
+        except Exception:
+            return None
+    import json
+    with open(index_path, "r") as f:
+        return json.load(f).get("weight_map", None)
+
+
+def _resolve_fp8_shard(model_name, shard, local_files_only, token):
+    """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, shard)
+    from huggingface_hub import hf_hub_download
+    return hf_hub_download(model_name, shard, local_files_only = local_files_only, token = token)
+
+
+def _restore_dropped_fp8_scales(model, model_name, *, local_files_only = False, token = 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)
+        weight_map = _load_fp8_weight_map(model_name, local_files_only, token)
+        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)
+
+        bs0, bs1 = block
+        restored = 0
+        skipped = 0
+        shard_cache = {}
+        for scale_key, shard in scale_keys.items():
+            module_name = scale_key[: -len(".weight_scale_inv")]
+            try:
+                module = model.get_submodule(module_name)
+            except AttributeError:
+                continue
+            weight = getattr(module, "weight", None)
+            if not isinstance(weight, torch.Tensor) or weight.device.type == "meta":
+                continue
+            if weight.dtype in _FP8_DTYPES:
+                # Correctly converted fp8 module: the scale is handled by the fp8 path already.
+                skipped += 1
+                continue
+
+            if shard not in shard_cache:
+                from safetensors import safe_open
+                shard_path = _resolve_fp8_shard(model_name, shard, local_files_only, token)
+                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
+            scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
+            scale_expanded = scale_expanded[:out_features, :in_features].to(weight.device)
+            with torch.no_grad():
+                module.weight.data = (weight.to(torch.float32) * scale_expanded).to(weight.dtype)
+            restored += 1
+
+        if restored > 0:
+            print(f"Unsloth: Restored {restored} dropped FP8 weight_scale_inv tensor(s) on load")
+        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 0a68a49fee6..ff3d447d3ee 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -41,7 +41,7 @@
     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 +1192,11 @@ def from_pretrained(
                     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,
+                )
                 if hasattr(model, "generate"):
                     model.fast_generate = make_fast_generate_wrapper(model.generate)
                     model.fast_generate_batches = error_out_no_vllm

From 23c998de8b48fbfe57119c9b8cc9013a0f2de3b0 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 8 Jul 2026 12:45:29 +0000
Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_fp8_restore_dropped_scale.py | 21 +++++++++++++--------
 unsloth/models/loader_utils.py          | 19 +++++++++++++++----
 unsloth/models/vision.py                | 12 +++++++++---
 3 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/tests/test_fp8_restore_dropped_scale.py b/tests/test_fp8_restore_dropped_scale.py
index c43fa91f721..a44d11eb60e 100644
--- a/tests/test_fp8_restore_dropped_scale.py
+++ b/tests/test_fp8_restore_dropped_scale.py
@@ -37,7 +37,11 @@
 _SHARD = "model-00001-of-00001.safetensors"
 
 
-def _write_checkpoint(path, tensors, include_index = True):
+def _write_checkpoint(
+    path,
+    tensors,
+    include_index = True,
+):
     save_file(tensors, os.path.join(path, _SHARD))
     if include_index:
         weight_map = {name: _SHARD for name in tensors}
@@ -73,10 +77,13 @@ def test_restore_dequantizes_orphaned_scale():
         model.layer.weight.copy_(raw)
 
     with tempfile.TemporaryDirectory() as d:
-        _write_checkpoint(d, {
-            "layer.weight": raw.to(torch.float32),
-            "layer.weight_scale_inv": scale,
-        })
+        _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 and skipped == 0
@@ -146,9 +153,7 @@ 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.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})
diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py
index 3e1c37e3367..e3ed3a6f1f9 100644
--- a/unsloth/models/loader_utils.py
+++ b/unsloth/models/loader_utils.py
@@ -363,7 +363,8 @@ def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
 
 
 _FP8_DTYPES = tuple(
-    dtype for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
+    dtype
+    for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
     if dtype is not None
 )
 
@@ -399,12 +400,15 @@ def _load_fp8_weight_map(model_name, local_files_only, token):
         from huggingface_hub import hf_hub_download
         try:
             index_path = hf_hub_download(
-                model_name, index_file,
-                local_files_only = local_files_only, token = token,
+                model_name,
+                index_file,
+                local_files_only = local_files_only,
+                token = token,
             )
         except Exception:
             return None
     import json
+
     with open(index_path, "r") as f:
         return json.load(f).get("weight_map", None)
 
@@ -414,10 +418,17 @@ def _resolve_fp8_shard(model_name, shard, local_files_only, token):
     if os.path.isdir(model_name):
         return os.path.join(model_name, shard)
     from huggingface_hub import hf_hub_download
+
     return hf_hub_download(model_name, shard, local_files_only = local_files_only, token = token)
 
 
-def _restore_dropped_fp8_scales(model, model_name, *, local_files_only = False, token = None):
+def _restore_dropped_fp8_scales(
+    model,
+    model_name,
+    *,
+    local_files_only = False,
+    token = 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
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index ff3d447d3ee..e8cb6732019 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -41,7 +41,11 @@
     set_task_config_attr,
 )
 from ._utils import *
-from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings, _restore_dropped_fp8_scales
+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 (
@@ -1194,8 +1198,10 @@ def from_pretrained(
                 )
                 # 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,
+                    model,
+                    model_name,
+                    local_files_only = local_files_only,
+                    token = token,
                 )
                 if hasattr(model, "generate"):
                     model.fast_generate = make_fast_generate_wrapper(model.generate)

From 8bd6f33bc292b13b8f55f6722cc220132ced4293 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 8 Jul 2026 13:40:29 +0000
Subject: [PATCH 3/6] 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.
---
 tests/test_fp8_restore_dropped_scale.py | 195 ++++++++++++++++++------
 unsloth/models/llama.py                 |  10 +-
 unsloth/models/loader_utils.py          | 171 +++++++++++++++------
 unsloth/models/vision.py                |  15 +-
 4 files changed, 286 insertions(+), 105 deletions(-)

diff --git a/tests/test_fp8_restore_dropped_scale.py b/tests/test_fp8_restore_dropped_scale.py
index a44d11eb60e..f7ae948fcab 100644
--- a/tests/test_fp8_restore_dropped_scale.py
+++ b/tests/test_fp8_restore_dropped_scale.py
@@ -31,20 +31,17 @@
 
 # Import unsloth first to set UNSLOTH_IS_PRESENT env var.
 import unsloth
-from unsloth.models.loader_utils import _restore_dropped_fp8_scales
+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,
-    include_index = True,
-):
-    save_file(tensors, os.path.join(path, _SHARD))
+def _write_checkpoint(path, tensors, filename = _SHARD, include_index = True):
+    save_file(tensors, os.path.join(path, filename))
     if include_index:
-        weight_map = {name: _SHARD for name in tensors}
+        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)
 
@@ -58,6 +55,20 @@ def _fp8_config(block = (2, 2)):
     )
 
 
+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)
@@ -66,37 +77,34 @@ def _expand(scale, block, shape):
 
 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.layer = nn.Linear(4, 4, bias = False).to(torch.bfloat16)
-    with torch.no_grad():
-        model.layer.weight.copy_(raw)
+    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,
-            },
-        )
+        _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 and skipped == 0
+    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."""
-    fp8_dtype = getattr(torch, "float8_e4m3fn", None)
-    if fp8_dtype is None:
+    if _FP8 is None:
         return
-    weight = torch.randn(4, 4).to(fp8_dtype)
+    weight = torch.randn(4, 4).to(_FP8)
     before = weight.clone()
 
     model = nn.Module()
@@ -112,16 +120,34 @@ def test_skips_already_fp8_weight():
     assert torch.equal(model.layer.weight.data.float(), before.float())
 
 
+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.layer = nn.Linear(4, 3, bias = False).to(torch.bfloat16)  # weight shape [3, 4]
-    with torch.no_grad():
-        model.layer.weight.copy_(raw)
+    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})
@@ -132,50 +158,125 @@ def test_non_block_divisible_shape():
     assert torch.equal(model.layer.weight.data, expected)
 
 
-def test_noop_without_scale_keys():
+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.layer = nn.Linear(4, 4, bias = False)
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(4, 2, raw)
+
     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)
+        _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
 
-def test_noop_without_index():
     model = nn.Module()
     model.config = _fp8_config((2, 2))
-    model.layer = nn.Linear(4, 4, bias = False)
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(4, 4, raw)
+
     with tempfile.TemporaryDirectory() as d:
-        assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
+        _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
 
-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)
+    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})
-        assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
+        restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert restored == 1
 
 
-def test_nested_submodule_restored():
-    """Nested module names (language_model.layers.N...) resolve via get_submodule."""
+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.language_model = nn.Module()
-    model.language_model.gate_proj = nn.Linear(2, 2, bias = False).to(torch.bfloat16)
-    with torch.no_grad():
-        model.language_model.gate_proj.weight.copy_(raw)
+    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:
-        _write_checkpoint(d, {"language_model.gate_proj.weight_scale_inv": scale})
-        restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+        # 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.language_model.gate_proj.weight.data, expected)
+    assert torch.equal(model.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 a1da0997589..3a81b630a8a 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -28,7 +28,7 @@
     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,
@@ -2697,6 +2697,14 @@ def from_pretrained(
                     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,
+                    revision = revision,
+                    subfolder = kwargs.get("subfolder"),
+                    cache_dir = kwargs.get("cache_dir"),
+                )
                 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 e3ed3a6f1f9..4cb10a585bd 100644
--- a/unsloth/models/loader_utils.py
+++ b/unsloth/models/loader_utils.py
@@ -363,8 +363,7 @@ def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
 
 
 _FP8_DTYPES = tuple(
-    dtype
-    for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
+    dtype for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
     if dtype is not None
 )
 
@@ -384,50 +383,101 @@ def _fp8_block_size_from_config(model):
     block = quant.get("weight_block_size")
     if not block:
         return None
-    if len(block) == 1:
-        block = [block[0], block[0]]
+    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):
-    """Return the safetensors index weight_map for a checkpoint, or None if there is no index."""
-    index_file = "model.safetensors.index.json"
-    if os.path.isdir(model_name):
-        index_path = os.path.join(model_name, index_file)
-        if not os.path.exists(index_path):
-            return None
-    else:
+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 = hf_hub_download(
-                model_name,
-                index_file,
-                local_files_only = local_files_only,
-                token = token,
-            )
+            index_path = _remote_path(index_file)
         except Exception:
-            return None
-    import json
+            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)
 
-    with open(index_path, "r") as f:
-        return json.load(f).get("weight_map", None)
+    # Unsharded single-file checkpoint: map every tensor to model.safetensors.
+    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):
+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, shard)
+        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,
+    )
 
-    return hf_hub_download(model_name, shard, 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 text-only key mappings.
+
+    VLM `text_only=True` loads strip the `language_model.` wrapper from checkpoint keys before
+    building the text decoder, so an index key like `model.language_model.layers.0.mlp.gate_proj`
+    must also match the remapped `model.layers.0.mlp.gate_proj` module.
+    """
+    if base in module_by_name:
+        return module_by_name[base]
+    if "language_model." in base:
+        candidate = base.replace("language_model.", "", 1)
+        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,
+    model, model_name, *,
+    local_files_only = False, token = None,
+    revision = None, subfolder = None, cache_dir = None,
 ):
     """Re-apply block-fp8 `weight_scale_inv` tensors that transformers dropped on load.
 
@@ -442,7 +492,12 @@ def _restore_dropped_fp8_scales(
         block = _fp8_block_size_from_config(model)
         if block is None or not _FP8_DTYPES:
             return (0, 0)
-        weight_map = _load_fp8_weight_map(model_name, local_files_only, token)
+        # A genuine fp8 load leaves most weights fp8. If nothing is fp8 the checkpoint was
+        # dequantized on purpose (e.g. load_in_16bit -> HF quantizer dequantize); 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)
 
@@ -450,39 +505,59 @@ def _restore_dropped_fp8_scales(
         if not scale_keys:
             return (0, 0)
 
+        module_by_name = dict(model.named_modules())
         bs0, bs1 = block
         restored = 0
         skipped = 0
+        failed = 0
         shard_cache = {}
         for scale_key, shard in scale_keys.items():
-            module_name = scale_key[: -len(".weight_scale_inv")]
-            try:
-                module = model.get_submodule(module_name)
-            except AttributeError:
+            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.device.type == "meta":
+            if not isinstance(weight, torch.Tensor) or weight.device.type == "meta" or weight.ndim != 2:
                 continue
             if weight.dtype in _FP8_DTYPES:
                 # Correctly converted fp8 module: the scale is handled by the fp8 path already.
                 skipped += 1
                 continue
 
-            if shard not in shard_cache:
-                from safetensors import safe_open
-                shard_path = _resolve_fp8_shard(model_name, shard, local_files_only, token)
-                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
-            scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
-            scale_expanded = scale_expanded[:out_features, :in_features].to(weight.device)
-            with torch.no_grad():
-                module.weight.data = (weight.to(torch.float32) * scale_expanded).to(weight.dtype)
-            restored += 1
+            # 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_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
+                scale_expanded = scale_expanded[:out_features, :in_features].to(weight.device)
+                with torch.no_grad():
+                    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")
         return (restored, skipped)
     except Exception:
         return (0, 0)
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index e8cb6732019..be7f1deb2bd 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -41,11 +41,7 @@
     set_task_config_attr,
 )
 from ._utils import *
-from .loader_utils import (
-    _exclude_rope_inv_freq_from_ddp,
-    _get_fp8_mode_and_check_settings,
-    _restore_dropped_fp8_scales,
-)
+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 (
@@ -1198,10 +1194,11 @@ def from_pretrained(
                 )
                 # 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,
+                    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"),
                 )
                 if hasattr(model, "generate"):
                     model.fast_generate = make_fast_generate_wrapper(model.generate)

From e4449e65cc739508d608b41ba4be0a1429e0e4c1 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 8 Jul 2026 13:41:02 +0000
Subject: [PATCH 4/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_fp8_restore_dropped_scale.py | 27 +++++---
 unsloth/models/llama.py                 | 12 +++-
 unsloth/models/loader_utils.py          | 89 ++++++++++++++++++++-----
 unsloth/models/vision.py                | 12 +++-
 4 files changed, 106 insertions(+), 34 deletions(-)

diff --git a/tests/test_fp8_restore_dropped_scale.py b/tests/test_fp8_restore_dropped_scale.py
index f7ae948fcab..bef910a964d 100644
--- a/tests/test_fp8_restore_dropped_scale.py
+++ b/tests/test_fp8_restore_dropped_scale.py
@@ -38,7 +38,12 @@
 _FP8 = _FP8_DTYPES[0] if _FP8_DTYPES else None
 
 
-def _write_checkpoint(path, tensors, filename = _SHARD, include_index = True):
+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}
@@ -89,10 +94,13 @@ def test_restore_dequantizes_orphaned_scale():
     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,
-        })
+        _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
@@ -193,8 +201,9 @@ def test_single_file_checkpoint_without_index():
     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)
+        _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
@@ -273,9 +282,7 @@ 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.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})
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 3a81b630a8a..2be18d59cb2 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -28,7 +28,11 @@
     is_bfloat16_supported,
     get_quant_type,
 )
-from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings, _restore_dropped_fp8_scales
+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,
@@ -2699,8 +2703,10 @@ def from_pretrained(
                 )
                 # 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,
+                    model,
+                    model_name,
+                    local_files_only = kwargs.get("local_files_only", False),
+                    token = token,
                     revision = revision,
                     subfolder = kwargs.get("subfolder"),
                     cache_dir = kwargs.get("cache_dir"),
diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py
index 4cb10a585bd..fd66949fc96 100644
--- a/unsloth/models/loader_utils.py
+++ b/unsloth/models/loader_utils.py
@@ -363,7 +363,8 @@ def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
 
 
 _FP8_DTYPES = tuple(
-    dtype for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
+    dtype
+    for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
     if dtype is not None
 )
 
@@ -395,21 +396,37 @@ def _fp8_block_size_from_config(model):
     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):
+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)
+        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,
+            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"
@@ -446,15 +463,32 @@ def _remote_path(filename):
         return None
 
 
-def _resolve_fp8_shard(model_name, shard, local_files_only, token, revision = None, subfolder = None, cache_dir = 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)
+        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,
+        model_name,
+        shard,
+        revision = revision,
+        subfolder = subfolder,
+        cache_dir = cache_dir,
+        local_files_only = local_files_only,
+        token = token,
     )
 
 
@@ -475,9 +509,14 @@ def _match_fp8_module(module_by_name, base):
 
 
 def _restore_dropped_fp8_scales(
-    model, model_name, *,
-    local_files_only = False, token = None,
-    revision = None, subfolder = None, cache_dir = None,
+    model,
+    model_name,
+    *,
+    local_files_only = False,
+    token = None,
+    revision = None,
+    subfolder = None,
+    cache_dir = None,
 ):
     """Re-apply block-fp8 `weight_scale_inv` tensors that transformers dropped on load.
 
@@ -497,7 +536,9 @@ def _restore_dropped_fp8_scales(
         # 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)
+        weight_map = _load_fp8_weight_map(
+            model_name, local_files_only, token, revision, subfolder, cache_dir
+        )
         if not weight_map:
             return (0, 0)
 
@@ -517,7 +558,11 @@ def _restore_dropped_fp8_scales(
             if module is None:
                 continue
             weight = getattr(module, "weight", None)
-            if not isinstance(weight, torch.Tensor) or weight.device.type == "meta" or weight.ndim != 2:
+            if (
+                not isinstance(weight, torch.Tensor)
+                or weight.device.type == "meta"
+                or weight.ndim != 2
+            ):
                 continue
             if weight.dtype in _FP8_DTYPES:
                 # Correctly converted fp8 module: the scale is handled by the fp8 path already.
@@ -529,7 +574,13 @@ def _restore_dropped_fp8_scales(
                 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,
+                        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)
@@ -548,7 +599,9 @@ def _restore_dropped_fp8_scales(
                 scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
                 scale_expanded = scale_expanded[:out_features, :in_features].to(weight.device)
                 with torch.no_grad():
-                    module.weight.data = (weight.to(torch.float32) * scale_expanded).to(weight.dtype)
+                    module.weight.data = (weight.to(torch.float32) * scale_expanded).to(
+                        weight.dtype
+                    )
                 restored += 1
             except Exception:
                 failed += 1
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index be7f1deb2bd..1f77e7f9e93 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -41,7 +41,11 @@
     set_task_config_attr,
 )
 from ._utils import *
-from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings, _restore_dropped_fp8_scales
+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 (
@@ -1194,8 +1198,10 @@ def from_pretrained(
                 )
                 # 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,
+                    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"),

From 98bdaa81c75b27129e8f68451050bc560e68a86c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 8 Jul 2026 14:07:18 +0000
Subject: [PATCH 5/6] 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.
---
 tests/test_fp8_restore_dropped_scale.py | 37 +++++++++++++++++++++
 unsloth/models/llama.py                 | 10 ++++++
 unsloth/models/loader_utils.py          | 43 +++++++++++++++++++------
 unsloth/models/vision.py                |  1 +
 4 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/tests/test_fp8_restore_dropped_scale.py b/tests/test_fp8_restore_dropped_scale.py
index bef910a964d..1b26a27425d 100644
--- a/tests/test_fp8_restore_dropped_scale.py
+++ b/tests/test_fp8_restore_dropped_scale.py
@@ -255,6 +255,43 @@ def test_text_only_prefix_mapping():
     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
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 2be18d59cb2..22b0715a49e 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2663,6 +2663,15 @@ def from_pretrained(
                     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,
+                    revision = revision,
+                    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
@@ -2710,6 +2719,7 @@ def from_pretrained(
                     revision = revision,
                     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
diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py
index fd66949fc96..cdc6aecf0be 100644
--- a/unsloth/models/loader_utils.py
+++ b/unsloth/models/loader_utils.py
@@ -493,16 +493,23 @@ def _resolve_fp8_shard(
 
 
 def _match_fp8_module(module_by_name, base):
-    """Resolve a checkpoint module name to a live module, allowing for text-only key mappings.
+    """Resolve a checkpoint module name to a live module, allowing for VLM key remappings.
 
-    VLM `text_only=True` loads strip the `language_model.` wrapper from checkpoint keys before
-    building the text decoder, so an index key like `model.language_model.layers.0.mlp.gate_proj`
-    must also match the remapped `model.layers.0.mlp.gate_proj` module.
+    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:
-        candidate = base.replace("language_model.", "", 1)
+        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
@@ -517,6 +524,7 @@ def _restore_dropped_fp8_scales(
     revision = None,
     subfolder = None,
     cache_dir = None,
+    variant = None,
 ):
     """Re-apply block-fp8 `weight_scale_inv` tensors that transformers dropped on load.
 
@@ -531,6 +539,10 @@ def _restore_dropped_fp8_scales(
         block = _fp8_block_size_from_config(model)
         if block is None or not _FP8_DTYPES:
             return (0, 0)
+        # A variant load (variant="fp8") reads variant-named files; do not risk applying default
+        # checkpoint scales to variant weights. Skip rather than resolve the wrong files.
+        if variant:
+            return (0, 0)
         # A genuine fp8 load leaves most weights fp8. If nothing is fp8 the checkpoint was
         # dequantized on purpose (e.g. load_in_16bit -> HF quantizer dequantize); re-applying
         # a scale would corrupt those already-correct 16bit weights, so do nothing.
@@ -596,12 +608,23 @@ def _restore_dropped_fp8_scales(
                 else:
                     # Shape does not match the block grid: skip rather than apply a wrong scale.
                     continue
-                scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
-                scale_expanded = scale_expanded[:out_features, :in_features].to(weight.device)
+                scale = scale.to(weight.device)
                 with torch.no_grad():
-                    module.weight.data = (weight.to(torch.float32) * scale_expanded).to(
-                        weight.dtype
-                    )
+                    if out_features % bs0 == 0 and in_features % bs1 == 0:
+                        # Memory-frugal path: multiply block views in place with the small fp32
+                        # scale broadcast, avoiding a full expanded scale and full fp32 copy so a
+                        # near-VRAM-limit load is not pushed into OOM by the repair. The in-place
+                        # multiply promotes to fp32 for the compute, 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
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 1f77e7f9e93..af79f79b278 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -1205,6 +1205,7 @@ def from_pretrained(
                     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)

From a11fe00c524f2b8897c9fbbe152e5a4672e25056 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 8 Jul 2026 14:07:59 +0000
Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_fp8_restore_dropped_scale.py | 4 +++-
 unsloth/models/llama.py                 | 6 ++++--
 unsloth/models/loader_utils.py          | 4 ++--
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/tests/test_fp8_restore_dropped_scale.py b/tests/test_fp8_restore_dropped_scale.py
index 1b26a27425d..378499413ab 100644
--- a/tests/test_fp8_restore_dropped_scale.py
+++ b/tests/test_fp8_restore_dropped_scale.py
@@ -283,7 +283,9 @@ def test_vlm_language_model_model_alias():
     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
+    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)
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 22b0715a49e..f8c7359e010 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -2665,8 +2665,10 @@ def from_pretrained(
                 )
                 # 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,
+                    model,
+                    model_name,
+                    local_files_only = kwargs.get("local_files_only", False),
+                    token = token,
                     revision = revision,
                     subfolder = kwargs.get("subfolder"),
                     cache_dir = kwargs.get("cache_dir"),
diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py
index cdc6aecf0be..00be4512e5b 100644
--- a/unsloth/models/loader_utils.py
+++ b/unsloth/models/loader_utils.py
@@ -504,11 +504,11 @@ def _match_fp8_module(module_by_name, base):
         return module_by_name[base]
     candidates = []
     if "language_model." in base:
-        candidates.append(base.replace("language_model.", "", 1))          # text-only: drop wrapper
+        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
+        candidates.append("model." + base)  # add model. prefix
     for candidate in candidates:
         if candidate in module_by_name:
             return module_by_name[candidate]