From 8259b16a5d9b4916e78a146cf970cabec9563591 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 26 Feb 2024 02:49:02 +1100 Subject: [PATCH] gemma --- unsloth/kernels/cross_entropy_loss.py | 205 +++++++++++++++----------- unsloth/models/gemma.py | 31 ++-- 2 files changed, 130 insertions(+), 106 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index f17a560200..a139fce49d 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,83 +36,102 @@ 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 @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 @@ -120,10 +139,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,25 +159,28 @@ 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) -pass + # 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 MAX_FUSED_SIZE = 65536 # 2**16 @@ -166,45 +188,45 @@ 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) + 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_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) 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") + logsumexp = torch.empty((n_rows, n_chunks), dtype = torch.float32, device = "cuda") - _large_cross_entropy_forward[(n_rows, n_splits,)]( + _chunked_cross_entropy_forward[(n_rows, n_chunks,)]( logits, logits.stride(0), losses, logsumexp, labels, n_rows, - n_cols, + VOCAB_SIZE = vocab_size, 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) @@ -214,16 +236,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 - _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 = n_cols, + BLOCK_SIZE = BLOCK_SIZE, num_warps = 8, ) return logits, None, None, diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 039e3e33af..d682e914cd 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -203,15 +203,14 @@ pass # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L590 def GemmaDecoderLayer_fast_forward( self, - hidden_states: torch.Tensor, - causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - output_attentions: Optional[bool] = False, - use_cache: Optional[bool] = False, - padding_mask: Optional[torch.LongTensor] = None, - *args, **kwargs, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, ): if False:#past_key_value is not None: do_prefill = not hasattr(self.self_attn, "paged_attention") @@ -239,13 +238,13 @@ def GemmaDecoderLayer_fast_forward( # hidden_states = self.input_layernorm(hidden_states) hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, - causal_mask=causal_mask, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, - padding_mask=padding_mask, + cache_position=cache_position, + **kwargs, ) hidden_states = residual + hidden_states @@ -541,12 +540,12 @@ class FastGemmaModel(FastLlamaModel): @staticmethod def pre_patch(): - GemmaAttention .forward = LlamaAttention_fast_forward - GemmaSdpaAttention .forward = LlamaAttention_fast_forward - GemmaFlashAttention2.forward = LlamaAttention_fast_forward + GemmaAttention .forward = GemmaAttention_fast_forward + GemmaSdpaAttention .forward = GemmaAttention_fast_forward + GemmaFlashAttention2.forward = GemmaAttention_fast_forward GemmaDecoderLayer .forward = GemmaDecoderLayer_fast_forward - GemmaModel .forward = LlamaModel_fast_forward - GemmaForCausalLM .forward = LlamaForCausalLM_fast_forward + GemmaModel .forward = GemmaModel_fast_forward + GemmaForCausalLM .forward = GemmaForCausalLM_fast_forward PeftModelForCausalLM.forward = PeftModelForCausalLM_fast_forward # Solves https://github.com/unslothai/unsloth/issues/168 # Static KV Cache was introduced in 4.38.0, causing training to be much slower.