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

chore: lint

Browse Source
This commit is contained in:
Wing Lian
2025-02-25 07:29:46 -05:00
parent e82268e580
commit a2e52a29e9
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/axolotl/integrations/kd/topk_logprob/forward_kl_triton.py
View File

@@ -338,6 +338,7 @@ class TopKKLDivergence(torch.autograd.Function):
) )
kd_loss = token_losses.sum() kd_loss = token_losses.sum()
# pylint: disable=duplicate-code
# Apply temperature scaling # Apply temperature scaling
if kd_temperature != 1.0: if kd_temperature != 1.0:
kd_loss = kd_loss * (kd_temperature**2) kd_loss = kd_loss * (kd_temperature**2)
@@ -426,7 +427,7 @@ class TopKKLDivergence(torch.autograd.Function):
target_mask.stride(0), target_mask.stride(0),
target_mask.stride(1), target_mask.stride(1),
target_mask.stride(2), target_mask.stride(2),
min(256, triton.next_power_of_2(top_k)), min(1024, triton.next_power_of_2(top_k)),
) )
else: else:
# Case 2: Softmax over full vocab # Case 2: Softmax over full vocab
@@ -457,7 +458,7 @@ class TopKKLDivergence(torch.autograd.Function):
target_mask.stride(0), target_mask.stride(0),
target_mask.stride(1), target_mask.stride(1),
target_mask.stride(2), target_mask.stride(2),
min(256, triton.next_power_of_2(top_k)), min(1024, triton.next_power_of_2(top_k)),
) )
# Return gradients for student_logits and None for other inputs # Return gradients for student_logits and None for other inputs