doc-to-lora/src/ctx_to_lora/training_utils.py

import gc
import json
import logging
from collections import defaultdict
from dataclasses import fields
from enum import Enum
from typing import Optional

import datasets
import torch
import numpy as np
from rouge_score import rouge_scorer
from torch.utils.data import RandomSampler, DataLoader
from transformers import (
    GenerationConfig,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    Trainer,
)
from transformers.trainer_utils import (
    get_last_checkpoint,
    has_length,
    seed_worker,
)
from transformers.trainer import _is_peft_model
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from transformers.trainer_pt_utils import LengthGroupedSampler
from transformers.utils import is_datasets_available


TRAINING_TASK = Enum("TRAINING_TASK", ["CAUSAL_LM", "COMPLETION"])

logger = logging.getLogger()


class ModulatedModelTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
        super().__init__(*args, **kwargs)

    def compute_loss(self, model, inputs, return_outputs=False):
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.

        Subclass and override for custom behavior.
        """
        if self.label_smoother is not None and "labels" in inputs:
            labels = inputs.pop("labels")
        else:
            labels = None

        ##### get generated LoRAs
        outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
        #####

        # Save past state if it exists
        # TODO: this needs to be fixed and made cleaner later.
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        if labels is not None:
            unwrapped_model = self.accelerator.unwrap_model(model)
            if _is_peft_model(unwrapped_model):
                model_name = unwrapped_model.base_model.model._get_name()
            else:
                model_name = unwrapped_model._get_name()
            if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
                loss = self.label_smoother(outputs, labels, shift_labels=True)
            else:
                loss = self.label_smoother(outputs, labels)
        else:
            if isinstance(outputs, dict) and "loss" not in outputs:
                raise ValueError(
                    "The model did not return a loss from the inputs, only the following keys: "
                    f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
                )
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

        ##### unpack gen lora dict and compute regularization loss
        if self.gen_lora_l1_reg_coef > 0.0:
            reg_loss = 0.0
            for module, lora in gen_loras.items():
                reg_loss += lora["A"].abs().mean() + lora["B"].abs().mean()
            loss += self.gen_lora_l1_reg_coef * reg_loss
        #####
        return (loss, outputs) if return_outputs else loss


def clear_gpu():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_max_memory_cached()


def train_model(
    model,
    # tokenizer,
    training_args,
    train_dataset=None,
    val_dataset=None,
    test_dataset=None,
    train_collator=None,
    # generation_collator=None,
    compute_metrics=None,
    train_sampler=None,
    # preprocess_logits_for_metrics=None,
    # max_new_tokens=2**13,
    # gen_per_device_eval_batch_size=1,
):
    checkpoint = None
    if training_args.resume_from_checkpoint is not None:
        checkpoint = training_args.resume_from_checkpoint
        logger.info(f"Resuming from the checkpoint: {checkpoint}")

    is_modulated_model = bool(getattr(model, "gen_lora_l1_reg_coef", 0))
    trainer_cls = Trainer if not is_modulated_model else ModulatedModelTrainer
    trainer_kwargs = dict(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=train_collator,
        compute_metrics=compute_metrics,
    )
    if is_modulated_model:
        logger.info(f"Training with modulated model. Using CustomTrainer.")
        trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.pop(
            "gen_lora_l1_reg_coef"
        )

    trainer = trainer_cls(**trainer_kwargs)

    # Trainer loads the best model after training
    # is done when load_best_model_at_end=True
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
    trainer.log_metrics("train", train_result.metrics)
    trainer.save_model()
    clear_gpu()

    # metrics = trainer.evaluate(dict(**val_dataset, test=test_dataset))
    # trainer.log_metrics("eval", metrics)
    # trainer.save_metrics("eval", metrics)
    # trainer.save_model()
    # clear_gpu()

    # ############## Evaluation
    # # TODO: eval does not work when using with deepspeed
    # # make a separate eval script

    # # max_input_len=2**13 # for input truncation

    # gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens)
    # # pad_token_id=tokenizer.pad_token_id,
    # # eos_token_id=?

    # eval_trainer_args = {}

    # # Copy only necessary attributes from training_args to eval_trainer_args
    # seq2seq_training_args_fields = {f.name for f in fields(Seq2SeqTrainingArguments)}
    # for attr, value in training_args.to_dict().items():
    #     if attr in seq2seq_training_args_fields:
    #         eval_trainer_args[attr] = value

    # eval_trainer_args["eval_strategy"] = "no"
    # eval_trainer_args["save_strategy"] = "no"
    # eval_trainer_args["overwrite_output_dir"] = True
    # eval_trainer_args["per_device_eval_batch_size"] = gen_per_device_eval_batch_size

    # # NOTE: could also set kv_cache implementation here
    # eval_trainer_args = Seq2SeqTrainingArguments(
    #     **eval_trainer_args,
    #     predict_with_generate=True,
    #     generation_config=GenerationConfig(**gen_kwargs),
    # )

    # # Seq2SeqTrainer is actually just the same as Trainer
    # # (although it uses a different data collator, i.e., explicit prompt/answer separation)
    # # it just allows `predict_with_generate`
    # # allowing us to compute metrics on the generated outputs
    # # no clue why they call this seq2seq...

    # logger.info("=" * 80 + "\n" + "Evaluating model..." + "\n" + "=" * 80)

    # model.eval()

    # eval_trainer = Seq2SeqTrainer(
    #     model=model,
    #     args=eval_trainer_args,
    #     # TODO: use a different collator for test, e.g., more max_len truncation
    #     # w/ left padding?
    #     # removing label part from input_ids
    #     data_collator=generation_collator,
    # )

    # for split, ds in zip(["eval", "test"], [val_dataset, test_dataset]):
    #     if ds is None:
    #         continue
    #     eval_generation(eval_trainer, tokenizer, ds, split, gen_kwargs)
    #     clear_gpu()