diff --git a/examples/open_llama/README.md b/examples/open_llama/README.md index acef240ded..1cb83f70c2 100644 --- a/examples/open_llama/README.md +++ b/examples/open_llama/README.md @@ -73,11 +73,21 @@ The relevant config file is [open_llama_sparsegpt_gpu.json](open_llama_sparsegpt Requirements file: [requirements-sparsegpt.txt](requirements-sparsegpt.txt) -### Fine-tune Llama Model using QLoRA +### Fine-tune Llama Model on a chatbot dataset using QLoRA This workflow fine-tunes LLaMA model using [QLoRA](https://arxiv.org/abs/2305.14314). The output model is still the input transformers model along with a quantization config and LoRA adapters that were fine-tuned on the training dataset. -The relevant configfile is [llama_qlora.json](llama_qlora.json). It corresponds to the [guqnaco 7b example in the original qlora implementation](https://github.com/artidoro/qlora/blob/main/scripts/finetune_guanaco_7b.sh). +The relevant config file is [llama_qlora.json](llama_qlora.json). It corresponds to the [guqnaco 7b example in the original qlora implementation](https://github.com/artidoro/qlora/blob/main/scripts/finetune_guanaco_7b.sh). + +Requirements file: [requirements-qlora.txt](requirements-qlora.txt) + +### Fine-tune Open Llama Model on a code generation dataset using QLoRA +This workflow fine-tunes Open LLaMA model using [QLoRA] to generate code given a prompt. + +The relevant config file is [open_llama_qlora_tinycodes.json](open_llama_qlora_tinycodes.json). The code language is set to `Python` but can be changed to other languages by changing the `language` field in the config file. +Supported languages are Python, TypeScript, JavaScript, Ruby, Julia, Rust, C++, Bash, Java, C#, and Go. Refer to the [dataset card](https://huggingface.co/datasets/nampdn-ai/tiny-codes) for more details on the dataset. + +Note: You must be logged in to HuggingFace using `huggingface-cli login` to download the dataset or update `token` field in the config file with your HuggingFace token. Requirements file: [requirements-qlora.txt](requirements-qlora.txt) diff --git a/examples/open_llama/open_llama_qlora_tinycodes.json b/examples/open_llama/open_llama_qlora_tinycodes.json new file mode 100644 index 0000000000..b41c742b7f --- /dev/null +++ b/examples/open_llama/open_llama_qlora_tinycodes.json @@ -0,0 +1,86 @@ +{ + "input_model":{ + "type": "PyTorchModel", + "config": { + "hf_config": { + "model_name": "openlm-research/open_llama_7b_v2", + "task": "text-generation" + } + } + }, + "data_configs": { + "tiny-codes-train": { + "name": "tiny-codes-train", + "type": "HuggingfaceContainer", + "user_script": "qlora_user_script.py", + "components": { + "load_dataset": { + "type": "load_tiny_code_dataset" + } + }, + "params_config": { + "dataset_name": "nampdn-ai/tiny-codes", + "split": "train", + "component_kwargs": { + "load_dataset": { + "language": "Python", + "token": true + }, + "pre_process_data": { + "dataset_type": "corpus", + "corpus_strategy": "join", + "text_template": "### Question: {prompt} \n### Answer: {response}", + "source_max_len": 1024 + } + } + } + } + }, + "passes": { + "qlora": { + "type": "QLoRA", + "config": { + "compute_dtype": "bfloat16", + "quant_type": "nf4", + "double_quant": true, + "lora_r": 64, + "lora_alpha": 16, + "lora_dropout": 0.1, + "train_data_config": "tiny-codes-train", + "eval_dataset_size": 1024, + "training_args": { + "seed": 0, + "data_seed": 42, + "per_device_train_batch_size": 16, + "per_device_eval_batch_size": 16, + "gradient_accumulation_steps": 1, + "gradient_checkpointing": true, + "learning_rate": 0.0002, + "max_steps": 1500, + "logging_steps": 10, + "evaluation_strategy": "steps", + "eval_steps": 100, + "save_steps": 500, + "group_by_length": true, + "adam_beta2": 0.999, + "max_grad_norm": 0.3, + "load_best_model_at_end": true + } + } + } + }, + "engine": { + "log_severity_level": 0, + "search_strategy": false, + "evaluate_input_model": false, + "target": { + "type": "LocalSystem", + "config": { + "accelerators": ["gpu"] + } + }, + "execution_providers": ["CPUExecutionProvider"], + "cache_dir": "cache", + "output_dir" : "models/qlora-tiny-codes" + } +} diff --git a/examples/open_llama/open_llama_sparsegpt_gpu.json b/examples/open_llama/open_llama_sparsegpt_gpu.json index 5b8ed4ac11..68fd972c73 100644 --- a/examples/open_llama/open_llama_sparsegpt_gpu.json +++ b/examples/open_llama/open_llama_sparsegpt_gpu.json @@ -22,6 +22,7 @@ "dataset_type": "corpus", "text_cols": ["text"], "corpus_strategy": "join-random", + "add_special_tokens": false, "source_max_len": 2048, "max_samples": 128, "random_seed": 42 @@ -41,6 +42,7 @@ "dataset_type": "corpus", "text_cols": ["text"], "corpus_strategy": "join", + "add_special_tokens": false, "source_max_len": 2048 } } diff --git a/examples/open_llama/qlora_user_script.py b/examples/open_llama/qlora_user_script.py new file mode 100644 index 0000000000..25704ecdc4 --- /dev/null +++ b/examples/open_llama/qlora_user_script.py @@ -0,0 +1,17 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from typing import Union + +from datasets import load_dataset + +from olive.data.registry import Registry + + +# TODO: remove custom dataset component once default dataset component supports filter, tokens and split +@Registry.register_dataset() +def load_tiny_code_dataset(dataset_name: str, split: str, language: str, token: Union[bool, str] = True): + dataset = load_dataset(dataset_name, split=split, token=token) + dataset = dataset.filter(lambda x: x["programming_language"] == language) + return dataset diff --git a/olive/data/component/text_generation.py b/olive/data/component/text_generation.py index 508f423608..42050672a4 100644 --- a/olive/data/component/text_generation.py +++ b/olive/data/component/text_generation.py @@ -5,6 +5,7 @@ from enum import Enum from random import Random +from typing import List, Union import torch from pydantic import validator @@ -65,7 +66,7 @@ class TextGenParams(ConfigBase): # if false, cannot gaurantee all sequences are same length. data loader will have to handle this during collation pad_to_max_len: bool = True # pad sequences to max_len, ignored for JOIN corpus strategy drop_short_sequences: bool = False # drop sequences shorter than max_len. Mutually exclusive with pad_to_max_len - add_special_tokens: bool = True # add bos and eos tokens + add_special_tokens: bool = True # add bos and eos tokens to each sequence @validator("drop_short_sequences", always=True) def _check_padding(cls, v, values): @@ -79,15 +80,37 @@ def _check_padding(cls, v, values): class TextGenCorpusParams(TextGenParams): """Parameters for text generation task with 'corpus' dataset type.""" - text_cols: list # list of text columns + # TODO: Add support for formatting function: formatting_func > text_template > text_cols + # TODO: formatting function support requires the data container to provide user module from the parent data config + # one of text_template or text_cols must be provided + # a python f-string template for the text with {column_name} as placeholders + text_template: str = None + # list of text columns, columns are concatenated together using a space + text_cols: Union[str, List[str]] = None + # in JOIN strategies, the rows of text_cols are concatenated together corpus_strategy: TextGenCorpusStrategy = TextGenCorpusStrategy.LINE_BY_LINE stride: int = None # required when corpus_strategy is JOIN_SLIDING_WINDOW - joiner: str = " " # delimiter to use when joining the rows of the input columns. + # text to join the rows of input columns when corpus_strategy is JOIN + # add_special_tokens: "{bos_token} {text_col1} {eos_token} {joiner} {bos_token} {text_col2} {eos_token}..." + # no add_special_tokens: "{text_col1} {joiner} {text_col2}..." + # if None, joined with a space + joiner: str = None + processing_batch_size: int = 1024 # number of examples to process at a time random_seed: int = None # random seed for LINE_BY_LINE_RANDOM and JOIN_RANDOM random_retries: int = ( 10 # number of resamples to try before giving up when a sample is too short for RANDOM strategies ) + @validator("text_cols", always=True) + def _check_text_cols(cls, v, values): + if "text_template" not in values: + raise ValueError("Invalid text_template") + if not (values["text_template"] or v): + raise ValueError("One of text_template or text_cols must be specified") + if v is not None and isinstance(v, str): + v = [v] + return v + @validator("stride", always=True) def _check_stride(cls, v, values): if "corpus_strategy" not in values: @@ -113,12 +136,19 @@ def _check_random(cls, v, values): return v +# TODO: absorb pair format into corpus format and drop dataset_type +# This is because pair format is just a special case of corpus format class TextGenPairParams(TextGenParams): """Parameters for text generation task with 'pair' dataset type.""" pair_format: TextGenPairFormat = TextGenPairFormat.DEFAULT - input_col: str = None # required when pair_format is CUSTOM - output_col: str = None # required when pair_format is CUSTOM + # TODO: Add support for formatting functions: formatting_func > template > col + # for custom pair_format, one of input_template or input_col must be provided + input_template: str = None # a python f-string template for the input with {column_name} as placeholders + input_col: str = None # column name for input + # for custom pair_format, one of output_template or output_col must be provided + output_template: str = None # a python f-string template for the output with {column_name} as placeholders + output_col: str = None # column name for output target_max_len: int # max length of target sequence ignore_source_in_labels: bool = True # set source tokens to ignore_index in labels @@ -126,12 +156,15 @@ class TextGenPairParams(TextGenParams): def _check_custom(cls, v, field, values): if "pair_format" not in values: raise ValueError("Invalid pair_format") - if values["pair_format"] == TextGenPairFormat.CUSTOM and v is None: - raise ValueError(f"{field.name} must be specified when pair_format is CUSTOM") + template_name = f"{field.name.split('_')[0]}_template" + if template_name not in values: + raise ValueError(f"Invalid {template_name}") + if values["pair_format"] == TextGenPairFormat.CUSTOM and not (v or values[template_name]): + raise ValueError(f"{field.name} or {template_name} must be specified when pair_format is CUSTOM") return v -def text_gen_corpus_pre_process(_dataset, tokenizer, all_kwargs): +def text_gen_corpus_pre_process(dataset, tokenizer, all_kwargs): """ Pre-process data for text generation task with 'corpus' dataset type. @@ -143,10 +176,18 @@ def text_gen_corpus_pre_process(_dataset, tokenizer, all_kwargs): args = validate_config(all_kwargs, TextGenCorpusParams, warn_unused_keys=True) - # gather text from all input columns - text_list = [] - for input_col in args.text_cols: - text_list += _dataset[input_col] + # template for joining text columns + text_template = args.text_template + if text_template is None: + text_template = " ".join(["{" + col + "}" for col in args.text_cols]) + if args.add_special_tokens: + # add bos and eos tokens before tokenizing + # some tokenizers like LlamaTokenizer do not add eos token + text_template = f"{tokenizer.bos_token} {text_template} {tokenizer.eos_token}" + # apply text_template + dataset = apply_template(dataset, "text", text_template, remove_cols=True) + text_list = dataset["text"] + total_examples = len(text_list) # total number of examples tokenized_inputs = { "input_ids": [], @@ -154,79 +195,93 @@ def text_gen_corpus_pre_process(_dataset, tokenizer, all_kwargs): "labels": [], } if "join" in args.corpus_strategy: - # delimiter between the text sequences - text = args.joiner.join(text_list) - - # in order to make processing faster we will only tokenize as much as needed - # assumes that num words > num tokens - split_text = text.split(" ") - num_text = len(split_text) - - seqlen = args.source_max_len + joiner_tokens = tokenizer.encode(args.joiner, add_special_tokens=False) if args.joiner else [] if args.corpus_strategy != TextGenCorpusStrategy.JOIN_RANDOM: # no randomization, just use contiguous blocks of tokens if args.corpus_strategy == TextGenCorpusStrategy.JOIN_SLIDING_WINDOW: # we use the stride as both the step between sequences and the context size - step, context = args.stride, seqlen - args.stride + step, context = args.stride, args.source_max_len - args.stride else: # JOIN corpus_strategy # text is split into non-overlapping sequences and there is no context - step, context = seqlen, None - - # only take as much text as needed - # assumes that num words > num tokens, so we can use num tokens as an upper bound - max_text = args.max_samples * seqlen if args.max_samples is not None else num_text - max_text = min(max_text, num_text) - # tokenize the text - encodings = tokenizer(" ".join(split_text[:max_text]), add_special_tokens=False, return_tensors="pt") - - num_tokens = encodings.input_ids.shape[1] - # loop over the number of tokens - # all inputs must be seqlen long - for begin_loc in range(0, num_tokens - seqlen, step): - # end_loc is the beginning of the next sequence - end_loc = begin_loc + seqlen - # get the input sequence - input_ids = encodings.input_ids[0, begin_loc:end_loc].clone() - append_text_gen_input_ids(tokenized_inputs, input_ids, tokenizer, context=context) + step, context = args.source_max_len, None + + example_idx = 0 # index of the first example in the current batch + num_samples = 0 # samples processed so far + overflow = [] # tokens overflowed from the previous batch of examples + # we will process in batches to make tokenization faster + # better than joining all text together and tokenizing all at once + while True: + if args.max_samples is not None and num_samples >= args.max_samples: + # we have reached max_samples + break + if example_idx >= total_examples: + # we have reached the end of the text_list + break + + examples_to_get = min(args.processing_batch_size, total_examples - example_idx) + # batch tokenize + batched_input_ids = tokenizer( + text_list[example_idx : example_idx + examples_to_get], # noqa E203 + add_special_tokens=False, + truncation=False, + )["input_ids"] + + # join all the input_ids together with joiner_tokens + joined_input_ids = overflow + for input_ids in batched_input_ids: + joined_input_ids += input_ids + joiner_tokens + + end_loc = 0 # position of unused token in joined_input_ids + # '- args.source_max_len' is used to make sure we don't get a sequence that is too short + for begin_loc in range(0, len(joined_input_ids) - args.source_max_len, step): + # end_loc is the beginning of the next sequence + end_loc = begin_loc + args.source_max_len + # get the input sequence + input_ids = torch.tensor(joined_input_ids[begin_loc:end_loc]) + append_text_gen_input_ids(tokenized_inputs, input_ids, tokenizer, context=context) + num_samples += 1 + if args.max_samples is not None and num_samples >= args.max_samples: + # we have reached max_samples + break + # update counters + example_idx += examples_to_get + overflow = joined_input_ids[end_loc:] else: # randomization, sample random blocks of tokens rng = Random(args.random_seed) + # cache to store tokenized examples cache = {} for _ in range(args.max_samples): resamples = 0 - encodings = None + # will try to sample sequences random_retries times before giving up while resamples < args.random_retries: - # sample a random block of tokens by sampling a random starting location + # sample a beginning example # randint is inclusive, so we need to subtract 1 - begin_loc = rng.randint(0, num_text - seqlen - 1) - # heuristic to make sure we don't get a sequence that is too short - if begin_loc not in cache: - encodings = tokenizer( - " ".join(split_text[begin_loc : begin_loc + seqlen]), # noqa E203 - add_special_tokens=False, - return_tensors="pt", - ) - cache[begin_loc] = encodings - else: - encodings = cache[begin_loc] - if encodings.input_ids.shape[1] >= seqlen: - # found a good sample + begin_example_idx = rng.randint(0, total_examples - 1) + joined_input_ids = [] + # loop through the examples until we have enough tokens + for i in range(begin_example_idx, total_examples): + # get the input_ids + if i not in cache: + cache[i] = tokenizer( + text_list[i], + add_special_tokens=False, + truncation=False, + )["input_ids"] + joined_input_ids += cache[i] + joiner_tokens + # stop if we have enough tokens + if len(joined_input_ids) >= args.source_max_len: + break + # add to samples if we have enough tokens + if len(joined_input_ids) >= args.source_max_len: + # found a good example + input_ids = torch.tensor(joined_input_ids[: args.source_max_len]) + append_text_gen_input_ids(tokenized_inputs, input_ids, tokenizer) break resamples += 1 - if not encodings: - # could not find a good sample after resampling - continue - input_ids = encodings.input_ids[0, :seqlen] - append_text_gen_input_ids(tokenized_inputs, input_ids, tokenizer) - else: - if args.add_special_tokens: - # add bos and eos tokens before tokenizing - # some tokenizers like LlamaTokenizer do not add eos token - text_list = [f"{tokenizer.bos_token} {text} {tokenizer.eos_token}" for text in text_list] - # each line is a sequence if args.corpus_strategy == TextGenCorpusStrategy.LINE_BY_LINE: # batched tokenization might be faster so lets tokenize all the text at once @@ -236,25 +291,24 @@ def text_gen_corpus_pre_process(_dataset, tokenizer, all_kwargs): input_ids = torch.tensor(input_ids) append_text_gen_input_ids(tokenized_inputs, input_ids, tokenizer) else: - total_samples = len(text_list) + example_idx = 0 # index of the first example in the current batch num_samples = 0 - begin_loc = 0 while True: - if num_samples >= args.max_samples or begin_loc >= total_samples: + if num_samples >= args.max_samples or example_idx >= total_examples: # we have reached max_samples or the end of the text_list break - # get as many samples as possible without going over max_samples - samples_to_get = min(args.max_samples - num_samples, total_samples - begin_loc) + # get as many examples as possible without going over max_samples + examples_to_get = min(args.max_samples - num_samples, total_examples - example_idx) # batch tokenize batched_input_ids = batch_tokenize_text( - text_list[begin_loc : begin_loc + samples_to_get], tokenizer, args # noqa E203 + text_list[example_idx : example_idx + examples_to_get], tokenizer, args # noqa E203 ) for input_ids in batched_input_ids: input_ids = torch.tensor(input_ids) append_text_gen_input_ids(tokenized_inputs, input_ids, tokenizer) # update counters num_samples += len(batched_input_ids) - begin_loc += samples_to_get + example_idx += examples_to_get else: # randomization, sample random lines rng = Random(args.random_seed) @@ -311,27 +365,28 @@ def batch_tokenize_text(text_list, tokenizer, args): # based on https://github.com/artidoro/qlora/blob/main/qlora.py -def text_gen_pair_pre_process(_dataset, tokenizer, all_kwargs): +def text_gen_pair_pre_process(dataset, tokenizer, all_kwargs): """ Pre-process data for text generation task with 'pair' dataset type. Dataset is expected to have two text columns: input and output. An example is a dataset with pairs of prompts and completions. - The input (truncate to source_max_len) and output (truncate to target_max_len) are concatenated together. + The input (truncated to source_max_len) and output (truncated to target_max_len) are concatenated together. """ from datasets import Dataset as HFDataset args = validate_config(all_kwargs, TextGenPairParams, warn_unused_keys=True) - # format dataset based on pair_format - # the formatted dataset has two columns: input and output - dataset = format_pair_dataset(_dataset, args.pair_format, args.input_col, args.output_col) if args.max_samples is not None: # truncate dataset to max_samples # makes tokenization faster dataset = dataset.select(range(args.max_samples)) + # format dataset based on pair_format + # the formatted dataset has two columns: input and output + dataset = format_pair_dataset(dataset, args) + # extract elements sources = dataset["input"] targets = dataset["output"] @@ -361,6 +416,8 @@ def text_gen_pair_pre_process(_dataset, tokenizer, all_kwargs): # skip short sequences if drop_short_sequences is True continue if args.pad_to_max_len: + if not tokenizer.pad_token_id: + raise ValueError("Tokenizer does not have a pad token") # add padding to max_len input_ids = torch.nn.functional.pad( input_ids, (0, max_len - input_ids.shape[0]), value=tokenizer.pad_token_id @@ -402,7 +459,7 @@ def append_text_gen_input_ids(tokenized_inputs, input_ids, tokenizer, context: i # based on https://github.com/artidoro/qlora/blob/main/qlora.py -def format_pair_dataset(dataset, pair_format, input_col=None, output_col=None): +def format_pair_dataset(dataset, args): """Format dataset based on pair_format.""" # format for input in ALPACA pair format @@ -426,11 +483,12 @@ def extract_alpaca_dataset(example): prompt_format = ALPACA_PROMPT_DICT["prompt_input"] else: prompt_format = ALPACA_PROMPT_DICT["prompt_no_input"] - return {"input": prompt_format.format(**example)} + return apply_template(example, "input", prompt_format) - if pair_format == TextGenPairFormat.ALPACA: - dataset = dataset.map(extract_alpaca_dataset, remove_columns=["instruction"]) - elif pair_format == TextGenPairFormat.CHIP2: + if args.pair_format == TextGenPairFormat.ALPACA: + # extract new input from instruction and input + dataset = extract_alpaca_dataset(dataset) + elif args.pair_format == TextGenPairFormat.CHIP2: # separate the human and bot text into input and output dataset = dataset.map( lambda x: { @@ -438,7 +496,7 @@ def extract_alpaca_dataset(example): "output": x["text"].split("\n: ")[1], } ) - elif pair_format == TextGenPairFormat.SELF_INSTRUCT: + elif args.pair_format == TextGenPairFormat.SELF_INSTRUCT: # rename prompt and completion to input and output dataset = dataset.map( lambda x: { @@ -446,20 +504,43 @@ def extract_alpaca_dataset(example): "output": x["completion"], } ) - elif pair_format == TextGenPairFormat.CUSTOM: - # rename input_col and output_col to input and output + elif args.pair_format == TextGenPairFormat.CUSTOM: + + def _formatting_func(template, col, x): + if template: + return template.format(**x) + elif col: + return x[col] + else: + raise ValueError("Neither template nor column name is specified") + dataset = dataset.map( - lambda x: { - "input": x[input_col], - "output": x[output_col], + lambda y: { + "input": _formatting_func(args.input_template, args.input_col, y), + "output": _formatting_func(args.output_template, args.output_col, y), } ) - elif pair_format == TextGenPairFormat.DEFAULT: + elif args.pair_format == TextGenPairFormat.DEFAULT: # do nothing pass else: - raise ValueError(f"Invalid pair_format: {pair_format}") + raise ValueError(f"Invalid pair_format: {args.pair_format}") # remove unused columns, keep only input and output dataset = dataset.remove_columns([col for col in dataset.column_names if col not in ["input", "output"]]) return dataset + + +def apply_template(dataset, new_col: str, template: str, remove_cols: bool = False): + """ + Apply template to column in dataset. + + :param dataset: dataset to apply template to + :param new_col: name of new column + :param template: python f-string template with {column_name} as placeholders. The column names must be in dataset. + :param remove_cols: remove columns after applying template + """ + dataset = dataset.map( + lambda x: {new_col: template.format(**x)}, remove_columns=dataset.column_names if remove_cols else None + ) + return dataset