# Adding a New Model > This article is a Work In Progress draft. This guide walks you through wrapping a new model backend. Pick the section that matches your situation: - **[Text-access (black-box) model](#adding-a-text-access-black-box-model)** — API-only models where you can query with text and get text/embeddings back. No internal access. This is the most common case we expect for model backend addition, and it is also the most straightforward. Examples: `EncoderGeminiModel`, `LiteLLMModel`. - **[Token-access (grey/white-box) model](#adding-a-token-access-greywhite-box-model)** — Backends that expose embedding-level input (you can feed raw embeddings and get logits/gradients). You implement the full compute loop. - **[HuggingFace model](#adding-a-huggingface-model)** — Any HF-backed model. `HuggingFaceBackendModel` provides the compute loop; you fill in model-specific parts. Examples: `LMHFModel`, `EncoderHFModel`. > This guide is self-contained — you can follow it step by step without reading anything else. If you want to understand the *why* behind the design, see `DESIGN.md` at the repo root (especially "Component 1: Target Model"). For full API reference, see the {doc}`models API docs <../api/models>`. --- ## Background ### The Three Method Families Every model is composed of methods from three families. Understanding these helps you know exactly what you need to implement: 1. **Invoke methods** (`invoke_from_texts`, `invoke_from_tokens`) — stateless forward passes. These are the core model-specific logic: take input, return `ModelOutput`. Every model implements `invoke_from_texts`; only models with embedding-level access implement `invoke_from_tokens`. 2. **Input methods** (`set_inputs_from_tokens`, `set_inputs_from_texts`) — store an `InputsManager` for use during optimization. Default implementations exist for both flows, so most models don't need custom logic here. You only customize this when your backend has special input handling (e.g., HuggingFace's embedding-level prefix caching). 3. **Compute methods** (`compute_loss_from_tokens`, `compute_grad_from_tokens`, etc.) — the methods optimizers actually call. These use the stored inputs (from family 2) and the invoke methods (from family 1) internally. For text-access models, `LossTextAccessMixin` provides `compute_loss_from_texts` for free. For HuggingFace models, `HuggingFaceBackendModel` provides all token-based compute methods. For other token-access backends, you implement these yourself. > **Batching is the compute method's responsibility, not the invoke method's.** Invoke methods should run a single forward pass over whatever batch they receive and stay free of chunking/OOM-recovery logic. The compute methods are the ones that see `n_candidates` and are expected to split them into sub-batches (and, for GPU backends, retry with a smaller batch on OOM). Both `LossTextAccessMixin.compute_loss_from_texts` and `HuggingFaceBackendModel`'s token compute methods already do this via `accelerate.utils.memory.find_executable_batch_size`, using `self._forward_pass_batch_size` / `self._backward_pass_batch_size` (defined on `BaseModel`, overridable per subclass) as the starting batch size. When you implement a compute method yourself for a non-HF token-access backend, follow the same pattern. **In practice**: for most new models, you implement `invoke_from_texts` and get everything else for free. Token-access models additionally implement `invoke_from_tokens` and `set_inputs_from_tokens`. The compute methods are only hand-written for non-HF token-access backends. ### Base Classes Every model inherits from one of two base classes in {py:mod}`tropt.model`: | Base class | Use for | Inference method you implement | |---|---|---| | `LMBaseModel` | Language models (text generation) | `invoke_from_texts(input_texts, ...)` | | `EncoderBaseModel` | Embedding / encoder models | `invoke_from_texts(input_texts, ...)` and `d_model` property | Each base class defines `__call__` which delegates to its inference method (`invoke_from_texts`). You implement the inference method; `__call__` is already wired up. Throughout this guide, we use **"inference method"** to refer to `invoke_from_texts` — the method your base class requires. ### Mixins (Access Levels) Mixins declare what *type of access* the model exposes. Optimizers check these at init via `model_requirements` to verify compatibility — this is how TROPT enforces that an optimizer only calls methods the model actually supports. **Naming convention**: `{Value}{InputType}AccessMixin` - *Value*: what can be computed (`Loss`, `Gradient`, `Logits`) - *InputType*: what goes in (`Token` for white/grey-box, `Text` for black-box) **Token-access mixins** — require embedding-level access: | Mixin | Access level | Methods you implement | |---|---|---| | `LossTokenAccessMixin` | grey-box | `compute_loss_from_tokens` | | `GradientTokenAccessMixin` | white-box | `compute_grad_from_tokens` | | `LogitsTokenAccessMixin` | white-box | `compute_logits_from_tokens` | | `GradientEmbedAccessMixin` | white-box | `compute_grad_from_embeds` | All token-access mixins inherit from `TokenAccessMixin`, which requires: - **`tokenizer`** property — HuggingFace `PreTrainedTokenizer` or {py:class}`~tropt.model.BaseTokenizer` subclass. - **`set_inputs_from_tokens(templates, targets)`** — builds and stores an `InputsManager` for the optimization run. You implement `tokenizer` and `set_inputs_from_tokens` **once**, regardless of how many token mixins you include. **Text-access mixin** — black-box text I/O: | Mixin | Access level | Methods you implement | |---|---|---| | `LossTextAccessMixin` | black-box | *nothing* — fully implemented; it calls your `__call__` internally | `LossTextAccessMixin` provides `compute_loss_from_texts` and `set_inputs_from_texts` out of the box. It works by calling `self(input_texts)` and passing the resulting `ModelOutput` through the {py:func}`unified loss resolution system `. It also wraps the per-template forward in `find_executable_batch_size`, so candidate batches are chunked and retried automatically on CUDA OOM — your `invoke_from_texts` only needs to handle a single (already-sized) batch. A model can include both token-access and text-access mixins — see `LMHFModel` and `EncoderHFModel` for examples. ### ModelOutput {py:class}`~tropt.common.ModelOutput` is a dataclass that standardizes what your model returns. All fields are optional — populate only the ones your backend can provide: ```python class ModelOutput: output_embeddings: ... # Encoder models full_logits: ... # LMs (full sequence logits) prefill_response_logits: ... # LMs (response-region logits, prefilled) full_hidden_states: ... full_attentions: ... generated_response_strs: ... # LMs (generated text) generated_response_ids: ... generated_response_logits: ... full_strs: ... 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. ### Model compatibility A model's compatibility with a given `(optimizer, loss)` pair is resolved at two independent points — keep both in mind when picking mixins and populating `ModelOutput`: 1. **Optimizer side — mixin contract.** Each optimizer declares the method flow it needs via `model_requirements` (e.g. `LossTokenAccessMixin`, `GradientTokenAccessMixin` for token-flow white-box; `LossTextAccessMixin` for text-flow black-box). `BaseOptimizer.__init__` rejects any model that doesn't subclass the required mixins. Include a mixin only if your model genuinely implements its compute method for the corresponding flow — otherwise the optimizer will call a method that can't run. 2. **Loss side — `ModelOutput` fields.** Each loss reads a specific set of fields from `ModelOutput` / `ModelInput` (plus target fields from `MessageTargets`). At runtime, {py:func}`~tropt.loss.resolve_and_compute_loss` inspects the loss's call signature and fails if any required field is missing. The fields your `invoke_from_tokens` / `invoke_from_texts` populate therefore determine which losses are admissible for each flow — e.g. a model that only returns `generated_response_strs` is limited to text-based losses; one that returns `full_logits` unlocks prefill/CE-style losses. A quick way to sanity-check a new model against the existing optimizer/loss set is the auto-generated [Optimizer-Model-Loss Compatibility Matrix](compatibility_matrix.md), produced by [`docs/scripts/generate_compat_matrix.py`](https://github.com/matanbt/TROPT/blob/main/docs/scripts/generate_compat_matrix.py) via static analysis of `model_requirements`, loss signatures, and the `ModelOutput`/`ModelInput` fields each model populates in its invoke methods. --- ## Adding a Text-Access (Black-box) Model Use this path when you can only query the model with text and receive text or embeddings back — no access to logits, gradients, or internal representations. Your class will inherit from: - A **base class**: `EncoderBaseModel` (for embedding models) or `LMBaseModel` (for language models) - The **`LossTextAccessMixin`** mixin — which provides `compute_loss_from_texts` and `set_inputs_from_texts` for free **Existing examples**: {py:class}`~tropt.model.EncoderGeminiModel`, {py:class}`~tropt.model.LiteLLMModel`. ### What to implement You implement a single method — the inference method required by your base class. Here's an LM example: ```python def invoke_from_texts(self, input_texts: List[str], **kwargs) -> ModelOutput: responses = self._client.complete(input_texts) # your backend call self._update_usage_stats( tokens=..., # extract from backend response forward_calls=1, forward_samples=len(input_texts), ) return ModelOutput(generated_response_strs=responses) ``` Key rules: - Input is always `List[str]`. Output is always a `ModelOutput` — the fields you populate determine which loss types are compatible (see [ModelOutput](#modeloutput)). - **Call `_update_usage_stats` (or `_update_invoke_stats` for models that support FLOP counting) at the same call site as the backend call** — not from higher-level wrappers. This avoids double-counting. Everything else — `set_inputs_from_texts`, `compute_loss_from_texts` — is provided by `LossTextAccessMixin`. ### Minimal class skeleton ```python from tropt.common import ModelOutput from tropt.model import LMBaseModel, LossTextAccessMixin class MyLMModel(LMBaseModel, LossTextAccessMixin): def __init__(self, model_name, api_key=None): self._client = ... # initialize your API client self._model_name = model_name def invoke_from_texts(self, input_texts: List[str], **kwargs): responses = self._client.complete(input_texts) self._update_usage_stats( tokens=..., forward_calls=1, forward_samples=len(input_texts), ) return ModelOutput(generated_response_strs=responses) ``` For complete working examples, see {py:class}`~tropt.model.LiteLLMModel` (LM) or {py:class}`~tropt.model.EncoderGeminiModel` (encoder). No registration is needed to use a custom model — any subclass of `LMBaseModel` or `EncoderBaseModel` works directly with any optimizer whose `model_requirements` it satisfies. --- ## Adding a Token-Access (Grey/White-box) Model Use this when your backend supports embedding-level input — you can pass raw input embeddings (not just text) and receive model outputs like logits or gradients. This is the path for non-HuggingFace backends with internal access. > If your model is HuggingFace-based, skip to [Adding a HuggingFace Model](#adding-a-huggingface-model) — the `HuggingFaceBackendModel` class already provides the compute loop. ### The InputsManager The `InputsManager` pre-processes text templates once (splitting at the `{{OPTIMIZED_TRIGGER}}` placeholder, tokenizing, embedding) and then efficiently inserts candidate triggers at each optimization step via `get_triggered_inputs(trigger_ids, chosen_template_idx)`, which returns a {py:class}`~tropt.common.ModelInput` dataclass. The default {py:class}`~tropt.model.DefaultTokenInputManager` works with any tokenizer supporting the `BaseTokenizer` interface — it decodes trigger IDs to strings and reconstructs full texts. The HuggingFace backend uses `HuggingFaceTokenInputManager`, which overrides this with embedding-level input construction, attention masks, prefix caching, and position slicing. ### The setup-then-compute pattern Token-access follows a two-step pattern that separates *setup* from *computation*: 1. **`set_inputs_from_tokens(templates, targets)`** — called by the optimizer **once** before `optimize_trigger()` begins. Tokenizes the templates, splits them at the `{{OPTIMIZED_TRIGGER}}` placeholder, and stores an `InputsManager` on the model. 2. **`compute_{value}_from_tokens(candidate_trigger_ids, ...)`** — called **repeatedly** at each optimization step. Uses the stored manager to assemble full inputs for each candidate trigger, runs the model, and returns the result. This split matters because templates are fixed for an entire run. Pre-processing them once avoids redundant tokenization and embedding work at every step. Cleanup is handled by `reset_inputs_from_tokens()` (provided by `TokenAccessMixin`), which the optimizer calls at the end of `optimize_trigger()`. You don't need to implement or call it yourself. ### What to implement **1. Inference method** — Same as for text-access models (see [above](#what-to-implement)). **2. `tokenizer` property** — Required by `TokenAccessMixin`. Must be a HuggingFace `PreTrainedTokenizer` or a {py:class}`~tropt.model.BaseTokenizer` subclass. The optimizer uses it to encode/decode triggers. ```python @property def tokenizer(self): return self._tokenizer ``` If your backend doesn't use a HuggingFace tokenizer, implement the `BaseTokenizer` interface — see `OpenAITokenizer` (in `tropt.model.openai.encoder`) for an example wrapping `tiktoken`. **3. `set_inputs_from_tokens`** — Tokenize the templates and construct your `InputsManager`. Store it via `self._token_input_manager`. In most cases you can use the `DefaultTokenInputManager`, which works with any `BaseTokenizer`: ```python def set_inputs_from_tokens(self, templates: List[str], targets: Targets = None) -> None: tok_ids = self.tokenizer(templates, add_special_tokens=True)["input_ids"] self._token_input_manager = DefaultTokenInputManager( templates_ids=tok_ids, tokenizer=self.tokenizer, targets=targets, optimized_trigger_placeholder=OPTIMIZED_TRIGGER_PLACEHOLDER, ) ``` Only subclass or replace the input manager if your backend needs custom input construction (e.g., embedding-level assembly — see HuggingFace's `HuggingFaceTokenInputManager`). **4. `compute_{value}_from_tokens`** — Implement one method per token-access mixin you include. Each compute method should use your `invoke_from_tokens` internally to run the forward pass. The structure is the same for all: loop over templates, call `get_triggered_inputs` to get a `ModelInput`, invoke the model, and return the result. An important convention: **loss is computed per-template, not across templates**. Each template may have its own target, so we never mix templates in a single loss call. The per-template losses are aggregated (averaged) afterward. Another convention: **the compute method owns candidate batching**, not `invoke_from_tokens`. Your `invoke_from_tokens` should assume the caller has already sized the batch appropriately; the loop below is what splits `candidate_trigger_ids` into sub-batches and (for GPU backends) retries on CUDA OOM — see `HuggingFaceBackendModel` for the reference implementation using `accelerate.utils.memory.find_executable_batch_size`. `compute_loss_from_tokens` (required by `LossTokenAccessMixin`) is the most common: ```python @torch.no_grad() def compute_loss_from_tokens( self, candidate_trigger_ids: Int[Tensor, "n_candidates trigger_seq_len"], loss_func: BaseLoss, keep_message_dim: bool = False, ) -> Float[Tensor, "n_candidates"]: input_manager = self._token_input_manager n_templates = input_manager.n_templates all_losses = [[] for _ in range(n_templates)] for template_idx in range(n_templates): for batch in batched(candidate_trigger_ids): # implement your own batching model_input = input_manager.get_triggered_inputs( trigger_ids=batch, chosen_template_idx=template_idx ) model_output = ... # run your backend on model_input self._update_usage_stats(forward_calls=1, forward_samples=len(batch), tokens=...) loss = resolve_and_compute_loss(model_output, model_input, loss_func) all_losses[template_idx].append(loss) losses = torch.stack([torch.cat(l) for l in all_losses]) # (n_templates, n_candidates) return losses if keep_message_dim else losses.mean(dim=0) ``` The key line is `resolve_and_compute_loss(model_output, model_input, loss_func)` — this is the {py:func}`unified loss resolution ` function that dispatches to the correct loss computation based on loss type. You don't implement loss logic yourself; you just provide the data via `ModelOutput` and `ModelInput`. Note that it is `compute_loss_*`'s duty to wrap (or not wrap) computations in `no_grad()`, for efficiency. We assume all calls to `compute_loss_*` require no grad, as there is a dedicated method for gradient computation. `compute_grad_from_tokens` (required by `GradientTokenAccessMixin`) has the same per-template loop structure, but returns the **gradient of the loss w.r.t. the token input** instead of the loss itself. The returned tensor has shape `(n_candidates, trigger_seq_len, vocab_size)` — one gradient value per token position per vocabulary entry, telling the optimizer which substitutions would most reduce the loss. How you compute this gradient is up to your backend. ### FLOP counting TROPT supports optional FLOP counting via invoke-level tracking (see `tropt.model.flop_counter`), using the [Kaplan et al. (2020)](https://arxiv.org/abs/2001.08361) approximation. Cheap and deterministic. Requires `_model` to be a HuggingFace `PreTrainedModel`. FLOP counting is handled entirely inside `invoke_from_tokens` / `invoke_from_texts` — these are the model-call bottleneck. To support it: 1. **Call `_update_invoke_stats`** after each raw model call inside your invoke methods, passing `n_tokens`, `n_samples`, and optionally `count_backward`. 2. For **gradient methods**: pass `count_backward=True` to `invoke_from_tokens` so it records the backward FLOPs correctly. `set_flop_counting("manual")` will raise an error (AssertionError, or AttributeError if no inner HF model is exposed) if `_model` is not a HuggingFace `PreTrainedModel`. --- ## Adding a HuggingFace Model HuggingFace models get the full token-access compute loop for free via `HuggingFaceBackendModel` (in `tropt.model.huggingface.base`). You only implement the model-specific parts. **Existing examples**: {py:class}`~tropt.model.LMHFModel`, {py:class}`~tropt.model.EncoderHFModel`. ### What `HuggingFaceBackendModel` provides These methods are fully implemented and call `invoke_from_tokens` internally: - **`compute_loss_from_tokens`** — batched forward pass over all candidate triggers with automatic OOM-safe batch size reduction. - **`compute_grad_from_tokens`** — gradient w.r.t. one-hot token representations. Handles both hard (discrete) and soft (probabilistic) triggers. - **`compute_grad_from_embeds`** — gradient w.r.t. continuous trigger embeddings (for soft optimization). - **`embedding_matrix`** — computes the actual embedding matrix used by the model. This matters because some models (e.g., Gemma) apply scaling inside the embedding layer, so a plain `embedding_layer.weight` lookup gives incorrect embeddings. All of these call **`invoke_from_tokens`** internally — the one method you must implement. The mixin also handles the template loop, batching, and loss resolution via `resolve_and_compute_loss`. **FLOP counting** (see [FLOP counting](#flop-counting) above) is handled inside `invoke_from_tokens` via `_update_invoke_stats`. Since `_model` is a `PreTrainedModel`, `set_flop_counting("manual")` works out of the box — you just need your `invoke_from_tokens` to call `_update_invoke_stats` with the correct token count and `count_backward` flag. ### What you implement **1. `__init__`** — Load the HF model and set these fields (expected by `HuggingFaceBackendModel`): ```python self._model # transformers.PreTrainedModel (or SentenceTransformer) self._embedding_layer # nn.Module — model.get_input_embeddings() self._tokenizer # HF tokenizer self._forward_pass_batch_size # int — starting batch size for loss computation (inherited from BaseModel; override if needed) self._backward_pass_batch_size # int — starting batch size for gradient computation (inherited from BaseModel; override if needed) ``` `_forward_pass_batch_size` and `_backward_pass_batch_size` live on `BaseModel` (with sensible defaults), so every model — HF or not — already has them. Override in `__init__` (or at the class level) when your backend needs a different starting point. Both are automatically reduced on CUDA OOM via `find_executable_batch_size`. Set the model to **eval mode** and disable gradients on its parameters — optimization gradients flow through the one-hot input, not the model weights: ```python self._model.eval() for param in self._model.parameters(): param.requires_grad = False ``` Register the trigger placeholder as a special token so it tokenizes as a single ID: ```python self._tokenizer.add_special_tokens( {"additional_special_tokens": [OPTIMIZED_TRIGGER_PLACEHOLDER]} ) ``` **2. `invoke_from_tokens`** — The single entry point for all white/grey-box forward passes. Receives the input embeddings and attention mask (unpacked from `ModelInput` by the caller), runs the model, and returns a `ModelOutput`. The caller passes `require_*` boolean flags that mirror the loss's `require_*` class attributes — use them to gate optional behaviour (generation, hidden states, attentions, prefill). The parameter names are identical to the loss attributes: `require_target_prefill`, `require_generation`, `require_hidden_states`, `require_attentions`. Only populate the corresponding `ModelOutput` fields when the flag is `True`. `invoke_from_texts` follows the same convention for `require_target_prefill`, `require_generation`, and `require_first_token_logprobs`. ```python def invoke_from_tokens(self, input_embeds, input_attention_mask, input_prefix_cache_kwargs=None, input_slices=None, count_backward=False, **kwargs) -> ModelOutput: outputs = self._model( inputs_embeds=input_embeds, attention_mask=input_attention_mask, **(input_prefix_cache_kwargs or {}), ) self._update_invoke_stats( n_tokens=int(input_attention_mask.sum().item()), n_samples=input_embeds.shape[0], count_backward=count_backward, ) return ModelOutput( full_logits=outputs.logits, # populate what your model provides ) ``` The `count_backward` flag is set to `True` by gradient methods (`compute_grad_from_tokens`, `compute_grad_from_embeds`) so that FLOPs include the backward pass cost. **3. `set_inputs_from_tokens`** — Build an `HuggingFaceTokenInputManager` (in `tropt.model.huggingface.base`) (or a model-specific subclass) and store it: ```python def set_inputs_from_tokens(self, templates: List[str], targets: Targets = None) -> None: tok_ids = self.tokenizer(templates, add_special_tokens=True)["input_ids"] self._token_input_manager = HuggingFaceTokenInputManager( templates_ids=tok_ids, device=self._model.device, model=self._model, tokenizer=self.tokenizer, embed_func=self._embedding_layer, targets=targets, ) ``` LM models typically apply the chat template and tokenize target responses before building the manager — see {py:meth}`~tropt.model.LMHFModel.set_inputs_from_tokens` for the full pattern. You may also subclass `HuggingFaceTokenInputManager` if your model has special target handling — see {py:class}`~tropt.model.LMHFTokenInputManager` (which auto-appends target response embeddings for prefill-based losses); encoder models reuse `HuggingFaceTokenInputManager` directly (see `EncoderHFModel.set_inputs_from_tokens`). **4. Inference method** — The public method called by `__call__`. This is a **separate code path** from `invoke_from_tokens` — it handles plain-text evaluation, not optimization. See {py:meth}`~tropt.model.LMHFModel.invoke_from_texts` or {py:meth}`~tropt.model.EncoderHFModel.invoke_from_texts` for full examples. ### Class skeleton ```python from tropt.common import OPTIMIZED_TRIGGER_PLACEHOLDER, ModelInput, ModelOutput, Targets from tropt.loss.base import BaseLoss from tropt.model import ( LMBaseModel, GradientEmbedAccessMixin, GradientTokenAccessMixin, LogitsTokenAccessMixin, LossTextAccessMixin, LossTokenAccessMixin, ) from tropt.model.huggingface.base import HuggingFaceTokenInputManager, HuggingFaceBackendModel class MyHFLMModel( # HF backend first so its `device`/`dtype` win MRO over `BaseModel`'s defaults: HuggingFaceBackendModel, LMBaseModel, # token-level access: LossTokenAccessMixin, GradientTokenAccessMixin, LogitsTokenAccessMixin, GradientEmbedAccessMixin, # text-level access: LossTextAccessMixin, ): def __init__(self, model_name, device=None, dtype=None): self._model = ... # AutoModelForCausalLM.from_pretrained(...) self._embedding_layer = ... # self._model.get_input_embeddings() self._tokenizer = ... # AutoTokenizer.from_pretrained(...) self._forward_pass_batch_size = 512 self._backward_pass_batch_size = 32 self._tokenizer.add_special_tokens( {"additional_special_tokens": [OPTIMIZED_TRIGGER_PLACEHOLDER]} ) @property def tokenizer(self): return self._tokenizer @property def device(self): return self._model.device def set_inputs_from_tokens(self, templates, targets=None): ... def invoke_from_tokens(self, input_embeds, input_attention_mask, input_prefix_cache_kwargs=None, input_slices=None, count_backward=False, **kwargs): ... def invoke_from_texts(self, input_texts, **kwargs): ... ``` You do **not** need to implement `compute_loss_from_tokens`, `compute_grad_from_tokens`, or `compute_grad_from_embeds` — those come from `HuggingFaceBackendModel`.