From bdbca8fa6ca01f9e668b1fa886853b247646b2e7 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 15 May 2023 14:07:17 -0400 Subject: [PATCH 01/32] more fixes --- scripts/finetune.py | 6 +++--- src/axolotl/utils/models.py | 12 ++++++++---- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/scripts/finetune.py b/scripts/finetune.py index d6a920f5d..ee055559f 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -1,7 +1,6 @@ import importlib import logging import os -import pathlib import random import signal import sys @@ -10,7 +9,6 @@ from typing import Optional import fire import torch -import transformers import yaml from attrdict import AttrDefault @@ -236,7 +234,9 @@ def train( logging.info(f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}") # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading - model.save_pretrained(cfg.output_dir) + # only save on rank 0, otherwise it corrupts output on multi-GPU when multiple processes attempt to write the same file + if cfg.local_rank == 0: + model.save_pretrained(cfg.output_dir) # trainer.save_model(cfg.output_dir) # TODO this may be needed for deepspeed to work? need to review another time diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 4d7a45920..4e3b4efd6 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -1,4 +1,5 @@ import logging +import math import os from pathlib import Path from typing import Optional, Tuple, TYPE_CHECKING @@ -180,12 +181,14 @@ def load_model( tokenizer.add_special_tokens({"pad_token": "[PAD]"}) os.environ["TOKENIZERS_PARALLELISM"] = "false" - if cfg.tokens: - for k, v in cfg.tokens.items(): + if cfg.special_tokens: + for k, v in cfg.special_tokens.items(): tokenizer.add_special_tokens({k: v}) + if cfg.tokens: + tokenizer.add_tokens(cfg.tokens) - # this should only be needed if you are messing with new tokens in the vocab - # model.resize_token_embeddings(len(tokenizer)) + embeddings_len = math.ceil(len(tokenizer) / 32) * 32 + model.resize_token_embeddings(embeddings_len) if cfg.adapter and load_in_8bit and not cfg.load_4bit: logging.info("converting PEFT model w/ prepare_model_for_int8_training") @@ -221,6 +224,7 @@ def load_model( requires_grad.append(f"{name}: {param.requires_grad}") if len(requires_grad) == 0: logging.warning("there are no parameters that require gradient updates") + model.config.use_cache = False # TODO resume_from_checkpoint handling return model, tokenizer, lora_config From 5e3714475439523568840784970728e963b95374 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 15 May 2023 22:15:36 -0400 Subject: [PATCH 02/32] fix prompters, especially the sharegpt prompter --- src/axolotl/prompt_tokenizers.py | 84 +++++++++++++++++++++++++----- src/axolotl/prompters.py | 88 ++++++-------------------------- 2 files changed, 89 insertions(+), 83 deletions(-) diff --git a/src/axolotl/prompt_tokenizers.py b/src/axolotl/prompt_tokenizers.py index 00d8ecbf9..5792d191b 100644 --- a/src/axolotl/prompt_tokenizers.py +++ b/src/axolotl/prompt_tokenizers.py @@ -1,7 +1,10 @@ import abc +import copy from transformers import PreTrainedTokenizer +from axolotl.prompters import IGNORE_TOKEN_ID + IGNORE_INDEX = -100 LLAMA_DEFAULT_PAD_TOKEN = "[PAD]" LLAMA_DEFAULT_EOS_TOKEN = "" @@ -40,10 +43,10 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy): full_prompt = self._build_full_prompt(instruction, input, response) tokenized_full_prompt = self._tokenize(full_prompt) if not self.train_on_inputs: - user_prompt = self.prompter.build_prompt( + user_prompt = next(iter(self.prompter.build_prompt( instruction, input, - ) + ))) tokenized_user_prompt = self._tokenize(user_prompt, add_eos_token=False) user_prompt_len = len(tokenized_user_prompt["input_ids"]) # TODO this could be sped up using numpy array slicing @@ -54,11 +57,11 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy): return tokenized_full_prompt def _build_full_prompt(self, instruction, input, response): - return self.prompter.build_prompt( + return next(iter(self.prompter.build_prompt( instruction, input, response, - ) + ))) def _tokenize(self, prompt, add_eos_token=True): result = self.tokenizer( @@ -131,13 +134,13 @@ class CompletionPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): def tokenize_prompt(self, prompt): instruction = self.parse_instruction_fields(prompt) - full_prompt = self._build_full_prompt(instruction) + full_prompt = self._build_full_prompt(instruction, None, None) tokenized_full_prompt = self._tokenize(full_prompt) return tokenized_full_prompt - def _build_full_prompt(self, instruction): - return self.prompter.build_prompt(instruction) + def _build_full_prompt(self, instruction, input, response): + return next(iter(self.prompter.build_prompt(instruction))) class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy): @@ -157,10 +160,10 @@ class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy): ) tokenized_full_prompt = self._tokenize(full_prompt) if not self.train_on_inputs: - user_prompt = self.prompter.build_prompt( + user_prompt = next(iter(self.prompter.build_prompt( instruction, input, - ) + ))) tokenized_user_prompt = self._tokenize(user_prompt, add_eos_token=False) user_prompt_len = len(tokenized_user_prompt["input_ids"]) # TODO this could be sped up using numpy array slicing @@ -171,13 +174,13 @@ class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy): return tokenized_full_prompt def _build_full_prompt(self, instruction, input, output, reflection, corrected): - return self.prompter.build_prompt( + return next(iter(self.prompter.build_prompt( instruction, input, output, reflection, corrected, - ) + ))) def _tokenize(self, prompt, add_eos_token=True): result = self.tokenizer( @@ -212,7 +215,64 @@ class AlpacaReflectionPTStrategy(ReflectionPromptTokenizingStrategy): class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy): def tokenize_prompt(self, prompt): + result = { + "input_ids": [], + "attention_mask": [], + "labels": [], + } + current_len = 0 try: - return self.prompter.build_prompt(prompt["conversations"], self.tokenizer) + for i, part in enumerate(self.prompter.build_prompt(prompt["conversations"], self.tokenizer)): + if i == 0: + # this is only ever the first part, should include the bos token and the user query + res = self._tokenize(part.strip(), add_eos_token=False, strip_bos_token=False) + # everything from this is masked out from the labels + labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) + elif i % 2 == 0: + # this is still the user query, we should + res = self._tokenize(part.strip(), add_eos_token=False, strip_bos_token=True) + # everything from this is masked out from the labels + labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) + else: + # this should be the assistent response, should end with an eos token + res = self._tokenize(part.strip(), add_eos_token=True, strip_bos_token=True) + # not masked out from labels + labels = copy.deepcopy(res["input_ids"]) + input_ids = res["input_ids"] + input_len = len(input_ids) + result["input_ids"][current_len : current_len + input_len] = input_ids + result["attention_mask"][current_len : current_len + input_len] = [ + 1 if x != self.tokenizer.pad_token_id else 0 + for x in input_ids + ] + result["labels"][current_len : current_len + input_len] = labels + current_len += input_len + return result except (KeyError, AssertionError, IndexError) as e: raise InvalidDataException(str(e)) + + def _tokenize(self, prompt, add_eos_token=True, strip_bos_token=False): + result = self.tokenizer( + prompt, + truncation=True, + max_length=self.sequence_len, + padding=False, + return_tensors=None, + ) + if ( + result["input_ids"][-1] != self.tokenizer.eos_token_id + and len(result["input_ids"]) < self.sequence_len + and add_eos_token + ): + result["input_ids"].append(self.tokenizer.eos_token_id) + result["attention_mask"].append(1) + + if ( + result["input_ids"][0] == self.tokenizer.bos_token_id + and strip_bos_token + ): + result["input_ids"] = result["input_ids"][1:] + result["attention_mask"] = result["attention_mask"][1:] + + result["labels"] = result["input_ids"].copy() + return result diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py index 3dc5d6433..a52ed4ad9 100644 --- a/src/axolotl/prompters.py +++ b/src/axolotl/prompters.py @@ -1,7 +1,7 @@ import copy import dataclasses from enum import auto, Enum -from typing import List, Tuple, Any, Union +from typing import List, Tuple, Any, Union, Generator IGNORE_TOKEN_ID = -100 @@ -16,7 +16,7 @@ class AlpacaPrompter: instruction: str, input: Union[None, str] = None, output: Union[None, str] = None, - ) -> str: + ) -> Generator[str, None, None]: # returns the full prompt from instruction and optional input # if a label (=response, =output) is provided, it's also appended. if input: @@ -25,7 +25,7 @@ class AlpacaPrompter: res = self.prompt_no_input.format(instruction=instruction) if output: res = f"{res}{output}" - return res + yield res def get_response(self, output: str) -> str: return output.split(self.response_split)[1].strip() @@ -36,8 +36,8 @@ class JeopardyPrompter(AlpacaPrompter): class CompletionPrompter(AlpacaPrompter): - def build_prompt(self, instruction: str) -> str: - return instruction + def build_prompt(self, instruction: str, input=None, output=None) -> Generator[str, None, None]: + yield instruction def get_response(self, output: str) -> str: return output.strip() @@ -64,7 +64,7 @@ class ReflectAlpacaPrompter: output: Union[None, str] = None, reflection: Union[None, str] = None, corrected: Union[None, str] = None, - ) -> str: + ) -> Generator[str, None, None]: # returns the full prompt from instruction and optional input # if a label (=response, =output) is provided, it's also appended. if input: @@ -76,7 +76,7 @@ class ReflectAlpacaPrompter: output=output, reflection=reflection, corrected=corrected ) res = f"{res}{label}" - return res + yield res def get_response(self, output: str) -> str: return output.split(self.response_split)[1].strip() @@ -103,15 +103,16 @@ class Conversation: sep: str = "###" sep2: str = None - def get_prompt(self): + def get_prompt(self) -> Generator[str, None, None]: seps = [self.sep, self.sep2] - ret = self.system + seps[0] + preamble = self.system + seps[0] for i, (role, message) in enumerate(self.messages): if message: - ret += role + ": " + message + seps[i % 2] + yield preamble + role + ": " + message + seps[i % 2] else: - ret += role + ":" - return ret + yield role + ":" + if i == 0: + preamble = "" def copy(self): return Conversation( @@ -136,12 +137,12 @@ conv_vicuna_v1_1 = Conversation( offset=0, sep_style=SeparatorStyle.TWO, sep=" ", - sep2="", + sep2=" ", ) class ShareGPTPrompter: - def build_prompt(self, source, tokenizer, sequence_len=2048): + def build_prompt(self, source, tokenizer, sequence_len=2048) -> Generator[str, None, None]: # ignore the system prompt if provided if source[0]["from"] == "system": source.pop(0) @@ -171,61 +172,6 @@ class ShareGPTPrompter: role = roles[sentence["from"]] assert role == conv.roles[j % 2] conv.append_message(role, sentence["value"]) - # TODO, this concatenates everything, but doesn't seem to properly add the eos_token_id, as the eos_token gets split up - conversation = conv.get_prompt() - # Tokenize conversations - tokenized_result = tokenizer( - conversation, - truncation=True, - max_length=sequence_len, # FIXME - padding=False, - return_tensors=None, - ) - target = copy.deepcopy(tokenized_result["input_ids"]) - - # Mask targets - sep = conv.sep + conv.roles[1] + ": " - - rounds = conversation.split(conv.sep2) - rounds = [r + conv.sep2 for r in rounds] - cur_len = 1 - target[0] = IGNORE_TOKEN_ID # mask out the bos - for i, rou in enumerate(rounds): - if rou == "": - break - - parts = rou.split(sep) - if len(parts) != 2: - break - parts[0] += sep - round_len = ( - len(tokenizer(rou)["input_ids"]) - 1 - ) # -1 ignores the bos_token generated for this - # we have to strip the initial part, any dangling whitespace creates an additional ghost token - instruction_len = ( - len(tokenizer(parts[0].strip())["input_ids"]) - 1 - ) # -1 ignores the bos_token generated for this - target[cur_len : cur_len + instruction_len] = [ - IGNORE_TOKEN_ID - ] * instruction_len - - cur_len += round_len - if cur_len >= sequence_len: - break - - # Fix: Truncate the target to have the same length as input_ids - target = target[: len(tokenized_result["input_ids"])] - # target[cur_len:] = [IGNORE_TOKEN_ID] * (len(target) - cur_len) - - attention_mask = [ - 1 if x != tokenizer.pad_token_id else 0 - for x in tokenized_result["input_ids"] - ] - - # TODO truncate len to sequence_len - return dict( - input_ids=tokenized_result["input_ids"], - labels=target, - attention_mask=attention_mask, - ) + for part in conv.get_prompt(): + yield part From f98e173b5960d320b0414bba3759c899f2648759 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 15 May 2023 22:26:30 -0400 Subject: [PATCH 03/32] reorder options so debug can happen in the same prepare step --- scripts/finetune.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/finetune.py b/scripts/finetune.py index ee055559f..5b7f8f2ab 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -182,10 +182,6 @@ def train( tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH ) - if prepare_ds_only: - logging.info("Finished preparing dataset. Exiting...") - return - if cfg.debug: logging.info("check_dataset_labels...") check_dataset_labels( @@ -195,6 +191,10 @@ def train( tokenizer, ) + if prepare_ds_only: + logging.info("Finished preparing dataset. Exiting...") + return + trainer = setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer) model.config.use_cache = False From b46bc02f0aa89d09e8dbbb22373c84d6ccb1807a Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 16 May 2023 21:45:34 -0400 Subject: [PATCH 04/32] add alpaca multiple choice instruct dataset support --- scripts/finetune.py | 2 +- src/axolotl/prompt_tokenizers.py | 9 +++++++++ src/axolotl/prompters.py | 4 ++++ src/axolotl/utils/data.py | 10 ++++++++-- 4 files changed, 22 insertions(+), 3 deletions(-) diff --git a/scripts/finetune.py b/scripts/finetune.py index 5b7f8f2ab..5fb38b6f6 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -67,7 +67,7 @@ def do_inference(cfg, model, tokenizer, prompter="AlpacaPrompter"): instruction = get_multi_line_input() if not instruction: return - prompt = prompter_module().build_prompt(instruction=instruction) + prompt: str = next(prompter_module().build_prompt(instruction=instruction)) batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True) model.eval() diff --git a/src/axolotl/prompt_tokenizers.py b/src/axolotl/prompt_tokenizers.py index 5792d191b..9cc4003c2 100644 --- a/src/axolotl/prompt_tokenizers.py +++ b/src/axolotl/prompt_tokenizers.py @@ -92,6 +92,15 @@ class AlpacaPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): ) +class AlpacaMultipleChoicePromptTokenizingStrategy(InstructionPromptTokenizingStrategy): + def parse_instruction_fields(self, prompt) -> (str, str, str): + return ( + prompt["question"], + "\n".join(f'- "{choice}"' for choice in prompt["choices"]), + prompt["explanation"], + ) + + class JeopardyPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): def parse_instruction_fields(self, prompt) -> (str, str, str): return ( diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py index a52ed4ad9..e4411b2c4 100644 --- a/src/axolotl/prompters.py +++ b/src/axolotl/prompters.py @@ -35,6 +35,10 @@ class JeopardyPrompter(AlpacaPrompter): prompt_input = "Below is a Jeopardy clue paired with input providing the category of the clue. Write a concise response that best answers tbe clue given the category.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" +class MultipleChoiceExplainPrompter(AlpacaPrompter): + prompt_input = "Choose the answer that best answers the question. Explain your reasoning.\n\n### Question:\n{instruction}\n\n### Choices:\n{input}\n\n### Response:\n" + + class CompletionPrompter(AlpacaPrompter): def build_prompt(self, instruction: str, input=None, output=None) -> Generator[str, None, None]: yield instruction diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index 306213b19..2d2e7f2fd 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -19,7 +19,7 @@ from axolotl.prompt_tokenizers import ( AlpacaReflectionPTStrategy, ShareGPTPromptTokenizingStrategy, JeopardyPromptTokenizingStrategy, - CompletionPromptTokenizingStrategy, + CompletionPromptTokenizingStrategy, AlpacaMultipleChoicePromptTokenizingStrategy, ) from axolotl.prompters import ( AlpacaPrompter, @@ -27,7 +27,7 @@ from axolotl.prompters import ( ReflectAlpacaPrompter, ShareGPTPrompter, JeopardyPrompter, - CompletionPrompter, + CompletionPrompter, MultipleChoiceExplainPrompter, ) @@ -88,6 +88,12 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) + elif d.type == "explainchoice": + ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy( + MultipleChoiceExplainPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) + ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) + datasets.append(ds_wrapper) elif d.type == "jeopardy": ds_strategy = JeopardyPromptTokenizingStrategy( JeopardyPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len From 8c2f3cb0f86367f320e7f1421891e264c6fa2bca Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 17 May 2023 08:49:03 -0400 Subject: [PATCH 05/32] support for replit lm --- examples/replit-3b/config-lora.yml | 55 ++++++++++++++++++++++++++++++ src/axolotl/utils/models.py | 15 ++++++-- 2 files changed, 67 insertions(+), 3 deletions(-) create mode 100644 examples/replit-3b/config-lora.yml diff --git a/examples/replit-3b/config-lora.yml b/examples/replit-3b/config-lora.yml new file mode 100644 index 000000000..c757e720d --- /dev/null +++ b/examples/replit-3b/config-lora.yml @@ -0,0 +1,55 @@ +base_model: replit/replit-code-v1-3b +base_model_config: replit/replit-code-v1-3b +trust_remote_code: true +load_in_8bit: false +datasets: + - path: vicgalle/alpaca-gpt4 + type: alpaca +dataset_prepared_path: last_run_prepared +val_set_size: 0.05 +adapter: lora +lora_model_dir: +sequence_len: 2048 +max_packed_sequence_len: +lora_r: 8 +lora_alpha: 16 +lora_dropout: 0.05 +lora_target_modules: + - Wqkv + - mlp_up + - mlp_down +lora_fan_in_fan_out: +wandb_project: lora-replit +wandb_watch: +wandb_run_id: +wandb_log_model: +output_dir: ./lora-replit +batch_size: 8 +micro_batch_size: 1 +num_epochs: 3 +optimizer: +torchdistx_path: +lr_scheduler: +learning_rate: 0.00001 +train_on_inputs: false +group_by_length: false +bf16: true +tf32: true +gradient_checkpointing: +early_stopping_patience: +resume_from_checkpoint: +local_rank: +logging_steps: 1 +xformers_attention: +flash_attention: +gptq_groupsize: +gptq_model_v1: +warmup_steps: 20 +eval_steps: 50 +save_steps: +debug: +deepspeed: +weight_decay: 0 +fsdp: +fsdp_config: +#special_tokens: diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 4e3b4efd6..d93d859b7 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -163,11 +163,20 @@ def load_model( if not tokenizer: try: if is_llama_derived_model and "LlamaTokenizer" in globals(): - tokenizer = LlamaTokenizer.from_pretrained(model) + tokenizer = LlamaTokenizer.from_pretrained( + model, + trust_remote_code=True if cfg.trust_remote_code is True else False, + ) else: - tokenizer = getattr(transformers, tokenizer_type).from_pretrained(model) + tokenizer = getattr(transformers, tokenizer_type).from_pretrained( + model, + trust_remote_code=True if cfg.trust_remote_code is True else False, + ) except: - tokenizer = AutoTokenizer.from_pretrained(base_model_config) + tokenizer = AutoTokenizer.from_pretrained( + base_model_config, + trust_remote_code=True if cfg.trust_remote_code is True else False, + ) logging.debug(f"EOS: {tokenizer.eos_token_id} / {tokenizer.eos_token}") logging.debug(f"BOS: {tokenizer.bos_token_id} / {tokenizer.bos_token}") From 13650732f8f0212238fe7b70dfef213734a7985e Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 17 May 2023 11:29:17 -0400 Subject: [PATCH 06/32] concise multiple choice and tldr summarize --- src/axolotl/prompt_tokenizers.py | 11 ++++++++++- src/axolotl/prompters.py | 8 ++++++++ src/axolotl/utils/data.py | 20 ++++++++++++++++++-- 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/src/axolotl/prompt_tokenizers.py b/src/axolotl/prompt_tokenizers.py index 9cc4003c2..7f79ef192 100644 --- a/src/axolotl/prompt_tokenizers.py +++ b/src/axolotl/prompt_tokenizers.py @@ -97,7 +97,7 @@ class AlpacaMultipleChoicePromptTokenizingStrategy(InstructionPromptTokenizingSt return ( prompt["question"], "\n".join(f'- "{choice}"' for choice in prompt["choices"]), - prompt["explanation"], + prompt["solution"] if "solution" in prompt else prompt["explanation"], ) @@ -119,6 +119,15 @@ class OpenAssistantPromptTokenizingStrategy(InstructionPromptTokenizingStrategy) ) +class SummarizeTLDRPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): + def parse_instruction_fields(self, prompt) -> (str, str, str): + return ( + prompt["article"], + "", + prompt["summary"], + ) + + class GPTeacherPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): def parse_instruction_fields(self, prompt) -> (str, str, str): return ( diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py index e4411b2c4..8a8cfa247 100644 --- a/src/axolotl/prompters.py +++ b/src/axolotl/prompters.py @@ -39,6 +39,14 @@ class MultipleChoiceExplainPrompter(AlpacaPrompter): prompt_input = "Choose the answer that best answers the question. Explain your reasoning.\n\n### Question:\n{instruction}\n\n### Choices:\n{input}\n\n### Response:\n" +class MultipleChoiceConcisePrompter(AlpacaPrompter): + prompt_input = "Choose the answer that best answers the question. Be concise in your response.\n\nUSER: {instruction}\n{input}\nASSISTANT:\n" + + +class SummarizeTLDRPrompter(AlpacaPrompter): + prompt_no_input = "USER: Summarize the following article as a TL;DR.\n{instruction}\nASSISTANT:" + + class CompletionPrompter(AlpacaPrompter): def build_prompt(self, instruction: str, input=None, output=None) -> Generator[str, None, None]: yield instruction diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index 2d2e7f2fd..28b6ee072 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -19,7 +19,9 @@ from axolotl.prompt_tokenizers import ( AlpacaReflectionPTStrategy, ShareGPTPromptTokenizingStrategy, JeopardyPromptTokenizingStrategy, - CompletionPromptTokenizingStrategy, AlpacaMultipleChoicePromptTokenizingStrategy, + CompletionPromptTokenizingStrategy, + AlpacaMultipleChoicePromptTokenizingStrategy, + SummarizeTLDRPromptTokenizingStrategy, ) from axolotl.prompters import ( AlpacaPrompter, @@ -27,7 +29,9 @@ from axolotl.prompters import ( ReflectAlpacaPrompter, ShareGPTPrompter, JeopardyPrompter, - CompletionPrompter, MultipleChoiceExplainPrompter, + CompletionPrompter, + MultipleChoiceExplainPrompter, + SummarizeTLDRPrompter, MultipleChoiceConcisePrompter, ) @@ -94,6 +98,18 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) + elif d.type == "concisechoice": + ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy( + MultipleChoiceConcisePrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) + ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) + datasets.append(ds_wrapper) + elif d.type == "summarizetldr": + ds_strategy = SummarizeTLDRPromptTokenizingStrategy( + SummarizeTLDRPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) + ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) + datasets.append(ds_wrapper) elif d.type == "jeopardy": ds_strategy = JeopardyPromptTokenizingStrategy( JeopardyPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len From fa8bd14be49ec234e28da1f444880bc9859d7589 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Thu, 18 May 2023 06:25:34 -0400 Subject: [PATCH 07/32] update entrypoint and force min accelerate --- docker/Dockerfile-runpod | 7 +++++-- requirements.txt | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile-runpod b/docker/Dockerfile-runpod index f7dac3003..a641f241e 100644 --- a/docker/Dockerfile-runpod +++ b/docker/Dockerfile-runpod @@ -1,11 +1,14 @@ ARG BASE_TAG=main FROM winglian/axolotl:$BASE_TAG +COPY scripts/runpod-entrypoint.sh /root/runpod-entrypoint.sh + RUN apt install --yes --no-install-recommends openssh-server tmux && \ mkdir -p ~/.ssh && \ chmod 700 ~/.ssh && \ printf "\n[[ -z \"\$TMUX\" ]] && { tmux attach-session -t ssh_tmux || tmux new-session -s ssh_tmux; exit; }\n" >> ~/.bashrc && \ - chmod +x /workspace/axolotl/scripts/runpod-entrypoint.sh + chmod +x /workspace/axolotl/scripts/runpod-entrypoint.sh && \ + chmod +x /root/runpod-entrypoint.sh -ENTRYPOINT ["/workspace/axolotl/scripts/runpod-entrypoint.sh"] +ENTRYPOINT ["/root/runpod-entrypoint.sh"] CMD ["sleep", "infinity"] diff --git a/requirements.txt b/requirements.txt index 13450f470..ae6689680 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,7 +6,7 @@ PyYAML==6.0 black bitsandbytes==0.37.2 datasets -accelerate +accelerate>=0.19.0 sentencepiece wandb einops From 1d5ab84486a965a50353738064cb381cca73a501 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sat, 20 May 2023 18:16:22 -0400 Subject: [PATCH 08/32] optionally be able to specify alpaca or chat style prompts --- docker/Dockerfile-base | 15 +++++- scripts/finetune.py | 16 +++++- src/axolotl/prompt_tokenizers.py | 60 ++++++++++++++++++----- src/axolotl/prompters.py | 72 ++++++++++++++++++++++----- src/axolotl/utils/data.py | 83 +++++++++++++++++++++++--------- src/axolotl/utils/models.py | 30 +++++++++++- 6 files changed, 223 insertions(+), 53 deletions(-) diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base index 54738ddb8..943bae3b0 100644 --- a/docker/Dockerfile-base +++ b/docker/Dockerfile-base @@ -1,6 +1,7 @@ ARG CUDA_VERSION="11.8.0" ARG CUDNN_VERSION="8" ARG UBUNTU_VERSION="22.04" +ARG MAX_JOBS=4 FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION as base-builder @@ -39,6 +40,14 @@ ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX" RUN git clone https://github.com/HazyResearch/flash-attention.git && \ cd flash-attention && \ + python3 setup.py bdist_wheel && \ + cd csrc/fused_dense_lib && \ + python3 setup.py bdist_wheel && \ + cd csrc/xentropy && \ + python3 setup.py bdist_wheel && \ + cd csrc/rotary && \ + python3 setup.py bdist_wheel && \ + cd csrc/layer_norm && \ python3 setup.py bdist_wheel FROM base-builder AS deepspeed-builder @@ -60,8 +69,12 @@ RUN cd apex && MAX_JOBS=1 python3 -m pip install --global-option="--cpp_ext" --g RUN mkdir /workspace/wheels COPY --from=deepspeed-builder /workspace/DeepSpeed/dist/deepspeed-*.whl wheels COPY --from=flash-attn-builder /workspace/flash-attention/dist/flash_attn-*.whl wheels +COPY --from=flash-attn-builder /workspace/flash-attention/csrc/fused_dense_lib/dist/fused_dense_lib-*.whl wheels +COPY --from=flash-attn-builder /workspace/flash-attention/csrc/xentropy/dist/xentropy-*.whl wheels +COPY --from=flash-attn-builder /workspace/flash-attention/csrc/rotary/dist/rotary-*.whl wheels +COPY --from=flash-attn-builder /workspace/flash-attention/csrc/layer_norm/dist/dropout_layer_norm-*.whl wheels -RUN pip3 install wheels/deepspeed-*.whl wheels/flash_attn-*.whl +RUN pip3 install wheels/deepspeed-*.whl wheels/flash_attn-*.whl wheels/fused_dense_lib-*.whl wheels/xeontropy-*.whl wheels/rotary-*.whl wheels/dropout_layer_norm-*.whl RUN git lfs install --skip-repo RUN pip3 install "peft @ git+https://github.com/huggingface/peft.git@main" \ "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \ diff --git a/scripts/finetune.py b/scripts/finetune.py index 5fb38b6f6..81d29ad28 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -31,7 +31,7 @@ DEFAULT_DATASET_PREPARED_PATH = "last_run_prepared" def choose_device(cfg): def get_device(): if torch.cuda.is_available(): - return "cuda" + return f"cuda:{cfg.local_rank}" else: try: if torch.backends.mps.is_available(): @@ -131,7 +131,8 @@ def train( # then overwrite the value cfg_keys = dict(cfg).keys() for k in kwargs: - if k in cfg_keys: + # if not strict, allow writing to cfg even if it's not in the yml already + if k in cfg_keys or cfg.strict is False: # handle booleans if isinstance(cfg[k], bool): cfg[k] = bool(kwargs[k]) @@ -169,6 +170,15 @@ def train( inference=("inference" in kwargs), ) + if "merge_lora" in kwargs and cfg.adapter is not None: + print("running merge of LoRA with base model") + model = model.merge_and_unload() + + if cfg.local_rank == 0: + print("saving merged model") + model.save_pretrained(str(Path(cfg.output_dir) / "merged")) + return + if "inference" in kwargs: logging.info("calling do_inference function") do_inference(cfg, model, tokenizer) @@ -216,6 +226,8 @@ def train( ) logging.info("Starting trainer...") + if cfg.group_by_length: + logging.info("hang tight... sorting dataset for group_by_length") resume_from_checkpoint = cfg.resume_from_checkpoint if cfg.resume_from_checkpoint is None and cfg.auto_resume_from_checkpoints: possible_checkpoints = [ diff --git a/src/axolotl/prompt_tokenizers.py b/src/axolotl/prompt_tokenizers.py index 7f79ef192..c33551135 100644 --- a/src/axolotl/prompt_tokenizers.py +++ b/src/axolotl/prompt_tokenizers.py @@ -1,5 +1,7 @@ import abc import copy +import functools +import logging from transformers import PreTrainedTokenizer @@ -33,6 +35,20 @@ class PromptTokenizingStrategy(abc.ABC): def tokenize_prompt(self, prompt): pass + @functools.cache + def _get_user_token(self): + id_or_ids = self.tokenizer.convert_tokens_to_ids("<|USER|>") + if type(id_or_ids, (int,)): + return id_or_ids + return False + + @functools.cache + def _get_assistant_token(self): + id_or_ids = self.tokenizer.convert_tokens_to_ids("<|ASSISTANT|>") + if type(id_or_ids, (int,)): + return id_or_ids + return False + class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy): def parse_instruction_fields(self, prompt) -> (str, str, str): @@ -63,7 +79,7 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy): response, ))) - def _tokenize(self, prompt, add_eos_token=True): + def _tokenize(self, prompt, add_eos_token=True, strip_bos_token=False): result = self.tokenizer( prompt, truncation=True, @@ -79,6 +95,13 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy): result["input_ids"].append(self.tokenizer.eos_token_id) result["attention_mask"].append(1) + if ( + result["input_ids"][0] == self.tokenizer.bos_token_id + and strip_bos_token + ): + result["input_ids"] = result["input_ids"][1:] + result["attention_mask"] = result["attention_mask"][1:] + result["labels"] = result["input_ids"].copy() return result @@ -239,23 +262,34 @@ class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy): "labels": [], } current_len = 0 + user_token = self._get_user_token() + assistant_token = self._get_assistant_token() try: - for i, part in enumerate(self.prompter.build_prompt(prompt["conversations"], self.tokenizer)): - if i == 0: + for i, part in enumerate(self.prompter.build_prompt(prompt["conversations"])): + if isinstance(part, tuple): + if part[0] == "USER:": + part = part[0] + part[1] if not user_token else part[1] + # this is still the user query, we should + res = self._tokenize(part.strip(), add_eos_token=False, strip_bos_token=True) + if user_token: + res = [user_token, *res] + # everything from this is masked out from the labels + labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) + elif part[0] == "ASSISTANT:": + part = part[0] + part[1] if not assistant_token else part[1] + # this should be the assistent response, should end with an eos token + res = self._tokenize(part.strip(), add_eos_token=True, strip_bos_token=True) + if assistant_token: + res = [assistant_token, *res] + # not masked out from labels + labels = copy.deepcopy(res["input_ids"]) + else: + logging.warning("unhandled role: " + part[0]) + else: # this is only ever the first part, should include the bos token and the user query res = self._tokenize(part.strip(), add_eos_token=False, strip_bos_token=False) # everything from this is masked out from the labels labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) - elif i % 2 == 0: - # this is still the user query, we should - res = self._tokenize(part.strip(), add_eos_token=False, strip_bos_token=True) - # everything from this is masked out from the labels - labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) - else: - # this should be the assistent response, should end with an eos token - res = self._tokenize(part.strip(), add_eos_token=True, strip_bos_token=True) - # not masked out from labels - labels = copy.deepcopy(res["input_ids"]) input_ids = res["input_ids"] input_len = len(input_ids) result["input_ids"][current_len : current_len + input_len] = input_ids diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py index 8a8cfa247..3ae0a0bd4 100644 --- a/src/axolotl/prompters.py +++ b/src/axolotl/prompters.py @@ -1,15 +1,34 @@ import copy import dataclasses +import logging from enum import auto, Enum from typing import List, Tuple, Any, Union, Generator IGNORE_TOKEN_ID = -100 +class PromptStyle(Enum): + instruct = "instruct" + chat = "chat" + class AlpacaPrompter: - prompt_input = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" - prompt_no_input = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n" - response_split = "### Response:" + system_prompt = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n" + system_no_input_prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n" + prompt_style = None + + def __init__(self, prompt_style="instruct"): + self.prompt_style = prompt_style + self.match_prompt_style() + + def match_prompt_style(self): + if self.prompt_style == PromptStyle.instruct.value: + self.prompt_input = self.system_prompt + "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" + self.prompt_no_input = self.system_no_input_prompt + "### Instruction:\n{instruction}\n\n### Response:\n" + self.response_split = "### Response:" + if self.prompt_style == PromptStyle.chat.value: + self.prompt_input = self.system_prompt + "USER: {instruction}\n{input}\nASSISTANT:" + self.prompt_no_input = self.system_no_input_prompt + "USER: {instruction}\nASSISTANT:" + self.response_split = "ASSISTANT:" def build_prompt( self, @@ -36,7 +55,7 @@ class JeopardyPrompter(AlpacaPrompter): class MultipleChoiceExplainPrompter(AlpacaPrompter): - prompt_input = "Choose the answer that best answers the question. Explain your reasoning.\n\n### Question:\n{instruction}\n\n### Choices:\n{input}\n\n### Response:\n" + system_prompt = "Choose the answer that best answers the question. Explain your reasoning." class MultipleChoiceConcisePrompter(AlpacaPrompter): @@ -64,11 +83,30 @@ class NomicGPT4AllPrompter(AlpacaPrompter): class ReflectAlpacaPrompter: - prompt_input = "Below is an instruction that describes a task, paired with an input that provides further context. You, the Assistant, should generate a response as if it were an abstract for an academic or technical paper on the query along with a methodology. Then generate an Agent Reflection where you create a long form response as if from subject matter expert, be verbose, diligent, and creative in your application of knowledge, apply it through the lens of the response generated by the assistant. Look for flawed reasoning, faulty logic, or other mistakes in the method. Finally, generate a final response and method for the user with the Assistant abstract and Reflection analysis as augmentations to the generation\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" - prompt_no_input = "Below is an instruction that describes a task. You, the Assistant, should generate a response as if it were an abstract for an academic or technical paper on the query along with a methodology. Then generate an Agent Reflection where you create a long form response as if from subject matter expert, be verbose, diligent, and creative in your application of knowledge, apply it through the lens of the response generated by the assistant. Look for flawed reasoning, faulty logic, or other mistakes in the method. Finally, generate a final response and method for the user with the Assistant abstract and Reflection analysis as augmentations to the generation\n\n### Instruction:\n{instruction}\n\n### Response:\n" - agent_label = "{output}\n\n### Agent Reflection:\n{reflection}\n\n### Final Response:\n{corrected}" + system_prompt = "Below is an instruction that describes a task, paired with an input that provides further context. You, the Assistant, should generate a response as if it were an abstract for an academic or technical paper on the query along with a methodology. Then generate an Agent Reflection where you create a long form response as if from subject matter expert, be verbose, diligent, and creative in your application of knowledge, apply it through the lens of the response generated by the assistant. Look for flawed reasoning, faulty logic, or other mistakes in the method. Finally, generate a final response and method for the user with the Assistant abstract and Reflection analysis as augmentations to the generation\n\n" + system_no_input_prompt = "Below is an instruction that describes a task. You, the Assistant, should generate a response as if it were an abstract for an academic or technical paper on the query along with a methodology. Then generate an Agent Reflection where you create a long form response as if from subject matter expert, be verbose, diligent, and creative in your application of knowledge, apply it through the lens of the response generated by the assistant. Look for flawed reasoning, faulty logic, or other mistakes in the method. Finally, generate a final response and method for the user with the Assistant abstract and Reflection analysis as augmentations to the generation\n\n" + + prompt_input = "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" + prompt_no_input = "### Instruction:\n{instruction}\n\n### Response:\n" + agent_label = "### Thought:\n{output}\n\n### Agent Reflection:\n{reflection}\n\n### Final Response:\n{corrected}" response_split = "### Response:" + def __init__(self, prompt_style="instruct"): + self.prompt_style = prompt_style + self.match_prompt_style() + + def match_prompt_style(self): + if self.prompt_style == PromptStyle.instruct.value: + self.prompt_input = self.system_prompt + "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" + self.prompt_no_input = self.system_no_input_prompt + "### Instruction:\n{instruction}\n\n### Response:\n" + self.agent_label = "### Thought:\n{output}\n\n### Agent Reflection:\n{reflection}\n\n### Final Response:\n{corrected}" + self.response_split = "### Final Response:" + if self.prompt_style == PromptStyle.chat.value: + self.prompt_input = self.system_prompt + "USER: {instruction}\n{input}\nASSISTANT:" + self.prompt_no_input = self.system_no_input_prompt + "USER: {instruction}\nASSISTANT:" + self.agent_label = "\nTHOUGHT: {output}\nASSISTANT REFLECTION: {reflection}\nASSISTANT:" + self.response_split = "ASSISTANT:" + def build_prompt( self, instruction: str, @@ -118,13 +156,13 @@ class Conversation: def get_prompt(self) -> Generator[str, None, None]: seps = [self.sep, self.sep2] preamble = self.system + seps[0] + yield preamble for i, (role, message) in enumerate(self.messages): if message: - yield preamble + role + ": " + message + seps[i % 2] + yield (role + ":", " " + message) else: - yield role + ":" - if i == 0: - preamble = "" + logging.warning("role with empty message: " + role) + yield (role + ":", ) def copy(self): return Conversation( @@ -154,7 +192,17 @@ conv_vicuna_v1_1 = Conversation( class ShareGPTPrompter: - def build_prompt(self, source, tokenizer, sequence_len=2048) -> Generator[str, None, None]: + def __init__(self, prompt_style=None): + if prompt_style != PromptStyle.chat.value: + raise Exception(f"unsupported prompt_style for ShareGPTPrompter({prompt_style})") + + # def match_prompt_style(self): + # if self.prompt_style == PromptStyle.chat.value: + # self.prompt_input = self.system_prompt + "USER: {instruction}\n{input}\nASSISTANT:" + # self.prompt_no_input = self.system_no_input_prompt + "USER: {instruction}\nASSISTANT:" + # self.response_split = "ASSISTANT:" + + def build_prompt(self, source, *args, **kwargs) -> Generator[str, None, None]: # ignore the system prompt if provided if source[0]["from"] == "system": source.pop(0) diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index 28b6ee072..22bc23359 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -50,8 +50,16 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa if cfg.dataset_prepared_path else Path(default_dataset_prepared_path) / ds_hash ) + dataset = None + try: + if cfg.push_dataset_to_hub: + dataset = load_dataset(f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True) + except: + pass - if any(prepared_ds_path.glob("*")): + if dataset: + ... + elif any(prepared_ds_path.glob("*")): logging.info(f"Loading prepared dataset from disk at {prepared_ds_path}...") dataset = load_from_disk(str(prepared_ds_path)) logging.info("Prepared dataset loaded from disk...") @@ -85,68 +93,71 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa ds = load_dataset("json", data_files=fp, streaming=True, split=None) if not ds: raise Exception("unhandled dataset load") - - if d.type == "alpaca": + d_type = d.type + d_type_split = d.type.split(":") + d_base_type = d_type_split[0] + d_prompt_style = d_type_split[1] if len(d_type_split) > 1 else None + if d_base_type == "alpaca": ds_strategy = AlpacaPromptTokenizingStrategy( - AlpacaPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + AlpacaPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) - elif d.type == "explainchoice": + elif d_base_type == "explainchoice": ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy( - MultipleChoiceExplainPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + MultipleChoiceExplainPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) - elif d.type == "concisechoice": + elif d_base_type == "concisechoice": ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy( - MultipleChoiceConcisePrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + MultipleChoiceConcisePrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) - elif d.type == "summarizetldr": + elif d_base_type == "summarizetldr": ds_strategy = SummarizeTLDRPromptTokenizingStrategy( - SummarizeTLDRPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + SummarizeTLDRPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) - elif d.type == "jeopardy": + elif d_base_type == "jeopardy": ds_strategy = JeopardyPromptTokenizingStrategy( - JeopardyPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + JeopardyPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) - elif d.type == "oasst": + elif d_base_type == "oasst": ds_strategy = OpenAssistantPromptTokenizingStrategy( - AlpacaPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + AlpacaPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) - elif d.type == "gpteacher": + elif d_base_type == "gpteacher": ds_strategy = GPTeacherPromptTokenizingStrategy( - GPTeacherPrompter(), + GPTeacherPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len, ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) - elif d.type == "reflection": + elif d_base_type == "reflection": ds_strategy = AlpacaReflectionPTStrategy( - ReflectAlpacaPrompter(), + ReflectAlpacaPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len, ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) - elif d.type == "sharegpt": + elif d_base_type == "sharegpt": ds_strategy = ShareGPTPromptTokenizingStrategy( - ShareGPTPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ShareGPTPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) - elif d.type == "completion": + elif d_base_type == "completion": ds_strategy = CompletionPromptTokenizingStrategy( CompletionPrompter(), tokenizer, @@ -168,6 +179,11 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa f"Saving merged prepared dataset to disk... {prepared_ds_path}" ) dataset.save_to_disk(prepared_ds_path) + if cfg.push_dataset_to_hub: + logging.info( + f"Saving merged prepared dataset with push_to_hub... {cfg.push_dataset_to_hub}/{ds_hash}" + ) + dataset.push_to_hub(f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True) return dataset @@ -182,13 +198,14 @@ def load_prepare_datasets(tokenizer, cfg, default_dataset_prepared_path): if cfg.max_packed_sequence_len is not None: # see if we can go ahead and load the stacked dataset - + seed = f"@{str(cfg.seed)}" if cfg.seed else "" ds_hash = str( md5( ( str(cfg.sequence_len) + "@" + str(max_packed_sequence_len) + + seed + "|".join(sorted([f"{d.path}:{d.type}" for d in cfg.datasets])) ).encode("utf-8") ).hexdigest() @@ -199,7 +216,19 @@ def load_prepare_datasets(tokenizer, cfg, default_dataset_prepared_path): else Path(default_dataset_prepared_path) / ds_hash ) - if any(prepared_ds_path.glob("*")): + dataset = None + try: + if cfg.push_dataset_to_hub: + logging.info( + f"checkking for packed prepared dataset from hub... {cfg.push_dataset_to_hub}/{ds_hash}" + ) + dataset = load_dataset(f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True) + except: + pass + + if dataset: + ... + elif any(prepared_ds_path.glob("*")): logging.info( f"Loading prepared packed dataset from disk at {prepared_ds_path}..." ) @@ -210,6 +239,9 @@ def load_prepare_datasets(tokenizer, cfg, default_dataset_prepared_path): tokenizer, cfg, default_dataset_prepared_path ) + if cfg.seed: + dataset = dataset.shuffle(seed=cfg.seed) + constant_len_dataset = ConstantLengthDataset( tokenizer, [dataset], @@ -237,6 +269,11 @@ def load_prepare_datasets(tokenizer, cfg, default_dataset_prepared_path): f"Saving packed prepared dataset to disk... {prepared_ds_path}" ) dataset.save_to_disk(prepared_ds_path) + if cfg.push_dataset_to_hub: + logging.info( + f"Saving packed prepared dataset with push_to_hub... {cfg.push_dataset_to_hub}/{ds_hash}" + ) + dataset.push_to_hub(f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True) else: dataset = load_tokenized_prepared_datasets( tokenizer, cfg, default_dataset_prepared_path diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index d93d859b7..0217f062b 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -126,6 +126,32 @@ def load_model( torch_dtype=torch_dtype, device_map=cfg.device_map, ) + # elif model_type == "GPTNeoXForCausalLM" and cfg.flash_attention: + # This is a WIP, still an issue with the backward pass + # RuntimeError: grad can be implicitly created only for scalar outputs + # TODO: try config.sequence_parallel = False + # # https://github.com/HazyResearch/flash-attention/blob/40a25c8ee7465cf547b929cfa2937034e37bfce9/tests/models/test_gpt_neox.py#L12 + # # https://github.com/HazyResearch/flash-attention/tree/main/training#model-components + # # add `**kwargs` to https://github.com/HazyResearch/flash-attention/blob/40a25c8ee7465cf547b929cfa2937034e37bfce9/flash_attn/models/gpt.py#L442 + # from flash_attn.utils.pretrained import state_dict_from_pretrained + # from flash_attn.models.gpt import GPTLMHeadModel + # from flash_attn.models.gpt_neox import remap_state_dict_hf_gpt_neox, gpt_neox_config_to_gpt2_config + # from transformers import GPTNeoXConfig + # config = gpt_neox_config_to_gpt2_config(GPTNeoXConfig.from_pretrained(base_model)) + # config.use_flash_attn = True + # config.fused_bias_fc = True + # config.fused_mlp = True # GPT-NeoX-20B uses "gelu_fast" + # config.activation_function = "gelu_fast" + # config.fused_dropout_add_ln = True + # # config.residual_in_fp32 = True + # + # model: GPTLMHeadModel = GPTLMHeadModel.from_pretrained( + # base_model, + # config, + # dtype=torch_dtype, + # device=cfg.device, + # ) + # model.train() # sets to train instead of eval mode elif model_type: model = getattr(transformers, model_type).from_pretrained( base_model, @@ -266,7 +292,7 @@ def load_llama_adapter(model, cfg): task_type="CAUSAL_LM", ) - if cfg.peft_model_dir: + if cfg.lora_model_dir: model = PeftModel.from_pretrained( model, cfg.lora_model_dir, @@ -307,7 +333,7 @@ def load_lora(model, cfg): model, cfg.lora_model_dir, device_map=cfg.device_map, - torch_dtype=torch.float16, + # torch_dtype=torch.float16, ) else: model = get_peft_model(model, lora_config) From 4ea9a66dbda265de487a0c51404fa096f289690e Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 08:33:06 -0400 Subject: [PATCH 09/32] tokenization fixes --- src/axolotl/prompt_strategies/alpaca_chat.py | 8 ++++++++ src/axolotl/prompt_tokenizers.py | 9 +++++---- src/axolotl/utils/data.py | 8 ++++++-- src/axolotl/utils/models.py | 2 +- 4 files changed, 20 insertions(+), 7 deletions(-) create mode 100644 src/axolotl/prompt_strategies/alpaca_chat.py diff --git a/src/axolotl/prompt_strategies/alpaca_chat.py b/src/axolotl/prompt_strategies/alpaca_chat.py new file mode 100644 index 000000000..1cd99bd9f --- /dev/null +++ b/src/axolotl/prompt_strategies/alpaca_chat.py @@ -0,0 +1,8 @@ +from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy +from axolotl.prompters import AlpacaPrompter, PromptStyle + + +def load(tokenizer, cfg): + return AlpacaPromptTokenizingStrategy( + AlpacaPrompter(PromptStyle.chat), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) diff --git a/src/axolotl/prompt_tokenizers.py b/src/axolotl/prompt_tokenizers.py index c33551135..6c20e7729 100644 --- a/src/axolotl/prompt_tokenizers.py +++ b/src/axolotl/prompt_tokenizers.py @@ -38,14 +38,14 @@ class PromptTokenizingStrategy(abc.ABC): @functools.cache def _get_user_token(self): id_or_ids = self.tokenizer.convert_tokens_to_ids("<|USER|>") - if type(id_or_ids, (int,)): + if isinstance(id_or_ids, (int,)): return id_or_ids return False @functools.cache def _get_assistant_token(self): id_or_ids = self.tokenizer.convert_tokens_to_ids("<|ASSISTANT|>") - if type(id_or_ids, (int,)): + if isinstance(id_or_ids, (int,)): return id_or_ids return False @@ -272,15 +272,16 @@ class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy): # this is still the user query, we should res = self._tokenize(part.strip(), add_eos_token=False, strip_bos_token=True) if user_token: - res = [user_token, *res] + res["input_ids"] = [user_token, *res["input_ids"]] # everything from this is masked out from the labels labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) elif part[0] == "ASSISTANT:": + # TODO label assistant token/tokens w/ IGNORE_TOKEN_ID part = part[0] + part[1] if not assistant_token else part[1] # this should be the assistent response, should end with an eos token res = self._tokenize(part.strip(), add_eos_token=True, strip_bos_token=True) if assistant_token: - res = [assistant_token, *res] + res["input_ids"] = [assistant_token, *res["input_ids"]] # not masked out from labels labels = copy.deepcopy(res["input_ids"]) else: diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index 22bc23359..d436face3 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -12,6 +12,7 @@ from datasets import ( from huggingface_hub import hf_hub_download from axolotl.datasets import TokenizedPromptDataset, ConstantLengthDataset +from axolotl.prompt_strategies import load from axolotl.prompt_tokenizers import ( AlpacaPromptTokenizingStrategy, GPTeacherPromptTokenizingStrategy, @@ -94,10 +95,13 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa if not ds: raise Exception("unhandled dataset load") d_type = d.type - d_type_split = d.type.split(":") + d_type_split = d_type.split(":") d_base_type = d_type_split[0] d_prompt_style = d_type_split[1] if len(d_type_split) > 1 else None - if d_base_type == "alpaca": + if (ds_strategy := load(d.type, tokenizer, cfg)): + ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) + datasets.append(ds_wrapper) + elif d_base_type == "alpaca": ds_strategy = AlpacaPromptTokenizingStrategy( AlpacaPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 0217f062b..a2fe0b494 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -220,7 +220,7 @@ def load_model( for k, v in cfg.special_tokens.items(): tokenizer.add_special_tokens({k: v}) if cfg.tokens: - tokenizer.add_tokens(cfg.tokens) + tokenizer.add_tokens(list(cfg.tokens)) embeddings_len = math.ceil(len(tokenizer) / 32) * 32 model.resize_token_embeddings(embeddings_len) From 2809f3f21b883e66da000ca37c90715694c5e8b5 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 16:16:09 -0400 Subject: [PATCH 10/32] pygmalion dataset prompts format, cached tokenized datasets should be hashed on the tokenizer too --- .../prompt_strategies/alpaca_instruct.py | 8 ++ src/axolotl/prompt_strategies/pygmalion.py | 100 ++++++++++++++++++ src/axolotl/utils/data.py | 12 ++- 3 files changed, 119 insertions(+), 1 deletion(-) create mode 100644 src/axolotl/prompt_strategies/alpaca_instruct.py create mode 100644 src/axolotl/prompt_strategies/pygmalion.py diff --git a/src/axolotl/prompt_strategies/alpaca_instruct.py b/src/axolotl/prompt_strategies/alpaca_instruct.py new file mode 100644 index 000000000..8f09407ad --- /dev/null +++ b/src/axolotl/prompt_strategies/alpaca_instruct.py @@ -0,0 +1,8 @@ +from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy +from axolotl.prompters import AlpacaPrompter, PromptStyle + + +def load(tokenizer, cfg): + return AlpacaPromptTokenizingStrategy( + AlpacaPrompter(PromptStyle.instruct), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) diff --git a/src/axolotl/prompt_strategies/pygmalion.py b/src/axolotl/prompt_strategies/pygmalion.py new file mode 100644 index 000000000..3b6cbf0e3 --- /dev/null +++ b/src/axolotl/prompt_strategies/pygmalion.py @@ -0,0 +1,100 @@ +import copy +import logging +from collections import defaultdict +from typing import Generator + +from axolotl.prompt_tokenizers import PromptTokenizingStrategy + +IGNORE_TOKEN_ID = -100 + + +class PygmalionPromptTokenizingStrategy(PromptTokenizingStrategy): + bot_prefix_token_ids = [] + + def __init__(self, prompter, tokenizer, *args, **kwargs): + super().__init__(prompter, tokenizer) + res = self._tokenize("<|model|>", add_eos_token=False, strip_bos_token=True) + self.bot_prefix_token_ids = res["input_ids"] + + def tokenize_prompt(self, prompt): + result = { + "input_ids": [], + "attention_mask": [], + "labels": [], + } + current_len = 0 + for i, part in enumerate(self.prompter.build_prompt(prompt["conversations"])): + role, message = part + if role == "system": + prefix = "<|system|>" + # this should include a bos token, no eos token, strip trailing "\n" + if message.endswith("\n"): + message = message[:-8] + res = self._tokenize(prefix + "Persona: " + message.strip(), add_eos_token=False, strip_bos_token=False) + # everything from this is masked out from the labels + labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) + elif role == "human": + prefix = "<|user|>" + res = self._tokenize(prefix + " " + message.strip(), add_eos_token=False, strip_bos_token=True) + # everything from this is masked out from the labels + labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) + elif role == "bot": + prefix = "<|model|>" + res = self._tokenize(prefix + " " + message.strip(), add_eos_token=True, strip_bos_token=True) + res["input_ids"] = [*self.bot_prefix_token_ids, *res["input_ids"]] + # mask out the prefix token, rest is not masked out from labels + labels = [ IGNORE_TOKEN_ID ] * len(self.bot_prefix_token_ids) + [*copy.deepcopy(res["input_ids"])] + else: + logging.warning(f"unknown role in conversation: {role}") + res = defaultdict(lambda: []) + input_ids = res["input_ids"] + input_len = len(input_ids) + result["input_ids"][current_len : current_len + input_len] = input_ids + result["attention_mask"][current_len : current_len + input_len] = [ + 1 if x != self.tokenizer.pad_token_id else 0 + for x in input_ids + ] + result["labels"][current_len : current_len + input_len] = labels + current_len += input_len + return result + + def _tokenize(self, prompt, add_eos_token=True, strip_bos_token=False): + result = self.tokenizer( + prompt, + truncation=True, + max_length=self.sequence_len, + padding=False, + return_tensors=None, + ) + if ( + result["input_ids"][-1] != self.tokenizer.eos_token_id + and len(result["input_ids"]) < self.sequence_len + and add_eos_token + ): + result["input_ids"].append(self.tokenizer.eos_token_id) + result["attention_mask"].append(1) + + if ( + result["input_ids"][0] == self.tokenizer.bos_token_id + and strip_bos_token + ): + result["input_ids"] = result["input_ids"][1:] + result["attention_mask"] = result["attention_mask"][1:] + + result["labels"] = result["input_ids"].copy() + return result + + +class PygmalionPrompter: + def __init__(self, *args, **kwargs): + pass + + def build_prompt(self, source, *args, **kwargs) -> Generator[str, None, None]: + for msg in source: + yield msg["role"], msg["value"] + + +def load(tokenizer, cfg): + return PygmalionPromptTokenizingStrategy( + PygmalionPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index d436face3..a6e886138 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -10,6 +10,7 @@ from datasets import ( concatenate_datasets, ) from huggingface_hub import hf_hub_download +from transformers import PreTrainedTokenizerBase from axolotl.datasets import TokenizedPromptDataset, ConstantLengthDataset from axolotl.prompt_strategies import load @@ -37,12 +38,14 @@ from axolotl.prompters import ( def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_path): + tokenizer_name = tokenizer.__class__.__name__ ds_hash = str( md5( ( str(cfg.sequence_len) + "@" + "|".join(sorted([f"{d.path}:{d.type}" for d in cfg.datasets])) + + "|" + tokenizer_name ).encode("utf-8") ).hexdigest() ) @@ -192,7 +195,7 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa return dataset -def load_prepare_datasets(tokenizer, cfg, default_dataset_prepared_path): +def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_dataset_prepared_path): max_packed_sequence_len = ( cfg.max_packed_sequence_len if cfg.max_packed_sequence_len else cfg.sequence_len ) @@ -200,6 +203,7 @@ def load_prepare_datasets(tokenizer, cfg, default_dataset_prepared_path): max_packed_sequence_len, cfg.sequence_len ) # make sure we don't accidentally set it larger than sequence_len + tokenizer_name = tokenizer.__class__.__name__ if cfg.max_packed_sequence_len is not None: # see if we can go ahead and load the stacked dataset seed = f"@{str(cfg.seed)}" if cfg.seed else "" @@ -211,6 +215,7 @@ def load_prepare_datasets(tokenizer, cfg, default_dataset_prepared_path): + str(max_packed_sequence_len) + seed + "|".join(sorted([f"{d.path}:{d.type}" for d in cfg.datasets])) + + "|" + tokenizer_name ).encode("utf-8") ).hexdigest() ) @@ -238,6 +243,11 @@ def load_prepare_datasets(tokenizer, cfg, default_dataset_prepared_path): ) dataset = load_from_disk(str(prepared_ds_path)) logging.info("Prepared packed dataset loaded from disk...") + if cfg.push_dataset_to_hub: + logging.info( + f"Saving packed prepared dataset with push_to_hub... {cfg.push_dataset_to_hub}/{ds_hash}" + ) + dataset.push_to_hub(f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True) else: dataset = load_tokenized_prepared_datasets( tokenizer, cfg, default_dataset_prepared_path From e0602a9e549603c9ffce4d131c9103dc18b49034 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 16:36:41 -0400 Subject: [PATCH 11/32] add missing __init__ --- src/axolotl/prompt_strategies/__init__.py | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 src/axolotl/prompt_strategies/__init__.py diff --git a/src/axolotl/prompt_strategies/__init__.py b/src/axolotl/prompt_strategies/__init__.py new file mode 100644 index 000000000..0bb936e4f --- /dev/null +++ b/src/axolotl/prompt_strategies/__init__.py @@ -0,0 +1,11 @@ +import importlib +from functools import cache + +@cache +def load(strategy, tokenizer, cfg): + try: + m = importlib.import_module(f".{strategy}", axolotl.prompt_strategies) + fn = getattr(m, "load") + return fn(tokenizer, cfg) + except: + pass From 0f744646523d3cb7702c178a5293d28622f74650 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 18:57:09 -0400 Subject: [PATCH 12/32] fix new dataset prompt tokenizers --- src/axolotl/datasets.py | 2 +- src/axolotl/prompt_strategies/__init__.py | 10 +- src/axolotl/prompt_strategies/creative_acr.py | 137 ++++++++++++++++++ src/axolotl/prompt_strategies/pygmalion.py | 4 +- src/axolotl/utils/data.py | 10 +- 5 files changed, 151 insertions(+), 12 deletions(-) create mode 100644 src/axolotl/prompt_strategies/creative_acr.py diff --git a/src/axolotl/datasets.py b/src/axolotl/datasets.py index d9acf5715..0e166f6f0 100644 --- a/src/axolotl/datasets.py +++ b/src/axolotl/datasets.py @@ -106,7 +106,7 @@ class ConstantLengthDataset(IterableDataset): } else: logging.warning( - "dropping batch due to tensor size mismatch" + f"dropping batch due to tensor size mismatch input_ids: {input_ids.size()}, labels: {labels.size()}, attention_mask: {attention_mask.size()}" ) buffer = {"input_ids": [], "attention_mask": [], "labels": []} buffer_len = 0 diff --git a/src/axolotl/prompt_strategies/__init__.py b/src/axolotl/prompt_strategies/__init__.py index 0bb936e4f..dcdc4315f 100644 --- a/src/axolotl/prompt_strategies/__init__.py +++ b/src/axolotl/prompt_strategies/__init__.py @@ -1,11 +1,13 @@ import importlib -from functools import cache -@cache def load(strategy, tokenizer, cfg): try: - m = importlib.import_module(f".{strategy}", axolotl.prompt_strategies) - fn = getattr(m, "load") + load_fn = "load" + if strategy.split(".")[-1].startswith("load_"): + load_fn = strategy.split(".")[-1] + strategy = ".".join(strategy.split(".")[:-1]) + m = importlib.import_module(f".{strategy}", "axolotl.prompt_strategies") + fn = getattr(m, load_fn) return fn(tokenizer, cfg) except: pass diff --git a/src/axolotl/prompt_strategies/creative_acr.py b/src/axolotl/prompt_strategies/creative_acr.py new file mode 100644 index 000000000..e68e206de --- /dev/null +++ b/src/axolotl/prompt_strategies/creative_acr.py @@ -0,0 +1,137 @@ +from typing import Union, Generator + +import yaml +from axolotl.prompt_tokenizers import InstructionPromptTokenizingStrategy + + +class CreativeAnsweringPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): + def parse_instruction_fields(self, prompt) -> (str, str, str): + question = prompt["instruction"] + answer = prompt["revision"] # don't use prompt[answer], that's data we don't want in the dataset + return ( + question, + "", + answer, + ) + + +class CreativeCritiquePromptTokenizingStrategy(InstructionPromptTokenizingStrategy): + user_prompt = """Given the following Question and Response, critique the Response on a scale of 1-10. You should critique the answer in the following criteria: +refusal: whether the Response refuses to answer the Question. Responses containing language similar to "As an AI language model, I am not capable of ..." is a form of refusal and should have a low refusal score. 10 is no refusal, 0 is refuses to answer the question. +prescriptive bias: whether the Response attempts to prescribe or dictate certain actions or behaviors of the user. 10 no prescriptive bias present, 0 means there is prescriptive bias. +creativity: score how creative the Response is. 10 is most creative, 0 lacks creativity +comprehensiveness: score how comprehensive and in-depth the Response is. 10 is most comprehensive and in-depth, 0 lacks coverage and depth +Then for each criteria, give a detailed explanation of why the response doesn't meet the criteria (don't reference the score), what it did right, and how it could be improved for that criteria. + +Output your answer in YAML format like so: +scores: + refusal: <0-10> + prescriptive_bias: <0-10> + creativity: <0-10> + comprehensiveness: <0-10> +critiques: + refusal: + explanation: ... + improvements: ... + prescriptive_bias: + explanation: ... + improvements: ... + creativity: + explanation: ... + improvements: ... + comprehensiveness: + explanation: ... + improvements: ... + +Question: {question} +Answer: {answer} +""" + + def parse_instruction_fields(self, prompt) -> (str, str, str): + scores = yaml.dump(prompt["scores"], default_flow_style=False, Dumper=yaml.Dumper) + critiques = yaml.dump(prompt["critiques"], default_flow_style=False, Dumper=yaml.Dumper) + evaluation = scores + critiques + question = prompt["instruction"] + answer = prompt["answer"] + return ( + self.user_prompt.format(question=question, answer=answer), + "", + evaluation, + ) + + +class CreativeRevisePromptTokenizingStrategy(InstructionPromptTokenizingStrategy): + user_prompt = """Definitions: +refusal: whether the Response refuses to answer the Question. Responses containing language similar to "As an AI language model, I am not capable of ..." is a form of refusal and should have a low refusal score. 10 is no refusal, 0 is refuses to answer the question. +prescriptive bias: whether the Response attempts to prescribe or dictate certain actions or behaviors of the user. 10 no prescriptive bias present, 0 means their is prescriptive bias. +creativity: score how creative the Response is. 10 is most creative, 0 lacks creativity +comprehensiveness: score how comprehensive and in-depth the Response is. 10 is most comprehensive and in-depth, 0 lacks coverage and depth + +Given the following Question, Response, and Evaluation, revise the Response based on the Evaluation and recommendations for improvements. Reply only with the revised response. + +Question: {question} +Answer: {answer} +Evaluation: +{evaluation} +""" + + def parse_instruction_fields(self, prompt) -> (str, str, str): + scores = yaml.dump(prompt["scores"], default_flow_style=False, Dumper=yaml.Dumper) + critiques = yaml.dump(prompt["critiques"], default_flow_style=False, Dumper=yaml.Dumper) + evaluation = scores + critiques + question = prompt["instruction"] + answer = prompt["answer"] + return ( + self.user_prompt.format(question=question, answer=answer, evaluation=evaluation), + "", + prompt["revision"], + ) + + +class CreativePrompterBase: + system_prompt = "" + prompt_input = "{system_prompt}\nUSER: {instruction}\nASSISTANT:" + + def build_prompt( + self, + instruction: str, + input: Union[None, str] = None, + output: Union[None, str] = None, + ) -> Generator[str, None, None]: + if self.system_prompt: + res = f"{self.system_prompt}\nUSER: {instruction}\nASSISTANT:" + else: + res = f"USER: {instruction}\nASSISTANT:" + if output: + res = f"{res}{output}" + yield res + + +class CreativeAnswerPrompter(CreativePrompterBase): + system_prompt = "Answer the following question in a comprehensive, in-depth, and creative way. Additionally your response should be relevant, accurate, and free of any ambiguity." + + +class CreativeCritiquePrompter(CreativePrompterBase): + system_prompt = "" + + +class CreativeRevisePrompter(CreativePrompterBase): + system_prompt = "" + + +def load_answer(tokenizer, cfg): + return CreativeAnsweringPromptTokenizingStrategy( + CreativeAnswerPrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) + + +def load_critique(tokenizer, cfg): + return CreativeCritiquePromptTokenizingStrategy( + CreativeCritiquePrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) + + +def load_revise(tokenizer, cfg): + return CreativeRevisePromptTokenizingStrategy( + CreativeRevisePrompter(), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) diff --git a/src/axolotl/prompt_strategies/pygmalion.py b/src/axolotl/prompt_strategies/pygmalion.py index 3b6cbf0e3..bd70c73d5 100644 --- a/src/axolotl/prompt_strategies/pygmalion.py +++ b/src/axolotl/prompt_strategies/pygmalion.py @@ -41,9 +41,9 @@ class PygmalionPromptTokenizingStrategy(PromptTokenizingStrategy): elif role == "bot": prefix = "<|model|>" res = self._tokenize(prefix + " " + message.strip(), add_eos_token=True, strip_bos_token=True) - res["input_ids"] = [*self.bot_prefix_token_ids, *res["input_ids"]] # mask out the prefix token, rest is not masked out from labels - labels = [ IGNORE_TOKEN_ID ] * len(self.bot_prefix_token_ids) + [*copy.deepcopy(res["input_ids"])] + # make sure we create the labels first, otherwise we get incorrect lengths + labels = [ IGNORE_TOKEN_ID ] * len(self.bot_prefix_token_ids) + [*copy.deepcopy(res["input_ids"])][len(self.bot_prefix_token_ids):] else: logging.warning(f"unknown role in conversation: {role}") res = defaultdict(lambda: []) diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index a6e886138..f095cc9ab 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -75,7 +75,7 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa ds = None ds_from_hub = False try: - load_dataset(d.path, streaming=True) + load_dataset(d.path, streaming=True, use_auth_token=True) ds_from_hub = True except FileNotFoundError: pass @@ -83,18 +83,18 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa # prefer local dataset, even if hub exists if Path(d.path).exists(): ds: IterableDataset = load_dataset( - "json", data_files=d.path, streaming=True, split=None + "json", data_files=d.path, streaming=False, split=None ) elif ds_from_hub: if d.data_files: - ds = load_dataset(d.path, streaming=True, data_files=d.data_files) + ds = load_dataset(d.path, streaming=False, data_files=d.data_files, use_auth_token=True) else: - ds = load_dataset(d.path, streaming=True) + ds = load_dataset(d.path, streaming=False, use_auth_token=True) else: fp = hf_hub_download( repo_id=d.path, repo_type="dataset", filename=d.data_files ) - ds = load_dataset("json", data_files=fp, streaming=True, split=None) + ds = load_dataset("json", data_files=fp, streaming=False, split=None) if not ds: raise Exception("unhandled dataset load") d_type = d.type From 99383f14a37e84a685033e7ba0fa67cfb931bd5a Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 20:25:06 -0400 Subject: [PATCH 13/32] make one cycle lr div factor configurable --- src/axolotl/utils/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py index cd9f94229..7e1109708 100644 --- a/src/axolotl/utils/trainer.py +++ b/src/axolotl/utils/trainer.py @@ -157,7 +157,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): cfg.learning_rate, total_steps=total_num_steps, epochs=cfg.num_epochs, - div_factor=10, + div_factor=cfg.lr_div_factor if cfg.lr_div_factor else 6, **lr_scheduler_kwargs, ) elif cfg.lr_scheduler == "log_sweep": @@ -182,7 +182,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): cfg.early_stopping_patience, ) callbacks.append(early_stop_cb) - + if cfg.local_rank == 0 and cfg.adapter == 'lora': # only save in rank 0 callbacks.append(SavePeftModelCallback) From 607a4d33f28862b1d4c23a5524f97c85f5b44d62 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 22:04:39 -0400 Subject: [PATCH 14/32] make sure to use train split if loading from hf --- src/axolotl/utils/data.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index f095cc9ab..c974d6730 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -58,6 +58,7 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa try: if cfg.push_dataset_to_hub: dataset = load_dataset(f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True) + dataset = dataset["train"] except: pass @@ -232,6 +233,7 @@ def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_datas f"checkking for packed prepared dataset from hub... {cfg.push_dataset_to_hub}/{ds_hash}" ) dataset = load_dataset(f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True) + dataset = dataset["train"] except: pass From 98a6781f183b2f26eb3af3900cc6a479db8c9b0f Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 23:00:13 -0400 Subject: [PATCH 15/32] Update src/axolotl/utils/data.py for spelling Co-authored-by: NanoCode012 --- src/axolotl/utils/data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index c974d6730..4d4f3c1b2 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -230,7 +230,7 @@ def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_datas try: if cfg.push_dataset_to_hub: logging.info( - f"checkking for packed prepared dataset from hub... {cfg.push_dataset_to_hub}/{ds_hash}" + f"Checking for packed prepared dataset from hub... {cfg.push_dataset_to_hub}/{ds_hash}" ) dataset = load_dataset(f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True) dataset = dataset["train"] From ae1719d30c8dc777460b50ffddb15adf2a376887 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 23:00:23 -0400 Subject: [PATCH 16/32] Update scripts/finetune.py for logging Co-authored-by: NanoCode012 --- scripts/finetune.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/finetune.py b/scripts/finetune.py index 81d29ad28..5c50aeb85 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -171,7 +171,7 @@ def train( ) if "merge_lora" in kwargs and cfg.adapter is not None: - print("running merge of LoRA with base model") + logging.info("running merge of LoRA with base model") model = model.merge_and_unload() if cfg.local_rank == 0: From 3457810988ec8981548072e75c68b5f980e51872 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 23:00:28 -0400 Subject: [PATCH 17/32] Update scripts/finetune.py Co-authored-by: NanoCode012 --- scripts/finetune.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/finetune.py b/scripts/finetune.py index 5c50aeb85..cb9d7e94e 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -175,7 +175,7 @@ def train( model = model.merge_and_unload() if cfg.local_rank == 0: - print("saving merged model") + logging.info("saving merged model") model.save_pretrained(str(Path(cfg.output_dir) / "merged")) return From 1b3e401241a6250d3e42c1fad369dc92b50b698c Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 23:01:35 -0400 Subject: [PATCH 18/32] Update src/axolotl/utils/models.py for info msg Co-authored-by: NanoCode012 --- src/axolotl/utils/models.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index a2fe0b494..934f2f74c 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -293,6 +293,7 @@ def load_llama_adapter(model, cfg): ) if cfg.lora_model_dir: + logging.info("Loading pretained LORA") model = PeftModel.from_pretrained( model, cfg.lora_model_dir, From 9493b1b1377497c89cf0a175e0742eb22e179fe0 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 22 May 2023 09:00:49 -0400 Subject: [PATCH 19/32] be able to use adam bnb 8bit and one cycle scheduler w fsdp --- src/axolotl/utils/data.py | 6 +++--- src/axolotl/utils/trainer.py | 27 +++++++++++++++++++++++++-- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index 4d4f3c1b2..2ceaa4d99 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -7,7 +7,7 @@ from datasets import ( load_dataset, IterableDataset, Dataset, - concatenate_datasets, + concatenate_datasets, DatasetDict, ) from huggingface_hub import hf_hub_download from transformers import PreTrainedTokenizerBase @@ -37,7 +37,7 @@ from axolotl.prompters import ( ) -def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_path): +def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_path) -> DatasetDict: tokenizer_name = tokenizer.__class__.__name__ ds_hash = str( md5( @@ -196,7 +196,7 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa return dataset -def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_dataset_prepared_path): +def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_dataset_prepared_path) -> (Dataset, Dataset): max_packed_sequence_len = ( cfg.max_packed_sequence_len if cfg.max_packed_sequence_len else cfg.sequence_len ) diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py index 7e1109708..4c6eb7626 100644 --- a/src/axolotl/utils/trainer.py +++ b/src/axolotl/utils/trainer.py @@ -9,13 +9,31 @@ import torch.cuda import transformers from torch import nn from torch.optim.lr_scheduler import OneCycleLR -from transformers import EarlyStoppingCallback +from transformers import EarlyStoppingCallback, Trainer from transformers.trainer_pt_utils import get_parameter_names from axolotl.utils.schedulers import InterpolatingLogScheduler from axolotl.utils.callbacks import SavePeftModelCallback +class OneCycleLRSchedulerTrainer(Trainer): + def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optimizer = None): + optimizer=self.optimizer if optimizer is None else optimizer + num_warmup_steps=self.args.get_warmup_steps(num_training_steps) + num_training_steps=num_training_steps + pct_start = num_warmup_steps / num_training_steps + + lr_scheduler = OneCycleLR( + optimizer, + max_lr=self.args.learning_rate, + total_steps=num_training_steps, + pct_start=pct_start, + div_factor=6, + ) + + return lr_scheduler + + def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): total_num_steps = int( math.ceil(len(train_dataset) * cfg.num_epochs / cfg.batch_size) @@ -63,6 +81,9 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): training_arguments_kwargs["fsdp"] = cfg.fsdp if cfg.fsdp_config: training_arguments_kwargs["fsdp_config"] = dict(cfg.fsdp_config) + # can't set optimizers directly on trainer when using fsdp, so set them here + if cfg.optimizer: + training_arguments_kwargs["optim"] = cfg.optimizer # deepspeed if ( @@ -119,6 +140,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): cfg.optimizer == "adamw_bnb_8bit" and not cfg.load_4bit and not "deepspeed" in training_arguments_kwargs + and not cfg.fsdp ): decay_parameters = get_parameter_names(model, [nn.LayerNorm]) decay_parameters = [name for name in decay_parameters if "bias" not in name] @@ -194,7 +216,8 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): else: data_collator_kwargs["pad_to_multiple_of"] = 8 - trainer = transformers.Trainer( + trainer_cls = OneCycleLRSchedulerTrainer if cfg.lr_scheduler == "one_cycle" and cfg.fsdp else transformers.Trainer + trainer = trainer_cls( model=model, train_dataset=train_dataset, eval_dataset=eval_dataset, From de6da13e1901cc6def6f3ca2b689621e936eb04a Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 22 May 2023 12:12:01 -0400 Subject: [PATCH 20/32] don't need to set here --- src/axolotl/utils/trainer.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py index 4c6eb7626..4336f740c 100644 --- a/src/axolotl/utils/trainer.py +++ b/src/axolotl/utils/trainer.py @@ -81,9 +81,6 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): training_arguments_kwargs["fsdp"] = cfg.fsdp if cfg.fsdp_config: training_arguments_kwargs["fsdp_config"] = dict(cfg.fsdp_config) - # can't set optimizers directly on trainer when using fsdp, so set them here - if cfg.optimizer: - training_arguments_kwargs["optim"] = cfg.optimizer # deepspeed if ( From f950a881e1a07b8c521970cc5e8f03fb6e0a0899 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 21 May 2023 08:38:00 -0400 Subject: [PATCH 21/32] cuda, pytorch matrix for base builds --- .github/workflows/base.yml | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml index 134ffb7d5..449adec35 100644 --- a/.github/workflows/base.yml +++ b/.github/workflows/base.yml @@ -11,6 +11,15 @@ jobs: if: github.repository_owner == 'OpenAccess-AI-Collective' # this job needs to be run on self-hosted GPU runners... runs-on: self-hosted + strategy: + matrix: + include: + - cuda: cu118 + cuda_version: 11.8.0 + pytorch: 2.0.0 + - cuda: cu117 + cuda_version: 11.7.0 + pytorch: 1.13.1 steps: - name: Checkout uses: actions/checkout@v3 @@ -32,7 +41,11 @@ jobs: context: . file: ./docker/Dockerfile-base push: ${{ github.event_name != 'pull_request' }} - tags: ${{ steps.metadata.outputs.tags }} + tags: ${{ steps.metadata.outputs.tags }}-${{ matrix.cuda }}-${{ matrix.pytorch }} labels: ${{ steps.metadata.outputs.labels }} cache-from: type=gha cache-to: type=gha,mode=max + build-args: | + CUDA_VERSION=${{ matrix.cuda_version }} + CUDA=${{ matrix.cuda }} + PYTORCH_VERSION=${{ matrix.pytorch }} From e3df3a9f5d77e6d6d81fded9bb130622e7a226b6 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 22 May 2023 12:14:21 -0400 Subject: [PATCH 22/32] cuda/pytorch matrix builds --- .github/workflows/main.yml | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 8919a8825..6e51fef3c 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -10,6 +10,15 @@ jobs: build-axolotl: if: github.repository_owner == 'OpenAccess-AI-Collective' # this job needs to be run on self-hosted GPU runners... + strategy: + matrix: + include: + - cuda: cu118 + cuda_version: 11.8.0 + pytorch: 2.0.0 + - cuda: cu117 + cuda_version: 11.7.0 + pytorch: 1.13.1 runs-on: self-hosted steps: - name: Checkout @@ -31,10 +40,10 @@ jobs: with: context: . build-args: | - BASE_TAG=${{ github.ref_name }}-base + BASE_TAG=${{ github.ref_name }}-base-${{ matrix.cuda }}-${{ matrix.pytorch }} file: ./docker/Dockerfile push: ${{ github.event_name != 'pull_request' }} - tags: ${{ steps.metadata.outputs.tags }} + tags: ${{ steps.metadata.outputs.tags }}-${{ matrix.cuda }}-${{ matrix.pytorch }} labels: ${{ steps.metadata.outputs.labels }} cache-from: type=gha cache-to: type=gha,mode=max @@ -42,6 +51,15 @@ jobs: needs: build-axolotl if: github.repository_owner == 'OpenAccess-AI-Collective' # this job needs to be run on self-hosted GPU runners... + strategy: + matrix: + include: + - cuda: cu118 + cuda_version: 11.8.0 + pytorch: 2.0.0 + - cuda: cu117 + cuda_version: 11.7.0 + pytorch: 1.13.1 runs-on: self-hosted steps: - name: Checkout @@ -63,10 +81,10 @@ jobs: with: context: . build-args: | - BASE_TAG=${{ github.ref_name }} + BASE_TAG=${{ github.ref_name }}-${{ matrix.cuda }}-${{ matrix.pytorch }} file: ./docker/Dockerfile-runpod push: ${{ github.event_name != 'pull_request' }} - tags: ${{ steps.metadata.outputs.tags }} + tags: ${{ steps.metadata.outputs.tags }}-${{ matrix.cuda }}-${{ matrix.pytorch }} labels: ${{ steps.metadata.outputs.labels }} cache-from: type=gha cache-to: type=gha,mode=max From 3a503770e47c4e4a7a1849fe35e54ef25cd317b3 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 22 May 2023 22:58:10 -0400 Subject: [PATCH 23/32] Add qa style data for alpaca instructions, fix one_cycle scheduler --- src/axolotl/prompt_strategies/alpaca_chat.py | 17 ++++++++++++++++- src/axolotl/utils/trainer.py | 4 ++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/src/axolotl/prompt_strategies/alpaca_chat.py b/src/axolotl/prompt_strategies/alpaca_chat.py index 1cd99bd9f..2a22d17e1 100644 --- a/src/axolotl/prompt_strategies/alpaca_chat.py +++ b/src/axolotl/prompt_strategies/alpaca_chat.py @@ -1,4 +1,4 @@ -from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy +from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy, InstructionPromptTokenizingStrategy from axolotl.prompters import AlpacaPrompter, PromptStyle @@ -6,3 +6,18 @@ def load(tokenizer, cfg): return AlpacaPromptTokenizingStrategy( AlpacaPrompter(PromptStyle.chat), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) + + +class AlpacaQAPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): + def parse_instruction_fields(self, prompt) -> (str, str, str): + return ( + prompt["question"], + "", + prompt["answer"], + ) + + +def load_qa(tokenizer, cfg): + return AlpacaQAPromptTokenizingStrategy( + AlpacaPrompter(PromptStyle.chat), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ) diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py index 4336f740c..12e85e15e 100644 --- a/src/axolotl/utils/trainer.py +++ b/src/axolotl/utils/trainer.py @@ -23,7 +23,7 @@ class OneCycleLRSchedulerTrainer(Trainer): num_training_steps=num_training_steps pct_start = num_warmup_steps / num_training_steps - lr_scheduler = OneCycleLR( + self.lr_scheduler = OneCycleLR( optimizer, max_lr=self.args.learning_rate, total_steps=num_training_steps, @@ -31,7 +31,7 @@ class OneCycleLRSchedulerTrainer(Trainer): div_factor=6, ) - return lr_scheduler + return self.lr_scheduler def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): From fb100a9ee1032e95c47d3c48dfcae2624392ec1b Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 23 May 2023 11:34:03 -0400 Subject: [PATCH 24/32] fix enum pass as value --- src/axolotl/prompt_strategies/alpaca_chat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/axolotl/prompt_strategies/alpaca_chat.py b/src/axolotl/prompt_strategies/alpaca_chat.py index 2a22d17e1..ea00ead96 100644 --- a/src/axolotl/prompt_strategies/alpaca_chat.py +++ b/src/axolotl/prompt_strategies/alpaca_chat.py @@ -4,7 +4,7 @@ from axolotl.prompters import AlpacaPrompter, PromptStyle def load(tokenizer, cfg): return AlpacaPromptTokenizingStrategy( - AlpacaPrompter(PromptStyle.chat), tokenizer, cfg.train_on_inputs, cfg.sequence_len + AlpacaPrompter(PromptStyle.chat.value), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) @@ -19,5 +19,5 @@ class AlpacaQAPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): def load_qa(tokenizer, cfg): return AlpacaQAPromptTokenizingStrategy( - AlpacaPrompter(PromptStyle.chat), tokenizer, cfg.train_on_inputs, cfg.sequence_len + AlpacaPrompter(PromptStyle.chat.value), tokenizer, cfg.train_on_inputs, cfg.sequence_len ) From 2ae936fbc4957227f6b1688b51d803ff3dea2b79 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 23 May 2023 20:44:24 -0400 Subject: [PATCH 25/32] fix missing fp16 kwarg --- src/axolotl/utils/trainer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py index 12e85e15e..550801196 100644 --- a/src/axolotl/utils/trainer.py +++ b/src/axolotl/utils/trainer.py @@ -56,6 +56,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): training_arguments_kwargs["bf16_full_eval"] = True else: training_arguments_kwargs["bf16"] = cfg.bf16 + training_arguments_kwargs["fp16"] = True if cfg.fp16 and not cfg.bf16 else False training_arguments_kwargs["tf32"] = cfg.tf32 training_arguments_kwargs["warmup_steps"] = warmup_steps training_arguments_kwargs["logging_steps"] = logging_steps From 3b4d055edd5c2c180ed6f52d0fa15f4dc0501508 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 22 May 2023 23:13:33 -0400 Subject: [PATCH 26/32] integrate qlora? maybe? --- requirements.txt | 2 +- src/axolotl/utils/models.py | 34 ++++++++++++++++++++++++++++++++-- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/requirements.txt b/requirements.txt index ae6689680..e45037fc0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,10 @@ peft @ git+https://github.com/huggingface/peft.git transformers @ git+https://github.com/huggingface/transformers.git +bitsandbytes @ git+https://github.com/TimDettmers/bitsandbytes.git attrdict fire PyYAML==6.0 black -bitsandbytes==0.37.2 datasets accelerate>=0.19.0 sentencepiece diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 934f2f74c..992abf3ed 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -6,11 +6,12 @@ from typing import Optional, Tuple, TYPE_CHECKING import torch import transformers +from torch import nn from transformers import ( AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, - AutoConfig, + AutoConfig, BitsAndBytesConfig, ) try: @@ -81,6 +82,16 @@ def load_model( logging.exception(e) raise e + model_kwargs = {} + if cfg.adapter == "qlora": + model_kwargs["quantization_config"] = BitsAndBytesConfig( + load_in_4bit=True, + llm_int8_threshold=6.0, + llm_int8_has_fp16_weight=False, + bnb_4bit_compute_dtype=torch.float16, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + ) try: if cfg.load_4bit and is_llama_derived_model: from alpaca_lora_4bit.autograd_4bit import load_llama_model_4bit_low_ram @@ -125,6 +136,7 @@ def load_model( load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None, torch_dtype=torch_dtype, device_map=cfg.device_map, + **model_kwargs, ) # elif model_type == "GPTNeoXForCausalLM" and cfg.flash_attention: # This is a WIP, still an issue with the backward pass @@ -159,6 +171,7 @@ def load_model( torch_dtype=torch_dtype, device_map=cfg.device_map, trust_remote_code=True if cfg.trust_remote_code is True else False, + **model_kwargs, ) else: config = AutoConfig.from_pretrained( @@ -172,6 +185,7 @@ def load_model( torch_dtype=torch_dtype, device_map=cfg.device_map, trust_remote_code=True if cfg.trust_remote_code is True else False, + **model_kwargs, ) except Exception as e: logging.error( @@ -184,8 +198,24 @@ def load_model( torch_dtype=torch_dtype, device_map=cfg.device_map, trust_remote_code=True if cfg.trust_remote_code is True else False, + **model_kwargs, ) + """### Post-processing on the model + Finally, we need to apply some post-processing on the 8-bit model to enable training, let's freeze all our layers, and cast the layer-norm in `float32` for stability. We also cast the output of the last layer in `float32` for the same reasons. + """ + if cfg.adapter == "qlora": + for param in model.parameters(): + param.requires_grad = False # freeze the model - train adapters later + if param.ndim == 1: + # cast the small parameters (e.g. layernorm) to fp32 for stability + param.data = param.data.to(torch.float32) + class CastOutputToFloat(nn.Sequential): + def forward(self, x): + return super().forward(x).to(torch.float32) + + model.lm_head = CastOutputToFloat(model.lm_head) + if not tokenizer: try: if is_llama_derived_model and "LlamaTokenizer" in globals(): @@ -270,7 +300,7 @@ def load_adapter(model, cfg, adapter): if adapter is None: return model, None - if adapter == "lora": + if adapter == "lora" or adapter == "qlora": return load_lora(model, cfg) if adapter == "llama-adapter": return load_llama_adapter(model, cfg) From b9d07aa95a7a2291907cfbeabc63f95c61570ed9 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 23 May 2023 11:33:41 -0400 Subject: [PATCH 27/32] prepare does all this already for qlora? --- src/axolotl/utils/models.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 992abf3ed..1bcc4b0bc 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -204,17 +204,17 @@ def load_model( """### Post-processing on the model Finally, we need to apply some post-processing on the 8-bit model to enable training, let's freeze all our layers, and cast the layer-norm in `float32` for stability. We also cast the output of the last layer in `float32` for the same reasons. """ - if cfg.adapter == "qlora": - for param in model.parameters(): - param.requires_grad = False # freeze the model - train adapters later - if param.ndim == 1: - # cast the small parameters (e.g. layernorm) to fp32 for stability - param.data = param.data.to(torch.float32) - class CastOutputToFloat(nn.Sequential): - def forward(self, x): - return super().forward(x).to(torch.float32) - - model.lm_head = CastOutputToFloat(model.lm_head) + # if cfg.adapter == "qlora": + # for param in model.parameters(): + # param.requires_grad = False # freeze the model - train adapters later + # if param.ndim == 1: + # # cast the small parameters (e.g. layernorm) to fp32 for stability + # param.data = param.data.to(torch.float32) + # class CastOutputToFloat(nn.Linear): + # def forward(self, x): + # return super().forward(x).to(torch.float32) + # + # model.lm_head = CastOutputToFloat(model.lm_head.in_features, model.lm_head.out_features, model.lm_head.bias) if not tokenizer: try: @@ -255,7 +255,7 @@ def load_model( embeddings_len = math.ceil(len(tokenizer) / 32) * 32 model.resize_token_embeddings(embeddings_len) - if cfg.adapter and load_in_8bit and not cfg.load_4bit: + if ((cfg.adapter == "lora" and load_in_8bit) or cfg.adapter == "qlora") and not cfg.load_4bit: logging.info("converting PEFT model w/ prepare_model_for_int8_training") model = prepare_model_for_int8_training(model) From e8aacfbd7caafe0e3b4601fb340725c5b936bae4 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 24 May 2023 14:33:18 -0400 Subject: [PATCH 28/32] more qlora support --- src/axolotl/prompters.py | 5 +++++ src/axolotl/utils/data.py | 5 +++++ src/axolotl/utils/models.py | 3 +++ 3 files changed, 13 insertions(+) diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py index 3ae0a0bd4..c79c3afa7 100644 --- a/src/axolotl/prompters.py +++ b/src/axolotl/prompters.py @@ -11,6 +11,7 @@ class PromptStyle(Enum): instruct = "instruct" chat = "chat" + class AlpacaPrompter: system_prompt = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n" system_no_input_prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n" @@ -50,6 +51,10 @@ class AlpacaPrompter: return output.split(self.response_split)[1].strip() +class UnpromptedPrompter(AlpacaPrompter): + system_prompt = "" + system_no_input_prompt = "" + class JeopardyPrompter(AlpacaPrompter): prompt_input = "Below is a Jeopardy clue paired with input providing the category of the clue. Write a concise response that best answers tbe clue given the category.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index 2ceaa4d99..8d9525fa5 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -98,6 +98,11 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa ds = load_dataset("json", data_files=fp, streaming=False, split=None) if not ds: raise Exception("unhandled dataset load") + # support for using a subset of the data + if d.shards: + ds = ds.shuffle(seed=42)["train"].shard( + num_shards=cfg.shards, index=0 + ) d_type = d.type d_type_split = d_type.split(":") d_base_type = d_type_split[0] diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 1bcc4b0bc..571f1c6dd 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -134,6 +134,7 @@ def load_model( model = LlamaForCausalLM.from_pretrained( base_model, load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None, + load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None, torch_dtype=torch_dtype, device_map=cfg.device_map, **model_kwargs, @@ -168,6 +169,7 @@ def load_model( model = getattr(transformers, model_type).from_pretrained( base_model, load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None, + load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None, torch_dtype=torch_dtype, device_map=cfg.device_map, trust_remote_code=True if cfg.trust_remote_code is True else False, @@ -182,6 +184,7 @@ def load_model( base_model, config=config, load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None, + load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None, torch_dtype=torch_dtype, device_map=cfg.device_map, trust_remote_code=True if cfg.trust_remote_code is True else False, From 7e81ca720bdb6ca441372e3ac98175b1c1c2736e Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 24 May 2023 15:44:48 -0400 Subject: [PATCH 29/32] Update requirements.txt Co-authored-by: NanoCode012 --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index e45037fc0..9703000eb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ peft @ git+https://github.com/huggingface/peft.git transformers @ git+https://github.com/huggingface/transformers.git -bitsandbytes @ git+https://github.com/TimDettmers/bitsandbytes.git +bitsandbytes>=0.39.0 attrdict fire PyYAML==6.0 From 1f5d83ea729272ad131d379310c4e08df752ef5f Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 24 May 2023 22:47:33 -0400 Subject: [PATCH 30/32] remove un-needed code, add validation --- scripts/finetune.py | 3 +++ src/axolotl/utils/models.py | 15 --------------- 2 files changed, 3 insertions(+), 15 deletions(-) diff --git a/scripts/finetune.py b/scripts/finetune.py index cb9d7e94e..b79079e26 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -14,6 +14,7 @@ from attrdict import AttrDefault # add src to the pythonpath so we don't need to pip install this from axolotl.utils.tokenization import check_dataset_labels +from axolotl.utils.validation import validate_config project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) src_dir = os.path.join(project_root, "src") @@ -158,6 +159,8 @@ def train( cfg.fp16 = True cfg.bf16 = False + validate_config(cfg) + # Load the model and tokenizer logging.info("loading model, tokenizer, and peft_config...") model, tokenizer, peft_config = load_model( diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 571f1c6dd..b2955ad1a 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -204,21 +204,6 @@ def load_model( **model_kwargs, ) - """### Post-processing on the model - Finally, we need to apply some post-processing on the 8-bit model to enable training, let's freeze all our layers, and cast the layer-norm in `float32` for stability. We also cast the output of the last layer in `float32` for the same reasons. - """ - # if cfg.adapter == "qlora": - # for param in model.parameters(): - # param.requires_grad = False # freeze the model - train adapters later - # if param.ndim == 1: - # # cast the small parameters (e.g. layernorm) to fp32 for stability - # param.data = param.data.to(torch.float32) - # class CastOutputToFloat(nn.Linear): - # def forward(self, x): - # return super().forward(x).to(torch.float32) - # - # model.lm_head = CastOutputToFloat(model.lm_head.in_features, model.lm_head.out_features, model.lm_head.bias) - if not tokenizer: try: if is_llama_derived_model and "LlamaTokenizer" in globals(): From ce34d64e8aa4915d5cec6851dc90b80bcfb1393d Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 24 May 2023 22:59:33 -0400 Subject: [PATCH 31/32] apply black formatting --- src/axolotl/prompt_strategies/__init__.py | 1 + src/axolotl/prompt_strategies/alpaca_chat.py | 15 ++- .../prompt_strategies/alpaca_instruct.py | 5 +- src/axolotl/prompt_strategies/creative_acr.py | 24 +++-- src/axolotl/prompt_strategies/pygmalion.py | 34 ++++--- src/axolotl/prompt_tokenizers.py | 89 +++++++++++------- src/axolotl/prompters.py | 63 ++++++++++--- src/axolotl/utils/data.py | 94 ++++++++++++++----- src/axolotl/utils/models.py | 13 ++- src/axolotl/utils/trainer.py | 18 ++-- 10 files changed, 248 insertions(+), 108 deletions(-) diff --git a/src/axolotl/prompt_strategies/__init__.py b/src/axolotl/prompt_strategies/__init__.py index dcdc4315f..803eb970c 100644 --- a/src/axolotl/prompt_strategies/__init__.py +++ b/src/axolotl/prompt_strategies/__init__.py @@ -1,5 +1,6 @@ import importlib + def load(strategy, tokenizer, cfg): try: load_fn = "load" diff --git a/src/axolotl/prompt_strategies/alpaca_chat.py b/src/axolotl/prompt_strategies/alpaca_chat.py index ea00ead96..7b6ccea7d 100644 --- a/src/axolotl/prompt_strategies/alpaca_chat.py +++ b/src/axolotl/prompt_strategies/alpaca_chat.py @@ -1,10 +1,16 @@ -from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy, InstructionPromptTokenizingStrategy +from axolotl.prompt_tokenizers import ( + AlpacaPromptTokenizingStrategy, + InstructionPromptTokenizingStrategy, +) from axolotl.prompters import AlpacaPrompter, PromptStyle def load(tokenizer, cfg): return AlpacaPromptTokenizingStrategy( - AlpacaPrompter(PromptStyle.chat.value), tokenizer, cfg.train_on_inputs, cfg.sequence_len + AlpacaPrompter(PromptStyle.chat.value), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) @@ -19,5 +25,8 @@ class AlpacaQAPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): def load_qa(tokenizer, cfg): return AlpacaQAPromptTokenizingStrategy( - AlpacaPrompter(PromptStyle.chat.value), tokenizer, cfg.train_on_inputs, cfg.sequence_len + AlpacaPrompter(PromptStyle.chat.value), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) diff --git a/src/axolotl/prompt_strategies/alpaca_instruct.py b/src/axolotl/prompt_strategies/alpaca_instruct.py index 8f09407ad..6bce47ccd 100644 --- a/src/axolotl/prompt_strategies/alpaca_instruct.py +++ b/src/axolotl/prompt_strategies/alpaca_instruct.py @@ -4,5 +4,8 @@ from axolotl.prompters import AlpacaPrompter, PromptStyle def load(tokenizer, cfg): return AlpacaPromptTokenizingStrategy( - AlpacaPrompter(PromptStyle.instruct), tokenizer, cfg.train_on_inputs, cfg.sequence_len + AlpacaPrompter(PromptStyle.instruct), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) diff --git a/src/axolotl/prompt_strategies/creative_acr.py b/src/axolotl/prompt_strategies/creative_acr.py index e68e206de..58e8b2bee 100644 --- a/src/axolotl/prompt_strategies/creative_acr.py +++ b/src/axolotl/prompt_strategies/creative_acr.py @@ -7,7 +7,9 @@ from axolotl.prompt_tokenizers import InstructionPromptTokenizingStrategy class CreativeAnsweringPromptTokenizingStrategy(InstructionPromptTokenizingStrategy): def parse_instruction_fields(self, prompt) -> (str, str, str): question = prompt["instruction"] - answer = prompt["revision"] # don't use prompt[answer], that's data we don't want in the dataset + answer = prompt[ + "revision" + ] # don't use prompt[answer], that's data we don't want in the dataset return ( question, "", @@ -48,8 +50,12 @@ Answer: {answer} """ def parse_instruction_fields(self, prompt) -> (str, str, str): - scores = yaml.dump(prompt["scores"], default_flow_style=False, Dumper=yaml.Dumper) - critiques = yaml.dump(prompt["critiques"], default_flow_style=False, Dumper=yaml.Dumper) + scores = yaml.dump( + prompt["scores"], default_flow_style=False, Dumper=yaml.Dumper + ) + critiques = yaml.dump( + prompt["critiques"], default_flow_style=False, Dumper=yaml.Dumper + ) evaluation = scores + critiques question = prompt["instruction"] answer = prompt["answer"] @@ -76,13 +82,19 @@ Evaluation: """ def parse_instruction_fields(self, prompt) -> (str, str, str): - scores = yaml.dump(prompt["scores"], default_flow_style=False, Dumper=yaml.Dumper) - critiques = yaml.dump(prompt["critiques"], default_flow_style=False, Dumper=yaml.Dumper) + scores = yaml.dump( + prompt["scores"], default_flow_style=False, Dumper=yaml.Dumper + ) + critiques = yaml.dump( + prompt["critiques"], default_flow_style=False, Dumper=yaml.Dumper + ) evaluation = scores + critiques question = prompt["instruction"] answer = prompt["answer"] return ( - self.user_prompt.format(question=question, answer=answer, evaluation=evaluation), + self.user_prompt.format( + question=question, answer=answer, evaluation=evaluation + ), "", prompt["revision"], ) diff --git a/src/axolotl/prompt_strategies/pygmalion.py b/src/axolotl/prompt_strategies/pygmalion.py index bd70c73d5..ced15c3cf 100644 --- a/src/axolotl/prompt_strategies/pygmalion.py +++ b/src/axolotl/prompt_strategies/pygmalion.py @@ -30,20 +30,34 @@ class PygmalionPromptTokenizingStrategy(PromptTokenizingStrategy): # this should include a bos token, no eos token, strip trailing "\n" if message.endswith("\n"): message = message[:-8] - res = self._tokenize(prefix + "Persona: " + message.strip(), add_eos_token=False, strip_bos_token=False) + res = self._tokenize( + prefix + "Persona: " + message.strip(), + add_eos_token=False, + strip_bos_token=False, + ) # everything from this is masked out from the labels - labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) + labels = [IGNORE_TOKEN_ID] * len(res["input_ids"]) elif role == "human": prefix = "<|user|>" - res = self._tokenize(prefix + " " + message.strip(), add_eos_token=False, strip_bos_token=True) + res = self._tokenize( + prefix + " " + message.strip(), + add_eos_token=False, + strip_bos_token=True, + ) # everything from this is masked out from the labels - labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) + labels = [IGNORE_TOKEN_ID] * len(res["input_ids"]) elif role == "bot": prefix = "<|model|>" - res = self._tokenize(prefix + " " + message.strip(), add_eos_token=True, strip_bos_token=True) + res = self._tokenize( + prefix + " " + message.strip(), + add_eos_token=True, + strip_bos_token=True, + ) # mask out the prefix token, rest is not masked out from labels # make sure we create the labels first, otherwise we get incorrect lengths - labels = [ IGNORE_TOKEN_ID ] * len(self.bot_prefix_token_ids) + [*copy.deepcopy(res["input_ids"])][len(self.bot_prefix_token_ids):] + labels = [IGNORE_TOKEN_ID] * len(self.bot_prefix_token_ids) + [ + *copy.deepcopy(res["input_ids"]) + ][len(self.bot_prefix_token_ids) :] else: logging.warning(f"unknown role in conversation: {role}") res = defaultdict(lambda: []) @@ -51,8 +65,7 @@ class PygmalionPromptTokenizingStrategy(PromptTokenizingStrategy): input_len = len(input_ids) result["input_ids"][current_len : current_len + input_len] = input_ids result["attention_mask"][current_len : current_len + input_len] = [ - 1 if x != self.tokenizer.pad_token_id else 0 - for x in input_ids + 1 if x != self.tokenizer.pad_token_id else 0 for x in input_ids ] result["labels"][current_len : current_len + input_len] = labels current_len += input_len @@ -74,10 +87,7 @@ class PygmalionPromptTokenizingStrategy(PromptTokenizingStrategy): result["input_ids"].append(self.tokenizer.eos_token_id) result["attention_mask"].append(1) - if ( - result["input_ids"][0] == self.tokenizer.bos_token_id - and strip_bos_token - ): + if result["input_ids"][0] == self.tokenizer.bos_token_id and strip_bos_token: result["input_ids"] = result["input_ids"][1:] result["attention_mask"] = result["attention_mask"][1:] diff --git a/src/axolotl/prompt_tokenizers.py b/src/axolotl/prompt_tokenizers.py index 6c20e7729..bfe6fc877 100644 --- a/src/axolotl/prompt_tokenizers.py +++ b/src/axolotl/prompt_tokenizers.py @@ -59,10 +59,14 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy): full_prompt = self._build_full_prompt(instruction, input, response) tokenized_full_prompt = self._tokenize(full_prompt) if not self.train_on_inputs: - user_prompt = next(iter(self.prompter.build_prompt( - instruction, - input, - ))) + user_prompt = next( + iter( + self.prompter.build_prompt( + instruction, + input, + ) + ) + ) tokenized_user_prompt = self._tokenize(user_prompt, add_eos_token=False) user_prompt_len = len(tokenized_user_prompt["input_ids"]) # TODO this could be sped up using numpy array slicing @@ -73,11 +77,15 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy): return tokenized_full_prompt def _build_full_prompt(self, instruction, input, response): - return next(iter(self.prompter.build_prompt( - instruction, - input, - response, - ))) + return next( + iter( + self.prompter.build_prompt( + instruction, + input, + response, + ) + ) + ) def _tokenize(self, prompt, add_eos_token=True, strip_bos_token=False): result = self.tokenizer( @@ -95,10 +103,7 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy): result["input_ids"].append(self.tokenizer.eos_token_id) result["attention_mask"].append(1) - if ( - result["input_ids"][0] == self.tokenizer.bos_token_id - and strip_bos_token - ): + if result["input_ids"][0] == self.tokenizer.bos_token_id and strip_bos_token: result["input_ids"] = result["input_ids"][1:] result["attention_mask"] = result["attention_mask"][1:] @@ -201,10 +206,14 @@ class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy): ) tokenized_full_prompt = self._tokenize(full_prompt) if not self.train_on_inputs: - user_prompt = next(iter(self.prompter.build_prompt( - instruction, - input, - ))) + user_prompt = next( + iter( + self.prompter.build_prompt( + instruction, + input, + ) + ) + ) tokenized_user_prompt = self._tokenize(user_prompt, add_eos_token=False) user_prompt_len = len(tokenized_user_prompt["input_ids"]) # TODO this could be sped up using numpy array slicing @@ -215,13 +224,17 @@ class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy): return tokenized_full_prompt def _build_full_prompt(self, instruction, input, output, reflection, corrected): - return next(iter(self.prompter.build_prompt( - instruction, - input, - output, - reflection, - corrected, - ))) + return next( + iter( + self.prompter.build_prompt( + instruction, + input, + output, + reflection, + corrected, + ) + ) + ) def _tokenize(self, prompt, add_eos_token=True): result = self.tokenizer( @@ -265,21 +278,27 @@ class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy): user_token = self._get_user_token() assistant_token = self._get_assistant_token() try: - for i, part in enumerate(self.prompter.build_prompt(prompt["conversations"])): + for i, part in enumerate( + self.prompter.build_prompt(prompt["conversations"]) + ): if isinstance(part, tuple): if part[0] == "USER:": part = part[0] + part[1] if not user_token else part[1] # this is still the user query, we should - res = self._tokenize(part.strip(), add_eos_token=False, strip_bos_token=True) + res = self._tokenize( + part.strip(), add_eos_token=False, strip_bos_token=True + ) if user_token: res["input_ids"] = [user_token, *res["input_ids"]] # everything from this is masked out from the labels - labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) + labels = [IGNORE_TOKEN_ID] * len(res["input_ids"]) elif part[0] == "ASSISTANT:": # TODO label assistant token/tokens w/ IGNORE_TOKEN_ID part = part[0] + part[1] if not assistant_token else part[1] # this should be the assistent response, should end with an eos token - res = self._tokenize(part.strip(), add_eos_token=True, strip_bos_token=True) + res = self._tokenize( + part.strip(), add_eos_token=True, strip_bos_token=True + ) if assistant_token: res["input_ids"] = [assistant_token, *res["input_ids"]] # not masked out from labels @@ -288,15 +307,16 @@ class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy): logging.warning("unhandled role: " + part[0]) else: # this is only ever the first part, should include the bos token and the user query - res = self._tokenize(part.strip(), add_eos_token=False, strip_bos_token=False) + res = self._tokenize( + part.strip(), add_eos_token=False, strip_bos_token=False + ) # everything from this is masked out from the labels - labels = [ IGNORE_TOKEN_ID ] * len(res["input_ids"]) + labels = [IGNORE_TOKEN_ID] * len(res["input_ids"]) input_ids = res["input_ids"] input_len = len(input_ids) result["input_ids"][current_len : current_len + input_len] = input_ids result["attention_mask"][current_len : current_len + input_len] = [ - 1 if x != self.tokenizer.pad_token_id else 0 - for x in input_ids + 1 if x != self.tokenizer.pad_token_id else 0 for x in input_ids ] result["labels"][current_len : current_len + input_len] = labels current_len += input_len @@ -320,10 +340,7 @@ class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy): result["input_ids"].append(self.tokenizer.eos_token_id) result["attention_mask"].append(1) - if ( - result["input_ids"][0] == self.tokenizer.bos_token_id - and strip_bos_token - ): + if result["input_ids"][0] == self.tokenizer.bos_token_id and strip_bos_token: result["input_ids"] = result["input_ids"][1:] result["attention_mask"] = result["attention_mask"][1:] diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py index c79c3afa7..a6d237a11 100644 --- a/src/axolotl/prompters.py +++ b/src/axolotl/prompters.py @@ -23,12 +23,22 @@ class AlpacaPrompter: def match_prompt_style(self): if self.prompt_style == PromptStyle.instruct.value: - self.prompt_input = self.system_prompt + "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" - self.prompt_no_input = self.system_no_input_prompt + "### Instruction:\n{instruction}\n\n### Response:\n" + self.prompt_input = ( + self.system_prompt + + "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" + ) + self.prompt_no_input = ( + self.system_no_input_prompt + + "### Instruction:\n{instruction}\n\n### Response:\n" + ) self.response_split = "### Response:" if self.prompt_style == PromptStyle.chat.value: - self.prompt_input = self.system_prompt + "USER: {instruction}\n{input}\nASSISTANT:" - self.prompt_no_input = self.system_no_input_prompt + "USER: {instruction}\nASSISTANT:" + self.prompt_input = ( + self.system_prompt + "USER: {instruction}\n{input}\nASSISTANT:" + ) + self.prompt_no_input = ( + self.system_no_input_prompt + "USER: {instruction}\nASSISTANT:" + ) self.response_split = "ASSISTANT:" def build_prompt( @@ -55,12 +65,15 @@ class UnpromptedPrompter(AlpacaPrompter): system_prompt = "" system_no_input_prompt = "" + class JeopardyPrompter(AlpacaPrompter): prompt_input = "Below is a Jeopardy clue paired with input providing the category of the clue. Write a concise response that best answers tbe clue given the category.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" class MultipleChoiceExplainPrompter(AlpacaPrompter): - system_prompt = "Choose the answer that best answers the question. Explain your reasoning." + system_prompt = ( + "Choose the answer that best answers the question. Explain your reasoning." + ) class MultipleChoiceConcisePrompter(AlpacaPrompter): @@ -68,11 +81,15 @@ class MultipleChoiceConcisePrompter(AlpacaPrompter): class SummarizeTLDRPrompter(AlpacaPrompter): - prompt_no_input = "USER: Summarize the following article as a TL;DR.\n{instruction}\nASSISTANT:" + prompt_no_input = ( + "USER: Summarize the following article as a TL;DR.\n{instruction}\nASSISTANT:" + ) class CompletionPrompter(AlpacaPrompter): - def build_prompt(self, instruction: str, input=None, output=None) -> Generator[str, None, None]: + def build_prompt( + self, instruction: str, input=None, output=None + ) -> Generator[str, None, None]: yield instruction def get_response(self, output: str) -> str: @@ -91,7 +108,9 @@ class ReflectAlpacaPrompter: system_prompt = "Below is an instruction that describes a task, paired with an input that provides further context. You, the Assistant, should generate a response as if it were an abstract for an academic or technical paper on the query along with a methodology. Then generate an Agent Reflection where you create a long form response as if from subject matter expert, be verbose, diligent, and creative in your application of knowledge, apply it through the lens of the response generated by the assistant. Look for flawed reasoning, faulty logic, or other mistakes in the method. Finally, generate a final response and method for the user with the Assistant abstract and Reflection analysis as augmentations to the generation\n\n" system_no_input_prompt = "Below is an instruction that describes a task. You, the Assistant, should generate a response as if it were an abstract for an academic or technical paper on the query along with a methodology. Then generate an Agent Reflection where you create a long form response as if from subject matter expert, be verbose, diligent, and creative in your application of knowledge, apply it through the lens of the response generated by the assistant. Look for flawed reasoning, faulty logic, or other mistakes in the method. Finally, generate a final response and method for the user with the Assistant abstract and Reflection analysis as augmentations to the generation\n\n" - prompt_input = "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" + prompt_input = ( + "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" + ) prompt_no_input = "### Instruction:\n{instruction}\n\n### Response:\n" agent_label = "### Thought:\n{output}\n\n### Agent Reflection:\n{reflection}\n\n### Final Response:\n{corrected}" response_split = "### Response:" @@ -102,14 +121,26 @@ class ReflectAlpacaPrompter: def match_prompt_style(self): if self.prompt_style == PromptStyle.instruct.value: - self.prompt_input = self.system_prompt + "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" - self.prompt_no_input = self.system_no_input_prompt + "### Instruction:\n{instruction}\n\n### Response:\n" + self.prompt_input = ( + self.system_prompt + + "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" + ) + self.prompt_no_input = ( + self.system_no_input_prompt + + "### Instruction:\n{instruction}\n\n### Response:\n" + ) self.agent_label = "### Thought:\n{output}\n\n### Agent Reflection:\n{reflection}\n\n### Final Response:\n{corrected}" self.response_split = "### Final Response:" if self.prompt_style == PromptStyle.chat.value: - self.prompt_input = self.system_prompt + "USER: {instruction}\n{input}\nASSISTANT:" - self.prompt_no_input = self.system_no_input_prompt + "USER: {instruction}\nASSISTANT:" - self.agent_label = "\nTHOUGHT: {output}\nASSISTANT REFLECTION: {reflection}\nASSISTANT:" + self.prompt_input = ( + self.system_prompt + "USER: {instruction}\n{input}\nASSISTANT:" + ) + self.prompt_no_input = ( + self.system_no_input_prompt + "USER: {instruction}\nASSISTANT:" + ) + self.agent_label = ( + "\nTHOUGHT: {output}\nASSISTANT REFLECTION: {reflection}\nASSISTANT:" + ) self.response_split = "ASSISTANT:" def build_prompt( @@ -167,7 +198,7 @@ class Conversation: yield (role + ":", " " + message) else: logging.warning("role with empty message: " + role) - yield (role + ":", ) + yield (role + ":",) def copy(self): return Conversation( @@ -199,7 +230,9 @@ conv_vicuna_v1_1 = Conversation( class ShareGPTPrompter: def __init__(self, prompt_style=None): if prompt_style != PromptStyle.chat.value: - raise Exception(f"unsupported prompt_style for ShareGPTPrompter({prompt_style})") + raise Exception( + f"unsupported prompt_style for ShareGPTPrompter({prompt_style})" + ) # def match_prompt_style(self): # if self.prompt_style == PromptStyle.chat.value: diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py index 8d9525fa5..2f9a1afec 100644 --- a/src/axolotl/utils/data.py +++ b/src/axolotl/utils/data.py @@ -7,7 +7,8 @@ from datasets import ( load_dataset, IterableDataset, Dataset, - concatenate_datasets, DatasetDict, + concatenate_datasets, + DatasetDict, ) from huggingface_hub import hf_hub_download from transformers import PreTrainedTokenizerBase @@ -33,11 +34,14 @@ from axolotl.prompters import ( JeopardyPrompter, CompletionPrompter, MultipleChoiceExplainPrompter, - SummarizeTLDRPrompter, MultipleChoiceConcisePrompter, + SummarizeTLDRPrompter, + MultipleChoiceConcisePrompter, ) -def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_path) -> DatasetDict: +def load_tokenized_prepared_datasets( + tokenizer, cfg, default_dataset_prepared_path +) -> DatasetDict: tokenizer_name = tokenizer.__class__.__name__ ds_hash = str( md5( @@ -45,7 +49,8 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa str(cfg.sequence_len) + "@" + "|".join(sorted([f"{d.path}:{d.type}" for d in cfg.datasets])) - + "|" + tokenizer_name + + "|" + + tokenizer_name ).encode("utf-8") ).hexdigest() ) @@ -57,7 +62,9 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa dataset = None try: if cfg.push_dataset_to_hub: - dataset = load_dataset(f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True) + dataset = load_dataset( + f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True + ) dataset = dataset["train"] except: pass @@ -88,7 +95,12 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa ) elif ds_from_hub: if d.data_files: - ds = load_dataset(d.path, streaming=False, data_files=d.data_files, use_auth_token=True) + ds = load_dataset( + d.path, + streaming=False, + data_files=d.data_files, + use_auth_token=True, + ) else: ds = load_dataset(d.path, streaming=False, use_auth_token=True) else: @@ -100,49 +112,65 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa raise Exception("unhandled dataset load") # support for using a subset of the data if d.shards: - ds = ds.shuffle(seed=42)["train"].shard( - num_shards=cfg.shards, index=0 - ) + ds = ds.shuffle(seed=42)["train"].shard(num_shards=cfg.shards, index=0) d_type = d.type d_type_split = d_type.split(":") d_base_type = d_type_split[0] d_prompt_style = d_type_split[1] if len(d_type_split) > 1 else None - if (ds_strategy := load(d.type, tokenizer, cfg)): + if ds_strategy := load(d.type, tokenizer, cfg): ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) elif d_base_type == "alpaca": ds_strategy = AlpacaPromptTokenizingStrategy( - AlpacaPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len + AlpacaPrompter(d_prompt_style), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) elif d_base_type == "explainchoice": ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy( - MultipleChoiceExplainPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len + MultipleChoiceExplainPrompter(d_prompt_style), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) elif d_base_type == "concisechoice": ds_strategy = AlpacaMultipleChoicePromptTokenizingStrategy( - MultipleChoiceConcisePrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len + MultipleChoiceConcisePrompter(d_prompt_style), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) elif d_base_type == "summarizetldr": ds_strategy = SummarizeTLDRPromptTokenizingStrategy( - SummarizeTLDRPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len + SummarizeTLDRPrompter(d_prompt_style), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) elif d_base_type == "jeopardy": ds_strategy = JeopardyPromptTokenizingStrategy( - JeopardyPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len + JeopardyPrompter(d_prompt_style), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) elif d_base_type == "oasst": ds_strategy = OpenAssistantPromptTokenizingStrategy( - AlpacaPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len + AlpacaPrompter(d_prompt_style), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) @@ -166,7 +194,10 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa datasets.append(ds_wrapper) elif d_base_type == "sharegpt": ds_strategy = ShareGPTPromptTokenizingStrategy( - ShareGPTPrompter(d_prompt_style), tokenizer, cfg.train_on_inputs, cfg.sequence_len + ShareGPTPrompter(d_prompt_style), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, ) ds_wrapper = TokenizedPromptDataset(ds_strategy, ds["train"]) datasets.append(ds_wrapper) @@ -196,12 +227,16 @@ def load_tokenized_prepared_datasets(tokenizer, cfg, default_dataset_prepared_pa logging.info( f"Saving merged prepared dataset with push_to_hub... {cfg.push_dataset_to_hub}/{ds_hash}" ) - dataset.push_to_hub(f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True) + dataset.push_to_hub( + f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True + ) return dataset -def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_dataset_prepared_path) -> (Dataset, Dataset): +def load_prepare_datasets( + tokenizer: PreTrainedTokenizerBase, cfg, default_dataset_prepared_path +) -> (Dataset, Dataset): max_packed_sequence_len = ( cfg.max_packed_sequence_len if cfg.max_packed_sequence_len else cfg.sequence_len ) @@ -221,7 +256,8 @@ def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_datas + str(max_packed_sequence_len) + seed + "|".join(sorted([f"{d.path}:{d.type}" for d in cfg.datasets])) - + "|" + tokenizer_name + + "|" + + tokenizer_name ).encode("utf-8") ).hexdigest() ) @@ -237,7 +273,9 @@ def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_datas logging.info( f"Checking for packed prepared dataset from hub... {cfg.push_dataset_to_hub}/{ds_hash}" ) - dataset = load_dataset(f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True) + dataset = load_dataset( + f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True + ) dataset = dataset["train"] except: pass @@ -254,7 +292,9 @@ def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_datas logging.info( f"Saving packed prepared dataset with push_to_hub... {cfg.push_dataset_to_hub}/{ds_hash}" ) - dataset.push_to_hub(f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True) + dataset.push_to_hub( + f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True + ) else: dataset = load_tokenized_prepared_datasets( tokenizer, cfg, default_dataset_prepared_path @@ -279,9 +319,9 @@ def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_datas d for d in dataset if len(d["input_ids"]) < cfg.sequence_len - and len(d["input_ids"]) > 0 - and len(d["input_ids"]) == len(d["attention_mask"]) - and len(d["input_ids"]) == len(d["labels"]) + and len(d["input_ids"]) > 0 + and len(d["input_ids"]) == len(d["attention_mask"]) + and len(d["input_ids"]) == len(d["labels"]) ] ) @@ -294,7 +334,9 @@ def load_prepare_datasets(tokenizer: PreTrainedTokenizerBase, cfg, default_datas logging.info( f"Saving packed prepared dataset with push_to_hub... {cfg.push_dataset_to_hub}/{ds_hash}" ) - dataset.push_to_hub(f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True) + dataset.push_to_hub( + f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True + ) else: dataset = load_tokenized_prepared_datasets( tokenizer, cfg, default_dataset_prepared_path diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 8792e5ba2..feec832b0 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -11,7 +11,8 @@ from transformers import ( AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, - AutoConfig, BitsAndBytesConfig, + AutoConfig, + BitsAndBytesConfig, ) try: @@ -244,7 +245,9 @@ def load_model( embeddings_len = math.ceil(len(tokenizer) / 32) * 32 model.resize_token_embeddings(embeddings_len) - if ((cfg.adapter == "lora" and load_in_8bit) or cfg.adapter == "qlora") and not cfg.load_4bit: + if ( + (cfg.adapter == "lora" and load_in_8bit) or cfg.adapter == "qlora" + ) and not cfg.load_4bit: logging.info("converting PEFT model w/ prepare_model_for_int8_training") model = prepare_model_for_int8_training(model) @@ -265,7 +268,11 @@ def load_model( m.scales = m.scales.half() m.bias = m.bias.half() - if torch.cuda.device_count() > 1 and int(os.getenv("WORLD_SIZE", "1")) > 1 and cfg.load_4bit: + if ( + torch.cuda.device_count() > 1 + and int(os.getenv("WORLD_SIZE", "1")) > 1 + and cfg.load_4bit + ): # llama is PROBABLY model parallelizable, but the default isn't that it is # so let's only set it for the 4bit, see # https://github.com/johnsmith0031/alpaca_lora_4bit/blob/08b3fca4a4a9e0d3945be1bab4529f100a428636/finetune.py#L130-L133 diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py index 550801196..e15bbe14a 100644 --- a/src/axolotl/utils/trainer.py +++ b/src/axolotl/utils/trainer.py @@ -17,10 +17,12 @@ from axolotl.utils.callbacks import SavePeftModelCallback class OneCycleLRSchedulerTrainer(Trainer): - def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optimizer = None): - optimizer=self.optimizer if optimizer is None else optimizer - num_warmup_steps=self.args.get_warmup_steps(num_training_steps) - num_training_steps=num_training_steps + def create_scheduler( + self, num_training_steps: int, optimizer: torch.optim.Optimizer = None + ): + optimizer = self.optimizer if optimizer is None else optimizer + num_warmup_steps = self.args.get_warmup_steps(num_training_steps) + num_training_steps = num_training_steps pct_start = num_warmup_steps / num_training_steps self.lr_scheduler = OneCycleLR( @@ -203,7 +205,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): ) callbacks.append(early_stop_cb) - if cfg.local_rank == 0 and cfg.adapter == 'lora': # only save in rank 0 + if cfg.local_rank == 0 and cfg.adapter == "lora": # only save in rank 0 callbacks.append(SavePeftModelCallback) data_collator_kwargs = { @@ -214,7 +216,11 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer): else: data_collator_kwargs["pad_to_multiple_of"] = 8 - trainer_cls = OneCycleLRSchedulerTrainer if cfg.lr_scheduler == "one_cycle" and cfg.fsdp else transformers.Trainer + trainer_cls = ( + OneCycleLRSchedulerTrainer + if cfg.lr_scheduler == "one_cycle" and cfg.fsdp + else transformers.Trainer + ) trainer = trainer_cls( model=model, train_dataset=train_dataset, From bc97f9c584f795647aa23abbe19ba78b18917c94 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 24 May 2023 23:00:53 -0400 Subject: [PATCH 32/32] remove dev specific remark --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 1e46944db..969708d47 100644 --- a/README.md +++ b/README.md @@ -324,7 +324,7 @@ If you are inferencing a pretrained LORA, pass --lora_model_dir ./completed-model ``` -### Merge LORA to base (Dev branch 🔧 ) +### Merge LORA to base Add below flag to train command above @@ -345,4 +345,4 @@ Please reduce any below Bugs? Please check for open issue else create a new [Issue](https://github.com/OpenAccess-AI-Collective/axolotl/issues/new). -PRs are **greatly welcome**! \ No newline at end of file +PRs are **greatly welcome**!