Decision: batch the prompts of a request with slice_logits_to_keep

src/transformers/decision/modeling_utils.py

@@ -19,17 +19,49 @@
 from pathlib import Path
 
 import torch
+from torch.nn.utils.rnn import pad_sequence
 
 from ..configuration_utils import PreTrainedConfig
 from ..utils import cached_file
 from .configuration_utils import DecisionConfig
 from .template_utils import Question
 
 
-def process_scores(decision_config, input_ids, marks, questions, scores) -> list[torch.Tensor]:
-    """The logits processors of the decision config on the option scores of each question of a prompt."""
-    tail = input_ids[0, int(marks[-1]) :]  # the tokens after the last mark, where the templates write the state
-    return [decision_config.logits_processors[q.type](tail, s.float()) for q, s in zip(questions, scores)]
+def batch_prompts(prompts: list[tuple[dict, list[Question]]]) -> dict[str, torch.Tensor]:
+    """
+    The prompts of a request as one right-padded batch, with `logits_to_keep` the mask of the positions read: the
+    token before each mark, so the read positions of a batch are its questions in order.
+    """
+    inputs = [prompt for prompt, _ in prompts]
+    reads = []
+    for prompt in inputs:
+        read = torch.zeros(prompt["input_ids"].shape[-1], dtype=torch.bool, device=prompt["input_ids"].device)
+        read[prompt["marks"] - 1] = True
+        reads.append(read)
+    batch = {
+        "attention_mask": pad_sequence([torch.ones_like(read, dtype=torch.long) for read in reads], batch_first=True),
+        "logits_to_keep": pad_sequence(reads, batch_first=True),
+    }
+    for key in inputs[0].keys() - {"marks", "attention_mask"}:
+        values = [prompt[key] for prompt in inputs]
+        if values[0].shape[-1] == inputs[0]["input_ids"].shape[-1]:  # per token: `input_ids`, `mm_token_type_ids`...
+            batch[key] = pad_sequence([value[0] for value in values], batch_first=True)
+        else:
+            batch[key] = torch.cat(values)
+    return batch
+
+
+def process_scores(decision_config, batch, questions, scores) -> list[list[torch.Tensor]]:
+    """
+    The logits processors of the decision config on the option scores of each question of each prompt of a batch from
+    [`batch_prompts`], `scores` being those of all its questions in order.
+    """
+    scores, processed = iter(scores), []
+    lengths = batch["attention_mask"].sum(-1).tolist()
+    for input_ids, read, length, group in zip(batch["input_ids"], batch["logits_to_keep"], lengths, questions):
+        tail = input_ids[int(read.nonzero()[-1]) + 1 : length]  # the tokens after the last mark: the state
+        processed.append([decision_config.logits_processors[q.type](tail, next(scores).float()) for q in group])
+    return processed
 
 
 def _hub_kwargs(kwargs: dict) -> dict:
@@ -53,6 +85,8 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, config=None, deci
             head_config, decision_config = cls.checkpoint_config(pretrained_model_name_or_path, **kwargs)
             if head_config is None and config is not None:
                 raise ValueError(f"{pretrained_model_name_or_path} has no {cls.__name__} head next to its weights.")
+            if head_config is not None and config is not None:  # the auto classes set `attn_implementation` on it
+                kwargs.setdefault("attn_implementation", config._attn_implementation_internal)
             config = head_config or config
         model = super().from_pretrained(pretrained_model_name_or_path, *args, config=config, **kwargs)
         if decision_config is None:  # the decision config of the repo, read now like `generation_config`
@@ -100,29 +134,28 @@ def save_pretrained(self, save_directory, *args, **kwargs):
     def decide(
         self,
         input_ids: torch.LongTensor,
-        marks: torch.LongTensor,
-        decision_questions: list[Question],
+        logits_to_keep: torch.BoolTensor,
+        decision_questions: list[list[Question]],
         decision_config: DecisionConfig | None = None,
         **kwargs,
-    ) -> list[torch.Tensor]:
+    ) -> list[list[torch.Tensor]]:
         """
-        Scores the options of each question of a prompt from `decision_config.prepare_inputs`: one forward pass,
-        then the logits processors of the decision config. Returns the scores of each question, which
+        Scores the options of each question of a batch of prompts from [`batch_prompts`]: one forward pass, then the
+        logits processors of the decision config. Returns the scores of each question of each prompt, which
         `Question.answer` turns into answers.
         """
         decision_config = decision_config or self.decision_config
-        inputs = {"input_ids": input_ids, **kwargs}
-        num_options = torch.tensor([[len(question.options) for question in decision_questions]])
-        question_types = [Question.questions_types.index(question.type) for question in decision_questions]
+        batch = {"input_ids": input_ids, "logits_to_keep": logits_to_keep, **kwargs}
+        batch = {key: value.to(self.device) if hasattr(value, "to") else value for key, value in batch.items()}
+        questions = [question for group in decision_questions for question in group]
+        num_options = [len(question.options) for question in questions]
+        question_types = [Question.questions_types.index(question.type) for question in questions]
         scores = self(
-            marks=marks[None].to(self.device),
-            num_options=num_options.to(self.device),
-            question_type_ids=torch.tensor([question_types], device=self.device),
-            **{key: value.to(self.device) if hasattr(value, "to") else value for key, value in inputs.items()},
-        )
-        return process_scores(
-            decision_config, inputs["input_ids"], marks, decision_questions, scores.split(num_options[0].tolist())
+            num_options=torch.tensor(num_options, device=self.device),
+            question_type_ids=torch.tensor(question_types, device=self.device),
+            **batch,
         )
+        return process_scores(decision_config, batch, decision_questions, scores.split(num_options))
 
 
 __all__ = ["DecisionMixin"]

src/transformers/models/jev/modeling_jev.py

@@ -16,7 +16,7 @@
 
 from ...decision import Question
 from ...decision.modeling_utils import DecisionMixin
-from ...modeling_utils import PreTrainedModel
+from ...modeling_utils import PreTrainedModel, slice_logits_to_keep
 from ...utils import auto_docstring
 from ..auto import AutoModel
 from .configuration_jev import JevConfig
@@ -46,15 +46,15 @@ def __init__(self, config: JevConfig):
     def forward(
         self,
         input_ids: torch.LongTensor,
-        marks: torch.LongTensor,
+        logits_to_keep: torch.BoolTensor,
         num_options: torch.LongTensor,
         question_type_ids: torch.LongTensor,
         **kwargs,
     ) -> torch.Tensor:
-        # A batch of one prompt
-        marks, num_options, question_type_ids = marks[0], num_options[0], question_type_ids[0]
-        hidden_states = self.model(input_ids=input_ids, use_cache=False, **kwargs).last_hidden_state[0]
-        logits = self.proj(hidden_states[marks - 1].to(self.proj.weight.dtype))
+        # The hidden states before the marks of the batch: one per question, in order
+        hidden_states = self.model(input_ids=input_ids, use_cache=False, **kwargs).last_hidden_state
+        hidden_states = slice_logits_to_keep(hidden_states, logits_to_keep)
+        logits = self.proj(hidden_states.to(self.proj.weight.dtype))
         if self.config.softcap is not None:
             logits = self.config.softcap * torch.tanh(logits / self.config.softcap)
         # The options of a question are scored by the rows from the offset of its type on, in order

src/transformers/pipelines/__init__.py

@@ -1052,8 +1052,8 @@ def pipeline(
             )
         task = get_task(model, token)
 
-    # Retrieve the task
-    if task in custom_tasks:
+    # Retrieve the task: the repo's own pipeline, unless remote code is refused for a task we implement
+    if task in custom_tasks and not (trust_remote_code is False and task in SUPPORTED_TASKS):
         targeted_task, task_options = clean_custom_task(custom_tasks[task])
         if pipeline_class is None:
             if not trust_remote_code:

src/transformers/pipelines/decision.py

@@ -8,7 +8,7 @@
 
 
 if is_torch_available():
-    from ..decision.modeling_utils import DecisionMixin, process_scores
+    from ..decision.modeling_utils import DecisionMixin, batch_prompts, process_scores
 
 
 @add_end_docstrings(build_pipeline_init_args(has_tokenizer=True))
@@ -81,21 +81,15 @@ def preprocess(self, request: dict[str, Any] | str):
         return {"prompts": self.decision_config.prepare_inputs(request, self.processor or self.tokenizer)}
 
     def _forward(self, model_inputs):
-        # TODO: batch / continious batching
         model, decision_config, prompts = self.model, self.decision_config, model_inputs["prompts"]
+        batch, questions = batch_prompts(prompts), [questions for _, questions in prompts]
         if isinstance(model, DecisionMixin):
-            scores = [
-                model.decide(**inputs, decision_questions=q, decision_config=decision_config) for inputs, q in prompts
-            ]
+            scores = model.decide(**batch, decision_questions=questions, decision_config=decision_config)
             return {"prompts": prompts, "scores": scores}
-        scores = []
-        for inputs, questions in prompts:  # the logits of the labels before each mark
-            inputs = dict(inputs)
-            marks = inputs.pop("marks")
-            logits = model(**inputs, use_cache=False, logits_to_keep=marks - 1).logits[0]
-            labels = [logits[i, self.label_token_ids(tuple(q.labels))] for i, q in enumerate(questions)]
-            scores.append(process_scores(decision_config, inputs["input_ids"], marks, questions, labels))
-        return {"prompts": prompts, "scores": scores}
+        logits = model(**batch, use_cache=False).logits  # the logits before each mark: one row per question
+        flat = [question for group in questions for question in group]
+        labels = [logits[i, self.label_token_ids(tuple(q.labels))] for i, q in enumerate(flat)]
+        return {"prompts": prompts, "scores": process_scores(decision_config, batch, questions, labels)}
 
     def postprocess(self, model_outputs):
         answers = {}

src/transformers/utils/auto_docstring.py

@@ -2232,23 +2232,16 @@ class ModelArgs:
         "shape": None,
     }
 
-    marks = {
-        "description": """
-    Positions recorded by the `<|mark|>` of a decision template: the number of tokens of the prompt before each of them.
-    """,
-        "shape": "of shape `(num_marks,)`",
-    }
-
     num_options = {
         "description": """
-    Number of options of each question of the prompt.
+    Number of options of each question of the batch, in order.
     """,
         "shape": "of shape `(num_questions,)`",
     }
 
     question_type_ids = {
         "description": """
-    Type of each question of the prompt, as its index in `Question.questions_types`.
+    Type of each question of the batch, as its index in `Question.questions_types`.
     """,
         "shape": "of shape `(num_questions,)`",
     }

tests/models/jev/test_modeling_jev.py

@@ -30,7 +30,7 @@
 if is_torch_available():
     import torch
 
-    from transformers import JevForDecision
+    from transformers import AutoModelForDecision, JevForDecision
 
 
 @require_torch
@@ -70,6 +70,7 @@ def test_from_a_vllm_decision_checkpoint(self):
             Path(tmp, "decision_config.json").write_text(json.dumps(settings), encoding="utf-8")
             save_file({"weight": readout}, f"{tmp}/readout.safetensors", metadata={"format": "pt"})
             model = JevForDecision.from_pretrained(tmp)
+            eager = AutoModelForDecision.from_pretrained(tmp, attn_implementation="eager")
             decide = pipeline("decision", model=tmp)
         with tempfile.TemporaryDirectory() as tmp:  # saved as a decision model, with its decision config
             model.save_pretrained(tmp)
@@ -78,6 +79,7 @@ def test_from_a_vllm_decision_checkpoint(self):
         torch.testing.assert_close(reloaded.proj.weight, readout)
         self.assertEqual(model.config.num_verbalizers, 3)
         self.assertFalse(model.config.text_config.is_causal)
+        self.assertEqual(eager.model.config._attn_implementation, "eager")  # not lost to the config of its head
         torch.testing.assert_close(model.proj.weight, readout)
         self.assertIsInstance(decide.model, JevForDecision)
         self.assertEqual(decide.decision_config.temperature, dict.fromkeys(["noul", "choice", "score"], 2.0))
@@ -127,9 +129,9 @@ def test_config_json_cannot_add_extra_files(self):
 class JevIntegrationTest(unittest.TestCase):
     def test_jev_9b(self):
         expected = {
-            "intent": {"refund": 0.9999, "track": 0.0000, "other": 0.0001},
-            "urgent": 0.5419,
-            "anger": {"0": 0.0165, "1": 0.1113, "2": 0.5163, "3": 0.3126, "4": 0.0419, "5": 0.0013},
+            "intent": {"refund": 0.9998, "track": 0.0000, "other": 0.0002},
+            "urgent": 0.5292,
+            "anger": {"0": 0.0216, "1": 0.1246, "2": 0.5156, "3": 0.2958, "4": 0.0410, "5": 0.0014},
         }
         check_decision(
             self, "hf-internal-testing/JEV-9B", "hf-internal-testing/JEV-9B", expected, atol=1e-2, dtype="bfloat16"
@@ -138,9 +140,9 @@ def test_jev_9b(self):
     def test_jev_35b(self):
         # System 1 merged into the weights: its head is rows of the `lm_head`
         expected = {
-            "intent": {"refund": 0.9981, "track": 0.0001, "other": 0.0018},
-            "urgent": 0.7539,
-            "anger": {"0": 0.0155, "1": 0.0859, "2": 0.4147, "3": 0.3896, "4": 0.0872, "5": 0.007},
+            "intent": {"refund": 0.9969, "track": 0.0001, "other": 0.0030},
+            "urgent": 0.7632,
+            "anger": {"0": 0.0145, "1": 0.0825, "2": 0.4165, "3": 0.3914, "4": 0.0878, "5": 0.0073},
         }
         check_decision(self, "hf-internal-testing/JEV-35B", "hf-internal-testing/JEV-35B", expected, atol=1e-2)
 

tests/test_decision_head_common.py

@@ -20,13 +20,13 @@
 from huggingface_hub import hf_hub_download
 
 from transformers import AutoConfig, AutoTokenizer, DecisionConfig, is_torch_available, pipeline
-from transformers.testing_utils import torch_device
 
 
 if is_torch_available():
     import torch
 
     from transformers import AutoModelForDecision, Qwen2Config, Qwen2ForCausalLM
+    from transformers.decision.modeling_utils import batch_prompts
 
 
 def render(processor, request, config=None):
@@ -89,14 +89,13 @@ def test_chat_template_reproduces_the_model_code(self):
         sequences = [{"input_ids": prompt.input_ids, "marks": prompt.marks} for prompt in self.prompts]
         self.assertEqual(sequences, self.test["sequences"])
 
-    def scores(self, model):
+    def scores(self, model, prompts=None):
+        """The scores of each question of each prompt, from one batch of `prompts` (all of them by default)."""
+        prompts = prompts or self.prompts
+        batch = batch_prompts([(prompt.inputs, prompt.questions) for prompt in prompts])
+        questions = [prompt.questions for prompt in prompts]
         with torch.no_grad():
-            return [
-                model.decide(
-                    **prompt.inputs, decision_questions=prompt.questions, decision_config=self.decision_config
-                )
-                for prompt in self.prompts
-            ]
+            return model.decide(**batch, decision_questions=questions, decision_config=self.decision_config)
 
     def test_scores_and_save_load(self):
         for model_class in self.all_model_classes:
@@ -108,6 +107,9 @@ def test_scores_and_save_load(self):
             for prompt, scores, loaded_scores in zip(self.prompts, self.scores(model), self.scores(loaded)):
                 torch.testing.assert_close(loaded_scores, scores)
                 self.assertEqual([len(s) for s in scores], [len(question.options) for question in prompt.questions])
+            # The padded batch scores each prompt like it alone
+            one_by_one = [scores for prompt in self.prompts for scores in self.scores(model, [prompt])]
+            torch.testing.assert_close(self.scores(model), one_by_one)
 
 
 def check_decision(
@@ -116,8 +118,15 @@ def check_decision(
     """The answers to the request of `decision_test.json` (probabilities in order, or `noul`) match `expected`."""
     test = json.loads(Path(hf_hub_download(decision_repo, "decision_test.json")).read_bytes())
     # By default the original checkpoint, which the decision files are converted from
+    # Some checkpoints ship their own remote-code decision pipeline: this one is ours
     decide = pipeline(
-        "decision", model=model or test["model"], decision_config=decision_repo, device=torch_device, **pipeline_kwargs
+        "decision",
+        model=model or test["model"],
+        decision_config=decision_repo,
+        # 27-35B checkpoints: sharded over the 24 GB CI GPUs
+        device_map="auto",
+        trust_remote_code=False,
+        **pipeline_kwargs,
     )
     answers = decide(test["request"])
     for name, value in expected.items():