Repository: huggingface/transformers
File: tests/models/gpt_oss/test_modeling_gpt_oss.py (Python)

Code before:
class GptOssModelTester(CausalLMModelTester):
    if is_torch_available():
        base_model_class = GptOssModel


@require_torch
class GptOssModelTest(CausalLMModelTest, unittest.TestCase):
    _is_stateful = True
    model_split_percents = [0.5, 0.6]
    model_tester_class = GptOssModelTester

    @require_kernels
    @require_flash_attn
    @pytest.mark.flash_attn_test
    @require_torch_gpu
    def test_initialization_raises_error_for_flash_attn(self):
        """
        Tests that initializing the model with unsupported Flash Attention implementations raises a ValueError,
        but allows the specific vllm kernel.
        """

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        config._attn_implementation = "flash_attention_2"
        with self.assertRaisesRegex(ValueError, "GPT-OSS model does not support"):
            model = GptOssModel(config)

        config._attn_implementation = "kernels-community/vllm-flash-attn3"
        model = GptOssModel(config).to(device=torch_device, dtype=torch.bfloat16)

        model.eval()
        with torch.no_grad():
            output = model(**inputs_dict)
        self.assertIsNotNone(output)

        model.config._attn_implementation = "flash_attention_2"
        with self.assertRaisesRegex(ValueError, "GPT-OSS model does not support"):
            with torch.no_grad():
                model(**inputs_dict)

        model.config._attn_implementation = "kernels-community/vllm-flash-attn3"
        with torch.no_grad():
            output = model(**inputs_dict)
        self.assertIsNotNone(output)

    @unittest.skip("GptOss's forcefully disables sdpa due to Sink")
    def test_sdpa_can_dispatch_non_composite_models(self):
        pass

    @unittest.skip("GptOss's eager attn/sdpa attn outputs are expected to be different")
    def test_eager_matches_sdpa_generate(self):

Review comment (suggestion):
And here added the dtype casting to half precision (otherwise it won't run because FA only works with half precision)

Code after:
class GptOssModelTester(CausalLMModelTester):
    if is_torch_available():
        base_model_class = GptOssModel


@require_torch
class GptOssModelTest(CausalLMModelTest, unittest.TestCase):
    _is_stateful = True
    model_split_percents = [0.5, 0.6]
    model_tester_class = GptOssModelTester

    @require_kernels
    @require_flash_attn
    @pytest.mark.flash_attn_test
    @require_torch_gpu
    def test_initialization_raises_error_for_flash_attn(self):
        """
        Tests that initializing the model with unsupported Flash Attention implementations raises a ValueError,
        but allows the specific vllm kernel.
        """

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        kernel_attn = "kernels-community/vllm-flash-attn3"

        # Checking each via `set_attn_implementation` and manually setting within the config
        model = GptOssModel(config).to(device=torch_device, dtype=torch.bfloat16)
        model.set_attn_implementation("kernels-community/vllm-flash-attn3")
        self.assertTrue(model.config._attn_implementation == kernel_attn)

        config._attn_implementation = kernel_attn
        self.assertTrue(model.config._attn_implementation == kernel_attn)

        with torch.no_grad():
            output = model(**inputs_dict)
        self.assertIsNotNone(output)

        with self.assertRaisesRegex(ValueError, "GPT-OSS model does not support"):
            model.set_attn_implementation("flash_attention_2")

        with self.assertRaisesRegex(ValueError, "GPT-OSS model does not support"):
            config._attn_implementation = "flash_attention_2"

    @unittest.skip("GptOss's forcefully disables sdpa due to Sink")
    def test_sdpa_can_dispatch_non_composite_models(self):
        pass

    @unittest.skip("GptOss's eager attn/sdpa attn outputs are expected to be different")
    def test_eager_matches_sdpa_generate(self):
        pass

    @unittest.skip("GptOss eager/FA2 attention outputs are expected to be different")