chunked ctx eval (#9)

* got burned by floating point precision again :)))))))))))

* split-ctx working eval (batch_size=1)

* batched lora aggregation (sum + mean) eval + finegrain eval len bins

* fix ctx_ids assert

README.md

@@ -204,7 +204,10 @@ WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/ru
 WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/checkpoint-105000/pytorch_model.bin --datasets negative_nq triviaqa_retrieved squad longbench_e --split test
 
 # squad only
-WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Jul08_15-52-27_slurm0-a3nodeset-7_71097_2c7c4d75/checkpoint-12000/pytorch_model.bin --datasets squad --split test
+WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad --split test
+
+# chunking
+WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad --split validation --max_ctx_chunk_len 100 --max_val_samples_per_ds 10 --lora_aggregation mean
 
 # base model
 WANDB_MODE=disabled uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --eval_batch_size 2

run_eval.py

@@ -51,6 +51,27 @@ if __name__ == "__main__":
         default=32,
         help="Eval batch size for generation",
     )
+    parser.add_argument(
+        "--max_val_samples_per_ds",
+        type=int,
+        default=-1,
+        help=(
+            "Maximum number of validation samples per dataset. "
+            "If -1, uses values from checkpoint config."
+        ),
+    )
+    parser.add_argument(
+        "--max_ctx_chunk_len",
+        type=int,
+        default=-1,
+        help="Maximum length of context chunk for evaluation",
+    )
+    parser.add_argument(
+        "--lora_aggregation",
+        choices=["mean", "sum"],
+        default="sum",
+        help="LoRA aggregation method",
+    )
     parser.add_argument(
         "--max_new_tokens",
         type=int,
@@ -68,16 +89,16 @@ if __name__ == "__main__":
 
     eval_batch_size_gen = cli_args.pop("eval_batch_size_gen")
     eval_batch_size = cli_args.pop("eval_batch_size")
-    run_eval(
-        **cli_args,
-        # cli_args.checkpoint_path,
-        # cli_args.model_name_or_path,
-        # cli_args.eval_batch_size,
-        # args,
-        # split=cli_args.split,
-        eval_batch_size=eval_batch_size,
-        generative=False,
-    )
+    # run_eval(
+    #     **cli_args,
+    #     # cli_args.checkpoint_path,
+    #     # cli_args.model_name_or_path,
+    #     # cli_args.eval_batch_size,
+    #     # args,
+    #     # split=cli_args.split,
+    #     eval_batch_size=eval_batch_size,
+    #     generative=False,
+    # )
     run_eval(
         **cli_args,
         # cli_args.checkpoint_path,

src/ctx_to_lora/configs.py

@@ -358,10 +358,6 @@ class CtxTrainingArguments:
         default=2**13,
         metadata={"help": "Maximum base length for training."},
     )
-    max_ctx_len: int | None = field(
-        default=2**13,
-        metadata={"help": "Maximum context length for training."},
-    )
     use_sequence_packing: bool = field(
         default=True,
         metadata={"help": "Whether to use sequence packing."},
@@ -384,6 +380,17 @@ class CtxTrainingArguments:
             "they will be split up into multiple samples."
         },
     )
+    max_ctx_chunk_len: int = field(
+        default=-1,
+        metadata={
+            "help": "Max context chunk length. If a context is longer than this value, "
+            "it will be split up into multiple chunks."
+        },
+    )
+    max_ctx_chunk_num: int = field(
+        default=-1,
+        metadata={"help": "Max number of context chunks per sample."},
+    )
     max_packed_inp_len: int | None = field(
         default=2**14,
         metadata={"help": "Maximum packed input length for training."},

src/ctx_to_lora/data/collator.py

@@ -5,7 +5,7 @@ from transformers.data import (
     default_data_collator,
 )
 
-from ctx_to_lora.utils import check_is_iterable
+from ctx_to_lora.utils import check_is_iterable, concat_list
 
 flattener = DataCollatorWithFlattening()
 
@@ -35,7 +35,13 @@ def flatten_if_not_packed(inp_list):
 
     packed_inputs["n_queries"] = n_queries
     if "ctx_ids" in sample:
-        ctx_ids = [{"input_ids": example["ctx_ids"]} for example in inp_list]
+        # HACK: assumes 1 ctx chunk here
+        # TODO: fix this after we have chunk ctx training
+        assert all(len(ctx_ids) == 1 for ctx_ids in sample["ctx_ids"]), (
+            "ctx_ids can only have one chunk for eval. "
+            "Please implement chunked ctx forward pass to handle this."
+        )
+        ctx_ids = [{"input_ids": example["ctx_ids"][0]} for example in inp_list]
         packed_ctx = flattener(ctx_ids, return_tensors="pt")
         packed_inputs["ctx_ids"] = packed_ctx["input_ids"]
         packed_inputs["ctx_position_ids"] = packed_ctx["position_ids"]
@@ -51,13 +57,14 @@ def flatten_if_not_packed(inp_list):
 def eval_collator(inp_list, tokenizer):
     # only used for teacher-forcing eval
     # input is a list of tokenized sequences
+    # TODO: handle list input (used to be torch.Tensor)
     padding_kwargs = dict(padding=True, padding_side="right", return_tensors="pt")
 
     has_ctx_ids = "ctx_ids" in inp_list[0]
     if has_ctx_ids:
         # pad to the longest ctx_len in the batch
         # which can have a different length from the input_ids, attn_mask, labels
-
+        # TODO: handle chunked ctx_ids
         ctx_ids = [example.pop("ctx_ids") for example in inp_list]
         ctx_attn_mask = [torch.ones_like(x) for x in ctx_ids]
         ctx_ids = torch.nn.utils.rnn.pad_sequence(
@@ -92,7 +99,7 @@ def eval_collator(inp_list, tokenizer):
 
 def generation_collator(inp_list, tokenizer):
     padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt")
-    input_ids = [x.pop("input_ids") for x in inp_list]
+    input_ids = [torch.tensor(x.pop("input_ids")) for x in inp_list]
     labels = [x.pop("labels") for x in inp_list]
     for i, label in enumerate(labels):
         # we don't include the labels in the output during generation
@@ -107,10 +114,15 @@ def generation_collator(inp_list, tokenizer):
     )
 
     if "ctx_ids" in inp_list[0]:
+        # TODO: handle chunked ctx_ids
         # pad to the longest ctx_len in the batch
         # which can have a different length from the input_ids, attn_mask, labels
         ctx_ids = [example.pop("ctx_ids") for example in inp_list]
+        n_chunks = [len(x) for x in ctx_ids]
+        ctx_ids = concat_list(ctx_ids)
+        ctx_ids = [torch.tensor(x) for x in ctx_ids]
         ctx_attn_mask = [torch.ones_like(x) for x in ctx_ids]
+
         ctx_ids = torch.nn.utils.rnn.pad_sequence(
             ctx_ids,
             batch_first=True,
@@ -123,4 +135,9 @@ def generation_collator(inp_list, tokenizer):
         )
         out["ctx_ids"] = ctx_ids
         out["ctx_attn_mask"] = ctx_attn_mask
+
+        out["n_ctx_chunks"] = torch.tensor(n_chunks, dtype=torch.int32)
+        # print(f"ctx_ids: {ctx_ids}")
+        # print(f"n_chunks: {n_chunks}")
+        # breakpoint()
     return out

src/ctx_to_lora/data/definitions.py

@@ -47,6 +47,18 @@ REPEAT_PROMPTS = [
     "Repeat precisely what is written in the information above.",
 ]
 
+# for chunking
+CTX_AFFIXES = {
+    "google/gemma-3-1b-it": {
+        "prefix": [2, 105, 2364, 109],  # <bos><start_of_turn>user\n\n\n
+        "suffix": [106, 107, 105, 4368, 107],  # <end_of_turn>\n<start_of_turn>model\n
+    },
+    "google/gemma-2-2b-it": {
+        "prefix": [2, 106, 1645, 110],  # <bos><start_of_turn>user\n\n\n
+        "suffix": [107, 108, 106, 2516, 108],  # <end_of_turn>\n<start_of_turn>model\n
+    },
+}
+
 # approximate length of the datasets (train split)
 # needed for streaming datasets
 DS_LEN = {

src/ctx_to_lora/data/processing.py

@@ -5,7 +5,7 @@ import random
 from collections.abc import Callable
 from glob import glob
 from hashlib import sha256
-from math import isclose
+from math import ceil, isclose
 from os import path
 from typing import Any
 
@@ -15,6 +15,7 @@ from datasets import Dataset, interleave_datasets, is_caching_enabled, load_data
 from transformers import PreTrainedTokenizerBase
 
 from ctx_to_lora.data.definitions import (
+    CTX_AFFIXES,
     DS_KWARGS,
     IGNORE_INDEX,
     RAW_DATA_DIR,
@@ -208,6 +209,8 @@ def get_tokenized_dataset(
     tokenizer: PreTrainedTokenizerBase,
     ctx_model_max_len: int,
     ctx_tokenizer: PreTrainedTokenizerBase,
+    max_ctx_chunk_len: int,
+    max_ctx_chunk_num: int,
     add_ctx_to_chat: bool,
     use_kl_loss: bool,
     max_new_tokens: int = 256,
@@ -230,7 +233,8 @@ def get_tokenized_dataset(
         base_model_max_len=base_model_max_len,
         ctx_model_max_len=ctx_model_max_len,
         add_ctx_to_chat=add_ctx_to_chat,
-        use_kl_loss=True,  # HACK: for consistent hash (to be removed)
+        max_ctx_chunk_len=max_ctx_chunk_len,
+        max_ctx_chunk_num=max_ctx_chunk_num,
         need_ctx_ids=need_ctx_ids,
         split=split,
         max_new_tokens=max_new_tokens,
@@ -269,9 +273,6 @@ def get_tokenized_dataset(
             )
     logger.info(f"Constructing and tokenizing {ds_name} with {split} split...")
 
-    # HACK
-    tokenize_kwargs.pop("use_kl_loss", None)
-
     tokenized_ds = construct_and_tokenize_ctx_qa(
         ds=ds,
         tokenizer=tokenizer,
@@ -302,12 +303,15 @@ def construct_and_tokenize_ctx_qa(
     ctx_tokenizer,
     add_ctx_to_chat,
     need_ctx_ids,
+    max_ctx_chunk_len,
+    max_ctx_chunk_num,
     ds,
     split,
     max_new_tokens,
     set_format=None,
     num_proc=None,
 ):
+    is_train = "train" in split
     # for sft + chat_model, we need to convert the dataset to chat format
     if "input_ids" in ds.column_names and "response_start_end" in ds.column_names:
         # already tokenized dataset (e.g., self-gen qa data)
@@ -349,15 +353,32 @@ def construct_and_tokenize_ctx_qa(
                 batch_size=100_000,
                 remove_columns=["context"],
             )
-        if "train" in split:
-            # TODO: do something if ctx length is longer than the ctx model length
-            # e.g., drop or split for multi-lora training
+
+        # TODO: this can be removed once we implement ctx chunking for training
+        if is_train:
+            # drop
             tokenized_ds = tokenized_ds.filter(
                 len_filter,
                 fn_kwargs={"max_length": ctx_model_max_len, "keys": ["ctx_ids"]},
                 num_proc=16,
             )
 
+        # ctx chunking
+        # Ideally we want to chunk the raw text directly
+        # however, since the contexts are tokenized during self-gen
+        # we can only chunk the tokenized context which requires some workaround
+        # e.g., removing/applying template to each chunk
+        # with some big caveats, e.g., losing order info
+        tokenized_ds = tokenized_ds.map(
+            split_too_long_ctx,
+            fn_kwargs={
+                "max_length": max_ctx_chunk_len,
+                "max_num_split": max_ctx_chunk_num,
+                "model_name_or_path": tokenizer.name_or_path,
+                "is_train": is_train,
+            },
+        )
+
     logging.info(f"Split too long QAs with max length {max_qas_len}")
     tokenized_ds = tokenized_ds.map(
         split_too_long_qas,
@@ -385,15 +406,16 @@ def construct_and_tokenize_ctx_qa(
             fn_kwargs={"columns": ["input_ids"]},
         )
         if "ctx_ids" in tokenized_ds.column_names:
-            tokenized_ds = tokenized_ds.map(
-                truncate_middle_if_too_long,
-                fn_kwargs={
-                    "max_length": ctx_model_max_len,
-                    "columns": ["ctx_ids"],
-                    # cxt encoder doesnt need to add new_tokens
-                    "max_new_tokens": 0,
-                },
-            )
+            # TODO: remove since we already have ctx chunking for eval
+            # tokenized_ds = tokenized_ds.map(
+            #     truncate_middle_if_too_long,
+            #     fn_kwargs={
+            #         "max_length": ctx_model_max_len,
+            #         "columns": ["ctx_ids"],
+            #         # cxt encoder doesnt need to add new_tokens
+            #         "max_new_tokens": 0,
+            #     },
+            # )
             tokenized_ds = tokenized_ds.map(
                 add_length_info,
                 fn_kwargs={"columns": ["ctx_ids"]},
@@ -468,6 +490,7 @@ def get_sft_prompt_formatting_fn(
             add_special_tokens=False,
             padding=False,
             truncation=False,
+            return_attention_mask=False,
             add_generation_prompt=False,
             return_assistant_tokens_mask=True,
             return_dict=True,
@@ -537,6 +560,73 @@ def convert_ctx_prompt_response_to_messages(
     return {"messages_list": messages_list}
 
 
+def split_too_long_ctx(
+    sample: dict[str, Any],
+    max_length: int,
+    max_num_split: int,
+    model_name_or_path: str,
+    is_train: bool,
+) -> dict[str, Any]:
+    """
+    Split context into smaller chunks if it exceeds the maximum length.
+
+    Args:
+        samples: Dictionary containing 'ctx_ids' and 'ctx_attn_mask'
+        max_length: Maximum length for each context chunk
+        max_num_split: Maximum number of splits allowed
+
+    Returns:
+        Dictionary with split context data
+    """
+
+    if is_train:
+        # TODO: for training, we might wanna sort the context by length since merging
+        # a batch of chunked loras need padded in the rank axes
+        # e.g., ctx1 has 5 chunks (rank-48), ctx2 has 10 chunks (rank-88)
+        #
+        # TODO: should be stochastic during training?
+        # e.g., even if the ctx is not too long, we still split it randomly?
+        # say 0-4k is max len for one split,
+        # for some ctx shorter than 4k, we leave it as is
+        # for some we split to smaller chunks
+        # For eval, always split at max len?
+        # raise NotImplementedError(
+        #     "Splitting contexts for training is not implemented yet. "
+        # )
+        # # Limit the number of splits for training?
+        # if len(chunks) > max_num_split > 0:
+        #     chunks = chunks[:max_num_split]
+        # TODO: training inputs right now is not nested
+        return sample
+
+    ctx_affixes = CTX_AFFIXES[model_name_or_path]
+    prefix = ctx_affixes["prefix"]
+    suffix = ctx_affixes["suffix"]
+    ctx_ids = sample["ctx_ids"]
+
+    if max_length < 0 and max_num_split < 0:
+        return {"ctx_ids": [ctx_ids]}
+
+    if len(ctx_ids) <= max_length:
+        return {"ctx_ids": [ctx_ids]}
+
+    # can we improve how we chunk?
+    # compute num chunks for more uniform chunking
+    n_chunks = ceil(len(ctx_ids) / max_length)
+    avg_len = ceil(len(ctx_ids) / n_chunks)
+
+    # Split the context into smaller chunks
+    chunks = [ctx_ids[i : i + avg_len] for i in range(0, len(ctx_ids), avg_len)]
+
+    # this would exceed the avg_len a bit
+    chunks[0] = chunks[0] + suffix
+    for i in range(1, len(chunks) - 1):
+        chunks[i] = prefix + chunks[i] + suffix
+    chunks[-1] = prefix + chunks[-1]
+
+    return {"ctx_ids": chunks}
+
+
 def split_too_long_qas(
     samples: dict[str, any], max_qas_len: int, max_qas_per_sample: int
 ):
@@ -722,7 +812,13 @@ def squeeze_tokens(sample: dict[str, Any]) -> dict[str, Any]:
 
 
 def add_length_info(sample: dict[str, any], columns: list[str]) -> dict[str, any]:
-    return {f"{k}_len": len(sample[k]) for k in columns}
+    out = {}
+    for k in columns:
+        if check_is_iterable(sample[k][0]):
+            out[f"{k}_len"] = sum([len(x) for x in sample[k]])
+        else:
+            out[f"{k}_len"] = len(sample[k])
+    return out
 
 
 def truncate_middle_if_too_long(
@@ -765,6 +861,7 @@ def tokenize_ctx_text(
             ],
             tokenize=True,
             add_generation_prompt=True,
+            return_attention_mask=False,
             padding=False,
             truncation=False,
             add_special_tokens=False,  # special tokens are already added by the chat template

src/ctx_to_lora/eval_utils.py

@@ -42,7 +42,7 @@ from ctx_to_lora.metrics import (
 from ctx_to_lora.model_loading import get_model, get_tokenizer
 from ctx_to_lora.modeling import hypernet
 from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
-from ctx_to_lora.utils import clear_gpu, setup_logging
+from ctx_to_lora.utils import clear_gpu, concat_list, setup_logging
 
 logger = logging.getLogger()
 
@@ -529,7 +529,7 @@ def decode_test_result(
         d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True).strip()
         if "ctx_ids" in sample:
             d["context"] = ctx_tokenizer.decode(
-                sample["ctx_ids"], skip_special_tokens=False
+                concat_list(sample["ctx_ids"]), skip_special_tokens=False
             )
         for k in sample:
             if k.endswith("_len"):
@@ -612,7 +612,7 @@ def eval_generation(
                     grouped_texts[group_key]["generated"].append(txt["generated"])
                     grouped_texts[group_key]["label"].append(txt["label"])
                     grouped_texts[group_key]["count"] += 1
-                    break
+                    # break
 
         for group_key, data in grouped_texts.items():
             if data["count"] > 0:
@@ -634,7 +634,7 @@ def eval_generation(
                                         )
                                     else:
                                         group_answers.append([txt["label"]])
-                                break
+                                # break
 
                     group_qa_f1_metric, _ = compute_qa_f1_score(
                         data["generated"], group_answers
@@ -727,6 +727,8 @@ def evaluate(
     eval_batch_size: int,
     args: Namespace,
     split: str,
+    max_ctx_chunk_len: int,
+    lora_aggregation: str,
     max_new_tokens: int,
     generative: bool,
 ) -> dict[str, dict]:
@@ -785,6 +787,8 @@ def evaluate(
         get_tokenized_dataset,
         max_qas_len=-1,
         max_qas_per_sample=1,
+        max_ctx_chunk_len=max_ctx_chunk_len,
+        max_ctx_chunk_num=-1,
         base_model_max_len=model.base_model.config.max_position_embeddings,
         tokenizer=tokenizer,
         ctx_model_max_len=ctx_model_max_len,
@@ -860,6 +864,11 @@ def evaluate(
 
     collator = generation_collator if generative else eval_collator
 
+    if max_ctx_chunk_len > 0:
+        model.generate = model.generate_with_multi_loras
+        print(f"Using {lora_aggregation} for aggregation generated LoRAs.")
+        model.generate = partial(model.generate, lora_aggregation=lora_aggregation)
+
     trainer_kwargs = {
         "model": model,
         "args": eval_trainer_args,
@@ -909,6 +918,9 @@ def run_eval(
     datasets: list[str] = None,
     split: str = "validation",
     eval_batch_size: int = 8,
+    max_val_samples_per_ds: int = -1,
+    max_ctx_chunk_len: int = -1,
+    lora_aggregation: str = "sum",
     remove_context: bool = False,
     max_new_tokens: int = 256,
     generative: bool = False,
@@ -960,6 +972,8 @@ def run_eval(
             test_ds_names=[],
             remove_context=remove_context,
         )
+    if max_val_samples_per_ds > 0:
+        args.max_val_samples_per_ds = max_val_samples_per_ds
     setup_logging(args.logging_dir)
 
     # Override dataset names if provided via CLI
@@ -975,6 +989,8 @@ def run_eval(
         eval_batch_size,
         args,
         split,
+        max_ctx_chunk_len,
+        lora_aggregation,
         max_new_tokens,
         generative=generative,
     )

src/ctx_to_lora/metrics.py

@@ -7,6 +7,11 @@ from rouge_score import rouge_scorer
 from transformers import EvalPrediction
 
 LENGTH_BINS = [
+    # finegrain bins
+    (0, 2**7 - 1),
+    (2**7, 2**8 - 1),
+    (2**8, 2**9 - 1),
+    # coarse bins
     (0, 2**9 - 1),
     (2**9, 2**10 - 1),
     (2**10, 2**11 - 1),

src/ctx_to_lora/modeling/hypernet.py

@@ -3,7 +3,7 @@ from collections.abc import Iterable
 from contextlib import contextmanager
 from dataclasses import dataclass
 from functools import partial
-from typing import Any
+from typing import Any, Literal
 
 import torch
 from einops import rearrange, unpack
@@ -50,6 +50,7 @@ from ctx_to_lora.modeling.lora_layer import (
     lora_forward,
     lora_forward_packed,
 )
+from ctx_to_lora.modeling.lora_merger import combine_lora
 from ctx_to_lora.utils import (
     get_layers,
     get_num_layers,
@@ -233,7 +234,7 @@ class HyperLoRA(nn.Module):
         self.lora_config = self.config.lora_config
 
         self.target_modules = (
-            self.lora_config.target_modules if self.lora_config else None
+            tuple(sorted(self.lora_config.target_modules)) if self.lora_config else None
         )
         self.num_modules = len(self.target_modules) if self.target_modules else 0
         self.extra_modules = (
@@ -526,7 +527,25 @@ class HyperLoRA(nn.Module):
         #             hidden_size=self.config.base_hidden_size,
         #         )
 
-    @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+    def get_head_bias(self):
+        bias_list = unpack(
+            self.head.bias,
+            [[] for _ in range(len(self.target_modules))],
+            "bs n_layers * r max_io_dim",
+        )
+        bias_dict = dict()
+        for module, bias in zip(self.target_modules, bias_list):
+            bias_A, bias_B = unpack(
+                bias[..., : self.d_in[module] + self.d_out[module]],
+                [[self.d_in[module]], [self.d_out[module]]],
+                "bs n_layers r *",
+            )
+            # transpose B
+            bias_B = rearrange(bias_B, "bs n_layers r d_out -> bs n_layers d_out r")
+            bias_dict[module] = dict(A=bias_A, B=bias_B)
+        return bias_dict
+
+    # @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
     def _to_lora_dict(
         self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"]
     ) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]:
@@ -582,7 +601,6 @@ class HyperLoRA(nn.Module):
         )
         return {k: v for k, v in zip(self.extra_modules, layernorms)}
 
-    @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
     def forward(
         self,
         features: Float[Tensor, "bs seq_len feature_dim"],
@@ -590,7 +608,9 @@ class HyperLoRA(nn.Module):
         position_ids: Integer[Tensor, "bs seq_len"] | None = None,
     ):
         # [bs, n_layers x n_total_modules x r, feature_dim]
-        lora_emb, extra_emb = self.aggregator(features, attn_mask, position_ids)
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            # OMG!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+            lora_emb, extra_emb = self.aggregator(features, attn_mask, position_ids)
         # lora_emb, extra_emb = unpack(
         #     emb,
         #     [[self.num_modules], [self.num_extra_modules]],
@@ -612,7 +632,6 @@ class HyperLoRA(nn.Module):
 
         return flat_loras, flat_layernorms
 
-    @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
     def generate_weights(
         self,
         features: Float[Tensor, "bs seq_len feature_dim"],
@@ -953,20 +972,22 @@ class ModulatedPretrainedModel(nn.Module):
         ctx_ids: Integer[Tensor, "bs ctx_length"],
         ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None,
         ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None,
+        n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None,
         n_queries: Integer[Tensor, "n_ctx"] | None = None,
+        lora_aggregation: Literal["mean", "sum"] = "sum",
         *model_inputs_args: Any,
         **model_inputs_kwargs: dict[str, Any],
     ):
-        # TODO: make this persistent
         generated_loras, _ = self.generate_weights(
             ctx_ids, ctx_attn_mask, ctx_position_ids
         )
-        # sum loras
-        for lora_at_module in generated_loras.values():
-            if len(lora_at_module["A"]) > 1:
-                print("Multiple LoRAs are generated, summing them up")
-            lora_at_module["A"] = torch.mean(lora_at_module["A"], dim=0, keepdim=True)
-            lora_at_module["B"] = torch.mean(lora_at_module["B"], dim=0, keepdim=True)
+
+        generated_loras = combine_lora(
+            generated_loras,
+            n_ctx_chunks,
+            aggregation=lora_aggregation,
+            lora_bias=self.hypernet.get_head_bias(),
+        )
 
         # apply lora hook to the base model
         position_ids = (
@@ -975,16 +996,10 @@ class ModulatedPretrainedModel(nn.Module):
             else None
         )
         if n_queries is None:
-            if ctx_position_ids is None:
-                n_queries = torch.ones(
-                    ctx_ids.shape[0], dtype=torch.int32, device=self.device
-                )
-            else:
-                # quite redundant (we do cu_seqlens many places)
-                # TODO: compute cu_seqlens here and propagate that
-                n_queries = torch.ones(
-                    (ctx_position_ids == 0).sum(), dtype=torch.int32, device=self.device
-                )
+            # assumes each group has one q
+            n_queries = torch.ones(
+                len(n_ctx_chunks), dtype=torch.int32, device=self.device
+            )
 
         apply_lora_to_layers(
             self.base_model,
@@ -1003,9 +1018,10 @@ class ModulatedPretrainedModel(nn.Module):
     def generate(
         self,
         # TODO: allow more than one LoRA per sample (multi-lora)
-        ctx_ids: Integer[Tensor, "bs ctx_length"] | None = None,
-        ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None,
-        ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None,
+        ctx_ids: Integer[Tensor, "n_chunks ctx_length"] | None = None,
+        ctx_attn_mask: Integer[Tensor, "n_chunks ctx_length"] | None = None,
+        ctx_position_ids: Integer[Tensor, "n_chunks ctx_length"] | None = None,
+        n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None,
         n_queries: Integer[Tensor, "n_ctx"] | None = None,
         *model_inputs_args: Any,
         **model_inputs_kwargs: dict[str, Any],
@@ -1053,7 +1069,18 @@ class ModulatedPretrainedModel(nn.Module):
                 ctx_ids, ctx_attn_mask, ctx_position_ids
             )
 
+        # generated_loras = combine_lora(
+        #     generated_loras,
+        #     n_chunks=n_ctx_chunks,
+        #     aggregation="sum",
+        #     lora_bias=self.hypernet.get_head_bias(),
+        # )
+
+        # # for generation the ctx_ids are batched not packed
+        # ctx_ids_list = torch.split(ctx_ids, n_ctx_chunks, dim=0)
+
         # apply lora hook to the base model
+        # TODO: we dont this position_ids for generation?
         position_ids = (
             model_inputs_kwargs["position_ids"]
             if "position_ids" in model_inputs_kwargs

src/ctx_to_lora/modeling/lora_merger.py

@@ -0,0 +1,200 @@
+"""
+Utilities for merging / aggregating LoRA adapters coming from multiple chunks.
+"""
+
+from typing import Literal
+
+import torch
+from einops import rearrange
+from jaxtyping import Integer
+from torch import Tensor
+
+
+def compute_rank(n_lora, rank, aggregation):
+    """Return the effective (max) rank required for a group.
+
+    aggregation == 'mean': we only concatenate the existing LoRA ranks (n_lora * r)
+    aggregation == 'sum' : an additional bias term may be appended (+1 * r)
+    """
+    if aggregation == "mean":
+        return n_lora * rank
+    elif aggregation == "sum":
+        return (n_lora + 1) * rank
+    else:
+        raise ValueError(f"Unknown aggregation: {aggregation}")
+
+
+def combine_lora(
+    generated_loras: dict[str, dict[str, Tensor]],
+    n_chunks: Integer[Tensor, "n_ctx"],
+    aggregation: Literal["mean", "sum"],
+    lora_bias: dict[str, dict[str, Tensor]] = None,
+) -> dict[str, dict[str, Tensor]]:
+    """Combine per-chunk LoRA adapter weights into grouped higher-rank adapters.
+
+    Parameters
+    ----------
+    generated_loras
+        Mapping: module_name -> {'A': Tensor, 'B': Tensor}
+        Each tensor shape: [total_chunks, n_layers, r, dim] for A, and
+        [total_chunks, n_layers, dim, r] for B (but we treat generically here).
+    n_chunks
+        1D tensor giving number of chunks in each aggregation group. len(n_chunks)=G.
+    aggregation
+        'mean' or 'sum'. Controls scaling and bias handling.
+    lora_bias
+        Optional mapping mirroring generated_loras providing bias LoRA tensors (same
+        per-chunk shape except missing the leading chunk dimension which is expanded).
+
+    Returns
+    -------
+    dict
+        module_name -> {'A': combined_A, 'B': combined_B}
+        combined_A shape: [G, n_layers, max_rank, dim]
+        combined_B shape: [G, n_layers, dim, max_rank]
+    """
+    assert aggregation in ["mean", "sum"]
+
+    # Precompute common quantities
+    num_groups = len(n_chunks)
+    total_chunks = int(n_chunks.sum())
+    sqrt_group_sizes = n_chunks**0.5  # [G]
+    # Repeat for each chunk within a group: e.g.,
+    # if n_chunks=[2,3] -> repeat factors [sqrt2, sqrt2, sqrt3, sqrt3, sqrt3]
+    sqrt_group_sizes_per_chunk = sqrt_group_sizes.repeat_interleave(n_chunks, dim=0)
+
+    # Assume all modules share same base rank r
+    # (second-to-last dim index for A, last dim for B)
+    # Use first module's A tensor to infer rank
+    first_module = next(iter(generated_loras))
+    base_rank = generated_loras[first_module]["A"].shape[-2]
+    max_rank_needed = max(compute_rank(n, base_rank, aggregation) for n in n_chunks)
+
+    # Initialize output container
+    combined_loras: dict[str, dict[str, Tensor]] = {
+        module: {"A": None, "B": None} for module in generated_loras.keys()
+    }
+
+    # Iterate over modules and LoRA matrices A / B.
+    # For A we concatenate along its rank dimension index=2; for B along index=3.
+    for module_name, module_loras in generated_loras.items():
+        for matrix_key, concat_dim in (("A", 2), ("B", 3)):
+            # Shape conventions (both stored with leading chunk axis):
+            #   loras: [total_chunks, n_layers, r, dim] (A)
+            #       or [total_chunks, n_layers, dim, r] (B)
+            loras = module_loras[matrix_key]
+
+            # Expand bias if provided
+            bias_tensor = None
+            if lora_bias is not None:
+                bias_tensor = lora_bias[module_name][matrix_key].expand(
+                    num_groups, *loras.shape[1:]
+                )
+
+            # For 'mean' aggregation, scale each chunk's LoRA by sqrt(group_size)
+            if aggregation == "mean":
+                loras = loras / sqrt_group_sizes_per_chunk.view(total_chunks, 1, 1, 1)
+
+            # Split the leading chunk dimension into groups per n_chunks specification
+            per_group_list = loras.split(n_chunks.tolist(), dim=0)
+
+            # Concat each group
+            if matrix_key == "A":
+                rearrange_pattern = "chunks n_layers r dim -> 1 n_layers (chunks r) dim"
+            else:
+                rearrange_pattern = "chunks n_layers dim r -> 1 n_layers dim (chunks r)"
+            per_group_deltas = [
+                rearrange(group_tensor, rearrange_pattern)
+                for group_tensor in per_group_list
+            ]
+
+            # Delegate to helper for combining rank expansion across groups.
+            combined_loras[module_name][matrix_key] = _combine_single_matrix(
+                per_group_deltas=per_group_deltas,
+                bias=bias_tensor,
+                n_chunks=n_chunks,
+                base_rank=base_rank,
+                max_rank_needed=max_rank_needed,
+                aggregation=aggregation,
+                matrix_key=matrix_key,
+                concat_dim=concat_dim,
+            )
+
+    return combined_loras
+
+
+def _combine_single_matrix(
+    *,
+    per_group_deltas: list[Tensor],
+    bias: Tensor | None,
+    n_chunks: Tensor,
+    base_rank: int,
+    max_rank_needed: int,
+    aggregation: str,
+    matrix_key: str,
+    concat_dim: int,
+) -> Tensor:
+    """Assemble a single LoRA matrix (either 'A' or 'B') across groups.
+
+    Parameters
+    ----------
+    loras
+        Tensor of shape (total_chunks, n_layers, r, dim) if 'A' else
+        (total_chunks, n_layers, dim, r) if 'B'.
+    per_group_deltas
+        A list of concat'd tensors
+    bias_tensor
+        Expanded bias tensor of shape (G, n_layers, r, dim) for 'A' or
+        (G, n_layers, dim, r) for 'B', or None.
+    n_chunks
+        1D tensor (G,) with number of chunks per group.
+    base_rank
+        Original (per-LoRA) rank r.
+    max_rank_needed
+        Maximum rank capacity allocated for any group (accounts for bias when summing).
+    aggregation
+        'mean' or 'sum'. For 'sum' we optionally append a bias-related term.
+    matrix_key
+        'A' or 'B'. Determines which axis is the rank axis and bias handling sign.
+    concat_dim
+        Dimension along which to concatenate within each group's tensors.
+
+    Returns
+    -------
+    Tensor
+        Combined LoRA matrix with shape (G, n_layers, max_rank, dim) for 'A'
+        or (G, n_layers, dim, max_rank) for 'B'.
+    """
+    assert matrix_key in ("A", "B"), f"matrix_key must be 'A' or 'B', got {matrix_key}"
+
+    num_groups = len(n_chunks)
+    combined_shape = list(bias.shape)
+    combined_shape[0] = num_groups
+
+    # Adjust the rank dimension size placeholder.
+    if matrix_key == "A":
+        rank_dim_index = -2  # (G, layers, rank, dim)
+    else:
+        rank_dim_index = -1  # (G, layers, dim, rank)
+    combined_shape[rank_dim_index] = max_rank_needed
+
+    combined = torch.zeros(*combined_shape, device=bias.device, dtype=bias.dtype)
+
+    for g, deltas in enumerate(per_group_deltas):
+        # we remove leading dim when rearrange
+        combined_rank = deltas.shape[concat_dim]
+
+        # Build slice pattern: rank slice up to combined_rank.
+        slice_pattern = [g, slice(None), slice(None), slice(None)]
+        slice_pattern[concat_dim] = slice(combined_rank)
+        combined[slice_pattern] = deltas
+        if bias is not None and aggregation == "sum":
+            if matrix_key == "A":
+                bias_to_add = -bias[g : g + 1] * (n_chunks[g] - 1)
+            else:  # 'B'
+                bias_to_add = bias[g : g + 1]
+
+            slice_pattern[concat_dim] = slice(combined_rank, combined_rank + base_rank)
+            combined[slice_pattern] = bias_to_add
+
+    return combined

train.py

@@ -259,6 +259,8 @@ def main():
         ctx_model_max_len=ctx_model_max_len,
         ctx_tokenizer=ctx_tokenizer,
         add_ctx_to_chat=add_ctx_to_chat,
+        max_ctx_chunk_len=ctx_args.max_ctx_chunk_len,
+        max_ctx_chunk_num=ctx_args.max_ctx_chunk_num,
         use_kl_loss=ctx_args.use_kl_loss,
     )
     splits = ["train"]