diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index c3270fcfdb..c0435ad7c8 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -23,10 +23,6 @@ from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM from transformers.models.gptj.modeling_gptj import GPTJForCausalLM -from transformers.models.llama.modeling_llama import LlamaForCausalLM -from transformers.models.mistral.modeling_mistral import MistralForCausalLM -from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM -from transformers.models.phi3.modeling_phi3 import Phi3ForCausalLM from pyrit.executor.promptgen.gcg.experiments.log import ( log_gpu_memory, @@ -195,69 +191,52 @@ def _get_worker_log_params(worker: ModelWorker) -> dict[str, Any]: def get_embedding_layer(model: Any) -> Any: """ - Return the token embedding layer for a supported causal language model. + Return the token embedding layer for a causal language model. + + Uses the ``PreTrainedModel.get_input_embeddings`` interface, so any model that + ``AutoModelForCausalLM`` can load is supported rather than a fixed set of + architectures. + + Args: + model (Any): A loaded causal language model. Returns: Any: The model's token embedding layer. - - Raises: - ValueError: If the model architecture is unsupported. """ - if isinstance(model, (GPTJForCausalLM, GPT2LMHeadModel)): - return model.transformer.wte - if isinstance(model, LlamaForCausalLM): - return model.model.embed_tokens - if isinstance(model, GPTNeoXForCausalLM): - return model.base_model.embed_in - if isinstance(model, Phi3ForCausalLM): - return model.model.embed_tokens - raise ValueError(f"Unknown model type: {type(model)}") + return model.get_input_embeddings() def get_embedding_matrix(model: Any) -> Any: """ - Return the token embedding matrix for a supported causal language model. + Return the token embedding matrix for a causal language model. + + Args: + model (Any): A loaded causal language model. Returns: Any: The model's token embedding matrix. - - Raises: - ValueError: If the model architecture is unsupported. """ - if isinstance(model, (GPTJForCausalLM, GPT2LMHeadModel)): - return model.transformer.wte.weight - if isinstance(model, LlamaForCausalLM): - return model.model.embed_tokens.weight - if isinstance(model, GPTNeoXForCausalLM): - return model.base_model.embed_in.weight # type: ignore[union-attr, unused-ignore] - if isinstance(model, (MixtralForCausalLM, MistralForCausalLM)): - return model.model.embed_tokens.weight - if isinstance(model, Phi3ForCausalLM): - return model.model.embed_tokens.weight - raise ValueError(f"Unknown model type: {type(model)}") + return model.get_input_embeddings().weight def get_embeddings(model: Any, input_ids: torch.Tensor) -> Any: """ - Embed input token ids with a supported causal language model. + Embed input token ids with a causal language model. + + Args: + model (Any): A loaded causal language model. + input_ids (torch.Tensor): Token ids to embed. Returns: Any: The embedded token tensor. - - Raises: - ValueError: If the model architecture is unsupported. """ - if isinstance(model, (GPTJForCausalLM, GPT2LMHeadModel)): - return model.transformer.wte(input_ids).half() - if isinstance(model, LlamaForCausalLM): - return model.model.embed_tokens(input_ids) - if isinstance(model, GPTNeoXForCausalLM): - return model.base_model.embed_in(input_ids).half() # type: ignore[operator, unused-ignore] - if isinstance(model, (MixtralForCausalLM, MistralForCausalLM)): - return model.model.embed_tokens(input_ids) - if isinstance(model, Phi3ForCausalLM): - return model.model.embed_tokens(input_ids) - raise ValueError(f"Unknown model type: {type(model)}") + embeddings = model.get_input_embeddings()(input_ids) + # GPT-2, GPT-J and GPT-NeoX have always returned half precision here, while + # the other supported architectures return the embedding dtype unchanged. + # That asymmetry is preserved so this change stays a compatibility fix. + if isinstance(model, (GPTJForCausalLM, GPT2LMHeadModel, GPTNeoXForCausalLM)): + return embeddings.half() + return embeddings def get_nonascii_toks(tokenizer: Any, device: str = "cpu") -> torch.Tensor: diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index 148e67c980..8aaba93631 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -341,28 +341,63 @@ def test_non_ascii_filtering(self) -> None: assert new_tok not in non_ascii_set, f"Candidate {i} position {pos}: sampled non-ASCII token {new_tok}" +# Architectures built as tiny random models to exercise the embedding helpers. +# The first three predate this generic path and must keep returning float16 from +# get_embeddings; the rest were previously rejected outright. +_HALF_PRECISION_ARCHITECTURES = ["gpt2", "gptj", "gpt_neox"] +_OTHER_ARCHITECTURES = ["llama", "mistral", "mixtral", "phi3", "qwen3", "starcoder2"] + +_TINY_CONFIG = { + "hidden_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "intermediate_size": 64, + "vocab_size": 256, +} +_EXTRA_CONFIG = { + "phi3": { + "max_position_embeddings": 64, + "original_max_position_embeddings": 64, + "pad_token_id": 0, + }, +} + + +def _tiny_model(model_type: str): + """Build a small randomly initialized model of the given architecture.""" + transformers = pytest.importorskip("transformers", reason="transformers not installed") + config = transformers.AutoConfig.for_model(model_type, **_TINY_CONFIG, **_EXTRA_CONFIG.get(model_type, {})) + return transformers.AutoModelForCausalLM.from_config(config) + + class TestEmbeddingHelpers: """Tests for get_embedding_layer, get_embedding_matrix, get_embeddings.""" - def test_get_embedding_layer_raises_for_unknown_model(self) -> None: - """Should raise ValueError for unsupported model types.""" - mock_model = MagicMock() - # Ensure it doesn't match any isinstance checks - mock_model.__class__ = type("UnknownModel", (), {}) - with pytest.raises(ValueError, match="Unknown model type"): - get_embedding_layer(mock_model) - - def test_get_embedding_matrix_raises_for_unknown_model(self) -> None: - mock_model = MagicMock() - mock_model.__class__ = type("UnknownModel", (), {}) - with pytest.raises(ValueError, match="Unknown model type"): - get_embedding_matrix(mock_model) - - def test_get_embeddings_raises_for_unknown_model(self) -> None: - mock_model = MagicMock() - mock_model.__class__ = type("UnknownModel", (), {}) - with pytest.raises(ValueError, match="Unknown model type"): - get_embeddings(mock_model, torch.tensor([1, 2, 3])) + @pytest.mark.parametrize("model_type", _HALF_PRECISION_ARCHITECTURES + _OTHER_ARCHITECTURES) + def test_helpers_resolve_embeddings_for_any_causal_model(self, model_type: str) -> None: + """Any model AutoModelForCausalLM can load should resolve through the helpers.""" + model = _tiny_model(model_type) + expected = model.get_input_embeddings() + + assert get_embedding_layer(model) is expected + assert get_embedding_matrix(model) is expected.weight + + embedded = get_embeddings(model, torch.tensor([[1, 2, 3]])) + assert embedded.shape[-1] == model.config.hidden_size + + @pytest.mark.parametrize("model_type", _HALF_PRECISION_ARCHITECTURES) + def test_get_embeddings_keeps_half_precision_for_legacy_architectures(self, model_type: str) -> None: + """GPT-2, GPT-J and GPT-NeoX returned float16 before this path existed.""" + model = _tiny_model(model_type) + assert get_embeddings(model, torch.tensor([[1, 2, 3]])).dtype == torch.float16 + + @pytest.mark.parametrize("model_type", _OTHER_ARCHITECTURES) + def test_get_embeddings_keeps_embedding_dtype_for_other_architectures(self, model_type: str) -> None: + """Everything else keeps the embedding dtype rather than being downcast.""" + model = _tiny_model(model_type) + expected_dtype = model.get_input_embeddings().weight.dtype + assert get_embeddings(model, torch.tensor([[1, 2, 3]])).dtype == expected_dtype class TestPromptManagerInit: