From ec955a935d72ec813d667b41c7bf14fc692b0d45 Mon Sep 17 00:00:00 2001 From: Rujikorn Charakorn Date: Mon, 11 Aug 2025 18:44:26 +0900 Subject: [PATCH] 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 | 5 +- run_eval.py | 41 +++-- src/ctx_to_lora/configs.py | 15 +- src/ctx_to_lora/data/collator.py | 25 ++- src/ctx_to_lora/data/definitions.py | 12 ++ src/ctx_to_lora/data/processing.py | 133 +++++++++++++--- src/ctx_to_lora/eval_utils.py | 24 ++- src/ctx_to_lora/metrics.py | 5 + src/ctx_to_lora/modeling/hypernet.py | 79 +++++++--- src/ctx_to_lora/modeling/lora_merger.py | 200 ++++++++++++++++++++++++ train.py | 2 + 11 files changed, 474 insertions(+), 67 deletions(-) create mode 100644 src/ctx_to_lora/modeling/lora_merger.py diff --git a/README.md b/README.md index e6b5bc8..1206eaf 100644 --- a/README.md +++ b/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 diff --git a/run_eval.py b/run_eval.py index d0f3301..8c59218 100644 --- a/run_eval.py +++ b/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, diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index b1832c8..db7d26c 100644 --- a/src/ctx_to_lora/configs.py +++ b/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."}, diff --git a/src/ctx_to_lora/data/collator.py b/src/ctx_to_lora/data/collator.py index 2b77c22..9cfe554 100644 --- a/src/ctx_to_lora/data/collator.py +++ b/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 diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py index fec08b7..9433b8c 100644 --- a/src/ctx_to_lora/data/definitions.py +++ b/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], # user\n\n\n + "suffix": [106, 107, 105, 4368, 107], # \nmodel\n + }, + "google/gemma-2-2b-it": { + "prefix": [2, 106, 1645, 110], # user\n\n\n + "suffix": [107, 108, 106, 2516, 108], # \nmodel\n + }, +} + # approximate length of the datasets (train split) # needed for streaming datasets DS_LEN = { diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index eddd00b..73e1c90 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/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 diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py index 404acdf..9d8afd1 100644 --- a/src/ctx_to_lora/eval_utils.py +++ b/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, ) diff --git a/src/ctx_to_lora/metrics.py b/src/ctx_to_lora/metrics.py index 974f6bc..ecd2679 100644 --- a/src/ctx_to_lora/metrics.py +++ b/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), diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index d2daff8..d954348 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/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 diff --git a/src/ctx_to_lora/modeling/lora_merger.py b/src/ctx_to_lora/modeling/lora_merger.py new file mode 100644 index 0000000..ae893fe --- /dev/null +++ b/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 diff --git a/train.py b/train.py index a923b29..94a8850 100755 --- a/train.py +++ b/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"]