From 48736a075563e1f00a0da87fd3bcf1a661dc5ff3 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 26 Feb 2024 03:35:45 +1100 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 112 +++++++++++++++----------- 1 file changed, 63 insertions(+), 49 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index d9b9bb3a1e..7a34e09bc3 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -20,13 +20,13 @@ from transformers.models.llama.modeling_llama import logger @triton.jit -def _small_cross_entropy_forward( +def _cross_entropy_forward( logits_ptr, logits_row_stride, loss_ptr, - lse_ptr, + logsumexp_ptr, labels_ptr, - n_cols, - BLOCK_SIZE: tl.constexpr, + VOCAB_SIZE : tl.constexpr, + BLOCK_SIZE : tl.constexpr, ): """ Cross Entropy Loss = 1/n sum [ -yi log(Pi) ] @@ -36,29 +36,38 @@ def _small_cross_entropy_forward( = y * (log[sum(exp(x))] - x) If y == 0: CE_i = 0 If y == 1: CE_i = logsumexp - x + + logsumexp is also stable + Take y = log[sum(exp(x))] + exp(y) = sum(exp(x)) + exp(y) = sum(exp(x - c)*exp(c)) Since e^(x-c)*e^c = e^x + exp(y) = exp(c)*sum(exp(x - c)) + y = log(exp(c)*sum(exp(x - c))) + y = c + log[sum(exp(x - c))] + This means we can set c = max(x) to make sure + exp(x - c) always is exp(x - max(x)). + 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 - loss_ptr += row_idx - lse_ptr += row_idx - labels_ptr += row_idx + logits_ptr += row_idx * logits_row_stride.to(tl.int64) + loss_ptr += row_idx + logsumexp_ptr += row_idx + labels_ptr += row_idx col_offsets = tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols + mask = col_offsets < VOCAB_SIZE - # TODO: Fixup int32 locations to 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)) if label_idx != -100: - logits_label = tl.load(logits_ptr + label_idx).to(tl.float32) - loss = lse - logits_label + x = tl.load(logits_ptr + label_idx).to(tl.float32) + loss = logsumexp - x else: loss = 0.0 + tl.store(logsumexp_ptr, logsumexp) tl.store(loss_ptr, loss) pass @@ -120,10 +129,10 @@ pass def _cross_entropy_backward( logits_ptr, logits_row_stride, dloss_ptr, dloss_row_stride, - lse_ptr, + logsumexp_ptr, labels_ptr, - n_cols, - BLOCK_SIZE: tl.constexpr, + VOCAB_SIZE : tl.constexpr, + BLOCK_SIZE : tl.constexpr, ): """ CE_i = -y log(P) = y * (log[sum(exp(x))] - x) @@ -140,24 +149,27 @@ def _cross_entropy_backward( If y == 1 and x == label: dC/dlabel = exp[x - logsumexp] - 1 If y == 1 and x != label: dC/dx = exp[x - logsumexp] """ - row_idx = tl.program_id(0) - col_idx = tl.program_id(1) + row_idx = tl.program_id(0) + block_idx = tl.program_id(1) + logits_ptr += row_idx * logits_row_stride.to(tl.int64) dloss_ptr += row_idx * dloss_row_stride - col_offsets = col_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols + col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < VOCAB_SIZE label_idx = tl.load(labels_ptr + row_idx).to(tl.int32) - if label_idx != -100: - dloss = tl.load(dloss_ptr) - else: - dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32) - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) + x = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32) + logsumexp = tl.load(logsumexp_ptr + row_idx) + y = tl.exp(x - logsumexp) + y = tl.where( + col_offsets == label_idx, + y - 1.0, # exp(x - logsumexp) - 1 + y, # exp(x - logsumexp) + ) - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(logits_ptr + col_offsets, dloss * probs, mask = mask) + # If y == 0: dC/dx = 0 ==> we already masked it to be = 0, so dloss = 0. + dloss = tl.load(dloss_ptr) if label_idx != -100 else 0.0 + tl.store(logits_ptr + col_offsets, dloss * y, mask = mask) pass @@ -166,38 +178,38 @@ MAX_FUSED_SIZE = 65536 # 2**16 class Fast_CrossEntropyLoss(torch.autograd.Function): @staticmethod def forward(ctx, logits, labels): - n_rows, n_cols = logits.shape + n_rows, vocab_size = logits.shape div, mod = divmod(n_cols, MAX_FUSED_SIZE) - n_splits = div + (mod != 0) + n_chunks = div + (mod != 0) - if n_splits == 1: + if n_chunks == 1: # For small vocabs <= 65336 like Llama, Mistral - BLOCK_SIZE, num_warps = calculate_settings(n_cols) - losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda") + 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") - _small_cross_entropy_forward[(n_rows,)]( + _cross_entropy_forward[(n_rows,)]( logits, logits.stride(0), losses, logsumexp, labels, - n_cols, + VOCAB_SIZE = vocab_size, BLOCK_SIZE = BLOCK_SIZE, num_warps = num_warps, ) else: # For large vocabs > 65336 like Gemma 256K - losses = torch.empty((n_splits, n_rows), dtype = torch.float32, device = "cuda") - logsumexp = torch.empty((n_splits, n_rows), dtype = torch.float32, device = "cuda") + losses = torch.empty((n_chunks, n_rows), dtype = torch.float32, device = "cuda") + logsumexp = torch.empty((n_chunks, n_rows), dtype = torch.float32, device = "cuda") - _large_cross_entropy_forward[(n_rows, n_splits,)]( + _large_cross_entropy_forward[(n_rows, n_chunks,)]( logits, logits.stride(0), losses, logsumexp, labels, n_rows, - n_cols, + vocab_size, BLOCK_SIZE = MAX_FUSED_SIZE, num_warps = 32, ) @@ -214,17 +226,19 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): @staticmethod def backward(ctx, dlosses): logits, logsumexp, labels = ctx.saved_tensors - n_rows, n_cols = logits.shape - grid = lambda meta: (n_rows, triton.cdiv(n_cols, meta["BLOCK_SIZE"])) + n_rows, vocab_size = logits.shape - print(logits.stride(), dlosses.stride(), dlosses.shape, dlosses) - _cross_entropy_backward[grid]( + BLOCK_SIZE = 4096 + div, mod = divmod(vocab_size, BLOCK_SIZE) + n_blocks = div + (mod != 0) + + _cross_entropy_backward[(n_rows, n_blocks,)]( logits, logits.stride(0), dlosses, dlosses.stride(0), logsumexp, labels, - n_cols, - BLOCK_SIZE = 4096, + VOCAB_SIZE = vocab_size, + BLOCK_SIZE = BLOCK_SIZE, num_warps = 8, ) return logits, None, None,