max_qas_per_sample and fix flash attn w/ padded ctx!!

intx_sft.py

@@ -79,6 +79,7 @@ def get_ds_prob(train_ds_len: list[int], total_len: int):
     for i, ds_len in enumerate(train_ds_len):
         if (ds_len / total_len) > 0.01:
             probs[i] = ds_len / res_total_len * res_probs
+    logger.debug(f"Dataset probabilities: {probs}")
     assert isclose(sum(probs), 1.0), (
         f"Probs sum to {sum(probs)} ({probs}), expected 1.0"
     )
@@ -270,6 +271,7 @@ def main():
     _get_tokenized_dataset = partial(
         get_tokenized_dataset,
         max_qas_len=ctx_args.max_qas_len,
+        max_qas_per_sample=ctx_args.max_qas_per_sample,
         base_model_max_len=model.base_model.config.max_position_embeddings,
         tokenizer=tokenizer,
         tokenizer_kwargs=tokenizer_kwargs,

src/ctx_to_lora/configs.py

@@ -322,6 +322,13 @@ class CtxTrainingArguments:
             "QA pairs that are longer than this value will be split up into multiple samples."
         },
     )
+    max_qas_per_sample: int = field(
+        default=-1,
+        metadata={
+            "help": "Max QA pair per context. If a context has more QA pairs than this value, "
+            "they will be split up into multiple samples."
+        },
+    )
     max_packed_inp_len: int | None = field(
         default=2**14,
         metadata={"help": "Maximum packed input length for training."},

src/ctx_to_lora/data/definitions.py

@@ -84,6 +84,20 @@ LONGBENCH_E_TASKS = [
 ]
 
 DS_KWARGS = {
+    "context_numbers_2_10": dict(
+        train=dict(
+            path="json",
+            data_files=glob("data/raw_datasets/context_numbers_[2-9]/train.jsonl")
+            + ["data/raw_datasets/context_numbers_10/train.jsonl"],
+            split="train",
+        ),
+        validation=dict(
+            path="json",
+            data_files=glob("data/raw_datasets/context_numbers_[2-9]/val.jsonl")
+            + ["data/raw_datasets/context_numbers_10/val.jsonl"],
+            split="train",
+        ),
+    ),
     "hotpot_qa": dict(
         train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:]"),
         validation=dict(
@@ -631,6 +645,6 @@ EVAL_INTX_TEMPLATES = {
 for ds_name in LONGBENCH_E_TASKS:
     EVAL_INTX_TEMPLATES[ds_name] = EVAL_INTX_TEMPLATES[ds_name[:-2]]
 
-for ds_name in DS_KWARGS:
-    if ds_name not in EVAL_INTX_TEMPLATES:
-        EVAL_INTX_TEMPLATES[ds_name] = "{input}"
+# for ds_name in DS_KWARGS:
+#     if ds_name not in EVAL_INTX_TEMPLATES:
+#         EVAL_INTX_TEMPLATES[ds_name] = "{input}"

src/ctx_to_lora/data/processing.py

@@ -42,7 +42,7 @@ def load_answers(ds_name, split):
     elif ds_name == "drop":
 
         def extract_ans(sample):
-            return {"answers": sample["answer_spans"]["spans"]}
+            return {"answers": sample["answers_spans"]["spans"]}
 
     ds_kwargs = get_ds_kwargs(ds_name, split)
     ds = load_dataset(**ds_kwargs, trust_remote_code=True)
@@ -174,7 +174,7 @@ def get_preprocessing_fn(
             return {
                 "context": sample["passage"],
                 "prompt": prompt,
-                "response": ", ".join(set(sample["answers_spans"]["spans"])),
+                "response": sample["answers_spans"]["spans"][0],
             }
 
     elif ds_name == "ropes":
@@ -284,10 +284,10 @@ def get_preprocessing_fn(
         def eval_intx_decorator(f):
             def g(sample):
                 sample = f(sample)
-                prompt_field = "prompt" if "prompt" in sample else "prompts"
-                sample[prompt_field] = prompt_template.format(
-                    input=sample[prompt_field]
+                assert "prompt" in sample, (
+                    f"Expected 'prompt' in sample, got {sample.keys()}"
                 )
+                sample["prompt"] = prompt_template.format(input=sample["prompt"])
                 return sample
 
             return g
@@ -599,6 +599,7 @@ def get_tokenized_dataset(
     ds_name: str,
     split: str,
     max_qas_len: int,
+    max_qas_per_sample: int,
     base_model_max_len: int,
     tokenizer: PreTrainedTokenizerBase,
     tokenizer_kwargs: dict[str, Any],
@@ -633,6 +634,7 @@ def get_tokenized_dataset(
     )
     tokenize_kwargs = dict(
         max_qas_len=max_qas_len,
+        max_qas_per_sample=max_qas_per_sample,
         base_model_max_len=base_model_max_len,
         tokenizer_kwargs=tokenizer_kwargs,
         ctx_model_max_len=ctx_model_max_len,
@@ -661,7 +663,7 @@ def get_tokenized_dataset(
         tokenized_ds = datasets.load_from_disk(ds_path)
         return tokenized_ds
 
-    num_proc = None if streaming and ("train" in split) else 8
+    num_proc = None if streaming and ("train" in split) else 4
     ds = load_and_process_dataset(
         **load_and_process_kwargs,
         num_proc=num_proc,
@@ -676,8 +678,8 @@ def get_tokenized_dataset(
     )
 
     if ("train" in split) and is_caching_enabled():
-        tokenized_ds.save_to_disk(ds_path, num_proc=num_proc)
         tokenized_ds = tokenized_ds.shuffle()
+        tokenized_ds.save_to_disk(ds_path, num_proc=num_proc)
         # force reload from disk for fingerprint consistency
         tokenized_ds = datasets.load_from_disk(ds_path)
     return tokenized_ds
@@ -685,6 +687,7 @@ def get_tokenized_dataset(
 
 def construct_and_tokenize_ctx_qa(
     max_qas_len,
+    max_qas_per_sample,
     base_model_max_len,
     tokenizer,
     tokenizer_kwargs,
@@ -826,7 +829,10 @@ def construct_and_tokenize_ctx_qa(
             logging.debug(f"Split too long QAs with max length {max_qas_len}")
             tokenized_ds = tokenized_ds.map(
                 split_too_long_qas,
-                fn_kwargs={"max_qas_len": max_qas_len},
+                fn_kwargs={
+                    "max_qas_len": max_qas_len,
+                    "max_qas_per_sample": max_qas_per_sample,
+                },
                 batched=True,
                 batch_size=100_000,
                 num_proc=4,
@@ -1010,13 +1016,15 @@ def convert_ctx_prompt_response_to_messages(
     return dict(data=data)
 
 
-def split_too_long_qas(samples: dict[str, any], max_qas_len: int):
+def split_too_long_qas(
+    samples: dict[str, any], max_qas_len: int, max_qas_per_sample: int
+):
     # samples keys: "input_ids", "attention_mask", "labels", "ctx_ids", "ctx_attn_mask"
     # split the qas into multiple samples if they are too long
     # e.g., if max_qas_len = 512, and qas is 1024 tokens long,
     # we split it such that each sample has at most 512 tokens
     # and the ctx_ids and ctx_attn_mask are the same for all samples
-    if max_qas_len < 0:
+    if max_qas_len < 0 and max_qas_per_sample < 0:
         return samples
     input_ids = samples["input_ids"]
     attention_mask = samples["attention_mask"]
@@ -1029,7 +1037,10 @@ def split_too_long_qas(samples: dict[str, any], max_qas_len: int):
     longest_old_qas_len = max(total_lengths) if total_lengths else 0
 
     # Early exit if no splitting needed
-    if all(length <= max_qas_len for length in total_lengths):
+    if (max_qas_len < 0 or all(length <= max_qas_len for length in total_lengths)) and (
+        max_qas_per_sample < 0
+        or all(len(seq) <= max_qas_per_sample for seq in input_ids)
+    ):
         logger.debug(f"Longest old qas len: {longest_old_qas_len}")
         logger.debug(f"Longest new qas len: {longest_old_qas_len}")
         return samples
@@ -1047,8 +1058,11 @@ def split_too_long_qas(samples: dict[str, any], max_qas_len: int):
         out["ctx_attn_mask"].append(ctx_attn)
 
     for i, tot_inp_len in enumerate(total_lengths):
-        if tot_inp_len <= max_qas_len:
+        if (max_qas_len < 0 or tot_inp_len <= max_qas_len) and (
+            max_qas_per_sample < 0 or len(input_ids[i]) <= max_qas_per_sample
+        ):
             # No need to split - add entire sample
+            # logger.debug(f"Sample {i} is within limits, adding as is.")
             for k in samples:
                 out[k].append(samples[k][i])
             continue
@@ -1070,7 +1084,12 @@ def split_too_long_qas(samples: dict[str, any], max_qas_len: int):
                 n_skip += 1
                 continue
 
-            if new_qas_len + inp_len <= max_qas_len:
+            # Check if we can add to current batch (both length and sample count limits)
+            can_add_to_current = (
+                max_qas_len < 0 or new_qas_len + inp_len <= max_qas_len
+            ) and (max_qas_per_sample < 0 or len(new_input_ids) < max_qas_per_sample)
+
+            if can_add_to_current:
                 # Add to current batch
                 new_qas_len += inp_len
                 new_input_ids.append(inp_ids)
@@ -1078,6 +1097,9 @@ def split_too_long_qas(samples: dict[str, any], max_qas_len: int):
                 new_labels.append(label)
             else:
                 # Current batch is full, save it and start new batch
+                # logger.debug(
+                #     f"sample {i}: adding batch with {len(new_input_ids)} sequences"
+                # )
                 if new_input_ids:  # Only add non-empty batches
                     add_batch(
                         new_input_ids,

src/ctx_to_lora/eval_utils.py

@@ -98,7 +98,8 @@ def normalize_answer(s: str) -> str:
 
     def remove_punc(text: str) -> str:
         exclude = set(string.punctuation)
-        return "".join(ch for ch in text if ch not in exclude)
+        # return "".join(ch for ch in text if ch not in exclude)
+        return " ".join(re.split(f"[{string.punctuation}]", text))
 
     def lower(text: str) -> str:
         return text.lower()
@@ -737,6 +738,7 @@ def evaluate(
     _get_tokenized_dataset = partial(
         get_tokenized_dataset,
         max_qas_len=-1,
+        max_qas_per_sample=1,
         base_model_max_len=model.base_model.config.max_position_embeddings,
         tokenizer=tokenizer,
         tokenizer_kwargs=tokenizer_kwargs,

src/ctx_to_lora/modeling/aggregator.py

@@ -104,7 +104,6 @@ class Perceiver(nn.Module):
         self.per_rank_gen = per_rank_gen
         self.r = lora_r if self.per_rank_gen else 1
         n_output_queries = num_layers * (num_modules * self.r + num_extra_modules)
-        self.tot_output_size = n_output_queries
         self.layer_to_layer = layer_to_layer_ctx_encoder
         if self.layer_to_layer:
             n_output_queries = num_modules * self.r + num_extra_modules

src/ctx_to_lora/modeling/hypernet.py

@@ -1,7 +1,6 @@
 import logging
 from collections.abc import Iterable
 from contextlib import contextmanager
-from copy import deepcopy
 from dataclasses import dataclass
 from functools import partial
 from typing import Any
@@ -23,7 +22,6 @@ from torch import Tensor, nn
 from transformers import (
     PretrainedConfig,
     PreTrainedModel,
-    PreTrainedTokenizerBase,
 )
 from transformers.modeling_outputs import ModelOutput
 from transformers.models.modernbert.modeling_modernbert import ModernBertModel
@@ -1035,91 +1033,91 @@ class ModulatedPretrainedModel(nn.Module):
         return model_outputs
 
 
-class ModulatedModelWithSharedInput(nn.Module):
-    def __init__(
-        self,
-        modulated_model: ModulatedPretrainedModel,
-        base_tokenizer: PreTrainedTokenizerBase,
-        ctx_tokenizer: PreTrainedTokenizerBase,
-        ctx_end_predicate: str | None = None,
-        remove_ctx_from_base_input: bool = False,
-    ):
-        super().__init__()
-        self.modulated_model = modulated_model
-        self.base_tokenizer = base_tokenizer
-        self.ctx_tokenizer = ctx_tokenizer
-        self.register_module("modulated_model", self.modulated_model)
-        self.ctx_end_predicate = ctx_end_predicate
-        self.remove_ctx_from_base_input = remove_ctx_from_base_input
+# class ModulatedModelWithSharedInput(nn.Module):
+#     def __init__(
+#         self,
+#         modulated_model: ModulatedPretrainedModel,
+#         base_tokenizer: PreTrainedTokenizerBase,
+#         ctx_tokenizer: PreTrainedTokenizerBase,
+#         ctx_end_predicate: str | None = None,
+#         remove_ctx_from_base_input: bool = False,
+#     ):
+#         super().__init__()
+#         self.modulated_model = modulated_model
+#         self.base_tokenizer = base_tokenizer
+#         self.ctx_tokenizer = ctx_tokenizer
+#         self.register_module("modulated_model", self.modulated_model)
+#         self.ctx_end_predicate = ctx_end_predicate
+#         self.remove_ctx_from_base_input = remove_ctx_from_base_input
 
-    # delegate to self.modulated_model
-    @property
-    def config(self):
-        return self.modulated_model.config
+#     # delegate to self.modulated_model
+#     @property
+#     def config(self):
+#         return self.modulated_model.config
 
-    @property
-    def device(self):
-        return self.modulated_model.device
+#     @property
+#     def device(self):
+#         return self.modulated_model.device
 
-    def tie_weights(self):
-        self.modulated_model.base_model.tie_weights()
+#     def tie_weights(self):
+#         self.modulated_model.base_model.tie_weights()
 
-    def state_dict(self, *args, **kwargs):
-        return self.modulated_model.state_dict(*args, **kwargs)
+#     def state_dict(self, *args, **kwargs):
+#         return self.modulated_model.state_dict(*args, **kwargs)
 
-    def forward(self, *args, **kwargs):
-        input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0]
-        input_txts = self.base_tokenizer.batch_decode(
-            input_ids,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=True,
-        )
-        ctx_txts = deepcopy(input_txts)
-        if self.ctx_end_predicate:
-            ctx_txts = [
-                txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts
-            ]
-        ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to(
-            self.device
-        )
-        ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask
-        if self.remove_ctx_from_base_input:
-            raise NotImplementedError("Not implemented")
-            # input_txts = [
-            #     txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts
-            # ]
-            # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True)
-            # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
-            # kwargs["input_ids"] = input_ids
-            # kwargs["attention_mask"] = attention_mask
-        return self.modulated_model(ctx_ids, ctx_attn_mask, *args, **kwargs)
+#     def forward(self, *args, **kwargs):
+#         input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0]
+#         input_txts = self.base_tokenizer.batch_decode(
+#             input_ids,
+#             skip_special_tokens=True,
+#             clean_up_tokenization_spaces=True,
+#         )
+#         ctx_txts = deepcopy(input_txts)
+#         if self.ctx_end_predicate:
+#             ctx_txts = [
+#                 txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts
+#             ]
+#         ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to(
+#             self.device
+#         )
+#         ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask
+#         if self.remove_ctx_from_base_input:
+#             raise NotImplementedError("Not implemented")
+#             # input_txts = [
+#             #     txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts
+#             # ]
+#             # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True)
+#             # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
+#             # kwargs["input_ids"] = input_ids
+#             # kwargs["attention_mask"] = attention_mask
+#         return self.modulated_model(ctx_ids, ctx_attn_mask, *args, **kwargs)
 
-    def generate(self, *args, **kwargs):
-        input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0]
-        input_txts = self.base_tokenizer.batch_decode(
-            input_ids,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=True,
-        )
-        ctx_txts = deepcopy(input_txts)
-        if self.ctx_end_predicate:
-            ctx_txts = [
-                txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts
-            ]
-        ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to(
-            self.device
-        )
-        ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask
-        if self.remove_ctx_from_base_input:
-            raise NotImplementedError("Not implemented")
-            # input_txts = [
-            #     txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts
-            # ]
-            # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True)
-            # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
-            # kwargs["input_ids"] = input_ids
-            # kwargs["attention_mask"] = attention_mask
-        return self.modulated_model.generate(ctx_ids, ctx_attn_mask, *args, **kwargs)
+#     def generate(self, *args, **kwargs):
+#         input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0]
+#         input_txts = self.base_tokenizer.batch_decode(
+#             input_ids,
+#             skip_special_tokens=True,
+#             clean_up_tokenization_spaces=True,
+#         )
+#         ctx_txts = deepcopy(input_txts)
+#         if self.ctx_end_predicate:
+#             ctx_txts = [
+#                 txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts
+#             ]
+#         ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to(
+#             self.device
+#         )
+#         ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask
+#         if self.remove_ctx_from_base_input:
+#             raise NotImplementedError("Not implemented")
+#             # input_txts = [
+#             #     txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts
+#             # ]
+#             # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True)
+#             # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
+#             # kwargs["input_ids"] = input_ids
+#             # kwargs["attention_mask"] = attention_mask
+#         return self.modulated_model.generate(ctx_ids, ctx_attn_mask, *args, **kwargs)
 
 
 @contextmanager

src/ctx_to_lora/modeling/idefics2.py

@@ -32,6 +32,7 @@ from transformers.utils import (
 )
 
 if is_flash_attn_2_available():
+    from flash_attn.bert_padding import unpad_input
     from transformers.modeling_flash_attention_utils import _flash_attention_forward
 
 logger = logging.get_logger(__name__)
@@ -706,10 +707,11 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
         # max_length_k = None
         if attention_mask is not None:
             logger.warning_once("Using attention mask for resampler")
-            seq_lens_k = attention_mask.sum(dim=-1, dtype=torch.int32)
-            cu_seq_lens_k = torch.cumsum(seq_lens_k, dim=0, dtype=torch.int32)
-            max_length_k = seq_lens_k.max().item()
-            cu_seq_lens_k = nn.functional.pad(cu_seq_lens_k, (1, 0))
+            context, _, cu_seq_lens_k, max_length_k, _ = unpad_input(
+                context, attention_mask
+            )
+            context = context.unsqueeze(0)
+            position_ids = True  # goes down flash attn path that uses cu_seq_lens
 
         elif position_ids is not None:
             logger.warning_once("Using position ids for resampler")
@@ -739,31 +741,26 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
 
         else:
             raise ValueError("either position_ids or attention_mask is required")
-
         x_attn_kwargs = dict(
-            attention_mask=attention_mask,
+            # attention_mask=attention_mask,
             position_ids=position_ids,
             cu_seq_lens_q=cu_seq_lens_q,
             cu_seq_lens_k=cu_seq_lens_k,
             max_length_q=max_length_q,
             max_length_k=max_length_k,
         )
-        self_attn_mask = (
-            torch.ones(
-                (bsz, self.n_latents), device=context.device, dtype=torch.float32
-            )
-            if attention_mask is not None
-            else None
-        )
-        self_attn_position_ids = (
-            torch.arange(
-                self.n_latents, device=position_ids.device, dtype=torch.int32
-            ).repeat(1, bsz)
-            if position_ids is not None
-            else None
-        )
+        # self_attn_mask = (
+        #     torch.ones(
+        #         (bsz, self.n_latents), device=context.device, dtype=torch.float32
+        #     )
+        #     if attention_mask is not None
+        #     else None
+        # )
+        self_attn_position_ids = torch.arange(
+            self.n_latents, device=context.device, dtype=torch.int32
+        ).repeat(1, bsz)
         self_attn_kwargs = dict(
-            attention_mask=self_attn_mask,
+            # attention_mask=self_attn_mask,
             position_ids=self_attn_position_ids,
             cu_seq_lens_q=cu_seq_lens_q,
             cu_seq_lens_k=cu_seq_lens_q,

watcher.py

@@ -93,7 +93,8 @@ if __name__ == "__main__":
                 generative=True,
             )
             metrics.update(gen_metrics)
-            wandb.log(metrics, step=curstep)
+            for k in metrics:
+                wandb.log(metrics[k], step=curstep)
             wandb.finish()
             print(f"Logged metrics: {metrics}")
             print("=" * 80)