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

sparse kl

Browse source
This commit is contained in:
51616 2025-09-05 17:51:37 +09:00
parent e2e7f45472
commit f02ae874f3
1 changed files with 12 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

18
src/ctx_to_lora/trainer.py
View file

@ -164,13 +164,19 @@ class DistillationTrainer(ModulatedModelTrainer):
##### KL loss
outputs_logits = outputs.logits[label_pos[0], label_pos[1] - 1] # shift back 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"
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)
p = teacher_logp.exp()
logq = nn.functional.log_softmax(outputs_logits, dim=-1)
loss = -torch.sum(p * logq, 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)