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

Code before:
    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

    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 = self.model_tester.get_config()
        config._attn_implementation = "flash_attention_2"
        with self.assertRaisesRegex(ValueError, "GPT-OSS models do not support"):
            GptOssModel(config)

        config._attn_implementation = "kernels-community/vllm-flash-attn3"
        try:
            GptOssModel(config)
        except ValueError:
            self.fail("GptOssModel raised ValueError unexpectedly with vllm-flash-attn3!")

    @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")
    def test_flash_attn_2_equivalence(self):
        pass

    @unittest.skip("Most probably because of the MOE, the moe and router does not ignore padding tokens")
    def test_eager_padding_matches_padding_free_with_position_ids(self):
        pass

    @unittest.skip("GptOss does not support flex officially")
    def test_flex_attention_with_grads(self):
        pass

    @unittest.skipIf(torch_device == "cpu", "GptOss does not support flex officially")
    def test_generate_compile_model_forward_fullgraph(self):
        return super().test_generate_compile_model_forward_fullgraph()


Review comment (suggestion):
I don't think we need this extra try catch, we should expect no error either way

Code after:
    NUM_GPUS = torch.cuda.device_count()


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):