Fix ignore_layers not working for FP8 models

auto_round/compressors/utils.py

@@ -864,14 +864,14 @@ def get_fp_layer_names(model: torch.nn.Module, ignore_layers: str):
         list: A list of layer names that match the specified FP layers or are
         subcomponents of those layers.
     """
-    from auto_round.utils import SUPPORTED_LAYER_TYPES
+    from auto_round.utils import INNER_SUPPORTED_LAYER_TYPES, SUPPORTED_LAYER_TYPES
 
     if not ignore_layers:
         return []
     ignore_layers = ignore_layers.replace(" ", "").split(",")
     all_layer_names = []
     for n, m in model.named_modules():
-        if type(m) in SUPPORTED_LAYER_TYPES:
+        if type(m) in SUPPORTED_LAYER_TYPES or m.__class__.__name__ in INNER_SUPPORTED_LAYER_TYPES:
             all_layer_names.append(n)
     not_to_quantized_layers = []
 

test/test_cpu/utils/test_compressor_utils.py

@@ -0,0 +1,111 @@
+"""
+Unit tests for auto_round.compressors.utils module.
+"""
+import torch
+import pytest
+
+from auto_round.compressors.utils import get_fp_layer_names
+from auto_round.utils import INNER_SUPPORTED_LAYER_TYPES
+
+
+class TestGetFpLayerNames:
+    """Test suite for get_fp_layer_names function."""
+
+    def test_regular_linear_layers(self):
+        """Test with regular Linear layers."""
+        class MockModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer1 = torch.nn.Linear(10, 10)
+                self.layer2 = torch.nn.Linear(10, 10)
+                self.mlp = torch.nn.Sequential(
+                    torch.nn.Linear(10, 10),
+                    torch.nn.Linear(10, 10)
+                )
+
+        model = MockModel()
+
+        # Test finding specific layer
+        result = get_fp_layer_names(model, 'layer1')
+        assert 'layer1' in result, "Should find layer1"
+
+        # Test finding layers with pattern
+        result = get_fp_layer_names(model, 'mlp')
+        assert len(result) == 2, "Should find 2 layers in mlp"
+        assert 'mlp.0' in result and 'mlp.1' in result
+
+    def test_fp8linear_layers(self):
+        """Test with FP8Linear layers (mocked by creating a proper class)."""
+        # Create a proper mock FP8Linear class
+        class FP8Linear(torch.nn.Linear):
+            """Mock FP8Linear class for testing."""
+            def __init__(self, in_features, out_features):
+                super().__init__(in_features, out_features)
+
+        class MockModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer1 = torch.nn.Linear(10, 10)
+                # Use proper FP8Linear mock
+                self.layer2 = FP8Linear(10, 10)
+
+                self.mlp = torch.nn.Sequential()
+                linear1 = torch.nn.Linear(10, 10)
+                self.mlp.add_module('0', linear1)
+                linear2 = FP8Linear(10, 10)
+                self.mlp.add_module('1', linear2)
+
+        model = MockModel()
+
+        # Test finding FP8Linear layer
+        result = get_fp_layer_names(model, 'layer2')
+        assert 'layer2' in result, "Should find FP8Linear layer (layer2)"
+
+        # Test finding mixed Linear and FP8Linear in mlp
+        result = get_fp_layer_names(model, 'mlp')
+        assert len(result) == 2, "Should find 2 layers in mlp (both Linear and FP8Linear)"
+        assert 'mlp.0' in result, "Should find regular Linear in mlp"
+        assert 'mlp.1' in result, "Should find FP8Linear in mlp"
+
+    def test_empty_ignore_layers(self):
+        """Test with empty ignore_layers string."""
+        class MockModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer1 = torch.nn.Linear(10, 10)
+
+        model = MockModel()
+        result = get_fp_layer_names(model, '')
+        assert len(result) == 0, "Empty ignore_layers should return empty list"
+
+    def test_multiple_ignore_patterns(self):
+        """Test with multiple ignore patterns."""
+        class MockModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer1 = torch.nn.Linear(10, 10)
+                self.layer2 = torch.nn.Linear(10, 10)
+                self.layer3 = torch.nn.Linear(10, 10)
+
+        model = MockModel()
+        result = get_fp_layer_names(model, 'layer1,layer3')
+        assert 'layer1' in result, "Should find layer1"
+        assert 'layer3' in result, "Should find layer3"
+        assert 'layer2' not in result, "Should not find layer2"
+
+    def test_pattern_with_digits(self):
+        """Test pattern matching with digits (special case in the code)."""
+        class MockModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList([
+                    torch.nn.Linear(10, 10),
+                    torch.nn.Linear(10, 10),
+                    torch.nn.Linear(10, 10)
+                ])
+
+        model = MockModel()
+        # Pattern ending with digit should get a dot appended for matching
+        result = get_fp_layer_names(model, 'layers.0')
+        # Should match 'layers.0'
+        assert 'layers.0' in result, "Should match layers.0"

test/test_cuda/advanced/test_fp8_input.py

@@ -141,3 +141,30 @@ def test_diff_datatype(self):
                 ar = AutoRound(model_name, iters=iters, scheme=scheme)
                 ar.quantize_and_save(output_dir=self.save_dir)
                 shutil.rmtree(self.save_dir, ignore_errors=True)
+
+    def test_ignore_layers_fp8(self):
+        """Test that ignore_layers works correctly with FP8 models."""
+        from auto_round.compressors.utils import get_fp_layer_names
+
+        model, tokenizer = self.tiny_fp8_model()
+
+        # Test that get_fp_layer_names can find FP8Linear layers
+        # Using "mlp" as the ignore pattern which should match mlp layers in the model
+        layer_names = get_fp_layer_names(model, "mlp")
+
+        # Verify that some layers were found
+        assert len(layer_names) > 0, "Should find layers matching 'mlp' pattern in FP8 model"
+
+        # Now test with AutoRound using ignore_layers
+        ar = AutoRound(
+            model=model, 
+            tokenizer=tokenizer, 
+            iters=0,
+            ignore_layers="mlp"
+        )
+        ar.quantize_and_save(output_dir=self.save_dir)
+
+        # Verify the model was saved successfully
+        assert os.path.exists(self.save_dir), "Model should be saved"
+
+        shutil.rmtree(self.save_dir, ignore_errors=True)