From 98d4b56629478d519559373b92312baff90716db Mon Sep 17 00:00:00 2001 From: 51616 Date: Fri, 20 Dec 2024 10:43:13 +0000 Subject: [PATCH] config parser + yaml config --- configs/default.yaml | 6 +++ hyperlora/configs.py | 109 +++++++++++++++++++++++++++++++++++++++++- hyperlora/intx_sft.py | 11 +++-- 3 files changed, 121 insertions(+), 5 deletions(-) create mode 100644 configs/default.yaml diff --git a/configs/default.yaml b/configs/default.yaml new file mode 100644 index 0000000..8f47298 --- /dev/null +++ b/configs/default.yaml @@ -0,0 +1,6 @@ +output_dir: train_outputs/ +bf16: true +target_modules: + - down_proj + - up_proj + - gate_proj diff --git a/hyperlora/configs.py b/hyperlora/configs.py index 9de62b5..40921e1 100644 --- a/hyperlora/configs.py +++ b/hyperlora/configs.py @@ -1,6 +1,7 @@ import dataclasses import os import sys +import yaml from dataclasses import dataclass, field from enum import Enum, auto from typing import Any, Dict, List, Literal, NewType, Optional, Tuple @@ -8,6 +9,112 @@ from typing import Any, Dict, List, Literal, NewType, Optional, Tuple from transformers import MODEL_FOR_CAUSAL_LM_MAPPING, HfArgumentParser +MODEL_CONFIG_CLASSES = list(MODEL_FOR_CAUSAL_LM_MAPPING.keys()) +MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES) + + +DataClassType = NewType("DataClassType", Any) + + +class ArgumentParser(HfArgumentParser): + def parse_yaml_and_args( + self, yaml_arg: str, other_args: Optional[List[str]] = None + ) -> List[dataclass]: + """ + Parse a YAML file and overwrite the default/loaded values with the values provided to the command line. + + Args: + yaml_arg (`str`): + The path to the config file used + other_args (`List[str]`, *optional`): + A list of strings to parse as command line arguments, e.g. ['--arg=val', '--arg2=val2']. + + Returns: + [`List[dataclass]`]: a list of dataclasses with the values from the YAML file and the command line + """ + arg_list = self.parse_yaml_file(os.path.abspath(yaml_arg)) + + outputs = [] + # strip other args list into dict of key-value pairs + other_args = { + arg.split("=")[0].strip("-"): arg.split("=")[1] for arg in other_args + } + used_args = {} + + # overwrite the default/loaded value with the value provided to the command line + # adapted from https://github.com/huggingface/transformers/blob/d0b5002378daabf62769159add3e7d66d3f83c3b/src/transformers/hf_argparser.py#L327 + print(arg_list) + for data_yaml, data_class in zip(arg_list, self.dataclass_types): + keys = {f.name for f in dataclasses.fields(data_yaml) if f.init} + inputs = {k: v for k, v in vars(data_yaml).items() if k in keys} + for arg, val in other_args.items(): + # add only if in keys + if arg in keys: + base_type = data_yaml.__dataclass_fields__[arg].type + inputs[arg] = val + + # cast type for ints, floats (default to strings) + if base_type in [int, float]: + inputs[arg] = base_type(val) + + if base_type == list[str]: + inputs[arg] = [str(v) for v in val.split(",")] + + # bool of a non-empty string is True, so we manually check for bools + if base_type == bool: + if val in ["true", "True"]: + inputs[arg] = True + else: + inputs[arg] = False + + if base_type == dict: + inputs[arg] = yaml.load(val, Loader=yaml.FullLoader) + + # add to used-args so we can check if double add + if arg not in used_args: + used_args[arg] = val + else: + raise ValueError( + f"Duplicate argument provided: {arg}, may cause unexpected behavior" + ) + # else: + # raise ValueError(f"Argument provided not found in dataclass: {arg}") + + obj = data_class(**inputs) + outputs.append(obj) + + return outputs + + def parse(self) -> DataClassType | Tuple[DataClassType]: + if len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"): + # If we pass only one argument to the script and it's the path to a YAML file, + # let's parse it to get our arguments. + output = self.parse_yaml_file(os.path.abspath(sys.argv[1].split("=")[-1])) + # parse command line args and yaml file + elif len(sys.argv) > 2 and sys.argv[1].endswith(".yaml"): + output = self.parse_yaml_and_args( + os.path.abspath(sys.argv[1].split("=")[-1]), sys.argv[2:] + ) + # parse --config for the yaml path and other command line args + elif any([arg.startswith("--config") for arg in sys.argv]): + yaml_arg = [ + arg + for arg in sys.argv[1:] + if arg.startswith("--config") and arg.endswith(".yaml") + ][0] + other_args = [arg for arg in sys.argv[1:] if arg != yaml_arg] + output = self.parse_yaml_and_args( + os.path.abspath(yaml_arg.split("=")[-1]), other_args + ) + # parse command line args only + else: + output = self.parse_args_into_dataclasses() + + if len(output) == 1: + output = output[0] + return output + + class ExperimentSetup(Enum): LORA = "lora" HYPER_LORA = "hyper_lora" @@ -20,7 +127,7 @@ class ModelArguments: Arguments for the base model. """ - model_name_or_path: Optional[str] = field( + model_name_or_path: str = field( default=None, metadata={"help": ("Base model name or path.")}, ) diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index 5865028..2b76300 100644 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -10,6 +10,8 @@ from data_utils import ( tokenize_chat_messages, ) from datasets import load_dataset + +from configs import ArgumentParser from model_loading import get_lora_config, get_model_and_tokenizer from modeling_utils import ModulatedPretrainedModel from training_utils import TRAINING_TASK, train_model @@ -48,10 +50,10 @@ def main(): # Set logging verbosity to INFO logging.basicConfig(level=logging.INFO) - parser = HfArgumentParser( + parser = ArgumentParser( (CtxTrainingArguments, ModelArguments, LoRAArguments, TrainingArguments) ) - ctx_args, model_args, lora_args, training_args = parser.parse_args_into_dataclasses() + ctx_args, model_args, lora_args, training_args = parser.parse() training_args.label_names = ["labels"] training_args.eval_on_start = True @@ -85,13 +87,14 @@ def main(): # activate LoRA model.set_adapter("default") + print(model) log_num_train_params(model) # max_seq_len = 1024 print("Loading dataset...") - train_file = "../data/raw_datasets/context_numbers/train.jsonl" - eval_file = "../data/raw_datasets/context_numbers/val.jsonl" + train_file = "data/raw_datasets/context_numbers/train.jsonl" + eval_file = "data/raw_datasets/context_numbers/val.jsonl" ds = load_dataset("json", data_files={"train": train_file, "eval": eval_file}) # preprocessing ds = ds.map(get_preprocessing_fn("context_numbers"))