import logging import torch from torch import nn from transformers import Trainer from transformers.trainer_pt_utils import get_parameter_names from transformers.trainer_utils import IntervalStrategy from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel logger = logging.getLogger() def per_ctx_loss_ce(inputs, labels, loss): # loss still has masked out elem (0 at labels=-100) n_queries_per_ctx = inputs["n_queries"].tolist() position_ids = inputs["position_ids"].squeeze(0) # account only label positions label_mask = labels.squeeze(0) != -100 label_pos_ids = label_mask * position_ids label_pos_ids_diff = label_pos_ids.diff( append=torch.tensor([0], device=position_ids.device) ) # assumes the input starts with non-assistant tokens start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0] end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0] label_seq_lens = end_label_pos - start_label_pos # these stack and split can be optimized but let's keep it simple # mean across tokens of each q qa_losses = torch.stack( [ loss[start : start + llen].mean() for start, llen in zip(start_label_pos, label_seq_lens) ] ) # mean across queries of each ctx per_ctx_losses = [ql.mean() for ql in torch.split(qa_losses, n_queries_per_ctx)] # per-ctx loss loss = torch.stack(per_ctx_losses) return loss def per_ctx_loss_kl(inputs, labels, loss): # loss is compact (label indices selected) n_queries_per_ctx = inputs["n_queries"].tolist() position_ids = inputs["position_ids"].squeeze(0) # account only label positions label_mask = labels.squeeze(0) != -100 label_pos_ids = label_mask * position_ids label_pos_ids_diff = label_pos_ids.diff( append=torch.tensor([0], device=position_ids.device) ) # assumes the input starts with non-assistant tokens start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0] end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0] label_seq_lens = end_label_pos - start_label_pos # find equiv start indices in the already sliced loss vector cu_label_seq_lens = torch.cumsum(label_seq_lens, dim=0) start_indices = torch.cat( ( torch.tensor([0], device=cu_label_seq_lens.device), cu_label_seq_lens[:-1], ) ) # these stack and split can be optimized but let's keep it simple # mean across tokens of each q qa_losses = torch.stack( [loss[start:end].mean() for start, end in zip(start_indices, cu_label_seq_lens)] ) # mean across queries of each ctx per_ctx_losses = [ql.mean() for ql in torch.split(qa_losses, n_queries_per_ctx)] # per-ctx loss loss = torch.stack(per_ctx_losses) return loss class ModulatedModelTrainer(Trainer): # modified from the base Trainer to support per-context average loss def get_batch_samples(self, epoch_iterator, num_batches, device): # only used with `use_per_ctx_average_loss=True` batch_samples = [] num_items_in_batch = None for _ in range(num_batches): try: batch_samples.append(next(epoch_iterator)) except StopIteration: break count_num_items_in_batch = ( len(batch_samples) > 0 and "labels" in batch_samples[0] and "n_ctx_chunks" in batch_samples[0] ) if count_num_items_in_batch: num_items_in_batch = dict() num_items_in_batch["ctx"] = torch.tensor( sum([batch["n_ctx_chunks"].numel() for batch in batch_samples]) ).to(device) # should we avg over num chunks? # num_items_in_batch["ctx"] = sum( # [(batch["ctx_position_ids"] == 0).sum() for batch in batch_samples] # ) num_items_in_batch["labels"] = sum( [(batch["labels"].ne(-100)).sum() for batch in batch_samples] ).to(device) if num_items_in_batch is not None: if self.args.average_tokens_across_devices: for k in num_items_in_batch: num_items_in_batch[k] = self.accelerator.gather( num_items_in_batch[k] ).sum() if torch.is_tensor(num_items_in_batch): num_items_in_batch = num_items_in_batch.to(device) if self.args.n_gpu > 1 and num_items_in_batch.dim() == 0: # In the DataParallel case, convert the scalar tensor into a 1-dim tensor num_items_in_batch = num_items_in_batch.unsqueeze(0) return batch_samples, num_items_in_batch class DistillationTrainer(ModulatedModelTrainer): def __init__(self, *args, **kwargs): self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0) self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False) super().__init__(*args, **kwargs) def compute_loss( self, model, inputs, return_outputs=False, num_items_in_batch=None ): # NOTE: the loss output from this fn will be ***added*** # meaning that we should always scale the loss wrt `num_items_in_batch` # (average over the number of items in the accumulated batch) is_train = num_items_in_batch is not None labels = inputs.pop("labels", None) label_pos = torch.where(labels != -100) outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True) if "logprobs_vals" not in inputs: return (torch.tensor(0.0), outputs) if return_outputs else torch.tensor(0.0) target_logp = inputs.pop("logprobs_vals").squeeze(0) indices = inputs.pop("logprobs_indices").squeeze(0) assert label_pos[0].shape[0] == target_logp.shape[0], ( "Label positions and target log probabilities should have the same # tokens." f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}" ) ##### KL loss outputs_logits = outputs.logits[label_pos[0], label_pos[1] - 1] # shift back 1 logq_full_denom = torch.logsumexp(outputs_logits, dim=-1, keepdim=True) # (N,1) selected_logits = outputs_logits.gather(1, indices) # (N,K) # log softmax at selected indices logq_selected = selected_logits - logq_full_denom p = target_logp.exp() loss = -(p * logq_selected).sum(dim=-1) # teacher_logp = torch.full_like(outputs_logits, -torch.inf) # teacher_logp.scatter_(1, indices, target_logp) # # reduction = "batchmean" if num_items_in_batch is None else "sum" # p = teacher_logp.exp() # logq = nn.functional.log_softmax(outputs_logits, dim=-1) # loss = -torch.sum(p * logq, dim=-1) if self.use_per_ctx_average_loss: loss = per_ctx_loss_kl(inputs, labels, loss) if is_train: if self.use_per_ctx_average_loss: loss = loss.sum() / num_items_in_batch["ctx"] else: loss = loss.sum() / num_items_in_batch["labels"] else: # eval loss = loss.mean() # if reduction == "batchmean": # loss = loss.mean() # elif reduction == "sum": # # loss does not scale with grad acc # # num_items_in_batch does # # this works for both token-avg and ctx-avg # # loss = loss.sum() / num_items_in_batch # `num_items_in_batch` is # tokens if `args.use_ctx_average_loss=False`` # loss = loss.sum() / num_items_in_batch ##### ##### unpack gen lora dict and compute regularization loss l1_norm = 0 n_modules = len(gen_loras) for module, lora in gen_loras.items(): l1_norm += lora["A"].abs().sum(0).mean() + lora["B"].abs().sum(0).mean() l1_norm /= n_modules if is_train: # during eval `num_items_in_batch` will be None l1_norm /= num_items_in_batch["ctx"] total_loss = loss + self.gen_lora_l1_reg_coef * l1_norm ##### scaler = self.args.gradient_accumulation_steps if is_train else 1 if self.args.average_tokens_across_devices and is_train: total_loss *= self.accelerator.num_processes scaler *= self.accelerator.num_processes # 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 ): # compensate `num_items_in_batch` division self.log( { "kl_loss": loss.item() * scaler, "gen_lora_l1_norm": l1_norm.item() * scaler, } ) return (total_loss, outputs) if return_outputs else total_loss def causal_lm_ce_loss( logits, labels, vocab_size: int, num_items_in_batch: torch.Tensor | None = None, ignore_index: int = -100, shift_labels: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: # Upcast to float if we need to compute the loss to avoid potential precision issues logits = logits.float() if shift_labels is None: # Shift so that tokens < n predict n labels = nn.functional.pad(labels, (0, 1), value=ignore_index) shift_labels = labels[..., 1:].contiguous() # Flatten the tokens logits = logits.view(-1, vocab_size) shift_labels = shift_labels.view(-1) # Enable model parallelism shift_labels = shift_labels.to(logits.device) # loss = fixed_cross_entropy( # logits, shift_labels, num_items_in_batch, ignore_index, **kwargs # ) loss = nn.functional.cross_entropy(logits, shift_labels, reduction="none") return loss class CrossEntropyTrainer(ModulatedModelTrainer): def __init__(self, *args, **kwargs): self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0) self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False) super().__init__(*args, **kwargs) 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. """ is_train = num_items_in_batch is not None labels = inputs.pop("labels", None) outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True) # [1, tot_seq_len] logits = outputs.logits # [tot_seq_len] loss = causal_lm_ce_loss(logits, labels, self.model.vocab_size) if self.use_per_ctx_average_loss: loss = per_ctx_loss_ce(inputs, labels, loss) if is_train: if self.use_per_ctx_average_loss: loss = loss.sum() / num_items_in_batch["ctx"] else: loss = loss.sum() / num_items_in_batch["labels"] else: # eval loss = loss.mean() ##### # if is_train: # if self.use_per_ctx_average_loss: # loss_kwargs["num_items_in_batch"] = num_items_in_batch["ctx"] # else: # loss_kwargs["num_items_in_batch"] = num_items_in_batch["labels"] # inputs = {**inputs, **loss_kwargs} # outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True) # # Save past state if it exists # 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() # # 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["labels"] # ) # 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) # 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 " # f"{','.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 l1_norm = 0 n_modules = len(gen_loras) for module, lora in gen_loras.items(): l1_norm += lora["A"].abs().sum(0).mean() + lora["B"].abs().sum(0).mean() l1_norm /= n_modules if is_train: # during eval `num_items_in_batch` will be None l1_norm /= num_items_in_batch["ctx"] total_loss = loss + self.gen_lora_l1_reg_coef * l1_norm ##### scaler = self.args.gradient_accumulation_steps if is_train else 1 if self.args.average_tokens_across_devices and is_train: total_loss *= self.accelerator.num_processes scaler *= self.accelerator.num_processes # 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 ): # compensate `num_items_in_batch` division self.log( { "ce_loss": loss.item() * scaler, "gen_lora_l1_norm": l1_norm.item() * scaler, } ) return (total_loss, outputs) if return_outputs else total_loss def get_decay_parameter_names(model) -> list[str]: """ Get all parameter names that weight decay will be applied to. This function filters out parameters in two ways: 1. By layer type (nn.Embedding) 2. By parameter name patterns (containing 'bias', 'layernorm', 'rmsnorm' or 'latents_q' [perceiver's latent queries]). """ decay_parameters = get_parameter_names( model, [nn.Embedding, nn.LayerNorm], ["scaler", "bias", "layernorm", "rmsnorm", "latents_q"], ) return decay_parameters def train_model( model, training_args, train_dataset=None, val_dataset=None, train_collator=None, compute_metrics=None, ): 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}") trainer_kwargs = dict( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, data_collator=train_collator, compute_metrics=compute_metrics, ) is_modulated_model = isinstance(model, ModulatedPretrainedModel) trainer_cls = Trainer if is_modulated_model: logger.info("Training with modulated model.") trainer_cls = CrossEntropyTrainer trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef trainer_kwargs["use_per_ctx_average_loss"] = ( training_args.use_per_ctx_average_loss ) del training_args.gen_lora_l1_reg_coef del training_args.use_per_ctx_average_loss if training_args.use_kl_loss: logger.info("Training with distillation loss. Using DistillationTrainer.") trainer_cls = DistillationTrainer del training_args.use_kl_loss if training_args.auto_find_batch_size: # set the batch size to some high number # which will be lowered by the Trainer training_args.per_device_train_batch_size = 128 trainer = trainer_cls(**trainer_kwargs) # if getattr(trainer, "use_per_ctx_average_loss", False): # trainer.get_batch_samples = trainer.get_batch_samples_ctx # MONKEY PATCH: remove embedding layers from weight decay trainer.get_decay_parameter_names = get_decay_parameter_names # 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() # TODO: add benchmark eval? # clear_gpu()