From e0e010301e8178e55189d5f2add09d1ddf671df1 Mon Sep 17 00:00:00 2001 From: 51616 Date: Fri, 20 Dec 2024 10:16:00 +0000 Subject: [PATCH] adding configs + loading + peft training --- hyperlora/configs.py | 54 ++++++++++++ hyperlora/intx_sft.py | 100 ++++++++++------------ hyperlora/model_loading.py | 168 +++++++++++++++++++++++++++++++++++++ hyperlora/utils.py | 28 +++++++ 4 files changed, 295 insertions(+), 55 deletions(-) create mode 100644 hyperlora/configs.py create mode 100644 hyperlora/model_loading.py create mode 100644 hyperlora/utils.py diff --git a/hyperlora/configs.py b/hyperlora/configs.py new file mode 100644 index 0000000..e9106b4 --- /dev/null +++ b/hyperlora/configs.py @@ -0,0 +1,54 @@ +import dataclasses +import os +import sys +from dataclasses import dataclass, field +from typing import Any, Dict, List, Literal, NewType, Optional, Tuple +from enum import Enum, auto + +from transformers import MODEL_FOR_CAUSAL_LM_MAPPING, HfArgumentParser + + +class ExperimentSetup(Enum): + LORA = "lora" + HYPER_LORA = "hyper_lora" + FULL_FINETUNE = "full_finetune" + + +@dataclass +class ModelArguments: + """ + Arguments for the base model. + """ + + model_name_or_path: Optional[str] = field( + default=None, + metadata={"help": ("Base model name or path.")}, + ) + # use_peft: bool = field( + # default=False, + # metadata={"help": ("Whether to use PEFT or not for training.")}, + # ) + + +@dataclass +class LoRAArguments: + r: Optional[int] = field( + default=8, + metadata={"help": ("LoRA R value.")}, + ) + lora_dropout: Optional[float] = field( + default=0.05, + metadata={"help": ("LoRA dropout.")}, + ) + target_modules: Optional[List[str]] = field( + default=None, + metadata={"help": ("LoRA target modules.")}, + ) + + +@dataclass +class CtxTrainingArguments: + exp_setup: ExperimentSetup = field( + default=ExperimentSetup.LORA, + metadata={"help": "Experiment setup - LoRA, HyperLoRA, or full finetuning"}, + ) diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index c49fa29..58ce7e9 100644 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -1,3 +1,4 @@ +import logging import numpy as np import torch @@ -12,6 +13,9 @@ from transformers import ( ) +from configs import CtxTrainingArguments, LoRAArguments, ModelArguments, ExperimentSetup +from utils import log_num_train_params +from model_loading import get_model_and_tokenizer, get_lora_config from modeling_utils import ModulatedPretrainedModel from data_utils import ( convert_ctx_prompt_response_to_messages, @@ -21,6 +25,8 @@ from data_utils import ( ) from training_utils import TRAINING_TASK, train_model +logger = logging.getLogger(__name__) + def compute_metrics(eval_pred: EvalPrediction) -> dict: """ @@ -41,9 +47,15 @@ def compute_metrics(eval_pred: EvalPrediction) -> dict: def main(): - parser = HfArgumentParser((TrainingArguments,)) - training_args, *_ = parser.parse_args_into_dataclasses() + # Set logging verbosity to INFO + logging.basicConfig(level=logging.INFO) + parser = HfArgumentParser( + (CtxTrainingArguments, ModelArguments, LoRAArguments, TrainingArguments) + ) + ctx_args, model_args, lora_args, training_args = parser.parse_args_into_dataclasses() + + training_args.label_names = ["labels"] training_args.eval_on_start = True training_args.eval_strategy = "steps" training_args.eval_steps = 500 @@ -60,61 +72,35 @@ def main(): "use_reentrant": False } # manually add this argument in the code - model = AutoModelForCausalLM.from_pretrained( - # "meta-llama/Llama-3.1-8B-Instruct", - "meta-llama/Llama-3.2-1B-Instruct", - torch_dtype=torch.bfloat16, - attn_implementation="flash_attention_2", + # "meta-llama/Llama-3.1-8B-Instruct", + + # model = AutoModelForCausalLM.from_pretrained( + # base_model_name, + # torch_dtype=torch.bfloat16, + # attn_implementation="flash_attention_2", + # ) + # tokenizer = AutoTokenizer.from_pretrained(base_model_name) + # tokenizer.pad_token_id = tokenizer.eos_token_id + # tokenizer.padding_side = "right" + + model_name = model_args.model_name_or_path + + model, tokenizer = get_model_and_tokenizer( + **vars(model_args), + train=True, + requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE, + peft_config=get_lora_config(model_name, **vars(lora_args)), ) - tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct") - tokenizer.pad_token_id = tokenizer.eos_token_id - tokenizer.padding_side = "right" + if ctx_args.exp_setup == ExperimentSetup.HYPER_LORA: + hypernet = ... + model = ModulatedPretrainedModel(model, hypernet) + else: + # activate LoRA + model.set_adapter("default") - max_seq_len = 1024 + log_num_train_params(model) - # if isinstance(model, ModulatedPretrainedModel): - - # # TODO: how to tokenize properly? - # def tokenize(example): - # model_inputs = tokenizer( - # example["prompt"], - # truncation=True, - # padding=False, - # ) - # model_inputs["ctx_ids"] = tokenizer(example["context"]).input_ids - # model_inputs["ctx_attention_mask"] = tokenizer( - # example["context"] - # ).attention_mask - # model_inputs["labels"] = ... - # return model_inputs - - # else: - - # def tokenize(example): - # inp = [ - # ctx + "\n" + prompt - # for ctx, prompt in zip(example["context"], example["prompt"]) - # ] - # model_inputs = tokenizer( - # inp, - # example["answer"], - # add_special_tokens=False, - # truncation=True, - # padding=False, - # max_length=max_seq_len, - # ) - - # input_ids = model_inputs["input_ids"] - # labels = [None] * len(input_ids) - - # for i in range(len(input_ids)): - # sequence_ids = model_inputs.sequence_ids(i) - # labels[i] = [ - # -100 if sequence_id == 0 else label - # for sequence_id, label in zip(sequence_ids, input_ids[i]) - # ] - # model_inputs["labels"] = labels - # return model_inputs + # max_seq_len = 1024 print("Loading dataset...") train_file = "../data/raw_datasets/context_numbers/train.jsonl" @@ -138,7 +124,7 @@ def main(): "tokenizer": tokenizer, "mask_assistant_inputs": True, "tokenizer_kwargs": { - "max_length": max_seq_len, + "max_length": None, }, }, remove_columns=ds["train"].column_names, @@ -158,8 +144,12 @@ def main(): # DataCollatorForSeq2Seq also pads the `labels` # useful when we're computing the labels manually # or masking the loss only on completion + # TODO: change to a faster collator? e.g., + # https://huggingface.co/blog/packing-with-FA2 data_collator = DataCollatorForSeq2Seq(tokenizer, model, pad_to_multiple_of=8) + # TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer + # TODO: use packing with SFTTrainer train_model( model, train_ds, diff --git a/hyperlora/model_loading.py b/hyperlora/model_loading.py new file mode 100644 index 0000000..ae50c14 --- /dev/null +++ b/hyperlora/model_loading.py @@ -0,0 +1,168 @@ +import logging +import os + +import torch +from peft import PeftModel, LoraConfig, VeraConfig, PeftConfig +from peft import get_peft_config as _get_peft_config +from peft.utils import PeftType +from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel + +from hyper_llm_modulator.utils.pooling import get_pooling_fn +from hyper_llm_modulator.utils.preprocessing import add_full_stop, apply_sfr_template + +logger = logging.getLogger() + + +def get_model_and_tokenizer( + model_name_or_path, + train, + requires_grad, + use_flash_attn=True, + peft_config=None, + model_kwargs=None, + tokenizer_kwargs=None, + device="cuda:0", + dtype=torch.bfloat16, +): + model = get_model( + model_name_or_path, + train, + requires_grad, + use_flash_attn, + peft_config, + model_kwargs, + device, + dtype, + ) + tokenizer = get_tokenizer(model_name_or_path, tokenizer_kwargs, peft_config, train) + return model, tokenizer + + +def get_tokenizer( + model_name_or_path, tokenizer_kwargs=None, peft_config=None, train=False +): + # tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side="left") + # NOTE: lora models don't have tokenizer config in the folder + + # left pad for generation, right pad for training (why?) + padding_side = "left" if not train else "right" + truncation_side = "left" + + # tokenizer = AutoTokenizer.from_pretrained( + # "models/Mistral-7B-v0.1/", padding_side=padding_side + # ) + if peft_config: + model_name_or_path = peft_config.base_model_name_or_path + + if tokenizer_kwargs is None: + tokenizer_kwargs = {} + + tokenizer = AutoTokenizer.from_pretrained( + model_name_or_path, + padding_side=padding_side, + truncation_side=truncation_side, + **tokenizer_kwargs, + ) + + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + # tokenizer.pad_token = tokenizer.eos_token + + # tokenizer.pad_token = tokenizer.eos_token + # tokenizer.add_bos_token = True + + # assert os.path.exists(f"{model_name_or_path}/chat_template.jinja"), ( + # f"Chat template not found in {model_name_or_path}\n\n" + # "We assume a specfic form of chat template for consistency between models. Please use the templates provided", + # ) + # if os.path.exists(f"{model_name_or_path}/chat_template.jinja"): + # chat_template = open(f"{model_name_or_path}/chat_template.jinja").read() + # chat_template = chat_template.replace(" ", "").replace("\n", "") + # tokenizer.chat_template = chat_template + + # tokenizer.add_eos_token = False + # if train: + # # NOTE: this correctly add an eos token that has attention = 1 + # # while padding eos tokens have attention = 0 + # tokenizer.add_eos_token = True + + # # shouldn't be needed as we manually compute the labels now + # # NOTE: this seems oddly important + # # and is not necessarily in the alignment handbook scripts + # # tokenizer.pad_token = tokenizer.unk_token # fix model not generating eos + + return tokenizer + + +def get_model( + model_name_or_path, + train, + requires_grad, + use_flash_attn=True, + peft_config=None, + model_kwargs=None, + device="cuda:0", + dtype=torch.bfloat16, +): + # TODO: we could use lower precision during finetuning + model_init_kwargs = dict( + pretrained_model_name_or_path=model_name_or_path, + device_map=device, + torch_dtype=dtype, + trust_remote_code=True, + # attn_implementation="flash_attention_2", + # load_in_4bit=True, + # load_in_8bit=True, + ) + if model_kwargs is not None: + model_init_kwargs.update(model_kwargs) + if use_flash_attn: + model_init_kwargs["attn_implementation"] = "flash_attention_2" + # for training disable cache + if train: + model_init_kwargs["use_cache"] = False + logger.debug(f"Model init kwargs: {model_init_kwargs}") + model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs) + if peft_config is not None: + model = PeftModel(model, peft_config) + model.train(train) + for param in model.parameters(): + param.requires_grad = requires_grad + return model + + +def get_lora_config(model_dir, **kwargs): + r = kwargs.get("r", 8) + peft_conf_kwargs = dict( + r=r, + peft_type=PeftType.LORA, + base_model_name_or_path=model_dir, + task_type="CAUSAL_LM", + lora_dropout=0.05, + lora_alpha=r ** (3 / 2) * 2, + ) + + peft_conf_kwargs.update(kwargs) + peft_config = _get_peft_config(peft_conf_kwargs) + return peft_config + + +# def get_emb_model_and_fns(emb_model_name, device): +# emb_model = AutoModel.from_pretrained( +# emb_model_name, +# device_map=device, +# torch_dtype=torch.float32 if "gte" in emb_model_name else torch.bfloat16, +# trust_remote_code=True, +# ).eval() +# emb_tokenizer = AutoTokenizer.from_pretrained(emb_model_name) +# if emb_tokenizer.pad_token_id is None: +# emb_tokenizer.pad_token_id = emb_tokenizer.eos_token_id +# # assert bool(args.query_intx), "query_intx must be provided for the emb_model_name" +# # partial(get_detailed_instruct, task_description=args.query_intx) +# task_desc_format_fn = add_full_stop +# if "SFR" in emb_model_name: +# task_desc_format_fn = apply_sfr_template +# pooling_fn = get_pooling_fn("last_token") +# elif "gte" in emb_model_name: +# pooling_fn = get_pooling_fn("cls") +# return emb_model, emb_tokenizer, task_desc_format_fn, pooling_fn diff --git a/hyperlora/utils.py b/hyperlora/utils.py new file mode 100644 index 0000000..a1446b3 --- /dev/null +++ b/hyperlora/utils.py @@ -0,0 +1,28 @@ +import logging + +logger = logging.getLogger(__name__) + + +def get_num_params(model): + total_params = 0 + trainable_params = 0 + for p in model.parameters(): + total_params += p.numel() + if p.requires_grad: + trainable_params += p.numel() + + return total_params, trainable_params + + +def log_num_train_params(model): + logger.debug("Trainable model parameters:") + for name, p in model.named_parameters(): + if p.requires_grad: + logger.debug(f"{name}, dtype:{p.dtype}") + + num_total_params, num_trainable_params = get_num_params(model) + logger.info( + f"trainable params: {num_trainable_params:,d} " + f"|| all params: {num_total_params:,d} " + f"|| trainable%: {100 * num_trainable_params / num_total_params:.4f}" + )