finetune.py

import torch
import evaluate
from datasets import load_dataset
from transformers import GPTNeoXForCausalLM, BloomForCausalLM, OPTForCausalLM
from transformers.models.gpt_neox.modeling_gpt_neox import RotaryEmbedding
from transformers.models.opt.modeling_opt import OPTLearnedPositionalEmbedding
from transformers.trainer_utils import get_last_checkpoint
from itertools import chain
from typing import Optional
from dataclasses import dataclass, field
from transformers import (
    AutoTokenizer,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    default_data_collator,
    set_seed,
)

from flash_attn_wrappers import FlashAttentionWrapper, FlashAttentionWrapperWithRotary, FlashAttentionWrapperWithAlibi

@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
    """

    model_name_or_path: Optional[str] = field(
        default="pythia-1.4b",
        metadata={
            "help": (
                "The model checkpoint for weights initialization. Don't set if you want to train a model from scratch."
            )
        },
    )

    max_positions: Optional[int] = field(
        default=8192,
        metadata={
            "help": (
                "The maximun sequence length of the model."
            )
        },
    )

@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    dataset_name: Optional[str] = field(
        default="pile", metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )

def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    last_checkpoint = get_last_checkpoint(training_args.output_dir)
    set_seed(training_args.seed)
    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
    tokenizer.pad_token = tokenizer.mask_token
    max_positions = model_args.max_positions
    tokenizer.model_max_length = max_positions
    
    if "pythia" in model_args.model_name_or_path or "gpt-neox" in model_args.model_name_or_path:
        model = GPTNeoXForCausalLM.from_pretrained(model_args.model_name_or_path)
        for each in model.gpt_neox.layers:
            each.attention.rotary_emb = RotaryEmbedding(each.attention.rotary_ndims,max_positions,10000)
            each.attention.bias = torch.tril(torch.ones((max_positions, max_positions), dtype=torch.uint8)).view(
                        1, 1, max_positions, max_positions
                    )
            each.attention = FlashAttentionWrapperWithRotary(each.attention, max_seqlen = max_positions)
    elif "bloom" in model_args.model_name_or_path:
        model = BloomForCausalLM.from_pretrained(model_args.model_name_or_path)
        for each in model.transformer.h:
            each.self_attention = FlashAttentionWrapperWithAlibi(each.self_attention, max_seqlen = max_positions)
    elif "opt" in model_args.model_name_or_path:
        model = OPTForCausalLM.from_pretrained(model_args.model_name_or_path)
        for each in model.model.decoder.layers:
            each.self_attn = FlashAttentionWrapper(each.self_attn, max_seqlen = max_positions)
        original_num_embeddings = model.model.decoder.embed_positions.num_embeddings - 2
        assert (max_positions + 2) % original_num_embeddings == 0
        original_embed_positions = model.model.decoder.embed_positions.weight.data
        duplicated_embed_positions =  torch.cat([original_embed_positions[:-2] * i for i in range(1, (max_positions + 2) // original_num_embeddings + 1)] +  [original_embed_positions[-2:]], dim = 0)
        model.model.decoder.embed_positions = OPTLearnedPositionalEmbedding((max_positions + 2), model.model.decoder.embed_positions.embedding_dim)
        model.model.decoder.embed_positions.weight.data = duplicated_embed_positions
    else:
        raise NotImplementedError

    # patching for the random contiguous tensors bug
    for p in model.parameters():
        p = p.contiguous()

    def merge_questions_and_answers(examples):
        out = tokenizer([question + " " + answer for question, answer in zip(examples["input"], examples["output"])])
        return out

    block_size = tokenizer.model_max_length
    def group_texts(examples):
        # Concatenate all texts.
        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
        total_length = len(concatenated_examples[list(examples.keys())[0]])
        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
        # customize this part to your needs.
        if total_length >= block_size:
            total_length = (total_length // block_size) * block_size
        # Split by chunks of max_len.
        result = {
            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
            for k, t in concatenated_examples.items()
        }
        result["labels"] = result["input_ids"].copy()
        return result
    

    if data_args.dataset_name == "pile":
        base_url = "https://the-eye.eu/public/AI/pile/"
        data_files = {
            "train": [base_url + "train/"+ f"{idx:02d}.jsonl.zst" for idx in range(30)],
            "validation": base_url + "val.jsonl.zst",
            "test": base_url + "test.jsonl.zst",
        }
        datasets = load_dataset("json", data_files=data_files, streaming=True)
        datasets = datasets.filter(lambda x: len(x["text"])>=max_positions)
        tokenized_datasets = datasets.map(
            lambda examples: tokenizer(examples["text"]),
            batched=True,
        )
        lm_datasets = tokenized_datasets.map(
            group_texts,
            batched=True,
        )
        lm_datasets = lm_datasets.filter(lambda x: len(x["input_ids"])>=max_positions)
    elif data_args.dataset_name == "qasper":
        datasets = load_dataset("tau/scrolls", "qasper")
        datasets.pop("test")
        tokenized_datasets = datasets.map(
           merge_questions_and_answers,
           batched=True,
           num_proc = 1,
           remove_columns = datasets["train"].column_names,
           desc="Running tokenizer on dataset",
        )

        lm_datasets = tokenized_datasets.map(
           group_texts,
           batched=True,
           num_proc=1,
           desc=f"Grouping texts in chunks of {block_size}",
        )
    else:
        raise Exception("Sorry, please the dataset specified can not be recognized")
    
    def preprocess_logits_for_metrics(logits, labels):
        if isinstance(logits, tuple):
            # Depending on the model and config, logits may contain extra tensors,
            # like past_key_values, but logits always come first
            logits = logits[0]
        return logits.argmax(dim=-1)

    metric = evaluate.load("accuracy")

    def compute_metrics(eval_pred):
        preds, labels = eval_pred
        labels = labels[:, 1:].reshape(-1)
        preds = preds[:, :-1].reshape(-1)
        return metric.compute(predictions=preds, references=labels)

    train_dataset = lm_datasets["train"]
    eval_dataset = lm_datasets["validation"]

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset= train_dataset,
        eval_dataset= eval_dataset,
        tokenizer = tokenizer,
        data_collator=default_data_collator,
        compute_metrics=compute_metrics,
        preprocess_logits_for_metrics=preprocess_logits_for_metrics
    )

    if training_args.resume_from_checkpoint is not None:
        checkpoint = training_args.resume_from_checkpoint
    elif last_checkpoint is not None:
        checkpoint = last_checkpoint
    else:
        checkpoint = None
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
    trainer.save_model()  # Saves the tokenizer too for easy upload

    metrics = train_result.metrics

    max_train_samples = (len(train_dataset)
    )
    metrics["train_samples"] = min(max_train_samples, len(train_dataset))

    trainer.log_metrics("train", metrics)
    trainer.save_metrics("train", metrics)
    trainer.save_state()

if __name__  == "__main__":
    main()