From 6d110468310183d430fabe698ee064f4745eebc0 Mon Sep 17 00:00:00 2001 From: "Md.Sadiq" Date: Fri, 21 Aug 2026 23:20:52 +0530 Subject: [PATCH 1/2] do not overwrite storage dtypes --- .../model_bridge/test_bridge_integration.py | 63 +++++++++++++++ tests/unit/utilities/test_multi_gpu_unit.py | 77 ++++++++++++++++++- .../model_bridge/sources/transformers.py | 7 +- transformer_lens/utilities/multi_gpu.py | 10 ++- 4 files changed, 154 insertions(+), 3 deletions(-) diff --git a/tests/integration/model_bridge/test_bridge_integration.py b/tests/integration/model_bridge/test_bridge_integration.py index 9ae02e4068..e61b24dc91 100644 --- a/tests/integration/model_bridge/test_bridge_integration.py +++ b/tests/integration/model_bridge/test_bridge_integration.py @@ -729,5 +729,68 @@ def test_TransformerBridge_gemma2_forward(): assert hasattr(bridge.rotary_emb, "hook_sin"), "rotary_emb should have hook_sin" +class TestCastFloatingParamsSkipsQuantizerOwnedScales: + """Regression tests for the MXFP4 bug: cast_floating_params_to_dtype must not + corrupt quantizer-owned FP8 scale tensors. + + The fix has two layers: + 1. cast_floating_params_to_dtype itself skips one-byte floats (FP8) + 2. The boot call site skips the cast entirely when model has quantization_config + """ + + def test_quantization_method_detects_mxfp4_config(self): + """Verify quantization_method correctly identifies MXFP4 configs.""" + from types import SimpleNamespace + + from transformer_lens.utilities.quantization import quantization_method + + config = SimpleNamespace(quantization_config=SimpleNamespace(quant_method="mxfp4")) + assert quantization_method(config) == "mxfp4" + + def test_quantization_method_returns_none_for_unquantized(self): + """Verify quantization_method returns None for unquantized models.""" + from types import SimpleNamespace + + from transformer_lens.utilities.quantization import quantization_method + + config = SimpleNamespace(quantization_config=None) + assert quantization_method(config) is None + + config = SimpleNamespace() # No quantization_config attr + assert quantization_method(config) is None + + def test_cast_skips_fp8_scales_even_if_called(self): + """Even if cast_floating_params_to_dtype is called on a mixed module, + FP8 scale tensors must be preserved. + """ + import torch.nn as nn + + from transformer_lens.utilities.multi_gpu import cast_floating_params_to_dtype + + class MixedQuantizedModule(nn.Module): + def __init__(self): + super().__init__() + # Packed int8 weights (untouched because not floating) + self.packed_weight = nn.Parameter( + torch.zeros(4, 4, dtype=torch.int8), requires_grad=False + ) + # FP8 scale (MUST be preserved) + self.fp8_scale = nn.Parameter( + torch.ones(4, dtype=torch.float8_e4m3fn), requires_grad=False + ) + # Normal float weight (should be cast) + self.normal_weight = nn.Parameter(torch.zeros(4, 4, dtype=torch.float32)) + + module = MixedQuantizedModule() + cast_floating_params_to_dtype(module, torch.bfloat16) + + # FP8 scale must NOT be cast + assert module.fp8_scale.dtype == torch.float8_e4m3fn + # Int8 packed weight must NOT be cast + assert module.packed_weight.dtype == torch.int8 + # Normal weight SHOULD be cast + assert module.normal_weight.dtype == torch.bfloat16 + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/unit/utilities/test_multi_gpu_unit.py b/tests/unit/utilities/test_multi_gpu_unit.py index 5924c74762..12e27dc53a 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,74 @@ 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. + + Issue: # — 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, + ], + ) + 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 diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index 405185f8e3..e3f3b0d003 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -847,7 +847,12 @@ 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.quantization import quantization_method + + if quantization_method(hf_model.config) is None: + cast_floating_params_to_dtype(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 e29023aedf..801576441a 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,6 +263,10 @@ 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) From 9fc650703ea769f6108f9000d0205b5335f6b00a Mon Sep 17 00:00:00 2001 From: "Md.Sadiq" Date: Sat, 22 Aug 2026 03:23:06 +0530 Subject: [PATCH 2/2] pipeline fix --- .../model_bridge/test_bridge_integration.py | 63 ------------------- tests/unit/utilities/test_multi_gpu_unit.py | 59 ++++++++++++++++- .../model_bridge/sources/transformers.py | 6 +- transformer_lens/utilities/multi_gpu.py | 15 +++++ 4 files changed, 73 insertions(+), 70 deletions(-) diff --git a/tests/integration/model_bridge/test_bridge_integration.py b/tests/integration/model_bridge/test_bridge_integration.py index e61b24dc91..9ae02e4068 100644 --- a/tests/integration/model_bridge/test_bridge_integration.py +++ b/tests/integration/model_bridge/test_bridge_integration.py @@ -729,68 +729,5 @@ def test_TransformerBridge_gemma2_forward(): assert hasattr(bridge.rotary_emb, "hook_sin"), "rotary_emb should have hook_sin" -class TestCastFloatingParamsSkipsQuantizerOwnedScales: - """Regression tests for the MXFP4 bug: cast_floating_params_to_dtype must not - corrupt quantizer-owned FP8 scale tensors. - - The fix has two layers: - 1. cast_floating_params_to_dtype itself skips one-byte floats (FP8) - 2. The boot call site skips the cast entirely when model has quantization_config - """ - - def test_quantization_method_detects_mxfp4_config(self): - """Verify quantization_method correctly identifies MXFP4 configs.""" - from types import SimpleNamespace - - from transformer_lens.utilities.quantization import quantization_method - - config = SimpleNamespace(quantization_config=SimpleNamespace(quant_method="mxfp4")) - assert quantization_method(config) == "mxfp4" - - def test_quantization_method_returns_none_for_unquantized(self): - """Verify quantization_method returns None for unquantized models.""" - from types import SimpleNamespace - - from transformer_lens.utilities.quantization import quantization_method - - config = SimpleNamespace(quantization_config=None) - assert quantization_method(config) is None - - config = SimpleNamespace() # No quantization_config attr - assert quantization_method(config) is None - - def test_cast_skips_fp8_scales_even_if_called(self): - """Even if cast_floating_params_to_dtype is called on a mixed module, - FP8 scale tensors must be preserved. - """ - import torch.nn as nn - - from transformer_lens.utilities.multi_gpu import cast_floating_params_to_dtype - - class MixedQuantizedModule(nn.Module): - def __init__(self): - super().__init__() - # Packed int8 weights (untouched because not floating) - self.packed_weight = nn.Parameter( - torch.zeros(4, 4, dtype=torch.int8), requires_grad=False - ) - # FP8 scale (MUST be preserved) - self.fp8_scale = nn.Parameter( - torch.ones(4, dtype=torch.float8_e4m3fn), requires_grad=False - ) - # Normal float weight (should be cast) - self.normal_weight = nn.Parameter(torch.zeros(4, 4, dtype=torch.float32)) - - module = MixedQuantizedModule() - cast_floating_params_to_dtype(module, torch.bfloat16) - - # FP8 scale must NOT be cast - assert module.fp8_scale.dtype == torch.float8_e4m3fn - # Int8 packed weight must NOT be cast - assert module.packed_weight.dtype == torch.int8 - # Normal weight SHOULD be cast - assert module.normal_weight.dtype == torch.bfloat16 - - if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/unit/utilities/test_multi_gpu_unit.py b/tests/unit/utilities/test_multi_gpu_unit.py index 12e27dc53a..f6b180b879 100644 --- a/tests/unit/utilities/test_multi_gpu_unit.py +++ b/tests/unit/utilities/test_multi_gpu_unit.py @@ -138,9 +138,10 @@ def test_cpu_device_is_returned_unchanged(self): class TestCastFloatingParamsToDtype: """Regression tests for cast_floating_params_to_dtype. - Issue: # — the function was casting quantizer-owned FP8 scale tensors - (float8_e8m0fnu) to bfloat16, which corrupts the weight/scale pair relationship - and breaks MXFP4 checkpoints. + 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): @@ -171,6 +172,12 @@ def test_skips_non_floating_point_params(self): [ 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): @@ -204,3 +211,49 @@ def __init__(self): 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 e3f3b0d003..25fa5cb1e6 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, @@ -849,10 +848,9 @@ def boot( # parameters before we touch them. # 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.quantization import quantization_method + from transformer_lens.utilities.multi_gpu import maybe_cast_floating_params - if quantization_method(hf_model.config) is None: - cast_floating_params_to_dtype(hf_model, dtype) + 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 801576441a..a4915088e6 100644 --- a/transformer_lens/utilities/multi_gpu.py +++ b/transformer_lens/utilities/multi_gpu.py @@ -270,6 +270,21 @@ def cast_floating_params_to_dtype(model: nn.Module, dtype: torch.dtype) -> None: 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.