do not overwrite storage dtypes

tests/unit/utilities/test_multi_gpu_unit.py

@@ -5,13 +5,17 @@
 
 import pytest
 import torch
+import torch.nn as nn
 
 from transformer_lens.utilities import (
     calculate_available_device_cuda_memory,
     determine_available_memory_for_available_devices,
     sort_devices_based_on_available_memory,
 )
-from transformer_lens.utilities.multi_gpu import get_device_for_block_index
+from transformer_lens.utilities.multi_gpu import (
+    cast_floating_params_to_dtype,
+    get_device_for_block_index,
+)
 
 
 def mock_available_devices(memory_stats: list[tuple[int, int]]):
@@ -129,3 +133,127 @@ def test_cpu_device_is_returned_unchanged(self):
         cfg = _cuda_cfg(n_layers=62, n_devices=8)
         result = get_device_for_block_index(30, cfg, device="cpu")
         assert result.type == "cpu"
+
+
+class TestCastFloatingParamsToDtype:
+    """Regression tests for cast_floating_params_to_dtype.
+
+    See: https://github.com/TransformerLensOrg/TransformerLens/issues/1713
+    The function was casting quantizer-owned FP8 scale tensors (float8_e8m0fnu)
+    to bfloat16, which corrupts the weight/scale pair relationship and breaks
+    MXFP4 checkpoints.
+    """
+
+    def test_casts_standard_floats_to_target_dtype(self):
+        """Positive control: standard float dtypes should be cast."""
+        model = nn.Linear(4, 4)
+        model.weight = nn.Parameter(torch.zeros(4, 4, dtype=torch.float32))
+        cast_floating_params_to_dtype(model, torch.bfloat16)
+        assert model.weight.dtype == torch.bfloat16
+
+    def test_skips_params_already_at_target_dtype(self):
+        """Params already at target dtype are left untouched."""
+        model = nn.Linear(4, 4)
+        original = torch.zeros(4, 4, dtype=torch.bfloat16)
+        model.weight = nn.Parameter(original)
+        cast_floating_params_to_dtype(model, torch.bfloat16)
+        assert model.weight.dtype == torch.bfloat16
+        assert model.weight.data_ptr() == original.data_ptr()
+
+    def test_skips_non_floating_point_params(self):
+        """Integer params (packed quantized weights) are left untouched."""
+        model = nn.Linear(4, 4, bias=False)
+        model.weight = nn.Parameter(torch.zeros(4, 4, dtype=torch.int8), requires_grad=False)
+        cast_floating_params_to_dtype(model, torch.bfloat16)
+        assert model.weight.dtype == torch.int8
+
+    @pytest.mark.parametrize(
+        "fp8_dtype",
+        [
+            torch.float8_e4m3fn,
+            torch.float8_e5m2,
+            pytest.param(
+                getattr(torch, "float8_e8m0fnu", None),
+                marks=pytest.mark.skipif(
+                    not hasattr(torch, "float8_e8m0fnu"), reason="torch < 2.7"
+                ),
+            ),
+        ],
+    )
+    def test_skips_one_byte_floats_fp8_scales(self, fp8_dtype):
+        """FP8 scale tensors must NOT be cast — they are quantizer-owned.
+
+        This is the key regression test for the MXFP4 bug: casting float8_e8m0fnu
+        scales to bfloat16 breaks the weight/scale pair relationship.
+        """
+        model = nn.Linear(4, 4, bias=False)
+        model.weight = nn.Parameter(torch.zeros(4, 4, dtype=fp8_dtype), requires_grad=False)
+        cast_floating_params_to_dtype(model, torch.bfloat16)
+        assert model.weight.dtype == fp8_dtype
+
+    def test_mixed_module_casts_selectively(self):
+        """A module with both standard and FP8 params: only standard params cast."""
+
+        class MixedModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.standard_weight = nn.Parameter(torch.zeros(4, 4, dtype=torch.float32))
+                self.fp8_scale = nn.Parameter(
+                    torch.zeros(4, 4, dtype=torch.float8_e4m3fn), requires_grad=False
+                )
+                self.packed_weight = nn.Parameter(
+                    torch.zeros(4, 4, dtype=torch.int8), requires_grad=False
+                )
+
+        model = MixedModule()
+        cast_floating_params_to_dtype(model, torch.bfloat16)
+
+        assert model.standard_weight.dtype == torch.bfloat16
+        assert model.fp8_scale.dtype == torch.float8_e4m3fn
+        assert model.packed_weight.dtype == torch.int8
+
+
+class TestMaybeCastFloatingParams:
+    """Tests for maybe_cast_floating_params helper.
+
+    See: https://github.com/TransformerLensOrg/TransformerLens/issues/1713
+    The helper wraps cast_floating_params_to_dtype with a quantization check,
+    skipping the cast entirely when the model has an active quantization_config.
+    """
+
+    def test_casts_unquantized_model(self):
+        """Unquantized models should have their params cast."""
+        from types import SimpleNamespace
+
+        from transformer_lens.utilities.multi_gpu import maybe_cast_floating_params
+
+        model = nn.Linear(4, 4)
+        model.weight = nn.Parameter(torch.zeros(4, 4, dtype=torch.float32))
+        model.config = SimpleNamespace(quantization_config=None)
+
+        maybe_cast_floating_params(model, torch.bfloat16)
+        assert model.weight.dtype == torch.bfloat16
+
+    def test_skips_quantized_model(self):
+        """Quantized models should NOT have their params cast."""
+        from types import SimpleNamespace
+
+        from transformer_lens.utilities.multi_gpu import maybe_cast_floating_params
+
+        model = nn.Linear(4, 4)
+        model.weight = nn.Parameter(torch.zeros(4, 4, dtype=torch.float32))
+        model.config = SimpleNamespace(quantization_config=SimpleNamespace(quant_method="mxfp4"))
+
+        maybe_cast_floating_params(model, torch.bfloat16)
+        assert model.weight.dtype == torch.float32  # NOT cast
+
+    def test_skips_model_without_config(self):
+        """Models without a config attribute should be cast (no quantization)."""
+        from transformer_lens.utilities.multi_gpu import maybe_cast_floating_params
+
+        model = nn.Linear(4, 4)
+        model.weight = nn.Parameter(torch.zeros(4, 4, dtype=torch.float32))
+        # No model.config attribute
+
+        maybe_cast_floating_params(model, torch.bfloat16)
+        assert model.weight.dtype == torch.bfloat16

transformer_lens/model_bridge/sources/transformers.py

@@ -750,7 +750,6 @@ def boot(
     # resolved values.
     from transformer_lens.utilities.multi_gpu import (
         MIXED_CPU_GPU_ERROR,
-        cast_floating_params_to_dtype,
         count_unique_devices,
         find_embedding_device,
         find_misplaced_modules,
@@ -847,7 +846,11 @@ def boot(
         # Cast params to dtype; preserve float32 buffers (e.g., RotaryEmbedding.inv_freq).
         # Use module-level alignment so Accelerate can temporarily materialize offloaded
         # parameters before we touch them.
-        cast_floating_params_to_dtype(hf_model, dtype)
+        # Skip dtype normalization entirely when model has an active quantizer: the
+        # quantizer owns specific dtypes (e.g., FP8 scales) that must not be overwritten.
+        from transformer_lens.utilities.multi_gpu import maybe_cast_floating_params
+
+        maybe_cast_floating_params(hf_model, dtype)
     # Derive cfg.device / cfg.n_devices from hf_device_map when present. This covers:
     #   - fresh loads with a resolved device_map (set above)
     #   - pre-loaded hf_model that the caller dispatched themselves (e.g., device_map="auto")

transformer_lens/utilities/multi_gpu.py

@@ -249,7 +249,11 @@ def is_mixed_cpu_gpu(values: Any) -> bool:
 
 
 def cast_floating_params_to_dtype(model: nn.Module, dtype: torch.dtype) -> None:
-    """Cast materialized floating parameters while preserving Accelerate offload hooks."""
+    """Cast materialized floating parameters while preserving Accelerate offload hooks.
+
+    Skips one-byte floats (FP8 dtypes like float8_e8m0fnu) which are quantizer-owned
+    scale parameters — casting them corrupts the quantization format.
+    """
     from accelerate.utils import align_module_device
 
     for module in model.modules():
@@ -259,9 +263,28 @@ def cast_floating_params_to_dtype(model: nn.Module, dtype: torch.dtype) -> None:
                     continue
                 if param.device.type == "meta":
                     continue
+                # Skip one-byte floats (FP8 scale tensors): they are quantizer-owned
+                # and casting them breaks the weight/scale pair relationship.
+                if param.dtype.itemsize < 2:
+                    continue
                 param.data = param.data.to(dtype=dtype)
 
 
+def maybe_cast_floating_params(model: nn.Module, dtype: torch.dtype) -> None:
+    """Cast floating params to dtype, skipping models with active quantization.
+
+    When a model has an active quantization_config, the quantizer owns specific
+    dtypes (e.g., FP8 scales) that must not be overwritten. This helper wraps
+    the cast with that check.
+
+    See: https://github.com/TransformerLensOrg/TransformerLens/issues/1713
+    """
+    from transformer_lens.utilities.quantization import quantization_method
+
+    if quantization_method(getattr(model, "config", None)) is None:
+        cast_floating_params_to_dtype(model, dtype)
+
+
 def find_embedding_device(hf_model: Any) -> Optional[torch.device]:
     """Return the device that input tokens should be placed on for a dispatched HF model.