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()