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

no where support

Browse Source
This commit is contained in:
Wing Lian
2024-12-21 13:21:54 -05:00
parent bc3326a808
commit 48ccf55752
1 changed files with 9 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/axolotl/integrations/kd/kernels/kd.py
View File

@@ -55,19 +55,20 @@ def fwd_kl_topk_kernel(
# load student logprobs # load student logprobs
s_lp = tl.load(s_row_ptr + k_offset * stride_sk, mask=mask, other=-float("inf")) s_lp = tl.load(s_row_ptr + k_offset * stride_sk, mask=mask, other=-float("inf"))
# load mask => bool # load mask => bool (0 or 1)
valid = tl.load( valid = tl.load(m_row_ptr + k_offset * stride_mk, mask=mask, other=0)
m_row_ptr + k_offset * stride_mk, mask=mask, other=0 valid_f32 = valid.to(tl.float32)
) # 0 or 1 => bool
valid_bool = valid.to(tl.int1)
# teacher probs # teacher probs
t_p = tl.exp(t_lp) t_p = tl.exp(t_lp)
# local_kl = p^T * (lp^T - lp^S) * valid # local_kl = p^T * (lp^T - lp^S)
local_kl = t_p * (t_lp - s_lp) local_kl = t_p * (t_lp - s_lp)
# sum only over valid positions # multiply by valid_f32 to ignore padded or invalid positions
kl_sum += tl.sum(local_kl, where=valid_bool) local_kl *= valid_f32
# sum over the chunk
kl_sum += tl.sum(local_kl, where=mask)
# store rowwise result # store rowwise result
tl.store(loss_out_ptr + row_id * stride_loss_n, kl_sum) tl.store(loss_out_ptr + row_id * stride_loss_n, kl_sum)