diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index 8ad8bf20be..561d974042 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -18,7 +18,6 @@ import torch from .utils import calculate_settings -@triton.heuristics({"ADD_ONE": lambda args: args["ADD_ONE"],}) @triton.jit def _rms_layernorm_forward( Y, Y_row_stride, @@ -26,8 +25,7 @@ def _rms_layernorm_forward( W, W_row_stride, r, r_row_stride, n_cols, eps, - ADD_ONE: tl.constexpr, - BLOCK_SIZE : tl.constexpr, + BLOCK_SIZE : tl.constexpr ): """ Fast RMS Layernorm kernel @@ -50,20 +48,12 @@ def _rms_layernorm_forward( tl.store(r, inv_var) normed = X_row * inv_var normed = normed.to(W_row.dtype) # Exact copy from HF - - # For Gemma - cannot do += 1 since float16 - maybe use FMADD - if not ADD_ONE: - output = normed * W_row - else: - # Error analysis shows we need to do +1 in float32 then downcast to float16 - output = normed * (W_row.to(tl.float32) + 1.0).to(W_row.dtype) - pass - + output = normed * W_row tl.store(Y + col_offsets, output, mask = mask) pass -@triton.heuristics({"ADD_ONE": lambda args: args["ADD_ONE"],}) +@triton.heuristics({"GEMMA": lambda args: args["GEMMA"],}) @triton.jit def _rms_layernorm_backward( dY, dY_row_stride, @@ -72,7 +62,7 @@ def _rms_layernorm_backward( r, r_row_stride, dW, dW_row_stride, n_cols, eps, - ADD_ONE: tl.constexpr, + GEMMA : tl.constexpr, BLOCK_SIZE : tl.constexpr, ): """ @@ -96,12 +86,8 @@ def _rms_layernorm_backward( inv_var = tl.load(r).to(tl.float32) normed = X_row * inv_var - # For Gemma - cannot do += 1 since float16 - maybe use FMADD - if not ADD_ONE: - dY_W = dY_row * W_row - else: - dY_W = dY_row * (W_row + 1.0) - pass + if GEMMA: dY_W = dY_row * (W_row + 1.0) + else: dY_W = dY_row * W_row rowsum_dY_normed = tl.sum(dY_W * normed, axis = 0) output = inv_var/n_cols * (n_cols*dY_W - normed*rowsum_dY_normed) @@ -109,9 +95,42 @@ def _rms_layernorm_backward( pass +@triton.jit +def _gemma_rms_layernorm_forward( + Y, Y_row_stride, + X, X_row_stride, + W, W_row_stride, + r, r_row_stride, + n_cols, eps, + BLOCK_SIZE : tl.constexpr, +): + # Copies https://github.com/google-deepmind/gemma/blob/main/gemma/layers.py#L31 + # and https://github.com/keras-team/keras-nlp/blob/v0.8.2/keras_nlp/models/gemma/rms_normalization.py#L33 + # exactly. Essentially all in float32! + row_idx = tl.program_id(0) + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < n_cols + + Y += row_idx * Y_row_stride + X += row_idx * X_row_stride + r += row_idx * r_row_stride + + X_row = tl.load(X + col_offsets, mask = mask, other = 0).to(tl.float32) + W_row = tl.load(W + col_offsets, mask = mask, other = 0).to(tl.float32) + + row_var = tl.sum(X_row * X_row, axis = 0) / n_cols + inv_var = tl.math.rsqrt(row_var + eps) + tl.store(r, inv_var) + normed = X_row * inv_var + output = normed * (W_row + 1.0) + + tl.store(Y + col_offsets, output, mask = mask) +pass + + class Fast_RMS_Layernorm(torch.autograd.Function): @staticmethod - def forward(ctx, X, W, eps, add_one = False): + def forward(ctx, X, W, eps, gemma = False): shape = X.shape dim = shape[-1] X = X.view(-1, dim) @@ -121,20 +140,20 @@ class Fast_RMS_Layernorm(torch.autograd.Function): Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda") r = torch.empty(n_rows, dtype = torch.float32, device = "cuda") - _rms_layernorm_forward[(n_rows,)]( + fx = _gemma_rms_layernorm_forward if gemma else _rms_layernorm_forward + fx[(n_rows,)]( Y, Y.stride(0), X, X.stride(0), W, W.stride(0), r, r.stride(0), n_cols, eps, - ADD_ONE = add_one, BLOCK_SIZE = BLOCK_SIZE, num_warps = num_warps, ) ctx.eps = eps ctx.BLOCK_SIZE = BLOCK_SIZE ctx.num_warps = num_warps - ctx.ADD_ONE = add_one + ctx.GEMMA = gemma ctx.save_for_backward(X, W, r) return Y.view(*shape) pass @@ -155,7 +174,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function): r, r .stride(0), dW, dW.stride(0), n_cols, ctx.eps, - ADD_ONE = ctx.ADD_ONE, + GEMMA = ctx.GEMMA, BLOCK_SIZE = ctx.BLOCK_SIZE, num_warps = ctx.num_warps, ) @@ -165,9 +184,9 @@ class Fast_RMS_Layernorm(torch.autograd.Function): pass -def fast_rms_layernorm(layernorm, X, add_one = False): +def fast_rms_layernorm(layernorm, X, gemma = False): W = layernorm.weight eps = layernorm.variance_epsilon - out = Fast_RMS_Layernorm.apply(X, W, eps, add_one) + out = Fast_RMS_Layernorm.apply(X, W, eps, gemma) return out pass diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 12001d7046..aee7ea38e3 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -58,16 +58,15 @@ def fast_geglu_inference(self, X): pass -def fast_rms_layernorm_inference_add_one(self, X, out_weight = None): - old_dtype = X.dtype +def fast_rms_layernorm_inference_gemma(self, X, out_weight): XX = X.to(torch.float32) variance = XX.square().mean(-1, keepdim = True) variance += self.variance_epsilon XX *= variance.rsqrt_() - X = XX.to(old_dtype) # Must preserve due to residual - out_weight = torch.add(self.weight, 1.0, out = out_weight) - X *= out_weight - return X + out_weight[:] = self.weight + out_weight += 1.0 + XX *= out_weight + return XX.to(X.dtype) pass @@ -86,10 +85,11 @@ def GemmaDecoderLayer_fast_forward( ): if past_key_value is not None: do_prefill = not hasattr(self.self_attn, "paged_attention") + out_weight = torch.empty(self.input_layernorm.weight.shape, dtype = torch.float32, dtype = "cuda") # Self Attention residual = hidden_states - hidden_states = fast_rms_layernorm_inference_add_one(self.input_layernorm, hidden_states) + hidden_states = fast_rms_layernorm_inference_gemma(self.input_layernorm, hidden_states, out_weight) hidden_states, present_key_value = LlamaAttention_fast_forward_inference( self.self_attn, hidden_states, @@ -101,12 +101,12 @@ def GemmaDecoderLayer_fast_forward( # Fully Connected residual = hidden_states - hidden_states = fast_rms_layernorm_inference_add_one(self.post_attention_layernorm, hidden_states) + hidden_states = fast_rms_layernorm_inference_gemma(self.post_attention_layernorm, hidden_states, out_weight) hidden_states = fast_geglu_inference(self.mlp, hidden_states) hidden_states += residual else: residual = hidden_states - hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states, add_one = True) + hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states, gemma = True) # hidden_states = self.input_layernorm(hidden_states) hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, @@ -122,7 +122,7 @@ def GemmaDecoderLayer_fast_forward( # Fully Connected residual = hidden_states - hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states, add_one = True) + hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states, gemma = True) # hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states @@ -151,16 +151,18 @@ def GemmaModel_fast_forward_inference( ): # Fix out of bounds tokenization input_ids = input_ids[:,:self.max_seq_length] - out_weight = torch.empty_like(self.layers[0].input_layernorm.weight) + out_weight = torch.empty_like(self.layers[0].input_layernorm.weight, dtype = torch.float32, device = "cuda") hidden_states = self.embed_tokens(input_ids) - hidden_states *= math_sqrt(self.config.hidden_size) + # 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32 + # 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32 + inputs_embeds *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype) next_decoder_cache = [] for idx, decoder_layer in enumerate(self.layers): # Self Attention residual = hidden_states - hidden_states = fast_rms_layernorm_inference_add_one(decoder_layer.input_layernorm, hidden_states, out_weight) + hidden_states = fast_rms_layernorm_inference_gemma(decoder_layer.input_layernorm, hidden_states, out_weight) hidden_states, present_key_value = LlamaAttention_fast_forward_inference( decoder_layer.self_attn, hidden_states, @@ -171,13 +173,13 @@ def GemmaModel_fast_forward_inference( # Fully Connected residual = hidden_states - hidden_states = fast_rms_layernorm_inference_add_one(decoder_layer.post_attention_layernorm, hidden_states, out_weight) + hidden_states = fast_rms_layernorm_inference_gemma(decoder_layer.post_attention_layernorm, hidden_states, out_weight) hidden_states = fast_geglu_inference(decoder_layer.mlp, hidden_states) hidden_states += residual next_decoder_cache.append(present_key_value) pass - hidden_states = fast_rms_layernorm_inference_add_one(self.norm, hidden_states, out_weight) + hidden_states = fast_rms_layernorm_inference_gemma(self.norm, hidden_states, out_weight) return BaseModelOutputWithPast( last_hidden_state = hidden_states, diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 05cd1c7426..d7743216fb 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -519,7 +519,11 @@ def LlamaModel_fast_forward( elif inputs_requires_grad: inputs_embeds.requires_grad_(False) pass - inputs_embeds *= math_sqrt(self.config.hidden_size) + # Match Gemma exactly by casting to bfloat16 / float16 + # inputs_embeds *= math_sqrt(self.config.hidden_size) + # Ie 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32 + # & 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32 + inputs_embeds *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype) if inputs_requires_grad: inputs_embeds.requires_grad_(True) pass @@ -621,7 +625,7 @@ def LlamaModel_fast_forward( all_self_attns += (layer_outputs[1],) pass - hidden_states = fast_rms_layernorm(self.norm, hidden_states, add_one = IS_GEMMA) + hidden_states = fast_rms_layernorm(self.norm, hidden_states, gemma = IS_GEMMA) # add hidden states from the last decoder layer if output_hidden_states: