tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

don't shift student logits for kd

Browse Source
This commit is contained in:
Wing Lian
2025-01-15 01:07:48 -05:00
parent 483defb9ae
commit 04efcb102f
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/integrations/kd/trainer.py
View File

@@ -70,7 +70,8 @@ class AxolotlKDTrainer(AxolotlTrainer):
student_logits = outputs["logits"][:, :seq_len, :].contiguous() student_logits = outputs["logits"][:, :seq_len, :].contiguous()
if shift_targets: if shift_targets:
shift_logits = student_logits[..., :-1, :].contiguous() # shift_logits = student_logits[..., :-1, :].contiguous()
shift_logits = student_logits.contiguous()
target_logprobs_for_loss = target_logprobs[..., 1:, :].contiguous() target_logprobs_for_loss = target_logprobs[..., 1:, :].contiguous()
target_token_ids_for_loss = target_token_ids[..., 1:, :].contiguous() target_token_ids_for_loss = target_token_ids[..., 1:, :].contiguous()
target_mask_for_loss = target_mask[..., 1:, :].contiguous() target_mask_for_loss = target_mask[..., 1:, :].contiguous()