From 4f4aecdb00b69807a53a7936da03e4d7cbbd24e6 Mon Sep 17 00:00:00 2001 From: 51616 Date: Sat, 18 Jan 2025 12:36:12 +0000 Subject: [PATCH] add squad and gemma --- chat_templates/google/gemma-2-2b-it.jinja | 22 +++++++ configs/fw_qa_large_ctx_pwc_hotpot_squad.yaml | 58 +++++++++++++++++++ src/ctx_to_lora/data_utils.py | 2 +- src/ctx_to_lora/eval.py | 17 ++++-- src/ctx_to_lora/model_loading.py | 15 +++-- src/ctx_to_lora/modeling_utils.py | 14 +++-- 6 files changed, 111 insertions(+), 17 deletions(-) create mode 100644 chat_templates/google/gemma-2-2b-it.jinja create mode 100644 configs/fw_qa_large_ctx_pwc_hotpot_squad.yaml diff --git a/chat_templates/google/gemma-2-2b-it.jinja b/chat_templates/google/gemma-2-2b-it.jinja new file mode 100644 index 0000000..e5eaf57 --- /dev/null +++ b/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 %} + {{ '' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '\n' }} + {%- else %} + {{ '' + role + '\n' + message['content'] | trim + '\n' }} + {%- endif %} +{% endfor %} +{% if add_generation_prompt %} + {{'model\n'}} +{% endif %} diff --git a/configs/fw_qa_large_ctx_pwc_hotpot_squad.yaml b/configs/fw_qa_large_ctx_pwc_hotpot_squad.yaml new file mode 100644 index 0000000..0c2d194 --- /dev/null +++ b/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 diff --git a/src/ctx_to_lora/data_utils.py b/src/ctx_to_lora/data_utils.py index 731510c..93708c4 100644 --- a/src/ctx_to_lora/data_utils.py +++ b/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 diff --git a/src/ctx_to_lora/eval.py b/src/ctx_to_lora/eval.py index 7da43c9..cc2d2d8 100644 --- a/src/ctx_to_lora/eval.py +++ b/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) diff --git a/src/ctx_to_lora/model_loading.py b/src/ctx_to_lora/model_loading.py index 159b7ab..c421072 100644 --- a/src/ctx_to_lora/model_loading.py +++ b/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) diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 0f60b8d..da6e487 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/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