Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce Knowledge Distillation Base #432

Merged
merged 11 commits into from
Dec 9, 2024
Merged

Introduce Knowledge Distillation Base #432

Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
0769e97
Add liger and naive distill base
austin362667 Dec 2, 2024
a81c959
Format
austin362667 Dec 2, 2024
e13994a
Refactor beta
austin362667 Dec 2, 2024
720b5cb
Remove imports
austin362667 Dec 2, 2024
17c5b33
Fix distill base `chunk_size` scaling
austin362667 Dec 4, 2024
e3dada0
Fix chunk division
austin362667 Dec 4, 2024
5662554
Remove chunk arg
austin362667 Dec 4, 2024
7acb5ca
Fix `distillation_loss` arg typo
austin362667 Dec 5, 2024
e381569
use torch no grad and change normalization term
shivam15s Dec 7, 2024
8aa842a
rearrange fns for readability
shivam15s Dec 7, 2024
3561525
add no grad in tests
shivam15s Dec 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions src/liger_kernel/chunked_loss/fused_linear_distillation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,10 @@ def chunk_forward(
student_log_probs_chunk = F.log_softmax(student_logits_chunk.float(), dim=-1)

# Teacher
teacher_logits_chunk = teacher_input_chunk @ teacher_weight.t()
if teacher_bias is not None:
teacher_logits_chunk += teacher_bias
with torch.no_grad():
teacher_logits_chunk = teacher_input_chunk @ teacher_weight.t()
if teacher_bias is not None:
teacher_logits_chunk += teacher_bias

# The hard/task loss
ce_loss = 0.0
Expand All @@ -63,7 +64,7 @@ def forward(
student_bias=None,
teacher_bias=None,
loss_fn=None,
chunk_size=1,
chunk_size=1024,
ignore_index=-100,
weight_hard_loss=0.5,
weight_soft_loss=0.5,
Expand Down Expand Up @@ -243,6 +244,7 @@ def _compute_loss(
hard_loss /= full_target.shape[0]

soft_loss = distillation_loss_fn(student_logits_chunk, teacher_logits_chunk, temperature)
soft_loss /= (full_target.shape[0] // student_input_chunk.shape[0])
soft_loss /= full_target.shape[0]
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@shivam15s Could you help me understand why this normalization term was modified? 😀


loss = weight_hard_loss * hard_loss + weight_soft_loss * soft_loss
return loss, (soft_loss, hard_loss, student_logits_chunk, teacher_logits_chunk)
Loading