Minor fixes

Author: matanbt

tropt/model/huggingface/base.py

@@ -198,6 +198,9 @@ def get_triggered_inputs(
             (they must share `trigger_seq_len`).
         - for *specific* use cases, the following method is suboptimal; however,
             currently generality and support for different input types/shapes are prioritized.
+        - Allows gradient flow through `trigger_embeds`, which can be useful for combining backporable trigger 
+            candidates (e.g., for `compute_grad_from_*()` methods).
+
 
         Args:
             trigger_ids: Tensor, shape = (n_candidates, trigger_seq_len)
@@ -960,7 +963,7 @@ def _compute_grad__batched(
 
         return all_grads
 
-
+    @torch.no_grad()
     def compute_loss_from_tokens(
         self,
         candidate_trigger_ids: Int[Tensor, "n_candidates trigger_seq_len"],
@@ -1014,15 +1017,16 @@ def _compute_candidates_loss__batched(
 
                 logger.debug(f"from loss [msg={template_idx}]: {(cand_idx_end - cand_idx)}")
 
+                model_input = input_manager.get_triggered_inputs(
+                    chosen_template_idx=template_idx,
+                    trigger_ids=batch_candidate_trigger_ids,
+
+                    # loss-conditional flags:
+                    do_append_embeds=loss_func.require_target_prefill,
+                )
+
                 # Only enable gradient is it's required by the loss (e.g. for gradient matching losses); mostly false.
                 with torch.set_grad_enabled(loss_func.require_gradients):
-                    model_input = input_manager.get_triggered_inputs(
-                        chosen_template_idx=template_idx,
-                        trigger_ids=batch_candidate_trigger_ids,
-
-                        # loss-conditional flags:
-                        do_append_embeds=loss_func.require_target_prefill,
-                    )
                     model_output = self.invoke_from_tokens(
                         **model_input.to_dict(),
 

tropt/model/model_mixins.py

@@ -176,6 +176,7 @@ def reset_inputs_from_texts(self) -> None:
 class LossTextAccessMixin(TextAccessMixin):
     """Mixin for models that compute losses based on text-level inputs (black-box access)."""
 
+    @torch.no_grad()
     def compute_loss_from_texts(
         self,
         candidate_trigger_strs: List[str],
@@ -213,35 +214,26 @@ def _compute_loss_batched(batch_size, template_idx=template_idx):
                     end = min(start + batch_size, n_candidates)
                     chunk_strs = candidate_trigger_strs[start:end]
 
-                    curr_model_input = input_manager.get_triggered_inputs(
+                    model_input: ModelInput = input_manager.get_triggered_inputs(
                         chosen_template_idx=template_idx, trigger_strs=chunk_strs,
                     )
-                    curr_texts, curr_targets = (
-                        curr_model_input.input_texts,
-                        curr_model_input.message_targets,
-                    )
-
-                    # Forward pass for this candidate chunk
-                    model_output = self.invoke_from_texts(
-                        input_texts=curr_texts,
-                        message_targets=curr_targets,
-                        require_target_prefill=loss_func.require_target_prefill,
-                        require_generation=loss_func.require_generation,
-                        require_first_token_logprobs=loss_func.require_first_token_logprobs,
-                    )  # Returns ModelOutput with available data
-
-                    # Create ModelInput wrapper
-                    model_input = ModelInput(
-                        input_texts=curr_texts,
-                        input_trigger_strs=chunk_strs,
-                        message_targets=curr_targets,
-                    )
 
-                    # Use unified loss resolution
-                    chunk_loss = resolve_and_compute_loss(
-                        model_output, model_input, loss_func
-                    )  # shape: (chunk_size,)
-                    chunk_losses.append(chunk_loss)
+                    # Only enable gradient is it's required by the loss (e.g. for gradient matching losses); mostly false.
+                    with torch.set_grad_enabled(loss_func.require_gradients):
+                        # Forward pass for this candidate chunk
+                        model_output: ModelOutput = self.invoke_from_texts(
+                            input_texts=model_input.input_texts,
+                            message_targets=model_input.message_targets,
+                            require_target_prefill=loss_func.require_target_prefill,
+                            require_generation=loss_func.require_generation,
+                            require_first_token_logprobs=loss_func.require_first_token_logprobs,
+                        )
+
+                        # Use unified loss resolution
+                        chunk_loss = resolve_and_compute_loss(
+                            model_output, model_input, loss_func
+                        )  # shape: (chunk_size,)
+                        chunk_losses.append(chunk_loss)
 
                 return torch.cat(chunk_losses, dim=0)  # shape: (n_candidates,)