layernorm

unsloth/kernels/rms_layernorm.py

@@ -123,7 +123,6 @@ def _gemma_rms_layernorm_forward(
     tl.store(r, inv_var)
     normed = X_row * inv_var
     output = normed * (W_row + 1.0)
-    output = output.to(X_row.dtype)
 
     tl.store(Y + col_offsets, output, mask = mask)
 pass
@@ -141,7 +140,6 @@ def forward(ctx, X, W, eps, gemma = False):
         Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0")
         r = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
 
-        torch.cuda.synchronize()
         fx = _gemma_rms_layernorm_forward if gemma else _rms_layernorm_forward
         fx[(n_rows,)](
             Y, Y.stride(0),
@@ -152,7 +150,6 @@ def forward(ctx, X, W, eps, gemma = False):
             BLOCK_SIZE = BLOCK_SIZE,
             num_warps  = num_warps,
         )
-        torch.cuda.synchronize()
         ctx.eps = eps
         ctx.BLOCK_SIZE = BLOCK_SIZE
         ctx.num_warps  = num_warps
@@ -170,7 +167,6 @@ def backward(ctx, dY):
         n_rows, n_cols = dY.shape
         dW = X
 
-        torch.cuda.synchronize()
         _rms_layernorm_backward[(n_rows,)](
             dY, dY.stride(0),
             X,  X .stride(0),
@@ -182,7 +178,6 @@ def backward(ctx, dY):
             BLOCK_SIZE = ctx.BLOCK_SIZE,
             num_warps  = ctx.num_warps,
         )
-        torch.cuda.synchronize()
         dX = dY.view(*shape)
         return dX, None, None, None
     pass

unsloth/models/_utils.py

@@ -350,7 +350,7 @@ def is_big_gpu(index):
     "epilogue_fusion"   : True,
     "max_autotune"      : True,
     "shape_padding"     : True,
-    "trace.enabled"     : False, # Output Triton kernel outputs!
+    "trace.enabled"     : True, # Output Triton kernel outputs!
     "triton.cudagraphs" : False,
 }
 # =============================================

unsloth/models/gemma2.py

@@ -18,7 +18,6 @@
     GemmaFixedRotaryEmbedding,
     GemmaFixedLinearScalingRotaryEmbedding,
     fast_geglu_inference,
-    fast_rms_layernorm,
 )
 try:
     from transformers.models.gemma2.modeling_gemma2 import (
@@ -205,7 +204,7 @@ def Gemma2DecoderLayer_fast_forward(
         hidden_states += residual
     else:
         residual = hidden_states
-        hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states, gemma = True)
+        hidden_states = fast_rms_layernorm_gemma2_compiled(self.input_layernorm, hidden_states, gemma = True)
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states=hidden_states,
             causal_mask=causal_mask,
@@ -216,14 +215,14 @@ def Gemma2DecoderLayer_fast_forward(
             use_cache=use_cache,
             padding_mask=padding_mask,
         )
-        hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states, gemma = True)
+        hidden_states = fast_rms_layernorm_gemma2_compiled(self.post_attention_layernorm, hidden_states, gemma = True)
         hidden_states = residual + hidden_states
 
         # Fully Connected
         residual = hidden_states
-        hidden_states = fast_rms_layernorm(self. pre_feedforward_layernorm, hidden_states, gemma = True)
+        hidden_states = fast_rms_layernorm_gemma2_compiled(self. pre_feedforward_layernorm, hidden_states, gemma = True)
         hidden_states = self.mlp(hidden_states)
-        hidden_states = fast_rms_layernorm(self.post_feedforward_layernorm, hidden_states, gemma = True)
+        hidden_states = fast_rms_layernorm_gemma2_compiled(self.post_feedforward_layernorm, hidden_states, gemma = True)
         hidden_states = residual + hidden_states
     pass