everything :)

README.md

@@ -22,4 +22,20 @@ WANDB_MODE=disabled run python hyperlora/intx_sft.py configs/context_numbers_128
 python process_fineweb.py
 python generate_fw_qa.py
 python post_process_fw_qa.py
+```
+
+
+### GSM8k LoRA
+```bash
+run python intx_sft.py configs/gsm8k.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10 --per_device_train_batch_size=16 --gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=lora --target_modules=up_proj,down_proj --eval_steps=5000 --save_steps=5000 --learning_rate=1e-4 --neftune_noise_alpha=5 --load_best_model_at_end=True --metric_for_best_model=gsm8k_loss
+```
+
+### Multitask LoRA
+```bash
+run python intx_sft.py configs/math_and_code.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=4 --gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=lora --target_modules=up_proj,down_proj --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --neftune_noise_alpha=5 --load_best_model_at_end=True --metric_for_best_model=gsm8k_loss
+```
+
+### HyperLoRA finetune on GSM8K
+```bash
+run python intx_sft.py configs/gsm8k.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=16 --gradient_accumulation_steps=1 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=1 --num_self_attends_per_block=16 --self_attention_widening_factor=1 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=512 --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False --add_repeat_prompt=False --ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-11B-Vision-Instruct
 ```

batch_gemma_llama_vision.sh

@@ -1,9 +1,8 @@
 #!/bin/bash
-#SBATCH --job-name=gemma_llama_vision
+#SBATCH --job-name=gemma_pixtral
 #SBATCH --partition=a3
 #SBATCH --nodes=1
 #SBATCH --gpus=8
-#SBATCH --exclude=slurm0-a3nodeset-7
 #SBATCH --output=outputs/%x-%j.out
 #SBATCH --error=outputs/%x-%j.out
 
@@ -19,12 +18,12 @@
 conda activate /home/rujikorn_sakana_ai/.conda/envs/llm-modulator
 # eval "$@"
 
-accelerate launch --num_processes=8 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
+accelerate launch --num_processes=8 --gradient_accumulation_steps=4 --gradient_clipping=1.0 \
 --gpu_ids all --main_process_port 29554 intx_sft.py configs/pretrain_all.yaml \
---model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=8 \
---gradient_accumulation_steps=2 --per_device_eval_batch_size=4 --exp_setup=hyper_lora --aggregator_type=perceiver \
+--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=4 \
+--gradient_accumulation_steps=4 --per_device_eval_batch_size=4 --exp_setup=hyper_lora --aggregator_type=perceiver \
 --target_modules=down_proj,up_proj --num_blocks=1 --num_self_attends_per_block=16 \
---self_attention_widening_factor=4 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 \
+--self_attention_widening_factor=1 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 \
 --neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=512 \
 --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \
 --add_repeat_prompt=False --ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-11B-Vision-Instruct

batch_llama_2b_.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+#SBATCH --job-name=gemma_pixtral
+#SBATCH --partition=a3
+#SBATCH --nodes=1
+#SBATCH --gpus=8
+#SBATCH --output=outputs/%x-%j.out
+#SBATCH --error=outputs/%x-%j.out
+
+# module load
+# module load cuda/12.1
+# module load cudnn/8.9.7
+# module load nccl/cuda-12.1/2.18.3
+# module load hpcx/2.20
+
+# export OMP_NUM_THREADS=24
+# export TRITON_CACHE_DIR=/tmp/.triton/
+. ~/miniconda3/etc/profile.d/conda.sh
+conda activate /home/rujikorn_sakana_ai/.conda/envs/llm-modulator
+# eval "$@"
+
+accelerate launch --num_processes=8 --gradient_accumulation_steps=4 --gradient_clipping=1.0 \
+--gpu_ids all --main_process_port 29554 intx_sft.py configs/pretrain_all.yaml \
+--model_name_or_path=meta-llama/Llama-3.2.3B-Instruct --num_train_epochs=5.1 --per_device_train_batch_size=4 \
+--gradient_accumulation_steps=4 --per_device_eval_batch_size=4 --exp_setup=hyper_lora --aggregator_type=perceiver \
+--target_modules=down_proj --num_blocks=1 --num_self_attends_per_block=16 \
+--self_attention_widening_factor=1 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 \
+--neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=1024 \
+--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \
+--add_repeat_prompt=False --ctx_encoder_model_name_or_path=AIDC-AI/Ovis1.6-Gemma2-9B
+

batch_llama_3_8b_1k_latent.sh

@@ -20,10 +20,10 @@ conda activate /home/rujikorn_sakana_ai/.conda/envs/llm-modulator
 
 accelerate launch --num_processes=8 --gradient_accumulation_steps=4 --gradient_clipping=1.0 \
 --gpu_ids all --main_process_port 29555 intx_sft.py configs/pretrain_all.yaml \
---model_name_or_path=meta-llama/Llama-3.1-8B-Instruct --num_train_epochs=10.1 --per_device_train_batch_size=8 \
---gradient_accumulation_steps=4 --per_device_eval_batch_size=4 --exp_setup=hyper_lora --aggregator_type=perceiver \
---target_modules=down_proj,up_proj --num_blocks=1 --num_self_attends_per_block=32 \
---self_attention_widening_factor=4 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 \
+--model_name_or_path=meta-llama/Llama-3.1-8B-Instruct --num_train_epochs=10.1 --per_device_train_batch_size=4 \
+--gradient_accumulation_steps=4 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver \
+--target_modules=down_proj --num_blocks=1 --num_self_attends_per_block=32 \
+--self_attention_widening_factor=1 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 \
 --neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=1024 \
 --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \
 --add_repeat_prompt=False --ctx_encoder_model_name_or_path=Alibaba-NLP/gte-large-en-v1.5

configs/math_and_code.yaml

@@ -0,0 +1,50 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+label_names: ["labels"]
+# eval_on_start: True
+# eval_strategy: "steps"
+# eval_steps: 500
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+max_val_samples_per_ds: 1000
+# optim: schedule_free_adamw
+
+learning_rate: 0.00002
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 5
+weight_decay: 0.01
+# warmup_ratio: 0.1
+warmup_steps: 100
+
+dataloader_prefetch_factor: 8
+dataloader_num_workers: 8
+# LoRA
+lora_r: 8
+lora_dropout: 0.05
+target_modules:
+  - down_proj
+  - up_proj
+
+# data
+train_ds_names:
+- openmathintx-2
+- opencoder-edu
+
+val_ds_names:
+- gsm8k
+- opencoder-edu
+
+load_best_model_at_end: true
+metric_for_best_model: eval_gsm8k_loss

install.sh

@@ -1,3 +1,9 @@
 pip install -e . 
 pip install flash-attn --no-build-isolation
 # huggingface-cli login
+
+# lm-evaluation-harness
+git clone https://github.com/EleutherAI/lm-evaluation-harness.git
+cd lm-evaluation-harness
+git checkout f724be699e8adf7ca8004ea0e519dfac83a06f18
+pip install -e .

intx_sft.py

@@ -46,6 +46,7 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     AutoConfig,
+    PretrainedConfig,
     DataCollatorForSeq2Seq,
     EvalPrediction,
     HfArgumentParser,
@@ -357,6 +358,10 @@ def main():
             raise ValueError("base model contains trainable parameters")
     else:
         # activate LoRA
+        base_model_config = AutoConfig.from_pretrained(
+            model_args.model_name_or_path, trust_remote_code=True
+        )
+        base_model_config.save_pretrained(output_dir)
         logger.info("Using LoRA")
         model.set_adapter("default")
 

src/ctx_to_lora/data_utils.py

@@ -134,7 +134,19 @@ DS_KWARGS = {
     "openmathintx-2": dict(
         train=dict(
             path="nvidia/OpenMathInstruct-2",
-            split="train_1M[:500000]",
+            split="train_1M[:100000]",
+        ),
+    ),
+    "opencoder-edu": dict(
+        train=dict(
+            path="OpenCoder-LLM/opc-sft-stage2",
+            name="educational_instruct",
+            split="train[900:]",
+        ),
+        validation=dict(
+            path="OpenCoder-LLM/opc-sft-stage2",
+            name="educational_instruct",
+            split="train[:900]",
         ),
     ),
 }
@@ -280,7 +292,7 @@ def get_tokenized_dataset(
     logger.debug(f"Loading dataset {ds_name} with split {split}...")
     need_ctx_ids = not add_ctx_to_chat
     try:
-        ds = load_dataset(**get_ds_kwargs(ds_name, split))
+        ds = load_dataset(**get_ds_kwargs(ds_name, split), trust_remote_code=True)
     except ValueError as e:
         logger.info(
             f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..."
@@ -564,6 +576,15 @@ def get_preprocessing_fn(ds_name: str) -> Callable[[dict[str, Any]], dict[str, A
                 "response": sample["answer"],
             }
 
+    elif "opencoder-edu" in ds_name:
+
+        def f(sample):
+            return {
+                "context": sample["instruction"],
+                "prompt": sample["instruction"],
+                "response": "```python\n" + sample["code"].strip() + "\n```",
+            }
+
     return f
 
 

src/ctx_to_lora/eval.py

@@ -464,10 +464,6 @@ if __name__ == "__main__":
     args.output_dir = f"{run_dir}/eval-results-{cur_it}"
     args.logging_dir = f"{run_dir}/eval-results-{cur_it}"
     args.run_name = run_dir.split("/")[-1]
-    if os.path.exists(args.output_dir):
-        # remove the existing output dir
-        print(f"Removing existing output dir: {args.output_dir}")
-        shutil.rmtree(args.output_dir)
 
     evaluate(checkpoint_path, args, split="validation", generative=False)
     evaluate(checkpoint_path, args, split="validation", generative=True)

src/ctx_to_lora/modeling_utils.py

@@ -1,4 +1,5 @@
 import logging
+from copy import deepcopy
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import Enum
@@ -31,6 +32,7 @@ from transformers import (
     PerceiverModel,
     PreTrainedModel,
     PretrainedConfig,
+    PreTrainedTokenizerBase,
 )
 from transformers.models.perceiver.modeling_perceiver import (
     PerceiverBasicDecoder,
@@ -326,10 +328,10 @@ def maybe_add_batch_dim(kwargs):
     try:
         batched_input = False
         batched_attn_mask = False
-        if len(kwargs["input_ids"].shape) == 1:
+        if "input_ids" in kwargs and len(kwargs["input_ids"].shape) == 1:
             kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0)
             batched_input = True
-        if len(kwargs["attention_mask"].shape) == 1:
+        if "attention_mask" in kwargs and len(kwargs["attention_mask"].shape) == 1:
             kwargs["attention_mask"] = kwargs["attention_mask"].unsqueeze(0)
             batched_attn_mask = True
         yield batched_input, batched_attn_mask
@@ -350,6 +352,10 @@ class EarlyExit(nn.Module):
             self.base_model.layers = base_model.layers[:exit_layer]
         # self.exit_layer = exit_layer
 
+    @property
+    def config(self):
+        return self.base_model.config
+
     @torch.no_grad()
     def forward(self, **kwargs):
         # if len(kwargs["input_ids"].shape) == 1:
@@ -568,12 +574,14 @@ class ModulatedPretrainedModel(nn.Module):
         ctx_encoder_args: CtxEncoderArguments,
         # use_kl_loss is only used for training
         use_kl_loss: bool = False,
+        use_base_input_as_ctx: bool = False,
     ):
         super().__init__()
         self.device = base_model.device
         self.hypernet_config = hypernet_config
         self.ctx_encoder_args = ctx_encoder_args
         self.use_kl_loss = use_kl_loss
+        self.use_base_input_as_ctx = use_base_input_as_ctx
 
         self.register_module("base_model", base_model)
         self._init_model()
@@ -776,10 +784,12 @@ class ModulatedPretrainedModel(nn.Module):
         self,
         ctx_ids: Integer[Tensor, "bs ctx_len"],
         ctx_attn_mask: Integer[Tensor, "bs ctx_len"],
+        *args: Any,
+        **kwargs: Any,
     ):
         with torch.no_grad():
             ctx_features = self.ctx_encoder(
-                input_ids=ctx_ids, attention_mask=ctx_attn_mask
+                input_ids=ctx_ids, attention_mask=ctx_attn_mask, *args, **kwargs
             )
         return self.hypernet.generate_loras(ctx_features, ctx_attn_mask)
 
@@ -788,15 +798,16 @@ class ModulatedPretrainedModel(nn.Module):
         # ctx_features: Optional[Float[Tensor, "bs ctx_length feature_dim"]] = None,
         ctx_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None,
         ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_len"]] = None,
-        chat_ids: Optional[Integer[Tensor, "bs chat_len"]] = None,
-        chat_attn_mask: Optional[Integer[Tensor, "bs chat_len"]] = None,
-        chat_labels: Optional[Integer[Tensor, "bs chat_len"]] = None,
+        # chat_ids: Optional[Integer[Tensor, "bs chat_len"]] = None,
+        # chat_attn_mask: Optional[Integer[Tensor, "bs chat_len"]] = None,
+        # chat_labels: Optional[Integer[Tensor, "bs chat_len"]] = None,
+        *model_inputs_args: Any,
         **model_inputs_kwargs: dict[str, Any],
     ) -> Union[tuple, ModelOutput]:
         """Forward pass of the modulated model."""
 
         generated_loras = None
-        if ctx_ids is None:
+        if ctx_ids is None and not self.use_base_input_as_ctx:
             logger.warning(
                 (
                     "*" * 100,
@@ -809,6 +820,17 @@ class ModulatedPretrainedModel(nn.Module):
             # return model_outputs
 
         else:
+            if self.use_base_input_as_ctx:
+                ctx_ids = (
+                    model_inputs_kwargs["input_ids"]
+                    if "input_ids" in model_inputs_kwargs
+                    else model_inputs_args[0]
+                )
+                ctx_attn_mask = (
+                    model_inputs_kwargs["attention_mask"]
+                    if "attention_mask" in model_inputs_kwargs
+                    else None
+                )
             with torch.no_grad():
                 ctx_features = self.ctx_encoder(
                     input_ids=ctx_ids, attention_mask=ctx_attn_mask
@@ -821,41 +843,44 @@ class ModulatedPretrainedModel(nn.Module):
         # - select only the position of the label tokens (but we're actually not using the labels themselves)
         # - compute kl loss between the two logits
         if self.use_kl_loss and "labels" in model_inputs_kwargs:
-            bs = chat_ids.shape[0]
-            labels = model_inputs_kwargs.pop("labels")  # [bs, prompt_res_len]
-            first_label_pos = torch.argmax((labels != -100).float(), dim=-1)
-            # just a place holder, we dont use labels as groundtruth anyway
-            labels[torch.arange(bs), first_label_pos - 1] = 1
-            # same for chat_labels
-            first_label_pos = torch.argmax((chat_labels != -100).float(), dim=-1)
-            chat_labels[torch.arange(bs), first_label_pos - 1] = 1
+            raise NotImplementedError("Not implemented")
+            # bs = chat_ids.shape[0]
+            # labels = model_inputs_kwargs.pop("labels")  # [bs, prompt_res_len]
+            # first_label_pos = torch.argmax((labels != -100).float(), dim=-1)
+            # # just a place holder, we dont use labels as groundtruth anyway
+            # labels[torch.arange(bs), first_label_pos - 1] = 1
+            # # same for chat_labels
+            # first_label_pos = torch.argmax((chat_labels != -100).float(), dim=-1)
+            # chat_labels[torch.arange(bs), first_label_pos - 1] = 1
 
-            with torch.no_grad():
-                chat_outputs = self.base_model(
-                    input_ids=chat_ids, attention_mask=chat_attn_mask
-                )
+            # with torch.no_grad():
+            #     chat_outputs = self.base_model(
+            #         input_ids=chat_ids, attention_mask=chat_attn_mask
+            #     )
 
-            with apply_generated_loras(
-                self.base_model,
-                generated_loras,
-                self.hypernet.layer_indices,
-                self.training,
-            ):
-                model_outputs = self.base_model(**model_inputs_kwargs)
-                pred_logits = model_outputs.logits  # [bs, prompt_res_len, vocab_size]
+            # with apply_generated_loras(
+            #     self.base_model,
+            #     generated_loras,
+            #     self.hypernet.layer_indices,
+            #     self.training,
+            # ):
+            #     model_outputs = self.base_model(
+            #         *model_inputs_args, **model_inputs_kwargs
+            #     )
+            #     pred_logits = model_outputs.logits  # [bs, prompt_res_len, vocab_size]
 
-            # [bs, ctx_prompt_res_len, vocab_size]
-            base_chat_logits = chat_outputs.logits
-            base_chat_logits = base_chat_logits[torch.where(chat_labels != -100)]
-            pred_logits = pred_logits[torch.where(labels != -100)]
-            kl_loss = F.kl_div(
-                F.log_softmax(pred_logits, dim=-1),
-                F.softmax(base_chat_logits, dim=-1),
-                reduction="none",
-            )
-            # only compute kl loss on tokens where labels != -100
-            kl_loss = kl_loss.sum(dim=-1).mean()
-            model_outputs = ModelOutput(loss=kl_loss, **model_outputs)
+            # # [bs, ctx_prompt_res_len, vocab_size]
+            # base_chat_logits = chat_outputs.logits
+            # base_chat_logits = base_chat_logits[torch.where(chat_labels != -100)]
+            # pred_logits = pred_logits[torch.where(labels != -100)]
+            # kl_loss = F.kl_div(
+            #     F.log_softmax(pred_logits, dim=-1),
+            #     F.softmax(base_chat_logits, dim=-1),
+            #     reduction="none",
+            # )
+            # # only compute kl loss on tokens where labels != -100
+            # kl_loss = kl_loss.sum(dim=-1).mean()
+            # model_outputs = ModelOutput(loss=kl_loss, **model_outputs)
         else:
             # input_ids in model_inputs_kwargs contains only
             # prompt + response (for hypernet training)
@@ -865,7 +890,9 @@ class ModulatedPretrainedModel(nn.Module):
                 self.hypernet.layer_indices,
                 self.training,
             ):
-                model_outputs = self.base_model(**model_inputs_kwargs)
+                model_outputs = self.base_model(
+                    *model_inputs_args, **model_inputs_kwargs
+                )
 
         return model_outputs
 
@@ -874,10 +901,11 @@ class ModulatedPretrainedModel(nn.Module):
         self,
         ctx_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None,
         ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_length"]] = None,
+        *model_inputs_args: Any,
         **model_inputs_kwargs: dict[str, Any],
     ):
         generated_loras = None
-        if ctx_ids is None:
+        if ctx_ids is None and not self.use_base_input_as_ctx:
             logger.warning(
                 (
                     "*" * 100,
@@ -889,6 +917,17 @@ class ModulatedPretrainedModel(nn.Module):
             # model_outputs.generated_loras = None
             # return model_outputs
         else:
+            if self.use_base_input_as_ctx:
+                ctx_ids = (
+                    model_inputs_kwargs["input_ids"]
+                    if "input_ids" in model_inputs_kwargs
+                    else model_inputs_args[0]
+                )
+                ctx_attn_mask = (
+                    model_inputs_kwargs["attention_mask"]
+                    if "attention_mask" in model_inputs_kwargs
+                    else None
+                )
             ctx_features = self.ctx_encoder(
                 input_ids=ctx_ids, attention_mask=ctx_attn_mask
             )
@@ -902,10 +941,99 @@ class ModulatedPretrainedModel(nn.Module):
             self.hypernet.layer_indices,
             self.training,
         ):
-            model_outputs = self.base_model.generate(**model_inputs_kwargs)
+            model_outputs = self.base_model.generate(
+                *model_inputs_args, **model_inputs_kwargs
+            )
         return model_outputs
 
 
+class ModulatedModelWithSharedInput(nn.Module):
+    def __init__(
+        self,
+        modulated_model: ModulatedPretrainedModel,
+        base_tokenizer: PreTrainedTokenizerBase,
+        ctx_tokenizer: PreTrainedTokenizerBase,
+        ctx_end_predicate: Optional[str] = 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
+
+    @property
+    def device(self):
+        return self.modulated_model.device
+
+    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 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)
+
+
 @contextmanager
 def apply_generated_loras(
     base_model: nn.Module,

src/ctx_to_lora/run_lm_eval.py

@@ -0,0 +1,188 @@
+import sys
+import os
+import json
+import torch
+import lm_eval
+from lm_eval.models.huggingface import HFLM
+from lm_eval.models.vllm_causallms import VLLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig
+import argparse
+
+from ctx_to_lora.modeling_utils import (
+    ModulatedModelWithSharedInput,
+    ModulatedPretrainedModel,
+)
+
+
+def safe_serialize(obj):
+    default = lambda o: f"<<non-serializable: {type(o).__qualname__}>>"
+    return json.dumps(obj, default=default)
+
+
+def get_model_max_length(model_name_or_path):
+    if "Llama" in model_name_or_path:
+        return 131072
+    elif "gemma" in model_name_or_path:
+        return 2**13
+    else:
+        raise NotImplementedError(f"Unknown model: {model_name_or_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("inp", type=str, help="Path to the model or checkpoint")
+    parser.add_argument("tasks", nargs="+", help="Tasks to evaluate on")
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        help="Limit the number of examples to evaluate on",
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=16, help="Batch size for evaluation"
+    )
+    parser.add_argument(
+        "--ctx_end_predicate",
+        type=str,
+        default=None,
+        help="Split context predicate",
+    )
+    parser.add_argument(
+        "--remove_ctx_from_base_input",
+        action="store_true",
+        help="Remove context from base input",
+    )
+    args = parser.parse_args()
+
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+    os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
+    os.environ["WANDB_MODE"] = "disabled"
+    torch.use_deterministic_algorithms(True, warn_only=True)
+    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+    inp = args.inp
+    tasks = args.tasks
+    limit = args.limit
+    batch_size = args.batch_size
+    print(
+        f"Evaluating {inp} with tasks {tasks} and limit {limit} and batch size {batch_size}"
+    )
+    if "pytorch_model.bin" in inp:
+        checkpoint_path = inp
+        state_dict = torch.load(checkpoint_path)
+        print(f"Evaluating {checkpoint_path}")
+
+        if "checkpoint" in checkpoint_path:
+            cur_it = int(checkpoint_path.split("checkpoint-")[1].split("/")[0])
+            run_dir = "/".join(checkpoint_path.split("/")[:-2])
+            out_dir = f"{run_dir}/eval-results-{cur_it}"
+        else:
+            run_dir = "/".join(checkpoint_path.split("/")[:-1])
+            out_dir = f"{run_dir}/eval-results"
+
+        model = ModulatedPretrainedModel.from_state_dict(state_dict, train=False)
+        tokenizer = AutoTokenizer.from_pretrained(model.base_model.config.name_or_path)
+        ctx_model_name = model.ctx_encoder_args.ctx_encoder_model_name_or_path
+        if not ctx_model_name:
+            ctx_model_name = model.base_model.config.name_or_path
+        ctx_tokenizer = AutoTokenizer.from_pretrained(ctx_model_name)
+
+        model = ModulatedModelWithSharedInput(
+            model,
+            tokenizer,
+            ctx_tokenizer,
+            ctx_end_predicate=args.ctx_end_predicate,
+            remove_ctx_from_base_input=args.remove_ctx_from_base_input,
+        )
+
+        lm_obj = HFLM(
+            model,
+            device="cuda",
+            max_length=get_model_max_length(ctx_model_name),
+            tokenizer=tokenizer,
+            batch_size=batch_size,
+        )
+
+    elif "adapter_model.bin" in inp:
+        adapter_dir = "/".join(inp.split("/")[:-1])
+        peft_config = PeftConfig.from_pretrained(adapter_dir)
+        if "checkpoint" in inp:
+            cur_it = int(inp.split("checkpoint-")[1].split("/")[0])
+            run_dir = "/".join(inp.split("/")[:-2])
+            out_dir = f"{run_dir}/eval-results-{cur_it}"
+        else:
+            run_dir = "/".join(inp.split("/")[:-1])
+            out_dir = f"{run_dir}/eval-results"
+
+        model = AutoModelForCausalLM.from_pretrained(
+            peft_config.base_model_name_or_path, device_map="cuda"
+        )
+        model = PeftModel.from_pretrained(model, adapter_dir)
+        model.set_adapter("default")
+        model.merge_and_unload()
+        tokenizer = AutoTokenizer.from_pretrained(model.base_model.config.name_or_path)
+        # lm_obj = VLLM(
+        #     pretrained=peft_config.base_model_name_or_path,
+        #     tokenizer=peft_config.base_model_name_or_path,
+        #     lora_local_path=adapter_dir,
+        #     enable_lora=True,
+        #     max_length=get_model_max_length(peft_config.base_model_name_or_path),
+        #     trust_remote_code=True,
+        #     batch_size=batch_size,
+        # )
+
+    else:
+        model = AutoModelForCausalLM.from_pretrained(inp, device_map="cuda")
+        tokenizer = AutoTokenizer.from_pretrained(inp)
+        lm_obj = HFLM(
+            model,
+            tokenizer=tokenizer,
+            device="cuda",
+            max_length=get_model_max_length(inp),
+            batch_size=batch_size,
+        )
+        out_dir = f"eval_results/{inp}"
+        # lm_obj = VLLM(
+        #     pretrained=inp,
+        #     trust_remote_code=True,
+        #     dtype="bfloat16",
+        #     batch_size=batch_size,
+        # )
+
+    os.makedirs(out_dir, exist_ok=True)
+    # instantiate an LM subclass that takes your initialized model and can run
+    # - `Your_LM.loglikelihood()`
+    # - `Your_LM.loglikelihood_rolling()`
+    # - `Your_LM.generate_until()`
+
+    # indexes all tasks from the `lm_eval/tasks` subdirectory.
+    # Alternatively, you can set `TaskManager(include_path="path/to/my/custom/task/configs")`
+    # to include a set of tasks in a separate directory.
+    task_manager = lm_eval.tasks.TaskManager()
+
+    # Setting `task_manager` to the one above is optional and should generally be done
+    # if you want to include tasks from paths other than ones in `lm_eval/tasks`.
+    # `simple_evaluate` will instantiate its own task_manager if it is set to None here.
+    for task in tasks:
+        results = lm_eval.simple_evaluate(  # call simple_evaluate
+            model=lm_obj,
+            tasks=[task],
+            device="cuda",
+            batch_size=batch_size,
+            apply_chat_template=True,
+            limit=limit,
+            num_fewshot=0,
+            write_out=True,
+            log_samples=True,
+            task_manager=task_manager,
+        )
+
+        print(f"saving to {out_dir}/eval_results_{task}.json")
+        with open(f"{out_dir}/eval_results_{task}.json", "w") as f:
+            json.dump(results, f, indent=2, default=safe_serialize)