Repository: huggingface/transformers
File: tests/quantization/finegrained_fp8/test_fp8.py (Python)

Code before:
        cls.tokenizer = AutoTokenizer.from_pretrained(cls.model_name)
        cls.quantized_model = AutoModelForCausalLM.from_pretrained(
            cls.model_name, device_map=cls.device_map, quantization_config=cls.quantization_config
        )

    def setup(self):
        """
        Clear also on each setup (e.g. if a different model is used than the base cls one)
        """
        gc.collect()
        backend_empty_cache(torch_device)
        gc.collect()

    def tearDown(self):
        gc.collect()
        backend_empty_cache(torch_device)
        gc.collect()

    @parameterized.expand(
        [
            "hf-internal-testing/tiny-random-Qwen3MoeForCausalLM",
            "hf-internal-testing/tiny-random-MixtralForCausalLM",
        ]
    )
    def test_moe_conversion_doesnt_raise(self, model_id):
        quantization_config = FineGrainedFP8Config()
        AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config)

    def test_quantized_model_conversion(self):
        """
        Simple test that checks if the quantized model has been converted properly
        """

        from transformers.integrations import FP8Linear, replace_with_fp8_linear

        model_id = "facebook/opt-350m"
        config = AutoConfig.from_pretrained(model_id, revision="cb32f77e905cccbca1d970436fb0f5e6b58ee3c5")
        quantization_config = FineGrainedFP8Config()

        with torch.device("meta"):
            model = OPTForCausalLM(config)

        nb_linears = 0
        for module in model.modules():
            if isinstance(module, torch.nn.Linear):
                nb_linears += 1
        model = replace_with_fp8_linear(model, quantization_config=quantization_config)
        nb_fp8_linear = 0
        for module in model.modules():
            if isinstance(module, FP8Linear):
                nb_fp8_linear += 1

Review comment (suggestion):
```suggestion
        quantization_config = FineGrainedFP8Config(weight_block_size=(32, 32))
```

Code after:
        cls.tokenizer = AutoTokenizer.from_pretrained(cls.model_name)
        cls.quantized_model = AutoModelForCausalLM.from_pretrained(
            cls.model_name, device_map=cls.device_map, quantization_config=cls.quantization_config
        )

    def setup(self):
        """
        Clear also on each setup (e.g. if a different model is used than the base cls one)
        """
        gc.collect()
        backend_empty_cache(torch_device)
        gc.collect()

    def tearDown(self):
        gc.collect()
        backend_empty_cache(torch_device)
        gc.collect()

    @parameterized.expand(
        [
            "hf-internal-testing/tiny-random-Qwen3MoeForCausalLM",
            "hf-internal-testing/tiny-random-MixtralForCausalLM",
        ]
    )
    def test_moe_conversion_doesnt_raise(self, model_id):
        quantization_config = FineGrainedFP8Config(weight_block_size=(32, 32))
        AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config)

    def test_quantized_model_conversion(self):
        """
        Simple test that checks if the quantized model has been converted properly
        """

        from transformers.integrations import FP8Linear, replace_with_fp8_linear

        model_id = "facebook/opt-350m"
        config = AutoConfig.from_pretrained(model_id, revision="cb32f77e905cccbca1d970436fb0f5e6b58ee3c5")
        quantization_config = FineGrainedFP8Config()

        with torch.device("meta"):
            model = OPTForCausalLM(config)

        nb_linears = 0
        for module in model.modules():
            if isinstance(module, torch.nn.Linear):
                nb_linears += 1
        model = replace_with_fp8_linear(model, quantization_config=quantization_config)
        nb_fp8_linear = 0
        for module in model.modules():
            if isinstance(module, FP8Linear):
                nb_fp8_linear += 1