diff --git a/tests/unit/utilities/test_multi_gpu_unit.py b/tests/unit/utilities/test_multi_gpu_unit.py index 5924c7476..f6b180b87 100644 --- a/tests/unit/utilities/test_multi_gpu_unit.py +++ b/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 diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index 405185f8e..25fa5cb1e 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/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") diff --git a/transformer_lens/utilities/multi_gpu.py b/transformer_lens/utilities/multi_gpu.py index e29023aed..a4915088e 100644 --- a/transformer_lens/utilities/multi_gpu.py +++ b/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.