add squad and gemma

chat_templates/google/gemma-2-2b-it.jinja

@@ -0,0 +1,22 @@
+{{- bos_token }}
+{%- if messages[0]['role'] == 'system' %}
+    {%- set system_message = messages[0]['content'] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{% for message in messages %}    
+    {% if (message['role'] == 'assistant') %}
+        {% set role = 'model' %}
+    {% else %}
+        {% set role = message['role'] %}
+    {% endif %}
+    {%- if message['role'] == 'user' and loop.first and system_message is defined %}
+        {{ '<start_of_turn>' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '<end_of_turn>\n' }}
+    {%- else %}
+        {{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}
+    {%- endif %}
+{% endfor %}
+{% if add_generation_prompt %}
+    {{'<start_of_turn>model\n'}}
+{% endif %}

configs/fw_qa_large_ctx_pwc_hotpot_squad.yaml

@@ -0,0 +1,58 @@
+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:
+- fw_qa
+- fw_qa_large
+- ctx_qa
+- pwc
+- hotpot_qa
+- squad
+
+val_ds_names:
+- fw_qa
+- fw_qa_large
+- ctx_qa
+- pwc
+- hotpot_qa
+- squad
+
+load_best_model_at_end: true
+metric_for_best_model: eval_pwc_loss

src/ctx_to_lora/data_utils.py

@@ -470,7 +470,7 @@ def get_preprocessing_fn(ds_name: str) -> Callable[[dict[str, Any]], dict[str, A
             return {
                 "context": sample["context"],
                 "prompt": sample["question"],
-                "response": sample["answer"]["text"][0],
+                "response": sample["answers"]["text"][0],
             }
 
     return f

src/ctx_to_lora/eval.py

@@ -4,6 +4,7 @@ import gc
 import json
 import os
 import random
+import shutil
 from argparse import Namespace
 from dataclasses import fields
 from functools import partial
@@ -326,11 +327,6 @@ def generation_collator(inp_list, tokenizer):
 
 def evaluate(checkpoint_path, args, split, generative):
     assert split in ["validation", "test"]
-    set_seed(42)
-    torch.use_deterministic_algorithms(True, warn_only=True)
-    torch.backends.cudnn.benchmark = False
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
     model = ModulatedPretrainedModel.from_state_dict(
         torch.load(open(checkpoint_path, "rb")),
         train=False,
@@ -386,7 +382,7 @@ def evaluate(checkpoint_path, args, split, generative):
     eval_trainer_args["eval_strategy"] = "no"
     eval_trainer_args["save_strategy"] = "no"
     eval_trainer_args["overwrite_output_dir"] = True
-    eval_trainer_args["per_device_eval_batch_size"] = 16
+    eval_trainer_args["per_device_eval_batch_size"] = 4
 
     eval_trainer_args = Seq2SeqTrainingArguments(
         **eval_trainer_args,
@@ -443,6 +439,11 @@ if __name__ == "__main__":
     os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
     os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
     os.environ["WANDB_MODE"] = "disabled"
+    set_seed(42)
+    torch.use_deterministic_algorithms(True, warn_only=True)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
     checkpoint_path = sys.argv[1]
     checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1])
     run_dir = "/".join(checkpoint_path.split("/")[:-2])
@@ -455,6 +456,10 @@ 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/model_loading.py

@@ -119,7 +119,8 @@ def get_model(
         device_map=device,
         torch_dtype=dtype,
         trust_remote_code=True,
-        # attn_implementation="flash_attention_2",
+        attn_implementation="eager",
+        use_cache=False,
         # load_in_4bit=True,
         # load_in_8bit=True,
     )
@@ -127,12 +128,14 @@ def get_model(
     if model_kwargs is not None:
         model_init_kwargs.update(model_kwargs)
     if use_flash_attn:
-        model_init_kwargs["attn_implementation"] = "flash_attention_2"
-        if is_vision_model:
-            model_init_kwargs["attn_implementation"] = "sdpa"
+        if "gemma" not in model_name_or_path:
+            model_init_kwargs["attn_implementation"] = "flash_attention_2"
+    if is_vision_model:
+        model_init_kwargs["attn_implementation"] = "sdpa"
+        model_init_kwargs.pop("use_cache")
         # for training disable cache
-    if train and not is_vision_model:
-        model_init_kwargs["use_cache"] = False
+    # if train and not is_vision_model:
+    #     model_init_kwargs["use_cache"] = False
     logger.debug(f"Model init kwargs: {model_init_kwargs}")
     if not is_vision_model:
         model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs)

src/ctx_to_lora/modeling_utils.py

@@ -37,7 +37,11 @@ from transformers.models.perceiver.modeling_perceiver import (
 )
 from transformers.modeling_outputs import ModelOutput
 
-from ctx_to_lora.configs import AggregatorArguments, HypernetArguments, CtxEncoderArguments
+from ctx_to_lora.configs import (
+    AggregatorArguments,
+    HypernetArguments,
+    CtxEncoderArguments,
+)
 from ctx_to_lora.hooks import add_generated_lora_hook, remove_hook_handles
 from ctx_to_lora.model_loading import get_lora_config, get_model, get_model_and_tokenizer
 from ctx_to_lora.pooling import POOL_FN, get_pooling_fn
@@ -215,15 +219,15 @@ class MLPResidualBlock(nn.Module):
         if pre_layer_norm:
             layers.append(nn.LayerNorm(input_size))
         layers += [
-            nn.Dropout(0.1),
+            nn.Dropout(0.05),
             nn.Linear(input_size, hidden_size),
             nn.SiLU(),
-            nn.Dropout(0.1),
+            nn.Dropout(0.05),
             nn.Linear(hidden_size, output_size),
             nn.SiLU(),
         ]
         if post_dropout:
-            layers.append(nn.Dropout(0.1))
+            layers.append(nn.Dropout(0.05))
         self.mlp = nn.Sequential(*layers)
 
     def forward(self, x):
@@ -627,6 +631,8 @@ class ModulatedPretrainedModel(nn.Module):
             ctx_model_name,
             train=True,
             requires_grad=False,
+            use_flash_attn=self.base_model.config._attn_implementation
+            == "flash_attention_2",
         )
         self.ctx_encoder = EarlyExit(
             get_base_model(encoder_model), self.ctx_encoder_args.layer_idx