qa data exp

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

@@ -0,0 +1,24 @@
+{{- 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 loop_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' }}
+    {%- elif message['role'] == 'assistant' %}
+        {{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] | trim + '<end_of_turn>\n' }}{% endgeneration %}
+    {%- 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 %}

chat_templates/google/gemma-3-4b-it.jinja

@@ -0,0 +1,30 @@
+{{ bos_token }}
+{%- if messages[0]['role'] == 'system' -%}
+    {%- if messages[0]['content'] is string -%}
+        {%- set first_user_prefix = messages[0]['content'] + '\n\n' -%}
+    {%- else -%}
+        {%- set first_user_prefix = messages[0]['content'][0]['text'] + '\n\n' -%}
+    {%- endif -%}
+    {%- set loop_messages = messages[1:] -%}
+{%- else -%}
+    {%- set first_user_prefix = "" -%}
+    {%- set loop_messages = messages -%}
+{%- endif -%}
+{%- for message in loop_messages -%}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
+        {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
+    {%- endif -%}
+    {%- if (message['role'] == 'assistant') -%}
+        {%- set role = "model" -%}
+    {%- else -%}
+        {%- set role = message['role'] -%}
+    {%- endif -%}
+    {%- if (message['role'] == 'user') -%}
+        {{ '<start_of_turn>' + role + '\n' + (first_user_prefix if loop.first else "") + message['content'] | trim + '<end_of_turn>\n' }}
+    {%- elif (message['role'] == 'assistant') -%}
+        {{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] | trim + '<end_of_turn>\n' }}{% endgeneration %}
+    {%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{'<start_of_turn>model\n'}}
+{%- endif -%}

configs/fw_qa_v2_level_0.yaml

@@ -0,0 +1,46 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: google/gemma-3-1b-it
+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.00004
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 1
+weight_decay: 0.01
+
+warmup_steps: 100
+
+dataloader_prefetch_factor: 16
+dataloader_num_workers: 8
+
+# LoRA
+lora_r: 8
+lora_dropout: 0.0
+target_modules:
+  - down_proj
+
+# data
+train_ds_names:
+  - fw_qa_v2_2k_len_level_0
+
+val_ds_names:
+  - fw_qa_v2_2k_len_level_0
+  - squad

configs/fw_qa_v2_level_1.yaml

@@ -0,0 +1,46 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: google/gemma-3-1b-it
+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.00004
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 1
+weight_decay: 0.01
+
+warmup_steps: 100
+
+dataloader_prefetch_factor: 16
+dataloader_num_workers: 8
+
+# LoRA
+lora_r: 8
+lora_dropout: 0.0
+target_modules:
+  - down_proj
+
+# data
+train_ds_names:
+  - fw_qa_v2_2k_len_level_1
+
+val_ds_names:
+  - fw_qa_v2_2k_len_level_0
+  - squad

configs/fw_qa_v2_level_2.yaml

@@ -0,0 +1,46 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: google/gemma-3-1b-it
+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.00004
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 1
+weight_decay: 0.01
+
+warmup_steps: 100
+
+dataloader_prefetch_factor: 16
+dataloader_num_workers: 8
+
+# LoRA
+lora_r: 8
+lora_dropout: 0.0
+target_modules:
+  - down_proj
+
+# data
+train_ds_names:
+  - fw_qa_v2_2k_len_level_2
+
+val_ds_names:
+  - fw_qa_v2_2k_len_level_0
+  - squad

scripts/short_ctx/gemma_fw_qa_level_0.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+#SBATCH --job-name=ctxlora
+#SBATCH --nodes=1
+#SBATCH --partition=sakura-gpu
+#SBATCH --gpus=4
+#SBATCH --output=slurm_logs/%x-%j.out
+#SBATCH --error=slurm_logs/%x-%j.out
+
+uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
+--gpu_ids all --main_process_port 29562 intx_sft.py configs/fw_qa_v2_level_0.yaml \
+--model_name_or_path=google/gemma-2-2b-it \
+--num_train_epochs=5 \
+--per_device_train_batch_size=-1 \
+--gradient_accumulation_steps=2 \
+--per_device_eval_batch_size=64 \
+--target_modules=down_proj \
+--num_self_attends_per_block=8 \
+--num_latent_factor=1 \
+--num_pre_head_layers=1 \
+--lora_r=8 \
+--eval_steps=1000 \
+--save_steps=1000 \
+--learning_rate=4e-5 \
+--lora_dropout=0.0 \
+--neftune_noise_alpha=5 \
+--add_negative_prompt=False \
+--add_repeat_prompt=False \
+--use_sequence_packing=True \
+--max_packed_inp_len=16000 \
+--max_packed_ctx_len=32000 \
+--per_rank_gen=True \
+--per_layer_processing=True \
+--gen_lora_l1_reg_coef=0.1 \
+--logging_steps=50

scripts/short_ctx/gemma_fw_qa_level_0_tiny.sh

@@ -8,7 +8,7 @@
 
 uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
 --gpu_ids all --main_process_port 29562 intx_sft.py configs/fw_qa_v2_level_0_tiny.yaml \
---model_name_or_path=google/gemma-3-1b-it \
+--model_name_or_path=google/gemma-2-2b-it \
 --num_train_epochs=5 \
 --per_device_train_batch_size=-1 \
 --gradient_accumulation_steps=2 \

scripts/short_ctx/gemma_fw_qa_level_1.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+#SBATCH --job-name=ctxlora
+#SBATCH --nodes=1
+#SBATCH --partition=sakura-gpu
+#SBATCH --gpus=4
+#SBATCH --output=slurm_logs/%x-%j.out
+#SBATCH --error=slurm_logs/%x-%j.out
+
+uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
+--gpu_ids all --main_process_port 29563 intx_sft.py configs/fw_qa_v2_level_1.yaml \
+--model_name_or_path=google/gemma-2-2b-it \
+--num_train_epochs=5 \
+--per_device_train_batch_size=-1 \
+--gradient_accumulation_steps=2 \
+--per_device_eval_batch_size=64 \
+--target_modules=down_proj \
+--num_self_attends_per_block=8 \
+--num_latent_factor=1 \
+--num_pre_head_layers=1 \
+--lora_r=8 \
+--eval_steps=1000 \
+--save_steps=1000 \
+--learning_rate=4e-5 \
+--lora_dropout=0.0 \
+--neftune_noise_alpha=5 \
+--add_negative_prompt=False \
+--add_repeat_prompt=False \
+--use_sequence_packing=True \
+--max_packed_inp_len=16000 \
+--max_packed_ctx_len=32000 \
+--per_rank_gen=True \
+--per_layer_processing=True \
+--gen_lora_l1_reg_coef=0.1 \
+--logging_steps=50

scripts/short_ctx/gemma_fw_qa_level_2.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+#SBATCH --job-name=ctxlora
+#SBATCH --nodes=1
+#SBATCH --partition=sakura-gpu
+#SBATCH --gpus=4
+#SBATCH --output=slurm_logs/%x-%j.out
+#SBATCH --error=slurm_logs/%x-%j.out
+
+uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
+--gpu_ids all --main_process_port 29564 intx_sft.py configs/fw_qa_v2_level_2.yaml \
+--model_name_or_path=google/gemma-2-2b-it \
+--num_train_epochs=5 \
+--per_device_train_batch_size=-1 \
+--gradient_accumulation_steps=2 \
+--per_device_eval_batch_size=64 \
+--target_modules=down_proj \
+--num_self_attends_per_block=8 \
+--num_latent_factor=1 \
+--num_pre_head_layers=1 \
+--lora_r=8 \
+--eval_steps=1000 \
+--save_steps=1000 \
+--learning_rate=4e-5 \
+--lora_dropout=0.0 \
+--neftune_noise_alpha=5 \
+--add_negative_prompt=False \
+--add_repeat_prompt=False \
+--use_sequence_packing=True \
+--max_packed_inp_len=16000 \
+--max_packed_ctx_len=32000 \
+--per_rank_gen=True \
+--per_layer_processing=True \
+--gen_lora_l1_reg_coef=0.1 \
+--logging_steps=50

scripts/short_ctx/gemma_fw_qa_level_3.sh

@@ -7,8 +7,8 @@
 #SBATCH --error=slurm_logs/%x-%j.out
 
 uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
---gpu_ids all --main_process_port 29564 intx_sft.py configs/fw_qa_v2_level_3.yaml \
---model_name_or_path=google/gemma-3-1b-it \
+--gpu_ids all --main_process_port 29565 intx_sft.py configs/fw_qa_v2_level_3.yaml \
+--model_name_or_path=google/gemma-2-2b-it \
 --num_train_epochs=5 \
 --per_device_train_batch_size=-1 \
 --gradient_accumulation_steps=2 \

scripts/short_ctx/gemma_fw_qa_level_3_tiny.sh

@@ -8,7 +8,7 @@
 
 uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
 --gpu_ids all --main_process_port 29563 intx_sft.py configs/fw_qa_v2_level_3_tiny.yaml \
---model_name_or_path=google/gemma-3-1b-it \
+--model_name_or_path=google/gemma-2-2b-it \
 --num_train_epochs=5 \
 --per_device_train_batch_size=-1 \
 --gradient_accumulation_steps=2 \

scripts/short_ctx/gemma_squad.sh

@@ -8,7 +8,7 @@
 
 uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
 --gpu_ids all --main_process_port 29561 intx_sft.py configs/squad.yaml \
---model_name_or_path=google/gemma-3-1b-it \
+--model_name_or_path=google/gemma-2-2b-it \
 --num_train_epochs=5 \
 --per_device_train_batch_size=-1 \
 --gradient_accumulation_steps=2 \

scripts/short_ctx/gemma_squad_and_fw_qa_level_3.sh

@@ -7,8 +7,8 @@
 #SBATCH --error=slurm_logs/%x-%j.out
 
 uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
---gpu_ids all --main_process_port 29565 intx_sft.py configs/squad.yaml \
---model_name_or_path=google/gemma-3-1b-it \
+--gpu_ids all --main_process_port 29566 intx_sft.py configs/squad_fw_qa_v2_level_3.yaml \
+--model_name_or_path=google/gemma-2-2b-it \
 --num_train_epochs=5 \
 --per_device_train_batch_size=-1 \
 --gradient_accumulation_steps=2 \

src/ctx_to_lora/data/definitions.py

@@ -150,6 +150,34 @@ DS_KWARGS = {
             split="train",
         ),
     ),
+    "fw_qa_v2_2k_len_level_1": dict(
+        train=dict(
+            path="parquet",
+            data_files="data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_1.parquet",
+            split="train",
+        ),
+        # validation=dict(
+        #     path="parquet",
+        #     data_files=glob(
+        #         "data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_1_val.parquet"
+        #     ),
+        #     split="train",
+        # ),
+    ),
+    "fw_qa_v2_2k_len_level_2": dict(
+        train=dict(
+            path="parquet",
+            data_files="data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_2.parquet",
+            split="train",
+        ),
+        # validation=dict(
+        #     path="parquet",
+        #     data_files=glob(
+        #         "data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_2_val.parquet"
+        #     ),
+        #     split="train",
+        # ),
+    ),
     "fw_qa_v2_2k_len_level_3_tiny": dict(
         train=dict(
             path="parquet",

src/ctx_to_lora/model_loading.py

@@ -9,11 +9,21 @@ from transformers import (
     AutoModel,
     AutoModelForCausalLM,
     AutoTokenizer,
-    MllamaForConditionalGeneration,
+    Gemma3ForConditionalGeneration,
 )
 
 logger = logging.getLogger()
 
+GEMMA_VISION_MODELS = [
+    "google/gemma-3-4b-it",
+    "google/gemma-3-12b-it",
+    "google/gemma-3-27b-it",
+]
+
+
+def check_is_vision_model(model_name):
+    return model_name in GEMMA_VISION_MODELS
+
 
 def get_model_and_tokenizer(
     model_name_or_path,
@@ -96,7 +106,7 @@ def get_model(
         attn_implementation="eager",
         use_cache=False,
     )
-    is_vision_model = "Llama" in model_name_or_path and "Vision" in model_name_or_path
+    is_vision_model = check_is_vision_model(model_name_or_path)
     if model_kwargs is not None:
         model_init_kwargs.update(model_kwargs)
 
@@ -112,7 +122,7 @@ def get_model(
 
     if is_vision_model:
         # always use sdpa for vision models
-        model_init_kwargs["attn_implementation"] = "sdpa"
+        # model_init_kwargs["attn_implementation"] = "sdpa"
         model_init_kwargs.pop("use_cache")
     elif is_bidir_model:
         model_init_kwargs["torch_dtype"] = torch.float32
@@ -125,7 +135,7 @@ def get_model(
         else:
             model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs)
     else:
-        model = MllamaForConditionalGeneration.from_pretrained(**model_init_kwargs)
+        model = Gemma3ForConditionalGeneration.from_pretrained(**model_init_kwargs)
         model = model.language_model
     if peft_config is not None:
         model = PeftModel(model, peft_config)