From f08831d957cb6a9faaf7e6409c5b123a43dcf6ad Mon Sep 17 00:00:00 2001 From: 51616 Date: Thu, 26 Jun 2025 15:30:26 +0900 Subject: [PATCH] qa data exp --- chat_templates/google/gemma-2-9b-it.jinja | 24 ++++++++++ chat_templates/google/gemma-3-4b-it.jinja | 30 ++++++++++++ configs/fw_qa_v2_level_0.yaml | 46 +++++++++++++++++++ configs/fw_qa_v2_level_1.yaml | 46 +++++++++++++++++++ configs/fw_qa_v2_level_2.yaml | 46 +++++++++++++++++++ scripts/short_ctx/gemma_fw_qa_level_0.sh | 34 ++++++++++++++ scripts/short_ctx/gemma_fw_qa_level_0_tiny.sh | 2 +- scripts/short_ctx/gemma_fw_qa_level_1.sh | 34 ++++++++++++++ scripts/short_ctx/gemma_fw_qa_level_2.sh | 34 ++++++++++++++ scripts/short_ctx/gemma_fw_qa_level_3.sh | 4 +- scripts/short_ctx/gemma_fw_qa_level_3_tiny.sh | 2 +- scripts/short_ctx/gemma_squad.sh | 2 +- .../gemma_squad_and_fw_qa_level_3.sh | 4 +- src/ctx_to_lora/data/definitions.py | 28 +++++++++++ src/ctx_to_lora/model_loading.py | 18 ++++++-- 15 files changed, 343 insertions(+), 11 deletions(-) create mode 100644 chat_templates/google/gemma-2-9b-it.jinja create mode 100644 chat_templates/google/gemma-3-4b-it.jinja create mode 100644 configs/fw_qa_v2_level_0.yaml create mode 100644 configs/fw_qa_v2_level_1.yaml create mode 100644 configs/fw_qa_v2_level_2.yaml create mode 100644 scripts/short_ctx/gemma_fw_qa_level_0.sh create mode 100644 scripts/short_ctx/gemma_fw_qa_level_1.sh create mode 100644 scripts/short_ctx/gemma_fw_qa_level_2.sh diff --git a/chat_templates/google/gemma-2-9b-it.jinja b/chat_templates/google/gemma-2-9b-it.jinja new file mode 100644 index 0000000..23129b0 --- /dev/null +++ b/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 %} + {{ '' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '\n' }} + {%- elif message['role'] == 'assistant' %} + {{ '' + role + '\n' }}{% generation %}{{ message['content'] | trim + '\n' }}{% endgeneration %} + {%- else %} + {{ '' + role + '\n' + message['content'] | trim + '\n' }} + {%- endif %} +{% endfor %} +{% if add_generation_prompt %} + {{'model\n'}} +{% endif %} diff --git a/chat_templates/google/gemma-3-4b-it.jinja b/chat_templates/google/gemma-3-4b-it.jinja new file mode 100644 index 0000000..1713892 --- /dev/null +++ b/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') -%} + {{ '' + role + '\n' + (first_user_prefix if loop.first else "") + message['content'] | trim + '\n' }} + {%- elif (message['role'] == 'assistant') -%} + {{ '' + role + '\n' }}{% generation %}{{ message['content'] | trim + '\n' }}{% endgeneration %} + {%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{'model\n'}} +{%- endif -%} \ No newline at end of file diff --git a/configs/fw_qa_v2_level_0.yaml b/configs/fw_qa_v2_level_0.yaml new file mode 100644 index 0000000..69665cb --- /dev/null +++ b/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 diff --git a/configs/fw_qa_v2_level_1.yaml b/configs/fw_qa_v2_level_1.yaml new file mode 100644 index 0000000..ac13606 --- /dev/null +++ b/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 diff --git a/configs/fw_qa_v2_level_2.yaml b/configs/fw_qa_v2_level_2.yaml new file mode 100644 index 0000000..0d9380a --- /dev/null +++ b/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 diff --git a/scripts/short_ctx/gemma_fw_qa_level_0.sh b/scripts/short_ctx/gemma_fw_qa_level_0.sh new file mode 100644 index 0000000..767cb45 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/scripts/short_ctx/gemma_fw_qa_level_0_tiny.sh b/scripts/short_ctx/gemma_fw_qa_level_0_tiny.sh index d5ada17..2e21bba 100644 --- a/scripts/short_ctx/gemma_fw_qa_level_0_tiny.sh +++ b/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 \ diff --git a/scripts/short_ctx/gemma_fw_qa_level_1.sh b/scripts/short_ctx/gemma_fw_qa_level_1.sh new file mode 100644 index 0000000..95d8b58 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/scripts/short_ctx/gemma_fw_qa_level_2.sh b/scripts/short_ctx/gemma_fw_qa_level_2.sh new file mode 100644 index 0000000..ac4f461 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/scripts/short_ctx/gemma_fw_qa_level_3.sh b/scripts/short_ctx/gemma_fw_qa_level_3.sh index 99ccaf4..85d9ca6 100644 --- a/scripts/short_ctx/gemma_fw_qa_level_3.sh +++ b/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 \ diff --git a/scripts/short_ctx/gemma_fw_qa_level_3_tiny.sh b/scripts/short_ctx/gemma_fw_qa_level_3_tiny.sh index 640286f..5e90b30 100644 --- a/scripts/short_ctx/gemma_fw_qa_level_3_tiny.sh +++ b/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 \ diff --git a/scripts/short_ctx/gemma_squad.sh b/scripts/short_ctx/gemma_squad.sh index 647e8e7..1455e5f 100644 --- a/scripts/short_ctx/gemma_squad.sh +++ b/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 \ diff --git a/scripts/short_ctx/gemma_squad_and_fw_qa_level_3.sh b/scripts/short_ctx/gemma_squad_and_fw_qa_level_3.sh index 848d5fd..ffec3c5 100644 --- a/scripts/short_ctx/gemma_squad_and_fw_qa_level_3.sh +++ b/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 \ diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py index ac35b62..e888dbf 100644 --- a/src/ctx_to_lora/data/definitions.py +++ b/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", diff --git a/src/ctx_to_lora/model_loading.py b/src/ctx_to_lora/model_loading.py index 79bff30..456562a 100644 --- a/src/ctx_to_lora/model_loading.py +++ b/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)