From d6c649e13cdfb02b5b321fbd3f7220ee92858328 Mon Sep 17 00:00:00 2001 From: Rujikorn Charakorn Date: Mon, 4 Aug 2025 23:04:15 +0900 Subject: [PATCH] fix use_per_ctx_average_loss w/ grad accum (#8) --- src/ctx_to_lora/trainer.py | 62 +++++++++++++++++++++++++++++++++++--- 1 file changed, 57 insertions(+), 5 deletions(-) diff --git a/src/ctx_to_lora/trainer.py b/src/ctx_to_lora/trainer.py index d28d690..7685502 100644 --- a/src/ctx_to_lora/trainer.py +++ b/src/ctx_to_lora/trainer.py @@ -19,6 +19,55 @@ class DistillationTrainer(Trainer): self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False) super().__init__(*args, **kwargs) + # modified from the base Trainer to support per-context average loss + def get_batch_samples_ctx(self, epoch_iterator, num_batches, device): + 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 "ctx_ids" in batch_samples[0] + and ( + # num_items_in_batch is passed to model forward + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3757 + self.model_accepts_loss_kwargs + # num_items_in_batch is passed to compute_loss_func + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3773 + or self.compute_loss_func is not None + # num_items_in_batch is also verified if (self.model_accepts_loss_kwargs or self.compute_loss_func) + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3790 + ) + ) + + if count_num_items_in_batch: + # For now we don't support object detection + try: + num_items_in_batch = sum( + [(batch["ctx_position_ids"] == 0).sum() for batch in batch_samples] + ) + except (TypeError, AttributeError): + pass + + if num_items_in_batch is not None: + if self.args.average_tokens_across_devices: + num_items_in_batch = self.accelerator.gather(num_items_in_batch).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 + def compute_loss( self, model, inputs, return_outputs=False, num_items_in_batch=None ): @@ -81,7 +130,7 @@ class DistillationTrainer(Trainer): ) # these stack and split can be optimized but let's keep it simple - # Calculate mean loss for each QA pair + # mean across tokens of each q qa_losses = torch.stack( [ loss[start:end].mean() @@ -89,15 +138,15 @@ class DistillationTrainer(Trainer): ] ) - # Split losses by context and calculate per-context average loss + # mean across queries of each ctx per_ctx_losses = [ ctx.mean() for ctx in torch.split(qa_losses, n_queries_per_ctx) ] - # Final loss is the average of per-context losses - loss = torch.stack(per_ctx_losses).mean() + # mean across contexts + loss = torch.stack(per_ctx_losses) - elif reduction == "batchmean": + if reduction == "batchmean": loss = loss.mean() elif reduction == "sum" and num_items_in_batch is not None: loss = loss.sum() / num_items_in_batch @@ -270,6 +319,9 @@ def train_model( training_args.per_device_train_batch_size = 128 trainer = trainer_cls(**trainer_kwargs) + if trainer.use_per_ctx_average_loss: + 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