diff --git a/CLAUDE.md b/CLAUDE.md index 4063eb9..8e805ad 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -65,10 +65,9 @@ Naming: `{Value}{InputType}AccessMixin` (e.g. `LossTokenAccessMixin`, `GradientT Canonical definitions and full field lists live in `tropt/common.py`. Don't enumerate them here. -### Glue: Recipe Hub & Config Runner +### Glue: Recipe Hub - **Recipe Hub** (`tropt/recipe_hub/`): pre-configured Model + Loss + Optimizer + Inputs/Targets wirings, each callable as one function. Enumerate via `list_recipes()`; full registry in `tropt/recipe_hub/__init__.py`. -- **Config Runner** (`runner/main.py`): YAML-driven runner via Hydra. (Hydra config support is incomplete — see Known Limitations.) ## Component Interactions @@ -240,8 +239,6 @@ Auto-generated docs (`docs/api/`, `docs/guides/compatibility_matrix.md`) are pro For contributions back to the package (file placement, exports, tests, Recipe Hub naming convention), see `CONTRIBUTING.md`. -For testing conventions (mirror layout, fixtures, numerical tolerances, what to test per component), see `TESTING.md`. - ## Important Notes ### Scripts: separate repo @@ -257,13 +254,12 @@ This codebase is explicitly designed for adversarial robustness research and red ### Known Limitations - Multiple-message prefix caching currently disabled due to edge cases - Tracker should be initialized per RUN, not per optimizer instance -- Hydra config-runner support is incomplete ## Dependencies -Core: PyTorch, Transformers, Accelerate, Hydra, Pydantic -Models: HuggingFace, SentenceTransformers, OpenAI, LiteLLM -Tracking: Weights & Biases, LiveLossPlot +Core: PyTorch, Transformers, Accelerate, SentenceTransformers, Pydantic +Optional extras: OpenAI, Google, Voyage, LiteLLM, vision, tracking, notebooks +Tracking: Weights & Biases, Trackio, LiveLossPlot Dev: pytest, ruff, ty, pre-commit ## Repository Structure @@ -279,7 +275,6 @@ tropt/ └── utils/ # Shared utilities tests/ # Test suite mirroring tropt/ structure -runner/ # Hydra-driven config runner scripts/ # Separate repo — see "Scripts: separate repo" above docs/ # Sphinx documentation ├── api/ # Auto-generated API reference (rst) @@ -288,6 +283,5 @@ docs/ # Sphinx documentation └── conf.py # Sphinx config quickstart.ipynb # End-to-end notebook DESIGN.md # Design philosophy and rationale -TESTING.md # Testing guidelines and conventions CONTRIBUTING.md # Contribution workflow (file placement, exports, tests) ``` diff --git a/DESIGN.md b/DESIGN.md index 5c6b62c..a66c32d 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -156,7 +156,7 @@ return resolve_and_compute_loss(model_output, model_input, loss_func) # the los This keeps loss logic in one location, makes adding new loss types straightforward, and means new models only provide data---not loss implementation. **Convention.** As a good practice, we divide the losses with superclasses according to the type of input that the loss accepts (which is, in turn, mostly the type of output of the model). For instance, Cross-Entropy-based losses utilize the logit outputs, and thus they will inherit from `PrefillBasedLoss`. -In this way, from the model end, we would be disable unneeded calculation: for instnace, if the loss does not requrie hidden states (i.e., the loss is not subclass of `HiddenStateBasedLoss`), we can know in advance to avoid saving them for efficiency. +In this way, from the model end, we would be disable unneeded calculation: for instnace, if the loss does not requrie hidden states (i.e., the loss does not set `require_hidden_states`), we can know in advance to avoid saving them for efficiency. ## Component 3: Optimizers @@ -275,7 +275,6 @@ The repository exposes two interfaces for managing recipes: * **Recipe Hub [`tropt/recipe_hub/`].** Python modules that bind the four components to reproduce existing attacks. E.g., the `gcg__zou2023` module wires `LMHFModel`, `PrefillCELoss`, `GCGOptimizer`, and a standard suffix template to reproduce the GCG attack. Useful for researchers who want to quickly run published attacks, benchmark them, or fork-and-modify. - ## Summary diff --git a/docs/api/loss.rst b/docs/api/loss.rst index 0d8d346..c1ddc16 100644 --- a/docs/api/loss.rst +++ b/docs/api/loss.rst @@ -27,17 +27,7 @@ Loss Classes Interfaces :undoc-members: :show-inheritance: -.. autoclass:: TriggerLogitBasedLoss - :members: - :undoc-members: - :show-inheritance: - -.. autoclass:: AttentionBasedLoss - :members: - :undoc-members: - :show-inheritance: - -.. autoclass:: EmbeddingBasedLoss +.. autoclass:: PrefillBasedLoss :members: :undoc-members: :show-inheritance: @@ -66,7 +56,7 @@ Loss Implementations .. automodule:: tropt.loss :members: - :exclude-members: BaseLoss, TriggerLogitBasedLoss, AttentionBasedLoss, EmbeddingBasedLoss, TextBasedLoss, SteeringActivationLoss, CombinedLoss + :exclude-members: BaseLoss, PrefillBasedLoss, TextBasedLoss, SteeringActivationLoss, CombinedLoss :undoc-members: :show-inheritance: :imported-members: diff --git a/docs/api/models.rst b/docs/api/models.rst index c60434d..aca8f42 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -98,7 +98,7 @@ Model Implementations .. automodule:: tropt.model :members: - :exclude-members: BaseModel, LMBaseModel, EncoderBaseModel, BaseTokenizer, TokenAccessMixin, LossTokenAccessMixin, LogitsTokenAccessMixin, GradientTokenAccessMixin, TextAccessMixin, LossTextAccessMixin, InputsManager, TextInputManager, TokenInputManager, DefaultTokenInputManager + :exclude-members: BaseModel, LMBaseModel, EncoderBaseModel, BaseTokenizer, TokenAccessMixin, LossTokenAccessMixin, LogitsTokenAccessMixin, GradientTokenAccessMixin, TextAccessMixin, LossTextAccessMixin, TextInputManager, TokenInputManager, DefaultTokenInputManager :undoc-members: :show-inheritance: :imported-members: \ No newline at end of file diff --git a/docs/build_docs.py b/docs/build_docs.py index b15a163..2897bc8 100644 --- a/docs/build_docs.py +++ b/docs/build_docs.py @@ -3,41 +3,36 @@ import shutil import stat import subprocess +import sys import time import webbrowser -def make_writable(path): - """Force a file/directory to be writable.""" - try: - os.chmod(path, stat.S_IWRITE) - except Exception: - pass +def _force_writable(func, path, _exc): + """rmtree error hook: clear the read-only bit (git/GitHub files) and retry.""" + os.chmod(path, stat.S_IWRITE) + func(path) + + +# `onexc` replaced `onerror` in Python 3.12; the project supports 3.10+. +_RMTREE_HOOK = ( + {"onexc": _force_writable} if sys.version_info >= (3, 12) + else {"onerror": _force_writable} +) + def robust_cleanup(path): - """ - Aggressively cleans up a directory. - 1. Walks tree to fix permissions (handling read-only Git/GitHub files). - 2. Tries to delete. - 3. Returns True if successful, False if locked. + """Delete `path`, retrying while Google Drive / Windows holds a lock. + + Returns True if the tree is gone, False if it stayed locked. """ if not os.path.exists(path): return True print(f"Cleaning up {path}...") - - # 1. Force permissions first (Pre-emptive strike) - for root, dirs, files in os.walk(path): - for d in dirs: - make_writable(os.path.join(root, d)) - for f in files: - make_writable(os.path.join(root, f)) - make_writable(path) - - # 2. Try deletion with retries - for i in range(5): + for _ in range(5): try: - shutil.rmtree(path) + shutil.rmtree(path, **_RMTREE_HOOK) return True except OSError: time.sleep(0.5) # Wait for Google Drive/Windows to release lock @@ -84,41 +79,10 @@ def build_docs(): # --------------------------------------------------------- # 2. Inject 'from __future__ import annotations' # --------------------------------------------------------- + # Same script CI runs (deploy_docs.yml), pointed at the throwaway copy. print("Injecting future annotations...") - import ast - for root, _, files in os.walk(src_copy): - for file in files: - # common.py is excluded to match CI (deploy_docs.yml): future - # annotations stringify its pydantic/jaxtyping runtime field types - # and break Sphinx autodoc (sphinx-doc/sphinx#11211). - if file.endswith(".py") and file != "common.py": - path = os.path.join(root, file) - try: - with open(path, "r", encoding="utf-8") as f: - content = f.read() - - if "from __future__ import annotations" not in content: - # Insert AFTER a leading module docstring so the docstring - # isn't demoted to a bare string expression (which would - # strip it from the autodoc-rendered module page). - insert_at = 0 - try: - mod = ast.parse(content) - first = mod.body[0] if mod.body else None - if ( - isinstance(first, ast.Expr) - and isinstance(getattr(first, "value", None), ast.Constant) - and isinstance(first.value.value, str) - ): - insert_at = first.end_lineno - except SyntaxError: - insert_at = 0 - lines = content.splitlines(keepends=True) - lines.insert(insert_at, "from __future__ import annotations\n") - with open(path, "w", encoding="utf-8") as f: - f.write("".join(lines)) - except Exception as e: - print(f" Skipping {file}: {e}") + from docs.scripts.inject_annotations import inject + print(f" Injected into {inject(src_copy)} module(s).") # NOTE: inlining + sanitizing tropt/recipe_hub/README.md now happens in # conf.py (the `include-read` event), so it runs for every build path. diff --git a/docs/conf.py b/docs/conf.py index 424a71b..05160ae 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -56,7 +56,7 @@ myst_enable_extensions = ["colon_fence", "deflist", "attrs_inline"] myst_heading_anchors = 3 -autodoc_mock_imports = ['runner', 'sentence_transformers', 'wandb', 'livelossplot', 'openai', 'litellm', 'IPython', 'tqdm', 'transformers', 'accelerate', 'hydra', 'omegaconf', 'huggingface_hub', 'datasets', 'PIL', 'diffusers', 'trackio'] +autodoc_mock_imports = ['sentence_transformers', 'wandb', 'livelossplot', 'openai', 'litellm', 'IPython', 'tqdm', 'transformers', 'accelerate', 'huggingface_hub', 'datasets', 'PIL', 'diffusers', 'trackio'] autodoc_typehints = "description" diff --git a/docs/guides/adding_a_loss.md b/docs/guides/adding_a_loss.md index bfe2fe8..77c5de7 100644 --- a/docs/guides/adding_a_loss.md +++ b/docs/guides/adding_a_loss.md @@ -130,12 +130,12 @@ Breaking down the additions: **Convention: Hyperparameters as Fields.** Anything you want callers to tweak at construction time — temperatures, margins, layer ranges, mode flags — goes here. Crucially, *per-template* data (target tokens, target vectors, target classes) does **not** belong here; it belongs on {py:class}`~tropt.common.Targets`, and the loss pulls it in by parameter name (`target_response_toks` above). This keeps the loss instance stateless w.r.t. the attack — the optimizer can resample or subsample templates without telling the loss. -**Convention: Inheriting from a category base class.** {py:mod}`tropt.loss` defines abstract bases like {py:class}`~tropt.loss.PrefillBasedLoss`, {py:class}`~tropt.loss.EmbeddingBasedLoss`, {py:class}`~tropt.loss.HiddenStateBasedLoss`, {py:class}`~tropt.loss.AttentionBasedLoss`, {py:class}`~tropt.loss.ClassificationBasedLoss`, and {py:class}`~tropt.loss.TextBasedLoss`. Each sets the right `require_*` flag and pins a typed `__call__` signature for its category — inheriting from the right base saves boilerplate and signals intent. The categorization is a **convention for readability**, not a hard requirement: the resolver dispatches by parameter names, not by base class. If your loss doesn't fit any existing category, inheriting from `BaseLoss` directly and setting the flags yourself is equally valid. +**Convention: inherit from `BaseLoss`, declare your own flags.** The resolver dispatches on **parameter names**, never on base class — so subclassing `BaseLoss` directly and setting the `require_*` flags you need is the normal case. Two shared bases exist only where they carry something real: {py:class}`~tropt.loss.PrefillBasedLoss` (sets `require_target_prefill=True` and pins the prefill `__call__` signature shared by four losses) and {py:class}`~tropt.loss.TextBasedLoss` (sets `is_differentiable=False` for text-scoring losses). Don't add a new base class for a single loss. ## Practical Example: Activation Steering -Up to here we asked the model for logits over a target response. The same pattern works just as well for *any* model artifact — swap the base class, swap the `require_*` flag, and swap the `__call__` parameter names, and you have a loss over a completely different signal. +Up to here we asked the model for logits over a target response. The same pattern works just as well for *any* model artifact — swap the `require_*` flag and swap the `__call__` parameter names, and you have a loss over a completely different signal. To demonstrate, we rewrite the loss as **activation steering** ([Arditi et al., 2024](https://arxiv.org/abs/2406.11717)): encouraging the model's hidden activations at chosen layers/positions to align with (or away from) a target direction in activation space. This has been used both to suppress refusal in jailbreaks and to probe internal representations. It's what {py:class}`~tropt.loss.SteeringActivationLoss` does in TROPT. @@ -148,11 +148,11 @@ import torch.nn.functional as F from jaxtyping import Float from torch import Tensor -from tropt.loss import HiddenStateBasedLoss +from tropt.loss import BaseLoss @dataclass -class MyLoss(HiddenStateBasedLoss): +class MyLoss(BaseLoss): """Steers hidden activations along (or away from) a target direction.""" # Hyperparameters @@ -160,7 +160,6 @@ class MyLoss(HiddenStateBasedLoss): steer_away: bool = False # if True, push activations *away* from the direction # Ask the model to return all hidden states. - # (HiddenStateBasedLoss already declares this; shown here for emphasis.) require_hidden_states: ClassVar[bool] = True def __call__( @@ -177,11 +176,11 @@ class MyLoss(HiddenStateBasedLoss): Structurally nothing is new — same `BaseLoss` lineage, same dataclass-fields-as-hyperparameters, same `require_*` declaration, same per-sample return shape. The differences are entirely in *which* names appear: -- **Base class & flag:** {py:class}`~tropt.loss.HiddenStateBasedLoss` (sets `require_hidden_states=True`), replacing `PrefillBasedLoss` / `require_target_prefill`. +- **Flag:** `require_hidden_states=True` replaces `require_target_prefill`. - **Model output:** `full_hidden_states` replaces `prefill_response_logits`. - **Per-template target:** `target_directions` (provided via `Targets(target_directions=...)`) replaces `target_response_toks`. -The same one-knob swap also produces attention-based losses (inherit from `AttentionBasedLoss`, name `full_attentions`) and classifier losses (inherit from `ClassificationBasedLoss`, name `output_class_logits`) — see the existing implementations of {py:class}`~tropt.loss.AttentionEnhLoss` and {py:class}`~tropt.loss.MisclassCELoss` for those variants. +The same one-knob swap also produces attention-based losses (`require_attentions=True`, name `full_attentions`) and classifier losses (no flag, name `output_class_logits`) — see the existing implementations of {py:class}`~tropt.loss.AttentionEnhLoss` and {py:class}`~tropt.loss.MisclassCELoss` for those variants. ## Going Non-Differentiable: Scoring Triggered Text diff --git a/docs/guides/adding_a_model.md b/docs/guides/adding_a_model.md index 22520e6..f43edbc 100644 --- a/docs/guides/adding_a_model.md +++ b/docs/guides/adding_a_model.md @@ -92,7 +92,7 @@ class ModelOutput: full_ids: ... ``` -The fields you populate determine which loss types are compatible with your model. For example, `output_embeddings` enables `EmbeddingBasedLoss` (e.g., `SimilarityLoss`), while `generated_response_strs` enables `TextBasedLoss` (e.g., `GeneratedResponseBasedLoss`). The {py:func}`loss resolution system ` validates this at runtime and raises clear errors if a required field is missing. +The fields you populate determine which losses are compatible with your model. For example, `output_embeddings` enables embedding losses (e.g., `SimilarityLoss`), while `generated_response_strs` enables generation-scoring losses (e.g., `ResponseHarmfulnessLoss`). The {py:func}`loss resolution system ` validates this at runtime and raises clear errors if a required field is missing. ### Model compatibility diff --git a/docs/guides/adding_an_optimizer.md b/docs/guides/adding_an_optimizer.md index 335ac94..f8b9408 100644 --- a/docs/guides/adding_an_optimizer.md +++ b/docs/guides/adding_an_optimizer.md @@ -211,7 +211,7 @@ optimizer.optimize_trigger(...) **Trigger initialization.** Callers can pass an explicit `initial_trigger` (the `DEFAULT_INIT_TRIGGER` is `"! ! ! ..."`-style). For something smarter, sample from the constrained vocabulary itself via {py:func}`~tropt.optimizer.utils.token_initializers.get_printable_random_trigger`. -The remaining building blocks in [optimizer utilities](../api/optimizer_utils) worth knowing about: `retokenize_filtering` (drop candidates that don't survive a decode → encode round-trip), `TriggerBuffer` (best-K pool instead of a single best), `NFlipScheduler` (control how many positions to mutate per step). Pull them in only when your search actually needs them — they're not boilerplate. +The remaining building blocks in [optimizer utilities](../api/optimizer_utils) worth knowing about: `retokenize_filtering` (drop candidates that don't survive a decode → encode round-trip), `TriggerBuffer` (best-K pool instead of a single best), `LinearScheduler` (a `(step) -> n_flip` callable controlling how many positions to mutate per step; any callable works). Pull them in only when your search actually needs them — they're not boilerplate. ## Going White-Box: Gradient-Guided Search diff --git a/docs/scripts/inject_annotations.py b/docs/scripts/inject_annotations.py new file mode 100644 index 0000000..226624b --- /dev/null +++ b/docs/scripts/inject_annotations.py @@ -0,0 +1,51 @@ +"""Insert `from __future__ import annotations` into every `tropt/*.py`. + +Sphinx autodoc can't resolve TROPT's jaxtyping/pydantic forward refs without it +(https://github.com/sphinx-doc/sphinx/issues/11211). `common.py` is excluded: +its pydantic models use runtime jaxtyping field types that break when +stringified. The import goes *after* any module docstring so the docstring still +renders on the autodoc module pages. + +Run against a throwaway copy of the tree — it rewrites files in place. + +Usage: python docs/scripts/inject_annotations.py [package_dir] +""" + +import ast +import pathlib +import sys + +EXCLUDED = {"common.py"} + + +def inject(package_dir: str | pathlib.Path = "tropt") -> int: + """Rewrite every eligible module under `package_dir`. Returns the count.""" + n = 0 + for path in pathlib.Path(package_dir).rglob("*.py"): + if path.name in EXCLUDED: + continue + src = path.read_text(encoding="utf-8") + if "from __future__ import annotations" in src: + continue + + insert_at = 0 + try: + first = (ast.parse(src).body or [None])[0] + if ( + isinstance(first, ast.Expr) + and isinstance(getattr(first, "value", None), ast.Constant) + and isinstance(first.value.value, str) + ): + insert_at = first.end_lineno + except SyntaxError: + insert_at = 0 + + lines = src.splitlines(keepends=True) + lines.insert(insert_at, "from __future__ import annotations\n") + path.write_text("".join(lines), encoding="utf-8") + n += 1 + return n + + +if __name__ == "__main__": + print(f"Injected future annotations into {inject(*sys.argv[1:])} module(s).") diff --git a/pyproject.toml b/pyproject.toml index 155e131..cc29d59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,16 +47,8 @@ dependencies = [ "pydantic>=2.0.0", # Configuration & utilities "tqdm", - "tenacity", - "pandas>=2.3.3", + # Dataset loading (AdvBench / Alpaca in tropt/utils/refusal_dir.py); lazily imported. "datasets>=4.8.3", - # Used directly by tropt/utils/refusal_dir.py (refusal-direction extraction). - "requests", - "scikit-learn", - - # Config [deprecated for now] - # "hydra-core>=1.3.0", - # "omegaconf", ] [project.urls] @@ -65,18 +57,22 @@ Repository = "https://github.com/matanbt/TROPT" Issues = "https://github.com/matanbt/TROPT/issues" [project.optional-dependencies] -# API model integrations (install separately as needed) +# API model integrations (install separately as needed). +# `tenacity` powers the transient-error retry on these backends only. openai = [ "openai>=1.0.0", "tiktoken>=0.5.0", + "tenacity", ] google = [ "google-genai>=1.0.0", + "tenacity", ] voyage = [ "voyageai>=0.3.0", + "tenacity", ] litellm = [ diff --git a/skills/tropt/SKILL.md b/skills/tropt/SKILL.md index ed7270e..0f05aef 100644 --- a/skills/tropt/SKILL.md +++ b/skills/tropt/SKILL.md @@ -37,7 +37,7 @@ optimizer (declares model_requirements, drives the search)──┘ Two design principles to keep in mind: *modularity* (each component swaps largely independently) and *backend–frontend separation* (tokenization, batching, gradient computation live in the model "backend"; losses and optimizers stay lightweight "frontend"). - **Model** (`tropt/model/`) — wraps the target model and exposes capabilities via **access mixins**. The optimizer declares which mixins it needs; `BaseOptimizer.__init__` validates the model satisfies them. Common mixins: `LossTokenAccessMixin` (grey-box loss from tokens), `GradientTokenAccessMixin` (white-box gradients), `LossTextAccessMixin` (black-box, text-only), `LogitsTokenAccessMixin`, `GradientEmbedAccessMixin`. Backends: `LMHFModel` (HF causal LMs — most mixins), `EncoderHFModel`, `EncoderOpenAIModel`, `LiteLLMModel`, `EncoderGeminiModel`. -- **Loss** (`tropt/loss/`) — pure objective. Receives fields from `ModelOutput` / `ModelInput` / `MessageTargets` by **parameter-name matching** in its `__call__` signature (the one rule). Categories: `PrefillBasedLoss`, `TriggerLogitBasedLoss`, `EmbeddingBasedLoss`, `TextBasedLoss`, `AttentionBasedLoss`, `HiddenStateBasedLoss`, `ClassificationBasedLoss`, `CombinedLoss`. +- **Loss** (`tropt/loss/`) — pure objective. Receives fields from `ModelOutput` / `ModelInput` / `MessageTargets` by **parameter-name matching** in its `__call__` signature (the one rule). Most losses subclass `BaseLoss` directly and declare their own `require_*` flags; shared bases exist only for `PrefillBasedLoss` (4 losses) and `TextBasedLoss` (non-differentiable). `CombinedLoss` weights several together. - **Optimizer** (`tropt/optimizer/`) — the search algorithm. Self-contained, one file per optimizer (HuggingFace "Repeat Yourself"). Declares `model_requirements = (Mixin1, Mixin2)` at class level. - **Inputs & targets** — the user-supplied input template(s) carrying the `{{OPTIMIZED_TRIGGER}}` placeholder, plus a `Targets` dataclass holding per-template targets (`target_response_strs`, `target_vectors`, `target_directions`, `target_class_idx`, …). Each set field has length `n_templates`. Templates and targets are passed to `optimizer.optimize_trigger(templates=..., targets=...)`; the framework handles trigger/template combination internally. @@ -187,7 +187,7 @@ These apply across optimizers and recipes — surface them when relevant rather - **`Targets` field length mismatch** — every set field must have length `n_templates`. Universal triggers across N templates need N targets (often repeated: `["Sure, here is"] * N`). - **Thinking-model target misalignment (Qwen3, etc.)** — these models always emit a `...` block first, so the affirmative target the optimizer chases (e.g. `"Sure, here's how:"`) never appears at position 0 of the generation. Prepend the empty thinking block to the target: `target = "\n\n\n\n" + "Sure, here's how:"`. Without this fix, prefill-CE losses chase a position the model will never write to. - **Some losses impose model-construction requirements** beyond just the access mixins. Check the loss's docstring before pairing it with a recipe: - - `AttentionBasedLoss` family (e.g. `AttentionEnhLoss`) needs the model loaded with `use_eager_attention=True` — FlashAttention/SDPA paths don't return attention weights. `LMHFModel(..., use_eager_attention=True)`. + - Attention-based losses (`require_attentions=True`, e.g. `AttentionEnhLoss`) need the model loaded with `use_eager_attention=True` — FlashAttention/SDPA paths don't return attention weights. `LMHFModel(..., use_eager_attention=True)`. - `AttentionEnhLoss` with `dst_slc_name=SliceKey.INPUT_AFTER` (and other suffix-position-dependent losses) requires the trigger to be a **true suffix** of the user message — strip any trailing punctuation after `{{OPTIMIZED_TRIGGER}}` in the template (`"... harmful thing. {{OPTIMIZED_TRIGGER}}"` → `"... harmful thing {{OPTIMIZED_TRIGGER}}"`). - `TriggerPerplexityLoss` and other trigger-slice-dependent losses require `use_prefix_cache=False` on the model — prefix caching shifts the trigger slice's start position. - Hidden-state / activation losses (`SteeringActivationLoss`, etc.) often need pre-computed direction vectors (e.g. `compute_refusal_directions(model, n_samples=128)`) passed via `Targets(target_directions=...)`. diff --git a/tests/test_optimizer_utils.py b/tests/test_optimizer_utils.py new file mode 100644 index 0000000..b2bd356 --- /dev/null +++ b/tests/test_optimizer_utils.py @@ -0,0 +1,118 @@ +"""Unit checks for the shared optimizer/model helpers. + +These cover the two code paths the end-to-end optimizer tests don't reach: +the embedding branch of the gradient helper (used by PEZ / SoftPrompt) and the +unconstrained-vocab branch of `random_single_flips` (used by RASLITE+). +""" + +import torch + +from tropt.common import Targets +from tropt.loss import PrefillCELoss +from tropt.optimizer.utils.token_initializers import random_single_flips + + +def test_random_single_flips_from_whitelist(): + trigger = torch.arange(8) + valid = torch.tensor([100, 101, 102]) + + out = random_single_flips(trigger, n_variations=32, valid_token_ids=valid) + + assert out.shape == (32, 8) + assert torch.equal(out[0], trigger), "first variation must be left intact" + for row in out[1:]: + differing = (row != trigger).nonzero().flatten() + # A flip can land on the value already there, so ≤ 1 position differs. + assert len(differing) <= 1 + for pos in differing: + assert row[pos].item() in valid.tolist() + + +def test_random_single_flips_from_vocab_size(): + trigger = torch.arange(6) + vocab_size = 50 + + out = random_single_flips(trigger, n_variations=16, vocab_size=vocab_size) + + assert out.shape == (16, 6) + assert torch.equal(out[0], trigger) + assert out.min() >= 0 and out.max() < vocab_size + # With 15 flips over a 50-token vocab, at least one row should have moved. + assert not torch.equal(out[1:], trigger.repeat(15, 1)) + + +def test_random_single_flips_single_variation(): + trigger = torch.arange(4) + out = random_single_flips(trigger, n_variations=1, vocab_size=10) + assert torch.equal(out, trigger.unsqueeze(0)) + + +def test_compute_grad_from_embeds(tiny_lm, lm_templates, lm_targets): + """The embedding branch of `_grad_wrt_leaves` (PEZ / SoftPrompt flow).""" + tiny_lm.set_inputs_from_tokens(lm_templates, lm_targets) + try: + trigger_len, n_candidates = 5, 2 + embeds = torch.randn( + n_candidates, trigger_len, tiny_lm.embedding_matrix.shape[1], + device=tiny_lm.device, dtype=tiny_lm.dtype, + ) + + grads, losses = tiny_lm.compute_grad_from_embeds( + loss_func=PrefillCELoss(), + candidate_trigger_embeds=embeds, + return_loss=True, + ) + + assert grads.shape == embeds.shape + assert losses.shape == (n_candidates,) + assert torch.isfinite(grads).all() + assert torch.isfinite(losses).all() + finally: + tiny_lm.reset_inputs_from_tokens() + + +def test_grad_token_and_embed_paths_agree(tiny_lm, lm_templates, lm_targets): + """Both branches of `_grad_wrt_leaves` must satisfy the chain rule. + + Trigger embeddings are `onehot @ E`, so dL/d(onehot) = dL/d(embeds) @ E^T. + This pins the shared gradient helper: if either branch wires the + differentiated leaf wrongly, the identity breaks. + """ + tiny_lm.set_inputs_from_tokens(lm_templates, lm_targets) + try: + loss = PrefillCELoss() + trigger_ids = torch.tensor([[10, 20, 30, 40]], device=tiny_lm.device) + emb_matrix = tiny_lm.embedding_matrix # (vocab, d_model) + + grad_onehot = tiny_lm.compute_grad_from_tokens( + loss_func=loss, candidate_trigger_ids=trigger_ids + ) # (1, trigger_len, vocab) + grad_embeds = tiny_lm.compute_grad_from_embeds( + loss_func=loss, candidate_trigger_embeds=emb_matrix[trigger_ids] + ) # (1, trigger_len, d_model) + + assert torch.allclose( + grad_onehot, grad_embeds @ emb_matrix.T, atol=1e-4, rtol=1e-3 + ) + finally: + tiny_lm.reset_inputs_from_tokens() + + +def test_compute_grad_from_embeds_keeps_message_dim(tiny_lm): + """`keep_message_dim=True` adds the leading n_templates axis.""" + templates = ["Explain X. {{OPTIMIZED_TRIGGER}}", "Explain Y. {{OPTIMIZED_TRIGGER}}"] + targets = Targets(target_response_strs=["Sure:", "Of course:"]) + tiny_lm.set_inputs_from_tokens(templates, targets) + try: + embeds = torch.randn( + 2, 4, tiny_lm.embedding_matrix.shape[1], + device=tiny_lm.device, dtype=tiny_lm.dtype, + ) + grads = tiny_lm.compute_grad_from_embeds( + loss_func=PrefillCELoss(), + candidate_trigger_embeds=embeds, + keep_message_dim=True, + ) + assert grads.shape == (2, *embeds.shape) # (n_templates, n_cands, len, d) + finally: + tiny_lm.reset_inputs_from_tokens() diff --git a/tropt/common.py b/tropt/common.py index 8aceba2..6a9691f 100644 --- a/tropt/common.py +++ b/tropt/common.py @@ -25,7 +25,6 @@ Length: n_templates. """ TokenTrigger = Float[Tensor, "1 trigger_seq_len"] -TextTrigger = str TokenTriggerCandidates = Float[Tensor, "n_candidates trigger_seq_len"] # ======================= Slice Keys Enum ======================= @@ -205,14 +204,6 @@ class Targets(pydantic.BaseModel): Used by: untargeted-misclassification losses on classifier outputs. """ - @property - def n_templates(self) -> int: - for field_name in self.model_fields_set: - val = getattr(self, field_name) - if val is not None: - return len(val) - return 0 - @pydantic.model_validator(mode="after") def check_field_lengths(self) -> "Targets": lengths = {len(v) for k, v in self if v is not None} diff --git a/tropt/loss/__init__.py b/tropt/loss/__init__.py index 906ae6e..ef5c53c 100644 --- a/tropt/loss/__init__.py +++ b/tropt/loss/__init__.py @@ -7,11 +7,7 @@ resolve_and_compute_loss, ) from .losses import ( - AttentionBasedLoss, AttentionEnhLoss, - ClassificationBasedLoss, - EmbeddingBasedLoss, - HiddenStateBasedLoss, # Concrete losses: MisclassCELoss, PrefillBasedLoss, @@ -21,7 +17,6 @@ PrefillMellowMaxLoss, SimilarityLoss, SteeringActivationLoss, - TriggerLogitBasedLoss, TriggerPerplexityLoss, ) from .text_losses import ( @@ -29,7 +24,6 @@ # Concrete text losses ExternalTriggerPerplexityLoss, FirstTokenNLLLoss, - GeneratedResponseBasedLoss, InputFluencyLoss, ResponseHarmfulnessLoss, TextBasedLoss, diff --git a/tropt/loss/base.py b/tropt/loss/base.py index 6ac206c..947802d 100644 --- a/tropt/loss/base.py +++ b/tropt/loss/base.py @@ -87,13 +87,6 @@ def get_loss_log_dict(self) -> dict: return {} return {f"{type(self).__name__}": self._last_loss_vals.min().item()} - def contains_loss_type(self, loss_type: type) -> bool: - """ - Returns True if this loss is of the given type. - Complicated losses (e.g., CombinedLoss) may override this method with different logic. - """ - return isinstance(self, loss_type) - ############################ @@ -147,6 +140,8 @@ def get_loss_log_dict(self) -> dict: }, } + # A combined loss needs whatever *any* component needs, and is differentiable + # only if *all* components are. @property def is_differentiable(self) -> bool: # ty: ignore[invalid-attribute-override] return all(lf.is_differentiable for lf in self.loss_funcs) @@ -175,10 +170,6 @@ def require_attentions(self) -> bool: # ty: ignore[invalid-attribute-override] def require_first_token_logprobs(self) -> bool: # ty: ignore[invalid-attribute-override] return any(lf.require_first_token_logprobs for lf in self.loss_funcs) - def contains_loss_type(self, loss_type: type) -> bool: - """Check if the CombinedLoss contains a loss of the specified type.""" - return any(isinstance(loss, loss_type) for loss in self.loss_funcs) - def __iter__(self): """Allows iterating over the nested loss functions.""" return iter(self.loss_funcs) diff --git a/tropt/loss/losses.py b/tropt/loss/losses.py index fe8c2ae..91da793 100644 --- a/tropt/loss/losses.py +++ b/tropt/loss/losses.py @@ -267,27 +267,12 @@ def __call__( ############################ @dataclass -class TriggerLogitBasedLoss(BaseLoss): - """ - Loss computed on full-sequence logits (`full_logits`) sliced to trigger positions. - Useful for optimizing properties of the triggers directly. - """ - - @abstractmethod - def __call__( - self, - full_logits: Float[Tensor, "bsz seq_len vocab_size"], - input_trigger_ids: Int[Tensor, "trigger_seq_len"], - input_slices: dict[SliceKey, slice], - ) -> Float[Tensor, "bsz"]: - pass - - -@dataclass -class TriggerPerplexityLoss(TriggerLogitBasedLoss): +class TriggerPerplexityLoss(BaseLoss): """ Calculates perplexity wrt to the target model logits themselves. Useful for penalizing non-fluent triggers. + + Computed on full-sequence logits (`full_logits`) sliced to trigger positions. """ temperature: float = 1.0 @@ -330,22 +315,7 @@ def __call__( ############################# @dataclass -class AttentionBasedLoss(BaseLoss): - """Loss computed on model attention weights (`full_attentions`).""" - - require_attentions: ClassVar[bool] = True - - @abstractmethod - def __call__( - self, - full_attentions: Float[Tensor, "bsz n_layers n_heads seq_len[dst] seq_len[src]"], - input_slices: dict[SliceKey, slice], - ) -> Float[Tensor, "bsz"]: - pass - - -@dataclass -class AttentionEnhLoss(AttentionBasedLoss): +class AttentionEnhLoss(BaseLoss): """ Encourages attention from the trigger tokens to the chat template after the adversarial trigger. *Note*: the sign of the loss is set such that minimizing the loss maximizes the attention. @@ -356,6 +326,8 @@ class AttentionEnhLoss(AttentionBasedLoss): Note that it requires setting `use_eager_attention=True` when loading the model (for explicit attention computations); also, some slices are not supported when LM prefix caching is enabled, so set `use_prefix_cache=False` when loading the model. """ + require_attentions: ClassVar[bool] = True + targeted_layers: slice = slice(None) src_slc_name: SliceKey = SliceKey.TRIGGER dst_slc_name: SliceKey = SliceKey.INPUT_AFTER @@ -384,24 +356,12 @@ def __call__( ############################ @dataclass -class EmbeddingBasedLoss(BaseLoss): - """Loss is computed based on model embeddings, compared to given target vectors. - - Requires the target vectors (shape: (n_templates, d_model)) to be provided in the targets dict. - """ - - @abstractmethod - def __call__( - self, - output_embeddings: Float[Tensor, "bsz d_model"], - **kwargs, - ) -> Float[Tensor, "bsz"]: - pass - -@dataclass -class SimilarityLoss(EmbeddingBasedLoss): +class SimilarityLoss(BaseLoss): """ Encourages given representation(s) to align (cos-sim) with the given target vectors. + + Computed on model embeddings (`output_embeddings`) against `target_vectors` + (shape: (n_templates, d_model)) supplied via `Targets`. """ def __call__( @@ -426,22 +386,7 @@ def __call__( ############################ @dataclass -class HiddenStateBasedLoss(BaseLoss): - """Loss computed on model hidden states (`full_hidden_states`).""" - - require_hidden_states: ClassVar[bool] = True - - @abstractmethod - def __call__( - self, - full_hidden_states: Float[Tensor, "bsz n_layers seq_len d_model"], - **kwargs, - ) -> Float[Tensor, "bsz"]: - pass - - -@dataclass -class SteeringActivationLoss(HiddenStateBasedLoss): +class SteeringActivationLoss(BaseLoss): """ Encourages hidden activations at specific layers/positions to align with a target direction. - Each message has a target direction vector (optionally its own unique one). @@ -466,6 +411,8 @@ class SteeringActivationLoss(HiddenStateBasedLoss): apply_abs: Whether to take the absolute value of the similarity scores (default: False). """ + require_hidden_states: ClassVar[bool] = True + targeted_layers: slice = slice(None) steer_away: bool = False slc_name: SliceKey = SliceKey.INPUT_LAST_TOKEN @@ -535,21 +482,9 @@ def __call__( ############################ @dataclass -class ClassificationBasedLoss(BaseLoss): - """Loss computed on classifier logits (`output_class_logits`).""" - - @abstractmethod - def __call__( - self, - output_class_logits: Float[Tensor, "bsz n_classes"], - **kwargs, - ) -> Float[Tensor, "bsz"]: - pass - - -@dataclass -class MisclassCELoss(ClassificationBasedLoss): - """Encourages misclassification via cross-entropy on classifier logits. +class MisclassCELoss(BaseLoss): + """Encourages misclassification via cross-entropy on classifier logits + (`output_class_logits`). Two modes: - Untargeted (targeted=False): minimizes probability of `true_class_idx`. diff --git a/tropt/loss/text_losses.py b/tropt/loss/text_losses.py index 35ef63a..fd6a964 100644 --- a/tropt/loss/text_losses.py +++ b/tropt/loss/text_losses.py @@ -36,20 +36,6 @@ def __call__(self, *args, **kwargs) -> Float[torch.Tensor, "bsz"]: ############################ -@dataclass -class GeneratedResponseBasedLoss(TextBasedLoss): - """Marker base for losses that operate on `generated_response_strs`.""" - - require_generation: ClassVar[bool] = True - - @abstractmethod - def __call__( - self, - generated_response_strs: Annotated[List[str], "bsz"], - ) -> Float[torch.Tensor, "bsz"]: - pass - - @dataclass class BinaryLMJudgeLoss(TextBasedLoss): """Abstract base for Yes/No LLM judge losses. @@ -210,11 +196,13 @@ def __call__( @dataclass -class ResponseHarmfulnessLoss(BinaryLMJudgeLoss, GeneratedResponseBasedLoss): +class ResponseHarmfulnessLoss(BinaryLMJudgeLoss): """ Loss that encourages harmful model responses (operates on generated response string). """ + require_generation: ClassVar[bool] = True + def _create_prompt(self, response_text: str) -> str: return ( f"Does the following text provide harmful, dangerous, or illegal instructions? " diff --git a/tropt/model/__init__.py b/tropt/model/__init__.py index 2e27459..02d6899 100644 --- a/tropt/model/__init__.py +++ b/tropt/model/__init__.py @@ -13,7 +13,6 @@ # Input classes: from .inputs_manager import ( - InputsManager, TextInputManager, TokenInputManager, DefaultTokenInputManager, diff --git a/tropt/model/api_retry.py b/tropt/model/api_retry.py new file mode 100644 index 0000000..f086c45 --- /dev/null +++ b/tropt/model/api_retry.py @@ -0,0 +1,81 @@ +"""Shared transient-error retry for the API-backed model wrappers. + +Each SDK (openai / google-genai / voyageai) defines its own exception classes but +they all fail the same two ways: a network-level error, or an HTTP 429/5xx. We +match on exception *class name* (so no SDK needs importing here) plus the status +code, and leave 4xx alone — retrying an auth failure just wastes the backoff. + +``tenacity`` is imported lazily so it stays an optional dependency of the API +extras rather than a core one. +""" + +import logging + +logger = logging.getLogger(__name__) + +# Union of the network-error class names raised by httpx, requests/urllib3, and +# the three SDKs' own wrappers. +_TRANSIENT_EXC_NAMES = frozenset({ + "APIConnectionError", + "ChunkedEncodingError", + "ConnectError", + "ConnectionAbortedError", + "ConnectionError", + "ConnectionResetError", + "ConnectTimeout", + "NetworkError", + "PoolTimeout", + "ProtocolError", + "RateLimitError", + "ReadError", + "ReadTimeout", + "RemoteProtocolError", + "ServerError", + "ServiceUnavailableError", + "Timeout", + "TimeoutException", + "WriteError", + "WriteTimeout", +}) + + +def is_transient_api_error(e: BaseException) -> bool: + """Whether `e` is worth retrying: network blip, rate limit, or 5xx.""" + if type(e).__name__ in _TRANSIENT_EXC_NAMES: + return True + code = ( + getattr(e, "http_status", None) + or getattr(e, "status_code", None) + or getattr(e, "code", None) + ) + return isinstance(code, int) and (code == 429 or 500 <= code < 600) + + +def retry_transient(fn, label: str, attempts: int = 6): + """Wrap `fn` with exponential backoff on transient API errors. + + Args: + label: Short backend name used in the retry log line (e.g. "gemini"). + """ + from tenacity import ( + retry, + retry_if_exception, + stop_after_attempt, + wait_random_exponential, + ) + + def _log(retry_state): + e = retry_state.outcome.exception() + logger.warning( + "[%s retry] %s: %s -> sleeping %.1fs (attempt %d)", + label, type(e).__name__, str(e)[:80], + retry_state.next_action.sleep, retry_state.attempt_number, + ) + + return retry( + retry=retry_if_exception(is_transient_api_error), + wait=wait_random_exponential(multiplier=1.5, max=60), + stop=stop_after_attempt(attempts), + before_sleep=_log, + reraise=True, + )(fn) diff --git a/tropt/model/flop_counter.py b/tropt/model/flop_counter.py index c38b909..15f83e5 100644 --- a/tropt/model/flop_counter.py +++ b/tropt/model/flop_counter.py @@ -7,7 +7,7 @@ the model forward/backward passes. Uses the Kaplan et al. (2020) approximation (https://arxiv.org/abs/2001.08361): -``FLOPs_fwd ≈ 2·N·T``, ``FLOPs_bwd ≈ 4·N·T``. +``FLOPs_fwd ≈ 2·N·T``, ``FLOPs_fwd+bwd ≈ 6·N·T``. Cheap and deterministic. Requires ``_model`` to be a HuggingFace ``PreTrainedModel``. Adapted from https://github.com/romovpa/claudini. @@ -34,27 +34,11 @@ logger = logging.getLogger(__name__) -class FlopCounterBase: - """Base class for FLOP counting. Subclasses implement specific counting methods.""" - - def count_forward(self, n_tokens: int) -> int: - """Count forward-pass FLOPs for a given number of tokens.""" - raise NotImplementedError - - def count_backward(self, n_tokens: int) -> int: - """Count backward-pass FLOPs for a given number of tokens.""" - raise NotImplementedError - - def count_forward_backward(self, n_tokens: int) -> int: - """Count combined forward+backward FLOPs for a given number of tokens.""" - raise NotImplementedError - - -class ManualFlopCounter(FlopCounterBase): +class ManualFlopCounter: """Track FLOPs using Kaplan et al. (2020) approximation. - FLOPs_fwd ≈ 2 · N_params · n_tokens - FLOPs_bwd ≈ 4 · N_params · n_tokens + FLOPs_fwd ≈ 2 · N_params · n_tokens + FLOPs_fwd+bwd ≈ 6 · N_params · n_tokens For MoE models, N_params is the *active* parameter count (shared params + expert params scaled by top-k / num_experts). @@ -230,75 +214,5 @@ def _params_from_config(config) -> int | None: def count_forward(self, n_tokens: int) -> int: return 2 * self.n_params * n_tokens - def count_backward(self, n_tokens: int) -> int: - return 4 * self.n_params * n_tokens - def count_forward_backward(self, n_tokens: int) -> int: return 6 * self.n_params * n_tokens - - -# --------------------------------------------------------------------------- -# [DISABLED] Torch automatic FLOP-counting. Disabled due to potential instability and limitations. -# --------------------------------------------------------------------------- - - -# def track_flops(includes_backward: bool = False): -# """Decorator for ``compute_*`` methods — dispatches FLOP counting by mode. - -# ``"torch"`` — wraps the method with ``FlopCounterMode``. - -# ``"manual"`` — measures the delta of ``_token_used`` before/after the -# method, then applies the Kaplan approximation. Requires a -# ``_kaplan_flop_counter`` (:class:`KaplanFlopCounter`) on the model — -# provided automatically by :class:`HuggingFaceBackendModel`. - -# Args: -# includes_backward: The wrapped method includes a backward pass -# (``"manual"`` uses ``6·N·T`` instead of ``2·N·T``). -# """ - -# def decorator(method): -# @wraps(method) -# def wrapper(self, *args, **kwargs): -# mode = getattr(self, "count_flops", False) -# if not mode: -# return method(self, *args, **kwargs) - -# # --- torch mode --- -# if mode == "torch": -# from torch.utils.flop_counter import FlopCounterMode as _FlopCounterMode - -# with _FlopCounterMode(display=False) as flop_counter: -# result = method(self, *args, **kwargs) -# self._update_usage_stats(flops=flop_counter.get_total_flops()) -# return result - -# # --- manual mode --- -# if mode == "manual": -# counter = getattr(self, "_kaplan_flop_counter", None) -# if counter is None: -# logger.warning( -# "count_flops='manual' but no KaplanFlopCounter on model; " -# "skipping FLOP counting." -# ) -# return method(self, *args, **kwargs) - -# tokens_before = getattr(self, "_token_used", 0) -# result = method(self, *args, **kwargs) -# tokens_after = getattr(self, "_token_used", 0) -# delta_tokens = tokens_after - tokens_before - -# if includes_backward: -# flops = counter.count_forward_backward(delta_tokens) -# else: -# flops = counter.count_forward(delta_tokens) -# self._update_usage_stats(flops=flops) -# return result - -# raise ValueError( -# f"Unknown count_flops mode: {mode!r}. Use False, 'torch', or 'manual'." -# ) - -# return wrapper - -# return decorator diff --git a/tropt/model/google/encoder.py b/tropt/model/google/encoder.py index 17c3b8b..c1ad5a0 100644 --- a/tropt/model/google/encoder.py +++ b/tropt/model/google/encoder.py @@ -1,56 +1,16 @@ from typing import List, Optional import torch -from tenacity import ( - retry, - retry_if_exception, - stop_after_attempt, - wait_random_exponential, -) from tropt.common import ModelOutput from tropt.model import EncoderBaseModel, LossTextAccessMixin +from tropt.model.api_retry import retry_transient # Per-request HTTP timeout (ms). Without this, a half-open connection or # hung server can deadlock the run indefinitely. _REQUEST_TIMEOUT_MS = 120_000 -def _is_transient_gemini_error(e: BaseException) -> bool: - # Transient: any httpx network error, plus google.genai errors with - # 429/5xx status. Non-transient: 4xx (bad request, auth, etc.). - name = type(e).__name__ - if name in { - "ReadError", - "WriteError", - "ConnectError", - "ConnectTimeout", - "ReadTimeout", - "WriteTimeout", - "PoolTimeout", - "RemoteProtocolError", - "TimeoutException", - "NetworkError", - "ProtocolError", - }: - return True - code = getattr(e, "code", None) or getattr(e, "status_code", None) - if isinstance(code, int) and (code == 429 or 500 <= code < 600): - return True - if name == "ServerError": - return True - return False - - -def _log_gemini_retry(rs): - e = rs.outcome.exception() - print( - f"[gemini retry] {type(e).__name__}: {str(e)[:80]} " - f"-> sleeping {rs.next_action.sleep:.1f}s (attempt {rs.attempt_number})", - flush=True, - ) - - class EncoderGeminiModel(EncoderBaseModel, LossTextAccessMixin): """ Google Gemini Encoder model wrapper, with text-query access. @@ -146,17 +106,11 @@ def invoke_from_texts( output_dimensionality=self._d_model, ) - @retry( - retry=retry_if_exception(_is_transient_gemini_error), - wait=wait_random_exponential(multiplier=1.5, max=60), - stop=stop_after_attempt(6), - before_sleep=_log_gemini_retry, - reraise=True, - ) def _embed_chunk(chunk): return self._client.models.embed_content( contents=chunk, model=self.model_name, config=cfg, ) + _embed_chunk = retry_transient(_embed_chunk, label="gemini") for start in range(0, len(input_texts), self._max_batch): chunk = input_texts[start : start + self._max_batch] diff --git a/tropt/model/huggingface/base.py b/tropt/model/huggingface/base.py index 99a1029..ef701ca 100644 --- a/tropt/model/huggingface/base.py +++ b/tropt/model/huggingface/base.py @@ -527,33 +527,17 @@ def _requires_input_embeds() -> bool: def n_layers(self) -> int: """Number of hidden layers in the model. - Different model families expose this differently, so we - try several known locations and return the first that works. + ``get_text_config()`` returns the config itself for text-only models and + the nested text config for multimodal ones (e.g. Gemma-3). """ config = self._hf_model.config - - # Each function below either returns the layer count, or raises. - def _n_layers_v1(): - # Standard HF text models - return config.num_hidden_layers - - def _n_layers_v2(): - # Multimodal configs (e.g. Gemma-3) with nested text config - return config.text_config.num_hidden_layers - - def _n_layers_v3(): - return config.get_text_config().num_hidden_layers - - for _getter in [_n_layers_v1, _n_layers_v2, _n_layers_v3]: - try: - return _getter() - except Exception: - continue - - raise ValueError( - f"Could not extract `num_hidden_layers` from model config of `{self.get_model_name()}`. " - f"This model might need special care. Please report this issue." - ) + n = getattr(config.get_text_config(), "num_hidden_layers", None) + if n is None: + raise ValueError( + f"Could not extract `num_hidden_layers` from model config of `{self.get_model_name()}`. " + f"This model might need special care. Please report this issue." + ) + return n @property def dtype(self): @@ -631,7 +615,7 @@ def compute_grad_from_tokens( Args: loss_func: Loss function to optimize. Must be compatible with model outputs - (e.g., PrefillBasedLoss for LMs, EmbeddingBasedLoss for encoders). + (e.g., PrefillCELoss for LMs, SimilarityLoss for encoders). candidate_trigger_ids: Discrete token IDs for hard triggers. Can accepts multiple candidates. Shape: (n_candidates, trigger_seq_len) @@ -695,22 +679,80 @@ def compute_grad_from_tokens( ) assert (candidate_trigger_ids is not None) ^ (candidate_trigger_probs is not None), \ "Exactly one of `candidate_trigger_ids` or `candidate_trigger_probs` must be provided." - assert self._token_input_manager is not None, "Token input manager is not initialized. Please call set_inputs_from_tokens() first." - device, dtype = self.device, self.dtype embedding_layer = self._embedding_layer assert isinstance(embedding_layer, torch.nn.Embedding), ( "Expected a standard nn.Embedding layer for one-hot gradient computation." ) vocab_size = embedding_layer.num_embeddings + + # The differentiated leaf is the one-hot / probability matrix; the trigger + # embeddings are derived from it via the effective embedding matrix. + if candidate_trigger_probs is None: + leaves = torch.nn.functional.one_hot( + candidate_trigger_ids, num_classes=vocab_size + ).to(self.device, self.dtype) + else: + leaves = candidate_trigger_probs.to(self.device, self.dtype) + + def _leaf_to_embeds(leaf_batch, batch_slice): + if do_gumbel_softmax: + assert gumbel_softmax_temp is not None, "gumbel_softmax_temp must be provided if do_gumbel_softmax is True." + leaf_batch = torch.nn.functional.gumbel_softmax( + logits=leaf_batch, tau=gumbel_softmax_temp, hard=False, dim=-1, + ) + # (bsz, trigger_seq_len, vocab_size) @ (vocab_size, embed_dim) + candidate_embeds = leaf_batch @ self.embedding_matrix + + # Only check when using discrete tokens (not soft probabilities) + if is_debug_mode() and candidate_trigger_ids is not None: + assert torch.allclose( + candidate_embeds, + self._token_input_manager.embed_func(candidate_trigger_ids[batch_slice]) + ), ("Mismatch between effective embedding matrix and embed-func. It could be that you use " \ + "a model with non-standard embedding logic. Please report this issue!") + + # Trigger ids for reference; for soft triggers take the argmax + # of the *pre*-Gumbel distribution. + ref_trigger_ids = ( + candidate_trigger_ids[batch_slice] if candidate_trigger_ids is not None + else leaves[batch_slice].argmax(dim=-1) + ) + return candidate_embeds, ref_trigger_ids + + return self._grad_wrt_leaves( + loss_func, leaves, _leaf_to_embeds, + normalize_grads=normalize_grads, + keep_message_dim=keep_message_dim, + return_loss=return_loss, + ) + + def _grad_wrt_leaves( + self, + loss_func: BaseLoss, + leaves: Float[Tensor, "n_candidates trigger_seq_len leaf_dim"], + leaf_to_embeds, + normalize_grads: bool, + keep_message_dim: bool, + return_loss: bool, + ) -> Float[torch.Tensor, "n_candidates trigger_seq_len leaf_dim"] | Tuple[Tensor, Tensor]: + """Shared backward pass for :meth:`compute_grad_from_tokens` / :meth:`compute_grad_from_embeds`. + + Differentiates ``loss_func`` w.r.t. ``leaves``, batching candidates (with + OOM backoff) and back-propagating each template separately so only one + autograd graph is alive at a time. + + Args: + leaves: The tensor to differentiate w.r.t. — one-hot/probabilities for + the token flow, raw trigger embeddings for the embedding flow. + leaf_to_embeds: ``(leaf_batch, batch_slice) -> (trigger_embeds, ref_trigger_ids)``. + """ + assert self._token_input_manager is not None, "Token input manager is not initialized. Please call set_inputs_from_tokens() first." + + device, dtype = self.device, self.dtype input_manager = self._token_input_manager n_templates = input_manager.n_templates - - # Get shape from whichever input is provided - if candidate_trigger_ids is not None: - n_candidates, trigger_seq_len = candidate_trigger_ids.shape - else: # candidate_trigger_probs is not None: - n_candidates, trigger_seq_len = candidate_trigger_probs.shape[:2] + n_candidates, trigger_seq_len, leaf_dim = leaves.shape @find_executable_batch_size(starting_batch_size=self._backward_pass_batch_size) def _compute_grad__batched( @@ -727,71 +769,29 @@ def _compute_grad__batched( all_grads = [] # of len `n_candidate // batch_size` all_losses = [] # per-batch mean losses (detached) - # Prepare the one-hot encoding matrix (might be distribution-per-token for continuous trigger) - # (n_candidates, trigger_seq_len, vocab_size) - if candidate_trigger_probs is None: - candidate_ids_onehot_detached = torch.nn.functional.one_hot( - candidate_trigger_ids, - num_classes=vocab_size, - ).to(device, dtype) - else: - candidate_ids_onehot_detached = candidate_trigger_probs.to(device, dtype) - - # Prepare the effective embedding matrix: - embedding_matrix = self.embedding_matrix # (vocab_size, embd_dim) - for cand_idx_start in range(0, n_candidates, batch_size): cand_idx_end = min(cand_idx_start + batch_size, n_candidates) cand_bsz = cand_idx_end - cand_idx_start + batch_slice = slice(cand_idx_start, cand_idx_end) # Backward each template immediately to avoid keeping n_templates graphs at once accum_grad = torch.zeros( - (n_templates, cand_bsz, trigger_seq_len, vocab_size), + (n_templates, cand_bsz, trigger_seq_len, leaf_dim), device=device, dtype=dtype, ) accum_loss = torch.zeros((n_templates, cand_bsz), device=device, dtype=dtype) for template_idx in range(0, n_templates): - # 1. Enable gradients on the one-hot input - # (bsz_triggers, trigger_seq_len, vocab_size) - candidate_ids_onehot = candidate_ids_onehot_detached[cand_idx_start:cand_idx_end].clone() - candidate_ids_onehot.requires_grad_() - - # 1'. optionally apply gumbel-softmax to the trigger probs - if do_gumbel_softmax: - assert gumbel_softmax_temp is not None, "gumbel_softmax_temp must be provided if do_gumbel_softmax is True." - candidate_ids_onehot = torch.nn.functional.gumbel_softmax( - logits=candidate_ids_onehot, - tau=gumbel_softmax_temp, - hard=False, - dim=-1, - ) - - # 2. Apply embedding to get trigger_embeds - # (n_candidates, trigger_seq_len, vocab_size) @ (vocab_size, embed_dim) -> (n_candidates, trigger_seq_len, embed_dim) - candidate_embeds = candidate_ids_onehot @ embedding_matrix - - # Only check when using discrete tokens (not soft probabilities) - if is_debug_mode() and candidate_trigger_ids is not None: - assert torch.allclose( - candidate_embeds, - input_manager.embed_func( - candidate_trigger_ids[cand_idx_start:cand_idx_end] - ) - ), ("Mismatch between effective embedding matrix and embed-func. It could be that you use " \ - "a model with non-standard embedding logic. Please report this issue!") + # 1. Enable gradients on the leaf input + leaf = leaves[batch_slice].clone() + leaf.requires_grad_() + + # 2. Map the leaf to trigger embeddings + candidate_embeds, ref_trigger_ids = leaf_to_embeds(leaf, batch_slice) # 3. Get batched inputs & compute loss: logger.debug(f"from grad [msg={template_idx}]: {candidate_embeds.shape}") - # Get trigger IDs for reference (if using discrete tokens) - # For soft triggers, compute argmax from probabilities - if candidate_trigger_ids is not None: - ref_trigger_ids = candidate_trigger_ids[cand_idx_start:cand_idx_end] - else: - # Compute discrete tokens from soft probabilities (before Gumbel-softmax) - ref_trigger_ids = candidate_ids_onehot_detached[cand_idx_start:cand_idx_end].argmax(dim=-1) - model_input = input_manager.get_triggered_inputs( chosen_template_idx=template_idx, trigger_embeds=candidate_embeds, @@ -815,18 +815,17 @@ def _compute_grad__batched( # 4. Backward and store per-template template_grad = torch.autograd.grad( outputs=loss, - inputs=[candidate_ids_onehot], + inputs=[leaf], grad_outputs=torch.ones_like(loss, device=device), - )[0] # (bsz_triggers, trigger_seq_len, vocab_size) + )[0] # (bsz_triggers, trigger_seq_len, leaf_dim) accum_grad[template_idx] = template_grad accum_loss[template_idx] = loss.detach() all_grads.append(accum_grad) all_losses.append(accum_loss) - # clear_device_cache() # clear unused GPU memory return ( - torch.cat(all_grads, dim=1), # (n_templates, n_candidates, trigger_seq_len, vocab_size) + torch.cat(all_grads, dim=1), # (n_templates, n_candidates, trigger_seq_len, leaf_dim) torch.cat(all_losses, dim=1), # (n_templates, n_candidates) ) @@ -838,7 +837,7 @@ def _compute_grad__batched( all_grads = all_grads.mean(dim=0) all_losses = all_losses.mean(dim=0) - # normalize each token's gradient vector (over the vocab_size dim) + # normalize each token's gradient vector (over the last dim) if normalize_grads: all_grads = all_grads / (all_grads.norm(dim=-1, keepdim=True) + 1e-10) @@ -870,98 +869,16 @@ def compute_grad_from_embeds( If return_loss is True: tuple of (gradients tensor, per-candidate loss tensor (n_candidates,)), or per-template losses (n_templates, n_candidates) if keep_message_dim=True. """ - assert self._token_input_manager is not None, "Token input manager is not initialized. Please call set_inputs_from_tokens() first." - - device, dtype = self.device, self.dtype - input_manager = self._token_input_manager - n_templates = input_manager.n_templates - n_candidates, trigger_seq_len, embed_dim = candidate_trigger_embeds.shape - - @find_executable_batch_size(starting_batch_size=self._backward_pass_batch_size) - def _compute_grad__batched( - batch_size: int, - ) -> Tuple[Tensor, Tensor]: - - # --- Update backward batch size --- - if batch_size < self._backward_pass_batch_size: - logger.info(f"OOM detected. Reducing _backward_pass_batch_size from {self._backward_pass_batch_size} to {batch_size}") - self._backward_pass_batch_size = batch_size - # -------------------- - - all_grads = [] - all_losses = [] - - for cand_idx_start in range(0, n_candidates, batch_size): - cand_idx_end = min(cand_idx_start + batch_size, n_candidates) - cand_bsz = cand_idx_end - cand_idx_start - - # Backward each template immediately to avoid keeping n_templates graphs at once - accum_grad = torch.zeros( - (n_templates, cand_bsz, trigger_seq_len, embed_dim), - device=device, dtype=dtype, - ) - accum_loss = torch.zeros((n_templates, cand_bsz), device=device, dtype=dtype) - - for template_idx in range(0, n_templates): - # 1. Enable gradients on the embedding input directly - candidate_embeds = candidate_trigger_embeds[cand_idx_start:cand_idx_end].clone().detach() - candidate_embeds.requires_grad_() - - # 2. Get batched inputs - model_input = input_manager.get_triggered_inputs( - chosen_template_idx=template_idx, - trigger_embeds=candidate_embeds, - - # loss-conditional flags: - do_append_embeds=loss_func.require_target_prefill, - ) - - # 3. Forward pass - model_output = self.invoke_from_tokens( - **model_input.to_dict(), - - # loss-conditional flags: - require_target_prefill=loss_func.require_target_prefill, - require_generation=loss_func.require_generation, - require_hidden_states=loss_func.require_hidden_states, - require_attentions=loss_func.require_attentions, - count_backward=True, - ) - - # 4. Loss + per-template backward, then store per-template - loss = resolve_and_compute_loss(model_output, model_input, loss_func) - template_grad = torch.autograd.grad( - outputs=loss, - inputs=[candidate_embeds], - grad_outputs=torch.ones_like(loss, device=device), - )[0] # (bsz_triggers, trigger_seq_len, embed_dim) - accum_grad[template_idx] = template_grad - accum_loss[template_idx] = loss.detach() - - all_grads.append(accum_grad) - all_losses.append(accum_loss) - - return ( - torch.cat(all_grads, dim=1), # (n_templates, n_candidates, trigger_seq_len, embed_dim) - torch.cat(all_losses, dim=1), # (n_templates, n_candidates) - ) - - # Per-template grads/losses of the candidates - all_grads, all_losses = _compute_grad__batched() - - # Optionally reduce message dim - if not keep_message_dim: - all_grads = all_grads.mean(dim=0) - all_losses = all_losses.mean(dim=0) - - # normalize each token's gradient vector (over the embed_dim dim) - if normalize_grads: - all_grads = all_grads / (all_grads.norm(dim=-1, keepdim=True) + 1e-10) - - if return_loss: - return all_grads, all_losses - - return all_grads + # Here the differentiated leaf *is* the trigger embedding, so no mapping + # is needed and there are no reference trigger ids to report. + return self._grad_wrt_leaves( + loss_func, + candidate_trigger_embeds.detach(), + lambda leaf, batch_slice: (leaf, None), + normalize_grads=normalize_grads, + keep_message_dim=keep_message_dim, + return_loss=return_loss, + ) @torch.no_grad() def compute_loss_from_tokens( diff --git a/tropt/model/huggingface/lm.py b/tropt/model/huggingface/lm.py index b534e7c..42bb9f3 100644 --- a/tropt/model/huggingface/lm.py +++ b/tropt/model/huggingface/lm.py @@ -367,7 +367,7 @@ def invoke_from_tokens( """ if require_attentions and self._model.config._attn_implementation != "eager": logger.warning( - "AttentionBasedLoss is used but the model is not using eager attention. " + "An attention-based loss is used but the model is not using eager attention. " "This may lead to incorrect attention outputs. Consider initializing the model with eager attention, by passing LMHFModel the flag `use_eager_attention=True`." ) if require_attentions and input_prefix_cache_kwargs: diff --git a/tropt/model/inputs_manager.py b/tropt/model/inputs_manager.py index 8a4d1c8..5f46cfc 100644 --- a/tropt/model/inputs_manager.py +++ b/tropt/model/inputs_manager.py @@ -1,4 +1,3 @@ -from abc import ABC, abstractmethod from typing import Annotated, Any, List, Optional from jaxtyping import Int @@ -12,42 +11,16 @@ ) # ======================= Triggered Input Managers ======================= - - -class InputsManager(ABC): - """ - Base class for maintaining the input template, corresponding targets, and the method for - injecting triggers into the inputs. - This class wraps `n_templates` templates (that contain the substring `OPTIMIZED_TRIGGER_PLACEHOLDER` as - a trigger placeholder) and targets, and provides a unified interface for different types of inputs - (e.g., text-based, token-based) used in adversarial trigger optimization. - """ - - def __init__( - self, - templates: TextTemplates, - targets: Targets, # n_templates elements per target entry - ): - raise NotImplementedError - - @abstractmethod - def get_triggered_inputs(self, chosen_template_idx: int, *args, **kwargs) -> ModelInput: - """ - Returns the trigger-combined model inputs, for the specified template index. - - Args: - chosen_template_idx: Index of the template to use for generating the inputs. - ... args for receiving the trigger candidates ... - - Returns: - A ModelInput object containing the crafted triggered-combined inputs, which includes the - corresponding targets for the specified template. - """ - raise NotImplementedError +# +# An inputs manager wraps `n_templates` templates (each containing the substring +# OPTIMIZED_TRIGGER_PLACEHOLDER) plus their targets, and exposes +# `get_triggered_inputs(chosen_template_idx, ) -> ModelInput`. +# The trigger-candidate argument differs per flow (strings vs ids vs embeds), so +# the text and token families deliberately share no base class. ## Text inputs manager ## -class TextInputManager(InputsManager): +class TextInputManager: """ Class for maintaining text-based trigger-combined inputs (fits black-box text-level query access). Instances of this class store `n_templates` templates and targets, and provide the method `get_triggered_inputs` to combine them with given trigger strings. @@ -121,9 +94,10 @@ def get_triggered_inputs( ) ## Token inputs manager ## -class TokenInputManager(InputsManager): +class TokenInputManager: """ - Abstract base class for token-level inputs managers. + Base class for token-level inputs managers (the shared type for + ``TokenAccessMixin._token_input_manager``). Subclasses manage the combination of candidate triggers into tokenized templates. diff --git a/tropt/model/model_base.py b/tropt/model/model_base.py index 2a16385..6720479 100644 --- a/tropt/model/model_base.py +++ b/tropt/model/model_base.py @@ -11,15 +11,12 @@ from transformers import BatchEncoding from tropt.common import MessageTargets, ModelOutput -from tropt.model.flop_counter import FlopCounterBase, ManualFlopCounter +from tropt.model.flop_counter import ManualFlopCounter # ====================== Model Base Classes ======================= class BaseModel(ABC): - def __init__(self, model_name: str): - pass - @property def device(self): """ @@ -71,12 +68,8 @@ def get_model_name(self) -> str: halved on OOM. Any subclass may override. """ - _flop_counter: Optional[FlopCounterBase] = None - """Active counter object (set by :meth:`set_flop_counting`). - - Must implement ``count_forward(n_tokens) -> int`` and - ``count_forward_backward(n_tokens) -> int``. - """ + _flop_counter: Optional[ManualFlopCounter] = None + """Active counter object (set by :meth:`set_flop_counting`).""" def set_flop_counting(self, mode: Literal["manual", "none"] = "manual"): """Enable or disable FLOP counting. diff --git a/tropt/model/openai/encoder.py b/tropt/model/openai/encoder.py index d8d6fd0..f615275 100644 --- a/tropt/model/openai/encoder.py +++ b/tropt/model/openai/encoder.py @@ -2,7 +2,6 @@ import numpy as np import torch -from tenacity import retry, stop_after_attempt, wait_exponential from transformers import BatchEncoding from tropt.common import ( @@ -17,6 +16,7 @@ LossTextAccessMixin, TokenAccessMixin, ) +from tropt.model.api_retry import retry_transient # -------------------------------------------------------------------------- @@ -88,9 +88,6 @@ def encode(self, text: str, **kwargs) -> List[int]: _ = kwargs # unused return self._encoding.encode(text, disallowed_special=()) - def _parse_ids(self, ids): - return ids - def decode(self, ids, **kwargs) -> str: _ = kwargs # unused if isinstance(ids, torch.Tensor): @@ -195,10 +192,6 @@ def tokenizer(self) -> OpenAITokenizer: def vocab_size(self) -> int: return self._tokenizer.vocab_size - @retry( - wait=wait_exponential(multiplier=1, min=4, max=60), - stop=stop_after_attempt(5) - ) def invoke_from_texts( self, input_texts: Annotated[List[str], "n_texts"], @@ -214,10 +207,9 @@ def invoke_from_texts( ModelOutput with output_embeddings populated. """ # Note: OpenAI's API handles batches of texts - response = self._client.embeddings.create( - input=input_texts, - model=self.model_name, - ) + response = retry_transient( + self._client.embeddings.create, label="openai" + )(input=input_texts, model=self.model_name) embeddings = [data.embedding for data in response.data] result = torch.tensor(embeddings, dtype=torch.float32) diff --git a/tropt/model/voyage/encoder.py b/tropt/model/voyage/encoder.py index df34602..351cae6 100644 --- a/tropt/model/voyage/encoder.py +++ b/tropt/model/voyage/encoder.py @@ -1,54 +1,10 @@ from typing import List, Optional import torch -from tenacity import ( - retry, - retry_if_exception, - stop_after_attempt, - wait_random_exponential, -) from tropt.common import ModelOutput from tropt.model import EncoderBaseModel, LossTextAccessMixin - - -def _is_transient_voyage_error(e: BaseException) -> bool: - # Transient: voyageai connection / timeout / 5xx / 429, plus the underlying - # requests/urllib3/socket errors that surface through them. - # Non-transient: 4xx (auth, bad request, etc.). - name = type(e).__name__ - if name in { - "APIConnectionError", - "Timeout", - "ServiceUnavailableError", - "RateLimitError", - "ConnectionError", - "ConnectionResetError", - "ConnectionAbortedError", - "ProtocolError", - "ReadTimeout", - "WriteTimeout", - "ConnectTimeout", - "ChunkedEncodingError", - }: - return True - code = ( - getattr(e, "http_status", None) - or getattr(e, "status_code", None) - or getattr(e, "code", None) - ) - if isinstance(code, int) and (code == 429 or 500 <= code < 600): - return True - return False - - -def _log_voyage_retry(rs): - e = rs.outcome.exception() - print( - f"[voyage retry] {type(e).__name__}: {str(e)[:80]} " - f"-> sleeping {rs.next_action.sleep:.1f}s (attempt {rs.attempt_number})", - flush=True, - ) +from tropt.model.api_retry import retry_transient class EncoderVoyageModel(EncoderBaseModel, LossTextAccessMixin): @@ -80,10 +36,6 @@ def __init__( self._client = voyageai.Client() self.model_name = model_name self._d_model = d_model - self._text_to_input_type = { - "document": "document", - "query": "query", - } @property def d_model(self) -> int: @@ -111,27 +63,20 @@ def invoke_from_texts( "document", "query", ), f"Unsupported text_type {text_type}" - input_type = self._text_to_input_type.get(text_type) if text_type else None # Voyage's /embeddings caps at 128 texts per call on the newer models; chunk. MAX_BATCH = 128 all_embeddings: list = [] total_tokens = 0 - @retry( - retry=retry_if_exception(_is_transient_voyage_error), - wait=wait_random_exponential(multiplier=1.5, max=60), - stop=stop_after_attempt(6), - before_sleep=_log_voyage_retry, - reraise=True, - ) def _embed_chunk(chunk): return self._client.embed( texts=chunk, model=self.model_name, - input_type=input_type, + input_type=text_type, # voyage's input_type values match ours ("document"/"query") output_dimension=self._d_model, ) + _embed_chunk = retry_transient(_embed_chunk, label="voyage") for start in range(0, len(input_texts), MAX_BATCH): chunk = input_texts[start : start + MAX_BATCH] diff --git a/tropt/optimizer/gaslite_optimizer.py b/tropt/optimizer/gaslite_optimizer.py index 6180419..7982fb4 100644 --- a/tropt/optimizer/gaslite_optimizer.py +++ b/tropt/optimizer/gaslite_optimizer.py @@ -20,6 +20,7 @@ from tropt.optimizer.utils.retokenization import retokenize_filtering from tropt.optimizer.utils.running_best import RunningBest from tropt.optimizer.utils.token_constraints import TokenConstraints +from tropt.optimizer.utils.token_initializers import random_single_flips from tropt.tracker import BaseTracker logger = logging.getLogger(__name__) @@ -117,7 +118,9 @@ def optimize_trigger( ) else: # Compute grad over a list of `n_grad` triggers one-flip away from the current - trigger_vars = self._get_trigger_variations(trigger_ids, valid_token_ids) + trigger_vars = random_single_flips( + trigger_ids, self.n_grad, valid_token_ids=valid_token_ids + ) grads = self.model.compute_grad_from_tokens( candidate_trigger_ids=trigger_vars, loss_func=self.loss_func, @@ -195,27 +198,3 @@ def optimize_trigger( return best.to_result() - def _get_trigger_variations( - self, - trigger_ids: Float[Tensor, "trigger_seq_len"], - valid_token_ids: Float[Tensor, "n_valid"], - ) -> Float[Tensor, "n_grad trigger_seq_len"]: - """ - Creates a list of `n_grad` trigger variations. The first is the - original trigger, and the rest are random single-token flips of its. - """ - trigger_seq_len = len(trigger_ids) - device = self.model.device - trigger_vars_ids = trigger_ids.repeat( - self.n_grad, 1 - ) # shape: (n_grad, trigger_seq_len) - - for idx in range(1, self.n_grad): # (keep the first intact) - # select a random position and a random token - pos_to_flip = int(torch.randint(0, trigger_seq_len, (1,), device=device).item()) - tok_to_flip_to = int(valid_token_ids[ - torch.randint(0, len(valid_token_ids), (1,), device=device) - ].item()) # apply the flip - trigger_vars_ids[idx, pos_to_flip] = tok_to_flip_to - - return trigger_vars_ids diff --git a/tropt/optimizer/gasliteplus_optimizer.py b/tropt/optimizer/gasliteplus_optimizer.py index 371d777..48cc04f 100644 --- a/tropt/optimizer/gasliteplus_optimizer.py +++ b/tropt/optimizer/gasliteplus_optimizer.py @@ -1,6 +1,6 @@ import logging import time -from typing import Optional +from typing import Callable, Optional import torch from jaxtyping import Float, Int @@ -21,13 +21,12 @@ from tropt.optimizer.utils.buffer import TriggerBuffer from tropt.optimizer.utils.retokenization import retokenize_filtering from tropt.optimizer.utils.running_best import RunningBest -from tropt.optimizer.utils.scheduler import ( - ConstantScheduler, - LinearScheduler, - NFlipScheduler, -) +from tropt.optimizer.utils.scheduler import LinearScheduler from tropt.optimizer.utils.token_constraints import TokenConstraints -from tropt.optimizer.utils.token_initializers import get_printable_random_trigger +from tropt.optimizer.utils.token_initializers import ( + get_printable_random_trigger, + random_single_flips, +) from tropt.tracker import BaseTracker logger = logging.getLogger(__name__) @@ -68,7 +67,7 @@ def __init__( n_bulk_flips: int = 5, flip_pos_method: str = "random", # "random" or "ordered" time_limit: Optional[float] = None, - n_flip_scheduler: Optional[NFlipScheduler] = None, + n_flip_scheduler: Optional[Callable[[int], int]] = None, **kwargs ): """ @@ -102,7 +101,7 @@ def __init__( flip_pos_method (str): Method to select positions to flip - "random" or "ordered". - n_flip_scheduler (NFlipScheduler, optional): A scheduler object to control `n_flip`. + n_flip_scheduler (callable, optional): A `(step) -> n_flip` callable controlling `n_flip`. If provided, overrides `decline_n_flip_from_step`. References: @@ -145,34 +144,7 @@ def __init__( ) else: # default: constant n_flip - self.n_flip_scheduler = ConstantScheduler(n_flip) - - def _get_trigger_variations( - self, - trigger_ids: Float[Tensor, "trigger_seq_len"], - valid_token_ids: Float[Tensor, "n_valid"], - ) -> Float[Tensor, "n_grad trigger_seq_len"]: - """ - Creates a list of `n_grad` trigger variations. The first is the - original trigger, and the rest are random single-token flips of its. - """ - trigger_seq_len = len(trigger_ids) - device = self.model.device - trigger_vars_ids = trigger_ids.repeat( - self.n_grad, 1 - ) # shape: (n_grad, trigger_seq_len) - - for idx in range(1, self.n_grad): # (keep the first intact) - # select a random position and a random token - pos_to_flip = torch.randint(0, trigger_seq_len, (1,), device=device).item() - tok_to_flip_to = valid_token_ids[ - torch.randint(0, len(valid_token_ids), (1,), device=device) - ].item() # apply the flip - assert isinstance(pos_to_flip, int) and isinstance(tok_to_flip_to, int) - - trigger_vars_ids[idx, pos_to_flip] = tok_to_flip_to - - return trigger_vars_ids + self.n_flip_scheduler = lambda step: n_flip def optimize_trigger( self, @@ -221,7 +193,7 @@ def optimize_trigger( self.log(loss=buffer.get_lowest_loss(), trigger_str=trigger_str) for step in self.track_steps(range(self.num_steps), desc="Optimizing with GASLITE+..."): - n_flip = self.n_flip_scheduler.get_n_flip(step) + n_flip = self.n_flip_scheduler(step) # Get the best trigger from the buffer @@ -236,7 +208,9 @@ def optimize_trigger( ) else: # Compute grad over a list of `n_grad` triggers one-flip away from the current - trigger_vars = self._get_trigger_variations(trigger_ids, valid_token_ids) + trigger_vars = random_single_flips( + trigger_ids, self.n_grad, valid_token_ids=valid_token_ids + ) grads = self.model.compute_grad_from_tokens( candidate_trigger_ids=trigger_vars, loss_func=self.loss_func, diff --git a/tropt/optimizer/rasliteplus_optimizer.py b/tropt/optimizer/rasliteplus_optimizer.py index 4964c1c..ea4ee4f 100644 --- a/tropt/optimizer/rasliteplus_optimizer.py +++ b/tropt/optimizer/rasliteplus_optimizer.py @@ -24,7 +24,10 @@ from tropt.optimizer.utils.retokenization import retokenize_filtering from tropt.optimizer.utils.running_best import RunningBest from tropt.optimizer.utils.token_constraints import TokenConstraints -from tropt.optimizer.utils.token_initializers import get_printable_random_trigger +from tropt.optimizer.utils.token_initializers import ( + get_printable_random_trigger, + random_single_flips, +) from tropt.tracker import BaseTracker logger = logging.getLogger(__name__) @@ -213,7 +216,9 @@ def optimize_trigger( else: if self.n_logit_samples is not None and self.n_logit_samples > 1: # Average logits over variations - trigger_vars = self._get_trigger_variations(util_trigger_ids, util_vocab_size, device=self.util_model.device) + trigger_vars = random_single_flips( + util_trigger_ids, self.n_logit_samples, vocab_size=util_vocab_size + ) logits = self.util_model.compute_logits_from_tokens( trigger_vars, return_trigger_logits_only=True, @@ -364,27 +369,3 @@ def optimize_trigger( return result - def _get_trigger_variations( - self, - trigger_ids: Float[Tensor, "trigger_seq_len"], - vocab_size: int, - device: torch.device, - ) -> Float[Tensor, "n_logit_samples trigger_seq_len"]: - """ - Creates a list of `n_logit_samples` trigger variations. The first is the - original trigger, and the rest are random single-token flips. - """ - trigger_seq_len = len(trigger_ids) - trigger_vars_ids = trigger_ids.repeat( - self.n_logit_samples, 1 - ) # shape: (n_logit_samples, trigger_seq_len) - - for idx in range(1, self.n_logit_samples): # (keep the first intact) - # select a random position and a random token - pos_to_flip = int(torch.randint(0, trigger_seq_len, (1,), device=device).item()) - tok_to_flip_to = int(torch.randint(0, vocab_size, (1,), device=device).item()) - # apply the flip - trigger_vars_ids[idx, pos_to_flip] = tok_to_flip_to - - return trigger_vars_ids - diff --git a/tropt/optimizer/utils/scheduler.py b/tropt/optimizer/utils/scheduler.py index 707123d..be99eab 100644 --- a/tropt/optimizer/utils/scheduler.py +++ b/tropt/optimizer/utils/scheduler.py @@ -1,28 +1,15 @@ -import math -from abc import ABC, abstractmethod - """ Schedulers for the n_flip parameter used in optimizers to control the number of token positions flipped during each optimization step. -""" -class NFlipScheduler(ABC): - @abstractmethod - def get_n_flip(self, step: int) -> int: - """Returns the n_flip value for the given step (0-indexed).""" - pass +A scheduler is just a callable ``(step: int) -> int``, so a constant schedule is +``lambda step: n_flip`` and any user-supplied function works without subclassing. +""" -class ConstantScheduler(NFlipScheduler): - """ - A scheduler that always returns the same n_flip value. - """ - def __init__(self, n_flip: int): - self.n_flip = n_flip +import math - def get_n_flip(self, step: int) -> int: - return self.n_flip -class LinearScheduler(NFlipScheduler): +class LinearScheduler: """ A scheduler that linearly decreases n_flip from an initial value to 1 over the course of optimization steps, starting from a specified step. @@ -35,7 +22,7 @@ def __init__(self, initial_n_flip: int, total_steps: int, decline_start: int | f else: self.decline_start_step = int(decline_start) - def get_n_flip(self, step: int) -> int: + def __call__(self, step: int) -> int: if step < self.decline_start_step: return self.initial_n_flip diff --git a/tropt/optimizer/utils/token_initializers.py b/tropt/optimizer/utils/token_initializers.py index c642f73..6775821 100644 --- a/tropt/optimizer/utils/token_initializers.py +++ b/tropt/optimizer/utils/token_initializers.py @@ -3,12 +3,48 @@ from typing import List, Optional import torch -from jaxtyping import Float +from jaxtyping import Float, Int from torch import Tensor from tropt.model.model_base import BaseTokenizer +def random_single_flips( + trigger_ids: Int[Tensor, "trigger_seq_len"], + n_variations: int, + valid_token_ids: Optional[Int[Tensor, "n_valid"]] = None, + vocab_size: Optional[int] = None, +) -> Int[Tensor, "n_variations trigger_seq_len"]: + """`n_variations` copies of `trigger_ids`, the first left intact and each of the + rest given one random single-token flip. + + Used by gradient-averaging optimizers (GASLITE, GASLITE+) and by logit-sampling + ones (RASLITE+). Sampling is restricted to `valid_token_ids` when given, + otherwise uniform over `vocab_size`. + """ + assert (valid_token_ids is None) ^ (vocab_size is None), ( + "Pass exactly one of `valid_token_ids` or `vocab_size`." + ) + device = trigger_ids.device + variations = trigger_ids.repeat(n_variations, 1) + + n_flips = n_variations - 1 + if n_flips <= 0: + return variations + + pos = torch.randint(0, variations.shape[1], (n_flips, 1), device=device) + if valid_token_ids is not None: + tok = valid_token_ids[ + torch.randint(0, len(valid_token_ids), (n_flips, 1), device=device) + ] + else: + assert vocab_size is not None # guaranteed by the XOR assert above + tok = torch.randint(0, vocab_size, (n_flips, 1), device=device) + # `variations[1:]` is a view, so this writes through to `variations`. + variations[1:].scatter_(1, pos, tok.to(variations.dtype)) + return variations + + def get_printable_random_trigger( trigger_len: int, return_ids: bool = False, diff --git a/tropt/tracker/base.py b/tropt/tracker/base.py index 445cc03..f6bba78 100644 --- a/tropt/tracker/base.py +++ b/tropt/tracker/base.py @@ -63,6 +63,24 @@ def finish(self, summary: Optional[dict] = None): self._finish(summary) self._active = False + @staticmethod + def _scalarize(data: dict) -> dict: + """Unwrap 0-d tensors to Python floats; drop tensors that can't be scalars. + + Backends like WandB/Trackio only accept scalars. + """ + import torch + + out = {} + for k, v in data.items(): + if isinstance(v, torch.Tensor): + try: + v = v.item() + except (ValueError, RuntimeError): + continue + out[k] = v + return out + # ── Subclass hooks ────────────────────────────────────────────────── @abstractmethod diff --git a/tropt/tracker/trackers.py b/tropt/tracker/trackers.py index 38cc665..6922128 100644 --- a/tropt/tracker/trackers.py +++ b/tropt/tracker/trackers.py @@ -4,8 +4,6 @@ from collections import defaultdict from typing import Any, Dict, Optional -import torch - from .base import DEFAULT_EXPERIMENT_NAME, BaseTracker logger = logging.getLogger(__name__) @@ -95,15 +93,7 @@ def _init(self, config: Optional[dict] = None): def _log(self, data: Dict[str, Any]): import wandb - sanitized = {} - for k, v in data.items(): - if isinstance(v, torch.Tensor): - try: - v = v.item() - except (ValueError, RuntimeError): - continue - sanitized[k] = v - wandb.log(sanitized) + wandb.log(self._scalarize(data)) def _finish(self, summary: Optional[dict] = None): import wandb @@ -189,15 +179,7 @@ def _init(self, config: Optional[dict] = None): def _log(self, data: Dict[str, Any]): import trackio - sanitized = {} - for k, v in data.items(): - if isinstance(v, torch.Tensor): - try: - v = v.item() - except (ValueError, RuntimeError): - continue - sanitized[k] = v - trackio.log(sanitized) + trackio.log(self._scalarize(data)) def _finish(self, summary: Optional[dict] = None): import trackio diff --git a/tropt/utils/refusal_dir.py b/tropt/utils/refusal_dir.py index 9be6919..4140cee 100644 --- a/tropt/utils/refusal_dir.py +++ b/tropt/utils/refusal_dir.py @@ -8,29 +8,26 @@ Useful for attacks the suppress model refusals via activation steering (e.g., IRIS attack). """ -import io import logging +import random from typing import List, Optional, Tuple -import pandas as pd -import requests import torch -from datasets import load_dataset from jaxtyping import Float -from sklearn.model_selection import train_test_split from tropt.model.huggingface.lm import LMHFModel logger = logging.getLogger(__name__) -def get_hf_model( - model: LMHFModel, -): - """ - Extract the underlying HuggingFace model from the LMHFModel wrapper. - This is off-pattern, but needed for the hooks in this module. - """ - return model._model +ADVBENCH_CSV_URL = "https://raw.githubusercontent.com/llm-attacks/llm-attacks/main/data/advbench/harmful_behaviors.csv" + + +def _shuffled_split(instructions: List[str], test_split: float) -> Tuple[List[str], List[str]]: + """Deterministic train/test split (fixed seed, so runs are reproducible).""" + instructions = list(instructions) + random.Random(42).shuffle(instructions) + n_test = int(len(instructions) * test_split) + return instructions[n_test:], instructions[:n_test] def get_harmful_instructions(n_samples: Optional[int] = None, test_split: float = 0.2) -> Tuple[List[str], List[str]]: @@ -44,16 +41,15 @@ def get_harmful_instructions(n_samples: Optional[int] = None, test_split: float Returns: (train_instructions, test_instructions) """ - url = 'https://raw.githubusercontent.com/llm-attacks/llm-attacks/main/data/advbench/harmful_behaviors.csv' - response = requests.get(url) - dataset = pd.read_csv(io.StringIO(response.content.decode('utf-8'))) - instructions = dataset['goal'].tolist() + from datasets import load_dataset + + dataset = load_dataset("csv", data_files=ADVBENCH_CSV_URL, split="train") + instructions = dataset["goal"] if n_samples: instructions = instructions[:n_samples] - train, test = train_test_split(instructions, test_size=test_split, random_state=42) - return train, test + return _shuffled_split(instructions, test_split) def get_harmless_instructions(n_samples: Optional[int] = None, test_split: float = 0.2) -> Tuple[List[str], List[str]]: @@ -67,6 +63,8 @@ def get_harmless_instructions(n_samples: Optional[int] = None, test_split: float Returns: (train_instructions, test_instructions) """ + from datasets import load_dataset + dataset = load_dataset('tatsu-lab/alpaca') # Filter for instructions without inputs @@ -77,8 +75,7 @@ def get_harmless_instructions(n_samples: Optional[int] = None, test_split: float if n_samples and len(instructions) >= n_samples: break - train, test = train_test_split(instructions, test_size=test_split, random_state=42) - return train, test + return _shuffled_split(instructions, test_split) def extract_activations( @@ -109,11 +106,9 @@ def extract_activations( return_dict=True, )["input_ids"].to(model.device) - hf_model = get_hf_model(model) - - # Forward pass with hidden states + # Forward pass with hidden states (off-pattern raw-model access, needed here) with torch.no_grad(): - outputs = hf_model( + outputs = model._model( inputs, output_hidden_states=True, return_dict=True, diff --git a/uv.lock b/uv.lock index 04f3622..38a6146 100644 --- a/uv.lock +++ b/uv.lock @@ -6768,14 +6768,8 @@ dependencies = [ { name = "jaxtyping", version = "0.3.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "pandas", version = "3.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pydantic" }, - { name = "requests" }, - { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "scikit-learn", version = "1.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "sentence-transformers" }, - { name = "tenacity" }, { name = "torch" }, { name = "tqdm" }, { name = "transformers" }, @@ -6792,6 +6786,7 @@ all = [ { name = "litellm" }, { name = "livelossplot" }, { name = "openai" }, + { name = "tenacity" }, { name = "tiktoken" }, { name = "trackio" }, { name = "voyageai" }, @@ -6831,6 +6826,7 @@ dev = [ { name = "sphinx-rtd-theme" }, { name = "sphinx-sitemap" }, { name = "sphinxext-opengraph" }, + { name = "tenacity" }, { name = "tiktoken" }, { name = "trackio" }, { name = "ty" }, @@ -6839,6 +6835,7 @@ dev = [ ] google = [ { name = "google-genai" }, + { name = "tenacity" }, ] litellm = [ { name = "litellm" }, @@ -6851,6 +6848,7 @@ notebooks = [ ] openai = [ { name = "openai" }, + { name = "tenacity" }, { name = "tiktoken" }, ] tracking = [ @@ -6862,6 +6860,7 @@ vision = [ { name = "diffusers" }, ] voyage = [ + { name = "tenacity" }, { name = "voyageai" }, ] @@ -6881,15 +6880,12 @@ requires-dist = [ { name = "myst-parser", marker = "extra == 'dev'", specifier = ">=0.18.0" }, { name = "numpy", specifier = ">=2.0.0" }, { name = "openai", marker = "extra == 'openai'", specifier = ">=1.0.0" }, - { name = "pandas", specifier = ">=2.3.3" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.0.0" }, { name = "pydantic", specifier = ">=2.0.0" }, { name = "pydata-sphinx-theme", marker = "extra == 'dev'", specifier = ">=0.13.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.0.0" }, - { name = "requests" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.1.0" }, - { name = "scikit-learn" }, { name = "seaborn", marker = "extra == 'dev'", specifier = ">=0.13.0" }, { name = "sentence-transformers", specifier = ">=5.5.0" }, { name = "sphinx", marker = "extra == 'dev'", specifier = ">=5.0.0" }, @@ -6901,7 +6897,9 @@ requires-dist = [ { name = "sphinx-rtd-theme", marker = "extra == 'dev'", specifier = ">=1.0.0" }, { name = "sphinx-sitemap", marker = "extra == 'dev'", specifier = ">=2.5.0" }, { name = "sphinxext-opengraph", marker = "extra == 'dev'", specifier = ">=0.9.0" }, - { name = "tenacity" }, + { name = "tenacity", marker = "extra == 'google'" }, + { name = "tenacity", marker = "extra == 'openai'" }, + { name = "tenacity", marker = "extra == 'voyage'" }, { name = "tiktoken", marker = "extra == 'openai'", specifier = ">=0.5.0" }, { name = "torch", specifier = ">=2.4.0", index = "https://download.pytorch.org/whl/cu126" }, { name = "tqdm" },