From 4997ae90eddf7e1bf87d01a0445f6c5c67e3198a Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 22 Apr 2025 10:47:41 +0000 Subject: [PATCH] logging l1 norm --- src/ctx_to_lora/training_utils.py | 72 +++++++++++++++---------------- 1 file changed, 35 insertions(+), 37 deletions(-) diff --git a/src/ctx_to_lora/training_utils.py b/src/ctx_to_lora/training_utils.py index f2783c8..6b9c0a1 100644 --- a/src/ctx_to_lora/training_utils.py +++ b/src/ctx_to_lora/training_utils.py @@ -1,32 +1,14 @@ 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 import Trainer from transformers.trainer import _is_peft_model +from transformers.trainer_utils import IntervalStrategy 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 +from ctx_to_lora.modeling_utils import ModulatedPretrainedModel TRAINING_TASK = Enum("TRAINING_TASK", ["CAUSAL_LM", "COMPLETION"]) @@ -38,21 +20,24 @@ class ModulatedModelTrainer(Trainer): 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): + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): """ 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: + if ( + self.label_smoother is not None or self.compute_loss_func is not None + ) and "labels" in inputs: labels = inputs.pop("labels") else: labels = None - - ##### get generated LoRAs + if self.model_accepts_loss_kwargs: + loss_kwargs = {} + if num_items_in_batch is not None: + loss_kwargs["num_items_in_batch"] = num_items_in_batch + inputs = {**inputs, **loss_kwargs} 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: @@ -64,7 +49,12 @@ class ModulatedModelTrainer(Trainer): 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(): + # User-defined compute_loss function + if self.compute_loss_func is not None: + loss = self.compute_loss_func( + outputs, labels, num_items_in_batch=num_items_in_batch + ) + elif 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) @@ -78,12 +68,21 @@ class ModulatedModelTrainer(Trainer): 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 + l1_norm = 0 + for module, lora in gen_loras.items(): + l1_norm += lora["A"].abs().mean() + lora["B"].abs().mean() + # rough estimate of the losses (we only log the values from one step) + if (self.state.global_step == 1 and self.args.logging_first_step) or ( + self.args.logging_strategy == IntervalStrategy.STEPS + and self.state.global_step % self.state.logging_steps == 0 + ): + self.log({"ce_loss": loss.item(), "gen_lora_reg_loss": l1_norm.item()}) + loss += self.gen_lora_l1_reg_coef * l1_norm ##### + + if self.args.average_tokens_across_devices and self.model_accepts_loss_kwargs: + loss *= self.accelerator.num_processes + return (loss, outputs) if return_outputs else loss @@ -114,7 +113,7 @@ def train_model( 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)) + is_modulated_model = isinstance(model, ModulatedPretrainedModel) trainer_cls = Trainer if not is_modulated_model else ModulatedModelTrainer trainer_kwargs = dict( model=model, @@ -126,9 +125,8 @@ def train_model( ) 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_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef + del training_args.gen_lora_l1_reg_coef trainer = trainer_cls(**trainer_kwargs)