From a5abe39ded5d7a221379da075d6f95d88cd76289 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 26 Feb 2024 03:53:49 +1100 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 102 ++++++++++++++------------ 1 file changed, 56 insertions(+), 46 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 171ae6259c..a1099735d8 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -73,55 +73,65 @@ pass @triton.jit -def _large_cross_entropy_forward( +def _chunked_cross_entropy_forward( logits_ptr, logits_row_stride, loss_ptr, - lse_ptr, + logsumexp_ptr, labels_ptr, - n_rows, - n_cols, - BLOCK_SIZE: tl.constexpr, + VOCAB_SIZE : tl.constexpr, + N_CHUNKS : tl.constexpr, + BLOCK_SIZE : tl.constexpr, ): """ - Cross Entropy Loss = 1/n sum [ -yi log(Pi) ] - Pi = exp(xi) / sum(exp(xi)) - CE_i = -y log(p) = -y log[ exp(x) / sum(exp(x)) ] - = -y [ x - log[sum(exp(x))] ] - = y * (log[sum(exp(x))] - x) + 256K vocab divided in 4 chunks + + |-65536-| |-65536-| |-65536-| |-65536-| + |-------| |-------| |-------| |-------| + |-------| |-------| |-------| |-------| + If y == 0: CE_i = 0 If y == 1: CE_i = logsumexp - x + + Notice we can do logsumexp for each chunk and then + logsumexp[chunk_sum(logsumexp)] == logsumexp + + chunk_sum = log[chunk_sum(logsumexp)] + = log[exp(logsumexp(a)) + ... + exp(logsumexp(z))] + = log[exp(log[sum(exp(a))]) + ... + exp(log[sum(exp(z))])] + = log[sum(exp(a)) + ... + sum(exp(z))] + = logsumexp(x) + + This means we can perform a logsumexp for each chunk, then do a + final logsumexp reduction! + + Ie do: logsumexp(chunked_logsumexp) - x """ - row_idx = tl.program_id(0) - col_idx = tl.program_id(1) - logits_ptr += row_idx * logits_row_stride.to(tl.int64) - loss_ptr += row_idx + col_idx*n_rows - lse_ptr += row_idx + col_idx*n_rows - labels_ptr += row_idx + row_idx = tl.program_id(0) + chunk_idx = tl.program_id(1) + logits_ptr += row_idx * logits_row_stride.to(tl.int64) + loss_ptr += row_idx + logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx + labels_ptr += row_idx - col_offsets = col_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols + col_offsets = chunk_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < VOCAB_SIZE - # Get labels and logits - label_idx = tl.load(labels_ptr).to(tl.int64) + label_idx = tl.load(labels_ptr).to(tl.int32) logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32) - max_logits = tl.max(logits, 0) - # Maximum stops overflow - lse = tl.log(tl.sum(tl.exp(logits - max_logits), 0)) + max_logits - tl.store(lse_ptr, lse) + c = tl.max(logits, 0) + logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0)) - loss = 0.0 - # chained boolean operators (A or B or C) are not supported; use parentheses to split the chain. - if (label_idx != -100): - if (label_idx >= (col_idx+0)*BLOCK_SIZE) and \ - (label_idx < min((col_idx+1)*BLOCK_SIZE, n_cols)): - - logits_label = tl.load(logits_ptr + label_idx).to(tl.float32) - lse = 0.0 - loss = lse - logits_label # We add the final logsumexp after a reduction - pass + if chunk_idx == 0: + # logsumexp(chunked_logsumexp) - x + # Do the -x separately + if label_idx != -100: + x = tl.load(logits_ptr + label_idx).to(tl.float32) + loss = -1.0 * x + else: + loss = 0.0 + tl.store(loss_ptr, loss) pass - - tl.store(loss_ptr, loss) + tl.store(logsumexp_ptr, logsumexp) pass @@ -186,11 +196,11 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): div, mod = divmod(vocab_size, MAX_FUSED_SIZE) n_chunks = div + (mod != 0) + losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda") if n_chunks == 1: # For small vocabs <= 65336 like Llama, Mistral BLOCK_SIZE, num_warps = calculate_settings(vocab_size) - losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda") logsumexp = torch.empty(n_rows, dtype = torch.float32, device = "cuda") _cross_entropy_forward[(n_rows,)]( @@ -204,23 +214,23 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): ) else: # For large vocabs > 65336 like Gemma 256K - losses = torch.empty((n_chunks, n_rows), dtype = torch.float32, device = "cuda") - logsumexp = torch.empty((n_chunks, n_rows), dtype = torch.float32, device = "cuda") + logsumexp = torch.empty((n_rows, n_chunks,), dtype = torch.float32, device = "cuda") - _large_cross_entropy_forward[(n_rows, n_chunks,)]( + _chunked_cross_entropy_forward[(n_rows, n_chunks,)]( logits, logits.stride(0), losses, logsumexp, labels, - n_rows, - vocab_size, + VOCAB_SIZE = vocab_size, + N_CHUNKS = n_chunks, BLOCK_SIZE = MAX_FUSED_SIZE, num_warps = 32, ) - logsumexp = torch.logsumexp(logsumexp, dim = 0) # Column sum - losses = losses.sum(dim = 0) # Column sum - losses += logsumexp # loss = lse - logits_label - losses.masked_fill_(labels == -100, 0) # Padding tokens + # logsumexp(chunked_logsumexp) - x + # Do the -x separately + logsumexp = torch.logsumexp(logsumexp, dim = 1) # Row sum + losses += logsumexp + losses.masked_fill_(labels == -100, 0) # Don't forget to mask padding out! pass ctx.save_for_backward(logits, logsumexp, labels)