apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix use_per_ctx_average_loss w/ grad accum (#8)

Browse source
This commit is contained in:
Rujikorn Charakorn 2025-08-04 23:04:15 +09:00 • committed by GitHub
parent 891c0bd256
commit d6c649e13c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 57 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

62
src/ctx_to_lora/trainer.py
View file

@ -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