ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update cross_entropy_loss.py

Browse source
This commit is contained in:
Daniel Han 2024-11-03 21:35:27 -08:00
commit e03a0dacdf
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/kernels/cross_entropy_loss.py
View file

@ -64,7 +64,7 @@ def _cross_entropy_forward(
This ensures exp(x - max(x))'s maximum is 1 as exp(0) = 1.
"""
row_idx = tl.program_id(0)
logits_ptr += row_idx * logits_row_stride
logits_ptr += row_idx * tl.int64(logits_row_stride)
loss_ptr += row_idx
logsumexp_ptr += row_idx
labels_ptr += row_idx
@ -143,7 +143,7 @@ def _chunked_cross_entropy_forward(
"""
row_idx = tl.program_id(0)
chunk_idx = tl.program_id(1)
logits_ptr += row_idx * logits_row_stride
logits_ptr += row_idx * tl.int64(logits_row_stride)
loss_ptr += row_idx
logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx
labels_ptr += row_idx
@ -218,7 +218,7 @@ def _cross_entropy_backward(
row_idx = tl.program_id(0)
block_idx = tl.program_id(1)
logits_ptr += row_idx * logits_row_stride
logits_ptr += row_idx * tl.int64(logits_row_stride)
dloss_ptr += row_idx * dloss_row_stride
col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = col_offsets < VOCAB_SIZE