Upload _gptqlora.py

_gptqlora.py

@@ -0,0 +1,613 @@
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from safetensors import safe_open
+from safetensors.torch import load_model, save_model, load_file
+
+
+from collections import defaultdict
+import copy
+import json
+import os
+from os.path import exists, join, isdir
+from dataclasses import dataclass, field
+import sys
+from typing import Optional, Dict, Sequence
+import numpy as np
+from tqdm import tqdm
+import logging
+
+import torch
+import transformers
+from torch.nn.utils.rnn import pad_sequence
+import argparse
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM, 
+    LineByLineTextDataset,
+    set_seed, 
+    Seq2SeqTrainer,
+    Trainer,
+    LlamaTokenizerFast
+)
+
+from trl import SFTTrainer
+from datasets import load_dataset
+import evaluate
+
+from peft import (
+    LoraConfig,
+    get_peft_model_state_dict,
+    set_peft_model_state_dict,
+    PeftModel
+)
+from peft.tuners.lora import LoraLayer
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+from auto_gptq.utils.peft_utils import get_gptq_peft_model, GPTQLoraConfig
+from auto_gptq import AutoGPTQForCausalLM
+from auto_gptq.nn_modules.qlinear import GeneralQuantLinear
+
+torch.backends.cuda.matmul.allow_tf32 = True
+
+logger = logging.getLogger(__name__)
+
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN = "[PAD]"
+
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+def prepare_model_for_int8_training(model, use_gradient_checkpointing=True):
+    r"""
+    This method wraps the entire protocol for preparing a model before running a training. This includes:
+        1- Cast the layernorm in fp32 2- making output embedding layer require grads 3- Add the upcasting of the lm
+        head to fp32
+
+    Args:
+        model, (`transformers.PreTrainedModel`):
+            The loaded model from `transformers`
+    """
+    for name, param in model.named_parameters():
+        # freeze base model's layers
+        param.requires_grad = False
+        
+    if use_gradient_checkpointing:
+        # For backward compatibility
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        else:
+
+            def make_inputs_require_grad(module, input, output):
+                output.requires_grad_(True)
+
+            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+
+        # enable gradient checkpointing for memory efficiency
+        model.gradient_checkpointing_enable()
+
+    return model
+
+@dataclass
+class ModelArguments:
+    model_path: Optional[str] = field(
+        default="./src/"
+    )
+    src_lora_path: Optional[str] = field(
+        default=None,
+    )
+    trust_remote_code: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Enable unpickling of arbitrary code in AutoModelForCausalLM#from_pretrained."}
+    )
+
+@dataclass
+class DataArguments:
+    eval_dataset_size: int = field(
+        default=1024, metadata={"help": "Size of validation dataset."}
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        },
+    )
+    offload_folder: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "Offload folder "
+            "value if set."
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+            "value if set."
+        },
+    )
+    source_max_len: int = field(
+        default=1024,
+        metadata={"help": "Maximum source sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+    target_max_len: int = field(
+        default=1024,
+        metadata={"help": "Maximum target sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+    dataset: str = field(
+        default='alpaca',
+        metadata={"help": "Which dataset to finetune on. See datamodule for options."}
+    )
+
+@dataclass
+class TrainingArguments(transformers.Seq2SeqTrainingArguments):
+    cache_dir: Optional[str] = field(
+        default=None
+    )
+    train_on_source: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether to train on the input in addition to the target text."}
+    )
+    mmlu_split: Optional[str] = field(
+        default='eval',
+        metadata={"help": "The MMLU split to run on"}
+    )
+    mmlu_dataset: Optional[str] = field(
+        default='mmlu-fs',
+        metadata={"help": "MMLU dataset to use: options are `mmlu-zs` for zero-shot or `mmlu-fs` for few shot."}
+    )
+    do_mmlu_eval: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether to run the MMLU evaluation."}
+    )
+    max_mmlu_samples: Optional[int] = field(
+        default=None,
+        metadata={"help": "If set, only evaluates on `max_mmlu_samples` of the MMMLU dataset."}
+    )
+    mmlu_source_max_len: int = field(
+        default=2048,
+        metadata={"help": "Maximum source sequence length for mmlu."}
+    )
+    full_finetune: bool = field(
+        default=False,
+        metadata={"help": "Finetune the entire model without adapters."}
+    )
+    adam8bit: bool = field(
+        default=False,
+        metadata={"help": "Use 8-bit adam."}
+    )
+    lora_r: int = field(
+        default=64,
+        metadata={"help": "Lora R dimension."}
+    )
+    lora_alpha: float = field(
+        default=16,
+        metadata={"help": " Lora alpha."}
+    )
+    lora_dropout: float = field(
+        default=0.0,
+        metadata={"help":"Lora dropout."}
+    )
+    max_memory_MB: int = field(
+        default=24000,
+        metadata={"help": "Free memory per gpu."}
+    )
+    report_to: str = field(
+        default='none',
+        metadata={"help": "To use wandb or something else for reporting."}
+    )
+    output_dir: str = field(default='./output', metadata={"help": 'The output dir for logs and checkpoints'})
+    optim: str = field(default='paged_adamw_32bit', metadata={"help": 'The optimizer to be used'})
+    per_device_train_batch_size: int = field(default=1, metadata={"help": 'The training batch size per GPU. Increase for better speed.'})
+    gradient_accumulation_steps: int = field(default=16, metadata={"help": 'How many gradients to accumulate before to perform an optimizer step'})
+    max_steps: int = field(default=0, metadata={"help": 'How many optimizer update steps to take'})
+    weight_decay: float = field(default=0.0, metadata={"help": 'The L2 weight decay rate of AdamW'}) # use lora dropout instead for regularization if needed
+    learning_rate: float = field(default=0.0002, metadata={"help": 'The learnign rate'})
+    remove_unused_columns: bool = field(default=False, metadata={"help": 'Removed unused columns. Needed to make this codebase work.'})
+    max_grad_norm: float = field(default=0.3, metadata={"help": 'Gradient clipping max norm. This is tuned and works well for all models tested.'})
+    gradient_checkpointing: bool = field(default=True, metadata={"help": 'Use gradient checkpointing. You want to use this.'})
+    do_train: bool = field(default=True, metadata={"help": 'To train or not to train, that is the question?'})
+    lr_scheduler_type: str = field(default='constant', metadata={"help": 'Learning rate schedule. Constant a bit better than cosine, and has advantage for analysis'})
+    warmup_ratio: float = field(default=0.03, metadata={"help": 'Fraction of steps to do a warmup for'})
+    logging_steps: int = field(default=10, metadata={"help": 'The frequency of update steps after which to log the loss'})
+    group_by_length: bool = field(default=True, metadata={"help": 'Group sequences into batches with same length. Saves memory and speeds up training considerably.'})
+    save_strategy: str = field(default='steps', metadata={"help": 'When to save checkpoints'})
+    save_steps: int = field(default=250, metadata={"help": 'How often to save a model'})
+    save_total_limit: int = field(default=40, metadata={"help": 'How many checkpoints to save before the oldest is overwritten'})
+
+def find_all_linear_names(args, model):
+    cls = GeneralQuantLinear if not(args.full_finetune) else torch.nn.Linear
+    lora_module_names = set()
+    for name, module in model.named_modules():
+        if isinstance(module, cls):
+            names = name.split('.')
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+
+
+    if 'lm_head' in lora_module_names: # needed for 16-bit
+        lora_module_names.remove('lm_head')
+    return list(lora_module_names)
+
+
+class SavePeftModelCallback(transformers.TrainerCallback):
+    def save_model(self, args, state, kwargs):
+        print('Saving PEFT checkpoint...')
+        if state.best_model_checkpoint is not None:
+            checkpoint_folder = os.path.join(state.best_model_checkpoint, "adapter_model")
+        else:
+            checkpoint_folder = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")
+
+        peft_model_path = os.path.join(checkpoint_folder, "adapter_model")
+        kwargs["model"].save_pretrained(peft_model_path)
+
+        pytorch_model_path = os.path.join(checkpoint_folder, "pytorch_model.bin")
+        if os.path.exists(pytorch_model_path):
+            os.remove(pytorch_model_path)
+
+    def on_save(self, args, state, control, **kwargs):
+        self.save_model(args, state, kwargs)
+        return control
+
+    def on_train_end(self, args, state, control, **kwargs):
+        def touch(fname, times=None):
+            with open(fname, 'a'):
+                os.utime(fname, times)
+
+        touch(join(args.output_dir, 'completed'))
+        self.save_model(args, state, kwargs)
+
+def get_accelerate_model(args, checkpoint_dir):
+
+    n_gpus = torch.cuda.device_count()
+    max_memory = f'{args.max_memory_MB}MB'
+    max_memory = {i: max_memory for i in range(n_gpus)}
+
+    if args.full_finetune: assert args.bits in [16, 32]
+
+    print(f'loading base model {args.model_path}...')
+    model = AutoGPTQForCausalLM.from_quantized(
+        args.model_path,
+        low_cpu_mem_usage=True,
+        device_map='auto',
+        max_memory=max_memory,
+        trust_remote_code=args.trust_remote_code,
+        inject_fused_attention = True,
+        inject_fused_mlp = False,
+        use_triton=False,
+        warmup_triton=False,
+        offload_folder='offload',
+        trainable=True
+    )
+    model.model.quantize_config = model.quantize_config
+    model.train()
+
+    setattr(model, 'model_parallel', True)
+    setattr(model, 'is_parallelizable', True)
+    modules = find_all_linear_names(args, model)
+
+    print("Modules: ", modules)
+
+    model.config.torch_dtype=torch.float16 #if args.fp16 else (torch.bfloat16 if args.bf16 else torch.float32))
+
+    if not args.full_finetune:
+        model = prepare_model_for_int8_training(model, use_gradient_checkpointing=args.gradient_checkpointing)
+    if args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+
+    config = GPTQLoraConfig(
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        target_modules=modules,
+        lora_dropout=args.lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    if not args.full_finetune:
+        if checkpoint_dir is not None:
+            print("Loading adapters from checkpoint.")
+            model = PeftModel.from_pretrained(model, join(checkpoint_dir, 'adapter_model'))
+            for name, p in model.named_parameters():
+                if 'lora' in name:
+                    print(name, p.sum())
+        else:
+            print(f'adding LoRA modules...')
+            model = get_gptq_peft_model(model, config, auto_find_all_linears=True, train_mode=True)
+
+    if args.gradient_checkpointing:
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        else:
+            def make_inputs_require_grad(module, input, output):
+                output.requires_grad_(True)
+            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+
+
+    for name, module in model.named_modules():
+        if isinstance(module, LoraLayer):
+            if args.bf16:
+                module = module.to(torch.bfloat16)
+        if 'norm' in name:
+            module = module.to(torch.float32)
+        if 'lm_head' in name or 'embed_tokens' in name:
+            if hasattr(module, 'weight'):
+                if args.bf16 and module.weight.dtype == torch.float32:
+                    module = module.to(torch.bfloat16)
+    return model
+
+def print_trainable_parameters(args, model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    try:
+        trainable_params /= (32//model.quantize_config.bits)
+    except:
+        pass
+    print(f"trainable params: {trainable_params} || all params: {all_param} || trainable: {100 * trainable_params / all_param}")
+
+def smart_tokenizer_and_embedding_resize(
+    special_tokens_dict: Dict,
+    tokenizer: transformers.PreTrainedTokenizer,
+    model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+@dataclass
+class DataCollatorForCausalLM(object):
+    tokenizer: transformers.PreTrainedTokenizer
+    source_max_len: int
+    target_max_len: int
+    train_on_source: bool
+    predict_with_generate: bool
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        # Extract elements
+        sources = [example['input'] for example in instances]
+        targets = [f"{example['output']}{self.tokenizer.eos_token}" for example in instances]
+        # Tokenize
+        tokenized_sources_with_prompt = self.tokenizer(
+            sources,
+            max_length=self.source_max_len,
+            truncation=True,
+        )
+        tokenized_targets = self.tokenizer(
+            targets,
+            max_length=self.target_max_len,
+            truncation=True,
+            add_special_tokens=False,
+        )
+        # Build the input and labels for causal LM
+        input_ids = []
+        labels = [] 
+        for tokenized_source, tokenized_target in zip(
+            tokenized_sources_with_prompt['input_ids'], 
+            tokenized_targets['input_ids']
+        ):
+            if not self.predict_with_generate:
+                input_ids.append(torch.tensor(tokenized_source + tokenized_target))
+                if not self.train_on_source:
+                    labels.append(
+                        torch.tensor([IGNORE_INDEX for _ in range(len(tokenized_source))] + copy.deepcopy(tokenized_target))
+                    )
+                else:
+                    labels.append(torch.tensor(copy.deepcopy(tokenized_source + tokenized_target)))
+            else:
+                input_ids.append(torch.tensor(tokenized_source))
+        # Apply padding
+        input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
+        labels = pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX) if not self.predict_with_generate else None
+        data_dict = {
+            'input_ids': input_ids,
+            'attention_mask':input_ids.ne(self.tokenizer.pad_token_id),
+        }
+        if labels is not None:
+            data_dict['labels'] = labels
+        return data_dict
+
+def extract_unnatural_instructions_data(examples, extract_reformulations=False):
+    out = {
+        'input': [],
+        'output': [],
+    }
+    for example_instances in examples['instances']:
+        for instance in example_instances:
+            out['input'].append(instance['instruction_with_input'])
+            out['output'].append(instance['output'])
+    if extract_reformulations:
+        for example_reformulations in examples['reformulations']:
+            if example_reformulations is not None:
+                for instance in example_reformulations:
+                    out['input'].append(instance['instruction_with_input'])
+                    out['output'].append(instance['output'])
+    return out
+
+def make_data_module(tokenizer: transformers.PreTrainedTokenizer, args) -> Dict:
+    # Load dataset.
+    print(args.dataset)
+
+    if args.dataset == 'txt':
+        from transformers import TextDataset
+        with open("txt.txt","r",encoding="utf-8") as f:
+            data = f.readlines()
+
+        tmp = ''
+        gdata = []
+        current_length = 0
+        print("Creating groups...")
+        for s in data:
+            if current_length + len(s) <= 512:
+                tmp = tmp + s + "\n"
+                current_length += len(s)
+            else:
+                gdata.append(tmp)
+                tmp = s
+                current_length = len(s)
+
+        l = list(map(lambda x: {
+                  'input': '',
+                  'output': x
+              }, gdata))
+        from datasets import Dataset
+        dataset=Dataset.from_list(l)
+
+    elif args.dataset == 'dataset':
+        dataset = load_dataset("json", data_files='dataset.json')
+
+    if args.do_train:
+        if args.dataset == 'txt':
+            train_dataset = dataset
+        else:
+            train_dataset = dataset['train']
+        if args.max_train_samples is not None and len(train_dataset) > args.max_train_samples:
+            train_dataset = train_dataset.select(range(args.max_train_samples))
+        if args.group_by_length:
+            train_dataset = train_dataset.map(lambda x: {'length': len(x['input']) + len(x['output'])})
+
+    data_collator = DataCollatorForCausalLM(
+        tokenizer=tokenizer, 
+        source_max_len=args.source_max_len,
+        target_max_len=args.target_max_len,
+        train_on_source=args.train_on_source,
+        predict_with_generate=args.predict_with_generate,
+    )
+    return dict(
+        train_dataset=train_dataset if args.do_train else None, 
+        eval_dataset=eval_dataset if args.do_eval else None,
+        predict_dataset=eval_dataset if args.do_predict else None,
+        data_collator=data_collator
+    )
+
+def get_last_checkpoint(checkpoint_dir):
+    if isdir(checkpoint_dir):
+        is_completed = exists(join(checkpoint_dir, 'completed'))
+        if is_completed: return None, True # already finished
+        max_step = 0
+        for filename in os.listdir(checkpoint_dir):
+            if isdir(join(checkpoint_dir, filename)) and filename.startswith('checkpoint'):
+                max_step = max(max_step, int(filename.replace('checkpoint-', '')))
+                print("MX: ", max_step, " - ", filename)
+        if max_step == 0: return None, is_completed # training started, but no checkpoint
+        checkpoint_dir = join(checkpoint_dir, f'checkpoint-{max_step}')
+        print(f"Found a previous checkpoint at: {checkpoint_dir}")
+        return checkpoint_dir, is_completed # checkpoint found!
+    return None, False # first training
+
+def train():
+    hfparser = transformers.HfArgumentParser((
+        ModelArguments, DataArguments, TrainingArguments
+    ))
+    model_args, data_args, training_args, extra_args = \
+        hfparser.parse_args_into_dataclasses(return_remaining_strings=True)
+#    training_args.generation_config = transformers.GenerationConfig(**vars(generation_args))
+    args = argparse.Namespace(
+        **vars(model_args), **vars(data_args), **vars(training_args)
+    )
+    
+    checkpoint_dir, completed_training = get_last_checkpoint(args.output_dir)
+
+    if completed_training:
+        print('Detected that training was already completed!')
+
+    model = get_accelerate_model(args, checkpoint_dir)
+    training_args.skip_loading_checkpoint_weights=True
+
+    load_existing_lora = os.path.exists('src_lora/adapter_model.safetensors')
+
+    if load_existing_lora:
+        print(f"Loading existing LoRA")
+        adapters_weights = load_file('src_lora/adapter_model.safetensors')
+        set_peft_model_state_dict(model, adapters_weights)
+
+    model.config.use_cache = False
+    print_trainable_parameters(args, model)
+    print('loaded model')
+    set_seed(args.seed)
+
+    # Tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model_path,
+        cache_dir=args.cache_dir,
+        padding_side="right",
+        use_fast=True,
+    )
+    
+    if tokenizer.pad_token is None:
+        smart_tokenizer_and_embedding_resize(
+            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
+            tokenizer=tokenizer,
+            model=model,
+        )
+        
+    if isinstance(tokenizer, LlamaTokenizerFast):
+        # LLaMA tokenizer may not have correct special tokens set.
+        # Check and add them if missing to prevent them from being parsed into different tokens.
+        # Note that these are present in the vocabulary. 
+        # Note also that `model.config.pad_token_id` is 0 which corresponds to `<unk>` token.
+        tokenizer.add_special_tokens(
+            {
+                "eos_token": tokenizer.convert_ids_to_tokens(model.config.eos_token_id),
+                "bos_token": tokenizer.convert_ids_to_tokens(model.config.bos_token_id),
+                "unk_token": tokenizer.convert_ids_to_tokens(model.config.pad_token_id),
+            }
+        )
+    
+    data_module = make_data_module(tokenizer=tokenizer, args=args)
+    trainer = Seq2SeqTrainer(
+#    trainer = SFTTrainer(
+        model=model, 
+        tokenizer=tokenizer,
+        args=training_args,
+        **{k:v for k,v in data_module.items() if k != 'predict_dataset'},
+    )
+
+    # Callbacks
+    if not args.full_finetune:
+        trainer.add_callback(SavePeftModelCallback)
+
+    # Verifying the datatypes.
+    dtypes = {}
+    for _, p in model.named_parameters():
+        dtype = p.dtype
+        if dtype not in dtypes: dtypes[dtype] = 0
+        dtypes[dtype] += p.numel()
+    total = 0
+    for k, v in dtypes.items(): total+= v
+    for k, v in dtypes.items():
+        print(k, v, v/total)
+
+    all_metrics = {"run_name": args.run_name}
+    # Training
+    if args.do_train:
+        train_result = trainer.train(resume_from_checkpoint=False)
+        metrics = train_result.metrics
+        trainer.log_metrics("train", metrics)
+        trainer.save_metrics("train", metrics)
+        trainer.save_state()
+        all_metrics.update(metrics)
+
+    if (args.do_train):
+        with open(os.path.join(args.output_dir, "metrics.json"), "w") as fout:
+            fout.write(json.dumps(all_metrics))
+
+if __name__ == "__main__":
+    train()