remove awq fuse test

Author: LRL2-ModelCloud

tests/quantization/autoawq/test_awq.py

@@ -287,233 +287,6 @@ def test_quantized_model_no_k_proj_quantized(self):
         output = quantized_model.generate(dummy_input, max_new_tokens=10)
         self.assertTrue((EXPECTED_OUTPUT == output).all())
 
-
-@slow
-@require_torch_accelerator
-@require_auto_awq
-@require_accelerate
-class AwqFusedTest(unittest.TestCase):
-    model_name = "TheBloke/Mistral-7B-OpenOrca-AWQ"
-    model_revision = "7048b2af77d0dd1c81b000b19d73f9cc8950b510"
-
-    custom_mapping_model_id = "TheBloke/Mistral-7B-v0.1-AWQ"
-    custom_model_revision = "f186bcfa9edbe2a4334262ec1e67f23e53ed1ae7"
-
-    mixtral_model_name = "casperhansen/mixtral-instruct-awq"
-    mixtral_model_revision = "87dd4ec502dde74fb3a624835c776b000d190c3b"
-
-    multi_modal_model_name = "ybelkada/llava-1.5-7b-hf-awq"
-    multi_modal_model_code_revision = "ad108a50f5b9e681bdd7378409f57b7fa59a7442"
-
-    prompt = (
-        "You're standing on the surface of the Earth. "
-        "You walk one mile south, one mile west and one mile north. "
-        "You end up exactly where you started. Where are you?"
-    )
-
-    EXPECTED_GENERATION = prompt + "\n\nYou're at the center of a square."
-    EXPECTED_GENERATION_CUSTOM_MODEL = "Hello,\n\nI have a problem with my 20"
-    EXPECTED_GENERATION_MIXTRAL = prompt + " You're on the North Pole.\n\nThe"
-
-    def tearDown(self):
-        gc.collect()
-        backend_empty_cache(torch_device)
-        gc.collect()
-
-    def _check_fused_modules(self, model):
-        has_fused_modules = False
-        fused_modules_name = ["QuantAttentionFused", "QuantFusedMLP", "FasterTransformerRMSNorm"]
-
-        for _, module in model.named_modules():
-            if module.__class__.__name__ in fused_modules_name:
-                has_fused_modules = True
-                break
-
-        self.assertTrue(has_fused_modules, "Modules fusing not performed correctly!")
-
-    def test_raise_save_pretrained(self):
-        """
-        Test that `save_pretrained` is effectively blocked for fused models
-        """
-        quantization_config = AwqConfig(bits=4, fuse_max_seq_len=128, do_fuse=True)
-
-        model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
-            quantization_config=quantization_config,
-            revision=self.model_revision,
-        ).to(torch_device)
-
-        self._check_fused_modules(model)
-
-        with self.assertRaises(ValueError), tempfile.TemporaryDirectory() as tmpdirname:
-            model.save_pretrained(tmpdirname)
-
-    def test_fused_modules_to_not_convert(self):
-        """
-        Test if fused + modules to_not_convert work as expected
-        """
-        model_id = "hf-internal-testing/Mixtral-tiny-AWQ"
-
-        quantization_config = AwqConfig(bits=4, fuse_max_seq_len=128, do_fuse=True)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            quantization_config=quantization_config,
-        ).to(torch_device)
-
-        # Check if model has been correctly fused
-        self._check_fused_modules(model)
-        # Checks if the modules_to_not_convert (here gate layer) is a Linear
-        self.assertTrue(isinstance(model.model.layers[0].block_sparse_moe.gate, torch.nn.Linear))
-
-    @unittest.skipIf(
-        get_device_properties()[0] == "cuda" and get_device_properties()[1] < 8,
-        "Skipping because RuntimeError: FlashAttention only supports Ampere GPUs or newer, so not supported on GPU with capability < 8.0",
-    )
-    @require_flash_attn
-    @require_torch_gpu
-    @pytest.mark.flash_attn_test
-    def test_generation_fused(self):
-        """
-        Test generation quality for fused models - single batch case
-        """
-        quantization_config = AwqConfig(bits=4, fuse_max_seq_len=128, do_fuse=True)
-
-        model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
-            quantization_config=quantization_config,
-            revision=self.model_revision,
-        ).to(torch_device)
-
-        self._check_fused_modules(model)
-
-        tokenizer = AutoTokenizer.from_pretrained(self.model_name, revision=self.model_revision)
-
-        inputs = tokenizer(self.prompt, return_tensors="pt").to(torch_device)
-
-        outputs = model.generate(**inputs, max_new_tokens=12)
-
-        self.assertEqual(tokenizer.decode(outputs[0], skip_special_tokens=True), self.EXPECTED_GENERATION)
-
-    @pytest.mark.flash_attn_test
-    @require_flash_attn
-    @require_torch_gpu
-    @unittest.skipIf(
-        get_device_properties()[0] == "cuda" and get_device_properties()[1] < 8,
-        "Skipping because RuntimeError: FlashAttention only supports Ampere GPUs or newer, so not supported on GPU with capability < 8.0",
-    )
-    def test_generation_fused_batched(self):
-        """
-        Test generation quality for fused models - multi batch case
-        """
-        quantization_config = AwqConfig(bits=4, fuse_max_seq_len=128, do_fuse=True)
-
-        model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
-            quantization_config=quantization_config,
-            revision=self.model_revision,
-        ).to(torch_device)
-
-        self._check_fused_modules(model)
-
-        tokenizer = AutoTokenizer.from_pretrained(self.model_name, revision=self.model_revision)
-
-        tokenizer.pad_token_id = tokenizer.eos_token_id
-        inputs = tokenizer([self.prompt, self.prompt], return_tensors="pt", padding=True).to(torch_device)
-
-        outputs = model.generate(**inputs, max_new_tokens=12)
-
-        self.assertEqual(tokenizer.decode(outputs[0], skip_special_tokens=True), self.EXPECTED_GENERATION)
-
-    def test_generation_llava_fused(self):
-        from transformers import pipeline
-
-        quantization_config = AwqConfig(do_fuse=True, fuse_max_seq_len=2048)
-
-        pipe = pipeline(
-            "image-to-text",
-            model=self.multi_modal_model_name,
-            device=0,
-            model_kwargs={
-                "quantization_config": quantization_config,
-            },
-            revision=self.multi_modal_model_code_revision,
-        )
-        url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/compel-neg.png"
-
-        prompt = "USER: <image>\nCan you please describe this image?\nASSISTANT:"
-
-        outputs = pipe(url, prompt=prompt, generate_kwargs={"max_new_tokens": 100})
-        EXPECTED_OUTPUT = "USER:  \nCan you please describe this image?\nASSISTANT: The image features a brown and white cat sitting on a green surface, possibly a carpet or a grassy area. The cat is holding a red ball in its paws, seemingly playing with it. The cat appears to be focused on the ball, possibly preparing to play or just enjoying the toy."
-
-        self.assertEqual(outputs[0]["generated_text"], EXPECTED_OUTPUT)
-
-    @pytest.mark.flash_attn_test
-    @require_flash_attn
-    @require_torch_multi_gpu
-    @unittest.skipIf(
-        get_device_properties()[0] == "cuda" and get_device_properties()[1] < 8,
-        "Skipping because RuntimeError: FlashAttention only supports Ampere GPUs or newer, so not supported on GPU with capability < 8.0",
-    )
-    def test_generation_custom_model(self):
-        """
-        Test generation quality for fused models using custom fused map.
-        """
-        quantization_config = AwqConfig(
-            bits=4,
-            fuse_max_seq_len=512,
-            modules_to_fuse={
-                "attention": ["q_proj", "k_proj", "v_proj", "o_proj"],
-                "mlp": ["gate_proj", "up_proj", "down_proj"],
-                "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"],
-                "use_alibi": False,
-                "hidden_size": 4096,
-                "num_attention_heads": 32,
-                "num_key_value_heads": 8,
-            },
-        )
-
-        model = AutoModelForCausalLM.from_pretrained(
-            self.custom_mapping_model_id,
-            quantization_config=quantization_config,
-            device_map="balanced",
-            revision=self.custom_model_revision,
-        )
-
-        self._check_fused_modules(model)
-
-        tokenizer = AutoTokenizer.from_pretrained(self.custom_mapping_model_id, revision=self.custom_model_revision)
-
-        prompt = "Hello"
-        inputs = tokenizer(prompt, return_tensors="pt").to(torch_device)
-
-        outputs = model.generate(**inputs, max_new_tokens=12)
-        self.assertEqual(tokenizer.decode(outputs[0], skip_special_tokens=True), self.EXPECTED_GENERATION_CUSTOM_MODEL)
-
-    @pytest.mark.flash_attn_test
-    @require_flash_attn
-    @require_torch_multi_gpu
-    @unittest.skip(reason="Not enough GPU memory on CI runners")
-    def test_generation_mixtral_fused(self):
-        """
-        Text generation test for Mixtral + AWQ + fused
-        """
-        quantization_config = AwqConfig(bits=4, fuse_max_seq_len=1024, do_fuse=True)
-        model = AutoModelForCausalLM.from_pretrained(
-            self.mixtral_model_name,
-            quantization_config=quantization_config,
-            device_map="auto",
-            revision=self.mixtral_model_revision,
-        )
-
-        tokenizer = AutoTokenizer.from_pretrained(self.mixtral_model_name)
-        tokenizer.pad_token = tokenizer.eos_token
-
-        inputs = tokenizer([self.prompt, self.prompt], return_tensors="pt", padding=True).to(torch_device)
-
-        outputs = model.generate(**inputs, max_new_tokens=12)
-        self.assertEqual(tokenizer.decode(outputs[0], skip_special_tokens=True), self.EXPECTED_GENERATION_MIXTRAL)
-
-
 @slow
 @require_torch_accelerator
 @require_auto_awq